[SWP] attempt to remove a workaround for a triton llvm codegen bug

[manman-ren]

Triton LLVM codegen has a bug where local_loads from #shared to #mma layout can lead to invalid code if the loaded shape is smaller than the mma tile. Remove the workaround.

See #3561.

Verified that with test case: https://pastebin.com/xxP3cFmy (test.mlir), running
triton-opt test.mlir -tritongpu-pipeline=num-stages=3 --convert-scf-to-cf --allocate-shared-memory --convert-triton-gpu-to-llvm
has no issue.

Unit test case added in #4798 also shows no issue.

[manman-ren]

@jlebar I am not sure if I am checking the failure correctly, I tried running
triton-opt test.mlir -tritongpu-pipeline=num-stages=3 --convert-scf-to-cf --allocate-shared-memory --convert-triton-gpu-to-llvm
but didn't hit any issue. From #3549, sounds like we should hit an assertion.
Thanks!

CC @pawelszczerbuk

[jlebar]

From #3549, sounds like we should hit an assertion.

I believe that assertion was removed by @Jokeren when he fixed the relevant codepath.

[Jokeren]

One of our recent PRs may have fixed this problem. Do you have a test case BTW? If so, happy to dig deeper

[manman-ren]

I am using dont_pipeline_128x1 from test/TritonGPU/loop-pipeline.mlir (https://pastebin.com/xxP3cFmy), if we remove the HACK from the compiler, we will start pipelining the load in the test case, which should trigger the issue of invalid code for local_loads from #shared to #mma layout if the loaded shape is smaller than the mma tile.

[Jokeren]

I am using dont_pipeline_128x1 from test/TritonGPU/loop-pipeline.mlir (https://pastebin.com/xxP3cFmy), if we remove the HACK from the compiler, we will start pipelining the load in the test case, which should trigger the issue of invalid code for local_loads from #shared to #mma layout if the loaded shape is smaller than the mma tile.

To confirm, are you saying that the bug is triggered at the following IR?

%161 = triton_gpu.convert_layout %151 : tensor<128x1xi32, #blocked> -> tensor<128x1xi32, #mma>

[Jokeren]

Honestly I'm not sure if constructing an MMA layout greater than the tensor size makes sense?

In theory we could repeat values in registers and threads. Would like to get your thoughts

[manman-ren]

I am using dont_pipeline_128x1 from test/TritonGPU/loop-pipeline.mlir (https://pastebin.com/xxP3cFmy), if we remove the HACK from the compiler, we will start pipelining the load in the test case, which should trigger the issue of invalid code for local_loads from #shared to #mma layout if the loaded shape is smaller than the mma tile.

To confirm, are you saying that the bug is triggered at the following IR?

%161 = triton_gpu.convert_layout %151 : tensor<128x1xi32, #blocked> -> tensor<128x1xi32, #mma>

https://pastebin.com/eN5HP8XW has the mlir after SWP, if we pipeline the load.

Specifically, this op should trigger the issue
%14 = triton_gpu.local_load %arg5 : !tt.memdesc<128x1xi32, #shared, #triton_gpu.shared_memory, mutable> -> tensor<128x1xi32, #mma>
#mma = #triton_gpu.nvidia_mma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [4, 1], instrShape = [16, 8]}>

i.e loading a 128x1 tensor for an MMAv2 dot with tile {16,8} is bad because 1 < 8

This is how I understand it from the comments :]

[Jokeren]

#4798
@manman-ren