diff --git a/mlir/include/mlir/Dialect/AffineOps/AffineOps.td b/mlir/include/mlir/Dialect/AffineOps/AffineOps.td --- a/mlir/include/mlir/Dialect/AffineOps/AffineOps.td +++ b/mlir/include/mlir/Dialect/AffineOps/AffineOps.td @@ -271,6 +271,76 @@ }]; } +def AffineParallelOp : Affine_Op<"parallel", [ImplicitAffineTerminator]> { + let summary = "multi-index parallel band operation"; + let description = [{ + The "affine.parallel" operation represents a hyper-rectangular affine + parallel band, defining multiple SSA values for its induction variables. It + has one region capturing the parallel band body. The induction variables are + represented as arguments of this region. These SSA values always have type + index, which is the size of the machine word. The strides, represented by + steps, are positive constant integers which defaults to "1" if not present. + The lower and upper bounds specify a half-open range: the range includes the + lower bound but does not include the upper bound. + + The body region must contain exactly one block that terminates with + "affine.terminator". Calling AffineParallelOp::build will create such + region and insert the terminator, so will the parsing even in cases if it is + absent from the custom format. + + The lower and upper bounds of a parallel operation are represented as an + application of an affine mapping to a list of SSA values passed to the map. + The same restrictions hold for these SSA values as for all bindings of SSA + values to dimensions and symbols. + + Example: + + affine.parallel [%i, %j] = [0, 0] to [10, 10] step [1, 1] { + ... + } + }]; + + let arguments = (ins + AffineMapAttr:$lowerBoundsMap, + AffineMapAttr:$upperBoundsMap, + I64ArrayAttr:$steps, + Variadic:$mapOperands); + let regions = (region SizedRegion<1>:$region); + + let builders = [ + OpBuilder<"Builder* builder, OperationState& result," + "ArrayRef ranges">, + OpBuilder<"Builder* builder, OperationState& result, AffineMap lbMap," + "ValueRange lbArgs, AffineMap ubMap, ValueRange ubArgs">, + OpBuilder<"Builder* builder, OperationState& result, AffineMap lbMap," + "ValueRange lbArgs, AffineMap ubMap, ValueRange ubArgs," + "ArrayRef steps"> + ]; + + let extraClassDeclaration = [{ + /// Get the number of dimensions. + unsigned getNumDims(); + + operand_range getLowerBoundsOperands(); + operand_range getUpperBoundsOperands(); + + AffineValueMap getLowerBoundsValueMap(); + AffineValueMap getUpperBoundsValueMap(); + AffineValueMap getRangesValueMap(); + + /// Get ranges as constants, may fail in dynamic case. + llvm::Optional> getConstantRanges(); + + mlir::Block *getBody(); + mlir::OpBuilder getBodyBuilder(); + void setSteps(ArrayRef newSteps); + + static StringRef getLowerBoundsMapAttrName() { return "lowerBoundsMap"; } + static StringRef getUpperBoundsMapAttrName() { return "upperBoundsMap"; } + static StringRef getStepsAttrName() { return "steps"; } + }]; +} + def AffinePrefetchOp : Affine_Op<"prefetch"> { let summary = "affine prefetch operation"; let description = [{ diff --git a/mlir/lib/Dialect/AffineOps/AffineOps.cpp b/mlir/lib/Dialect/AffineOps/AffineOps.cpp --- a/mlir/lib/Dialect/AffineOps/AffineOps.cpp +++ b/mlir/lib/Dialect/AffineOps/AffineOps.cpp @@ -135,9 +135,11 @@ return isTopLevelValue(dimOp.getOperand()); return false; } - // This value has to be a block argument for a FuncOp or an affine.for. + // This value has to be a block argument of a FuncOp, an 'affine.for', or an + // 'affine.parallel'. auto *parentOp = value.cast().getOwner()->getParentOp(); - return isa(parentOp) || isa(parentOp); + return isa(parentOp) || isa(parentOp) || + isa(parentOp); } /// Returns true if the 'index' dimension of the `memref` defined by @@ -2151,6 +2153,230 @@ return foldMemRefCast(*this); } +//===----------------------------------------------------------------------===// +// AffineParallelOp +//===----------------------------------------------------------------------===// + +void AffineParallelOp::build(Builder *builder, OperationState &result, + ArrayRef ranges) { + // Default initalize empty maps + auto lbMap = AffineMap::get(builder->getContext()); + auto ubMap = AffineMap::get(builder->getContext()); + // If ranges, set to [0, N) for each range + if (ranges.size()) { + SmallVector lbExprs; + SmallVector ubExprs; + // Make range expressions for each range + for (int64_t range : ranges) { + lbExprs.push_back(builder->getAffineConstantExpr(0)); + ubExprs.push_back(builder->getAffineConstantExpr(range)); + } + lbMap = AffineMap::get(0, 0, lbExprs); + ubMap = AffineMap::get(0, 0, ubExprs); + } + // Fall through + build(builder, result, lbMap, {}, ubMap, {}); +} + +void AffineParallelOp::build(Builder *builder, OperationState &result, + AffineMap lbMap, ValueRange lbArgs, + AffineMap ubMap, ValueRange ubArgs) { + // Verify sizes + auto dims = lbMap.getNumResults(); + assert(dims == ubMap.getNumResults()); + // Make default step sizes of 1 + SmallVector steps(dims, 1); + // Call through + build(builder, result, lbMap, lbArgs, ubMap, ubArgs, steps); +} + +void AffineParallelOp::build(Builder *builder, OperationState &result, + AffineMap lbMap, ValueRange lbArgs, + AffineMap ubMap, ValueRange ubArgs, + ArrayRef steps) { + // Verify sizes + auto dims = lbMap.getNumResults(); + assert(dims == ubMap.getNumResults()); + assert(dims == steps.size()); + // Set all of the attributes + result.addAttribute(getLowerBoundsMapAttrName(), AffineMapAttr::get(lbMap)); + result.addAttribute(getUpperBoundsMapAttrName(), AffineMapAttr::get(ubMap)); + result.addAttribute(getStepsAttrName(), builder->getI64ArrayAttr(steps)); + result.addOperands(lbArgs); + result.addOperands(ubArgs); + // Create a region and a block for the body. + auto bodyRegion = result.addRegion(); + auto body = new Block(); + // Add all the args + for (unsigned i = 0; i < dims; ++i) { + body->addArgument(IndexType::get(builder->getContext())); + } + bodyRegion->push_back(body); + // Terminate + ensureTerminator(*bodyRegion, *builder, result.location); +} + +AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() { + return {operand_begin(), operand_begin() + lowerBoundsMap().getNumInputs()}; +} + +AffineParallelOp::operand_range AffineParallelOp::getUpperBoundsOperands() { + return {operand_begin() + lowerBoundsMap().getNumInputs(), operand_end()}; +} + +unsigned AffineParallelOp::getNumDims() { return steps().size(); } + +AffineValueMap AffineParallelOp::getLowerBoundsValueMap() { + return AffineValueMap(lowerBoundsMap(), getLowerBoundsOperands()); +} + +AffineValueMap AffineParallelOp::getUpperBoundsValueMap() { + return AffineValueMap(upperBoundsMap(), getUpperBoundsOperands()); +} + +AffineValueMap AffineParallelOp::getRangesValueMap() { + AffineValueMap out; + AffineValueMap::difference(getUpperBoundsValueMap(), getLowerBoundsValueMap(), + &out); + return out; +} + +Optional> AffineParallelOp::getConstantRanges() { + SmallVector out; + // Get the ranges + AffineValueMap rangesValueMap = getRangesValueMap(); + // Try to convert them to constants + for (unsigned i = 0, e = rangesValueMap.getNumResults(); i < e; ++i) { + auto expr = rangesValueMap.getResult(i); + if (auto cst = expr.dyn_cast()) { + out.push_back(cst.getValue()); + } else { + return llvm::None; + } + } + return out; +} + +mlir::Block *AffineParallelOp::getBody() { return ®ion().front(); } + +mlir::OpBuilder AffineParallelOp::getBodyBuilder() { + return mlir::OpBuilder(getBody(), std::prev(getBody()->end())); +} + +void AffineParallelOp::setSteps(ArrayRef newSteps) { + assert(newSteps.size() == getNumDims()); + setAttr("steps", getBodyBuilder().getI64ArrayAttr(newSteps)); +} + +static LogicalResult verify(AffineParallelOp op) { + auto dims = op.getNumDims(); + if (op.lowerBoundsMap().getNumResults() != dims || + op.upperBoundsMap().getNumResults() != dims || + op.steps().size() != dims || op.getBody()->getNumArguments() != dims) { + return op.emitOpError("region argument count and num results of upper " + "bounds, lower bounds, and steps must all match"); + } + return success(); +} + +static void print(OpAsmPrinter &p, AffineParallelOp op) { + p << op.getOperationName() << " ["; + p.printOperands(op.getBody()->getArguments()); + p << "] = ["; + p.printAffineMapOfSSAIds(op.lowerBoundsMapAttr(), + op.getLowerBoundsOperands()); + p << "] to ["; + p.printAffineMapOfSSAIds(op.upperBoundsMapAttr(), + op.getUpperBoundsOperands()); + p << ']'; + bool elideSteps = true; + for (auto attr : op.steps()) { + auto step = attr.cast().getValue(); + if (step != 1) { + elideSteps = false; + break; + } + } + if (!elideSteps) { + p << " step " << op.steps(); + } + p.printRegion(op.region(), /*printEntryBlockArgs=*/false, + /*printBlockTerminators=*/false); + p.printOptionalAttrDict( + op.getAttrs(), + /*elidedAttrs=*/{AffineParallelOp::getLowerBoundsMapAttrName(), + AffineParallelOp::getUpperBoundsMapAttrName(), + AffineParallelOp::getStepsAttrName()}); +} + +// +// operation ::= `affine.parallel` `[` ssa-ids `]` `=` `[` map-of-ssa-ids `]` +// `to` `[` map-of-ssa-ids `]` steps? region attr-dict? +// steps ::= `steps` `[` integer-literals `]` +// +static ParseResult parseAffineParallelOp(OpAsmParser &parser, + OperationState &result) { + auto &builder = parser.getBuilder(); + auto indexType = builder.getIndexType(); + AffineMapAttr lowerBoundsAttr, upperBoundsAttr; + SmallVector ivs, lowerBoundsMapOperands, + upperBoundsMapOperands; + if (parser.parseRegionArgumentList(ivs, /*requiredOperandCount=*/-1, + OpAsmParser::Delimiter::Square) || + parser.parseEqual() || + parser.parseAffineMapOfSSAIds( + lowerBoundsMapOperands, lowerBoundsAttr, + AffineParallelOp::getLowerBoundsMapAttrName(), result.attributes) || + parser.resolveOperands(lowerBoundsMapOperands, indexType, + result.operands) || + parser.parseKeyword("to") || + parser.parseAffineMapOfSSAIds( + upperBoundsMapOperands, upperBoundsAttr, + AffineParallelOp::getUpperBoundsMapAttrName(), result.attributes) || + parser.resolveOperands(upperBoundsMapOperands, indexType, + result.operands)) + return failure(); + + AffineMapAttr stepsMapAttr; + SmallVector stepsAttrs; + SmallVector stepsMapOperands; + if (parser.parseOptionalKeyword("step")) { + SmallVector steps(ivs.size(), 1); + result.addAttribute(AffineParallelOp::getStepsAttrName(), + builder.getI64ArrayAttr(steps)); + } else { + if (parser.parseAffineMapOfSSAIds(stepsMapOperands, stepsMapAttr, + AffineParallelOp::getStepsAttrName(), + stepsAttrs)) + return failure(); + + // Convert steps from an AffineMap into an I64ArrayAttr. + SmallVector steps; + auto stepsMap = stepsMapAttr.getValue(); + for (const auto &result : stepsMap.getResults()) { + auto constExpr = result.dyn_cast(); + if (!constExpr) + return parser.emitError(parser.getNameLoc(), + "steps must be constant integers"); + steps.push_back(constExpr.getValue()); + } + result.addAttribute(AffineParallelOp::getStepsAttrName(), + builder.getI64ArrayAttr(steps)); + } + + // Now parse the body. + Region *body = result.addRegion(); + SmallVector types(ivs.size(), indexType); + if (parser.parseRegion(*body, ivs, types) || + parser.parseOptionalAttrDict(result.attributes)) + return failure(); + + // Add a terminator if none was parsed. + AffineParallelOp::ensureTerminator(*body, builder, result.location); + + return success(); +} + //===----------------------------------------------------------------------===// // TableGen'd op method definitions //===----------------------------------------------------------------------===// diff --git a/mlir/test/AffineOps/invalid.mlir b/mlir/test/AffineOps/invalid.mlir --- a/mlir/test/AffineOps/invalid.mlir +++ b/mlir/test/AffineOps/invalid.mlir @@ -197,4 +197,40 @@ %0 = affine.max affine_map<(d0) -> (d0)> () return -} \ No newline at end of file +} + +// ----- + +// CHECK-LABEL: @affine_parallel +func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { + // expected-error@+1 {{region argument count and num results of upper bounds, lower bounds, and steps must all match}} + affine.parallel [%i0] = [0, 0] to [100, 100] step [10, 10] { + } +} + +// ----- + +// CHECK-LABEL: @affine_parallel +func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { + // expected-error@+1 {{region argument count and num results of upper bounds, lower bounds, and steps must all match}} + affine.parallel [%i0, %j0] = [0] to [100, 100] step [10, 10] { + } +} + +// ----- + +// CHECK-LABEL: @affine_parallel +func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { + // expected-error@+1 {{region argument count and num results of upper bounds, lower bounds, and steps must all match}} + affine.parallel [%i0, %j0] = [0, 0] to [100] step [10, 10] { + } +} + +// ----- + +// CHECK-LABEL: @affine_parallel +func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { + // expected-error@+1 {{region argument count and num results of upper bounds, lower bounds, and steps must all match}} + affine.parallel [%i0, %j0] = [0, 0] to [100, 100] step [10] { + } +} diff --git a/mlir/test/AffineOps/ops.mlir b/mlir/test/AffineOps/ops.mlir --- a/mlir/test/AffineOps/ops.mlir +++ b/mlir/test/AffineOps/ops.mlir @@ -112,3 +112,17 @@ } return } + +// ----- + +// CHECK-LABEL: @parallel +// CHECK-SAME: (%[[N:.*]]: index) +func @parallel(%N : index) { + // CHECK: affine.parallel [%[[I0:.*]], %[[J0:.*]]] = [0, 0] to [symbol(%[[N]]), 100] step [10, 10] + affine.parallel [%i0, %j0] = [0, 0] to [symbol(%N), 100] step [10, 10] { + // CHECK-NEXT: affine.parallel [%{{.*}}, %{{.*}}] = [%[[I0]], %[[J0]]] to [%[[I0]] + 10, %[[J0]] + 10] + affine.parallel [%i1, %j1] = [%i0, %j0] to [%i0 + 10, %j0 + 10] { + } + } + return +}