diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td --- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td @@ -249,6 +249,30 @@ 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 //===----------------------------------------------------------------------===// diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir --- a/mlir/test/Target/LLVMIR/nvvmir.mlir +++ b/mlir/test/Target/LLVMIR/nvvmir.mlir @@ -30,6 +30,35 @@ %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 }