[AMD] Support warp specialization on gfx1250 (triton-lang#8947)

Add warp specialization support for AMD gfx1250.

Author: PMylon

bin/RegisterTritonDialects.h

@@ -97,6 +97,7 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
 
   // TritonAMDGPUToLLVM passes
   mlir::triton::registerAllocateAMDGPUSharedMemory();
+  mlir::triton::registerTritonAMDGPUConvertWarpSpecializeToLLVM();
   mlir::triton::registerConvertTritonAMDGPUToLLVM();
   mlir::triton::registerConvertBuiltinFuncToLLVM();
   mlir::triton::registerConvertWarpPipeline();

python/triton/experimental/gluon/language/amd/gfx1250/mbarrier.py

@@ -53,7 +53,7 @@ def arrive(mbarrier, *, count=1, _semantic=None):
     Arrive at an mbarrier with a specified count. The operation requires a `count` attribute
     of at least 1, and decreases the pending arrival count of the mbarrier by the specific count.
     If the pending count reaches zero, the phase changes (is decremented in a wraparound manner) and the
-    pending count is reloaded with the init count value. Returns the mbarrier's phase prior to the "arrive" operation.
+    pending count is reloaded with the init count value. Returns the mbarrier's phase parity (0 for even, 1 for odd) prior to the "arrive" operation.
 
     Args:
         mbarrier (shared_memory_descriptor): Barrier to be signalled.

test/Conversion/warp_specialize_to_llvm.mlir

@@ -278,6 +278,7 @@ def gluon_to_ttgir(src, metadata, options):
         passes.gluon.add_canonicalizer(pm)
         passes.ttgpuir.add_combine_tensor_select_and_if(pm)
         amd.passes.ttgpuir.add_warp_pipeline(pm)
+        passes.ttgpuir.add_allocate_warp_groups(pm)
 
         pm.run(mod, 'gluon_to_ttgir')
         metadata["tensordesc_meta"] = mod.get_tensordesc_metadata()
@@ -308,6 +309,7 @@ def make_llir(src, metadata, options):
         ##    For now it is used as a controller for developers only.
         __HIP_FTZ = True
         amd.passes.ttgpuir.add_to_llvmir(pm, options.arch, __HIP_FTZ)
+        amd.passes.ttgpuir.add_warp_specialize_to_llvm(pm, options.arch)
         passes.common.add_canonicalizer(pm)
         passes.common.add_cse(pm)
 
@@ -371,7 +373,12 @@ def make_llir(src, metadata, options):
         fns = [fn for fn in llvm_mod.get_functions() if not fn.is_declaration()]
         # The public kernel should be kernel 0.
         fns[0].set_calling_conv(amd.CALLING_CONV_AMDGPU_KERNEL)
-        fns[0].add_fn_attr("amdgpu-flat-work-group-size", f"1,{options.num_warps*options.warp_size}")
+        # warp-specialization mutates num_warps
+        total_warps_num = options.num_warps
+        total_num_warps = src.get_int_attr("ttg.total-num-warps")
+        if total_num_warps is not None:
+            total_warps_num = total_num_warps
+        fns[0].add_fn_attr("amdgpu-flat-work-group-size", f"1,{total_warps_num*options.warp_size}")
         if "memory-bound-attention" in options.schedule_hint.split(','):
             fns[0].add_fn_attr("amdgpu-sched-strategy", "iterative-ilp")
         fns[0].add_fn_attr("uniform-work-group-size", "true")
@@ -425,6 +432,7 @@ def make_llir(src, metadata, options):
             amd.add_scalarize_packed_fops_llvm_pass(fns[0])
 
         # Get some metadata
+        metadata["num_warps"] = total_warps_num
         metadata["shared"] = src.get_int_attr("ttg.shared")
         metadata["profile_scratch_size"] = src.get_int_attr("ttg.profile_scratch_memory_size") or 0
         metadata["profile_scratch_align"] = src.get_int_attr("ttg.profile_scratch_memory_alignment") or 1

third_party/amd/backend/compiler.py

@@ -893,7 +893,7 @@ def ArriveBarrierOp : TT_AMDGPU_Op<"arrive_barrier"> {
     Performs the "arrive" operation on an mbarrier object in shared memory. The operation requires a `count` attribute
     of at least 1, and decreases the pending arrival count of the mbarrier by the specific count. If the pending count reaches
     zero, the phase changes (is decremented in a wraparound manner) and the pending count is reloaded with the init count value. Returns the phase
-    of the mbarrier object prior to the "arrive" operation.
+    parity (0 for even, 1 for odd) of the mbarrier object prior to the "arrive" operation.
 
     Example:
 

third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOps.td

@@ -24,7 +24,10 @@ namespace mlir::triton {
 } // namespace mlir::triton
 
 namespace mlir::triton::AMD {
+
 std::unique_ptr<OperationPass<ModuleOp>> createConvertWarpPipelinePass();
+std::unique_ptr<OperationPass<ModuleOp>>
+createTritonAMDGPUConvertWarpSpecializeToLLVMPass(StringRef arch);
 void runScalarizePackedFOpsPass(llvm::Function &F);
 
 } // namespace mlir::triton::AMD

third_party/amd/include/TritonAMDGPUToLLVM/Passes.h

@@ -84,7 +84,24 @@ def ConvertWarpPipeline : Pass<"convert-warp-pipeline", "mlir::ModuleOp"> {
                              "mlir::gpu::GPUDialect",
                              "mlir::ROCDL::ROCDLDialect",
                              "mlir::triton::amdgpu::TritonAMDGPUDialect"];
+}
+
+def TritonAMDGPUConvertWarpSpecializeToLLVM : Pass<"triton-amdgpu-convert-warp-specialize-to-llvm", "mlir::ModuleOp"> {
+  let summary = "lower `ttg.warp_specialize` to LLVM";
+  let constructor = "mlir::triton::AMD::createTritonAMDGPUConvertWarpSpecializeToLLVMPass(\"\")";
+  let description = [{
+    The `triton-amdgpu-convert-warp-specialize-to-llvm` pass performs codegen for warp
+    specialization. It is a function-level transformation that rewrites
+    warp-specialized kernels by using shared memory and barriers to communicate
+    states between the default warpgroup and the worker warps.
+  }];
+
+  let dependentDialects = ["mlir::LLVM::LLVMDialect", "mlir::ROCDL::ROCDLDialect"];
 
+  let options = [
+    Option<"arch", "arch", "std::string", /*default*/"\"\"",
+           "target device architecture, e.g., gfx1250">,
+  ];
 }
 
 #endif

third_party/amd/include/TritonAMDGPUToLLVM/Passes.td

@@ -9,7 +9,10 @@ using namespace mlir;
 using namespace mlir::triton;
 
 constexpr int kBarrierCountBitWidth = 29;
-constexpr int kBarrierPhaseMask = ((1ULL << (32 - kBarrierCountBitWidth)) - 1);
+// NOTE: We only care for the parity of the phase (0: even, 1: odd), so use 1
+// bit constexpr int kBarrierPhaseMask = ((1ULL << (32 - kBarrierCountBitWidth))
+// - 1);
+constexpr int kBarrierPhaseMask = 1;
 constexpr int kInitCountPos = 32;
 
 namespace {

third_party/amd/lib/TritonAMDGPUToLLVM/BarrierOpToLLVM.cpp

@@ -6,6 +6,7 @@ add_triton_library(TritonAMDGPUToLLVM
     TensorPtrOpsToLLVM.cpp
     ConvertLayoutOpToLLVM.cpp
     ConvertWarpPipeline.cpp
+    ConvertWarpSpecializeToLLVM.cpp
     MemoryOpToLLVM.cpp
     MaskedOpsToLLVM.cpp
     DotOpToLLVM/FMA.cpp
@@ -35,6 +36,7 @@ add_triton_library(TritonAMDGPUToLLVM
     LLVMIRIncGen
 
     LINK_LIBS PUBLIC
+    MLIRReconcileUnrealizedCasts
     TritonGPUToLLVM
     TritonAMDGPUIR
     LLVMCore