[MLIR][OpenMP] Add omp.loopnest operation

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

@@ -506,6 +506,66 @@ def SingleOp : OpenMP_Op<"single", [AttrSizedOperandSegments]> {
  let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// Loop Nest
+//===----------------------------------------------------------------------===//
+
+def LoopNestOp : OpenMP_Op<"loopnest", [SameVariadicOperandSize,
+ AllTypesMatch<["lowerBound", "upperBound", "step"]>,
+ ParentOneOf<["DistributeOp", "SimdLoopOp", "TaskLoopOp",
+ "WsLoopOp"]>,
+ RecursiveMemoryEffects]> {
+ let summary = "rectangular loop nest";
+ let description = [{
+ This operation represents a collapsed rectangular loop nest. For each
+ rectangular loop of the nest represented by an instance of this operation,
+ lower and upper bounds, as well as a step variable, must be defined.
+
+ The lower and upper bounds specify a half-open range: the range includes the
+ lower bound but does not include the upper bound. If the `inclusive`
+ attribute is specified then the upper bound is also included.
+
+ The body region can contain any number of blocks. The region is terminated
+ by "omp.yield" instruction without operands. The induction variables,
+ represented as entry block arguments to the loop nest operation's single
+ region, match the types of the `lowerBound`, `upperBound` and `step`
+ arguments.
+
+ ```mlir
+ omp.loopnest (%i1, %i2) : i32 = (%c0, %c0) to (%c10, %c10) step (%c1, %c1) {
+ %a = load %arrA[%i1, %i2] : memref<?x?xf32>
+ %b = load %arrB[%i1, %i2] : memref<?x?xf32>
+ %sum = arith.addf %a, %b : f32
+ store %sum, %arrC[%i1, %i2] : memref<?x?xf32>
+ omp.yield
+ }
+ ```
+
+ This is a temporary simplified definition of a loop based on existing OpenMP
+ loop operations intended to serve as a stopgap solution until the long-term
+ representation of canonical loops is defined. Specifically, this approach is
+ not intended to help with the addition of support for loop transformations.
+ }];
+
+ let arguments = (ins Variadic<IntLikeType>:$lowerBound,
+ Variadic<IntLikeType>:$upperBound,
+ Variadic<IntLikeType>:$step,
+ UnitAttr:$inclusive);
+
+ let regions = (region AnyRegion:$region);
+
+ let extraClassDeclaration = [{
+ /// Returns the number of loops in the loop nest.
+ unsigned getNumLoops() { return getLowerBound().size(); }
+
+ /// Returns the induction variables of the loop nest.
+ ArrayRef<BlockArgument> getIVs() { return getRegion().getArguments(); }
+ }];
+
+ let hasCustomAssemblyFormat = 1;
+ let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // 2.9.2 Workshare Loop Construct
 //===----------------------------------------------------------------------===//
@@ -714,7 +774,7 @@ def SimdLoopOp : OpenMP_Op<"simdloop", [AttrSizedOperandSegments,
 
 def YieldOp : OpenMP_Op<"yield",
  [Pure, ReturnLike, Terminator,
- ParentOneOf<["WsLoopOp", "ReductionDeclareOp",
+ ParentOneOf<["LoopNestOp", "WsLoopOp", "ReductionDeclareOp",
  "AtomicUpdateOp", "SimdLoopOp", "PrivateClauseOp"]>]> {
  let summary = "loop yield and termination operation";
  let description = [{

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

@@ -1487,6 +1487,77 @@ LogicalResult SingleOp::verify() {
  getCopyprivateFuncs());
 }
 
+//===----------------------------------------------------------------------===//
+// LoopNestOp
+//===----------------------------------------------------------------------===//
+
+ParseResult LoopNestOp::parse(OpAsmParser &parser, OperationState &result) {
+ // Parse an opening `(` followed by induction variables followed by `)`
+ SmallVector<OpAsmParser::Argument> ivs;
+ SmallVector<OpAsmParser::UnresolvedOperand> lbs, ubs;
+ Type loopVarType;
+ if (parser.parseArgumentList(ivs, OpAsmParser::Delimiter::Paren) ||
+ parser.parseColonType(loopVarType) ||
+ // Parse loop bounds.
+ parser.parseEqual() ||
+ parser.parseOperandList(lbs, ivs.size(), OpAsmParser::Delimiter::Paren) ||
+ parser.parseKeyword("to") ||
+ parser.parseOperandList(ubs, ivs.size(), OpAsmParser::Delimiter::Paren))
+ return failure();
+
+ for (auto &iv : ivs)
+ iv.type = loopVarType;
+
+ // Parse "inclusive" flag.
+ if (succeeded(parser.parseOptionalKeyword("inclusive")))
+ result.addAttribute("inclusive",
+ UnitAttr::get(parser.getBuilder().getContext()));
+
+ // Parse step values.
+ SmallVector<OpAsmParser::UnresolvedOperand> steps;
+ if (parser.parseKeyword("step") ||
+ parser.parseOperandList(steps, ivs.size(), OpAsmParser::Delimiter::Paren))
+ return failure();
+
+ // Parse the body.
+ Region *region = result.addRegion();
+ if (parser.parseRegion(*region, ivs))
+ return failure();
+
+ // Resolve operands.
+ if (parser.resolveOperands(lbs, loopVarType, result.operands) ||
+ parser.resolveOperands(ubs, loopVarType, result.operands) ||
+ parser.resolveOperands(steps, loopVarType, result.operands))
+ return failure();
+
+ // Parse the optional attribute list.
+ return parser.parseOptionalAttrDict(result.attributes);
+}
+
+void LoopNestOp::print(OpAsmPrinter &p) {
+ Region &region = getRegion();
+ auto args = region.getArguments();
+ p << " (" << args << ") : " << args[0].getType() << " = (" << getLowerBound()
+ << ") to (" << getUpperBound() << ") ";
+ if (getInclusive())
+ p << "inclusive ";
+ p << "step (" << getStep() << ") ";
+ p.printRegion(region, /*printEntryBlockArgs=*/false);
+}
+
+LogicalResult LoopNestOp::verify() {
+ if (getLowerBound().size() != getIVs().size())
+ return emitOpError() << "number of range arguments and IVs do not match";
+
+ for (auto [lb, iv] : llvm::zip_equal(getLowerBound(), getIVs())) {
+ if (lb.getType() != iv.getType())
+ return emitOpError()
+ << "range argument type does not match corresponding IV type";
+ }
+
+ return success();
+}
+
 //===----------------------------------------------------------------------===//
 // WsLoopOp
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/OpenMP/invalid.mlir

@@ -87,6 +87,43 @@ func.func @proc_bind_once() {
 
 // -----
 
+func.func @invalid_parent(%lb : index, %ub : index, %step : index) {
+ // expected-error@+1 {{op expects parent op to be one of 'omp.distribute, omp.simdloop, omp.taskloop, omp.wsloop'}}
+ omp.loopnest (%iv) : index = (%lb) to (%ub) step (%step) {
+ omp.yield
+ }
+}
+
+// -----
+
+func.func @type_mismatch(%lb : index, %ub : index, %step : index) {
+ // TODO Remove induction variables from omp.wsloop.
+ omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+ // expected-error@+1 {{range argument type does not match corresponding IV type}}
+ "omp.loopnest" (%lb, %ub, %step) ({
+ ^bb0(%iv2: i32):
+ omp.yield
+ }) : (index, index, index) -> ()
+ omp.yield
+ }
+}
+
+// -----
+
+func.func @iv_number_mismatch(%lb : index, %ub : index, %step : index) {
+ // TODO Remove induction variables from omp.wsloop.
+ omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+ // expected-error@+1 {{number of range arguments and IVs do not match}}
+ "omp.loopnest" (%lb, %ub, %step) ({
+ ^bb0(%iv1 : index, %iv2 : index):
+ omp.yield
+ }) : (index, index, index) -> ()
+ omp.yield
+ }
+}
+
+// -----
+
 func.func @inclusive_not_a_clause(%lb : index, %ub : index, %step : index) {
  // expected-error @below {{expected 'for'}}
  omp.wsloop nowait inclusive

mlir/test/Dialect/OpenMP/ops.mlir

@@ -133,6 +133,87 @@ func.func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_thre
  return
 }
 
+// CHECK-LABEL: omp_loopnest
+func.func @omp_loopnest(%lb : index, %ub : index, %step : index) -> () {
+
+ // TODO Remove induction variables from omp.wsloop.
+ omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+ // CHECK: omp.loopnest
+ // CHECK-SAME: (%{{.*}}) : index =
+ // CHECK-SAME: (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+ "omp.loopnest" (%lb, %ub, %step) ({
+ ^bb0(%iv2: index):
+ omp.yield
+ }) : (index, index, index) -> ()
+ omp.yield
+ }
+
+ // TODO Remove induction variables from omp.wsloop.
+ omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+ // CHECK: omp.loopnest
+ // CHECK-SAME: (%{{.*}}) : index =
+ // CHECK-SAME: (%{{.*}}) to (%{{.*}}) inclusive step (%{{.*}})
+ "omp.loopnest" (%lb, %ub, %step) ({
+ ^bb0(%iv2: index):
+ omp.yield
+ }) {inclusive} : (index, index, index) -> ()
+ omp.yield
+ }
+
+ // TODO Remove induction variables from omp.wsloop.
+ omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+ // CHECK: omp.loopnest
+ // CHECK-SAME: (%{{.*}}, %{{.*}}) : index =
+ // CHECK-SAME: (%{{.*}}, %{{.*}}) to (%{{.*}}, %{{.*}}) step (%{{.*}}, %{{.*}})
+ "omp.loopnest" (%lb, %lb, %ub, %ub, %step, %step) ({
+ ^bb0(%iv2: index, %iv3: index):
+ omp.yield
+ }) : (index, index, index, index, index, index) -> ()
+ omp.yield
+ }
+
+ return
+}
+
+// CHECK-LABEL: omp_loopnest_pretty
+func.func @omp_loopnest_pretty(%lb : index, %ub : index, %step : index) -> () {
+
+ // TODO Remove induction variables from omp.wsloop.
+ omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+ // CHECK: omp.loopnest
+ // CHECK-SAME: (%{{.*}}) : index =
+ // CHECK-SAME: (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+ omp.loopnest (%iv2) : index = (%lb) to (%ub) step (%step) {
+ omp.yield
+ }
+ omp.yield
+ }
+
+ // TODO Remove induction variables from omp.wsloop.
+ omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+ // CHECK: omp.loopnest
+ // CHECK-SAME: (%{{.*}}) : index =
+ // CHECK-SAME: (%{{.*}}) to (%{{.*}}) inclusive step (%{{.*}})
+ omp.loopnest (%iv2) : index = (%lb) to (%ub) inclusive step (%step) {
+ omp.yield
+ }
+ omp.yield
+ }
+
+ // TODO Remove induction variables from omp.wsloop.
+ omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+ // CHECK: omp.loopnest
+ // CHECK-SAME: (%{{.*}}) : index =
+ // CHECK-SAME: (%{{.*}}, %{{.*}}) to (%{{.*}}, %{{.*}}) step (%{{.*}}, %{{.*}})
+ omp.loopnest (%iv2, %iv3) : index = (%lb, %lb) to (%ub, %ub) step (%step, %step) {
+ omp.yield
+ }
+ omp.yield
+ }
+
+ return
+}
+
 // CHECK-LABEL: omp_wsloop
 func.func @omp_wsloop(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32) -> () {