[OpenMP Dialect] Add omp.canonical_loop operation.

mlir/include/mlir/Dialect/OpenMP/CMakeLists.txt

@@ -2,18 +2,15 @@ set(LLVM_TARGET_DEFINITIONS ${LLVM_MAIN_INCLUDE_DIR}/llvm/Frontend/OpenMP/OMP.td
 mlir_tablegen(OmpCommon.td --gen-directive-decl --directives-dialect=OpenMP)
 add_public_tablegen_target(omp_common_td)
 
+add_mlir_dialect(OpenMPOps omp)
 set(LLVM_TARGET_DEFINITIONS OpenMPOps.td)
-mlir_tablegen(OpenMPOpsDialect.h.inc -gen-dialect-decls -dialect=omp)
-mlir_tablegen(OpenMPOpsDialect.cpp.inc -gen-dialect-defs -dialect=omp)
-mlir_tablegen(OpenMPOps.h.inc -gen-op-decls)
-mlir_tablegen(OpenMPOps.cpp.inc -gen-op-defs)
 mlir_tablegen(OpenMPOpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(OpenMPOpsEnums.cpp.inc -gen-enum-defs)
 mlir_tablegen(OpenMPOpsAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=omp)
 mlir_tablegen(OpenMPOpsAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=omp)
 add_mlir_doc(OpenMPOps OpenMPDialect Dialects/ -gen-dialect-doc -dialect=omp)
-add_public_tablegen_target(MLIROpenMPOpsIncGen)
 add_dependencies(OpenMPDialectDocGen omp_common_td)
+
 add_mlir_interface(OpenMPOpsInterfaces)
 
 set(LLVM_TARGET_DEFINITIONS OpenMPTypeInterfaces.td)

mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h

@@ -25,6 +25,9 @@
 #include "mlir/Dialect/OpenMP/OpenMPOpsEnums.h.inc"
 #include "mlir/Dialect/OpenMP/OpenMPTypeInterfaces.h.inc"
 
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/OpenMP/OpenMPOpsTypes.h.inc"
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.h.inc"
 

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

@@ -28,6 +28,7 @@ def OpenMP_Dialect : Dialect {
   let cppNamespace = "::mlir::omp";
   let dependentDialects = ["::mlir::LLVM::LLVMDialect, ::mlir::func::FuncDialect"];
   let useDefaultAttributePrinterParser = 1;
+  let useDefaultTypePrinterParser = 1;
 }
 
 // OmpCommon requires definition of OpenACC_Dialect.
@@ -78,6 +79,10 @@ def TargetAttr : OpenMP_Attr<"Target", "target"> {
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
+class OpenMP_Type<string name, string typeMnemonic> : 
+      TypeDef<OpenMP_Dialect, name> {
+  let mnemonic = typeMnemonic;
+}
 
 class OpenMP_Op<string mnemonic, list<Trait> traits = []> :
       Op<OpenMP_Dialect, mnemonic, traits>;
@@ -399,6 +404,150 @@ def SingleOp : OpenMP_Op<"single", [AttrSizedOperandSegments]> {
   let hasVerifier = 1;
 }
 
+//===---------------------------------------------------------------------===//
+// OpenMP Canonical Loop Operation					  	
+//===---------------------------------------------------------------------===//
+
+def CanonicalLoopInfoType : OpenMP_Type<"CanonicalLoopInfo", "cli"> {
+  let summary = "Type for representing a reference to a canonical loop";
+  let description = [{  
+    A variable of type CanonicalLoopInfo refers to an OpenMP-compatible
+    canonical loop in the same function. Variables of this type are not
+    available at runtime and therefore cannot be used by the program itself,
+    i.e. an opaque type. It is similar to the transform dialect's
+    `!transform.interface` type, but instead of implementing an interface
+    for each transformation, the OpenMP dialect itself defines possible
+    operations on this type.
+
+    A CanonicalLoopInfo variable can
+
+    1. be passed to omp.yield to be accessible outside the loop.
+    2. passed to omp operations that take a CanonicalLoopInfo argument,
+       such as `omp.unroll`.
+
+    A CanonicalLoopInfo variable can not
+
+    1. be returned from a function,
+    2. passed to operations that are not specifically designed to take a
+       CanonicalLoopInfo, including AnyType.
+
+    A CanonicalLoopInfo variable directly corresponds to an object of
+    OpenMPIRBuilder's CanonicalLoopInfo struct when lowering to LLVM-IR.
+  }];
+}
+
+def CanonicalLoopOp : OpenMP_Op<"canonical_loop", [SingleBlockImplicitTerminator<"omp::YieldOp">]> {
+  let summary = "OpenMP Canonical Loop Operation"; 
+  let description = [{
+    All loops that conform to OpenMP's definition of a canonical loop can be
+    simplified to a CanonicalLoopOp. In particular, there are no loop-carried
+    variables and the number of iterations it will execute is know before the
+    operation. This allows e.g. to determine the number of threads and chunks
+    the iterations space is split into before executing any iteration. More
+    restrictions may apply in cases such as (collapsed) loop nests, doacross
+    loops, etc.
+
+    The induction variable is always of the same type as the tripcount argument.
+    Since it can never be negative, tripcount is always interpreted as an
+    unsigned integer. It is the caller's responsbility to ensure the tripcount
+    is not negative when its interpretation is signed, i.e.
+    `%tripcount = max(0,%tripcount)`.
+
+    In contrast to other loop operations such as `scf.for`, the number of
+    iterations is determined by only a single variable, the trip-count. The
+    induction variable value is the logical iteration number of that iteration,
+    which OpenMP defines to be between 0 and the trip-count (exclusive).
+    Loop representation having lower-bound, upper-bound, and step-size operands,
+    require passes to do more work than necessary, including handling special
+    cases such as upper-bound smaller than lower-bound, upper-bound equal to
+    the integer type's maximal value, negative step size, etc. This complexity
+    is better only handled once by the front-end and can apply its semantics
+    for such cases while still being able to represent any kind of loop, which
+    kind of the point of a mid-end intermediate representation. User-defined
+    types such as random-access iterators in C++ could not directly be
+    represented anyway.
+
+    The return value of a omp.canonical_loop is a CanonicalLoopInfo that can be
+    used to refer to the canonical loop to apply transformations -- such as
+    tiling, unrolling, or work-sharing -- to the loop, similar to the transform
+    dialect but with OpenMP-specific semantics. To refer to nested canonical
+    loops, the CanonicalLoopInfo can be passed to omp.yield and becomes an
+    additional return value of the the outer `omp.canonical_loop`.
+    Every `omp.yield` on the loop body must be passed the same CanonicalLoopInfo
+    since nesting is a static/compile-time property.
+
+    A CanonicalLoopOp can be lowered to LLVM-IR using OpenMPIRBuilder's
+    createCanonicalLoop method.
+
+    #### Examples
+
+    Translation from lower-bound, upper-bount, step-size to trip-count.
+    ```c
+    for (int i = 3; i < 42; i+=2) {
+      B[i] = A[i];
+    }
+    ```
+
+    ```mlir
+    %lb = arith.constant 3 : i32
+    %ub = arith.constant 42 : i32
+    %step = arith.constant 2 : i32
+    %range = arith.sub %ub, %lb : i32
+    %tc = arith.div %range, %step : i32
+    %cli = omp.canonical_loop %iv : i32 in [0, %tc) {
+      %offset = arith.mul %iv, %step : i32
+      %i = arith.add %offset, %lb : i32
+      %a = load %arrA[%i] : memref<?xf32>
+      store %a, %arrB[%i] : memref<?xf32>
+    }
+    ```
+
+    Nested canonical loop with transformation.
+    ```mlir
+    %outer,%inner = omp.canonical_loop %iv1 : i32 in [0, %tripcount) {
+      %inner = omp.canonical_loop %iv2 : i32 in [0, %tc) {
+        %a = load %arrA[%iv1, %iv2] : memref<?x?xf32>
+        store %a, %arrB[%iv1, %iv2] : memref<?x?xf32>
+      }
+      omp.yield(%inner : !omp.cli)
+    }
+    omp.tile(%outer, %inner : !omp.cli, !omp.cli)
+    ```
+
+    Nested canonical loop with other constructs. The `omp.distribute`
+    operation has not been added yet, so this is suggested use with other
+    constructs.
+    ```mlir
+    omp.target {
+      %outer,%inner = omp.canonical_loop %iv1 : i32 in [0, %tripcount) {
+        %inner = omp.canonical_loop %iv2 : i32 in [0, %tc) {
+          %a = load %arrA[%iv1, %iv2] : memref<?x?xf32>
+          store %a, %arrB[%iv1, %iv2] : memref<?x?xf32>
+        }
+        omp.yield(%inner : !omp.cli)
+      }
+      %collapsed_loopinfo = omp.collapse(%outer, %inner)
+      omp.teams {
+        call @foo() : () -> ()
+        omp.distribute(%collapsed_loopinfo)
+      }
+    }
+    ```
+
+  }];
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+
+  let arguments = (ins IntLikeType:$tripCount);
+  let regions = (region AnyRegion:$region);
+  let results = (outs Variadic<CanonicalLoopInfoType>:$loopInfo);
+
+  let extraClassDeclaration = [{
+    ::mlir::Value getInductionVar();
+  }];
+}
+
+
 //===----------------------------------------------------------------------===//
 // 2.9.2 Workshare Loop Construct
 //===----------------------------------------------------------------------===//
@@ -613,7 +762,7 @@ def SimdLoopOp : OpenMP_Op<"simdloop", [AttrSizedOperandSegments,
 def YieldOp : OpenMP_Op<"yield",
     [Pure, ReturnLike, Terminator,
      ParentOneOf<["WsLoopOp", "ReductionDeclareOp",
-     "AtomicUpdateOp", "SimdLoopOp"]>]> {
+     "AtomicUpdateOp", "SimdLoopOp", "CanonicalLoopOp"]>]> {
   let summary = "loop yield and termination operation";
   let description = [{
     "omp.yield" yields SSA values from the OpenMP dialect op region and

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

@@ -1,4 +1,4 @@
-//===- OpenMPDialect.cpp - MLIR Dialect for OpenMP implementation ---------===//
+//===- OpenMPDialect.cpp - MLIR Dialect for OpenMP implementation ---------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -65,6 +65,10 @@ void OpenMPDialect::initialize() {
 #define GET_ATTRDEF_LIST
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
       >();
+  addTypes<
+#define GET_TYPEDEF_LIST
+#include "mlir/Dialect/OpenMP/OpenMPOpsTypes.cpp.inc"
+      >();
 
   addInterface<OpenMPDialectFoldInterface>();
   LLVM::LLVMPointerType::attachInterface<
@@ -1524,8 +1528,123 @@ LogicalResult CancellationPointOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// CanonicaLoopOp
+//===----------------------------------------------------------------------===//
+
+Value mlir::omp::CanonicalLoopOp::getInductionVar() {
+  return getRegion().getArgument(0);
+}
+
+void mlir::omp::CanonicalLoopOp::print(OpAsmPrinter &p) {
+  p << " " << getInductionVar() << " : " << getInductionVar().getType()
+    << " in [0, " << getTripCount() << ") ";
+
+  // omp.yield is implicit if no arguments passed to it.
+  p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
+                /*printBlockTerminators=*/getResultTypes().size() > 1);
+
+  p.printOptionalAttrDict((*this)->getAttrs());
+}
+
+mlir::ParseResult
+mlir::omp::CanonicalLoopOp::parse(::mlir::OpAsmParser &parser,
+                                  ::mlir::OperationState &result) {
+  Builder &builder = parser.getBuilder();
+  MLIRContext *context = parser.getContext();
+
+  // We derive the type of tripCount from inductionVariable. Unfortunatelty we
+  // cannot do the other way around because MLIR requires the type of tripCount
+  // to be known when calling resolveOperand.
+  OpAsmParser::Argument inductionVariable;
+  if (parser.parseArgument(inductionVariable, /*allowType*/ true) ||
+      parser.parseKeyword("in") || parser.parseLSquare())
+    return failure();
+
+  int zero = -1;
+  SMLoc zeroLoc = parser.getCurrentLocation();
+  if (parser.parseInteger(zero))
+    return failure();
+  if (zero != 0) {
+    parser.emitError(zeroLoc, "Logical iteration space starts with zero");
+    return failure();
+  }
+
+  OpAsmParser::UnresolvedOperand tripcount;
+  if (parser.parseComma() || parser.parseOperand(tripcount) ||
+      parser.parseRParen() ||
+      parser.resolveOperand(tripcount, inductionVariable.type, result.operands))
+    return failure();
+
+  // Parse the loop body.
+  Region *region = result.addRegion();
+  if (parser.parseRegion(*region, {inductionVariable}))
+    return failure();
+  CanonicalLoopOp::ensureTerminator(*region, builder, result.location);
+
+  // Return the CanonicalLoopInfo for this loop, plus the CanonicalLoopInfos
+  // passed to omp.yield.
+  int numResults = 0;
+  for (Block &block : *region) {
+    if (auto yield = dyn_cast<YieldOp>(block.getTerminator())) {
+      numResults = yield.getNumOperands();
+      break;
+    }
+  }
+  numResults += 1;
+  for (int i = 0; i < numResults; ++i)
+    result.types.push_back(CanonicalLoopInfoType::get(context));
+
+  // Parse the optional attribute list.
+  if (parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+
+  return mlir::success();
+}
+
+LogicalResult CanonicalLoopOp::verify() {
+  Value indVar = getInductionVar();
+  Value tripCount = getTripCount();
+  Block *body = getBody();
+  Region &region = getRegion();
+
+  if (indVar.getType() != tripCount.getType())
+    return emitOpError(
+        "Region argument must be the same type as the trip count");
+
+  auto numResults = getResultTypes().size();
+  if (numResults <= 0)
+    return emitOpError(
+        "omp.canonical_loop must return at least one CanonicalLoopInfo");
+
+  // Check arguments to omp.yield operations
+  YieldOp *firstYield = nullptr;
+  for (Block &block : region) {
+    if (auto yield = dyn_cast<YieldOp>(block.getTerminator())) {
+      if (yield.getNumOperands() != numResults - 1)
+        return emitOpError("omp.yield arguments must match number of return "
+                           "CanonicalLoopInfo's");
+
+      if (!firstYield) {
+        firstYield = &yield;
+        continue;
+      }
+
+      for (int i = 0; i < numResults - 1; ++i) {
+        if (yield.getOperand(i) != firstYield->getOperand(i))
+          return emitOpError("Each omp.yield must return the same values");
+      }
+    }
+  }
+
+  return success();
+}
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"
+
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/OpenMP/OpenMPOpsTypes.cpp.inc"