[SYCL-MLIR] Define SYCL index space operations (#8089)

Define high-level operations to retrieve information from the SYCL index
space.

---------

Signed-off-by: Victor Perez <[email protected]>

Author: victor-eds

mlir-sycl/include/mlir/Dialect/SYCL/IR/SYCLOpInterfaces.td

@@ -95,5 +95,7 @@ def SYCLGetID : SYCLOpTrait<"SYCLGetID">;
 def SYCLGetComponent : SYCLOpTrait<"SYCLGetComponent">;
 def SYCLGetRange : SYCLOpTrait<"SYCLGetRange">;
 def SYCLGetGroup : SYCLOpTrait<"SYCLGetGroup">;
+def SYCLIndexSpaceGetID : SYCLOpTrait<"SYCLIndexSpaceGetID">;
+def SYCLIndexSpaceGetRange : SYCLOpTrait<"SYCLIndexSpaceGetRange">;
 
 #endif // SYCL_OP_INTERFACES

mlir-sycl/include/mlir/Dialect/SYCL/IR/SYCLOpTraits.h

@@ -17,6 +17,8 @@ LogicalResult verifySYCLGetIDTrait(Operation *Op);
 LogicalResult verifySYCLGetComponentTrait(Operation *Op);
 LogicalResult verifySYCLGetRangeTrait(Operation *Op);
 LogicalResult verifySYCLGetGroupTrait(Operation *Op);
+LogicalResult verifySYCLIndexSpaceGetIDTrait(Operation *Op);
+LogicalResult verifySYCLIndexSpaceGetRangeTrait(Operation *Op);
 
 ///  This interface describes an SYCLMethodOpInterface that returns a range if
 ///  called with a single argument and a size_t if called with two arguments.
@@ -64,6 +66,28 @@ class SYCLGetGroup : public OpTrait::TraitBase<ConcreteType, SYCLGetGroup> {
     return verifySYCLGetGroupTrait(Op);
   }
 };
+
+///  This interface describes an operation returning either a SYCL ID type (for
+///  cardinality 0) or an MLIR index type (for cardinality 1).
+template <typename ConcreteType>
+class SYCLIndexSpaceGetID
+    : public OpTrait::TraitBase<ConcreteType, SYCLIndexSpaceGetID> {
+public:
+  static LogicalResult verifyTrait(Operation *Op) {
+    return verifySYCLIndexSpaceGetIDTrait(Op);
+  }
+};
+
+///  This interface describes an operation returning either a SYCL range type
+///  (for cardinality 0) or an MLIR index type (for cardinality 1).
+template <typename ConcreteType>
+class SYCLIndexSpaceGetRange
+    : public OpTrait::TraitBase<ConcreteType, SYCLIndexSpaceGetRange> {
+public:
+  static LogicalResult verifyTrait(Operation *Op) {
+    return verifySYCLIndexSpaceGetRangeTrait(Op);
+  }
+};
 } // namespace sycl
 } // namespace mlir
 

mlir-sycl/include/mlir/Dialect/SYCL/IR/SYCLOps.td

