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