diff --git a/mlir/include/mlir/Conversion/AffineToStandard/AffineToStandard.h b/mlir/include/mlir/Conversion/AffineToStandard/AffineToStandard.h --- a/mlir/include/mlir/Conversion/AffineToStandard/AffineToStandard.h +++ b/mlir/include/mlir/Conversion/AffineToStandard/AffineToStandard.h @@ -44,6 +44,11 @@ void populateAffineToStdConversionPatterns(OwningRewritePatternList &patterns, MLIRContext *ctx); +/// Collect a set of patterns to convert vector-related Affine ops to the Vector +/// dialect. +void populateAffineToVectorConversionPatterns( + OwningRewritePatternList &patterns, MLIRContext *ctx); + /// Emit code that computes the lower bound of the given affine loop using /// standard arithmetic operations. Value lowerAffineLowerBound(AffineForOp op, OpBuilder &builder); diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td @@ -370,7 +370,50 @@ let hasFolder = 1; } -def AffineLoadOp : Affine_Op<"load", []> { +class AffineLoadOpBase traits = []> : + Affine_Op { + let arguments = (ins Arg:$memref, + Variadic:$indices); + + let extraClassDeclaration = [{ + /// Returns the operand index of the memref. + unsigned getMemRefOperandIndex() { return 0; } + + /// Get memref operand. + Value getMemRef() { return getOperand(getMemRefOperandIndex()); } + void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); } + MemRefType getMemRefType() { + return getMemRef().getType().cast(); + } + + /// Get affine map operands. + operand_range getMapOperands() { return llvm::drop_begin(getOperands(), 1); } + + /// Returns the affine map used to index the memref for this operation. + AffineMap getAffineMap() { return getAffineMapAttr().getValue(); } + AffineMapAttr getAffineMapAttr() { + return getAttr(getMapAttrName()).cast(); + } + + /// Returns the AffineMapAttr associated with 'memref'. + NamedAttribute getAffineMapAttrForMemRef(Value memref) { + assert(memref == getMemRef()); + return {Identifier::get(getMapAttrName(), getContext()), + getAffineMapAttr()}; + } + + static StringRef getMapAttrName() { return "map"; } + + // TODO: Move getVectorType to AffineVLoadOp when ODS supports its + // (Bugzilla/45878). + VectorType getVectorType() { + return result().getType().cast(); + } + }]; +} + +def AffineLoadOp : AffineLoadOpBase<"load", []> { let summary = "affine load operation"; let description = [{ The "affine.load" op reads an element from a memref, where the index @@ -393,9 +436,6 @@ ``` }]; - let arguments = (ins Arg:$memref, - Variadic:$indices); let results = (outs AnyType:$result); let builders = [ @@ -410,36 +450,6 @@ "AffineMap map, ValueRange mapOperands"> ]; - let extraClassDeclaration = [{ - /// Returns the operand index of the memref. - unsigned getMemRefOperandIndex() { return 0; } - - /// Get memref operand. - Value getMemRef() { return getOperand(getMemRefOperandIndex()); } - void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); } - MemRefType getMemRefType() { - return getMemRef().getType().cast(); - } - - /// Get affine map operands. - operand_range getMapOperands() { return llvm::drop_begin(getOperands(), 1); } - - /// Returns the affine map used to index the memref for this operation. - AffineMap getAffineMap() { return getAffineMapAttr().getValue(); } - AffineMapAttr getAffineMapAttr() { - return getAttr(getMapAttrName()).cast(); - } - - /// Returns the AffineMapAttr associated with 'memref'. - NamedAttribute getAffineMapAttrForMemRef(Value memref) { - assert(memref == getMemRef()); - return {Identifier::get(getMapAttrName(), getContext()), - getAffineMapAttr()}; - } - - static StringRef getMapAttrName() { return "map"; } - }]; - let hasCanonicalizer = 1; let hasFolder = 1; } @@ -659,42 +669,8 @@ let hasFolder = 1; } -def AffineStoreOp : Affine_Op<"store", []> { - let summary = "affine store operation"; - let description = [{ - The "affine.store" op writes an element to a memref, where the index - for each memref dimension is an affine expression of loop induction - variables and symbols. The 'affine.store' op stores a new value which is the - same type as the elements of the memref. An affine expression of loop IVs - and symbols must be specified for each dimension of the memref. The keyword - 'symbol' can be used to indicate SSA identifiers which are symbolic. - - Example 1: - - ```mlir - affine.store %v0, %0[%i0 + 3, %i1 + 7] : memref<100x100xf32> - ``` - - Example 2: Uses 'symbol' keyword for symbols '%n' and '%m'. - - ```mlir - affine.store %v0, %0[%i0 + symbol(%n), %i1 + symbol(%m)] : memref<100x100xf32> - ``` - }]; - let arguments = (ins AnyType:$value, - Arg:$memref, - Variadic:$indices); - - - let skipDefaultBuilders = 1; - let builders = [ - OpBuilder<"OpBuilder &builder, OperationState &result, " - "Value valueToStore, Value memref, ValueRange indices">, - OpBuilder<"OpBuilder &builder, OperationState &result, " - "Value valueToStore, Value memref, AffineMap map, " - "ValueRange mapOperands"> - ]; +class AffineStoreOpBase traits = []> : + Affine_Op { let extraClassDeclaration = [{ /// Get value to be stored by store operation. @@ -728,7 +704,50 @@ } static StringRef getMapAttrName() { return "map"; } + + // TODO: Move getVectorType to AffineVStoreOp when ODS supports its + // (Bugzilla/45878). + VectorType getVectorType() { + return value().getType().cast(); + } }]; +} + +def AffineStoreOp : AffineStoreOpBase<"store", []> { + let summary = "affine store operation"; + let description = [{ + The "affine.store" op writes an element to a memref, where the index + for each memref dimension is an affine expression of loop induction + variables and symbols. The 'affine.store' op stores a new value which is the + same type as the elements of the memref. An affine expression of loop IVs + and symbols must be specified for each dimension of the memref. The keyword + 'symbol' can be used to indicate SSA identifiers which are symbolic. + + Example 1: + + ```mlir + affine.store %v0, %0[%i0 + 3, %i1 + 7] : memref<100x100xf32> + ``` + + Example 2: Uses 'symbol' keyword for symbols '%n' and '%m'. + + ```mlir + affine.store %v0, %0[%i0 + symbol(%n), %i1 + symbol(%m)] : memref<100x100xf32> + ``` + }]; + let arguments = (ins AnyType:$value, + Arg:$memref, + Variadic:$indices); + + let skipDefaultBuilders = 1; + let builders = [ + OpBuilder<"OpBuilder &builder, OperationState &result, " + "Value valueToStore, Value memref, ValueRange indices">, + OpBuilder<"OpBuilder &builder, OperationState &result, " + "Value valueToStore, Value memref, AffineMap map, " + "ValueRange mapOperands"> + ]; let hasCanonicalizer = 1; let hasFolder = 1; @@ -765,4 +784,75 @@ let verifier = ?; } +def AffineVLoadOp : AffineLoadOpBase<"vload", []> { + let summary = "affine vector load operation"; + let description = [{ + The "affine.vload" is the vector counterpart of + [affine.load](#affineload-operation). It performs a blocking read from a + slice within a [MemRef](../LangRef.md#memref-type) supplied as its first + operand into a [vector](../LangRef.md#vector-type) of the same base + elemental type. + The index for each memref dimension is an affine expression of loop + induction variables and symbols. The output of 'affine.vload' is a new + vector value with the same base element type as the elements of the memref. + An affine expression of loop IVs and symbols must be specified for each + dimension of the memref. The keyword 'symbol' can be used to indicate SSA + identifiers which are symbolic. + + Example 1: 8-wide f32 vector load. + + ```mlir + %1 = affine.vload %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>, vector<8xf32> + ``` + + Example 2: 4-wide f32 vector load. Uses 'symbol' keyword for symbols '%n' and '%m'. + + ```mlir + %1 = affine.load %0[%i0 + symbol(%n), %i1 + symbol(%m)] : memref<100x100xf32>, vector<4xf32> + ``` + }]; + + let results = (outs AnyVector:$result); + + // Fully specified by traits. + let verifier = ?; +} + +def AffineVStoreOp : AffineStoreOpBase<"vstore", []> { + let summary = "affine vector store operation"; + let description = [{ + The "affine.vstore" is the vector counterpart of + [affine.store](#affinestore-affinestoreop). It performs a blocking write + from a [vector](../LangRef.md#vector-type), supplied as its first operand, + into a slice within a [MemRef](../LangRef.md#memref-type) of the same base + elemental type, supplied as its second operand. + The index for each memref dimension is an affine expression of loop + induction variables and symbols. + The 'affine.vstore' op stores a new vector value with the same base element + type as the elements of the memref. An affine expression of loop IVs and + symbols must be specified for each dimension of the memref. The keyword + 'symbol' can be used to indicate SSA identifiers which are symbolic. + + Example 1: 8-wide f32 vector store. + + ```mlir + affine.vstore %v0, %0[%i0 + 3, %i1 + 7] : vector<8xf32>, memref<100x100xf32> + ``` + + Example 2: 4-wide f32 vector store. Uses 'symbol' keyword for symbols '%n' and '%m'. + + ```mlir + affine.vstore %v0, %0[%i0 + symbol(%n), %i1 + symbol(%m)] : vector<4xf32>, memref<100x100xf32> + ``` + }]; + + let arguments = (ins AnyVector:$value, + Arg:$memref, + Variadic:$indices); + + // Fully specified by traits. + let verifier = ?; +} + #endif // AFFINE_OPS diff --git a/mlir/include/mlir/Dialect/Vector/VectorOps.td b/mlir/include/mlir/Dialect/Vector/VectorOps.td --- a/mlir/include/mlir/Dialect/Vector/VectorOps.td +++ b/mlir/include/mlir/Dialect/Vector/VectorOps.td @@ -994,6 +994,21 @@ ``` }]; + let builders = [ + // Builder that sets permunation map and padding to 'getMinorIdentityMap' + // and zero, respectively, by default. + OpBuilder<"OpBuilder &builder, OperationState &result, Type vector, " + "Value memref, ValueRange indices", [{ + auto permMap = AffineMap::getMinorIdentityMap(memref.getType() + .cast().getRank(), /*results=*/1, builder.getContext()); + Type elemType = vector.cast().getElementType(); + Value padding = builder.create(result.location, elemType, + builder.getZeroAttr(elemType)); + + build(builder, result, vector, memref, indices, permMap, padding); + }]> + ]; + let extraClassDeclaration = [{ MemRefType getMemRefType() { return memref().getType().cast(); @@ -1058,6 +1073,17 @@ ``` }]; + let builders = [ + // Builder that sets permunation map and padding to 'getMinorIdentityMap' + // by default. + OpBuilder<"OpBuilder &builder, OperationState &result, Value vector, " + "Value memref, ValueRange indices", [{ + auto permMap = AffineMap::getMinorIdentityMap(memref.getType() + .cast().getRank(), /*results=*/1, builder.getContext()); + build(builder, result, vector, memref, indices, permMap); + }]> + ]; + let extraClassDeclaration = [{ VectorType getVectorType() { return vector().getType().cast(); diff --git a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp --- a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp +++ b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp @@ -17,6 +17,7 @@ #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/LoopOps/LoopOps.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/Dialect/Vector/VectorOps.h" #include "mlir/IR/AffineExprVisitor.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/IR/Builders.h" @@ -27,6 +28,7 @@ #include "mlir/Transforms/Passes.h" using namespace mlir; +using namespace mlir::vector; namespace { /// Visit affine expressions recursively and build the sequence of operations @@ -556,6 +558,62 @@ } }; +/// Apply the affine map from an 'affine.vload' operation to its operands, and +/// feed the results to a newly created 'vector.transfer_read' operation (which +/// replaces the original 'affine.vload'). +class AffineVLoadLowering : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(AffineVLoadOp op, + PatternRewriter &rewriter) const override { + // Expand affine map from 'affineVLoadOp'. + SmallVector indices(op.getMapOperands()); + auto resultOperands = + expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices); + if (!resultOperands) + return failure(); + + // Build vector.transfer_read memref[expandedMap.results]. + auto permMap = AffineMap::getMinorIdentityMap(op.getMemRefType().getRank(), + 1, rewriter.getContext()); + Type elemType = op.getVectorType().getElementType(); + Value padding = rewriter.create(op.getLoc(), elemType, + rewriter.getZeroAttr(elemType)); + + rewriter.replaceOpWithNewOp( + op, op.getVectorType(), op.getMemRef(), *resultOperands, + AffineMapAttr::get(permMap), padding); + return success(); + } +}; + +/// Apply the affine map from an 'affine.vstore' operation to its operands, and +/// feed the results to a newly created 'vector.transfer_write' operation (which +/// replaces the original 'affine.vstore'). +class AffineVStoreLowering : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(AffineVStoreOp op, + PatternRewriter &rewriter) const override { + // Expand affine map from 'affineVStoreOp'. + SmallVector indices(op.getMapOperands()); + auto maybeExpandedMap = + expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices); + if (!maybeExpandedMap) + return failure(); + + // Build std.store valueToStore, memref[expandedMap.results]. + auto permMap = AffineMap::getMinorIdentityMap(op.getMemRefType().getRank(), + 1, rewriter.getContext()); + rewriter.replaceOpWithNewOp( + op, op.getValueToStore(), op.getMemRef(), *maybeExpandedMap, + AffineMapAttr::get(permMap)); + return success(); + } +}; + } // end namespace void mlir::populateAffineToStdConversionPatterns( @@ -576,13 +634,24 @@ // clang-format on } +void mlir::populateAffineToVectorConversionPatterns( + OwningRewritePatternList &patterns, MLIRContext *ctx) { + // clang-format off + patterns.insert< + AffineVLoadLowering, + AffineVStoreLowering>(ctx); + // clang-format on +} + namespace { class LowerAffinePass : public ConvertAffineToStandardBase { void runOnFunction() override { OwningRewritePatternList patterns; populateAffineToStdConversionPatterns(patterns, &getContext()); + populateAffineToVectorConversionPatterns(patterns, &getContext()); ConversionTarget target(getContext()); - target.addLegalDialect(); + target.addLegalDialect(); if (failed(applyPartialConversion(getFunction(), target, patterns))) signalPassFailure(); } diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -2484,6 +2484,80 @@ return success(); } +//===----------------------------------------------------------------------===// +// AffineVLoadOp +//===----------------------------------------------------------------------===// + +ParseResult parseAffineVLoadOp(OpAsmParser &parser, OperationState &result) { + auto &builder = parser.getBuilder(); + auto indexTy = builder.getIndexType(); + + MemRefType memrefType; + VectorType resultType; + OpAsmParser::OperandType memrefInfo; + AffineMapAttr mapAttr; + SmallVector mapOperands; + return failure( + parser.parseOperand(memrefInfo) || + parser.parseAffineMapOfSSAIds(mapOperands, mapAttr, + AffineVLoadOp::getMapAttrName(), + result.attributes) || + parser.parseOptionalAttrDict(result.attributes) || + parser.parseColonType(memrefType) || parser.parseComma() || + parser.parseType(resultType) || + parser.resolveOperand(memrefInfo, memrefType, result.operands) || + parser.resolveOperands(mapOperands, indexTy, result.operands) || + parser.addTypeToList(resultType, result.types)); +} + +void print(OpAsmPrinter &p, AffineVLoadOp op) { + p << "affine.vload " << op.getMemRef() << '['; + if (AffineMapAttr mapAttr = + op.getAttrOfType(op.getMapAttrName())) + p.printAffineMapOfSSAIds(mapAttr, op.getMapOperands()); + p << ']'; + p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{op.getMapAttrName()}); + p << " : " << op.getMemRefType() << ", " << op.getType(); +} + +//===----------------------------------------------------------------------===// +// AffineVStoreOp +//===----------------------------------------------------------------------===// + +ParseResult parseAffineVStoreOp(OpAsmParser &parser, OperationState &result) { + auto indexTy = parser.getBuilder().getIndexType(); + + MemRefType memrefType; + VectorType resultType; + OpAsmParser::OperandType storeValueInfo; + OpAsmParser::OperandType memrefInfo; + AffineMapAttr mapAttr; + SmallVector mapOperands; + return failure( + parser.parseOperand(storeValueInfo) || parser.parseComma() || + parser.parseOperand(memrefInfo) || + parser.parseAffineMapOfSSAIds(mapOperands, mapAttr, + AffineVStoreOp::getMapAttrName(), + result.attributes) || + parser.parseOptionalAttrDict(result.attributes) || + parser.parseColonType(resultType) || parser.parseComma() || + parser.parseType(memrefType) || + parser.resolveOperand(storeValueInfo, resultType, result.operands) || + parser.resolveOperand(memrefInfo, memrefType, result.operands) || + parser.resolveOperands(mapOperands, indexTy, result.operands)); +} + +void print(OpAsmPrinter &p, AffineVStoreOp op) { + p << "affine.vstore " << op.getValueToStore(); + p << ", " << op.getMemRef() << '['; + if (AffineMapAttr mapAttr = + op.getAttrOfType(op.getMapAttrName())) + p.printAffineMapOfSSAIds(mapAttr, op.getMapOperands()); + p << ']'; + p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{op.getMapAttrName()}); + p << " : " << op.getValueToStore().getType() << ", " << op.getMemRefType(); +} + //===----------------------------------------------------------------------===// // TableGen'd op method definitions //===----------------------------------------------------------------------===// diff --git a/mlir/test/Conversion/AffineToStandard/lower-affine-to-vector.mlir b/mlir/test/Conversion/AffineToStandard/lower-affine-to-vector.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Conversion/AffineToStandard/lower-affine-to-vector.mlir @@ -0,0 +1,39 @@ +// RUN: mlir-opt -lower-affine --split-input-file %s | FileCheck %s + +// CHECK: #[[perm_map:.*]] = affine_map<(d0) -> (d0)> +// CHECK-LABEL: func @affine_vload +func @affine_vload(%arg0 : index) { + %0 = alloc() : memref<16xf32> + affine.for %i0 = 0 to 16 { + %1 = affine.vload %0[%i0 + symbol(%arg0) + 7] : memref<16xf32>, vector<8xf32> + } +// CHECK: %[[buf:.*]] = alloc +// CHECK: %[[a:.*]] = addi %{{.*}}, %{{.*}} : index +// CHECK-NEXT: %[[c7:.*]] = constant 7 : index +// CHECK-NEXT: %[[b:.*]] = addi %[[a]], %[[c7]] : index +// CHECK-NEXT: %[[pad:.*]] = constant 0.0 +// CHECK-NEXT: vector.transfer_read %[[buf]][%[[b]]], %[[pad]] {permutation_map = #[[perm_map]]} : memref<16xf32>, vector<8xf32> + return +} + +// ----- + +// CHECK: #[[perm_map:.*]] = affine_map<(d0) -> (d0)> +// CHECK-LABEL: func @affine_vstore +func @affine_vstore(%arg0 : index) { + %0 = alloc() : memref<16xf32> + %1 = constant dense<11.0> : vector<4xf32> + affine.for %i0 = 0 to 16 { + affine.vstore %1, %0[%i0 - symbol(%arg0) + 7] : vector<4xf32>, memref<16xf32> + } +// CHECK: %[[buf:.*]] = alloc +// CHECK: %[[val:.*]] = constant dense +// CHECK: %[[c_1:.*]] = constant -1 : index +// CHECK-NEXT: %[[a:.*]] = muli %arg0, %[[c_1]] : index +// CHECK-NEXT: %[[b:.*]] = addi %{{.*}}, %[[a]] : index +// CHECK-NEXT: %[[c7:.*]] = constant 7 : index +// CHECK-NEXT: %[[c:.*]] = addi %[[b]], %[[c7]] : index +// CHECK-NEXT: vector.transfer_write %[[val]], %[[buf]][%[[c]]] {permutation_map = #[[perm_map]]} : vector<4xf32>, memref<16xf32> + return +} + diff --git a/mlir/test/Dialect/Affine/load-store.mlir b/mlir/test/Dialect/Affine/load-store.mlir --- a/mlir/test/Dialect/Affine/load-store.mlir +++ b/mlir/test/Dialect/Affine/load-store.mlir @@ -214,3 +214,47 @@ } return } + +// ----- + +// CHECK: [[MAP0:#map[0-9]+]] = affine_map<(d0, d1) -> (d0, d1)> + +// Test with just loop IVs. +func @vload_vstore_iv(%arg0 : index, %arg1 : index) { + %0 = alloc() : memref<100x100xf32> + affine.for %i0 = 0 to 16 { + affine.for %i1 = 0 to 16 step 8 { + %1 = affine.vload %0[%i0, %i1] : memref<100x100xf32>, vector<8xf32> + affine.vstore %1, %0[%i0, %i1] : vector<8xf32>, memref<100x100xf32> +// CHECK: %[[buf:.*]] = alloc +// CHECK-NEXT: affine.for %[[i0:.*]] = 0 +// CHECK-NEXT: affine.for %[[i1:.*]] = 0 +// CHECK-NEXT: %[[val:.*]] = affine.vload %[[buf]][%[[i0]], %[[i1]]] : memref<100x100xf32>, vector<8xf32> +// CHECK-NEXT: affine.vstore %[[val]], %[[buf]][%[[i0]], %[[i1]]] : vector<8xf32>, memref<100x100xf32> + } + } + return +} + +// ----- + +// CHECK: [[MAP0:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 + 3, d1 + 7)> + +// Test with loop IVs and constants. +func @vload_vstore_iv_constant(%arg0 : index, %arg1 : index) { + %0 = alloc() : memref<100x100xf32> + affine.for %i0 = 0 to 10 { + affine.for %i1 = 0 to 16 step 4 { + %1 = affine.vload %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>, vector<4xf32> + affine.vstore %1, %0[%i0 + 3, %i1 + 7] : vector<4xf32>, memref<100x100xf32> +// CHECK: %[[buf:.*]] = alloc +// CHECK-NEXT: affine.for %[[i0:.*]] = 0 +// CHECK-NEXT: affine.for %[[i1:.*]] = 0 +// CHECK-NEXT: %[[val:.*]] = affine.vload %{{.*}}[%{{.*}} + 3, %{{.*}} + 7] : memref<100x100xf32>, vector<4xf32> +// CHECK-NEXT: affine.vstore %[[val]], %[[buf]][%[[i0]] + 3, %[[i1]] + 7] : vector<4xf32>, memref<100x100xf32> + } + } + return +} + +