@@ -413,10 +413,188 @@ def SYCLMemref : AnyTypeOf<[
 def IndexType : AnyTypeOf<[I32, I64, Index]>;
 def SYCLGetIDResult : AnyTypeOf<[I64, SYCL_IDType]>;
 def SYCLGetRangeResult : AnyTypeOf<[I64, SYCL_RangeType]>;
+def SYCLIndexSpaceGetIDResult : AnyTypeOf<[Index, SYCL_IDType]>;
+def SYCLIndexSpaceGetRangeResult : AnyTypeOf<[Index, SYCL_RangeType]>;
 def SYCLGetResult : AnyTypeOf<[I64, MemRefOf<[I64]>]>;
 def VectorEltTy : AnyTypeOf<[I1, I8, I16, I32, I64, F16, F32, F64]>;
 def VectorSplatArg : MemRefOf<[VectorEltTy]>;
 
+////////////////////////////////////////////////////////////////////////////////
+// SYCL GRID OPERATIONS
+////////////////////////////////////////////////////////////////////////////////
+
+def SYCLNumWorkItemsOp : SYCL_Op<"num_work_items", [SYCLIndexSpaceGetRange]> {
+  let summary = "Retrieve the number of work-items.";
+  let description = [{
+    This operation returns the number of work-items in the index space. If the
+    optional argument is passed, the number of work-items in the given dimension
+    is returned.
+  }];
+
+  let arguments = (ins Optional<I32>:$dimension);
+  let results = (outs SYCLIndexSpaceGetRangeResult:$result);
+  let assemblyFormat = [{
+    `(` operands `)` attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def SYCLGlobalIDOp : SYCL_Op<"global_id", [SYCLIndexSpaceGetID]> {
+  let summary = "Retrieve the global ID of the item.";
+  let description = [{
+    This operation returns the global ID of the item in the index space. If the
+    optional argument is passed, the global ID of the item in the given
+    dimension is returned.
+  }];
+
+  let arguments = (ins Optional<I32>:$dimension);
+  let results = (outs SYCLIndexSpaceGetIDResult:$result);
+  let assemblyFormat = [{
+    `(` operands `)` attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def SYCLLocalIDOp : SYCL_Op<"local_id", [SYCLIndexSpaceGetID]> {
+  let summary = "Retrieve the local ID of the item.";
+  let description = [{
+    This operation returns the local ID of the item within the work-group. If
+    the optional argument is passed, the local ID of the item in the given
+    dimension is returned.
+  }];
+
+  let arguments = (ins Optional<I32>:$dimension);
+  let results = (outs SYCLIndexSpaceGetIDResult:$result);
+  let assemblyFormat = [{
+    `(` operands `)` attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def SYCLGlobalOffsetOp : SYCL_Op<"global_offset", [SYCLIndexSpaceGetID]>,
+                       Deprecated<"deprecated in SYCL 2020"> {
+  let summary = "Retrieve the global offset of the item.";
+  let description = [{
+    This operation returns the global offset of the item in the index space. If
+    the optional argument is passed, the global offset of the item in the given
+    dimension is returned.
+
+    Note that this is deprecated in SYCL 2020.
+  }];
+
+  let arguments = (ins Optional<I32>:$dimension);
+  let results = (outs SYCLIndexSpaceGetIDResult:$result);
+  let assemblyFormat = [{
+    `(` operands `)` attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def SYCLNumWorkGroupsOp : SYCL_Op<"num_work_groups", [SYCLIndexSpaceGetRange]> {
+  let summary = "Retrieve the number of work-groups.";
+  let description = [{
+    This operation returns the number of work-groups in the index space. If the
+    optional argument is passed, the number of work-groups in the given
+    dimension is returned.
+  }];
+
+  let arguments = (ins Optional<I32>:$dimension);
+  let results = (outs SYCLIndexSpaceGetRangeResult:$result);
+  let assemblyFormat = [{
+    `(` operands `)` attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def SYCLWorkGroupSizeOp : SYCL_Op<"work_group_size", [SYCLIndexSpaceGetRange]> {
+  let summary = "Retrieve the number of work-items in a work-group.";
+  let description = [{
+    This operation returns the number of work-items per work-group. If the
+    optional argument is passed, the number of work-items per work-group in the
+    given dimension is returned.
+  }];
+
+  let arguments = (ins Optional<I32>:$dimension);
+  let results = (outs SYCLIndexSpaceGetRangeResult:$result);
+  let assemblyFormat = [{
+    `(` operands `)` attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def SYCLWorkGroupIDOp : SYCL_Op<"work_group_id", [SYCLIndexSpaceGetID]> {
+  let summary = "Retrieve the ID of the work-group.";
+  let description = [{
+    This operation returns the ID of the work-group. If the optional argument is
+    passed, the ID of the work-group in the given dimension is returned.
+  }];
+
+  let arguments = (ins Optional<I32>:$dimension);
+  let results = (outs SYCLIndexSpaceGetIDResult:$result);
+  let assemblyFormat = [{
+    `(` operands `)` attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def SYCLNumSubGroupsOp : SYCL_Op<"num_sub_groups"> {
+  let summary = "Retrieve the number of sub-groups.";
+  let description = [{
+    This operation returns the number of sub-groups.
+  }];
+
+  let arguments = (ins);
+  let results = (outs I32:$result);
+  let assemblyFormat = [{
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def SYCLSubGroupMaxSizeOp : SYCL_Op<"sub_group_max_size"> {
+  let summary = "Retrieve the maximum size of a sub-group.";
+  let description = [{
+    This operation returns the maximum size of a sub-group.
+  }];
+
+  let arguments = (ins);
+  let results = (outs I32:$result);
+  let assemblyFormat = [{
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def SYCLSubGroupSizeOp : SYCL_Op<"sub_group_size"> {
+  let summary = "Retrieve the sub-group size.";
+  let description = [{
+    This operation returns the sub-group size.
+  }];
+
+  let arguments = (ins);
+  let results = (outs I32:$result);
+  let assemblyFormat = [{
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def SYCLSubGroupIDOp : SYCL_Op<"sub_group_id"> {
+  let summary = "Retrieve the ID of the sub-group.";
+  let description = [{
+    This operation returns the ID of the sub-group.
+  }];
+
+  let arguments = (ins);
+  let results = (outs I32:$result);
+  let assemblyFormat = [{
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def SYCLSubGroupLocalIDOp : SYCL_Op<"sub_group_local_id"> {
+  let summary = "Retrieve the local ID of the sub-group.";
+  let description = [{
+    This operation returns the local ID of the sub-group.
+  }];
+
+  let arguments = (ins);
+  let results = (outs I32:$result);
+  let assemblyFormat = [{
+    attr-dict `:` functional-type(operands, results)
+  }];
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // CONSTRUCTOR OPERATION
 ////////////////////////////////////////////////////////////////////////////////

mlir-sycl/lib/Dialect/IR/SYCLTraits.cpp

@@ -10,10 +10,14 @@
 
 #include "mlir/Dialect/SYCL/IR/SYCLOpTraits.h"
 
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/SYCL/IR/SYCLOps.h"
 
 #include "llvm/ADT/TypeSwitch.h"
 
+using namespace mlir;
+using namespace mlir::sycl;
+
 static unsigned getDimensions(mlir::Type Type) {
   if (auto MemRefTy = Type.dyn_cast<mlir::MemRefType>()) {
     Type = MemRefTy.getElementType();
@@ -174,3 +178,37 @@ mlir::LogicalResult mlir::sycl::verifySYCLGetGroupTrait(Operation *Op) {
   return verifyGetSYCLTyOperation(cast<mlir::sycl::SYCLMethodOpInterface>(Op),
                                   "group");
 }
+
+static LogicalResult verifyIndexSpaceTrait(Operation *Op) {
+  const auto Ty = Op->getResultTypes();
+  assert(Ty.size() == 1 && "Expecting a single return value");
+  const auto IsIndex = Ty[0].isa<IndexType>();
+  switch (Op->getNumOperands()) {
+  case 0:
+    return !IsIndex ? success()
+                    : Op->emitOpError("Not expecting an index return value for "
+                                      "this cardinality");
+  case 1:
+    if (auto C = Op->getOperand(0).getDefiningOp<arith::ConstantOp>()) {
+      const auto Value = static_cast<arith::ConstantIntOp>(C).value();
+      if (!(0 <= Value && Value < 3)) {
+        return Op->emitOpError(
+            "The SYCL index space can only be 1, 2, or 3 dimensional");
+      }
+    }
+    return IsIndex
+               ? success()
+               : Op->emitOpError(
+                     "Expecting an index return value for this cardinality");
+  default:
+    llvm_unreachable("Invalid cardinality");
+  }
+}
+
+LogicalResult mlir::sycl::verifySYCLIndexSpaceGetIDTrait(Operation *Op) {
+  return verifyIndexSpaceTrait(Op);
+}
+
+LogicalResult mlir::sycl::verifySYCLIndexSpaceGetRangeTrait(Operation *Op) {
+  return verifyIndexSpaceTrait(Op);
+}