1 //===- AffineDataCopyGeneration.cpp - Explicit memref copying pass ------*-===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This file implements a pass to automatically promote accessed memref regions 10 // to buffers in a faster memory space that is explicitly managed, with the 11 // necessary data movement operations performed through either regular 12 // point-wise load/store's or DMAs. Such explicit copying (also referred to as 13 // array packing/unpacking in the literature), when done on arrays that exhibit 14 // reuse, results in near elimination of conflict misses, TLB misses, reduced 15 // use of hardware prefetch streams, and reduced false sharing. It is also 16 // necessary for hardware that explicitly managed levels in the memory 17 // hierarchy, and where DMAs may have to be used. This optimization is often 18 // performed on already tiled code. 19 // 20 //===----------------------------------------------------------------------===// 21 22 #include "PassDetail.h" 23 #include "mlir/Dialect/Affine/Analysis/Utils.h" 24 #include "mlir/Dialect/Affine/IR/AffineOps.h" 25 #include "mlir/Dialect/Affine/LoopUtils.h" 26 #include "mlir/Dialect/Affine/Passes.h" 27 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" 28 #include "mlir/Dialect/MemRef/IR/MemRef.h" 29 #include "mlir/Transforms/GreedyPatternRewriteDriver.h" 30 #include "llvm/ADT/MapVector.h" 31 #include "llvm/Support/CommandLine.h" 32 #include "llvm/Support/Debug.h" 33 #include <algorithm> 34 35 #define DEBUG_TYPE "affine-data-copy-generate" 36 37 using namespace mlir; 38 39 namespace { 40 41 /// Replaces all loads and stores on memref's living in 'slowMemorySpace' by 42 /// introducing copy operations to transfer data into `fastMemorySpace` and 43 /// rewriting the original load's/store's to instead load/store from the 44 /// allocated fast memory buffers. Additional options specify the identifier 45 /// corresponding to the fast memory space and the amount of fast memory space 46 /// available. The pass traverses through the nesting structure, recursing to 47 /// inner levels if necessary to determine at what depth copies need to be 48 /// placed so that the allocated buffers fit within the memory capacity 49 /// provided. 50 // TODO: We currently can't generate copies correctly when stores 51 // are strided. Check for strided stores. 52 struct AffineDataCopyGeneration 53 : public AffineDataCopyGenerationBase<AffineDataCopyGeneration> { 54 AffineDataCopyGeneration() = default; 55 explicit AffineDataCopyGeneration(unsigned slowMemorySpace, 56 unsigned fastMemorySpace, 57 unsigned tagMemorySpace, 58 int minDmaTransferSize, 59 uint64_t fastMemCapacityBytes) { 60 this->slowMemorySpace = slowMemorySpace; 61 this->fastMemorySpace = fastMemorySpace; 62 this->tagMemorySpace = tagMemorySpace; 63 this->minDmaTransferSize = minDmaTransferSize; 64 this->fastMemoryCapacity = fastMemCapacityBytes / 1024; 65 } 66 67 void runOnOperation() override; 68 void runOnBlock(Block *block, DenseSet<Operation *> ©Nests); 69 70 // Constant zero index to avoid too many duplicates. 71 Value zeroIndex = nullptr; 72 }; 73 74 } // namespace 75 76 /// Generates copies for memref's living in 'slowMemorySpace' into newly created 77 /// buffers in 'fastMemorySpace', and replaces memory operations to the former 78 /// by the latter. Only load op's handled for now. 79 /// TODO: extend this to store op's. 80 std::unique_ptr<OperationPass<func::FuncOp>> 81 mlir::createAffineDataCopyGenerationPass(unsigned slowMemorySpace, 82 unsigned fastMemorySpace, 83 unsigned tagMemorySpace, 84 int minDmaTransferSize, 85 uint64_t fastMemCapacityBytes) { 86 return std::make_unique<AffineDataCopyGeneration>( 87 slowMemorySpace, fastMemorySpace, tagMemorySpace, minDmaTransferSize, 88 fastMemCapacityBytes); 89 } 90 std::unique_ptr<OperationPass<func::FuncOp>> 91 mlir::createAffineDataCopyGenerationPass() { 92 return std::make_unique<AffineDataCopyGeneration>(); 93 } 94 95 /// Generate copies for this block. The block is partitioned into separate 96 /// ranges: each range is either a sequence of one or more operations starting 97 /// and ending with an affine load or store op, or just an affine.forop (which 98 /// could have other affine for op's nested within). 99 void AffineDataCopyGeneration::runOnBlock(Block *block, 100 DenseSet<Operation *> ©Nests) { 101 if (block->empty()) 102 return; 103 104 uint64_t fastMemCapacityBytes = 105 fastMemoryCapacity != std::numeric_limits<uint64_t>::max() 106 ? fastMemoryCapacity * 1024 107 : fastMemoryCapacity; 108 AffineCopyOptions copyOptions = {generateDma, slowMemorySpace, 109 fastMemorySpace, tagMemorySpace, 110 fastMemCapacityBytes}; 111 112 // Every affine.for op in the block starts and ends a block range for copying; 113 // in addition, a contiguous sequence of operations starting with a 114 // load/store op but not including any copy nests themselves is also 115 // identified as a copy block range. Straightline code (a contiguous chunk of 116 // operations excluding AffineForOp's) are always assumed to not exhaust 117 // memory. As a result, this approach is conservative in some cases at the 118 // moment; we do a check later and report an error with location info. 119 // TODO: An 'affine.if' operation is being treated similar to an 120 // operation. 'affine.if''s could have 'affine.for's in them; 121 // treat them separately. 122 123 // Get to the first load, store, or for op (that is not a copy nest itself). 124 auto curBegin = 125 std::find_if(block->begin(), block->end(), [&](Operation &op) { 126 return isa<AffineLoadOp, AffineStoreOp, AffineForOp>(op) && 127 copyNests.count(&op) == 0; 128 }); 129 130 // Create [begin, end) ranges. 131 auto it = curBegin; 132 while (it != block->end()) { 133 AffineForOp forOp; 134 // If you hit a non-copy for loop, we will split there. 135 if ((forOp = dyn_cast<AffineForOp>(&*it)) && copyNests.count(forOp) == 0) { 136 // Perform the copying up unti this 'for' op first. 137 (void)affineDataCopyGenerate(/*begin=*/curBegin, /*end=*/it, copyOptions, 138 /*filterMemRef=*/llvm::None, copyNests); 139 140 // Returns true if the footprint is known to exceed capacity. 141 auto exceedsCapacity = [&](AffineForOp forOp) { 142 Optional<int64_t> footprint = 143 getMemoryFootprintBytes(forOp, 144 /*memorySpace=*/0); 145 return (footprint.has_value() && 146 static_cast<uint64_t>(footprint.value()) > 147 fastMemCapacityBytes); 148 }; 149 150 // If the memory footprint of the 'affine.for' loop is higher than fast 151 // memory capacity (when provided), we recurse to copy at an inner level 152 // until we find a depth at which footprint fits in fast mem capacity. If 153 // the footprint can't be calculated, we assume for now it fits. Recurse 154 // inside if footprint for 'forOp' exceeds capacity, or when 155 // skipNonUnitStrideLoops is set and the step size is not one. 156 bool recurseInner = skipNonUnitStrideLoops ? forOp.getStep() != 1 157 : exceedsCapacity(forOp); 158 if (recurseInner) { 159 // We'll recurse and do the copies at an inner level for 'forInst'. 160 // Recurse onto the body of this loop. 161 runOnBlock(forOp.getBody(), copyNests); 162 } else { 163 // We have enough capacity, i.e., copies will be computed for the 164 // portion of the block until 'it', and for 'it', which is 'forOp'. Note 165 // that for the latter, the copies are placed just before this loop (for 166 // incoming copies) and right after (for outgoing ones). 167 168 // Inner loop copies have their own scope - we don't thus update 169 // consumed capacity. The footprint check above guarantees this inner 170 // loop's footprint fits. 171 (void)affineDataCopyGenerate(/*begin=*/it, /*end=*/std::next(it), 172 copyOptions, 173 /*filterMemRef=*/llvm::None, copyNests); 174 } 175 // Get to the next load or store op after 'forOp'. 176 curBegin = std::find_if(std::next(it), block->end(), [&](Operation &op) { 177 return isa<AffineLoadOp, AffineStoreOp, AffineForOp>(op) && 178 copyNests.count(&op) == 0; 179 }); 180 it = curBegin; 181 } else { 182 assert(copyNests.count(&*it) == 0 && 183 "all copy nests generated should have been skipped above"); 184 // We simply include this op in the current range and continue for more. 185 ++it; 186 } 187 } 188 189 // Generate the copy for the final block range. 190 if (curBegin != block->end()) { 191 // Can't be a terminator because it would have been skipped above. 192 assert(!curBegin->hasTrait<OpTrait::IsTerminator>() && 193 "can't be a terminator"); 194 // Exclude the affine.yield - hence, the std::prev. 195 (void)affineDataCopyGenerate(/*begin=*/curBegin, 196 /*end=*/std::prev(block->end()), copyOptions, 197 /*filterMemRef=*/llvm::None, copyNests); 198 } 199 } 200 201 void AffineDataCopyGeneration::runOnOperation() { 202 func::FuncOp f = getOperation(); 203 OpBuilder topBuilder(f.getBody()); 204 zeroIndex = topBuilder.create<arith::ConstantIndexOp>(f.getLoc(), 0); 205 206 // Nests that are copy-in's or copy-out's; the root AffineForOps of those 207 // nests are stored herein. 208 DenseSet<Operation *> copyNests; 209 210 // Clear recorded copy nests. 211 copyNests.clear(); 212 213 for (auto &block : f) 214 runOnBlock(&block, copyNests); 215 216 // Promote any single iteration loops in the copy nests and collect 217 // load/stores to simplify. 218 SmallVector<Operation *, 4> copyOps; 219 for (Operation *nest : copyNests) 220 // With a post order walk, the erasure of loops does not affect 221 // continuation of the walk or the collection of load/store ops. 222 nest->walk([&](Operation *op) { 223 if (auto forOp = dyn_cast<AffineForOp>(op)) 224 (void)promoteIfSingleIteration(forOp); 225 else if (isa<AffineLoadOp, AffineStoreOp>(op)) 226 copyOps.push_back(op); 227 }); 228 229 // Promoting single iteration loops could lead to simplification of 230 // contained load's/store's, and the latter could anyway also be 231 // canonicalized. 232 RewritePatternSet patterns(&getContext()); 233 AffineLoadOp::getCanonicalizationPatterns(patterns, &getContext()); 234 AffineStoreOp::getCanonicalizationPatterns(patterns, &getContext()); 235 FrozenRewritePatternSet frozenPatterns(std::move(patterns)); 236 (void)applyOpPatternsAndFold(copyOps, frozenPatterns, /*strict=*/true); 237 } 238