[Issue]: IR explosion in ConvertTritonAMDGPUToLLVM

[pabloantoniom]

Problem Description

In the Gluon program attached in this issue, if passed to --convert-triton-amdgpu-to-llvm, it will take several minutes to compile and will generate an output of ~14GB. This seems to be due an IR explosion generated by the --convert-triton-amdgpu-to-llvm pass. Note that I set arch to gfx942, but this seems to be arch-independent.

explode.py

Operating System

Ubuntu 22.04.5 LTS (Jammy Jellyfish)

CPU

Intel(R) Xeon(R) Platinum 8480C

GPU

AMD Instinct MI300X

ROCm Version

7.0

ROCm Component

No response

Steps to Reproduce

$ time python3 explode.py | triton-opt --convert-triton-amdgpu-to-llvm=arch=gfx942 | wc -c
14204339549

real    2m11.133s
user    2m8.865s
sys     0m9.593s

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

[ravil-mobile]

Hi @pabloantoniom

I had a look at your script. The first thing that strike my eyes was the blocking configuration

ARCH = "gfx942"                                                                                                                                                                                                     
WARP_SIZE = 64                                                                                                                                                                                                      
NUM_WARPS = 2                                                                                                                                                                                                       
SHAPE = (1024, 1024)                                                                                                                                                                                                
LAYOUT = gl.BlockedLayout(                                                                                                                                                                                          
    size_per_thread=[64, 16],                                                                                                                                                                                       
    threads_per_warp=[1, 64],                                                                                                                                                                                       
    warps_per_cta=[2, 1],                                                                                                                                                                                           
    order=[1, 0],                                                                                                                                                                                                   
)

According to this, each thread needs 64 x 16 x 8 == 8192 per tile where each element if fp32. That is a lot! You use SHAPE as the tile size. Are you sure about it?

The following config worked for me without an issue

ARCH = "gfx942"                                                                                                                                                                                                     
WARP_SIZE = 64                                                                                                                                                                                                      
NUM_WARPS = 2                                                                                                                                                                                                       
SHAPE = (32, 64)                                                                                                                                                                                                    
LAYOUT = gl.BlockedLayout(                                                                                                                                                                                          
    size_per_thread=[1, 16],                                                                                                                                                                                        
    threads_per_warp=[16, 4],                                                                                                                                                                                       
    warps_per_cta=[NUM_WARPS, 1],                                                                                                                                                                                   
    order=[1, 0],                                                                                                                                                                                                   
)                                                                                                                                                                                                                   

So, if you need to process 1024x1024 matrices and you are forced to use 2 warps per work_group then you tile the matrices by 32x64 tiles. You will have 512 work_groups where each one processes its own chunk of data. Triton delegates grid level parallelization to the user (similar to CUDA/HIP/SYCL).

Could you explain that specific choice of blocking and the layout in your example?

[ravil-mobile]

@zhanglx13 I talked to @pabloantoniom. He provided an extreme example on purpose. However, the main complain is the slowdown resulting from LLVM passes (not MLIR). I am waiting from @pabloantoniom to provide a better reproducible example.

[pabloantoniom]

Let me clarify that the slowdown does not come from Triton itself. It comes from LLVM actually.

I have analyzed the original IR that triggered the issue (which is slightly different compared to the reproducer I shared):

===-------------------------------------------------------------------------===
                         ... Execution time report ...
===-------------------------------------------------------------------------===
  Total Execution Time: 297.3944 seconds

  ----Wall Time----  ----Name----
    ...
    3.3226 (  1.1%)  ConvertTritonAMDGPUToLLVM
    1.0319 (  0.3%)  TritonAMDGPUConvertWarpSpecializeToLLVM
    3.3991 (  1.1%)  CanonicalizerPass
    ...
  283.4979 ( 95.3%)  TritonToHsacoPass
    ...
    0.0038 (  0.0%)  Rest
  297.3944 (100.0%)  Total

The summary is that 95% of compile time is in TritonToHsacoPass, not in ConvertTritonAMDGPUToLLVM. Within TritonToHsacoPass, --time-passes shows that 77 % is in the AMDGPU post‑RA machine scheduler and another 8 % in the pre‑RA machine scheduler. Register allocation is only ~9 %. So it's not that Triton loweing is slow, is that the IR it generates it's very inefficient.

This happens only in corner cases where the amount of elements per thread is big, like you noticed. The kernel isn't hand‑written; this config is generated by our autotuner. But I just hoped that Triton could do better in this kind of cases?

ElementwiseOpConversionBase::matchAndRewrite always unpacks !llvm.struct operands into scalar SSA values via extractvalue, applies the per‑element lowering, and re‑packs with insertvalue. I think this is the real trigger of this IR explosion. So maybe this could be an area of improvement?

[ravil-mobile]

Hi @pabloantoniom,

Could you attach a Triton kernel which consumes 95.3% in TritonToHsaco pass?

[pabloantoniom]

I was unable to reproduce this from a Triton kernel. I don't remember the details, but somehow the semantic.py (or some other file) decided to use a different config (different threads per warp, different warps, I'm not sure) that was "less extreme" and didn't trigger the issue in such an obvious way. That's why I gave a Gluon reproducer.

