[AMD] added slicing ttg.async_copy_global_to_local

[ravil-mobile]

> TRITON_HIP_USE_BLOCK_PINGPONG=1 TRITON_HIP_USE_ASYNC_COPY=1 pytest -s -v op_tests/triton_tests/test_gemm_afp4wfp4.py

...

================================================================================================= warnings summary =================================================================================================
op_tests/triton_tests/test_gemm_afp4wfp4.py::test_gemm_afp4_wfp4[dtype0-1024-1024-1024]
  /home/rdorozhi/work/aiter/op_tests/triton_tests/test_gemm_afp4wfp4.py:92: UserWarning: Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at /var/lib/jenkins/pytorch/aten/src/ATen/Context.cpp:328.)
    return torch.mm(x_f32, w_f32).to(dtype)

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
===================================================================================== 74 passed, 1 warning in 65.59s (0:01:05) =====================================================================================

[ravil-mobile]

@jungpark-mlir, @raikonenfnu. Thanks for your comments! I addressed all of them.

[raikonenfnu]

I have a silly Q for educational purposes, looks like we are only updating the threadsPerWarp here, does it not make more sense to updated the sizePerThread in the newEncoding?

[ravil-mobile]

No, we want to preserve the number of elements per thread (i.e., each holds 16 consecutive elements of a tensor). We just want to change the layout of threads; thus, to change which part of a tensor is held by CTA.

Each CTA holds 64x128 tile of a tensor in the following example

#blocked = #ttg.blocked<{sizePerThread = [1, 16], threadsPerWarp = [8, 8], warpsPerCTA = [8, 1], order = [1, 0]}>

But if we change threadsPerWarp to [32,2] then a CTA holds 256x32 tile of a tensor.

#blocked3 = #ttg.blocked<{sizePerThread = [1, 16], threadsPerWarp = [32, 2], warpsPerCTA = [8, 1], order = [1, 0]}>

In both cases, a thread holds 16 consecutive elements which determines the width of load instructions

[raikonenfnu]

Thanks for the answer, that's really interesting! A follow up question, if we change the which part of tensor the CTA will hold, wouldn't we need to have an extra global read (or potentially fused but wider global read) to get that part of tensor?

[ravil-mobile]

ConverLayout results in some machine ops if it cannot be fully optimized (propagated till the top of a function). If I remember correctly, the layout change is going to happen in LDS. In out case, it is going to be optimized

[raikonenfnu]

Nice! didn't realize this was pre-layout propagation phase, but this makes a lot of sense, thanks :)

[raikonenfnu]

Another silly Q, I know amdgpu.extract_slice src and dst layout need to match, is this the main purpose of ttg.convert_layout to set up layout for stuff like extract_slice where the layout is for sure to change? On top of that, does that mean ttg.convert_layout for slices will most likely not have a matching layout to it's shape?

[ravil-mobile]

dst layout is determined by the source layout. The problem comes from the following. Let's assume we want to slice 256x128 tensor into 4 pieces of 256x32 tiles. Let's also assume that the original layout is

# orig-layout
#blocked = #ttg.blocked<{sizePerThread = [1, 16], threadsPerWarp = [8, 8], warpsPerCTA = [8, 1], order = [1, 0]}>

which holds 64x128 tiles per CTA. ExtractSliceOp is the CTA level op - i.e., it can slice a tensor if a new size is proportional to the CTA tile - i.e.,64x128 in our case. Therefore, we cannot apply ExtractSliceOp to our tensor with orig-layout. Thus, we change the source layout to a new one - i.e.,

# new-layout
#blocked3 = #ttg.blocked<{sizePerThread = [1, 16], threadsPerWarp = [32, 2], warpsPerCTA = [8, 1], order = [1, 0]}>

which has its CTA-level slice equaled to 256x32. Now we can slice 256x128 tensor into 4 pieces. (note: it was determined by dot-slicing)

[raikonenfnu]

NIT: IIRC, RankedTensorType/Type/Value will default to null even if you don't explicitly set nullptr

[ravil-mobile]

You are right. I just like to be explicit in the code about initialization of all local variables (old habit).