diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td index 963078882c59d..0965979b7c39d 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td @@ -254,7 +254,7 @@ def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res", } //===----------------------------------------------------------------------===// -// Thread index and Block index +// Thread, Block and Cluster index //===----------------------------------------------------------------------===// def ROCDL_ThreadIdXOp : ROCDL_SpecialIdRegisterOp<"workitem.id.x">; @@ -265,6 +265,10 @@ def ROCDL_BlockIdXOp : ROCDL_SpecialIdRegisterOp<"workgroup.id.x">; def ROCDL_BlockIdYOp : ROCDL_SpecialIdRegisterOp<"workgroup.id.y">; def ROCDL_BlockIdZOp : ROCDL_SpecialIdRegisterOp<"workgroup.id.z">; +def ROCDL_ClusterIdXOp : ROCDL_SpecialIdRegisterOp<"cluster.id.x">; +def ROCDL_ClusterIdYOp : ROCDL_SpecialIdRegisterOp<"cluster.id.y">; +def ROCDL_ClusterIdZOp : ROCDL_SpecialIdRegisterOp<"cluster.id.z">; + def ROCDL_WavefrontSizeOp : ROCDL_SpecialIdRegisterOp<"wavefrontsize">; //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir index 9131c8ee9e598..675975ae597ac 100644 --- a/mlir/test/Dialect/LLVMIR/rocdl.mlir +++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir @@ -14,18 +14,24 @@ func.func @rocdl_special_regs() -> i32 { %4 = rocdl.workgroup.id.y : i32 // CHECK: rocdl.workgroup.id.z : i32 %5 = rocdl.workgroup.id.z : i32 + // CHECK: rocdl.cluster.id.x : i32 + %6 = rocdl.cluster.id.x : i32 + // CHECK: rocdl.cluster.id.y : i32 + %7 = rocdl.cluster.id.y : i32 + // CHECK: rocdl.cluster.id.z : i32 + %8 = rocdl.cluster.id.z : i32 // CHECK: rocdl.workgroup.dim.x : i32 - %6 = rocdl.workgroup.dim.x : i32 + %9 = rocdl.workgroup.dim.x : i32 // CHECK: rocdl.workgroup.dim.y : i32 - %7 = rocdl.workgroup.dim.y : i32 + %10 = rocdl.workgroup.dim.y : i32 // CHECK: rocdl.workgroup.dim.z : i32 - %8 = rocdl.workgroup.dim.z : i32 + %11 = rocdl.workgroup.dim.z : i32 // CHECK: rocdl.grid.dim.x : i32 - %9 = rocdl.grid.dim.x : i32 + %12 = rocdl.grid.dim.x : i32 // CHECK: rocdl.grid.dim.y : i32 - %10 = rocdl.grid.dim.y : i32 + %13 = rocdl.grid.dim.y : i32 // CHECK: rocdl.grid.dim.z : i32 - %11 = rocdl.grid.dim.z : i32 + %14 = rocdl.grid.dim.z : i32 llvm.return %0 : i32 } diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir index ca1620b81c318..2922665295cf3 100644 --- a/mlir/test/Target/LLVMIR/rocdl.mlir +++ b/mlir/test/Target/LLVMIR/rocdl.mlir @@ -14,30 +14,36 @@ llvm.func @rocdl_special_regs() -> i32 { %5 = rocdl.workgroup.id.y : i32 // CHECK: call i32 @llvm.amdgcn.workgroup.id.z() %6 = rocdl.workgroup.id.z : i32 + // CHECK: call i32 @llvm.amdgcn.cluster.id.x() + %7 = rocdl.cluster.id.x : i32 + // CHECK: call i32 @llvm.amdgcn.cluster.id.y() + %8 = rocdl.cluster.id.y : i32 + // CHECK: call i32 @llvm.amdgcn.cluster.id.z() + %9 = rocdl.cluster.id.z : i32 // CHECK: call i64 @__ockl_get_local_size(i32 0) - %7 = rocdl.workgroup.dim.x : i64 + %10 = rocdl.workgroup.dim.x : i64 // CHECK: call i64 @__ockl_get_local_size(i32 1) - %8 = rocdl.workgroup.dim.y : i64 + %11 = rocdl.workgroup.dim.y : i64 // CHECK: call i64 @__ockl_get_local_size(i32 2) - %9 = rocdl.workgroup.dim.z : i64 + %12 = rocdl.workgroup.dim.z : i64 // CHECK: call i64 @__ockl_get_num_groups(i32 0) - %10 = rocdl.grid.dim.x : i64 + %13 = rocdl.grid.dim.x : i64 // CHECK: call i64 @__ockl_get_num_groups(i32 1) - %11 = rocdl.grid.dim.y : i64 + %14 = rocdl.grid.dim.y : i64 // CHECK: call i64 @__ockl_get_num_groups(i32 2) - %12 = rocdl.grid.dim.z : i64 + %15 = rocdl.grid.dim.z : i64 // CHECK: call range(i32 0, 64) i32 @llvm.amdgcn.workitem.id.x() - %13 = rocdl.workitem.id.x range : i32 + %16 = rocdl.workitem.id.x range : i32 // CHECK: call range(i64 1, 65) i64 @__ockl_get_local_size(i32 0) - %14 = rocdl.workgroup.dim.x range : i64 + %17 = rocdl.workgroup.dim.x range : i64 // CHECK: call i32 @llvm.amdgcn.wavefrontsize() - %15 = rocdl.wavefrontsize : i32 + %18 = rocdl.wavefrontsize : i32 // CHECK: call range(i32 32, 65) i32 @llvm.amdgcn.wavefrontsize() - %16 = rocdl.wavefrontsize range : i32 + %19 = rocdl.wavefrontsize range : i32 llvm.return %1 : i32 }