[Triton][Gluon] Run the inliner after scf-to-cf (#8017)

The expectation is that functions with tensor-typed arguments will be
inlined, and this is enforced in the frontend. However, if the callee
has early returns it doesn't get inlined, hitting crashes in
interprocedural dataflow analyses later (AxisInfo). @ThomasRaoux's
suggestion is just to run the inliner again after scf-to-cf.

Author: Mogball

lib/Dialect/TritonNvidiaGPU/IR/Ops.cpp

@@ -30,6 +30,7 @@
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOpInterfaces.cpp.inc"
 #include "triton/Dialect/TritonNvidiaGPU/Transforms/Utility.h"
+#include "llvm/Support/ErrorHandling.h"
 
 using namespace mlir::triton::gpu;
 
@@ -277,7 +278,7 @@ static std::string strMMADTypeKind(MMADTypeKind kind) {
   case MMADTypeKind::i8:
     return "i8";
   }
-  __builtin_unreachable();
+  llvm_unreachable("unknown mma dtype kind");
 }
 
 static std::optional<std::pair<MMADTypeKind, SmallVector<Type>>>

python/test/gluon/test_core.py

@@ -783,3 +783,31 @@ def tmem_copy_no_scales(in_ptr, out_ptr, M: ttgl.constexpr, N: ttgl.constexpr, B
 
     tmem_copy_no_scales[(1, )](input, output, M, N, BLOCK_N, swizzle, num_warps=num_warps)
     assert (output == input).all()
+
+
+@gluon.jit
+def early_return_kernel(x):
+    if x.sum(0).sum(0):
+        return x
+    x = x + x
+    return x
+
+
+def test_2d_tensor_early_return():
+    warp_size = ttgl.constexpr(THREADS_PER_WARP)
+
+    @gluon.jit
+    def kernel(N, out):
+        layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, warp_size], [1, 4], [1, 0])
+        BLOCK: ttgl.constexpr = 32
+
+        x0 = ttgl.arange(0, BLOCK, layout=ttgl.SliceLayout(1, layout))
+        x1 = ttgl.arange(0, BLOCK, layout=ttgl.SliceLayout(0, layout))
+        x = x0[:, None] * x1[None, :]
+        for i in range(N):
+            x += early_return_kernel(x)
+        ttgl.store(out, x.sum(0).sum(0))
+
+    out = torch.empty(1, dtype=torch.int32, device="cuda")
+    compiled_kernel = kernel.warmup(N=100, out=out, grid=(1, ))
+    assert compiled_kernel.asm["llir"].count("define") == 1

third_party/amd/backend/compiler.py

@@ -288,6 +288,7 @@ def make_llir(src, metadata, options):
         custom_lds_size = 0
         amd.passes.ttgpuir.add_optimize_lds_usage(pm, options.arch, custom_lds_size)
         passes.convert.add_scf_to_cf(pm)
+        passes.gluon.add_inliner(pm)
         passes.convert.add_index_to_llvmir(pm)
 
         amd.passes.ttgpuir.add_allocate_shared_memory(pm)

third_party/nvidia/backend/compiler.py

@@ -349,6 +349,7 @@ def make_llir(self, src, metadata, options, capability):
         passes.ttgpuir.add_combine_tensor_select_and_if(pm)
         passes.ttgpuir.add_allocate_warp_groups(pm)
         passes.convert.add_scf_to_cf(pm)
+        passes.gluon.add_inliner(pm)
         nvidia.passes.ttgpuir.add_allocate_shared_memory_nv(pm, capability, ptx_version)
         nvidia.passes.ttnvgpuir.add_allocate_tensor_memory(pm)
         if knobs.compilation.enable_experimental_consan: