[Mlir-commits] [mlir] 66af942 - [MLIR][XeVM] Add XeVM special id ops. (#160735)

[llvmlistbot at llvm.org]

Author: Sang Ik Lee
Date: 2025-09-30T10:16:20-07:00
New Revision: 66af9423e882247ca2389d1d20c7ee9b21b50a82

URL: https://github.com/llvm/llvm-project/commit/66af9423e882247ca2389d1d20c7ee9b21b50a82
DIFF: https://github.com/llvm/llvm-project/commit/66af9423e882247ca2389d1d20c7ee9b21b50a82.diff

LOG: [MLIR][XeVM] Add XeVM special id ops. (#160735)

Add special GPU id, index ops.

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
    mlir/test/Dialect/LLVMIR/xevm.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
index 514b01a69fb9b..4f7a8421c07b9 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
@@ -634,4 +634,37 @@ def XeVM_TargetAttr : XeVM_Attr<"XeVMTarget", "target"> {
   let genVerifyDecl = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// XeVM special register op definitions
+//===----------------------------------------------------------------------===//
+
+class XeVM_SpecialIdRegisterOp<string mnemonic, list<Trait> traits = []>
+    : XeVM_Op<mnemonic, traits>,
+      Results<(outs AnyTypeOf<[I32, I64]>:$res)>,
+      Arguments<(ins OptionalAttr<LLVM_ConstantRangeAttr>:$range)> {
+  let assemblyFormat = "(`range` $range^)? attr-dict `:` type($res)";
+}
+
+multiclass XeVM_SpecialRegisterXYZ<string mnemonic, list<Trait> traits = []> {
+  def XOp : XeVM_SpecialIdRegisterOp<!strconcat(mnemonic, ".x"), traits>;
+  def YOp : XeVM_SpecialIdRegisterOp<!strconcat(mnemonic, ".y"), traits>;
+  def ZOp : XeVM_SpecialIdRegisterOp<!strconcat(mnemonic, ".z"), traits>;
+}
+
+//===----------------------------------------------------------------------===//
+// Workitem index and range
+defm XeVM_WorkitemId : XeVM_SpecialRegisterXYZ<"local_id">;
+defm XeVM_WorkgroupDim : XeVM_SpecialRegisterXYZ<"local_size">;
+
+//===----------------------------------------------------------------------===//
+// Workgroup index and range
+defm XeVM_WorkgroupId : XeVM_SpecialRegisterXYZ<"group_id">;
+defm XeVM_GridDim : XeVM_SpecialRegisterXYZ<"group_count">;
+
+//===----------------------------------------------------------------------===//
+// Lane, Subgroup index and range
+def XeVM_LaneIdOp : XeVM_SpecialIdRegisterOp<"lane_id">;
+def XeVM_SubgroupIdOp : XeVM_SpecialIdRegisterOp<"subgroup_id">;
+def XeVM_SubgroupSizeOp : XeVM_SpecialIdRegisterOp<"subgroup_size">;
+
 #endif // XEVMIR_OPS

diff  --git a/mlir/test/Dialect/LLVMIR/xevm.mlir b/mlir/test/Dialect/LLVMIR/xevm.mlir
index bb1f650a1cd12..66fb2949a270f 100644
--- a/mlir/test/Dialect/LLVMIR/xevm.mlir
+++ b/mlir/test/Dialect/LLVMIR/xevm.mlir
@@ -116,3 +116,39 @@ func.func @prefetch(%ptr: !llvm.ptr<1>) {
 // CHECK-LABEL: @xevm_module [#xevm.target<O = 3, chip = "pvc">] {
 gpu.module @xevm_module [#xevm.target<O = 3, chip = "pvc">]{
 }
+
+// -----
+// CHECK-LABEL: @xevm_special_ids
+llvm.func @xevm_special_ids() -> i32 {
+  // CHECK: xevm.local_id.x : i32
+  %1 = xevm.local_id.x : i32
+  // CHECK: xevm.local_id.y : i32
+  %2 = xevm.local_id.y : i32
+  // CHECK: xevm.local_id.z : i32
+  %3 = xevm.local_id.z : i32
+  // CHECK: xevm.local_size.x : i32
+  %4 = xevm.local_size.x : i32
+  // CHECK: xevm.local_size.y : i32
+  %5 = xevm.local_size.y : i32
+  // CHECK: xevm.local_size.z : i32
+  %6 = xevm.local_size.z : i32
+  // CHECK: xevm.group_id.x : i32
+  %7 = xevm.group_id.x : i32
+  // CHECK: xevm.group_id.y : i32
+  %8 = xevm.group_id.y : i32
+  // CHECK: xevm.group_id.z : i32
+  %9 = xevm.group_id.z : i32
+  // CHECK: xevm.group_count.x : i32
+  %10 = xevm.group_count.x : i32
+  // CHECK: xevm.group_count.y : i32
+  %11 = xevm.group_count.y : i32
+  // CHECK: xevm.group_count.z : i32
+  %12 = xevm.group_count.z : i32
+  // CHECK: xevm.lane_id : i32
+  %14 = xevm.lane_id : i32
+  // CHECK: xevm.subgroup_size : i32
+  %39 = xevm.subgroup_size : i32
+  // CHECK: xevm.subgroup_id : i32
+  %40 = xevm.subgroup_id : i32
+  llvm.return %1 : i32
+}