[MLIR][XeVM] Add XeVM special id ops.

[silee2]

Add special GPU id, index ops.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Sang Ik Lee (silee2)

Changes

Add special GPU id, index ops.

---

Full diff: https://github.com/llvm/llvm-project/pull/160735.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td (+33)
• (modified) mlir/test/Dialect/LLVMIR/xevm.mlir (+36)

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
+}

[silee2]

@akroviakov

[akroviakov]

Not sure whether Wg/Sg should really be used instead of full names, like we have in WorkitemId, this may need more consistency. Otherwise, LGTM.

[Jianhui-Li]

LGTM

[Jianhui-Li]

Not sure whether Wg/Sg should really be used instead of full names, like we have in WorkitemId, this may need more consistency. Otherwise, LGTM.

Agree on using full name to be consistent.

[silee2]

Not sure whether Wg/Sg should really be used instead of full names, like we have in WorkitemId, this may need more consistency. Otherwise, LGTM.

Renamed and using full names.