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

[guraypp]

Summary

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.

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

 def NVVM_BlockIdXOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.ctaid.x">;
 def NVVM_BlockIdYOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.ctaid.y">;
 def NVVM_BlockIdZOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.ctaid.z">;
 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
 //===----------------------------------------------------------------------===//
 
 def NVVM_RcpApproxFtzF32Op : NVVM_IntrOp<"rcp.approx.ftz.f", [Pure], 1> {
   let arguments = (ins F32:$arg);
   let results = (outs F32:$res);
   let assemblyFormat = "$arg attr-dict `:` type($res)";
 }

mlir/test/Target/LLVMIR/nvvmir.mlir

llvm.func @nvvm_special_regs() -> i32 {
   // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.nctaid.y()
   %11 = nvvm.read.ptx.sreg.nctaid.y : i32
   // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.nctaid.z()
   %12 = nvvm.read.ptx.sreg.nctaid.z : i32
   // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.warpsize()
   %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
 }
 
 // CHECK-LABEL: @nvvm_rcp
 llvm.func @nvvm_rcp(%0: f32) -> f32 {
   // CHECK: call float @llvm.nvvm.rcp.approx.ftz.f
   %1 = nvvm.rcp.approx.ftz.f %0 : f32
   llvm.return %1 : f32