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 @@ -274,7 +274,30 @@ def AffineParallelOp : Affine_Op<"parallel", [ImplicitAffineTerminator]> { let summary = "multi-index parallel band operation"; let description = [{ - TODO + The "affine.parallel" operation represents an 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 @@ -295,13 +318,12 @@ ]; let extraClassDeclaration = [{ + /// Get the number of dimensions. + unsigned getNumDims(); + operand_range getLowerBoundsOperands(); operand_range getUpperBoundsOperands(); - /// Get the number of dimensions (i.e. number of induction variables). - size_t getNumDims(); - - /// Get various AffineValueMaps AffineValueMap getLowerBoundsValueMap(); AffineValueMap getUpperBoundsValueMap(); AffineValueMap getRangesValueMap(); @@ -309,16 +331,14 @@ /// Get ranges as constants, may fail in dynamic case. llvm::Optional> getConstantRanges(); - /// Some helpers for rewrites. mlir::Block *getBody(); mlir::OpBuilder getBodyBuilder(); - void setSteps(ArrayRef newSteps); - }]; - // No custom parsing/printing form. - let parser = ?; - let printer = ?; + static StringRef getLowerBoundsMapAttrName() { return "lowerBoundsMap"; } + static StringRef getUpperBoundsMapAttrName() { return "upperBoundsMap"; } + static StringRef getStepsAttrName() { return "steps"; } + }]; } def AffinePrefetchOp : Affine_Op<"prefetch"> { 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 @@ -1340,7 +1340,7 @@ } static void print(OpAsmPrinter &p, AffineForOp op) { - p << "affine.for "; + p << op.getOperationName() << ' '; p.printOperand(op.getBody()->getArgument(0)); p << " = "; printBound(op.getLowerBoundMapAttr(), op.getLowerBoundOperands(), "max", p); @@ -2182,10 +2182,10 @@ AffineMap lbMap, ValueRange lbArgs, AffineMap ubMap, ValueRange ubArgs) { // Verify sizes - size_t idxCount = lbMap.getNumResults(); - assert(idxCount == ubMap.getNumResults()); + auto dims = lbMap.getNumResults(); + assert(dims == ubMap.getNumResults()); // Make default step sizes of 1 - SmallVector steps(idxCount, 1); + SmallVector steps(dims, 1); // Call through build(builder, result, lbMap, lbArgs, ubMap, ubArgs, steps); } @@ -2195,20 +2195,20 @@ AffineMap ubMap, ValueRange ubArgs, ArrayRef steps) { // Verify sizes - size_t idxCount = lbMap.getNumResults(); - assert(idxCount == ubMap.getNumResults()); - assert(idxCount == steps.size()); + auto dims = lbMap.getNumResults(); + assert(dims == ubMap.getNumResults()); + assert(dims == steps.size()); // Set all of the attributes - result.addAttribute("lowerBoundsMap", AffineMapAttr::get(lbMap)); - result.addAttribute("upperBoundsMap", AffineMapAttr::get(ubMap)); - result.addAttribute("steps", builder->getI64ArrayAttr(steps)); + 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 (size_t i = 0; i < idxCount; i++) { + for (unsigned i = 0; i < dims; ++i) { body->addArgument(IndexType::get(builder->getContext())); } bodyRegion->push_back(body); @@ -2224,24 +2224,14 @@ return {operand_begin() + lowerBoundsMap().getNumInputs(), operand_end()}; } -size_t AffineParallelOp::getNumDims() { - return steps().cast().size(); -} +unsigned AffineParallelOp::getNumDims() { return steps().size(); } AffineValueMap AffineParallelOp::getLowerBoundsValueMap() { - SmallVector ops; - for (auto op : getLowerBoundsOperands()) { - ops.push_back(op); - } - return AffineValueMap(lowerBoundsMap(), ops); + return AffineValueMap(lowerBoundsMap(), getLowerBoundsOperands()); } AffineValueMap AffineParallelOp::getUpperBoundsValueMap() { - SmallVector ops; - for (auto op : getUpperBoundsOperands()) { - ops.push_back(op); - } - return AffineValueMap(upperBoundsMap(), ops); + return AffineValueMap(upperBoundsMap(), getUpperBoundsOperands()); } AffineValueMap AffineParallelOp::getRangesValueMap() { @@ -2256,7 +2246,7 @@ // Get the ranges AffineValueMap rangesValueMap = getRangesValueMap(); // Try to convert them to constants - for (size_t i = 0; i < rangesValueMap.getNumResults(); i++) { + 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()); @@ -2274,27 +2264,116 @@ } void AffineParallelOp::setSteps(ArrayRef newSteps) { - assert(newSteps.size() == steps().size()); + assert(newSteps.size() == getNumDims()); setAttr("steps", getBodyBuilder().getI64ArrayAttr(newSteps)); } static LogicalResult verify(AffineParallelOp op) { - size_t idxCount = op.steps().size(); - if (op.lowerBoundsMap().getNumResults() != idxCount) { - return op.emitOpError("steps.size (") - << idxCount << ") must match lowerBoundsMap().getNumResults() (" - << op.lowerBoundsMap().getNumResults() << ")"; + 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 (op.upperBoundsMap().getNumResults() != idxCount) { - return op.emitOpError("steps.size (") - << idxCount << ") must match upperBoundsMap().getNumResults() (" - << op.upperBoundsMap().getNumResults() << ")"; + if (!elideSteps) { + p << " step " << op.steps(); } - if (op.getBody()->getNumArguments() != idxCount) { - return op.emitOpError("steps.size (") - << idxCount << ") must match getBody()->getNumArguments() (" - << op.getBody()->getNumArguments() << ")"; + 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(); } 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 +}