Skip to content

Commit

Permalink
[MLIR][OpenMP] Make omp.wsloop into a loop wrapper (1/5) (#89209)
Browse files Browse the repository at this point in the history 
This patch updates the definition of `omp.wsloop` to enforce the
restrictions of a loop wrapper operation.

Related tests are updated but this PR on its own will not pass premerge
tests. All patches in the stack are needed before it can be compiled and
passes tests.
  • Loading branch information
@skatrak
skatrak authored Apr 24, 2024
1 parent bd34bc6 commit 07e6c16
Show file tree
Hide file tree
Showing 8 changed files with 628 additions and 521 deletions.
7 changes: 3 additions & 4 deletions mlir/include/mlir/Dialect/OpenMP/OpenMPClauseOperands.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -295,10 +295,9 @@ using TeamsClauseOps =
PrivateClauseOps, ReductionClauseOps, ThreadLimitClauseOps>;

using WsloopClauseOps =
detail::Clauses<AllocateClauseOps, CollapseClauseOps, LinearClauseOps,
LoopRelatedOps, NowaitClauseOps, OrderClauseOps,
OrderedClauseOps, PrivateClauseOps, ReductionClauseOps,
ScheduleClauseOps>;
detail::Clauses<AllocateClauseOps, LinearClauseOps, NowaitClauseOps,
OrderClauseOps, OrderedClauseOps, PrivateClauseOps,
ReductionClauseOps, ScheduleClauseOps>;

} // namespace omp
} // namespace mlir
Expand Down
62 changes: 28 additions & 34 deletions mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -600,29 +600,30 @@ def LoopNestOp : OpenMP_Op<"loop_nest", [SameVariadicOperandSize,
//===----------------------------------------------------------------------===//

def WsloopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
AllTypesMatch<["lowerBound", "upperBound", "step"]>,
DeclareOpInterfaceMethods<LoopWrapperInterface>,
RecursiveMemoryEffects, ReductionClauseInterface]> {
RecursiveMemoryEffects, ReductionClauseInterface,
SingleBlockImplicitTerminator<"TerminatorOp">]> {
let summary = "worksharing-loop construct";
let description = [{
The worksharing-loop construct specifies that the iterations of the loop(s)
will be executed in parallel by threads in the current context. These
iterations are spread across threads that already exist in the enclosing
parallel region. 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.
parallel region.

The body region can contain any number of blocks. The region is terminated
by "omp.yield" instruction without operands.
The body region can only contain a single block which must contain a single
operation and a terminator. The operation must be another compatible loop
wrapper or an `omp.loop_nest`.

```
omp.wsloop <clauses>
for (%i1, %i2) : index = (%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
omp.wsloop <clauses> {
omp.loop_nest (%i1, %i2) : index = (%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
}
omp.terminator
}
```

Expand Down Expand Up @@ -665,10 +666,7 @@ def WsloopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
passed by reference.
}];

let arguments = (ins Variadic<IntLikeType>:$lowerBound,
Variadic<IntLikeType>:$upperBound,
Variadic<IntLikeType>:$step,
Variadic<AnyType>:$linear_vars,
let arguments = (ins Variadic<AnyType>:$linear_vars,
Variadic<I32>:$linear_step_vars,
Variadic<OpenMP_PointerLikeType>:$reduction_vars,
OptionalAttr<SymbolRefArrayAttr>:$reductions,
Expand All @@ -679,22 +677,16 @@ def WsloopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
UnitAttr:$nowait,
UnitAttr:$byref,
ConfinedAttr<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$ordered_val,
OptionalAttr<OrderKindAttr>:$order_val,
UnitAttr:$inclusive);
OptionalAttr<OrderKindAttr>:$order_val);

let builders = [
OpBuilder<(ins "ValueRange":$lowerBound, "ValueRange":$upperBound,
"ValueRange":$step,
CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
OpBuilder<(ins CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
OpBuilder<(ins CArg<"const WsloopClauseOps &">:$clauses)>
];

let regions = (region AnyRegion:$region);

let extraClassDeclaration = [{
/// Returns the number of loops in the worksharing-loop nest.
unsigned getNumLoops() { return getLowerBound().size(); }

/// Returns the number of reduction variables.
unsigned getNumReductionVars() { return getReductionVars().size(); }
}];
Expand All @@ -711,9 +703,8 @@ def WsloopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
|`byref` $byref
|`ordered` `(` $ordered_val `)`
|`order` `(` custom<ClauseAttr>($order_val) `)`
) custom<Wsloop>($region, $lowerBound, $upperBound, $step, type($step),
$reduction_vars, type($reduction_vars), $reductions,
$inclusive) attr-dict
) custom<Wsloop>($region, $reduction_vars, type($reduction_vars),
$reductions) attr-dict
}];
let hasVerifier = 1;
}
Expand All @@ -732,7 +723,7 @@ def SimdOp : OpenMP_Op<"simd", [AttrSizedOperandSegments,
transformed into a SIMD loop (that is, multiple iterations of the loop can
be executed concurrently using SIMD instructions).

The body region can contain a single block which must contain a single
The body region can only contain a single block which must contain a single
operation and a terminator. The operation must be another compatible loop
wrapper or an `omp.loop_nest`.

Expand Down Expand Up @@ -766,6 +757,7 @@ def SimdOp : OpenMP_Op<"simd", [AttrSizedOperandSegments,
store %sum, %arrC[%i1, %i2] : memref<?x?xf32>
omp.yield
}
omp.terminator
}
```
}];
Expand Down Expand Up @@ -805,8 +797,8 @@ def SimdOp : OpenMP_Op<"simd", [AttrSizedOperandSegments,

def YieldOp : OpenMP_Op<"yield",
[Pure, ReturnLike, Terminator,
ParentOneOf<["LoopNestOp", "WsloopOp", "DeclareReductionOp",
"AtomicUpdateOp", "PrivateClauseOp"]>]> {
ParentOneOf<["AtomicUpdateOp", "DeclareReductionOp", "LoopNestOp",
"PrivateClauseOp"]>]> {
let summary = "loop yield and termination operation";
let description = [{
"omp.yield" yields SSA values from the OpenMP dialect op region and
Expand Down Expand Up @@ -846,7 +838,7 @@ def DistributeOp : OpenMP_Op<"distribute", [AttrSizedOperandSegments,
iterations are spread across threads that already exist in the enclosing
region.

The body region can contain a single block which must contain a single
The body region can only contain a single block which must contain a single
operation and a terminator. The operation must be another compatible loop
wrapper or an `omp.loop_nest`.

Expand All @@ -864,6 +856,7 @@ def DistributeOp : OpenMP_Op<"distribute", [AttrSizedOperandSegments,
store %sum, %arrC[%i1, %i2] : memref<?x?xf32>
omp.yield
}
omp.terminator
}
```
// TODO: private_var, firstprivate_var, lastprivate_var, collapse
Expand Down Expand Up @@ -1029,7 +1022,7 @@ def TaskloopOp : OpenMP_Op<"taskloop", [AttrSizedOperandSegments,
iterations are distributed across tasks generated by the construct and
scheduled to be executed.

The body region can contain a single block which must contain a single
The body region can only contain a single block which must contain a single
operation and a terminator. The operation must be another compatible loop
wrapper or an `omp.loop_nest`.

Expand All @@ -1042,6 +1035,7 @@ def TaskloopOp : OpenMP_Op<"taskloop", [AttrSizedOperandSegments,
store %sum, %arrC[%i1, %i2] : memref<?x?xf32>
omp.yield
}
omp.terminator
}
```

Expand Down
150 changes: 50 additions & 100 deletions mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1484,86 +1484,72 @@ LogicalResult SingleOp::verify() {
// WsloopOp
//===----------------------------------------------------------------------===//

/// loop-control ::= `(` ssa-id-list `)` `:` type `=` loop-bounds
/// loop-bounds := `(` ssa-id-list `)` to `(` ssa-id-list `)` inclusive? steps
/// steps := `step` `(`ssa-id-list`)`
ParseResult
parseWsloop(OpAsmParser &parser, Region &region,
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &lowerBound,
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &upperBound,
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &steps,
SmallVectorImpl<Type> &loopVarTypes,
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &reductionOperands,
SmallVectorImpl<Type> &reductionTypes, ArrayAttr &reductionSymbols,
UnitAttr &inclusive) {

SmallVectorImpl<Type> &reductionTypes,
ArrayAttr &reductionSymbols) {
// Parse an optional reduction clause
llvm::SmallVector<OpAsmParser::Argument> privates;
bool hasReduction = succeeded(parser.parseOptionalKeyword("reduction")) &&
succeeded(parseClauseWithRegionArgs(
parser, region, reductionOperands, reductionTypes,
reductionSymbols, privates));

if (parser.parseKeyword("for"))
return failure();

// Parse an opening `(` followed by induction variables followed by `)`
SmallVector<OpAsmParser::Argument> ivs;
Type loopVarType;
if (parser.parseArgumentList(ivs, OpAsmParser::Delimiter::Paren) ||
parser.parseColonType(loopVarType) ||
// Parse loop bounds.
parser.parseEqual() ||
parser.parseOperandList(lowerBound, ivs.size(),
OpAsmParser::Delimiter::Paren) ||
parser.parseKeyword("to") ||
parser.parseOperandList(upperBound, ivs.size(),
OpAsmParser::Delimiter::Paren))
return failure();

if (succeeded(parser.parseOptionalKeyword("inclusive")))
inclusive = UnitAttr::get(parser.getBuilder().getContext());

// Parse step values.
if (parser.parseKeyword("step") ||
parser.parseOperandList(steps, ivs.size(), OpAsmParser::Delimiter::Paren))
return failure();

// Now parse the body.
loopVarTypes = SmallVector<Type>(ivs.size(), loopVarType);
for (auto &iv : ivs)
iv.type = loopVarType;

SmallVector<OpAsmParser::Argument> regionArgs{ivs};
if (hasReduction)
llvm::copy(privates, std::back_inserter(regionArgs));

return parser.parseRegion(region, regionArgs);
if (succeeded(parser.parseOptionalKeyword("reduction"))) {
if (failed(parseClauseWithRegionArgs(parser, region, reductionOperands,
reductionTypes, reductionSymbols,
privates)))
return failure();
}
return parser.parseRegion(region, privates);
}

void printWsloop(OpAsmPrinter &p, Operation *op, Region &region,
ValueRange lowerBound, ValueRange upperBound, ValueRange steps,
TypeRange loopVarTypes, ValueRange reductionOperands,
TypeRange reductionTypes, ArrayAttr reductionSymbols,
UnitAttr inclusive) {
ValueRange reductionOperands, TypeRange reductionTypes,
ArrayAttr reductionSymbols) {
if (reductionSymbols) {
auto reductionArgs =
region.front().getArguments().drop_front(loopVarTypes.size());
auto reductionArgs = region.front().getArguments();
printClauseWithRegionArgs(p, op, reductionArgs, "reduction",
reductionOperands, reductionTypes,
reductionSymbols);
}

p << " for ";
auto args = region.front().getArguments().drop_back(reductionOperands.size());
p << " (" << args << ") : " << args[0].getType() << " = (" << lowerBound
<< ") to (" << upperBound << ") ";
if (inclusive)
p << "inclusive ";
p << "step (" << steps << ") ";
p.printRegion(region, /*printEntryBlockArgs=*/false);
}

void WsloopOp::build(OpBuilder &builder, OperationState &state,
ArrayRef<NamedAttribute> attributes) {
build(builder, state, /*linear_vars=*/ValueRange(),
/*linear_step_vars=*/ValueRange(), /*reduction_vars=*/ValueRange(),
/*reductions=*/nullptr, /*schedule_val=*/nullptr,
/*schedule_chunk_var=*/nullptr, /*schedule_modifier=*/nullptr,
/*simd_modifier=*/false, /*nowait=*/false, /*byref=*/false,
/*ordered_val=*/nullptr, /*order_val=*/nullptr);
state.addAttributes(attributes);
}

void WsloopOp::build(OpBuilder &builder, OperationState &state,
const WsloopClauseOps &clauses) {
MLIRContext *ctx = builder.getContext();
// TODO: Store clauses in op: allocateVars, allocatorVars, privateVars,
// privatizers.
WsloopOp::build(
builder, state, clauses.linearVars, clauses.linearStepVars,
clauses.reductionVars, makeArrayAttr(ctx, clauses.reductionDeclSymbols),
clauses.scheduleValAttr, clauses.scheduleChunkVar,
clauses.scheduleModAttr, clauses.scheduleSimdAttr, clauses.nowaitAttr,
clauses.reductionByRefAttr, clauses.orderedAttr, clauses.orderAttr);
}

LogicalResult WsloopOp::verify() {
if (!isWrapper())
return emitOpError() << "must be a loop wrapper";

if (LoopWrapperInterface nested = getNestedWrapper()) {
// Check for the allowed leaf constructs that may appear in a composite
// construct directly after DO/FOR.
if (!isa<SimdOp>(nested))
return emitError() << "only supported nested wrapper is 'omp.simd'";
}

return verifyReductionVarList(*this, getReductions(), getReductionVars());
}

//===----------------------------------------------------------------------===//
// Simd construct [2.9.3.1]
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -1947,42 +1933,6 @@ void LoopNestOp::gatherWrappers(
}
}

//===----------------------------------------------------------------------===//
// WsloopOp
//===----------------------------------------------------------------------===//

void WsloopOp::build(OpBuilder &builder, OperationState &state,
ValueRange lowerBound, ValueRange upperBound,
ValueRange step, ArrayRef<NamedAttribute> attributes) {
build(builder, state, lowerBound, upperBound, step,
/*linear_vars=*/ValueRange(),
/*linear_step_vars=*/ValueRange(), /*reduction_vars=*/ValueRange(),
/*reductions=*/nullptr, /*schedule_val=*/nullptr,
/*schedule_chunk_var=*/nullptr, /*schedule_modifier=*/nullptr,
/*simd_modifier=*/false, /*nowait=*/false, /*byref=*/false,
/*ordered_val=*/nullptr,
/*order_val=*/nullptr, /*inclusive=*/false);
state.addAttributes(attributes);
}

void WsloopOp::build(OpBuilder &builder, OperationState &state,
const WsloopClauseOps &clauses) {
MLIRContext *ctx = builder.getContext();
// TODO Store clauses in op: allocateVars, allocatorVars, privateVars,
// privatizers.
WsloopOp::build(
builder, state, clauses.loopLBVar, clauses.loopUBVar, clauses.loopStepVar,
clauses.linearVars, clauses.linearStepVars, clauses.reductionVars,
makeArrayAttr(ctx, clauses.reductionDeclSymbols), clauses.scheduleValAttr,
clauses.scheduleChunkVar, clauses.scheduleModAttr,
clauses.scheduleSimdAttr, clauses.nowaitAttr, clauses.reductionByRefAttr,
clauses.orderedAttr, clauses.orderAttr, clauses.loopInclusiveAttr);
}

LogicalResult WsloopOp::verify() {
return verifyReductionVarList(*this, getReductions(), getReductionVars());
}

//===----------------------------------------------------------------------===//
// Critical construct (2.17.1)
//===----------------------------------------------------------------------===//
Expand Down
7 changes: 5 additions & 2 deletions mlir/test/CAPI/execution_engine.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,8 +99,11 @@ void testOmpCreation(void) {
" %1 = arith.constant 1 : i32 \n"
" %2 = arith.constant 2 : i32 \n"
" omp.parallel { \n"
" omp.wsloop for (%3) : i32 = (%0) to (%2) step (%1) { \n"
" omp.yield \n"
" omp.wsloop { \n"
" omp.loop_nest (%3) : i32 = (%0) to (%2) step (%1) { \n"
" omp.yield \n"
" } \n"
" omp.terminator \n"
" } \n"
" omp.terminator \n"
" } \n"
Expand Down
Loading

0 comments on commit 07e6c16

Please sign in to comment.