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