1 //===- Promotion.cpp - Implementation of linalg Promotion -----------------===// 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 the linalg dialect Promotion pass. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #include "PassDetail.h" 14 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" 15 #include "mlir/Dialect/Complex/IR/Complex.h" 16 #include "mlir/Dialect/Linalg/IR/Linalg.h" 17 #include "mlir/Dialect/Linalg/Passes.h" 18 #include "mlir/Dialect/Linalg/Transforms/Transforms.h" 19 #include "mlir/Dialect/Linalg/Utils/Utils.h" 20 #include "mlir/Dialect/SCF/SCF.h" 21 #include "mlir/IR/AffineExpr.h" 22 #include "mlir/IR/AffineExprVisitor.h" 23 #include "mlir/IR/AffineMap.h" 24 #include "mlir/IR/ImplicitLocOpBuilder.h" 25 #include "mlir/Support/LLVM.h" 26 #include "mlir/Transforms/FoldUtils.h" 27 #include "llvm/ADT/MapVector.h" 28 #include "llvm/ADT/SmallBitVector.h" 29 #include "llvm/ADT/TypeSwitch.h" 30 #include "llvm/Support/CommandLine.h" 31 #include "llvm/Support/Debug.h" 32 33 using namespace mlir; 34 using namespace mlir::linalg; 35 using namespace mlir::scf; 36 37 using llvm::MapVector; 38 39 #define DEBUG_TYPE "linalg-promotion" 40 41 /// Alloc a new buffer of `size` * `width` i8; where `width` is given by the 42 /// data `layout` for `elementType`. 43 /// Use AllocOp or AllocaOp depending on `options`. 44 /// Take an optional alignment. 45 static Value allocBuffer(ImplicitLocOpBuilder &b, 46 const LinalgPromotionOptions &options, 47 Type elementType, Value allocSize, DataLayout &layout, 48 Optional<unsigned> alignment = None) { 49 auto width = layout.getTypeSize(elementType); 50 51 IntegerAttr alignmentAttr; 52 if (alignment.hasValue()) 53 alignmentAttr = b.getI64IntegerAttr(alignment.getValue()); 54 55 // Static buffer. 56 if (auto cst = allocSize.getDefiningOp<arith::ConstantIndexOp>()) { 57 auto staticBufferType = 58 MemRefType::get(width * cst.value(), b.getIntegerType(8)); 59 if (options.useAlloca) { 60 return b.createOrFold<memref::AllocaOp>(staticBufferType, ValueRange{}, 61 alignmentAttr); 62 } 63 return b.createOrFold<memref::AllocOp>(staticBufferType, ValueRange{}, 64 alignmentAttr); 65 } 66 67 // Fallback dynamic buffer. 68 auto dynamicBufferType = MemRefType::get(-1, b.getIntegerType(8)); 69 Value mul = b.createOrFold<arith::MulIOp>( 70 b.create<arith::ConstantIndexOp>(width), allocSize); 71 if (options.useAlloca) 72 return b.create<memref::AllocaOp>(dynamicBufferType, mul, alignmentAttr); 73 return b.create<memref::AllocOp>(dynamicBufferType, mul, alignmentAttr); 74 } 75 76 /// Default allocation callback function. This allocates a promoted buffer when 77 /// no call back to do so is provided. The default is to allocate a 78 /// memref<..xi8> and return a view to get a memref type of shape 79 /// boundingSubViewSize. 80 static Optional<Value> 81 defaultAllocBufferCallBack(const LinalgPromotionOptions &options, 82 OpBuilder &builder, memref::SubViewOp subView, 83 ArrayRef<Value> boundingSubViewSize, 84 Optional<unsigned> alignment, DataLayout &layout) { 85 ShapedType viewType = subView.getType(); 86 ImplicitLocOpBuilder b(subView.getLoc(), builder); 87 auto zero = b.createOrFold<arith::ConstantIndexOp>(0); 88 auto one = b.createOrFold<arith::ConstantIndexOp>(1); 89 90 Value allocSize = one; 91 for (const auto &size : llvm::enumerate(boundingSubViewSize)) 92 allocSize = b.createOrFold<arith::MulIOp>(allocSize, size.value()); 93 Value buffer = allocBuffer(b, options, viewType.getElementType(), allocSize, 94 layout, alignment); 95 SmallVector<int64_t, 4> dynSizes(boundingSubViewSize.size(), 96 ShapedType::kDynamicSize); 97 Value view = b.createOrFold<memref::ViewOp>( 98 MemRefType::get(dynSizes, viewType.getElementType()), buffer, zero, 99 boundingSubViewSize); 100 return view; 101 } 102 103 /// Default implementation of deallocation of the buffer use for promotion. It 104 /// expects to get the same value that the default allocation method returned, 105 /// i.e. result of a ViewOp. 106 static LogicalResult 107 defaultDeallocBufferCallBack(const LinalgPromotionOptions &options, 108 OpBuilder &b, Value fullLocalView) { 109 if (!options.useAlloca) { 110 auto viewOp = cast<memref::ViewOp>(fullLocalView.getDefiningOp()); 111 b.create<memref::DeallocOp>(viewOp.source().getLoc(), viewOp.source()); 112 } 113 return success(); 114 } 115 116 namespace { 117 118 /// Helper struct that captures the information required to apply the 119 /// transformation on each op. This bridges the abstraction gap with the 120 /// user-facing API which exposes positional arguments to control which operands 121 /// are promoted. 122 struct LinalgOpInstancePromotionOptions { 123 LinalgOpInstancePromotionOptions(LinalgOp op, 124 const LinalgPromotionOptions &options); 125 /// SubViews to promote. 126 MapVector<int64_t, Value> subViews; 127 /// True if the full view should be used for the promoted buffer. 128 DenseMap<Value, bool> useFullTileBuffers; 129 130 /// Callback functions for allocation and deallocation of promoted buffers, as 131 /// well as to copy the data into and out of these buffers. 132 AllocBufferCallbackFn allocationFn; 133 DeallocBufferCallbackFn deallocationFn; 134 CopyCallbackFn copyInFn; 135 CopyCallbackFn copyOutFn; 136 137 /// Allow the use of dynamically-sized buffers. 138 bool dynamicBuffers; 139 140 /// Alignment of promoted buffer. 141 Optional<unsigned> alignment; 142 }; 143 } // namespace 144 145 LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions( 146 LinalgOp linalgOp, const LinalgPromotionOptions &options) 147 : subViews(), dynamicBuffers(options.dynamicBuffers), 148 alignment(options.alignment) { 149 assert(linalgOp.hasBufferSemantics() && "revisit usage of shaped operand"); 150 auto vUseFullTileBuffers = 151 options.useFullTileBuffers.getValueOr(llvm::SmallBitVector()); 152 vUseFullTileBuffers.resize(linalgOp.getNumInputsAndOutputs(), 153 options.useFullTileBuffersDefault); 154 155 for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) { 156 int64_t operandNumber = opOperand->getOperandNumber(); 157 if (options.operandsToPromote && 158 !options.operandsToPromote->count(operandNumber)) 159 continue; 160 Operation *op = opOperand->get().getDefiningOp(); 161 if (auto sv = dyn_cast_or_null<memref::SubViewOp>(op)) { 162 subViews[operandNumber] = sv; 163 useFullTileBuffers[sv] = vUseFullTileBuffers[operandNumber]; 164 } 165 } 166 167 if (options.allocationFn) { 168 allocationFn = *options.allocationFn; 169 } else { 170 allocationFn = [&](OpBuilder &b, memref::SubViewOp subViewOp, 171 ArrayRef<Value> boundingSubViewSize, 172 DataLayout &layout) -> Optional<Value> { 173 return defaultAllocBufferCallBack(options, b, subViewOp, 174 boundingSubViewSize, alignment, layout); 175 }; 176 } 177 178 if (options.deallocationFn) { 179 deallocationFn = *options.deallocationFn; 180 } else { 181 deallocationFn = [&](OpBuilder &b, Value buffer) { 182 return defaultDeallocBufferCallBack(options, b, buffer); 183 }; 184 } 185 186 // Save the loc because `linalgOp` goes out of scope. 187 Location loc = linalgOp.getLoc(); 188 auto defaultCopyCallBack = [loc](OpBuilder &b, Value src, 189 Value dst) -> LogicalResult { 190 b.create<memref::CopyOp>(loc, src, dst); 191 return success(); 192 }; 193 copyInFn = (options.copyInFn ? *(options.copyInFn) : defaultCopyCallBack); 194 copyOutFn = (options.copyOutFn ? *(options.copyOutFn) : defaultCopyCallBack); 195 } 196 197 // Performs promotion of a `subView` into a local buffer of the size of the 198 // *ranges* of the `subView`. This produces a buffer whose size may be bigger 199 // than the actual size of the `subView` at the boundaries. 200 // This is related to the full/partial tile problem. 201 // Returns a PromotionInfo containing a `buffer`, `fullLocalView` and 202 // `partialLocalView` such that: 203 // * `buffer` is always the size of the full tile. 204 // * `fullLocalView` is a dense contiguous view into that buffer. 205 // * `partialLocalView` is a dense non-contiguous slice of `fullLocalView` 206 // that corresponds to the size of `subView` and accounting for boundary 207 // effects. 208 // The point of the full tile buffer is that constant static tile sizes are 209 // folded and result in a buffer type with statically known size and alignment 210 // properties. 211 // To account for general boundary effects, padding must be performed on the 212 // boundary tiles. For now this is done with an unconditional `fill` op followed 213 // by a partial `copy` op. 214 FailureOr<PromotionInfo> mlir::linalg::promoteSubviewAsNewBuffer( 215 OpBuilder &b, Location loc, memref::SubViewOp subView, 216 const AllocBufferCallbackFn &allocationFn, DataLayout &layout) { 217 auto viewType = subView.getType(); 218 auto rank = viewType.getRank(); 219 SmallVector<Value, 4> fullSizes; 220 SmallVector<OpFoldResult> partialSizes; 221 fullSizes.reserve(rank); 222 partialSizes.reserve(rank); 223 llvm::SmallBitVector droppedDims = subView.getDroppedDims(); 224 int64_t resultDimIdx = 0; 225 for (const auto &en : llvm::enumerate(subView.getOrCreateRanges(b, loc))) { 226 if (droppedDims[en.index()]) 227 continue; 228 auto rangeValue = en.value(); 229 // Try to extract a tight constant. 230 LLVM_DEBUG(llvm::dbgs() << "Extract tightest: " << rangeValue.size << "\n"); 231 FailureOr<int64_t> upperBound = 232 getConstantUpperBoundForIndex(rangeValue.size); 233 Value size = 234 failed(upperBound) 235 ? rangeValue.size 236 : b.create<arith::ConstantIndexOp>(loc, upperBound.getValue()); 237 LLVM_DEBUG(llvm::dbgs() << "Extracted tightest: " << size << "\n"); 238 fullSizes.push_back(size); 239 partialSizes.push_back( 240 b.createOrFold<memref::DimOp>(loc, subView, resultDimIdx++)); 241 } 242 SmallVector<int64_t, 4> dynSizes(fullSizes.size(), -1); 243 // If a callback is not specified, then use the default implementation for 244 // allocating the promoted buffer. 245 Optional<Value> fullLocalView = allocationFn(b, subView, fullSizes, layout); 246 if (!fullLocalView) 247 return failure(); 248 SmallVector<OpFoldResult, 4> zeros(fullSizes.size(), b.getIndexAttr(0)); 249 SmallVector<OpFoldResult, 4> ones(fullSizes.size(), b.getIndexAttr(1)); 250 auto partialLocalView = b.createOrFold<memref::SubViewOp>( 251 loc, *fullLocalView, zeros, partialSizes, ones); 252 return PromotionInfo{*fullLocalView, partialLocalView}; 253 } 254 255 static FailureOr<MapVector<int64_t, PromotionInfo>> 256 promoteSubViews(ImplicitLocOpBuilder &b, 257 LinalgOpInstancePromotionOptions options, DataLayout &layout) { 258 if (options.subViews.empty()) 259 return failure(); 260 261 MapVector<int64_t, PromotionInfo> promotionInfoMap; 262 263 for (auto v : options.subViews) { 264 memref::SubViewOp subView = 265 cast<memref::SubViewOp>(v.second.getDefiningOp()); 266 auto promotionInfo = promoteSubviewAsNewBuffer( 267 b, b.getLoc(), subView, options.allocationFn, layout); 268 if (failed(promotionInfo)) 269 return failure(); 270 promotionInfoMap[v.first] = *promotionInfo; 271 272 // Only fill the buffer if the full local view is used 273 if (!options.useFullTileBuffers[v.second]) 274 continue; 275 Type subviewEltType = subView.getType().getElementType(); 276 Value fillVal = 277 llvm::TypeSwitch<Type, Value>(subviewEltType) 278 .Case([&](FloatType t) { 279 return b.create<arith::ConstantOp>(FloatAttr::get(t, 0.0)); 280 }) 281 .Case([&](IntegerType t) { 282 return b.create<arith::ConstantOp>(IntegerAttr::get(t, 0)); 283 }) 284 .Case([&](ComplexType t) { 285 Value tmp; 286 if (auto et = t.getElementType().dyn_cast<FloatType>()) 287 tmp = b.create<arith::ConstantOp>(FloatAttr::get(et, 0.0)); 288 else if (auto et = t.getElementType().cast<IntegerType>()) 289 tmp = b.create<arith::ConstantOp>(IntegerAttr::get(et, 0)); 290 return b.create<complex::CreateOp>(t, tmp, tmp); 291 }) 292 .Default([](auto) { return Value(); }); 293 if (!fillVal) 294 return failure(); 295 b.create<linalg::FillOp>(fillVal, promotionInfo->fullLocalView); 296 } 297 298 // Copy data into the promoted buffers. Use callback if provided. 299 for (auto v : options.subViews) { 300 auto info = promotionInfoMap.find(v.first); 301 if (info == promotionInfoMap.end()) 302 continue; 303 if (failed(options.copyInFn( 304 b, cast<memref::SubViewOp>(v.second.getDefiningOp()), 305 info->second.partialLocalView))) 306 return failure(); 307 } 308 return promotionInfoMap; 309 } 310 311 static FailureOr<LinalgOp> 312 promoteSubViews(ImplicitLocOpBuilder &b, LinalgOp op, 313 LinalgOpInstancePromotionOptions options, DataLayout &layout) { 314 assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics"); 315 316 // 1. Promote the specified views and use them in the new op. 317 auto promotedBuffersAndViews = promoteSubViews(b, options, layout); 318 if (failed(promotedBuffersAndViews) || 319 promotedBuffersAndViews->size() != options.subViews.size()) 320 return failure(); 321 322 // 2. Append all other operands as they appear, this enforces that such 323 // operands are not views. This is to support cases such as FillOp taking 324 // extra scalars etc. Keep a reference to output buffers; 325 SmallVector<Value, 8> opViews; 326 opViews.reserve(op.getNumInputsAndOutputs()); 327 SmallVector<std::pair<Value, Value>, 8> writebackViews; 328 writebackViews.reserve(promotedBuffersAndViews->size()); 329 for (OpOperand *opOperand : op.getInputAndOutputOperands()) { 330 int64_t operandNumber = opOperand->getOperandNumber(); 331 if (options.subViews.count(operandNumber) != 0) { 332 if (options.useFullTileBuffers[opOperand->get()]) 333 opViews.push_back( 334 (*promotedBuffersAndViews)[operandNumber].fullLocalView); 335 else 336 opViews.push_back( 337 (*promotedBuffersAndViews)[operandNumber].partialLocalView); 338 if (operandNumber >= op.getNumInputs()) 339 writebackViews.emplace_back(std::make_pair( 340 opOperand->get(), 341 (*promotedBuffersAndViews)[operandNumber].partialLocalView)); 342 } else { 343 opViews.push_back(opOperand->get()); 344 } 345 } 346 op->setOperands(0, opViews.size(), opViews); 347 348 OpBuilder::InsertionGuard guard(b); 349 b.setInsertionPointAfter(op); 350 // 3. Emit write-back for the promoted output views: copy the partial view. 351 for (auto viewAndPartialLocalView : writebackViews) { 352 if (failed(options.copyOutFn(b, viewAndPartialLocalView.second, 353 viewAndPartialLocalView.first))) 354 return failure(); 355 } 356 357 // 4. Dealloc all local buffers. 358 for (const auto &pi : *promotedBuffersAndViews) 359 (void)options.deallocationFn(b, pi.second.fullLocalView); 360 return op; 361 } 362 363 LogicalResult 364 mlir::linalg::promoteSubviewsPrecondition(Operation *op, 365 LinalgPromotionOptions options) { 366 LinalgOp linalgOp = dyn_cast<LinalgOp>(op); 367 // Transformation applies to buffers only. 368 if (!linalgOp || !linalgOp.hasBufferSemantics()) 369 return failure(); 370 // Check that at least one of the requested operands is indeed a subview. 371 for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) { 372 auto sv = 373 isa_and_nonnull<memref::SubViewOp>(opOperand->get().getDefiningOp()); 374 if (sv) { 375 if (!options.operandsToPromote.hasValue() || 376 options.operandsToPromote->count(opOperand->getOperandNumber())) 377 return success(); 378 } 379 } 380 // TODO: Check all subviews requested are bound by a static constant. 381 // TODO: Check that the total footprint fits within a given size. 382 return failure(); 383 } 384 385 FailureOr<LinalgOp> 386 mlir::linalg::promoteSubViews(OpBuilder &builder, LinalgOp linalgOp, 387 const LinalgPromotionOptions &options) { 388 LinalgOpInstancePromotionOptions linalgOptions(linalgOp, options); 389 auto layout = DataLayout::closest(linalgOp); 390 ImplicitLocOpBuilder b(linalgOp.getLoc(), builder); 391 auto res = ::promoteSubViews(b, linalgOp, linalgOptions, layout); 392 if (failed(res)) 393 return failure(); 394 return res; 395 } 396 397 namespace { 398 struct LinalgPromotionPass : public LinalgPromotionBase<LinalgPromotionPass> { 399 LinalgPromotionPass() = default; 400 LinalgPromotionPass(bool dynamicBuffers, bool useAlloca) { 401 this->dynamicBuffers = dynamicBuffers; 402 this->useAlloca = useAlloca; 403 } 404 405 void runOnOperation() override { 406 getOperation().walk([&](LinalgOp op) { 407 auto options = LinalgPromotionOptions() 408 .setDynamicBuffers(dynamicBuffers) 409 .setUseAlloca(useAlloca); 410 if (failed(promoteSubviewsPrecondition(op, options))) 411 return; 412 LLVM_DEBUG(llvm::dbgs() << "Promote: " << *(op.getOperation()) << "\n"); 413 ImplicitLocOpBuilder b(op.getLoc(), op); 414 // TODO: signalPassFailure() ? 415 (void)promoteSubViews(b, op, options); 416 }); 417 } 418 }; 419 } // namespace 420 421 // TODO: support more transformation options in the pass. 422 std::unique_ptr<OperationPass<func::FuncOp>> 423 mlir::createLinalgPromotionPass(bool dynamicBuffers, bool useAlloca) { 424 return std::make_unique<LinalgPromotionPass>(dynamicBuffers, useAlloca); 425 } 426 std::unique_ptr<OperationPass<func::FuncOp>> mlir::createLinalgPromotionPass() { 427 return std::make_unique<LinalgPromotionPass>(); 428 } 429