[MLIR][NVVM] Introduce special registers for CTA Cluster

This work introduces special registers such as cluster ID, dimensions, and more for managing CTA clusters, which are groups of CTAsthat can synchronize and communicate through shared memory. This is for Nvidia's sm_90 capability.

Differential Revision: https://reviews.llvm.org/D158588

Author: grypp

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

@@ -249,6 +249,30 @@ def NVVM_GridDimXOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.nctaid.x">;
 def NVVM_GridDimYOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.nctaid.y">;
 def NVVM_GridDimZOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.nctaid.z">;
 
+//===----------------------------------------------------------------------===//
+// CTA Cluster index and range
+def NVVM_ClusterIdXOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.clusterid.x">;
+def NVVM_ClusterIdYOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.clusterid.y">;
+def NVVM_ClusterIdZOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.clusterid.z">;
+def NVVM_ClusterDimXOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.nclusterid.x">;
+def NVVM_ClusterDimYOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.nclusterid.y">;
+def NVVM_ClusterDimZOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.nclusterid.z">;
+
+
+//===----------------------------------------------------------------------===//
+// CTA index and range within Cluster
+def NVVM_BlockInClusterIdXOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.cluster.ctaid.x">;
+def NVVM_BlockInClusterIdYOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.cluster.ctaid.y">;
+def NVVM_BlockInClusterIdZOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.cluster.ctaid.z">;
+def NVVM_GridInClusterDimXOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.cluster.nctaid.x">;
+def NVVM_GridInClusterDimYOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.cluster.nctaid.y">;
+def NVVM_GridInClusterDimZOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.cluster.nctaid.z">;
+
+//===----------------------------------------------------------------------===//
+// CTA index and across Cluster dimensions
+def NVVM_ClusterId : NVVM_SpecialRegisterOp<"read.ptx.sreg.cluster.ctarank">;
+def NVVM_ClusterDim : NVVM_SpecialRegisterOp<"read.ptx.sreg.cluster.nctarank">;
+
 //===----------------------------------------------------------------------===//
 // NVVM approximate op definitions
 //===----------------------------------------------------------------------===//

mlir/test/Target/LLVMIR/nvvmir.mlir

@@ -30,6 +30,35 @@ llvm.func @nvvm_special_regs() -> i32 {
   %13 = nvvm.read.ptx.sreg.warpsize : i32
   // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.laneid()
   %14 = nvvm.read.ptx.sreg.laneid : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.clusterid.x
+  %15 = nvvm.read.ptx.sreg.clusterid.x : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.clusterid.y
+  %16 = nvvm.read.ptx.sreg.clusterid.y : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.clusterid.z
+  %17 = nvvm.read.ptx.sreg.clusterid.z : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.nclusterid.x
+  %18 = nvvm.read.ptx.sreg.nclusterid.x : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.nclusterid.y
+  %19 = nvvm.read.ptx.sreg.nclusterid.y : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.nclusterid.z
+  %20 = nvvm.read.ptx.sreg.nclusterid.z : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid
+  %21 = nvvm.read.ptx.sreg.cluster.ctaid.x : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid
+  %22 = nvvm.read.ptx.sreg.cluster.ctaid.y : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid
+  %23 = nvvm.read.ptx.sreg.cluster.ctaid.z : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid
+  %24 = nvvm.read.ptx.sreg.cluster.nctaid.x : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid
+  %25 = nvvm.read.ptx.sreg.cluster.nctaid.y : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid
+  %26 = nvvm.read.ptx.sreg.cluster.nctaid.z : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.cluster.ctarank
+  %27 = nvvm.read.ptx.sreg.cluster.ctarank : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctarank
+  %28 = nvvm.read.ptx.sreg.cluster.nctarank : i32
+  
   llvm.return %1 : i32
 }