| //===- VectorToSCF.cpp - Conversion from Vector to mix of SCF and Std -----===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // This file implements target-dependent lowering of vector transfer operations. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include <type_traits> |
| |
| #include "mlir/Conversion/VectorToSCF/VectorToSCF.h" |
| |
| #include "../PassDetail.h" |
| #include "mlir/Dialect/Affine/EDSC/Intrinsics.h" |
| #include "mlir/Dialect/MemRef/EDSC/Intrinsics.h" |
| #include "mlir/Dialect/SCF/EDSC/Builders.h" |
| #include "mlir/Dialect/SCF/EDSC/Intrinsics.h" |
| #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h" |
| #include "mlir/Dialect/Vector/EDSC/Intrinsics.h" |
| #include "mlir/Dialect/Vector/VectorOps.h" |
| #include "mlir/Dialect/Vector/VectorUtils.h" |
| #include "mlir/IR/AffineExpr.h" |
| #include "mlir/IR/AffineMap.h" |
| #include "mlir/IR/Builders.h" |
| #include "mlir/IR/Matchers.h" |
| #include "mlir/Pass/Pass.h" |
| #include "mlir/Transforms/GreedyPatternRewriteDriver.h" |
| #include "mlir/Transforms/Passes.h" |
| |
| using namespace mlir; |
| using namespace mlir::edsc; |
| using namespace mlir::edsc::intrinsics; |
| using vector::TransferReadOp; |
| using vector::TransferWriteOp; |
| |
| // Return a list of Values that correspond to multiple AffineApplyOp, one for |
| // each result of `map`. Each `expr` in `map` is canonicalized and folded |
| // greedily according to its operands. |
| // TODO: factor out in a common location that both linalg and vector can use. |
| static SmallVector<Value, 4> |
| applyMapToValues(OpBuilder &b, Location loc, AffineMap map, ValueRange values) { |
| SmallVector<Value, 4> res; |
| res.reserve(map.getNumResults()); |
| unsigned numDims = map.getNumDims(), numSym = map.getNumSymbols(); |
| // For each `expr` in `map`, applies the `expr` to the values extracted from |
| // ranges. If the resulting application can be folded into a Value, the |
| // folding occurs eagerly. Otherwise, an affine.apply operation is emitted. |
| for (auto expr : map.getResults()) { |
| AffineMap map = AffineMap::get(numDims, numSym, expr); |
| SmallVector<Value, 4> operands(values.begin(), values.end()); |
| fullyComposeAffineMapAndOperands(&map, &operands); |
| canonicalizeMapAndOperands(&map, &operands); |
| res.push_back(b.createOrFold<AffineApplyOp>(loc, map, operands)); |
| } |
| return res; |
| } |
| |
| namespace { |
| /// Helper class captures the common information needed to lower N>1-D vector |
| /// transfer operations (read and write). |
| /// On construction, this class opens an edsc::ScopedContext for simpler IR |
| /// manipulation. |
| /// In pseudo-IR, for an n-D vector_transfer_read such as: |
| /// |
| /// ``` |
| /// vector_transfer_read(%m, %offsets, identity_map, %fill) : |
| /// memref<(leading_dims) x (major_dims) x (minor_dims) x type>, |
| /// vector<(major_dims) x (minor_dims) x type> |
| /// ``` |
| /// |
| /// where rank(minor_dims) is the lower-level vector rank (e.g. 1 for LLVM or |
| /// higher). |
| /// |
| /// This is the entry point to emitting pseudo-IR resembling: |
| /// |
| /// ``` |
| /// %tmp = alloc(): memref<(major_dims) x vector<minor_dim x type>> |
| /// for (%ivs_major, {0}, {vector_shape}, {1}) { // (N-1)-D loop nest |
| /// if (any_of(%ivs_major + %offsets, <, major_dims)) { |
| /// %v = vector_transfer_read( |
| /// {%offsets_leading, %ivs_major + %offsets_major, %offsets_minor}, |
| /// %ivs_minor): |
| /// memref<(leading_dims) x (major_dims) x (minor_dims) x type>, |
| /// vector<(minor_dims) x type>; |
| /// store(%v, %tmp); |
| /// } else { |
| /// %v = splat(vector<(minor_dims) x type>, %fill) |
| /// store(%v, %tmp, %ivs_major); |
| /// } |
| /// } |
| /// %res = load(%tmp, %0): memref<(major_dims) x vector<minor_dim x type>>): |
| // vector<(major_dims) x (minor_dims) x type> |
| /// ``` |
| /// |
| template <typename ConcreteOp> |
| class NDTransferOpHelper { |
| public: |
| NDTransferOpHelper(PatternRewriter &rewriter, ConcreteOp xferOp, |
| const VectorTransferToSCFOptions &options) |
| : rewriter(rewriter), options(options), loc(xferOp.getLoc()), |
| scope(std::make_unique<ScopedContext>(rewriter, loc)), xferOp(xferOp), |
| op(xferOp.getOperation()) { |
| vectorType = xferOp.getVectorType(); |
| // TODO: when we go to k > 1-D vectors adapt minorRank. |
| minorRank = 1; |
| majorRank = vectorType.getRank() - minorRank; |
| leadingRank = xferOp.getLeadingShapedRank(); |
| majorVectorType = |
| VectorType::get(vectorType.getShape().take_front(majorRank), |
| vectorType.getElementType()); |
| minorVectorType = |
| VectorType::get(vectorType.getShape().take_back(minorRank), |
| vectorType.getElementType()); |
| /// Memref of minor vector type is used for individual transfers. |
| memRefMinorVectorType = MemRefType::get( |
| majorVectorType.getShape(), minorVectorType, {}, |
| xferOp.getShapedType().template cast<MemRefType>().getMemorySpace()); |
| } |
| |
| LogicalResult doReplace(); |
| |
| private: |
| /// Creates the loop nest on the "major" dimensions and calls the |
| /// `loopBodyBuilder` lambda in the context of the loop nest. |
| void |
| emitLoops(llvm::function_ref<void(ValueRange, ValueRange, ValueRange, |
| ValueRange, const MemRefBoundsCapture &)> |
| loopBodyBuilder); |
| |
| /// Common state to lower vector transfer ops. |
| PatternRewriter &rewriter; |
| const VectorTransferToSCFOptions &options; |
| Location loc; |
| std::unique_ptr<ScopedContext> scope; |
| ConcreteOp xferOp; |
| Operation *op; |
| // A vector transfer copies data between: |
| // - memref<(leading_dims) x (major_dims) x (minor_dims) x type> |
| // - vector<(major_dims) x (minor_dims) x type> |
| unsigned minorRank; // for now always 1 |
| unsigned majorRank; // vector rank - minorRank |
| unsigned leadingRank; // memref rank - vector rank |
| VectorType vectorType; // vector<(major_dims) x (minor_dims) x type> |
| VectorType majorVectorType; // vector<(major_dims) x type> |
| VectorType minorVectorType; // vector<(minor_dims) x type> |
| MemRefType memRefMinorVectorType; // memref<vector<(minor_dims) x type>> |
| }; |
| |
| template <typename ConcreteOp> |
| void NDTransferOpHelper<ConcreteOp>::emitLoops( |
| llvm::function_ref<void(ValueRange, ValueRange, ValueRange, ValueRange, |
| const MemRefBoundsCapture &)> |
| loopBodyBuilder) { |
| /// Loop nest operates on the major dimensions |
| MemRefBoundsCapture memrefBoundsCapture(xferOp.source()); |
| |
| if (options.unroll) { |
| auto shape = majorVectorType.getShape(); |
| auto strides = computeStrides(shape); |
| unsigned numUnrolledInstances = computeMaxLinearIndex(shape); |
| ValueRange indices(xferOp.indices()); |
| for (unsigned idx = 0; idx < numUnrolledInstances; ++idx) { |
| SmallVector<int64_t, 4> offsets = delinearize(strides, idx); |
| SmallVector<Value, 4> offsetValues = |
| llvm::to_vector<4>(llvm::map_range(offsets, [](int64_t off) -> Value { |
| return std_constant_index(off); |
| })); |
| loopBodyBuilder(offsetValues, indices.take_front(leadingRank), |
| indices.drop_front(leadingRank).take_front(majorRank), |
| indices.take_back(minorRank), memrefBoundsCapture); |
| } |
| } else { |
| VectorBoundsCapture vectorBoundsCapture(majorVectorType); |
| auto majorLbs = vectorBoundsCapture.getLbs(); |
| auto majorUbs = vectorBoundsCapture.getUbs(); |
| auto majorSteps = vectorBoundsCapture.getSteps(); |
| affineLoopNestBuilder( |
| majorLbs, majorUbs, majorSteps, [&](ValueRange majorIvs) { |
| ValueRange indices(xferOp.indices()); |
| loopBodyBuilder(majorIvs, indices.take_front(leadingRank), |
| indices.drop_front(leadingRank).take_front(majorRank), |
| indices.take_back(minorRank), memrefBoundsCapture); |
| }); |
| } |
| } |
| |
| static Optional<int64_t> extractConstantIndex(Value v) { |
| if (auto cstOp = v.getDefiningOp<ConstantIndexOp>()) |
| return cstOp.getValue(); |
| if (auto affineApplyOp = v.getDefiningOp<AffineApplyOp>()) |
| if (affineApplyOp.getAffineMap().isSingleConstant()) |
| return affineApplyOp.getAffineMap().getSingleConstantResult(); |
| return None; |
| } |
| |
| // Missing foldings of scf.if make it necessary to perform poor man's folding |
| // eagerly, especially in the case of unrolling. In the future, this should go |
| // away once scf.if folds properly. |
| static Value onTheFlyFoldSLT(Value v, Value ub) { |
| using namespace mlir::edsc::op; |
| auto maybeCstV = extractConstantIndex(v); |
| auto maybeCstUb = extractConstantIndex(ub); |
| if (maybeCstV && maybeCstUb && *maybeCstV < *maybeCstUb) |
| return Value(); |
| return slt(v, ub); |
| } |
| |
| /// 1. Compute the indexings `majorIvs + majorOffsets` and save them in |
| /// `majorIvsPlusOffsets`. |
| /// 2. Return a value of i1 that determines whether the first |
| /// `majorIvs.rank()` |
| /// dimensions `majorIvs + majorOffsets` are all within `memrefBounds`. |
| static Value |
| emitInBoundsCondition(PatternRewriter &rewriter, |
| VectorTransferOpInterface xferOp, unsigned leadingRank, |
| ValueRange majorIvs, ValueRange majorOffsets, |
| const MemRefBoundsCapture &memrefBounds, |
| SmallVectorImpl<Value> &majorIvsPlusOffsets) { |
| Value inBoundsCondition; |
| majorIvsPlusOffsets.reserve(majorIvs.size()); |
| unsigned idx = 0; |
| SmallVector<Value, 4> bounds = |
| applyMapToValues(rewriter, xferOp.getLoc(), xferOp.permutation_map(), |
| memrefBounds.getUbs()); |
| for (auto it : llvm::zip(majorIvs, majorOffsets, bounds)) { |
| Value iv = std::get<0>(it), off = std::get<1>(it), ub = std::get<2>(it); |
| using namespace mlir::edsc::op; |
| majorIvsPlusOffsets.push_back(iv + off); |
| if (!xferOp.isDimInBounds(leadingRank + idx)) { |
| Value inBoundsCond = onTheFlyFoldSLT(majorIvsPlusOffsets.back(), ub); |
| if (inBoundsCond) |
| inBoundsCondition = (inBoundsCondition) |
| ? (inBoundsCondition && inBoundsCond) |
| : inBoundsCond; |
| } |
| ++idx; |
| } |
| return inBoundsCondition; |
| } |
| |
| // TODO: Parallelism and threadlocal considerations. |
| static Value setAllocAtFunctionEntry(MemRefType memRefMinorVectorType, |
| Operation *op) { |
| auto &b = ScopedContext::getBuilderRef(); |
| OpBuilder::InsertionGuard guard(b); |
| Operation *scope = |
| op->getParentWithTrait<OpTrait::AutomaticAllocationScope>(); |
| assert(scope && "Expected op to be inside automatic allocation scope"); |
| b.setInsertionPointToStart(&scope->getRegion(0).front()); |
| Value res = memref_alloca(memRefMinorVectorType); |
| return res; |
| } |
| |
| template <> |
| LogicalResult NDTransferOpHelper<TransferReadOp>::doReplace() { |
| Value alloc, result; |
| if (options.unroll) |
| result = std_splat(vectorType, xferOp.padding()); |
| else |
| alloc = setAllocAtFunctionEntry(memRefMinorVectorType, op); |
| |
| emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets, |
| ValueRange majorOffsets, ValueRange minorOffsets, |
| const MemRefBoundsCapture &memrefBounds) { |
| /// Lambda to load 1-D vector in the current loop ivs + offset context. |
| auto load1DVector = [&](ValueRange majorIvsPlusOffsets) -> Value { |
| SmallVector<Value, 8> indexing; |
| indexing.reserve(leadingRank + majorRank + minorRank); |
| indexing.append(leadingOffsets.begin(), leadingOffsets.end()); |
| indexing.append(majorIvsPlusOffsets.begin(), majorIvsPlusOffsets.end()); |
| indexing.append(minorOffsets.begin(), minorOffsets.end()); |
| Value memref = xferOp.source(); |
| auto map = |
| getTransferMinorIdentityMap(xferOp.getShapedType(), minorVectorType); |
| ArrayAttr inBounds; |
| if (xferOp.isDimInBounds(xferOp.getVectorType().getRank() - 1)) { |
| OpBuilder &b = ScopedContext::getBuilderRef(); |
| inBounds = b.getBoolArrayAttr({true}); |
| } |
| return vector_transfer_read(minorVectorType, memref, indexing, |
| AffineMapAttr::get(map), xferOp.padding(), |
| inBounds); |
| }; |
| |
| // 1. Compute the inBoundsCondition in the current loops ivs + offset |
| // context. |
| SmallVector<Value, 4> majorIvsPlusOffsets; |
| Value inBoundsCondition = emitInBoundsCondition( |
| rewriter, cast<VectorTransferOpInterface>(xferOp.getOperation()), |
| leadingRank, majorIvs, majorOffsets, memrefBounds, majorIvsPlusOffsets); |
| |
| if (inBoundsCondition) { |
| // 2. If the condition is not null, we need an IfOp, which may yield |
| // if `options.unroll` is true. |
| SmallVector<Type, 1> resultType; |
| if (options.unroll) |
| resultType.push_back(vectorType); |
| |
| // 3. If in-bounds, progressively lower to a 1-D transfer read, otherwise |
| // splat a 1-D vector. |
| ValueRange ifResults = conditionBuilder( |
| resultType, inBoundsCondition, |
| [&]() -> scf::ValueVector { |
| Value vector = load1DVector(majorIvsPlusOffsets); |
| // 3.a. If `options.unroll` is true, insert the 1-D vector in the |
| // aggregate. We must yield and merge with the `else` branch. |
| if (options.unroll) { |
| vector = vector_insert(vector, result, majorIvs); |
| return {vector}; |
| } |
| // 3.b. Otherwise, just go through the temporary `alloc`. |
| memref_store(vector, alloc, majorIvs); |
| return {}; |
| }, |
| [&]() -> scf::ValueVector { |
| Value vector = std_splat(minorVectorType, xferOp.padding()); |
| // 3.c. If `options.unroll` is true, insert the 1-D vector in the |
| // aggregate. We must yield and merge with the `then` branch. |
| if (options.unroll) { |
| vector = vector_insert(vector, result, majorIvs); |
| return {vector}; |
| } |
| // 3.d. Otherwise, just go through the temporary `alloc`. |
| memref_store(vector, alloc, majorIvs); |
| return {}; |
| }); |
| |
| if (!resultType.empty()) |
| result = *ifResults.begin(); |
| } else { |
| // 4. Guaranteed in-bounds, progressively lower to a 1-D transfer read. |
| Value loaded1D = load1DVector(majorIvsPlusOffsets); |
| // 5.a. If `options.unroll` is true, insert the 1-D vector in the |
| // aggregate. |
| if (options.unroll) |
| result = vector_insert(loaded1D, result, majorIvs); |
| // 5.b. Otherwise, just go through the temporary `alloc`. |
| else |
| memref_store(loaded1D, alloc, majorIvs); |
| } |
| }); |
| |
| assert((!options.unroll ^ (bool)result) && |
| "Expected resulting Value iff unroll"); |
| if (!result) |
| result = |
| memref_load(vector_type_cast(MemRefType::get({}, vectorType), alloc)); |
| rewriter.replaceOp(op, result); |
| |
| return success(); |
| } |
| |
| template <> |
| LogicalResult NDTransferOpHelper<TransferWriteOp>::doReplace() { |
| Value alloc; |
| if (!options.unroll) { |
| alloc = setAllocAtFunctionEntry(memRefMinorVectorType, op); |
| memref_store(xferOp.vector(), |
| vector_type_cast(MemRefType::get({}, vectorType), alloc)); |
| } |
| |
| emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets, |
| ValueRange majorOffsets, ValueRange minorOffsets, |
| const MemRefBoundsCapture &memrefBounds) { |
| // Lower to 1-D vector_transfer_write and let recursion handle it. |
| auto emitTransferWrite = [&](ValueRange majorIvsPlusOffsets) { |
| SmallVector<Value, 8> indexing; |
| indexing.reserve(leadingRank + majorRank + minorRank); |
| indexing.append(leadingOffsets.begin(), leadingOffsets.end()); |
| indexing.append(majorIvsPlusOffsets.begin(), majorIvsPlusOffsets.end()); |
| indexing.append(minorOffsets.begin(), minorOffsets.end()); |
| Value result; |
| // If `options.unroll` is true, extract the 1-D vector from the |
| // aggregate. |
| if (options.unroll) |
| result = vector_extract(xferOp.vector(), majorIvs); |
| else |
| result = memref_load(alloc, majorIvs); |
| auto map = |
| getTransferMinorIdentityMap(xferOp.getShapedType(), minorVectorType); |
| ArrayAttr inBounds; |
| if (xferOp.isDimInBounds(xferOp.getVectorType().getRank() - 1)) { |
| OpBuilder &b = ScopedContext::getBuilderRef(); |
| inBounds = b.getBoolArrayAttr({true}); |
| } |
| vector_transfer_write(result, xferOp.source(), indexing, |
| AffineMapAttr::get(map), inBounds); |
| }; |
| |
| // 1. Compute the inBoundsCondition in the current loops ivs + offset |
| // context. |
| SmallVector<Value, 4> majorIvsPlusOffsets; |
| Value inBoundsCondition = emitInBoundsCondition( |
| rewriter, cast<VectorTransferOpInterface>(xferOp.getOperation()), |
| leadingRank, majorIvs, majorOffsets, memrefBounds, majorIvsPlusOffsets); |
| |
| if (inBoundsCondition) { |
| // 2.a. If the condition is not null, we need an IfOp, to write |
| // conditionally. Progressively lower to a 1-D transfer write. |
| conditionBuilder(inBoundsCondition, |
| [&] { emitTransferWrite(majorIvsPlusOffsets); }); |
| } else { |
| // 2.b. Guaranteed in-bounds. Progressively lower to a 1-D transfer write. |
| emitTransferWrite(majorIvsPlusOffsets); |
| } |
| }); |
| |
| rewriter.eraseOp(op); |
| |
| return success(); |
| } |
| |
| } // namespace |
| |
| /// Analyzes the `transfer` to find an access dimension along the fastest remote |
| /// MemRef dimension. If such a dimension with coalescing properties is found, |
| /// `pivs` and `vectorBoundsCapture` are swapped so that the invocation of |
| /// LoopNestBuilder captures it in the innermost loop. |
| template <typename TransferOpTy> |
| static int computeCoalescedIndex(TransferOpTy transfer) { |
| // rank of the remote memory access, coalescing behavior occurs on the |
| // innermost memory dimension. |
| auto remoteRank = transfer.getShapedType().getRank(); |
| // Iterate over the results expressions of the permutation map to determine |
| // the loop order for creating pointwise copies between remote and local |
| // memories. |
| int coalescedIdx = -1; |
| auto exprs = transfer.permutation_map().getResults(); |
| for (auto en : llvm::enumerate(exprs)) { |
| auto dim = en.value().template dyn_cast<AffineDimExpr>(); |
| if (!dim) { |
| continue; |
| } |
| auto memRefDim = dim.getPosition(); |
| if (memRefDim == remoteRank - 1) { |
| // memRefDim has coalescing properties, it should be swapped in the last |
| // position. |
| assert(coalescedIdx == -1 && "Unexpected > 1 coalesced indices"); |
| coalescedIdx = en.index(); |
| } |
| } |
| return coalescedIdx; |
| } |
| |
| template <typename TransferOpTy> |
| VectorTransferRewriter<TransferOpTy>::VectorTransferRewriter( |
| VectorTransferToSCFOptions options, MLIRContext *context) |
| : RewritePattern(TransferOpTy::getOperationName(), 1, context), |
| options(options) {} |
| |
| /// Used for staging the transfer in a local buffer. |
| template <typename TransferOpTy> |
| MemRefType VectorTransferRewriter<TransferOpTy>::tmpMemRefType( |
| TransferOpTy transfer) const { |
| auto vectorType = transfer.getVectorType(); |
| return MemRefType::get(vectorType.getShape().drop_back(), |
| VectorType::get(vectorType.getShape().take_back(), |
| vectorType.getElementType()), |
| {}, 0); |
| } |
| |
| static void emitWithBoundsChecks( |
| PatternRewriter &rewriter, VectorTransferOpInterface transfer, |
| ValueRange ivs, const MemRefBoundsCapture &memRefBoundsCapture, |
| function_ref<void(ArrayRef<Value>)> inBoundsFun, |
| function_ref<void(ArrayRef<Value>)> outOfBoundsFun = nullptr) { |
| // Permute the incoming indices according to the permutation map. |
| SmallVector<Value, 4> indices = |
| applyMapToValues(rewriter, transfer.getLoc(), transfer.permutation_map(), |
| transfer.indices()); |
| |
| // Generate a bounds check if necessary. |
| SmallVector<Value, 4> majorIvsPlusOffsets; |
| Value inBoundsCondition = |
| emitInBoundsCondition(rewriter, transfer, 0, ivs, indices, |
| memRefBoundsCapture, majorIvsPlusOffsets); |
| |
| // Apply the permutation map to the ivs. The permutation map may not use all |
| // the inputs. |
| SmallVector<Value, 4> scalarAccessExprs(transfer.indices().size()); |
| for (unsigned memRefDim = 0; memRefDim < transfer.indices().size(); |
| ++memRefDim) { |
| // Linear search on a small number of entries. |
| int loopIndex = -1; |
| auto exprs = transfer.permutation_map().getResults(); |
| for (auto en : llvm::enumerate(exprs)) { |
| auto expr = en.value(); |
| auto dim = expr.dyn_cast<AffineDimExpr>(); |
| // Sanity check. |
| assert((dim || expr.cast<AffineConstantExpr>().getValue() == 0) && |
| "Expected dim or 0 in permutationMap"); |
| if (dim && memRefDim == dim.getPosition()) { |
| loopIndex = en.index(); |
| break; |
| } |
| } |
| |
| using namespace edsc::op; |
| auto i = transfer.indices()[memRefDim]; |
| scalarAccessExprs[memRefDim] = loopIndex < 0 ? i : i + ivs[loopIndex]; |
| } |
| |
| if (inBoundsCondition) |
| conditionBuilder( |
| /* scf.if */ inBoundsCondition, // { |
| [&] { inBoundsFun(scalarAccessExprs); }, |
| // } else { |
| outOfBoundsFun ? [&] { outOfBoundsFun(scalarAccessExprs); } |
| : function_ref<void()>() |
| // } |
| ); |
| else |
| inBoundsFun(scalarAccessExprs); |
| } |
| |
| namespace mlir { |
| |
| /// Lowers TransferReadOp into a combination of: |
| /// 1. local memory allocation; |
| /// 2. perfect loop nest over: |
| /// a. scalar load from local buffers (viewed as a scalar memref); |
| /// a. scalar store to original memref (with padding). |
| /// 3. vector_load from local buffer (viewed as a memref<1 x vector>); |
| /// 4. local memory deallocation. |
| /// |
| /// Lowers the data transfer part of a TransferReadOp while ensuring no |
| /// out-of-bounds accesses are possible. Out-of-bounds behavior is handled by |
| /// padding. |
| |
| /// Performs the rewrite. |
| template <> |
| LogicalResult VectorTransferRewriter<TransferReadOp>::matchAndRewrite( |
| Operation *op, PatternRewriter &rewriter) const { |
| using namespace mlir::edsc::op; |
| |
| TransferReadOp transfer = cast<TransferReadOp>(op); |
| if (transfer.mask()) |
| return failure(); |
| auto memRefType = transfer.getShapedType().dyn_cast<MemRefType>(); |
| if (!memRefType) |
| return failure(); |
| // Fall back to a loop if the fastest varying stride is not 1 or it is |
| // permuted. |
| int64_t offset; |
| SmallVector<int64_t, 4> strides; |
| auto successStrides = getStridesAndOffset(memRefType, strides, offset); |
| if (succeeded(successStrides) && strides.back() == 1 && |
| transfer.permutation_map().isMinorIdentity()) { |
| // If > 1D, emit a bunch of loops around 1-D vector transfers. |
| if (transfer.getVectorType().getRank() > 1) |
| return NDTransferOpHelper<TransferReadOp>(rewriter, transfer, options) |
| .doReplace(); |
| // If 1-D this is now handled by the target-specific lowering. |
| if (transfer.getVectorType().getRank() == 1) |
| return failure(); |
| } |
| |
| // Conservative lowering to scalar load / stores. |
| // 1. Setup all the captures. |
| ScopedContext scope(rewriter, transfer.getLoc()); |
| MemRefIndexedValue remote(transfer.source()); |
| MemRefBoundsCapture memRefBoundsCapture(transfer.source()); |
| VectorBoundsCapture vectorBoundsCapture(transfer.vector()); |
| int coalescedIdx = computeCoalescedIndex(transfer); |
| // Swap the vectorBoundsCapture which will reorder loop bounds. |
| if (coalescedIdx >= 0) |
| vectorBoundsCapture.swapRanges(vectorBoundsCapture.rank() - 1, |
| coalescedIdx); |
| |
| auto lbs = vectorBoundsCapture.getLbs(); |
| auto ubs = vectorBoundsCapture.getUbs(); |
| SmallVector<Value, 8> steps; |
| steps.reserve(vectorBoundsCapture.getSteps().size()); |
| for (auto step : vectorBoundsCapture.getSteps()) |
| steps.push_back(std_constant_index(step)); |
| |
| // 2. Emit alloc-copy-load-dealloc. |
| MLIRContext *ctx = op->getContext(); |
| Value tmp = setAllocAtFunctionEntry(tmpMemRefType(transfer), transfer); |
| MemRefIndexedValue local(tmp); |
| loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) { |
| auto ivsStorage = llvm::to_vector<8>(loopIvs); |
| // Swap the ivs which will reorder memory accesses. |
| if (coalescedIdx >= 0) |
| std::swap(ivsStorage.back(), ivsStorage[coalescedIdx]); |
| |
| ArrayRef<Value> ivs(ivsStorage); |
| Value pos = std_index_cast(IntegerType::get(ctx, 32), ivs.back()); |
| Value inVector = local(ivs.drop_back()); |
| auto loadValue = [&](ArrayRef<Value> indices) { |
| Value vector = vector_insert_element(remote(indices), inVector, pos); |
| local(ivs.drop_back()) = vector; |
| }; |
| auto loadPadding = [&](ArrayRef<Value>) { |
| Value vector = vector_insert_element(transfer.padding(), inVector, pos); |
| local(ivs.drop_back()) = vector; |
| }; |
| emitWithBoundsChecks( |
| rewriter, cast<VectorTransferOpInterface>(transfer.getOperation()), ivs, |
| memRefBoundsCapture, loadValue, loadPadding); |
| }); |
| Value vectorValue = memref_load(vector_type_cast(tmp)); |
| |
| // 3. Propagate. |
| rewriter.replaceOp(op, vectorValue); |
| return success(); |
| } |
| |
| /// Lowers TransferWriteOp into a combination of: |
| /// 1. local memory allocation; |
| /// 2. vector_store to local buffer (viewed as a memref<1 x vector>); |
| /// 3. perfect loop nest over: |
| /// a. scalar load from local buffers (viewed as a scalar memref); |
| /// a. scalar store to original memref (if in bounds). |
| /// 4. local memory deallocation. |
| /// |
| /// More specifically, lowers the data transfer part while ensuring no |
| /// out-of-bounds accesses are possible. |
| template <> |
| LogicalResult VectorTransferRewriter<TransferWriteOp>::matchAndRewrite( |
| Operation *op, PatternRewriter &rewriter) const { |
| using namespace edsc::op; |
| |
| TransferWriteOp transfer = cast<TransferWriteOp>(op); |
| if (transfer.mask()) |
| return failure(); |
| auto memRefType = transfer.getShapedType().template dyn_cast<MemRefType>(); |
| if (!memRefType) |
| return failure(); |
| |
| // Fall back to a loop if the fastest varying stride is not 1 or it is |
| // permuted. |
| int64_t offset; |
| SmallVector<int64_t, 4> strides; |
| auto successStrides = getStridesAndOffset(memRefType, strides, offset); |
| if (succeeded(successStrides) && strides.back() == 1 && |
| transfer.permutation_map().isMinorIdentity()) { |
| // If > 1D, emit a bunch of loops around 1-D vector transfers. |
| if (transfer.getVectorType().getRank() > 1) |
| return NDTransferOpHelper<TransferWriteOp>(rewriter, transfer, options) |
| .doReplace(); |
| // If 1-D this is now handled by the target-specific lowering. |
| if (transfer.getVectorType().getRank() == 1) |
| return failure(); |
| } |
| |
| // 1. Setup all the captures. |
| ScopedContext scope(rewriter, transfer.getLoc()); |
| MemRefIndexedValue remote(transfer.source()); |
| MemRefBoundsCapture memRefBoundsCapture(transfer.source()); |
| Value vectorValue(transfer.vector()); |
| VectorBoundsCapture vectorBoundsCapture(transfer.vector()); |
| int coalescedIdx = computeCoalescedIndex(transfer); |
| // Swap the vectorBoundsCapture which will reorder loop bounds. |
| if (coalescedIdx >= 0) |
| vectorBoundsCapture.swapRanges(vectorBoundsCapture.rank() - 1, |
| coalescedIdx); |
| |
| auto lbs = vectorBoundsCapture.getLbs(); |
| auto ubs = vectorBoundsCapture.getUbs(); |
| SmallVector<Value, 8> steps; |
| steps.reserve(vectorBoundsCapture.getSteps().size()); |
| for (auto step : vectorBoundsCapture.getSteps()) |
| steps.push_back(std_constant_index(step)); |
| |
| // 2. Emit alloc-store-copy-dealloc. |
| Value tmp = setAllocAtFunctionEntry(tmpMemRefType(transfer), transfer); |
| MemRefIndexedValue local(tmp); |
| Value vec = vector_type_cast(tmp); |
| memref_store(vectorValue, vec); |
| loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) { |
| auto ivsStorage = llvm::to_vector<8>(loopIvs); |
| // Swap the ivsStorage which will reorder memory accesses. |
| if (coalescedIdx >= 0) |
| std::swap(ivsStorage.back(), ivsStorage[coalescedIdx]); |
| |
| ArrayRef<Value> ivs(ivsStorage); |
| Value pos = |
| std_index_cast(IntegerType::get(op->getContext(), 32), ivs.back()); |
| auto storeValue = [&](ArrayRef<Value> indices) { |
| Value scalar = vector_extract_element(local(ivs.drop_back()), pos); |
| remote(indices) = scalar; |
| }; |
| emitWithBoundsChecks( |
| rewriter, cast<VectorTransferOpInterface>(transfer.getOperation()), ivs, |
| memRefBoundsCapture, storeValue); |
| }); |
| |
| // 3. Erase. |
| rewriter.eraseOp(op); |
| return success(); |
| } |
| |
| void populateVectorToSCFConversionPatterns( |
| RewritePatternSet &patterns, const VectorTransferToSCFOptions &options) { |
| patterns.add<VectorTransferRewriter<vector::TransferReadOp>, |
| VectorTransferRewriter<vector::TransferWriteOp>>( |
| options, patterns.getContext()); |
| } |
| |
| } // namespace mlir |
| |
| namespace { |
| |
| struct ConvertVectorToSCFPass |
| : public ConvertVectorToSCFBase<ConvertVectorToSCFPass> { |
| ConvertVectorToSCFPass() = default; |
| ConvertVectorToSCFPass(const VectorTransferToSCFOptions &options) { |
| this->fullUnroll = options.unroll; |
| } |
| |
| void runOnFunction() override { |
| RewritePatternSet patterns(getFunction().getContext()); |
| populateVectorToSCFConversionPatterns( |
| patterns, VectorTransferToSCFOptions().setUnroll(fullUnroll)); |
| (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns)); |
| } |
| }; |
| |
| } // namespace |
| |
| std::unique_ptr<Pass> |
| mlir::createConvertVectorToSCFPass(const VectorTransferToSCFOptions &options) { |
| return std::make_unique<ConvertVectorToSCFPass>(options); |
| } |
