[NVIDIA] Replace inline assembly for the lowering of ttn::ClusterCTAIdOp (#7512)

This PR refactors the ClusterCTAIdOp conversion from using inline PTX
assembly to a series of operations (including some NVVM ops that can
generate intrinsic call), preserving more semantic information at the
LLVM level. While the new implementation expands the computation into
separate multiply and add operations, the backend will typically
optimize them into `mad`, so there is no performance regression.

Author: Pecco-314

test/Conversion/nvgpu_to_llvm.mlir

@@ -2,11 +2,11 @@
 
 // CHECK-LABEL: @cluster_id
 llvm.func @cluster_id() -> i32 {
-  // CHECK:      %cluster_ctaid.x;
-  // CHECK-SAME: %cluster_ctaid.y;
-  // CHECK-SAME: %cluster_ctaid.z;
-  // CHECK-SAME: %cluster_nctaid.x;
-  // CHECK-SAME: %cluster_nctaid.y;
+  // CHECK: nvvm.read.ptx.sreg.cluster.ctaid.x
+  // CHECK: nvvm.read.ptx.sreg.cluster.ctaid.y
+  // CHECK: nvvm.read.ptx.sreg.cluster.ctaid.z
+  // CHECK: nvvm.read.ptx.sreg.cluster.nctaid.x
+  // CHECK: nvvm.read.ptx.sreg.cluster.nctaid.y
   %id = nvgpu.cluster_id
   llvm.return %id : i32
 }

third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp

@@ -23,17 +23,6 @@ namespace triton {
 
 namespace {
 
-const std::string kClusterCtaIdOp = "{\n"
-                                    ".reg .u32 a<5>;              \n"
-                                    "mov.u32 a0, %cluster_ctaid.x;\n"  // x
-                                    "mov.u32 a1, %cluster_ctaid.y;\n"  // y
-                                    "mov.u32 a2, %cluster_ctaid.z;\n"  // z
-                                    "mov.u32 a3, %cluster_nctaid.x;\n" // nx
-                                    "mov.u32 a4, %cluster_nctaid.y;\n" // ny
-                                    "mad.lo.u32 a1, a2, a4, a1;     \n"
-                                    "mad.lo.u32 $0, a1, a3, a0;     \n"
-                                    "}";
-
 bool isNumber(const std::string &s) {
   return !s.empty() && std::find_if(s.begin(), s.end(), [](unsigned char c) {
                          return !std::isdigit(c);
@@ -238,6 +227,26 @@ class WarpIdOpPattern : public OpRewritePattern<ttn::WarpIdOp> {
   }
 };
 
+class ClusterCTAIdOpPattern : public OpRewritePattern<ttn::ClusterCTAIdOp> {
+  using OpRewritePattern<ttn::ClusterCTAIdOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ttn::ClusterCTAIdOp op,
+                                PatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto a0 = rewriter.create<NVVM::BlockInClusterIdXOp>(loc, i32_ty);
+    auto a1 = rewriter.create<NVVM::BlockInClusterIdYOp>(loc, i32_ty);
+    auto a2 = rewriter.create<NVVM::BlockInClusterIdZOp>(loc, i32_ty);
+    auto a3 = rewriter.create<NVVM::ClusterDimBlocksXOp>(loc, i32_ty);
+    auto a4 = rewriter.create<NVVM::ClusterDimBlocksYOp>(loc, i32_ty);
+    auto p1 = rewriter.create<LLVM::MulOp>(loc, a2, a4);
+    auto s1 = rewriter.create<LLVM::AddOp>(loc, a1, p1);
+    auto p2 = rewriter.create<LLVM::MulOp>(loc, s1, a3);
+    auto res = rewriter.create<LLVM::AddOp>(loc, a0, p2);
+    rewriter.replaceOp(op, res);
+    return success();
+  }
+};
+
 // Base class for Matrix Operation Patterns
 template <typename MatrixOpType, typename ConcreteMatrixOpPattern>
 class MatrixOpPattern : public OpRewritePattern<MatrixOpType> {
@@ -736,11 +745,10 @@ class ConvertNVGPUToLLVM
     ModuleOp mod = getOperation();
     RewritePatternSet patterns(context);
 
-    patterns.add<NVGPUOpGenericPattern<ttn::ClusterCTAIdOp>>(
-        context, kClusterCtaIdOp, Constraints({"=r"}), Constraints());
-
-    patterns.add<LoadMatrixOpPattern, WGMMAOpPattern, LoadAcquireOpPattern,
-                 WGMMAWaitGroupOpPattern, WarpIdOpPattern>(context);
+    patterns
+        .add<ClusterCTAIdOpPattern, LoadMatrixOpPattern, WGMMAOpPattern,
+             LoadAcquireOpPattern, WGMMAWaitGroupOpPattern, WarpIdOpPattern>(
+            context);
 
     if (applyPatternsGreedily(mod, std::move(patterns)).failed())
       signalPassFailure();