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