//===- 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 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h" #include "mlir/Dialect/Affine/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/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/Builders.h" #include "mlir/IR/Location.h" #include "mlir/IR/Matchers.h" #include "mlir/IR/OperationSupport.h" #include "mlir/IR/PatternMatch.h" #include "mlir/IR/Types.h" using namespace mlir; using namespace mlir::edsc; using namespace mlir::edsc::intrinsics; using vector::TransferReadOp; using vector::TransferWriteOp; 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> /// 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>): // vector<(major_dims) x (minor_dims) x type> /// ``` /// template class NDTransferOpHelper { public: NDTransferOpHelper(PatternRewriter &rewriter, ConcreteOp xferOp) : rewriter(rewriter), loc(xferOp.getLoc()), scope(std::make_unique(rewriter, loc)), xferOp(xferOp), op(xferOp.getOperation()) { vectorType = xferOp.getVectorType(); // TODO(ntv, ajcbik): when we go to k > 1-D vectors adapt minorRank. minorRank = 1; majorRank = vectorType.getRank() - minorRank; leadingRank = xferOp.getMemRefType().getRank() - (majorRank + minorRank); 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.getMemRefType().getMemorySpace()); } LogicalResult doReplace(); private: /// Creates the loop nest on the "major" dimensions and calls the /// `loopBodyBuilder` lambda in the context of the loop nest. template void emitLoops(Lambda loopBodyBuilder); /// Operate within the body of `emitLoops` to: /// 1. Compute the indexings `majorIvs + majorOffsets`. /// 2. Compute a boolean that determines whether the first `majorIvs.rank()` /// dimensions `majorIvs + majorOffsets` are all within `memrefBounds`. /// 3. Create an IfOp conditioned on the boolean in step 2. /// 4. Call a `thenBlockBuilder` and an `elseBlockBuilder` to append /// operations to the IfOp blocks as appropriate. template void emitInBounds(ValueRange majorIvs, ValueRange majorOffsets, MemRefBoundsCapture &memrefBounds, LambdaThen thenBlockBuilder, LambdaElse elseBlockBuilder); /// Common state to lower vector transfer ops. PatternRewriter &rewriter; Location loc; std::unique_ptr 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> }; template template void NDTransferOpHelper::emitLoops(Lambda loopBodyBuilder) { /// Loop nest operates on the major dimensions MemRefBoundsCapture memrefBoundsCapture(xferOp.memref()); VectorBoundsCapture vectorBoundsCapture(majorVectorType); auto majorLbs = vectorBoundsCapture.getLbs(); auto majorUbs = vectorBoundsCapture.getUbs(); auto majorSteps = vectorBoundsCapture.getSteps(); SmallVector majorIvs(vectorBoundsCapture.rank()); AffineLoopNestBuilder(majorIvs, majorLbs, majorUbs, majorSteps)([&] { ValueRange indices(xferOp.indices()); loopBodyBuilder(majorIvs, indices.take_front(leadingRank), indices.drop_front(leadingRank).take_front(majorRank), indices.take_back(minorRank), memrefBoundsCapture); }); } template template void NDTransferOpHelper::emitInBounds( ValueRange majorIvs, ValueRange majorOffsets, MemRefBoundsCapture &memrefBounds, LambdaThen thenBlockBuilder, LambdaElse elseBlockBuilder) { Value inBounds; SmallVector majorIvsPlusOffsets; majorIvsPlusOffsets.reserve(majorIvs.size()); unsigned idx = 0; for (auto it : llvm::zip(majorIvs, majorOffsets, memrefBounds.getUbs())) { 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.isMaskedDim(leadingRank + idx)) { Value inBounds2 = majorIvsPlusOffsets.back() < ub; inBounds = (inBounds) ? (inBounds && inBounds2) : inBounds2; } ++idx; } if (inBounds) { auto ifOp = ScopedContext::getBuilderRef().create( ScopedContext::getLocation(), TypeRange{}, inBounds, /*withElseRegion=*/std::is_same()); BlockBuilder(&ifOp.thenRegion().front(), Append())([&] { thenBlockBuilder(majorIvsPlusOffsets); }); if (std::is_same()) BlockBuilder(&ifOp.elseRegion().front(), Append())([&] { elseBlockBuilder(majorIvsPlusOffsets); }); } else { // Just build the body of the then block right here. thenBlockBuilder(majorIvsPlusOffsets); } } template <> LogicalResult NDTransferOpHelper::doReplace() { Value alloc = std_alloc(memRefMinorVectorType); emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets, ValueRange majorOffsets, ValueRange minorOffsets, MemRefBoundsCapture &memrefBounds) { // If in-bounds, index into memref and lower to 1-D transfer read. auto thenBlockBuilder = [&](ValueRange majorIvsPlusOffsets) { SmallVector 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.memref(); auto map = TransferReadOp::getTransferMinorIdentityMap( xferOp.getMemRefType(), minorVectorType); ArrayAttr masked; if (xferOp.isMaskedDim(xferOp.getVectorType().getRank() - 1)) { OpBuilder &b = ScopedContext::getBuilderRef(); masked = b.getBoolArrayAttr({true}); } auto loaded1D = vector_transfer_read(minorVectorType, memref, indexing, AffineMapAttr::get(map), xferOp.padding(), masked); // Store the 1-D vector. std_store(loaded1D, alloc, majorIvs); }; // If out-of-bounds, just store a splatted vector. auto elseBlockBuilder = [&](ValueRange majorIvsPlusOffsets) { auto vector = std_splat(minorVectorType, xferOp.padding()); std_store(vector, alloc, majorIvs); }; emitInBounds(majorIvs, majorOffsets, memrefBounds, thenBlockBuilder, elseBlockBuilder); }); Value loaded = std_load(vector_type_cast(MemRefType::get({}, vectorType), alloc)); rewriter.replaceOp(op, loaded); return success(); } template <> LogicalResult NDTransferOpHelper::doReplace() { Value alloc = std_alloc(memRefMinorVectorType); std_store(xferOp.vector(), vector_type_cast(MemRefType::get({}, vectorType), alloc)); emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets, ValueRange majorOffsets, ValueRange minorOffsets, MemRefBoundsCapture &memrefBounds) { auto thenBlockBuilder = [&](ValueRange majorIvsPlusOffsets) { SmallVector indexing; indexing.reserve(leadingRank + majorRank + minorRank); indexing.append(leadingOffsets.begin(), leadingOffsets.end()); indexing.append(majorIvsPlusOffsets.begin(), majorIvsPlusOffsets.end()); indexing.append(minorOffsets.begin(), minorOffsets.end()); // Lower to 1-D vector_transfer_write and let recursion handle it. Value loaded1D = std_load(alloc, majorIvs); auto map = TransferWriteOp::getTransferMinorIdentityMap( xferOp.getMemRefType(), minorVectorType); ArrayAttr masked; if (xferOp.isMaskedDim(xferOp.getVectorType().getRank() - 1)) { OpBuilder &b = ScopedContext::getBuilderRef(); masked = b.getBoolArrayAttr({true}); } vector_transfer_write(loaded1D, xferOp.memref(), indexing, AffineMapAttr::get(map), masked); }; // Don't write anything when out of bounds. auto elseBlockBuilder = [&](ValueRange majorIvsPlusOffsets) {}; emitInBounds(majorIvs, majorOffsets, memrefBounds, thenBlockBuilder, elseBlockBuilder); }); 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 static int computeCoalescedIndex(TransferOpTy transfer) { // rank of the remote memory access, coalescing behavior occurs on the // innermost memory dimension. auto remoteRank = transfer.getMemRefType().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(); 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; } /// Emits remote memory accesses that are clipped to the boundaries of the /// MemRef. template static SmallVector clip(TransferOpTy transfer, MemRefBoundsCapture &bounds, ArrayRef ivs) { using namespace mlir::edsc; Value zero(std_constant_index(0)), one(std_constant_index(1)); SmallVector memRefAccess(transfer.indices()); SmallVector clippedScalarAccessExprs(memRefAccess.size()); // Indices accessing to remote memory are clipped and their expressions are // returned in clippedScalarAccessExprs. for (unsigned memRefDim = 0; memRefDim < clippedScalarAccessExprs.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.template dyn_cast(); // Sanity check. assert( (dim || expr.template cast().getValue() == 0) && "Expected dim or 0 in permutationMap"); if (dim && memRefDim == dim.getPosition()) { loopIndex = en.index(); break; } } // We cannot distinguish atm between unrolled dimensions that implement // the "always full" tile abstraction and need clipping from the other // ones. So we conservatively clip everything. using namespace edsc::op; auto N = bounds.ub(memRefDim); auto i = memRefAccess[memRefDim]; if (loopIndex < 0) { auto N_minus_1 = N - one; auto select_1 = std_select(i < N, i, N_minus_1); clippedScalarAccessExprs[memRefDim] = std_select(i < zero, zero, select_1); } else { auto ii = ivs[loopIndex]; auto i_plus_ii = i + ii; auto N_minus_1 = N - one; auto select_1 = std_select(i_plus_ii < N, i_plus_ii, N_minus_1); clippedScalarAccessExprs[memRefDim] = std_select(i_plus_ii < zero, zero, select_1); } } return clippedScalarAccessExprs; } namespace { /// Implements lowering of TransferReadOp and TransferWriteOp to a /// proper abstraction for the hardware. /// /// For now, we only emit a simple loop nest that performs clipped pointwise /// copies from a remote to a locally allocated memory. /// /// Consider the case: /// /// ```mlir /// // Read the slice `%A[%i0, %i1:%i1+256, %i2:%i2+32]` into /// // vector<32x256xf32> and pad with %f0 to handle the boundary case: /// %f0 = constant 0.0f : f32 /// scf.for %i0 = 0 to %0 { /// scf.for %i1 = 0 to %1 step %c256 { /// scf.for %i2 = 0 to %2 step %c32 { /// %v = vector.transfer_read %A[%i0, %i1, %i2], %f0 /// {permutation_map: (d0, d1, d2) -> (d2, d1)} : /// memref, vector<32x256xf32> /// }}} /// ``` /// /// The rewriters construct loop and indices that access MemRef A in a pattern /// resembling the following (while guaranteeing an always full-tile /// abstraction): /// /// ```mlir /// scf.for %d2 = 0 to %c256 { /// scf.for %d1 = 0 to %c32 { /// %s = %A[%i0, %i1 + %d1, %i2 + %d2] : f32 /// %tmp[%d2, %d1] = %s /// } /// } /// ``` /// /// In the current state, only a clipping transfer is implemented by `clip`, /// which creates individual indexing expressions of the form: /// /// ```mlir-dsc /// auto condMax = i + ii < N; /// auto max = std_select(condMax, i + ii, N - one) /// auto cond = i + ii < zero; /// std_select(cond, zero, max); /// ``` /// /// In the future, clipping should not be the only way and instead we should /// load vectors + mask them. Similarly on the write side, load/mask/store for /// implementing RMW behavior. /// /// Lowers TransferOp into a combination of: /// 1. local memory allocation; /// 2. perfect loop nest over: /// a. scalar load/stores from local buffers (viewed as a scalar memref); /// a. scalar store/load to original memref (with clipping). /// 3. vector_load/store /// 4. local memory deallocation. /// Minor variations occur depending on whether a TransferReadOp or /// a TransferWriteOp is rewritten. template struct VectorTransferRewriter : public RewritePattern { explicit VectorTransferRewriter(MLIRContext *context) : RewritePattern(TransferOpTy::getOperationName(), 1, context) {} /// Used for staging the transfer in a local scalar buffer. MemRefType tmpMemRefType(TransferOpTy transfer) const { auto vectorType = transfer.getVectorType(); return MemRefType::get(vectorType.getShape(), vectorType.getElementType(), {}, 0); } /// Performs the rewrite. LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override; }; /// 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 clipping). /// 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 /// clipping. This means that a given value in memory can be read multiple /// times and concurrently. /// /// Important notes about clipping and "full-tiles only" abstraction: /// ================================================================= /// When using clipping for dealing with boundary conditions, the same edge /// value will appear multiple times (a.k.a edge padding). This is fine if the /// subsequent vector operations are all data-parallel but **is generally /// incorrect** in the presence of reductions or extract operations. /// /// More generally, clipping is a scalar abstraction that is expected to work /// fine as a baseline for CPUs and GPUs but not for vector_load and DMAs. /// To deal with real vector_load and DMAs, a "padded allocation + view" /// abstraction with the ability to read out-of-memref-bounds (but still within /// the allocated region) is necessary. /// /// Whether using scalar loops or vector_load/DMAs to perform the transfer, /// junk values will be materialized in the vectors and generally need to be /// filtered out and replaced by the "neutral element". This neutral element is /// op-dependent so, in the future, we expect to create a vector filter and /// apply it to a splatted constant vector with the proper neutral element at /// each ssa-use. This filtering is not necessary for pure data-parallel /// operations. /// /// In the case of vector_store/DMAs, Read-Modify-Write will be required, which /// also have concurrency implications. Note that by using clipped scalar stores /// in the presence of data-parallel only operations, we generate code that /// writes the same value multiple time on the edge locations. /// /// TODO(ntv): implement alternatives to clipping. /// TODO(ntv): support non-data-parallel operations. /// Performs the rewrite. template <> LogicalResult VectorTransferRewriter::matchAndRewrite( Operation *op, PatternRewriter &rewriter) const { using namespace mlir::edsc::op; TransferReadOp transfer = cast(op); if (AffineMap::isMinorIdentity(transfer.permutation_map())) { // If > 1D, emit a bunch of loops around 1-D vector transfers. if (transfer.getVectorType().getRank() > 1) return NDTransferOpHelper(rewriter, transfer).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()); StdIndexedValue remote(transfer.memref()); MemRefBoundsCapture memRefBoundsCapture(transfer.memref()); 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 steps; steps.reserve(vectorBoundsCapture.getSteps().size()); for (auto step : vectorBoundsCapture.getSteps()) steps.push_back(std_constant_index(step)); // 2. Emit alloc-copy-load-dealloc. Value tmp = std_alloc(tmpMemRefType(transfer)); StdIndexedValue local(tmp); Value vec = vector_type_cast(tmp); loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) { auto ivs = llvm::to_vector<8>(loopIvs); // Swap the ivs which will reorder memory accesses. if (coalescedIdx >= 0) std::swap(ivs.back(), ivs[coalescedIdx]); // Computes clippedScalarAccessExprs in the loop nest scope (ivs exist). local(ivs) = remote(clip(transfer, memRefBoundsCapture, ivs)); }); Value vectorValue = std_load(vec); (std_dealloc(tmp)); // vexing parse // 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 (with clipping). /// 4. local memory deallocation. /// /// More specifically, lowers the data transfer part while ensuring no /// out-of-bounds accesses are possible. Out-of-bounds behavior is handled by /// clipping. This means that a given value in memory can be written to multiple /// times and concurrently. /// /// See `Important notes about clipping and full-tiles only abstraction` in the /// description of `readClipped` above. /// /// TODO(ntv): implement alternatives to clipping. /// TODO(ntv): support non-data-parallel operations. template <> LogicalResult VectorTransferRewriter::matchAndRewrite( Operation *op, PatternRewriter &rewriter) const { using namespace edsc::op; TransferWriteOp transfer = cast(op); if (AffineMap::isMinorIdentity(transfer.permutation_map())) { // If > 1D, emit a bunch of loops around 1-D vector transfers. if (transfer.getVectorType().getRank() > 1) return NDTransferOpHelper(rewriter, transfer) .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()); StdIndexedValue remote(transfer.memref()); MemRefBoundsCapture memRefBoundsCapture(transfer.memref()); 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 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 = std_alloc(tmpMemRefType(transfer)); StdIndexedValue local(tmp); Value vec = vector_type_cast(tmp); std_store(vectorValue, vec); loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) { auto ivs = llvm::to_vector<8>(loopIvs); // Swap the ivs which will reorder memory accesses. if (coalescedIdx >= 0) std::swap(ivs.back(), ivs[coalescedIdx]); // Computes clippedScalarAccessExprs in the loop nest scope (ivs exist). remote(clip(transfer, memRefBoundsCapture, ivs)) = local(ivs); }); (std_dealloc(tmp)); // vexing parse... rewriter.eraseOp(op); return success(); } } // namespace void mlir::populateVectorToSCFConversionPatterns( OwningRewritePatternList &patterns, MLIRContext *context) { patterns.insert, VectorTransferRewriter>(context); }