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 @@ -116,7 +116,8 @@ // 2.9.2 Workshare Loop Construct //===----------------------------------------------------------------------===// -def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments]> { +def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments, + AllTypesMatch<["lowerBound", "upperBound", "step"]>]> { let summary = "workshare loop construct"; let description = [{ The workshare loop construct specifies that the iterations of the loop(s) @@ -130,13 +131,13 @@ by "omp.yield" instruction without operands. ``` - omp.wsloop (%i1, %i2) = (%c0, %c0) to (%c10, %c10) step (%c1, %c1) { - %a = load %arrA[%i1, %i2] : memref - %b = load %arrB[%i1, %i2] : memref - %sum = addf %a, %b : f32 - store %sum, %arrC[%i1, %i2] : memref - omp.yield - } + omp.wsloop (%i1, %i2) : index = (%c0, %c0) to (%c10, %c10) step (%c1, %c1) { + %a = load %arrA[%i1, %i2] : memref + %b = load %arrB[%i1, %i2] : memref + %sum = addf %a, %b : f32 + store %sum, %arrC[%i1, %i2] : memref + omp.yield + } ``` `private_vars`, `firstprivate_vars`, `lastprivate_vars` and `linear_vars` @@ -181,10 +182,23 @@ OptionalAttr:$order_val, UnitAttr:$inclusive); + let skipDefaultBuilders = 1; + let builders = [ OpBuilder<(ins "ValueRange":$lowerBound, "ValueRange":$upperBound, - "ValueRange":$step, - CArg<"ArrayRef", "{}">:$attributes)> + "ValueRange":$step, + CArg<"ArrayRef", "{}">:$attributes)>, + OpBuilder<(ins "TypeRange":$resultTypes, "ValueRange":$lowerBound, + "ValueRange":$upperBound, "ValueRange":$step, + "ValueRange":$privateVars, "ValueRange":$firstprivateVars, + "ValueRange":$lastprivate_vars, "ValueRange":$linear_vars, + "ValueRange":$linear_step_vars, "StringAttr":$schedule_val, + "Value":$schedule_chunk_var, "IntegerAttr":$collapse_val, + "UnitAttr":$nowait, "IntegerAttr":$ordered_val, + "StringAttr":$order_val, "UnitAttr":$inclusive, CArg<"bool", + "true">:$buildBody)>, + OpBuilder<(ins "TypeRange":$resultTypes, "ValueRange":$operands, + CArg<"ArrayRef", "{}">:$attributes)> ]; let regions = (region AnyRegion:$region); @@ -193,6 +207,8 @@ /// Returns the number of loops in the workshape loop nest. unsigned getNumLoops() { return lowerBound().size(); } }]; + let parser = [{ return parseWsLoopOp(parser, result); }]; + let printer = [{ return printWsLoopOp(p, *this); }]; } def YieldOp : OpenMP_Op<"yield", [NoSideEffect, ReturnLike, Terminator, 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 @@ -17,6 +17,7 @@ #include "mlir/IR/OperationSupport.h" #include "llvm/ADT/SmallString.h" +#include "llvm/ADT/StringExtras.h" #include "llvm/ADT/StringRef.h" #include "llvm/ADT/StringSwitch.h" #include @@ -172,8 +173,8 @@ } /// Emit an error if the same clause is present more than once on an operation. -static ParseResult allowedOnce(OpAsmParser &parser, llvm::StringRef clause, - llvm::StringRef operation) { +static ParseResult allowedOnce(OpAsmParser &parser, StringRef clause, + StringRef operation) { return parser.emitError(parser.getNameLoc()) << " at most one " << clause << " clause can appear on the " << operation << " operation"; @@ -213,7 +214,7 @@ SmallVector allocators; SmallVector allocatorTypes; std::array segments{0, 0, 0, 0, 0, 0, 0, 0}; - llvm::StringRef keyword; + StringRef keyword; bool defaultVal = false; bool procBind = false; @@ -225,11 +226,11 @@ const int copyinClausePos = 5; const int allocateClausePos = 6; const int allocatorPos = 7; - const llvm::StringRef opName = result.name.getStringRef(); + const StringRef opName = result.name.getStringRef(); while (succeeded(parser.parseOptionalKeyword(&keyword))) { if (keyword == "if") { - // Fail if there was already another if condition + // Fail if there was already another if condition. if (segments[ifClausePos]) return allowedOnce(parser, "if", opName); if (parser.parseLParen() || parser.parseOperand(ifCond.first) || @@ -237,7 +238,7 @@ return failure(); segments[ifClausePos] = 1; } else if (keyword == "num_threads") { - // fail if there was already another num_threads clause + // Fail if there was already another num_threads clause. if (segments[numThreadsClausePos]) return allowedOnce(parser, "num_threads", opName); if (parser.parseLParen() || parser.parseOperand(numThreads.first) || @@ -245,35 +246,35 @@ return failure(); segments[numThreadsClausePos] = 1; } else if (keyword == "private") { - // fail if there was already another private clause + // Fail if there was already another private clause. if (segments[privateClausePos]) return allowedOnce(parser, "private", opName); if (parseOperandAndTypeList(parser, privates, privateTypes)) return failure(); segments[privateClausePos] = privates.size(); } else if (keyword == "firstprivate") { - // fail if there was already another firstprivate clause + // Fail if there was already another firstprivate clause. if (segments[firstprivateClausePos]) return allowedOnce(parser, "firstprivate", opName); if (parseOperandAndTypeList(parser, firstprivates, firstprivateTypes)) return failure(); segments[firstprivateClausePos] = firstprivates.size(); } else if (keyword == "shared") { - // fail if there was already another shared clause + // Fail if there was already another shared clause. if (segments[sharedClausePos]) return allowedOnce(parser, "shared", opName); if (parseOperandAndTypeList(parser, shareds, sharedTypes)) return failure(); segments[sharedClausePos] = shareds.size(); } else if (keyword == "copyin") { - // fail if there was already another copyin clause + // Fail if there was already another copyin clause. if (segments[copyinClausePos]) return allowedOnce(parser, "copyin", opName); if (parseOperandAndTypeList(parser, copyins, copyinTypes)) return failure(); segments[copyinClausePos] = copyins.size(); } else if (keyword == "allocate") { - // fail if there was already another allocate clause + // Fail if there was already another allocate clause. if (segments[allocateClausePos]) return allowedOnce(parser, "allocate", opName); if (parseAllocateAndAllocator(parser, allocates, allocateTypes, @@ -282,27 +283,27 @@ segments[allocateClausePos] = allocates.size(); segments[allocatorPos] = allocators.size(); } else if (keyword == "default") { - // fail if there was already another default clause + // Fail if there was already another default clause. if (defaultVal) return allowedOnce(parser, "default", opName); defaultVal = true; - llvm::StringRef defval; + StringRef defval; if (parser.parseLParen() || parser.parseKeyword(&defval) || parser.parseRParen()) return failure(); - llvm::SmallString<16> attrval; + SmallString<16> attrval; // The def prefix is required for the attribute as "private" is a keyword - // in C++ + // in C++. attrval += "def"; attrval += defval; auto attr = parser.getBuilder().getStringAttr(attrval); result.addAttribute("default_val", attr); } else if (keyword == "proc_bind") { - // fail if there was already another proc_bind clause + // Fail if there was already another proc_bind clause. if (procBind) return allowedOnce(parser, "proc_bind", opName); procBind = true; - llvm::StringRef bind; + StringRef bind; if (parser.parseLParen() || parser.parseKeyword(&bind) || parser.parseRParen()) return failure(); @@ -315,48 +316,48 @@ } } - // Add if parameter + // Add if parameter. if (segments[ifClausePos] && parser.resolveOperand(ifCond.first, ifCond.second, result.operands)) return failure(); - // Add num_threads parameter + // Add num_threads parameter. if (segments[numThreadsClausePos] && parser.resolveOperand(numThreads.first, numThreads.second, result.operands)) return failure(); - // Add private parameters + // Add private parameters. if (segments[privateClausePos] && parser.resolveOperands(privates, privateTypes, privates[0].location, result.operands)) return failure(); - // Add firstprivate parameters + // Add firstprivate parameters. if (segments[firstprivateClausePos] && parser.resolveOperands(firstprivates, firstprivateTypes, firstprivates[0].location, result.operands)) return failure(); - // Add shared parameters + // Add shared parameters. if (segments[sharedClausePos] && parser.resolveOperands(shareds, sharedTypes, shareds[0].location, result.operands)) return failure(); - // Add copyin parameters + // Add copyin parameters. if (segments[copyinClausePos] && parser.resolveOperands(copyins, copyinTypes, copyins[0].location, result.operands)) return failure(); - // Add allocate parameters + // Add allocate parameters. if (segments[allocateClausePos] && parser.resolveOperands(allocates, allocateTypes, allocates[0].location, result.operands)) return failure(); - // Add allocator parameters + // Add allocator parameters. if (segments[allocatorPos] && parser.resolveOperands(allocators, allocatorTypes, allocators[0].location, result.operands)) @@ -373,6 +374,335 @@ return success(); } +/// linear ::= `linear` `(` linear-list `)` +/// linear-list := linear-val | linear-val linear-list +/// linear-val := ssa-id-and-type `=` ssa-id-and-type +static ParseResult +parseLinearClause(OpAsmParser &parser, + SmallVectorImpl &vars, + SmallVectorImpl &types, + SmallVectorImpl &stepVars) { + if (parser.parseLParen()) + return failure(); + + do { + OpAsmParser::OperandType var; + Type type; + OpAsmParser::OperandType stepVar; + if (parser.parseOperand(var) || parser.parseEqual() || + parser.parseOperand(stepVar) || parser.parseColonType(type)) + return failure(); + + vars.push_back(var); + types.push_back(type); + stepVars.push_back(stepVar); + } while (succeeded(parser.parseOptionalComma())); + + if (parser.parseRParen()) + return failure(); + + return success(); +} + +/// schedule ::= `schedule` `(` sched-list `)` +/// sched-list ::= sched-val | sched-val sched-list +/// sched-val ::= sched-with-chunk | sched-wo-chunk +/// sched-with-chunk ::= sched-with-chunk-types (`=` ssa-id-and-type)? +/// sched-with-chunk-types ::= `static` | `dynamic` | `guided` +/// sched-wo-chunk ::= `auto` | `runtime` +static ParseResult +parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule, + Optional &chunkSize) { + if (parser.parseLParen()) + return failure(); + + StringRef keyword; + if (parser.parseKeyword(&keyword)) + return failure(); + + schedule = keyword; + if (keyword == "static" || keyword == "dynamic" || keyword == "guided") { + if (succeeded(parser.parseOptionalEqual())) { + chunkSize = OpAsmParser::OperandType{}; + if (parser.parseOperand(*chunkSize)) + return failure(); + } else { + chunkSize = llvm::NoneType::None; + } + } else if (keyword == "auto" || keyword == "runtime") { + chunkSize = llvm::NoneType::None; + } else { + return parser.emitError(parser.getNameLoc()) << " expected schedule kind"; + } + + if (parser.parseRParen()) + return failure(); + + return success(); +} + +/// Parses an OpenMP Workshare Loop operation +/// +/// operation ::= `omp.wsloop` loop-control clause-list +/// loop-control ::= `(` ssa-id-list `)` `:` type `=` loop-bounds +/// loop-bounds := `(` ssa-id-list `)` to `(` ssa-id-list `)` steps +/// steps := `step` `(`ssa-id-list`)` +/// clause-list ::= clause | empty | clause-list +/// clause ::= private | firstprivate | lastprivate | linear | schedule | +// collapse | nowait | ordered | order | inclusive +/// private ::= `private` `(` ssa-id-and-type-list `)` +/// firstprivate ::= `firstprivate` `(` ssa-id-and-type-list `)` +/// lastprivate ::= `lastprivate` `(` ssa-id-and-type-list `)` +/// linear ::= `linear` `(` linear-list `)` +/// schedule ::= `schedule` `(` sched-list `)` +/// collapse ::= `collapse` `(` ssa-id-and-type `)` +/// nowait ::= `nowait` +/// ordered ::= `ordered` `(` ssa-id-and-type `)` +/// order ::= `order` `(` `concurrent` `)` +/// inclusive ::= `inclusive` +/// +static ParseResult parseWsLoopOp(OpAsmParser &parser, OperationState &result) { + Type loopVarType; + int numIVs; + + // Parse an opening `(` followed by induction variables followed by `)` + SmallVector ivs; + if (parser.parseRegionArgumentList(ivs, /*requiredOperandCount=*/-1, + OpAsmParser::Delimiter::Paren)) + return failure(); + + numIVs = static_cast(ivs.size()); + + if (parser.parseColonType(loopVarType)) + return failure(); + + // Parse loop bounds. + SmallVector lower; + if (parser.parseEqual() || + parser.parseOperandList(lower, numIVs, OpAsmParser::Delimiter::Paren) || + parser.resolveOperands(lower, loopVarType, result.operands)) + return failure(); + + SmallVector upper; + if (parser.parseKeyword("to") || + parser.parseOperandList(upper, numIVs, OpAsmParser::Delimiter::Paren) || + parser.resolveOperands(upper, loopVarType, result.operands)) + return failure(); + + // Parse step values. + SmallVector steps; + if (parser.parseKeyword("step") || + parser.parseOperandList(steps, numIVs, OpAsmParser::Delimiter::Paren) || + parser.resolveOperands(steps, loopVarType, result.operands)) + return failure(); + + SmallVector privates; + SmallVector privateTypes; + SmallVector firstprivates; + SmallVector firstprivateTypes; + SmallVector lastprivates; + SmallVector lastprivateTypes; + SmallVector linears; + SmallVector linearTypes; + SmallVector linearSteps; + SmallString<8> schedule; + Optional scheduleChunkSize; + std::array segments{numIVs, numIVs, numIVs, 0, 0, 0, 0, 0, 0}; + + const StringRef opName = result.name.getStringRef(); + StringRef keyword; + + enum SegmentPos { + lbPos = 0, + ubPos, + stepPos, + privateClausePos, + firstprivateClausePos, + lastprivateClausePos, + linearClausePos, + linearStepPos, + scheduleClausePos, + }; + + while (succeeded(parser.parseOptionalKeyword(&keyword))) { + if (keyword == "private") { + if (segments[privateClausePos]) + return allowedOnce(parser, "private", opName); + if (parseOperandAndTypeList(parser, privates, privateTypes)) + return failure(); + segments[privateClausePos] = privates.size(); + } else if (keyword == "firstprivate") { + // fail if there was already another firstprivate clause + if (segments[firstprivateClausePos]) + return allowedOnce(parser, "firstprivate", opName); + if (parseOperandAndTypeList(parser, firstprivates, firstprivateTypes)) + return failure(); + segments[firstprivateClausePos] = firstprivates.size(); + } else if (keyword == "lastprivate") { + // fail if there was already another shared clause + if (segments[lastprivateClausePos]) + return allowedOnce(parser, "lastprivate", opName); + if (parseOperandAndTypeList(parser, lastprivates, lastprivateTypes)) + return failure(); + segments[lastprivateClausePos] = lastprivates.size(); + } else if (keyword == "linear") { + // fail if there was already another linear clause + if (segments[linearClausePos]) + return allowedOnce(parser, "linear", opName); + if (parseLinearClause(parser, linears, linearTypes, linearSteps)) + return failure(); + segments[linearClausePos] = linears.size(); + segments[linearStepPos] = linearSteps.size(); + } else if (keyword == "schedule") { + if (!schedule.empty()) + return allowedOnce(parser, "schedule", opName); + if (parseScheduleClause(parser, schedule, scheduleChunkSize)) + return failure(); + if (scheduleChunkSize) { + segments[scheduleClausePos] = 1; + } + } else if (keyword == "collapse") { + auto type = parser.getBuilder().getI64Type(); + mlir::IntegerAttr attr; + if (parser.parseLParen() || parser.parseAttribute(attr, type) || + parser.parseRParen()) + return failure(); + result.addAttribute("collapse_val", attr); + } else if (keyword == "nowait") { + auto attr = UnitAttr::get(parser.getBuilder().getContext()); + result.addAttribute("nowait", attr); + } else if (keyword == "ordered") { + mlir::IntegerAttr attr; + if (succeeded(parser.parseOptionalLParen())) { + auto type = parser.getBuilder().getI64Type(); + if (parser.parseAttribute(attr, type)) + return failure(); + if (parser.parseRParen()) + return failure(); + } else { + // Use 0 to represent no ordered parameter was specified + attr = parser.getBuilder().getI64IntegerAttr(0); + } + result.addAttribute("ordered_val", attr); + } else if (keyword == "order") { + StringRef order; + if (parser.parseLParen() || parser.parseKeyword(&order) || + parser.parseRParen()) + return failure(); + auto attr = parser.getBuilder().getStringAttr(order); + result.addAttribute("order", attr); + } else if (keyword == "inclusive") { + auto attr = UnitAttr::get(parser.getBuilder().getContext()); + result.addAttribute("inclusive", attr); + } + } + + if (segments[privateClausePos]) { + parser.resolveOperands(privates, privateTypes, privates[0].location, + result.operands); + } + + if (segments[firstprivateClausePos]) { + parser.resolveOperands(firstprivates, firstprivateTypes, + firstprivates[0].location, result.operands); + } + + if (segments[lastprivateClausePos]) { + parser.resolveOperands(lastprivates, lastprivateTypes, + lastprivates[0].location, result.operands); + } + + if (segments[linearClausePos]) { + parser.resolveOperands(linears, linearTypes, linears[0].location, + result.operands); + auto linearStepType = parser.getBuilder().getI32Type(); + SmallVector linearStepTypes(linearSteps.size(), linearStepType); + parser.resolveOperands(linearSteps, linearStepTypes, + linearSteps[0].location, result.operands); + } + + if (!schedule.empty()) { + schedule[0] = llvm::toUpper(schedule[0]); + auto attr = parser.getBuilder().getStringAttr(schedule); + result.addAttribute("schedule_val", attr); + if (scheduleChunkSize) { + auto chunkSizeType = parser.getBuilder().getI32Type(); + parser.resolveOperand(*scheduleChunkSize, chunkSizeType, result.operands); + } + } + + result.addAttribute("operand_segment_sizes", + parser.getBuilder().getI32VectorAttr(segments)); + + // Now parse the body. + Region *body = result.addRegion(); + SmallVector ivTypes(numIVs, loopVarType); + if (parser.parseRegion(*body, ivs, ivTypes)) + return failure(); + return success(); +} + +static void printWsLoopOp(OpAsmPrinter &p, WsLoopOp op) { + auto args = op.getRegion().front().getArguments(); + p << op.getOperationName() << " (" << args << ") : " << args[0].getType() + << " = (" << op.lowerBound() << ") to (" << op.upperBound() << ") step (" + << op.step() << ")"; + + // Print private, firstprivate, shared and copyin parameters + auto printDataVars = [&p](StringRef name, OperandRange vars) { + if (vars.empty()) + return; + + p << " " << name << "("; + llvm::interleaveComma( + vars, p, [&](const Value &v) { p << v << " : " << v.getType(); }); + p << ")"; + }; + printDataVars("private", op.private_vars()); + printDataVars("firstprivate", op.firstprivate_vars()); + printDataVars("lastprivate", op.lastprivate_vars()); + + auto linearVars = op.linear_vars(); + auto linearVarsSize = linearVars.size(); + if (linearVarsSize) { + p << " " + << "linear" + << "("; + for (unsigned i = 0; i < linearVarsSize; ++i) { + std::string separator = i == linearVarsSize - 1 ? ")" : ", "; + p << linearVars[i]; + if (op.linear_step_vars().size() > i) + p << " = " << op.linear_step_vars()[i]; + p << " : " << linearVars[i].getType() << separator; + } + } + + if (auto sched = op.schedule_val()) { + auto schedLower = sched->lower(); + p << " schedule(" << schedLower; + if (auto chunk = op.schedule_chunk_var()) { + p << " = " << chunk; + } + p << ")"; + } + + if (auto collapse = op.collapse_val()) + p << " collapse(" << collapse << ")"; + + if (op.nowait()) + p << " nowait"; + + if (auto ordered = op.ordered_val()) { + p << " ordered(" << ordered << ")"; + } + + if (op.inclusive()) { + p << " inclusive"; + } + + p.printRegion(op.region(), /*printEntryBlockArgs=*/false); +} + //===----------------------------------------------------------------------===// // WsLoopOp //===----------------------------------------------------------------------===// @@ -386,9 +716,71 @@ /*linear_vars=*/ValueRange(), /*linear_step_vars=*/ValueRange(), /*schedule_val=*/nullptr, /*schedule_chunk_var=*/nullptr, /*collapse_val=*/nullptr, - /*nowait=*/false, /*ordered_val=*/nullptr, /*order_val=*/nullptr, - /*inclusive=*/false); + /*nowait=*/nullptr, /*ordered_val=*/nullptr, /*order_val=*/nullptr, + /*inclusive=*/nullptr, /*buildBody=*/false); + state.addAttributes(attributes); +} + +void WsLoopOp::build(OpBuilder &, OperationState &state, TypeRange resultTypes, + ValueRange operands, ArrayRef attributes) { + state.addOperands(operands); state.addAttributes(attributes); + (void)state.addRegion(); + assert(resultTypes.size() == 0u && "mismatched number of return types"); + state.addTypes(resultTypes); +} + +void WsLoopOp::build(OpBuilder &builder, OperationState &result, + TypeRange typeRange, ValueRange lowerBounds, + ValueRange upperBounds, ValueRange steps, + ValueRange privateVars, ValueRange firstprivateVars, + ValueRange lastprivateVars, ValueRange linearVars, + ValueRange linearStepVars, StringAttr scheduleVal, + Value scheduleChunkVar, IntegerAttr collapseVal, + UnitAttr nowait, IntegerAttr orderedVal, + StringAttr orderVal, UnitAttr inclusive, bool buildBody) { + result.addOperands(lowerBounds); + result.addOperands(upperBounds); + result.addOperands(steps); + result.addOperands(privateVars); + result.addOperands(firstprivateVars); + result.addOperands(linearVars); + result.addOperands(linearStepVars); + if (scheduleChunkVar) + result.addOperands(scheduleChunkVar); + + if (scheduleVal) + result.addAttribute("schedule_val", scheduleVal); + if (collapseVal) + result.addAttribute("collapse_val", collapseVal); + if (nowait) + result.addAttribute("nowait", nowait); + if (orderedVal) + result.addAttribute("ordered_val", orderedVal); + if (orderVal) + result.addAttribute("order", orderVal); + if (inclusive) + result.addAttribute("inclusive", inclusive); + result.addAttribute( + WsLoopOp::getOperandSegmentSizeAttr(), + builder.getI32VectorAttr( + {static_cast(lowerBounds.size()), + static_cast(upperBounds.size()), + static_cast(steps.size()), + static_cast(privateVars.size()), + static_cast(firstprivateVars.size()), + static_cast(lastprivateVars.size()), + static_cast(linearVars.size()), + static_cast(linearStepVars.size()), + static_cast(scheduleChunkVar != nullptr ? 1 : 0)})); + + Region *bodyRegion = result.addRegion(); + if (buildBody) { + OpBuilder::InsertionGuard guard(builder); + unsigned numIVs = steps.size(); + SmallVector argTypes(numIVs, steps.getType().front()); + builder.createBlock(bodyRegion, {}, argTypes); + } } #define GET_OP_CLASSES diff --git a/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir b/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir --- a/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir +++ b/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir @@ -34,10 +34,8 @@ func @wsloop(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) { // CHECK: omp.parallel omp.parallel { - // CHECK: omp.wsloop - // CHECK: (%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG5]]) + // CHECK: omp.wsloop (%[[ARG6:.*]], %[[ARG7:.*]]) : i64 = (%[[ARG0]], %[[ARG1]]) to (%[[ARG2]], %[[ARG3]]) step (%[[ARG4]], %[[ARG5]]) { "omp.wsloop"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5) ( { - // CHECK: ^{{.*}}(%[[ARG6:.*]]: i64, %[[ARG7:.*]]: i64): ^bb0(%arg6: index, %arg7: index): // no predecessors // CHECK: "test.payload"(%[[ARG6]], %[[ARG7]]) : (i64, i64) -> () "test.payload"(%arg6, %arg7) : (index, index) -> () diff --git a/mlir/test/Conversion/SCFToOpenMP/scf-to-openmp.mlir b/mlir/test/Conversion/SCFToOpenMP/scf-to-openmp.mlir --- a/mlir/test/Conversion/SCFToOpenMP/scf-to-openmp.mlir +++ b/mlir/test/Conversion/SCFToOpenMP/scf-to-openmp.mlir @@ -4,9 +4,9 @@ func @parallel(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) { // CHECK: omp.parallel { - // CHECK: "omp.wsloop"({{.*}}) ( { + // CHECK: omp.wsloop (%[[LVAR1:.*]], %[[LVAR2:.*]]) : index = (%arg0, %arg1) to (%arg2, %arg3) step (%arg4, %arg5) { scf.parallel (%i, %j) = (%arg0, %arg1) to (%arg2, %arg3) step (%arg4, %arg5) { - // CHECK: test.payload + // CHECK: "test.payload"(%[[LVAR1]], %[[LVAR2]]) : (index, index) -> () "test.payload"(%i, %j) : (index, index) -> () // CHECK: omp.yield // CHECK: } @@ -20,12 +20,12 @@ func @nested_loops(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) { // CHECK: omp.parallel { - // CHECK: "omp.wsloop"({{.*}}) ( { + // CHECK: omp.wsloop (%[[LVAR_OUT1:.*]]) : index = (%arg0) to (%arg2) step (%arg4) { // CHECK-NOT: omp.parallel scf.parallel (%i) = (%arg0) to (%arg2) step (%arg4) { - // CHECK: "omp.wsloop"({{.*}}) ( { + // CHECK: omp.wsloop (%[[LVAR_IN1:.*]]) : index = (%arg1) to (%arg3) step (%arg5) { scf.parallel (%j) = (%arg1) to (%arg3) step (%arg5) { - // CHECK: test.payload + // CHECK: "test.payload"(%[[LVAR_OUT1]], %[[LVAR_IN1]]) : (index, index) -> () "test.payload"(%i, %j) : (index, index) -> () // CHECK: omp.yield // CHECK: } @@ -41,9 +41,9 @@ func @adjacent_loops(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) { // CHECK: omp.parallel { - // CHECK: "omp.wsloop"({{.*}}) ( { + // CHECK: omp.wsloop (%[[LVAR_AL1:.*]]) : index = (%arg0) to (%arg2) step (%arg4) { scf.parallel (%i) = (%arg0) to (%arg2) step (%arg4) { - // CHECK: test.payload1 + // CHECK: "test.payload1"(%[[LVAR_AL1]]) : (index) -> () "test.payload1"(%i) : (index) -> () // CHECK: omp.yield // CHECK: } @@ -52,9 +52,9 @@ // CHECK: } // CHECK: omp.parallel { - // CHECK: "omp.wsloop"({{.*}}) ( { + // CHECK: omp.wsloop (%[[LVAR_AL2:.*]]) : index = (%arg1) to (%arg3) step (%arg5) { scf.parallel (%j) = (%arg1) to (%arg3) step (%arg5) { - // CHECK: test.payload2 + // CHECK: "test.payload2"(%[[LVAR_AL2]]) : (index) -> () "test.payload2"(%j) : (index) -> () // CHECK: omp.yield // CHECK: } diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -89,77 +89,192 @@ } func @omp_parallel_pretty(%data_var : memref, %if_cond : i1, %num_threads : si32, %allocator : si32) -> () { - // CHECK: omp.parallel - omp.parallel { - omp.terminator - } + // CHECK: omp.parallel + omp.parallel { + omp.terminator + } + + // CHECK: omp.parallel num_threads(%{{.*}} : si32) + omp.parallel num_threads(%num_threads : si32) { + omp.terminator + } + + // CHECK: omp.parallel allocate(%{{.*}} : memref -> %{{.*}} : memref) + omp.parallel allocate(%data_var : memref -> %data_var : memref) { + omp.terminator + } + + // CHECK: omp.parallel private(%{{.*}} : memref, %{{.*}} : memref) firstprivate(%{{.*}} : memref) + omp.parallel private(%data_var : memref, %data_var : memref) firstprivate(%data_var : memref) { + omp.terminator + } + + // CHECK omp.parallel shared(%{{.*}} : memref) copyin(%{{.*}} : memref, %{{.*}} : memref) + omp.parallel shared(%data_var : memref) copyin(%data_var : memref, %data_var : memref) { + omp.parallel if(%if_cond: i1) { + omp.terminator + } + omp.terminator + } + + // CHECK omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref) proc_bind(close) + omp.parallel num_threads(%num_threads : si32) if(%if_cond: i1) + private(%data_var : memref) proc_bind(close) { + omp.terminator + } + + return +} - // CHECK: omp.parallel num_threads(%{{.*}} : si32) - omp.parallel num_threads(%num_threads : si32) { - omp.terminator - } +// CHECK-LABEL: omp_wsloop +func @omp_wsloop(%lb : index, %ub : index, %step : index, %data_var : memref, %linear_var : i32, %chunk_var : i32) -> () { - // CHECK: omp.parallel allocate(%{{.*}} : memref -> %{{.*}} : memref) - omp.parallel allocate(%data_var : memref -> %data_var : memref) { - omp.terminator + // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref, %{{.*}} : memref) collapse(2) ordered(1) + "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var) ({ + ^bb0(%iv: index): + omp.yield + }) {operand_segment_sizes = dense<[1,1,1,2,0,0,0,0,0]> : vector<9xi32>, collapse_val = 2, ordered_val = 1} : + (index, index, index, memref, memref) -> () + + // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref) schedule(static) + "omp.wsloop" (%lb, %ub, %step, %data_var, %linear_var) ({ + ^bb0(%iv: index): + omp.yield + }) {operand_segment_sizes = dense<[1,1,1,0,0,0,1,1,0]> : vector<9xi32>, schedule_val = "Static"} : + (index, index, index, memref, i32) -> () + + // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref, %{{.*}} = %{{.*}} : memref) schedule(static) + "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var, %linear_var, %linear_var) ({ + ^bb0(%iv: index): + omp.yield + }) {operand_segment_sizes = dense<[1,1,1,0,0,0,2,2,0]> : vector<9xi32>, schedule_val = "Static"} : + (index, index, index, memref, memref, i32, i32) -> () + + // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref) firstprivate(%{{.*}} : memref) lastprivate(%{{.*}} : memref) linear(%{{.*}} = %{{.*}} : memref) schedule(dynamic = %{{.*}}) collapse(3) ordered(2) + "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var, %data_var, %data_var, %linear_var, %chunk_var) ({ + ^bb0(%iv: index): + omp.yield + }) {operand_segment_sizes = dense<[1,1,1,1,1,1,1,1,1]> : vector<9xi32>, schedule_val = "Dynamic", collapse_val = 3, ordered_val = 2} : + (index, index, index, memref, memref, memref, memref, i32, i32) -> () + + // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref) schedule(auto) nowait + "omp.wsloop" (%lb, %ub, %step, %data_var) ({ + ^bb0(%iv: index): + omp.yield + }) {operand_segment_sizes = dense<[1,1,1,1,0,0,0,0,0]> : vector<9xi32>, nowait, schedule_val = "Auto"} : + (index, index, index, memref) -> () + + return +} + +// CHECK-LABEL: omp_wsloop_pretty +func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index, + %data_var : memref, %linear_var : i32, %chunk_var : i32) -> () { + + // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref) + omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) private(%data_var : memref) collapse(2) ordered(2) { + omp.yield } - // CHECK: omp.parallel private(%{{.*}} : memref, %{{.*}} : memref) firstprivate(%{{.*}} : memref) - omp.parallel private(%data_var : memref, %data_var : memref) firstprivate(%data_var : memref) { - omp.terminator + // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref) schedule(static) + omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) schedule(static) lastprivate(%data_var : memref) linear(%data_var = %linear_var : memref) { + omp.yield } - // CHECK omp.parallel shared(%{{.*}} : memref) copyin(%{{.*}} : memref, %{{.*}} : memref) - omp.parallel shared(%data_var : memref) copyin(%data_var : memref, %data_var : memref) { - omp.parallel if(%if_cond: i1) { - omp.terminator - } - omp.terminator + // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref) firstprivate(%{{.*}} : memref) lastprivate(%{{.*}} : memref) linear(%{{.*}} = %{{.*}} : memref) schedule(static = %{{.*}}) collapse(3) ordered(2) + omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) ordered(2) private(%data_var : memref) + firstprivate(%data_var : memref) lastprivate(%data_var : memref) linear(%data_var = %linear_var : memref) + schedule(static = %chunk_var) collapse(3) { + omp.yield } - // CHECK omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref) proc_bind(close) - omp.parallel num_threads(%num_threads : si32) if(%if_cond: i1) - private(%data_var : memref) proc_bind(close) { - omp.terminator + // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private({{.*}} : memref) + omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) private(%data_var : memref) { + omp.yield } return } -func @omp_wsloop(%lb : index, %ub : index, %step : index, - %data_var : memref, %linear_var : si32, %chunk_var : si32) -> () { +// CHECK-LABEL: omp_wsloop_pretty_multi_block +func @omp_wsloop_pretty_multi_block(%lb : index, %ub : index, %step : index, %data1 : memref, %data2 : memref) -> () { - // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) - "omp.wsloop" (%lb, %ub, %step, %data_var) ({ + // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) + omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) { + %1 = "test.payload"(%iv) : (index) -> (i32) + br ^bb1(%1: i32) + ^bb1(%arg: i32): + memref.store %arg, %data1[%iv] : memref omp.yield - }) {operand_segment_sizes = dense<[1,1,1,1,0,0,0,0,0]> : vector<9xi32>, collapse_val = 2, ordered_val = 1} : - (index, index, index, memref) -> () + } - // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) - "omp.wsloop" (%lb, %lb, %ub, %ub, %step, %step, %data_var) ({ + // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) + omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) { + %c = "test.condition"(%iv) : (index) -> (i1) + %v1 = "test.payload"(%iv) : (index) -> (i32) + cond_br %c, ^bb1(%v1: i32), ^bb2(%v1: i32) + ^bb1(%arg0: i32): + memref.store %arg0, %data1[%iv] : memref + br ^bb3 + ^bb2(%arg1: i32): + memref.store %arg1, %data2[%iv] : memref + br ^bb3 + ^bb3: omp.yield - }) {operand_segment_sizes = dense<[2,2,2,1,0,0,0,0,0]> : vector<9xi32>, collapse_val = 2, ordered_val = 1} : - (index, index, index, index, index, index, memref) -> () - + } - // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) - "omp.wsloop" (%lb, %ub, %step, %data_var, %linear_var) ({ + // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) + omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) { + %c = "test.condition"(%iv) : (index) -> (i1) + %v1 = "test.payload"(%iv) : (index) -> (i32) + cond_br %c, ^bb1(%v1: i32), ^bb2(%v1: i32) + ^bb1(%arg0: i32): + memref.store %arg0, %data1[%iv] : memref omp.yield - }) {operand_segment_sizes = dense<[1,1,1,0,0,0,1,1,0]> : vector<9xi32>, schedule_val = "Static"} : - (index, index, index, memref, si32) -> () + ^bb2(%arg1: i32): + memref.store %arg1, %data2[%iv] : memref + omp.yield + } - // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) - "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var, %data_var, %data_var, %linear_var, %chunk_var) ({ + return +} + +// CHECK-LABEL: omp_wsloop_pretty_non_index +func @omp_wsloop_pretty_non_index(%lb1 : i32, %ub1 : i32, %step1 : i32, %lb2 : i64, %ub2 : i64, %step2 : i64, + %data1 : memref, %data2 : memref) -> () { + + // CHECK: omp.wsloop (%{{.*}}) : i32 = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) + omp.wsloop (%iv1) : i32 = (%lb1) to (%ub1) step (%step1) { + %1 = "test.payload"(%iv1) : (i32) -> (index) + br ^bb1(%1: index) + ^bb1(%arg1: index): + memref.store %iv1, %data1[%arg1] : memref omp.yield - }) {operand_segment_sizes = dense<[1,1,1,1,1,1,1,1,1]> : vector<9xi32>, schedule_val = "Dynamic", collapse_val = 3, ordered_val = 2} : - (index, index, index, memref, memref, memref, memref, si32, si32) -> () + } - // CHECK: "omp.wsloop"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) - "omp.wsloop" (%lb, %ub, %step, %data_var) ({ + // CHECK: omp.wsloop (%{{.*}}) : i64 = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) + omp.wsloop (%iv2) : i64 = (%lb2) to (%ub2) step (%step2) { + %2 = "test.payload"(%iv2) : (i64) -> (index) + br ^bb1(%2: index) + ^bb1(%arg2: index): + memref.store %iv2, %data2[%arg2] : memref omp.yield - }) {operand_segment_sizes = dense<[1,1,1,1,0,0,0,0,0]> : vector<9xi32>, nowait, schedule_val = "Auto"} : - (index, index, index, memref) -> () + } + + return +} +// CHECK-LABEL: omp_wsloop_pretty_multiple +func @omp_wsloop_pretty_multiple(%lb1 : i32, %ub1 : i32, %step1 : i32, %lb2 : i32, %ub2 : i32, %step2 : i32, %data1 : memref) -> () { + + // CHECK: omp.wsloop (%{{.*}}, %{{.*}}) : i32 = (%{{.*}}, %{{.*}}) to (%{{.*}}, %{{.*}}) step (%{{.*}}, %{{.*}}) + omp.wsloop (%iv1, %iv2) : i32 = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) { + %1 = "test.payload"(%iv1) : (i32) -> (index) + %2 = "test.payload"(%iv2) : (i32) -> (index) + memref.store %iv1, %data1[%1] : memref + memref.store %iv2, %data1[%2] : memref + omp.yield + } return }