[Mlir-commits] [mlir] [ROCDL] Added workgroup cluster ids to ROCDL (PR #179897)

[Ravil Dorozhinskii]

https://github.com/ravil-mobile updated https://github.com/llvm/llvm-project/pull/179897

>From d1964f2d11057e038f409e05059a1f555748da27 Mon Sep 17 00:00:00 2001
From: ravil-mobile <ravil.aviva.com at gmail.com>
Date: Thu, 5 Feb 2026 09:38:37 +0000
Subject: [PATCH] [ROCDL] Added workgroup cluster ids to ROCDL

---
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td |  4 +++
 mlir/test/Dialect/LLVMIR/rocdl.mlir          | 20 +++++++-----
 mlir/test/Target/LLVMIR/rocdl.mlir           | 32 ++++++++++++--------
 3 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index a6be6dea4b08f..6a874aafdec38 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -351,6 +351,10 @@ 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_ClusterWorkgroupIdXOp : ROCDL_SpecialIdRegisterOp<"cluster.workgroup.id.x">;
+def ROCDL_ClusterWorkgroupIdYOp : ROCDL_SpecialIdRegisterOp<"cluster.workgroup.id.y">;
+def ROCDL_ClusterWorkgroupIdZOp : ROCDL_SpecialIdRegisterOp<"cluster.workgroup.id.z">;
+
 def ROCDL_WaveId : ROCDL_SpecialIdRegisterOp<"wave.id">;
 def ROCDL_WavefrontSizeOp : ROCDL_SpecialIdRegisterOp<"wavefrontsize">;
 
diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir
index 3258d86c1642e..ca599e8025348 100644
--- a/mlir/test/Dialect/LLVMIR/rocdl.mlir
+++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir
@@ -20,20 +20,26 @@ func.func @rocdl_special_regs() -> i32 {
   %7 = rocdl.cluster.id.y : i32
   // CHECK: rocdl.cluster.id.z : i32
   %8 = rocdl.cluster.id.z : i32
+  // CHECK: rocdl.cluster.workgroup.id.x : i32
+  %9 = rocdl.cluster.workgroup.id.x : i32
+  // CHECK: rocdl.cluster.workgroup.id.y : i32
+  %10 = rocdl.cluster.workgroup.id.y : i32
+  // CHECK: rocdl.cluster.workgroup.id.z : i32
+  %11 = rocdl.cluster.workgroup.id.z : i32
   // CHECK: rocdl.workgroup.dim.x : i32
-  %9 = rocdl.workgroup.dim.x : i32
+  %12 = rocdl.workgroup.dim.x : i32
   // CHECK: rocdl.workgroup.dim.y : i32
-  %10 = rocdl.workgroup.dim.y : i32
+  %13 = rocdl.workgroup.dim.y : i32
   // CHECK: rocdl.workgroup.dim.z : i32
-  %11 = rocdl.workgroup.dim.z : i32
+  %14 = rocdl.workgroup.dim.z : i32
   // CHECK: rocdl.grid.dim.x : i32
-  %12 = rocdl.grid.dim.x : i32
+  %15 = rocdl.grid.dim.x : i32
   // CHECK: rocdl.grid.dim.y : i32
-  %13 = rocdl.grid.dim.y : i32
+  %16 = rocdl.grid.dim.y : i32
   // CHECK: rocdl.grid.dim.z : i32
-  %14 = rocdl.grid.dim.z : i32
+  %17 = rocdl.grid.dim.z : i32
   // CHECK: rocdl.wave.id : i32
-  %15 = rocdl.wave.id : i32
+  %18 = rocdl.wave.id : i32
   llvm.return %0 : i32
 }
 
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index 46c6d253547a3..382bc0b9f8ff6 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -20,36 +20,44 @@ llvm.func @rocdl_special_regs() -> i32 {
   %8 = rocdl.cluster.id.y : i32
   // CHECK: call i32 @llvm.amdgcn.cluster.id.z()
   %9 = rocdl.cluster.id.z : i32
+  // CHECK: call range(i32 0, 16) i32 @llvm.amdgcn.cluster.workgroup.id.x()
+  %10 = rocdl.cluster.workgroup.id.x range <i32, 0, 16> : i32
+  // CHECK: call range(i32 0, 16) i32 @llvm.amdgcn.cluster.workgroup.id.y()
+  %11 = rocdl.cluster.workgroup.id.y range <i32, 0, 16> : i32
+  // CHECK: call range(i32 0, 16) i32 @llvm.amdgcn.cluster.workgroup.id.z()
+  %12 = rocdl.cluster.workgroup.id.z range <i32, 0, 16> : i32
+
   // CHECK: call i64 @__ockl_get_local_size(i32 0)
-  %10 = rocdl.workgroup.dim.x : i64
+  %13 = rocdl.workgroup.dim.x : i64
   // CHECK: call i64 @__ockl_get_local_size(i32 1)
-  %11 = rocdl.workgroup.dim.y : i64
+  %14 = rocdl.workgroup.dim.y : i64
   // CHECK: call i64 @__ockl_get_local_size(i32 2)
-  %12 = rocdl.workgroup.dim.z : i64
+  %15 = rocdl.workgroup.dim.z : i64
+
   // CHECK: call i64 @__ockl_get_num_groups(i32 0)
-  %13 = rocdl.grid.dim.x : i64
+  %16 = rocdl.grid.dim.x : i64
   // CHECK: call i64 @__ockl_get_num_groups(i32 1)
-  %14 = rocdl.grid.dim.y : i64
+  %17 = rocdl.grid.dim.y : i64
   // CHECK: call i64 @__ockl_get_num_groups(i32 2)
-  %15 = rocdl.grid.dim.z : i64
+  %18 = rocdl.grid.dim.z : i64
 
   // CHECK: call range(i32 0, 64) i32 @llvm.amdgcn.workitem.id.x()
-  %16 = rocdl.workitem.id.x range <i32, 0, 64> : i32
+  %19 = rocdl.workitem.id.x range <i32, 0, 64> : i32
 
   // CHECK: call range(i64 1, 65) i64 @__ockl_get_local_size(i32 0)
-  %17 = rocdl.workgroup.dim.x range <i32, 1, 65> : i64
+  %20 = rocdl.workgroup.dim.x range <i32, 1, 65> : i64
 
   // CHECK: call i32 @llvm.amdgcn.wave.id()
-  %18 = rocdl.wave.id : i32
+  %21 = rocdl.wave.id : i32
 
   // CHECK: call range(i32 32, 65) i32 @llvm.amdgcn.wave.id()
-  %19 = rocdl.wave.id range <i32, 32, 65> : i32
+  %22 = rocdl.wave.id range <i32, 32, 65> : i32
 
   // CHECK: call i32 @llvm.amdgcn.wavefrontsize()
-  %20 = rocdl.wavefrontsize : i32
+  %23 = rocdl.wavefrontsize : i32
 
   // CHECK: call range(i32 32, 65) i32 @llvm.amdgcn.wavefrontsize()
-  %21 = rocdl.wavefrontsize range <i32, 32, 65> : i32
+  %24 = rocdl.wavefrontsize range <i32, 32, 65> : i32
 
   llvm.return %1 : i32
 }