From c500759c3a7a1ac8277f87b9771b279d0ac4b062 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee@intel.com>
Date: Mon, 22 Sep 2025 20:16:27 +0000
Subject: [PATCH 1/2] Add XeVM special id ops.

---
 mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td | 33 +++++++++++++++++++
 mlir/test/Dialect/LLVMIR/xevm.mlir          | 36 +++++++++++++++++++++
 2 files changed, 69 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
index 514b01a69fb9b..ee45f479a926e 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_WgDim : XeVM_SpecialRegisterXYZ<"local_size">;
+
+//===----------------------------------------------------------------------===//
+// Workgroup index and range
+defm XeVM_WgId : 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_SgIdOp : XeVM_SpecialIdRegisterOp<"subgroup_id">;
+def XeVM_SgSizeOp : 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
+}

From 685dfd7890bbcef9748a8c989978f40061c98ef6 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee@intel.com>
Date: Mon, 29 Sep 2025 20:16:56 +0000
Subject: [PATCH 2/2] Use better naming.

---
 mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
index ee45f479a926e..4f7a8421c07b9 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
@@ -654,17 +654,17 @@ multiclass XeVM_SpecialRegisterXYZ<string mnemonic, list<Trait> traits = []> {
 //===----------------------------------------------------------------------===//
 // Workitem index and range
 defm XeVM_WorkitemId : XeVM_SpecialRegisterXYZ<"local_id">;
-defm XeVM_WgDim : XeVM_SpecialRegisterXYZ<"local_size">;
+defm XeVM_WorkgroupDim : XeVM_SpecialRegisterXYZ<"local_size">;
 
 //===----------------------------------------------------------------------===//
 // Workgroup index and range
-defm XeVM_WgId : XeVM_SpecialRegisterXYZ<"group_id">;
+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_SgIdOp : XeVM_SpecialIdRegisterOp<"subgroup_id">;
-def XeVM_SgSizeOp : XeVM_SpecialIdRegisterOp<"subgroup_size">;
+def XeVM_SubgroupIdOp : XeVM_SpecialIdRegisterOp<"subgroup_id">;
+def XeVM_SubgroupSizeOp : XeVM_SpecialIdRegisterOp<"subgroup_size">;
 
 #endif // XEVMIR_OPS