That said, I have a new Triton kernel that is able to reproduce this issue, but it only triggers on RDNA3/RDNA4. The reason is that, even though Triton frontend somehow avoids the original issue I shared, on Navi it seems like --convert-triton-amdgpu-to-llvm it's broken, so it will still trigger it. I think this is a better reproducer because it represents more faithfully the issues that we are having in our project.

triton_explosion.py

$ python3 triton_explosion.py > triton_explosion.mlir && time triton-opt triton_explosion.mlir --convert-scf-to-cf --convert-index-to-llvm --allocate-amdgpu-shared-memory --tritongpu-global-scratch-memory-allocation --convert-triton-amdgpu-to-llvm='arch=gfx1100' &> exploded.mlir

real    1m48.428s
user    1m26.107s
sys     0m21.562s
$ ls -lh exploded.mlir
-rw-r--r-- 1 pamartin pamartin 6.6G May 15 14:45 exploded.mlir
$ wc -l exploded.mlir
1937342 exploded.mlir

The problem seems to be specific to the Navi lowering. Using fp8 makes it worse, but it's not the root cause: switching IN_DTYPE from fp8 to fp16 on gfx1100 (no tt.fp_to_fp ops at all) still produces 6.6 GB (takes ~90 s).

For reference, running the exact same kernel through the same pipeline with arch=gfx942 produces 37 MB / 2.5 s at 256×256×64 fp8 with 2 warps.

Also sharing some experiments I made on gfx1100 with different configs (the script has a default value but allows you to override with env vars):

| Config      | DType | Warps | IR Size | Pass time |
|-------------|-------|-------|---------|-----------|
| 256×256×64  | fp8   | 2     | 6.6 GB  | ~90s       |
| 128×256×64  | fp8   | 2     | 1.4 GB  | ~18s       |
| 256×128×64  | fp8   | 2     | 1.4 GB  | ~14s       |
| 256×128×64  | fp8   | 4     | 440 MB  | ~6s        |
| 128×128×64  | fp16  | 4     | 92 MB   | ~2s        |
| 256x256x64  | fp16  | 2     | 6.6 GB  | ~90s       |

[ravil-mobile]

Hi @pabloantoniom,

I examine tile=256x256x64, dtype=fp16, warps=2 config (a similar thing happens with other configs) and compared the generated code on gfx1100, gfx1250 and gfx942.

So, in this particular case, the work per thread twice as large on gfx1* in comparison to gfx942 because the former has 64 threads per wave. The compilation slow-down correlates to the number of elements held by a thread. For example, 256x256x64 is going to be split into 128x256x64 between 2 warps. Each thread needs to hold 256 elements of A, 512 elements of B, and 1024 elements of C. The same config on gfx942 is going to result in each thread holding 128 elements of A, 256 elements of B, and 512 elements of C. You need to use 4 warps to match gfx942 regarding a work per thread. You will notice a similar effect if you halve tiles along M and N dimensions by a factor of 2 (it is even shown in your table).

---

GFX1100

time ARCH=gfx1100 IN_DTYPE=fp16 BLOCK_M=256 BLOCK_N=256 BLOCK_K=64 NUM_WARPS=2 python3  ./triton_explosion.py
32.93s user 0.30s system 99% cpu 33.239 total

time ARCH=gfx1100 IN_DTYPE=fp16 BLOCK_M=256 BLOCK_N=256 BLOCK_K=64 NUM_WARPS=4 python3  ./triton_explosion.py
11.44s user 0.12s system 99% cpu 11.567 total

GFX1250

time ARCH=gfx1250 IN_DTYPE=fp16 BLOCK_M=256 BLOCK_N=256 BLOCK_K=64 NUM_WARPS=2 python3  ./triton_explosion.py
38.35s user 0.42s system 99% cpu 38.771 total

time ARCH=gfx1250 IN_DTYPE=fp16 BLOCK_M=256 BLOCK_N=256 BLOCK_K=64 NUM_WARPS=4 python3  ./triton_explosion.py
8.58s user 0.18s system 99% cpu 8.753 total

---

Once, we enter the LLVM land we are dealing with SIMT level; elements are held in llvm.struct. Our lowering from TTGIR to LLVM resembles some sort code emission during post-order traversal of an AST - i.e., the code has redundant LLVM operations; mainly, packing/unpacking llvm.structs. We fire some clean-up passes to remove most of them. It feels like the codegen length growth quadratically (or even worse) relevant to llvm.struct sizes.

There are 2 points that I need to mention:

1. many configs that you shared are going to result in register spilling; thus, it just waste of time to compile and test them. A good solution for MIGraphX (or any graph compiler) would be to add or extend pruning of an auto-tuning config to avoid such scenarios where register spilling is obvious. A rough estimate could be a number of registers per thread required to hold a tile A, B and C for WMMAs/MFMAs. A more advanced approach could take num_stages into account because it dictates a structure of the pipelined loop

2. We need to take a closer look at smaller data types (mxfp4, mxfp8) because we usually tend to deal with larger tiles than for fp16. Despite targeting the same VGPR count as for FP16 kernels, sizes of llvm.struct gets larger because they are holding elements (not registers).

3. We may not notice the compilation slowdown for gfx1250 because of FFM/AM; we could think that it was a pure simulation slowness