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

Add special GPU id, index ops.

Author: silee2

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

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