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/Func/IR/FuncOps.h" 17 #include "mlir/Dialect/GPU/IR/GPUDialect.h" 18 #include "mlir/Dialect/Linalg/IR/Linalg.h" 19 #include "mlir/Dialect/Linalg/Passes.h" 20 #include "mlir/Dialect/Linalg/Transforms/Transforms.h" 21 #include "mlir/Dialect/SCF/IR/SCF.h" 22 #include "mlir/IR/AffineExpr.h" 23 #include "mlir/IR/AffineExprVisitor.h" 24 #include "mlir/IR/AffineMap.h" 25 #include "mlir/IR/ImplicitLocOpBuilder.h" 26 #include "mlir/Interfaces/ValueBoundsOpInterface.h" 27 #include "mlir/Support/LLVM.h" 28 #include "mlir/Transforms/FoldUtils.h" 29 #include "llvm/ADT/MapVector.h" 30 #include "llvm/ADT/SmallBitVector.h" 31 #include "llvm/ADT/SmallSet.h" 32 #include "llvm/ADT/TypeSwitch.h" 33 #include "llvm/Support/CommandLine.h" 34 #include "llvm/Support/Debug.h" 35 36 using namespace mlir; 37 using namespace mlir::linalg; 38 using namespace mlir::scf; 39 40 using llvm::MapVector; 41 42 #define DEBUG_TYPE "linalg-promotion" 43 44 /// Alloc a new buffer of `size` * `width` i8; where `width` is given by the 45 /// data `layout` for `elementType`. 46 /// Use AllocOp or AllocaOp depending on `options`. 47 /// Take an optional alignment. 48 static Value allocBuffer(ImplicitLocOpBuilder &b, 49 const LinalgPromotionOptions &options, 50 Type elementType, Value allocSize, DataLayout &layout, 51 std::optional<unsigned> alignment = std::nullopt) { 52 auto width = layout.getTypeSize(elementType); 53 54 IntegerAttr alignmentAttr; 55 if (alignment.has_value()) 56 alignmentAttr = b.getI64IntegerAttr(alignment.value()); 57 58 Attribute memorySpaceAttr; 59 if (options.memorySpace.has_value()) 60 memorySpaceAttr = *options.memorySpace; 61 62 // Static buffer. 63 if (std::optional<int64_t> cst = getConstantIntValue(allocSize)) { 64 auto staticBufferType = 65 MemRefType::get(width * cst.value(), b.getIntegerType(8)); 66 staticBufferType = 67 MemRefType::Builder(staticBufferType).setMemorySpace(memorySpaceAttr); 68 if (options.useAlloca) { 69 return b.create<memref::AllocaOp>(staticBufferType, ValueRange{}, 70 alignmentAttr); 71 } 72 return b.create<memref::AllocOp>(staticBufferType, ValueRange{}, 73 alignmentAttr); 74 } 75 76 // Fallback dynamic buffer. 77 auto dynamicBufferType = 78 MemRefType::get(ShapedType::kDynamic, b.getIntegerType(8)); 79 dynamicBufferType = 80 MemRefType::Builder(dynamicBufferType).setMemorySpace(memorySpaceAttr); 81 Value mul = b.createOrFold<arith::MulIOp>( 82 b.create<arith::ConstantIndexOp>(width), allocSize); 83 if (options.useAlloca) 84 return b.create<memref::AllocaOp>(dynamicBufferType, mul, alignmentAttr); 85 return b.create<memref::AllocOp>(dynamicBufferType, mul, alignmentAttr); 86 } 87 88 /// Default allocation callback function. This allocates a promoted buffer when 89 /// no call back to do so is provided. The default is to allocate a 90 /// memref<..xi8> and return a view to get a memref type of shape 91 /// boundingSubViewSize. 92 static std::optional<Value> defaultAllocBufferCallBack( 93 const LinalgPromotionOptions &options, OpBuilder &builder, 94 memref::SubViewOp subView, ArrayRef<Value> boundingSubViewSize, 95 std::optional<unsigned> alignment, DataLayout &layout) { 96 ShapedType viewType = subView.getType(); 97 ImplicitLocOpBuilder b(subView.getLoc(), builder); 98 auto zero = b.create<arith::ConstantIndexOp>(0); 99 auto one = b.create<arith::ConstantIndexOp>(1); 100 101 Attribute memorySpaceAttr; 102 if (options.memorySpace.has_value()) 103 memorySpaceAttr = *options.memorySpace; 104 105 Value allocSize = one; 106 for (const auto &size : llvm::enumerate(boundingSubViewSize)) 107 allocSize = b.createOrFold<arith::MulIOp>(allocSize, size.value()); 108 Value buffer = allocBuffer(b, options, viewType.getElementType(), allocSize, 109 layout, alignment); 110 SmallVector<int64_t, 4> dynSizes(boundingSubViewSize.size(), 111 ShapedType::kDynamic); 112 113 auto viewMemRefType = MemRefType::get(dynSizes, viewType.getElementType()); 114 viewMemRefType = 115 MemRefType::Builder(viewMemRefType).setMemorySpace(memorySpaceAttr); 116 Value view = b.createOrFold<memref::ViewOp>(viewMemRefType, buffer, zero, 117 boundingSubViewSize); 118 return view; 119 } 120 121 /// Default implementation of deallocation of the buffer use for promotion. It 122 /// expects to get the same value that the default allocation method returned, 123 /// i.e. result of a ViewOp. 124 static LogicalResult 125 defaultDeallocBufferCallBack(const LinalgPromotionOptions &options, 126 OpBuilder &b, Value fullLocalView) { 127 if (!options.useAlloca) { 128 auto viewOp = cast<memref::ViewOp>(fullLocalView.getDefiningOp()); 129 b.create<memref::DeallocOp>(viewOp.getSource().getLoc(), 130 viewOp.getSource()); 131 } 132 return success(); 133 } 134 135 namespace { 136 137 /// Helper struct that captures the information required to apply the 138 /// transformation on each op. This bridges the abstraction gap with the 139 /// user-facing API which exposes positional arguments to control which operands 140 /// are promoted. 141 struct LinalgOpInstancePromotionOptions { 142 LinalgOpInstancePromotionOptions(LinalgOp op, 143 const LinalgPromotionOptions &options); 144 /// SubViews to promote. 145 MapVector<int64_t, Value> subViews; 146 /// Subviews operand numbers to copy in using copyInFn. 147 llvm::SmallSet<int64_t, 4> operandsNumbersToCopyIn; 148 /// True if the full view should be used for the promoted buffer. 149 DenseMap<Value, bool> useFullTileBuffers; 150 151 /// Callback functions for allocation and deallocation of promoted buffers, as 152 /// well as to copy the data into and out of these buffers. 153 AllocBufferCallbackFn allocationFn; 154 DeallocBufferCallbackFn deallocationFn; 155 CopyCallbackFn copyInFn; 156 CopyCallbackFn copyOutFn; 157 158 /// Alignment of promoted buffer. 159 std::optional<unsigned> alignment; 160 }; 161 } // namespace 162 163 LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions( 164 LinalgOp linalgOp, const LinalgPromotionOptions &options) 165 : subViews(), alignment(options.alignment) { 166 assert(linalgOp.hasBufferSemantics() && "revisit usage of shaped operand"); 167 auto vUseFullTileBuffers = 168 options.useFullTileBuffers.value_or(llvm::SmallBitVector()); 169 vUseFullTileBuffers.resize(linalgOp->getNumOperands(), 170 options.useFullTileBuffersDefault); 171 172 for (OpOperand &opOperand : linalgOp->getOpOperands()) { 173 int64_t operandNumber = opOperand.getOperandNumber(); 174 if (options.operandsToPromote && 175 !options.operandsToPromote->count(operandNumber)) 176 continue; 177 Operation *op = opOperand.get().getDefiningOp(); 178 if (auto sv = dyn_cast_or_null<memref::SubViewOp>(op)) { 179 subViews[operandNumber] = sv; 180 // In case of linalg generic, copy in only if subview is used in linalg 181 // payload. 182 if (!isa<linalg::GenericOp>(linalgOp) || 183 linalgOp.payloadUsesValueFromOperand(&opOperand)) 184 operandsNumbersToCopyIn.insert(operandNumber); 185 useFullTileBuffers[sv] = vUseFullTileBuffers[operandNumber]; 186 } 187 } 188 189 if (options.allocationFn) { 190 allocationFn = *options.allocationFn; 191 } else { 192 allocationFn = [&](OpBuilder &b, memref::SubViewOp subViewOp, 193 ArrayRef<Value> boundingSubViewSize, 194 DataLayout &layout) -> std::optional<Value> { 195 return defaultAllocBufferCallBack(options, b, subViewOp, 196 boundingSubViewSize, alignment, layout); 197 }; 198 } 199 200 if (options.deallocationFn) { 201 deallocationFn = *options.deallocationFn; 202 } else { 203 deallocationFn = [&](OpBuilder &b, Value buffer) { 204 return defaultDeallocBufferCallBack(options, b, buffer); 205 }; 206 } 207 208 // Save the loc because `linalgOp` goes out of scope. 209 Location loc = linalgOp.getLoc(); 210 auto defaultCopyCallBack = [loc](OpBuilder &b, Value src, 211 Value dst) -> LogicalResult { 212 b.create<linalg::CopyOp>(loc, src, dst); 213 return success(); 214 }; 215 copyInFn = (options.copyInFn ? *(options.copyInFn) : defaultCopyCallBack); 216 copyOutFn = (options.copyOutFn ? *(options.copyOutFn) : defaultCopyCallBack); 217 } 218 219 // Performs promotion of a `subView` into a local buffer of the size of the 220 // *ranges* of the `subView`. This produces a buffer whose size may be bigger 221 // than the actual size of the `subView` at the boundaries. 222 // This is related to the full/partial tile problem. 223 // Returns a PromotionInfo containing a `buffer`, `fullLocalView` and 224 // `partialLocalView` such that: 225 // * `buffer` is always the size of the full tile. 226 // * `fullLocalView` is a dense contiguous view into that buffer. 227 // * `partialLocalView` is a dense non-contiguous slice of `fullLocalView` 228 // that corresponds to the size of `subView` and accounting for boundary 229 // effects. 230 // The point of the full tile buffer is that constant static tile sizes are 231 // folded and result in a buffer type with statically known size and alignment 232 // properties. 233 // To account for general boundary effects, padding must be performed on the 234 // boundary tiles. For now this is done with an unconditional `fill` op followed 235 // by a partial `copy` op. 236 FailureOr<PromotionInfo> mlir::linalg::promoteSubviewAsNewBuffer( 237 OpBuilder &b, Location loc, memref::SubViewOp subView, 238 const AllocBufferCallbackFn &allocationFn, DataLayout &layout) { 239 auto viewType = subView.getType(); 240 auto rank = viewType.getRank(); 241 SmallVector<Value, 4> fullSizes; 242 SmallVector<OpFoldResult> partialSizes; 243 fullSizes.reserve(rank); 244 partialSizes.reserve(rank); 245 llvm::SmallBitVector droppedDims = subView.getDroppedDims(); 246 int64_t resultDimIdx = 0; 247 for (const auto &en : llvm::enumerate(subView.getOrCreateRanges(b, loc))) { 248 if (droppedDims[en.index()]) 249 continue; 250 auto rangeValue = en.value(); 251 // Try to extract a tight constant. If the size is known statically, no need 252 // to look for the bound. 253 LLVM_DEBUG(llvm::dbgs() << "Extract tightest: " << rangeValue.size << "\n"); 254 Value size; 255 if (auto attr = llvm::dyn_cast_if_present<Attribute>(rangeValue.size)) { 256 size = getValueOrCreateConstantIndexOp(b, loc, rangeValue.size); 257 } else { 258 Value materializedSize = 259 getValueOrCreateConstantIndexOp(b, loc, rangeValue.size); 260 FailureOr<int64_t> upperBound = 261 ValueBoundsConstraintSet::computeConstantBound( 262 presburger::BoundType::UB, materializedSize, /*dim=*/std::nullopt, 263 /*stopCondition=*/nullptr, /*closedUB=*/true); 264 size = failed(upperBound) 265 ? materializedSize 266 : b.create<arith::ConstantIndexOp>(loc, *upperBound); 267 } 268 LLVM_DEBUG(llvm::dbgs() << "Extracted tightest: " << size << "\n"); 269 fullSizes.push_back(size); 270 partialSizes.push_back( 271 b.createOrFold<memref::DimOp>(loc, subView, resultDimIdx++)); 272 } 273 SmallVector<int64_t, 4> dynSizes(fullSizes.size(), ShapedType::kDynamic); 274 // If a callback is not specified, then use the default implementation for 275 // allocating the promoted buffer. 276 std::optional<Value> fullLocalView = 277 allocationFn(b, subView, fullSizes, layout); 278 if (!fullLocalView) 279 return failure(); 280 SmallVector<OpFoldResult, 4> zeros(fullSizes.size(), b.getIndexAttr(0)); 281 SmallVector<OpFoldResult, 4> ones(fullSizes.size(), b.getIndexAttr(1)); 282 auto partialLocalView = b.createOrFold<memref::SubViewOp>( 283 loc, *fullLocalView, zeros, partialSizes, ones); 284 return PromotionInfo{*fullLocalView, partialLocalView}; 285 } 286 287 static FailureOr<MapVector<int64_t, PromotionInfo>> 288 promoteSubViews(ImplicitLocOpBuilder &b, 289 LinalgOpInstancePromotionOptions options, DataLayout &layout) { 290 if (options.subViews.empty()) 291 return failure(); 292 293 MapVector<int64_t, PromotionInfo> promotionInfoMap; 294 295 for (auto v : options.subViews) { 296 memref::SubViewOp subView = 297 cast<memref::SubViewOp>(v.second.getDefiningOp()); 298 auto promotionInfo = promoteSubviewAsNewBuffer( 299 b, b.getLoc(), subView, options.allocationFn, layout); 300 if (failed(promotionInfo)) 301 return failure(); 302 promotionInfoMap[v.first] = *promotionInfo; 303 304 // Only fill the buffer if the full local view is used 305 if (!options.useFullTileBuffers[v.second]) 306 continue; 307 Type subviewEltType = subView.getType().getElementType(); 308 Value fillVal = 309 llvm::TypeSwitch<Type, Value>(subviewEltType) 310 .Case([&](FloatType t) { 311 return b.create<arith::ConstantOp>(FloatAttr::get(t, 0.0)); 312 }) 313 .Case([&](IntegerType t) { 314 return b.create<arith::ConstantOp>(IntegerAttr::get(t, 0)); 315 }) 316 .Case([&](ComplexType t) { 317 Value tmp; 318 if (auto et = dyn_cast<FloatType>(t.getElementType())) 319 tmp = b.create<arith::ConstantOp>(FloatAttr::get(et, 0.0)); 320 else if (auto et = cast<IntegerType>(t.getElementType())) 321 tmp = b.create<arith::ConstantOp>(IntegerAttr::get(et, 0)); 322 return b.create<complex::CreateOp>(t, tmp, tmp); 323 }) 324 .Default([](auto) { return Value(); }); 325 if (!fillVal) 326 return failure(); 327 b.create<linalg::FillOp>(fillVal, promotionInfo->fullLocalView); 328 } 329 330 // Copy data into the promoted buffers. Use callback if provided. 331 for (auto v : options.subViews) { 332 auto info = promotionInfoMap.find(v.first); 333 if (info == promotionInfoMap.end()) 334 continue; 335 if (options.operandsNumbersToCopyIn.count(v.first) == 0) 336 continue; 337 if (failed(options.copyInFn( 338 b, cast<memref::SubViewOp>(v.second.getDefiningOp()), 339 info->second.partialLocalView))) 340 return failure(); 341 } 342 return promotionInfoMap; 343 } 344 345 static FailureOr<LinalgOp> 346 promoteSubViews(ImplicitLocOpBuilder &b, LinalgOp op, 347 LinalgOpInstancePromotionOptions options, DataLayout &layout) { 348 assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics"); 349 350 // 1. Promote the specified views and use them in the new op. 351 auto promotedBuffersAndViews = promoteSubViews(b, options, layout); 352 if (failed(promotedBuffersAndViews) || 353 promotedBuffersAndViews->size() != options.subViews.size()) 354 return failure(); 355 356 // 2. Append all other operands as they appear, this enforces that such 357 // operands are not views. This is to support cases such as FillOp taking 358 // extra scalars etc. Keep a reference to output buffers; 359 SmallVector<Value, 8> opViews; 360 opViews.reserve(op->getNumOperands()); 361 SmallVector<std::pair<Value, Value>, 8> writebackViews; 362 writebackViews.reserve(promotedBuffersAndViews->size()); 363 for (OpOperand &opOperand : op->getOpOperands()) { 364 int64_t operandNumber = opOperand.getOperandNumber(); 365 if (options.subViews.count(operandNumber) != 0) { 366 if (options.useFullTileBuffers[opOperand.get()]) 367 opViews.push_back( 368 (*promotedBuffersAndViews)[operandNumber].fullLocalView); 369 else 370 opViews.push_back( 371 (*promotedBuffersAndViews)[operandNumber].partialLocalView); 372 if (operandNumber >= op.getNumDpsInputs()) 373 writebackViews.emplace_back(std::make_pair( 374 opOperand.get(), 375 (*promotedBuffersAndViews)[operandNumber].partialLocalView)); 376 } else { 377 opViews.push_back(opOperand.get()); 378 } 379 } 380 op->setOperands(0, opViews.size(), opViews); 381 382 OpBuilder::InsertionGuard guard(b); 383 b.setInsertionPointAfter(op); 384 // 3. Emit write-back for the promoted output views: copy the partial view. 385 for (auto viewAndPartialLocalView : writebackViews) { 386 if (failed(options.copyOutFn(b, viewAndPartialLocalView.second, 387 viewAndPartialLocalView.first))) 388 return failure(); 389 } 390 391 // 4. Dealloc all local buffers. 392 for (const auto &pi : *promotedBuffersAndViews) 393 (void)options.deallocationFn(b, pi.second.fullLocalView); 394 return op; 395 } 396 397 LogicalResult 398 mlir::linalg::promoteSubviewsPrecondition(Operation *op, 399 LinalgPromotionOptions options) { 400 LinalgOp linalgOp = dyn_cast<LinalgOp>(op); 401 // Transformation applies to buffers only. 402 if (!linalgOp || !linalgOp.hasBufferSemantics()) 403 return failure(); 404 // Check that at least one of the requested operands is indeed a subview. 405 for (OpOperand &opOperand : linalgOp->getOpOperands()) { 406 auto sv = 407 isa_and_nonnull<memref::SubViewOp>(opOperand.get().getDefiningOp()); 408 if (sv) { 409 if (!options.operandsToPromote || 410 options.operandsToPromote->count(opOperand.getOperandNumber())) 411 return success(); 412 } 413 } 414 // TODO: Check all subviews requested are bound by a static constant. 415 // TODO: Check that the total footprint fits within a given size. 416 return failure(); 417 } 418 419 FailureOr<LinalgOp> 420 mlir::linalg::promoteSubViews(OpBuilder &builder, LinalgOp linalgOp, 421 const LinalgPromotionOptions &options) { 422 LinalgOpInstancePromotionOptions linalgOptions(linalgOp, options); 423 auto layout = DataLayout::closest(linalgOp); 424 ImplicitLocOpBuilder b(linalgOp.getLoc(), builder); 425 auto res = ::promoteSubViews(b, linalgOp, linalgOptions, layout); 426 if (failed(res)) 427 return failure(); 428 return res; 429 } 430 431 /// Allocate the given subview to a memory address space in GPU by creating a 432 /// allocation operation and setting the memref type address space to desired 433 /// address space. 434 static std::optional<Value> allocateSubviewGPUMemoryInAddressSpace( 435 OpBuilder &builder, memref::SubViewOp subview, ArrayRef<Value> sizeBounds, 436 gpu::AddressSpace addressSpace) { 437 OpBuilder::InsertionGuard guard(builder); 438 439 func::FuncOp funcOp = subview->getParentOfType<func::FuncOp>(); 440 if (!funcOp) 441 return std::nullopt; 442 443 // The subview size bounds are expected to be constant; they specify the shape 444 // of the allocation. 445 SmallVector<int64_t> shape; 446 for (Value bound : sizeBounds) { 447 APInt value; 448 if (!matchPattern(bound, m_ConstantInt(&value))) 449 return std::nullopt; 450 shape.push_back(value.getSExtValue()); 451 } 452 453 builder.setInsertionPoint(&funcOp.front(), funcOp.front().begin()); 454 auto type = MemRefType::get( 455 shape, subview.getType().getElementType(), MemRefLayoutAttrInterface{}, 456 gpu::AddressSpaceAttr::get(builder.getContext(), addressSpace)); 457 Value buffer; 458 if (addressSpace == gpu::GPUDialect::getWorkgroupAddressSpace()) { 459 buffer = builder.create<memref::AllocOp>(funcOp.getLoc(), type); 460 } else if (addressSpace == gpu::GPUDialect::getPrivateAddressSpace()) { 461 buffer = builder.create<memref::AllocaOp>(funcOp.getLoc(), type); 462 } else { 463 return std::nullopt; 464 } 465 return buffer; 466 } 467 468 /// Allocate the subview in the GPU workgroup memory. 469 std::optional<Value> mlir::linalg::allocateWorkgroupMemory( 470 OpBuilder &builder, memref::SubViewOp subview, ArrayRef<Value> sizeBounds, 471 DataLayout &) { 472 return allocateSubviewGPUMemoryInAddressSpace( 473 builder, subview, sizeBounds, 474 gpu::GPUDialect::getWorkgroupAddressSpace()); 475 } 476 477 /// In case of GPU group memory there is no need to deallocate. 478 LogicalResult mlir::linalg::deallocateWorkgroupMemory(OpBuilder &, 479 Value /*buffer*/) { 480 return success(); 481 } 482 483 /// Create Memref copy operations and add gpu barrier guards before and after 484 /// the copy operation to ensure data integrity. 485 LogicalResult mlir::linalg::copyToWorkgroupMemory(OpBuilder &b, Value src, 486 Value dst) { 487 b.create<gpu::BarrierOp>(src.getLoc()); 488 Operation *copyOp = b.create<memref::CopyOp>(src.getLoc(), src, dst); 489 b.create<gpu::BarrierOp>(copyOp->getLoc()); 490 return success(); 491 } 492 493 /// Allocate the subview in the GPU private memory. 494 std::optional<Value> mlir::linalg::allocateGPUPrivateMemory( 495 OpBuilder &builder, memref::SubViewOp subview, ArrayRef<Value> sizeBounds, 496 DataLayout &) { 497 return allocateSubviewGPUMemoryInAddressSpace( 498 builder, subview, sizeBounds, gpu::GPUDialect::getPrivateAddressSpace()); 499 } 500 501 /// Normal copy to between src and dst. 502 LogicalResult mlir::linalg::copyToGPUPrivateMemory(OpBuilder &b, Value src, 503 Value dst) { 504 b.create<memref::CopyOp>(src.getLoc(), src, dst); 505 return success(); 506 } 507 508 /// In case of GPU private memory there is no need to deallocate since the 509 /// memory is freed when going outside of the scope. 510 LogicalResult mlir::linalg::deallocateGPUPrivateMemory(OpBuilder &, 511 Value /*buffer*/) { 512 return success(); 513 } 514