235 lines · cpp
1//===----------- MultiBuffering.cpp ---------------------------------------===//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-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements multi buffering transformation.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Affine/IR/AffineOps.h"14#include "mlir/Dialect/Arith/Utils/Utils.h"15#include "mlir/Dialect/MemRef/IR/MemRef.h"16#include "mlir/Dialect/MemRef/Transforms/Transforms.h"17#include "mlir/IR/AffineExpr.h"18#include "mlir/IR/BuiltinAttributes.h"19#include "mlir/IR/Dominance.h"20#include "mlir/IR/PatternMatch.h"21#include "mlir/IR/ValueRange.h"22#include "mlir/Interfaces/LoopLikeInterface.h"23#include "llvm/ADT/STLExtras.h"24#include "llvm/Support/Debug.h"25 26using namespace mlir;27 28#define DEBUG_TYPE "memref-transforms"29#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")30#define DBGSNL() (llvm::dbgs() << "\n")31 32/// Return true if the op fully overwrite the given `buffer` value.33static bool overrideBuffer(Operation *op, Value buffer) {34 auto copyOp = dyn_cast<memref::CopyOp>(op);35 if (!copyOp)36 return false;37 return copyOp.getTarget() == buffer;38}39 40/// Replace the uses of `oldOp` with the given `val` and for subview uses41/// propagate the type change. Changing the memref type may require propagating42/// it through subview ops so we cannot just do a replaceAllUse but need to43/// propagate the type change and erase old subview ops.44static void replaceUsesAndPropagateType(RewriterBase &rewriter,45 Operation *oldOp, Value val) {46 // Iterate with early_inc to erase current user inside the loop.47 for (OpOperand &use : llvm::make_early_inc_range(oldOp->getUses())) {48 Operation *user = use.getOwner();49 if (auto subviewUse = dyn_cast<memref::SubViewOp>(user)) {50 // `subview(old_op)` is replaced by a new `subview(val)`.51 OpBuilder::InsertionGuard g(rewriter);52 rewriter.setInsertionPoint(subviewUse);53 MemRefType newType = memref::SubViewOp::inferRankReducedResultType(54 subviewUse.getType().getShape(), cast<MemRefType>(val.getType()),55 subviewUse.getStaticOffsets(), subviewUse.getStaticSizes(),56 subviewUse.getStaticStrides());57 Value newSubview = memref::SubViewOp::create(58 rewriter, subviewUse->getLoc(), newType, val,59 subviewUse.getMixedOffsets(), subviewUse.getMixedSizes(),60 subviewUse.getMixedStrides());61 62 // Ouch recursion ... is this really necessary?63 replaceUsesAndPropagateType(rewriter, subviewUse, newSubview);64 65 // Safe to erase.66 rewriter.eraseOp(subviewUse);67 continue;68 }69 // Non-subview: replace with new value.70 rewriter.startOpModification(user);71 use.set(val);72 rewriter.finalizeOpModification(user);73 }74}75 76// Transformation to do multi-buffering/array expansion to remove dependencies77// on the temporary allocation between consecutive loop iterations.78// Returns success if the transformation happened and failure otherwise.79// This is not a pattern as it requires propagating the new memref type to its80// uses and requires updating subview ops.81FailureOr<memref::AllocOp>82mlir::memref::multiBuffer(RewriterBase &rewriter, memref::AllocOp allocOp,83 unsigned multiBufferingFactor,84 bool skipOverrideAnalysis) {85 LLVM_DEBUG(DBGS() << "Start multibuffering: " << allocOp << "\n");86 DominanceInfo dom(allocOp->getParentOp());87 LoopLikeOpInterface candidateLoop;88 for (Operation *user : allocOp->getUsers()) {89 auto parentLoop = user->getParentOfType<LoopLikeOpInterface>();90 if (!parentLoop) {91 if (isa<memref::DeallocOp>(user)) {92 // Allow dealloc outside of any loop.93 // TODO: The whole precondition function here is very brittle and will94 // need to rethought an isolated into a cleaner analysis.95 continue;96 }97 LLVM_DEBUG(DBGS() << "--no parent loop -> fail\n");98 LLVM_DEBUG(DBGS() << "----due to user: " << *user << "\n");99 return failure();100 }101 if (!skipOverrideAnalysis) {102 /// Make sure there is no loop-carried dependency on the allocation.103 if (!overrideBuffer(user, allocOp.getResult())) {104 LLVM_DEBUG(DBGS() << "--Skip user: found loop-carried dependence\n");105 continue;106 }107 // If this user doesn't dominate all the other users keep looking.108 if (llvm::any_of(allocOp->getUsers(), [&](Operation *otherUser) {109 return !dom.dominates(user, otherUser);110 })) {111 LLVM_DEBUG(112 DBGS() << "--Skip user: does not dominate all other users\n");113 continue;114 }115 } else {116 if (llvm::any_of(allocOp->getUsers(), [&](Operation *otherUser) {117 return !isa<memref::DeallocOp>(otherUser) &&118 !parentLoop->isProperAncestor(otherUser);119 })) {120 LLVM_DEBUG(121 DBGS()122 << "--Skip user: not all other users are in the parent loop\n");123 continue;124 }125 }126 candidateLoop = parentLoop;127 break;128 }129 130 if (!candidateLoop) {131 LLVM_DEBUG(DBGS() << "Skip alloc: no candidate loop\n");132 return failure();133 }134 135 std::optional<Value> inductionVar = candidateLoop.getSingleInductionVar();136 std::optional<OpFoldResult> lowerBound = candidateLoop.getSingleLowerBound();137 std::optional<OpFoldResult> singleStep = candidateLoop.getSingleStep();138 if (!inductionVar || !lowerBound || !singleStep ||139 !llvm::hasSingleElement(candidateLoop.getLoopRegions())) {140 LLVM_DEBUG(DBGS() << "Skip alloc: no single iv, lb, step or region\n");141 return failure();142 }143 144 if (!dom.dominates(allocOp.getOperation(), candidateLoop)) {145 LLVM_DEBUG(DBGS() << "Skip alloc: does not dominate candidate loop\n");146 return failure();147 }148 149 LLVM_DEBUG(DBGS() << "Start multibuffering loop: " << candidateLoop << "\n");150 151 // 1. Construct the multi-buffered memref type.152 ArrayRef<int64_t> originalShape = allocOp.getType().getShape();153 SmallVector<int64_t, 4> multiBufferedShape{multiBufferingFactor};154 llvm::append_range(multiBufferedShape, originalShape);155 LLVM_DEBUG(DBGS() << "--original type: " << allocOp.getType() << "\n");156 MemRefType mbMemRefType = MemRefType::Builder(allocOp.getType())157 .setShape(multiBufferedShape)158 .setLayout(MemRefLayoutAttrInterface());159 LLVM_DEBUG(DBGS() << "--multi-buffered type: " << mbMemRefType << "\n");160 161 // 2. Create the multi-buffered alloc.162 Location loc = allocOp->getLoc();163 OpBuilder::InsertionGuard g(rewriter);164 rewriter.setInsertionPoint(allocOp);165 auto mbAlloc = memref::AllocOp::create(rewriter, loc, mbMemRefType,166 ValueRange{}, allocOp->getAttrs());167 LLVM_DEBUG(DBGS() << "--multi-buffered alloc: " << mbAlloc << "\n");168 169 // 3. Within the loop, build the modular leading index (i.e. each loop170 // iteration %iv accesses slice ((%iv - %lb) / %step) % %mb_factor).171 rewriter.setInsertionPointToStart(172 &candidateLoop.getLoopRegions().front()->front());173 Value ivVal = *inductionVar;174 Value lbVal = getValueOrCreateConstantIndexOp(rewriter, loc, *lowerBound);175 Value stepVal = getValueOrCreateConstantIndexOp(rewriter, loc, *singleStep);176 AffineExpr iv, lb, step;177 bindDims(rewriter.getContext(), iv, lb, step);178 Value bufferIndex = affine::makeComposedAffineApply(179 rewriter, loc, ((iv - lb).floorDiv(step)) % multiBufferingFactor,180 {ivVal, lbVal, stepVal});181 LLVM_DEBUG(DBGS() << "--multi-buffered indexing: " << bufferIndex << "\n");182 183 // 4. Build the subview accessing the particular slice, taking modular184 // rotation into account.185 int64_t mbMemRefTypeRank = mbMemRefType.getRank();186 IntegerAttr zero = rewriter.getIndexAttr(0);187 IntegerAttr one = rewriter.getIndexAttr(1);188 SmallVector<OpFoldResult> offsets(mbMemRefTypeRank, zero);189 SmallVector<OpFoldResult> sizes(mbMemRefTypeRank, one);190 SmallVector<OpFoldResult> strides(mbMemRefTypeRank, one);191 // Offset is [bufferIndex, 0 ... 0 ].192 offsets.front() = bufferIndex;193 // Sizes is [1, original_size_0 ... original_size_n ].194 for (int64_t i = 0, e = originalShape.size(); i != e; ++i)195 sizes[1 + i] = rewriter.getIndexAttr(originalShape[i]);196 // Strides is [1, 1 ... 1 ].197 MemRefType dstMemref = memref::SubViewOp::inferRankReducedResultType(198 originalShape, mbMemRefType, offsets, sizes, strides);199 Value subview = memref::SubViewOp::create(rewriter, loc, dstMemref, mbAlloc,200 offsets, sizes, strides);201 LLVM_DEBUG(DBGS() << "--multi-buffered slice: " << subview << "\n");202 203 // 5. Due to the recursive nature of replaceUsesAndPropagateType , we need204 // to handle dealloc uses separately..205 for (OpOperand &use : llvm::make_early_inc_range(allocOp->getUses())) {206 auto deallocOp = dyn_cast<memref::DeallocOp>(use.getOwner());207 if (!deallocOp)208 continue;209 OpBuilder::InsertionGuard g(rewriter);210 rewriter.setInsertionPoint(deallocOp);211 auto newDeallocOp =212 memref::DeallocOp::create(rewriter, deallocOp->getLoc(), mbAlloc);213 (void)newDeallocOp;214 LLVM_DEBUG(DBGS() << "----Created dealloc: " << newDeallocOp << "\n");215 rewriter.eraseOp(deallocOp);216 }217 218 // 6. RAUW with the particular slice, taking modular rotation into account.219 replaceUsesAndPropagateType(rewriter, allocOp, subview);220 221 // 7. Finally, erase the old allocOp.222 rewriter.eraseOp(allocOp);223 224 return mbAlloc;225}226 227FailureOr<memref::AllocOp>228mlir::memref::multiBuffer(memref::AllocOp allocOp,229 unsigned multiBufferingFactor,230 bool skipOverrideAnalysis) {231 IRRewriter rewriter(allocOp->getContext());232 return multiBuffer(rewriter, allocOp, multiBufferingFactor,233 skipOverrideAnalysis);234}235