diff --git a/mlir/include/mlir/Dialect/OpenMP/CMakeLists.txt b/mlir/include/mlir/Dialect/OpenMP/CMakeLists.txt --- a/mlir/include/mlir/Dialect/OpenMP/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/OpenMP/CMakeLists.txt @@ -2,18 +2,10 @@ mlir_tablegen(OmpCommon.td --gen-directive-decl --directives-dialect=OpenMP) add_public_tablegen_target(omp_common_td) -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_dialect(OpenMPOps 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) diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h b/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h +++ b/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" diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -28,6 +28,7 @@ let cppNamespace = "::mlir::omp"; let dependentDialects = ["::mlir::LLVM::LLVMDialect, ::mlir::func::FuncDialect"]; let useDefaultAttributePrinterParser = 1; + let useDefaultTypePrinterParser = 1; let usePropertiesForAttributes = 1; } @@ -79,6 +80,10 @@ let assemblyFormat = "`<` struct(params) `>`"; } +class OpenMP_Type : + TypeDef { + let mnemonic = typeMnemonic; +} class OpenMP_Op traits = []> : Op; @@ -396,6 +401,129 @@ 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 = [{ + A CanonicalLoopOp represents a loop that is conforms to OpenMP's defintion + of a canonical loop. 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, incliding 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 + store %a, %arrB[%i] : memref + } + ``` + + 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 + store %a, %arrB[%iv1, %iv2] : memref + } + omp.yield(%inner : !omp.cli) + } + omp.tile(%outer, %inner : !omp.cli, !omp.cli) + ``` + }]; + let hasCustomAssemblyFormat = 1; + let hasVerifier = 1; + + let arguments = (ins IntLikeType:$tripCount); + let regions = (region AnyRegion:$region); + let results = (outs Variadic:$loopInfo); + + let extraClassDeclaration = [{ + ::mlir::Value getInductionVar(); + }]; +} + + //===----------------------------------------------------------------------===// // 2.9.2 Workshare Loop Construct //===----------------------------------------------------------------------===// @@ -610,7 +738,7 @@ 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 diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp --- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp +++ b/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 @@ #define GET_ATTRDEF_LIST #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc" >(); + addTypes< +#define GET_TYPEDEF_LIST +#include "mlir/Dialect/OpenMP/OpenMPOpsTypes.cpp.inc" + >(); addInterface(); LLVM::LLVMPointerType::attachInterface< @@ -1508,8 +1512,123 @@ 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(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 ®ion = 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(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" diff --git a/mlir/test/Dialect/OpenMP/cli.mlir b/mlir/test/Dialect/OpenMP/cli.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/OpenMP/cli.mlir @@ -0,0 +1,99 @@ +// RUN: mlir-opt %s | mlir-opt | FileCheck %s + +// CHECK-LABEL: @omp_canonloop_raw +func.func @omp_canonloop_raw(%tc : i32) -> () { + // CHECK: %0 = omp.canonical_loop %arg1 : i32 in [0, %arg0) { + %cli = "omp.canonical_loop" (%tc) ({ + ^bb0(%iv: i32): + // omp.yield without argument is implicit + // CHECK-NOT: omp.yield + omp.yield + }) : (i32) -> (!omp.cli) + return +} + + +// CHECK-LABEL: @omp_nested_canonloop_raw +func.func @omp_nested_canonloop_raw(%tc_outer : i32, %tc_inner : i32) -> () { + // CHECK: %0:2 = omp.canonical_loop %arg2 : i32 in [0, %arg0) { + %outer,%inner = "omp.canonical_loop" (%tc_outer) ({ + ^bb_outer(%iv_outer: i32): + // CHECK: %1 = omp.canonical_loop %arg3 : i32 in [0, %arg1) { + %inner = "omp.canonical_loop" (%tc_inner) ({ + ^bb_inner(%iv_inner: i32): + omp.yield + }) : (i32) -> (!omp.cli) + // CHECK: omp.yield(%1 : !omp.cli) + omp.yield (%inner : !omp.cli) + }) : (i32) -> (!omp.cli, !omp.cli) + return +} + + +// CHECK-LABEL: @omp_triple_nested_canonloop_raw +func.func @omp_triple_nested_canonloop_raw(%tc_outer : i32,%tc_middle : i32, %tc_inner : i32) -> () { + // CHECK: %0:3 = omp.canonical_loop %arg3 : i32 in [0, %arg0) { + %outer,%middle,%inner = "omp.canonical_loop" (%tc_outer) ({ + ^bb_outer(%iv_outer: i32): + // CHECK: %1:2 = omp.canonical_loop %arg4 : i32 in [0, %arg1) { + %middle, %inner= "omp.canonical_loop" (%tc_middle) ({ + ^bb_middle(%iv_middle: i32): + // CHECK: %2 = omp.canonical_loop %arg5 : i32 in [0, %arg2) { + %inner = "omp.canonical_loop" (%tc_inner) ({ + ^bb_inner(%iv_inner: i32): + omp.yield + }) : (i32) -> (!omp.cli) + // CHECK: omp.yield(%2 : !omp.cli) + omp.yield (%inner : !omp.cli) + }) : (i32) -> (!omp.cli,!omp.cli) +// CHECK: omp.yield(%1#0, %1#1 : !omp.cli, !omp.cli) + omp.yield (%middle , %inner : !omp.cli , !omp.cli) + }) : (i32) -> (!omp.cli, !omp.cli, !omp.cli) + + return +} + + +// CHECK-LABEL: @omp_canonloop_pretty +func.func @omp_canonloop_pretty(%tc : i32) -> () { + + // CHECK: %0 = omp.canonical_loop %arg1 : i32 in [0, %arg0) { + %cli = omp.canonical_loop %iv : i32 in [0, %tc) { + %newval = llvm.add %iv, %iv: i32 +// CHECK-NOT: omp.yield + } + + return +} + + +// CHECK-LABEL: @omp_canonloop_implicit_yield +func.func @omp_canonloop_implicit_yield(%tc : i32) -> () { + + // CHECK: %0 = omp.canonical_loop %arg1 : i32 in [0, %arg0) { + %cli = omp.canonical_loop %iv : i32 in [0, %tc) { +// CHECK-NOT: omp.yield + } + + return +} + + + + + +// CHECK-LABEL: @omp_canonloop_nested_pretty +func.func @omp_canonloop_nested_pretty(%tc : i32) -> () { + + // CHECK: %0:2 = omp.canonical_loop %arg1 : i32 in [0, %arg0) { + %outer,%inner = omp.canonical_loop %iv1 : i32 in [0, %tc) { + // CHECK: %1 = omp.canonical_loop %arg2 : i32 in [0, %arg0) { + %inner = omp.canonical_loop %iv2 : i32 in [0, %tc) { + } + // CHECK: omp.yield(%1 : !omp.cli) + omp.yield (%inner : !omp.cli) + } + + return +} +