diff --git a/mlir/include/mlir/Dialect/AffineOps/AffineOps.h b/mlir/include/mlir/Dialect/AffineOps/AffineOps.h --- a/mlir/include/mlir/Dialect/AffineOps/AffineOps.h +++ b/mlir/include/mlir/Dialect/AffineOps/AffineOps.h @@ -23,6 +23,7 @@ #include "mlir/Transforms/LoopLikeInterface.h" namespace mlir { +class AffineApplyOp; class AffineBound; class AffineDimExpr; class AffineValueMap; @@ -46,57 +47,6 @@ Location loc) override; }; -/// The "affine.apply" operation applies an affine map to a list of operands, -/// yielding a single result. The operand list must be the same size as the -/// number of arguments to the affine mapping. All operands and the result are -/// of type 'Index'. This operation requires a single affine map attribute named -/// "map". For example: -/// -/// %y = "affine.apply" (%x) { map: (d0) -> (d0 + 1) } : -/// (index) -> (index) -/// -/// equivalently: -/// -/// #map42 = (d0)->(d0+1) -/// %y = affine.apply #map42(%x) -/// -class AffineApplyOp : public Op { -public: - using Op::Op; - - /// Builds an affine apply op with the specified map and operands. - static void build(Builder *builder, OperationState &result, AffineMap map, - ValueRange operands); - - /// Returns the affine map to be applied by this operation. - AffineMap getAffineMap() { - return getAttrOfType("map").getValue(); - } - - /// Returns the affine value map computed from this operation. - AffineValueMap getAffineValueMap(); - - /// Returns true if the result of this operation can be used as dimension id. - bool isValidDim(); - - /// Returns true if the result of this operation is a symbol. - bool isValidSymbol(); - - static StringRef getOperationName() { return "affine.apply"; } - - operand_range getMapOperands() { return getOperands(); } - - // Hooks to customize behavior of this op. - static ParseResult parse(OpAsmParser &parser, OperationState &result); - void print(OpAsmPrinter &p); - LogicalResult verify(); - OpFoldResult fold(ArrayRef operands); - - static void getCanonicalizationPatterns(OwningRewritePatternList &results, - MLIRContext *context); -}; - /// AffineDmaStartOp starts a non-blocking DMA operation that transfers data /// from a source memref to a destination memref. The source and destination /// memref need not be of the same dimensionality, but need to have the same 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 @@ -40,6 +40,60 @@ def ImplicitAffineTerminator : SingleBlockImplicitTerminator<"AffineTerminatorOp">; +def AffineApplyOp : Affine_Op<"apply", [NoSideEffect]> { + let summary = "affine apply operation"; + let description = [{ + The affine.apply operation applies an affine mapping to a list of SSA + values, yielding a single SSA value. The number of dimension and symbol + arguments to affine.apply must be equal to the respective number of + dimensional and symbolic inputs to the affine mapping; the affine mapping + has to be one-dimensional, and so the affine.apply operation always returns + one value. The input operands and result must all have ‘index’ type. + + Example: + + ```mlir + #map10 = affine_map<(d0, d1) -> (d0 floordiv 8 + d1 floordiv 128)> + ... + %1 = affine.apply #map10 (%s, %t) + + // Inline example. + %2 = affine.apply affine_map<(i)[s0] -> (i+s0)> (%42)[%n] + ``` + }]; + let arguments = (ins AffineMapAttr:$map, Variadic:$mapOperands); + let results = (outs Index); + + // TODO: The auto-generated builders should check to see if the return type + // has a constant builder. That way we wouldn't need to explicitly specify the + // result types here. + let builders = [ + OpBuilder<"Builder *builder, OperationState &result, " + "AffineMap map, ValueRange mapOperands", [{ + build(builder, result, builder->getIndexType(), map, mapOperands); + }]> + ]; + + let extraClassDeclaration = [{ + /// Returns the affine map to be applied by this operation. + AffineMap getAffineMap() { return map(); } + + /// Returns the affine value map computed from this operation. + AffineValueMap getAffineValueMap(); + + /// Returns true if the result of this operation can be used as dimension id. + bool isValidDim(); + + /// Returns true if the result of this operation is a symbol. + bool isValidSymbol(); + + operand_range getMapOperands() { return getOperands(); } + }]; + + let hasCanonicalizer = 1; + let hasFolder = 1; +} + def AffineForOp : Affine_Op<"for", [ImplicitAffineTerminator, RecursiveSideEffects, DeclareOpInterfaceMethods]> { 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 @@ -69,8 +69,7 @@ AffineOpsDialect::AffineOpsDialect(MLIRContext *context) : Dialect(getDialectNamespace(), context) { - addOperations(); @@ -217,18 +216,12 @@ // AffineApplyOp //===----------------------------------------------------------------------===// -void AffineApplyOp::build(Builder *builder, OperationState &result, - AffineMap map, ValueRange operands) { - result.addOperands(operands); - result.types.append(map.getNumResults(), builder->getIndexType()); - result.addAttribute("map", AffineMapAttr::get(map)); -} - AffineValueMap AffineApplyOp::getAffineValueMap() { return AffineValueMap(getAffineMap(), getOperands(), getResult()); } -ParseResult AffineApplyOp::parse(OpAsmParser &parser, OperationState &result) { +static ParseResult parseAffineApplyOp(OpAsmParser &parser, + OperationState &result) { auto &builder = parser.getBuilder(); auto indexTy = builder.getIndexType(); @@ -250,39 +243,25 @@ return success(); } -void AffineApplyOp::print(OpAsmPrinter &p) { - p << "affine.apply " << getAttr("map"); - printDimAndSymbolList(operand_begin(), operand_end(), - getAffineMap().getNumDims(), p); - p.printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/{"map"}); +static void print(OpAsmPrinter &p, AffineApplyOp op) { + p << AffineApplyOp::getOperationName() << " " << op.mapAttr(); + printDimAndSymbolList(op.operand_begin(), op.operand_end(), + op.getAffineMap().getNumDims(), p); + p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"map"}); } -LogicalResult AffineApplyOp::verify() { - // Check that affine map attribute was specified. - auto affineMapAttr = getAttrOfType("map"); - if (!affineMapAttr) - return emitOpError("requires an affine map"); - +static LogicalResult verify(AffineApplyOp op) { // Check input and output dimensions match. - auto map = affineMapAttr.getValue(); + auto map = op.map(); // Verify that operand count matches affine map dimension and symbol count. - if (getNumOperands() != map.getNumDims() + map.getNumSymbols()) - return emitOpError( + if (op.getNumOperands() != map.getNumDims() + map.getNumSymbols()) + return op.emitOpError( "operand count and affine map dimension and symbol count must match"); - // Verify that all operands are of `index` type. - for (Type t : getOperandTypes()) { - if (!t.isIndex()) - return emitOpError("operands must be of type 'index'"); - } - - if (!getResult().getType().isIndex()) - return emitOpError("result must be of type 'index'"); - // Verify that the map only produces one result. if (map.getNumResults() != 1) - return emitOpError("mapping must produce one value"); + return op.emitOpError("mapping must produce one value"); return success(); } diff --git a/mlir/test/Dialect/AffineOps/invalid.mlir b/mlir/test/Dialect/AffineOps/invalid.mlir --- a/mlir/test/Dialect/AffineOps/invalid.mlir +++ b/mlir/test/Dialect/AffineOps/invalid.mlir @@ -5,7 +5,7 @@ func @affine_apply_operand_non_index(%arg0 : i32) { // Custom parser automatically assigns all arguments the `index` so we must // use the generic syntax here to exercise the verifier. - // expected-error@+1 {{operands must be of type 'index'}} + // expected-error@+1 {{op operand #0 must be index, but got 'i32'}} %0 = "affine.apply"(%arg0) {map = affine_map<(d0) -> (d0)>} : (i32) -> (index) return } @@ -15,7 +15,7 @@ func @affine_apply_resul_non_index(%arg0 : index) { // Custom parser automatically assigns `index` as the result type so we must // use the generic syntax here to exercise the verifier. - // expected-error@+1 {{result must be of type 'index'}} + // expected-error@+1 {{op result #0 must be index, but got 'i32'}} %0 = "affine.apply"(%arg0) {map = affine_map<(d0) -> (d0)>} : (index) -> (i32) return } diff --git a/mlir/test/IR/invalid-ops.mlir b/mlir/test/IR/invalid-ops.mlir --- a/mlir/test/IR/invalid-ops.mlir +++ b/mlir/test/IR/invalid-ops.mlir @@ -58,7 +58,7 @@ func @affine_apply_no_map() { ^bb0: %i = constant 0 : index - %x = "affine.apply" (%i) { } : (index) -> (index) // expected-error {{'affine.apply' op requires an affine map}} + %x = "affine.apply" (%i) { } : (index) -> (index) // expected-error {{'affine.apply' op requires attribute 'map'}} return }