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 traits = []> + : XeVM_Op, + Results<(outs AnyTypeOf<[I32, I64]>:$res)>, + Arguments<(ins OptionalAttr:$range)> { + let assemblyFormat = "(`range` $range^)? attr-dict `:` type($res)"; +} + +multiclass XeVM_SpecialRegisterXYZ traits = []> { + def XOp : XeVM_SpecialIdRegisterOp; + def YOp : XeVM_SpecialIdRegisterOp; + def ZOp : XeVM_SpecialIdRegisterOp; +} + +//===----------------------------------------------------------------------===// +// 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] { gpu.module @xevm_module [#xevm.target]{ } + +// ----- +// 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 +}