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,79 @@ }]; } +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". + + 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. + + Note: Calling AffineParallelOp::build will create the required region and + block, and insert the required terminator. Parsing will also create the + required region, block, and terminator, even when they are missing from the + textual representation. + + Example: + + ```mlir + 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. + Optional> getConstantRanges(); + + Block *getBody(); + 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/include/mlir/IR/OpImplementation.h b/mlir/include/mlir/IR/OpImplementation.h --- a/mlir/include/mlir/IR/OpImplementation.h +++ b/mlir/include/mlir/IR/OpImplementation.h @@ -105,7 +105,8 @@ if (types.begin() != types.end()) printArrowTypeList(types); } - template void printArrowTypeList(TypeRange &&types) { + template + void printArrowTypeList(TypeRange &&types) { auto &os = getStream() << " -> "; bool wrapped = !has_single_element(types) || @@ -517,7 +518,8 @@ virtual ParseResult parseAffineMapOfSSAIds(SmallVectorImpl &operands, Attribute &map, StringRef attrName, - SmallVectorImpl &attrs) = 0; + SmallVectorImpl &attrs, + Delimiter delimiter = Delimiter::Square) = 0; //===--------------------------------------------------------------------===// // Region Parsing @@ -579,7 +581,8 @@ virtual ParseResult parseType(Type &result) = 0; /// Parse a type of a specific type. - template ParseResult parseType(TypeT &result) { + template + ParseResult parseType(TypeT &result) { llvm::SMLoc loc = getCurrentLocation(); // Parse any kind of type. @@ -614,7 +617,8 @@ virtual ParseResult parseColonType(Type &result) = 0; /// Parse a colon followed by a type of a specific kind, e.g. a FunctionType. - template ParseResult parseColonType(TypeType &result) { + template + ParseResult parseColonType(TypeType &result) { llvm::SMLoc loc = getCurrentLocation(); // Parse any kind of type. 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 @@ -134,9 +134,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,236 @@ } //===----------------------------------------------------------------------===// +// 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 there are ranges, set each to [0, N). + if (ranges.size()) { + SmallVector lbExprs(ranges.size(), + builder->getAffineConstantExpr(0)); + lbMap = AffineMap::get(0, 0, lbExprs); + SmallVector ubExprs; + for (int64_t range : ranges) + ubExprs.push_back(builder->getAffineConstantExpr(range)); + ubMap = AffineMap::get(0, 0, ubExprs); + } + build(builder, result, lbMap, {}, ubMap, {}); +} + +void AffineParallelOp::build(Builder *builder, OperationState &result, + AffineMap lbMap, ValueRange lbArgs, + AffineMap ubMap, ValueRange ubArgs) { + auto numDims = lbMap.getNumResults(); + // Verify that the dimensionality of both maps are the same. + assert(numDims == ubMap.getNumResults() && + "num dims and num results mismatch"); + // Make default step sizes of 1. + SmallVector steps(numDims, 1); + 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) { + auto numDims = lbMap.getNumResults(); + // Verify that the dimensionality of the maps matches the number of steps. + assert(numDims == ubMap.getNumResults() && + "num dims and num results mismatch"); + assert(numDims == steps.size() && "num dims and num steps mismatch"); + 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 block arguments. + for (unsigned i = 0; i < numDims; ++i) + body->addArgument(IndexType::get(builder->getContext())); + bodyRegion->push_back(body); + ensureTerminator(*bodyRegion, *builder, result.location); +} + +unsigned AffineParallelOp::getNumDims() { return steps().size(); } + +AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() { + return getOperands().take_front(lowerBoundsMap().getNumInputs()); +} + +AffineParallelOp::operand_range AffineParallelOp::getUpperBoundsOperands() { + return getOperands().drop_front(lowerBoundsMap().getNumInputs()); +} + +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() { + // Try to convert all the ranges to constant expressions. + SmallVector out; + AffineValueMap rangesValueMap = getRangesValueMap(); + out.reserve(rangesValueMap.getNumResults()); + for (unsigned i = 0, e = rangesValueMap.getNumResults(); i < e; ++i) { + auto expr = rangesValueMap.getResult(i); + auto cst = expr.dyn_cast(); + if (!cst) + return llvm::None; + out.push_back(cst.getValue()); + } + return out; +} + +Block *AffineParallelOp::getBody() { return ®ion().front(); } + +OpBuilder AffineParallelOp::getBodyBuilder() { + return OpBuilder(getBody(), std::prev(getBody()->end())); +} + +void AffineParallelOp::setSteps(ArrayRef newSteps) { + assert(newSteps.size() == getNumDims() && "steps & num dims mismatch"); + setAttr(getStepsAttrName(), getBodyBuilder().getI64ArrayAttr(newSteps)); +} + +static LogicalResult verify(AffineParallelOp op) { + auto numDims = op.getNumDims(); + if (op.lowerBoundsMap().getNumResults() != numDims || + op.upperBoundsMap().getNumResults() != numDims || + op.steps().size() != numDims || + op.getBody()->getNumArguments() != numDims) { + return op.emitOpError("region argument count and num results of upper " + "bounds, lower bounds, and steps must all match"); + } + // Verify that the bound operands are valid dimension/symbols. + /// Lower bounds. + if (failed(verifyDimAndSymbolIdentifiers(op, op.getLowerBoundsOperands(), + op.lowerBoundsMap().getNumDims()))) + return failure(); + /// Upper bounds. + if (failed(verifyDimAndSymbolIdentifiers(op, op.getUpperBoundsOperands(), + op.upperBoundsMap().getNumDims()))) + return failure(); + return success(); +} + +static void print(OpAsmPrinter &p, AffineParallelOp op) { + p << op.getOperationName() << " (" << op.getBody()->getArguments() << ") = ("; + p.printAffineMapOfSSAIds(op.lowerBoundsMapAttr(), + op.getLowerBoundsOperands()); + p << ") to ("; + p.printAffineMapOfSSAIds(op.upperBoundsMapAttr(), + op.getUpperBoundsOperands()); + p << ')'; + SmallVector steps; + bool elideSteps = true; + for (auto attr : op.steps()) { + auto step = attr.cast().getInt(); + elideSteps &= (step == 1); + steps.push_back(step); + } + if (!elideSteps) { + p << " step ("; + interleaveComma(steps, p); + p << ')'; + } + 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; + SmallVector lowerBoundsMapOperands; + SmallVector upperBoundsMapOperands; + if (parser.parseRegionArgumentList(ivs, /*requiredOperandCount=*/-1, + OpAsmParser::Delimiter::Paren) || + parser.parseEqual() || + parser.parseAffineMapOfSSAIds( + lowerBoundsMapOperands, lowerBoundsAttr, + AffineParallelOp::getLowerBoundsMapAttrName(), result.attributes, + OpAsmParser::Delimiter::Paren) || + parser.resolveOperands(lowerBoundsMapOperands, indexType, + result.operands) || + parser.parseKeyword("to") || + parser.parseAffineMapOfSSAIds( + upperBoundsMapOperands, upperBoundsAttr, + AffineParallelOp::getUpperBoundsMapAttrName(), result.attributes, + OpAsmParser::Delimiter::Paren) || + parser.resolveOperands(upperBoundsMapOperands, indexType, + result.operands)) + return failure(); + + AffineMapAttr stepsMapAttr; + SmallVector stepsAttrs; + SmallVector stepsMapOperands; + if (failed(parser.parseOptionalKeyword("step"))) { + SmallVector steps(ivs.size(), 1); + result.addAttribute(AffineParallelOp::getStepsAttrName(), + builder.getI64ArrayAttr(steps)); + } else { + if (parser.parseAffineMapOfSSAIds(stepsMapOperands, stepsMapAttr, + AffineParallelOp::getStepsAttrName(), + stepsAttrs, + OpAsmParser::Delimiter::Paren)) + 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/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp --- a/mlir/lib/Parser/Parser.cpp +++ b/mlir/lib/Parser/Parser.cpp @@ -343,7 +343,8 @@ /// Parse an AffineMap where the dim and symbol identifiers are SSA ids. ParseResult parseAffineMapOfSSAIds(AffineMap &map, - function_ref parseElement); + function_ref parseElement, + OpAsmParser::Delimiter delimiter); private: /// The Parser is subclassed and reinstantiated. Do not add additional @@ -2431,7 +2432,8 @@ AffineMap parseAffineMapRange(unsigned numDims, unsigned numSymbols); ParseResult parseAffineMapOrIntegerSetInline(AffineMap &map, IntegerSet &set); IntegerSet parseIntegerSetConstraints(unsigned numDims, unsigned numSymbols); - ParseResult parseAffineMapOfSSAIds(AffineMap &map); + ParseResult parseAffineMapOfSSAIds(AffineMap &map, + OpAsmParser::Delimiter delimiter); void getDimsAndSymbolSSAIds(SmallVectorImpl &dimAndSymbolSSAIds, unsigned &numDims); @@ -2917,9 +2919,24 @@ } /// Parse an AffineMap where the dim and symbol identifiers are SSA ids. -ParseResult AffineParser::parseAffineMapOfSSAIds(AffineMap &map) { - if (parseToken(Token::l_square, "expected '['")) - return failure(); +ParseResult +AffineParser::parseAffineMapOfSSAIds(AffineMap &map, + OpAsmParser::Delimiter delimiter) { + Token::Kind rightToken; + switch (delimiter) { + case OpAsmParser::Delimiter::Square: + if (parseToken(Token::l_square, "expected '['")) + return failure(); + rightToken = Token::r_square; + break; + case OpAsmParser::Delimiter::Paren: + if (parseToken(Token::l_paren, "expected '('")) + return failure(); + rightToken = Token::r_paren; + break; + default: + return emitError("unexpected delimiter"); + } SmallVector exprs; auto parseElt = [&]() -> ParseResult { @@ -2931,7 +2948,7 @@ // Parse a multi-dimensional affine expression (a comma-separated list of // 1-d affine expressions); the list cannot be empty. Grammar: // multi-dim-affine-expr ::= `(` affine-expr (`,` affine-expr)* `) - if (parseCommaSeparatedListUntil(Token::r_square, parseElt, + if (parseCommaSeparatedListUntil(rightToken, parseElt, /*allowEmptyList=*/true)) return failure(); // Parsed a valid affine map. @@ -3081,9 +3098,10 @@ /// parse SSA value uses encountered while parsing affine expressions. ParseResult Parser::parseAffineMapOfSSAIds(AffineMap &map, - function_ref parseElement) { + function_ref parseElement, + OpAsmParser::Delimiter delimiter) { return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement) - .parseAffineMapOfSSAIds(map); + .parseAffineMapOfSSAIds(map, delimiter); } //===----------------------------------------------------------------------===// @@ -4182,10 +4200,10 @@ } /// Parse an AffineMap of SSA ids. - ParseResult - parseAffineMapOfSSAIds(SmallVectorImpl &operands, - Attribute &mapAttr, StringRef attrName, - SmallVectorImpl &attrs) override { + ParseResult parseAffineMapOfSSAIds(SmallVectorImpl &operands, + Attribute &mapAttr, StringRef attrName, + SmallVectorImpl &attrs, + Delimiter delimiter) override { SmallVector dimOperands; SmallVector symOperands; @@ -4201,7 +4219,7 @@ }; AffineMap map; - if (parser.parseAffineMapOfSSAIds(map, parseElement)) + if (parser.parseAffineMapOfSSAIds(map, parseElement, delimiter)) return failure(); // Add AffineMap attribute. if (map) { 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,68 @@ %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 (%i) = (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 (%i, %j) = (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 (%i, %j) = (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 (%i, %j) = (0, 0) to (100, 100) step (10) { + } +} + +// ----- + +// CHECK-LABEL: @affine_parallel + +func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { + affine.for %x = 0 to 7 { + %y = addi %x, %x : index + // expected-error@+1 {{operand cannot be used as a dimension id}} + affine.parallel (%i, %j) = (0, 0) to (%y, 100) step (10, 10) { + } + } + return +} + +// ----- + +// CHECK-LABEL: @affine_parallel + +func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { + affine.for %x = 0 to 7 { + %y = addi %x, %x : index + // expected-error@+1 {{operand cannot be used as a symbol}} + affine.parallel (%i, %j) = (0, 0) to (symbol(%y), 100) step (10, 10) { + } + } + return +} 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 +}