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 /// If `size` comes from an AffineMinOp and one of the values of AffineMinOp 48 /// is a constant then return a new value set to the smallest such constant. 49 /// Otherwise return size. 50 static Value extractSmallestConstantBoundingSize(OpBuilder &b, Location loc, 51 Value size) { 52 Optional<int64_t> boundingConst = {}; 53 if (auto affineMinOp = size.getDefiningOp<AffineMinOp>()) { 54 for (auto e : affineMinOp.getAffineMap().getResults()) 55 if (auto cst = e.dyn_cast<AffineConstantExpr>()) 56 boundingConst = boundingConst 57 ? std::min(boundingConst.getValue(), cst.getValue()) 58 : cst.getValue(); 59 } else if (auto constIndexOp = size.getDefiningOp<ConstantOp>()) { 60 if (constIndexOp.getType().isa<IndexType>()) 61 boundingConst = constIndexOp.value().cast<IntegerAttr>().getInt(); 62 } 63 return boundingConst && *boundingConst >= 0 64 ? b.create<ConstantIndexOp>(loc, *boundingConst) 65 : size; 66 } 67 68 /// Alloc a new buffer of `size`. If `dynamicBuffers` is true allocate exactly 69 /// the size needed, otherwise try to allocate a static bounding box. 70 static Value allocBuffer(const LinalgPromotionOptions &options, 71 Type elementType, Value size, bool dynamicBuffers, 72 OperationFolder *folder, 73 Optional<unsigned> alignment = None) { 74 auto *ctx = size.getContext(); 75 auto width = llvm::divideCeil(elementType.getIntOrFloatBitWidth(), 8); 76 IntegerAttr alignment_attr; 77 if (alignment.hasValue()) 78 alignment_attr = 79 IntegerAttr::get(IntegerType::get(64, ctx), alignment.getValue()); 80 if (!dynamicBuffers) 81 if (auto cst = size.getDefiningOp<ConstantIndexOp>()) 82 return options.useAlloca 83 ? std_alloca(MemRefType::get(width * cst.getValue(), 84 IntegerType::get(8, ctx)), 85 ValueRange{}, alignment_attr) 86 .value 87 : std_alloc(MemRefType::get(width * cst.getValue(), 88 IntegerType::get(8, ctx)), 89 ValueRange{}, alignment_attr) 90 .value; 91 Value mul = 92 folded_std_muli(folder, folded_std_constant_index(folder, width), size); 93 return options.useAlloca 94 ? std_alloca(MemRefType::get(-1, IntegerType::get(8, ctx)), mul, 95 alignment_attr) 96 .value 97 : std_alloc(MemRefType::get(-1, IntegerType::get(8, ctx)), mul, 98 alignment_attr) 99 .value; 100 } 101 102 /// Default allocation callback function. This allocates a promoted buffer when 103 /// no call back to do so is provided. The default is to allocate a 104 /// memref<..xi8> and return a view to get a memref type of shape 105 /// boundingSubViewSize. 106 static Optional<Value> defaultAllocBufferCallBack( 107 const LinalgPromotionOptions &options, OpBuilder &builder, 108 SubViewOp subView, ArrayRef<Value> boundingSubViewSize, bool dynamicBuffers, 109 Optional<unsigned> alignment, OperationFolder *folder) { 110 ShapedType viewType = subView.getType(); 111 int64_t rank = viewType.getRank(); 112 (void)rank; 113 assert(rank > 0 && boundingSubViewSize.size() == static_cast<size_t>(rank)); 114 auto zero = folded_std_constant_index(folder, 0); 115 auto one = folded_std_constant_index(folder, 1); 116 117 Value allocSize = one; 118 for (auto size : llvm::enumerate(boundingSubViewSize)) 119 allocSize = folded_std_muli(folder, allocSize, size.value()); 120 Value buffer = allocBuffer(options, viewType.getElementType(), allocSize, 121 dynamicBuffers, folder, alignment); 122 SmallVector<int64_t, 4> dynSizes(boundingSubViewSize.size(), 123 ShapedType::kDynamicSize); 124 Value view = folded_std_view( 125 folder, MemRefType::get(dynSizes, viewType.getElementType()), buffer, 126 zero, boundingSubViewSize); 127 return view; 128 } 129 130 /// Default implementation of deallocation of the buffer use for promotion. It 131 /// expects to get the same value that the default allocation method returned, 132 /// i.e. result of a ViewOp. 133 static LogicalResult 134 defaultDeallocBufferCallBack(const LinalgPromotionOptions &options, 135 OpBuilder &b, Value fullLocalView) { 136 auto viewOp = fullLocalView.getDefiningOp<ViewOp>(); 137 assert(viewOp && "expected full local view to be a ViewOp"); 138 if (!options.useAlloca) 139 std_dealloc(viewOp.source()); 140 return success(); 141 } 142 143 namespace { 144 145 /// Helper struct that captures the information required to apply the 146 /// transformation on each op. This bridges the abstraction gap with the 147 /// user-facing API which exposes positional arguments to control which operands 148 /// are promoted. 149 struct LinalgOpInstancePromotionOptions { 150 LinalgOpInstancePromotionOptions(LinalgOp op, 151 const LinalgPromotionOptions &options); 152 /// SubViews to promote. 153 MapVector<unsigned, Value> subViews; 154 /// True if the full view should be used for the promoted buffer. 155 DenseMap<Value, bool> useFullTileBuffers; 156 157 /// Callback functions for allocation and deallocation of promoted buffers, as 158 /// well as to copy the data into and out of these buffers. 159 AllocBufferCallbackFn allocationFn; 160 DeallocBufferCallbackFn deallocationFn; 161 CopyCallbackFn copyInFn; 162 CopyCallbackFn copyOutFn; 163 164 /// Allow the use of dynamicaly-sized buffers. 165 bool dynamicBuffers; 166 /// Alignment of promoted buffer. 167 Optional<unsigned> alignment; 168 }; 169 170 struct PromotionInfo { 171 Value fullLocalView; 172 Value partialLocalView; 173 }; 174 } // namespace 175 176 LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions( 177 LinalgOp linalgOp, const LinalgPromotionOptions &options) 178 : subViews(), dynamicBuffers(options.dynamicBuffers), 179 alignment(options.alignment) { 180 unsigned nBuffers = linalgOp.getNumInputsAndOutputBuffers(); 181 auto vUseFullTileBuffers = 182 options.useFullTileBuffers.getValueOr(llvm::SmallBitVector()); 183 vUseFullTileBuffers.resize(nBuffers, options.useFullTileBuffersDefault); 184 185 for (unsigned idx = 0; idx != nBuffers; ++idx) { 186 if (options.operandsToPromote && !options.operandsToPromote->count(idx)) 187 continue; 188 auto *op = linalgOp.getBuffer(idx).getDefiningOp(); 189 if (auto sv = dyn_cast_or_null<SubViewOp>(op)) { 190 subViews[idx] = sv; 191 useFullTileBuffers[sv] = vUseFullTileBuffers[idx]; 192 } 193 } 194 195 allocationFn = 196 (options.allocationFn ? *(options.allocationFn) 197 : [&](OpBuilder &builder, SubViewOp subViewOp, 198 ArrayRef<Value> boundingSubViewSize, 199 OperationFolder *folder) -> Optional<Value> { 200 return defaultAllocBufferCallBack(options, builder, subViewOp, 201 boundingSubViewSize, dynamicBuffers, 202 alignment, folder); 203 }); 204 deallocationFn = 205 (options.deallocationFn 206 ? *(options.deallocationFn) 207 : [&](OpBuilder &b, Value buffer) { 208 return defaultDeallocBufferCallBack(options, b, buffer); 209 }); 210 auto defaultCopyCallBack = [&](OpBuilder &builder, Value src, 211 Value dst) -> LogicalResult { 212 linalg_copy(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 static Optional<PromotionInfo> 237 promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, SubViewOp subView, 238 LinalgOpInstancePromotionOptions const &options, 239 OperationFolder *folder) { 240 auto viewType = subView.getType(); 241 auto rank = viewType.getRank(); 242 SmallVector<Value, 4> fullSizes, partialSizes; 243 fullSizes.reserve(rank); 244 partialSizes.reserve(rank); 245 for (auto en : llvm::enumerate(subView.getOrCreateRanges(b, loc))) { 246 auto rangeValue = en.value(); 247 // Try to extract a tight constant. 248 LLVM_DEBUG(llvm::dbgs() << "Extract tightest: " << rangeValue.size << "\n"); 249 Value size = extractSmallestConstantBoundingSize(b, loc, rangeValue.size); 250 LLVM_DEBUG(llvm::dbgs() << "Extracted tightest: " << size << "\n"); 251 fullSizes.push_back(size); 252 partialSizes.push_back(folded_std_dim(folder, subView, en.index())); 253 } 254 SmallVector<int64_t, 4> dynSizes(fullSizes.size(), -1); 255 // If a callback is not specified, then use the default implementation for 256 // allocating the promoted buffer. 257 Optional<Value> fullLocalView = 258 options.allocationFn(b, subView, fullSizes, folder); 259 if (!fullLocalView) 260 return {}; 261 auto zero = folded_std_constant_index(folder, 0); 262 auto one = folded_std_constant_index(folder, 1); 263 SmallVector<Value, 4> zeros(fullSizes.size(), zero); 264 SmallVector<Value, 4> ones(fullSizes.size(), one); 265 auto partialLocalView = 266 folded_std_subview(folder, *fullLocalView, zeros, partialSizes, ones); 267 return PromotionInfo{*fullLocalView, partialLocalView}; 268 } 269 270 static Optional<MapVector<unsigned, PromotionInfo>> 271 promoteSubViews(OpBuilder &b, Location loc, 272 LinalgOpInstancePromotionOptions options, 273 OperationFolder *folder) { 274 if (options.subViews.empty()) 275 return {}; 276 277 ScopedContext scope(b, loc); 278 MapVector<unsigned, PromotionInfo> promotionInfoMap; 279 280 for (auto v : options.subViews) { 281 SubViewOp subView = cast<SubViewOp>(v.second.getDefiningOp()); 282 Optional<PromotionInfo> promotionInfo = 283 promoteSubviewAsNewBuffer(b, loc, subView, options, folder); 284 if (!promotionInfo) 285 return {}; 286 promotionInfoMap[v.first] = *promotionInfo; 287 288 // Only fill the buffer if the full local view is used 289 if (!options.useFullTileBuffers[v.second]) 290 continue; 291 Value fillVal; 292 if (auto t = subView.getType().getElementType().dyn_cast<FloatType>()) 293 fillVal = folded_std_constant(folder, FloatAttr::get(t, 0.0)); 294 else if (auto t = 295 subView.getType().getElementType().dyn_cast<IntegerType>()) 296 fillVal = folded_std_constant_int(folder, 0, t); 297 linalg_fill(promotionInfo->fullLocalView, fillVal); 298 } 299 300 // Copy data into the promoted buffers. Use callback if provided. 301 for (auto v : options.subViews) { 302 auto info = promotionInfoMap.find(v.first); 303 if (info == promotionInfoMap.end()) 304 continue; 305 if (failed(options.copyInFn(b, cast<SubViewOp>(v.second.getDefiningOp()), 306 info->second.partialLocalView))) 307 return {}; 308 } 309 return promotionInfoMap; 310 } 311 312 static Optional<LinalgOp> 313 promoteSubViews(OpBuilder &b, LinalgOp op, 314 LinalgOpInstancePromotionOptions options, 315 OperationFolder *folder) { 316 assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics"); 317 318 if (auto convOp = dyn_cast<linalg::ConvOp>(op.getOperation())) { 319 // TODO(ntv): add a level of indirection to linalg.generic. 320 if (convOp.padding()) 321 return {}; 322 } 323 324 // 1. Promote the specified views and use them in the new op. 325 auto loc = op.getLoc(); 326 auto promotedBuffersAndViews = promoteSubViews(b, loc, options, folder); 327 if (!promotedBuffersAndViews || 328 promotedBuffersAndViews->size() != options.subViews.size()) 329 return {}; 330 331 // 2. Append all other operands as they appear, this enforces that such 332 // operands are not views. This is to support cases such as FillOp taking 333 // extra scalars etc. Keep a reference to output buffers; 334 SmallVector<Value, 8> opViews; 335 opViews.reserve(op.getNumInputsAndOutputs()); 336 SmallVector<std::pair<Value, Value>, 8> writebackViews; 337 writebackViews.reserve(promotedBuffersAndViews->size()); 338 for (auto view : llvm::enumerate(op.getInputsAndOutputBuffers())) { 339 if (options.subViews.count(view.index()) != 0) { 340 if (options.useFullTileBuffers[view.value()]) 341 opViews.push_back( 342 (*promotedBuffersAndViews)[view.index()].fullLocalView); 343 else 344 opViews.push_back( 345 (*promotedBuffersAndViews)[view.index()].partialLocalView); 346 if (view.index() >= op.getNumInputs()) 347 writebackViews.emplace_back(std::make_pair( 348 view.value(), 349 (*promotedBuffersAndViews)[view.index()].partialLocalView)); 350 } else { 351 opViews.push_back(view.value()); 352 } 353 } 354 op.getOperation()->setOperands(0, opViews.size(), opViews); 355 356 OpBuilder::InsertionGuard guard(b); 357 b.setInsertionPointAfter(op); 358 ScopedContext scope(b, loc); 359 // 3. Emit write-back for the promoted output views: copy the partial view. 360 for (auto viewAndPartialLocalView : writebackViews) { 361 if (failed(options.copyOutFn(b, viewAndPartialLocalView.second, 362 viewAndPartialLocalView.first))) 363 return {}; 364 } 365 366 // 4. Dealloc all local buffers. 367 for (const auto &pi : *promotedBuffersAndViews) 368 options.deallocationFn(b, pi.second.fullLocalView); 369 return op; 370 } 371 372 LogicalResult 373 mlir::linalg::promoteSubviewsPrecondition(Operation *op, 374 LinalgPromotionOptions options) { 375 LinalgOp linOp = dyn_cast<LinalgOp>(op); 376 // Transformation applies to buffers only. 377 if (!linOp || !linOp.hasBufferSemantics()) 378 return failure(); 379 // Check that at least one of the requested operands is indeed a subview. 380 for (auto en : llvm::enumerate(linOp.getInputsAndOutputBuffers())) { 381 auto sv = isa_and_nonnull<SubViewOp>(en.value().getDefiningOp()); 382 if (sv) { 383 if (!options.operandsToPromote.hasValue() || 384 options.operandsToPromote->count(en.index())) 385 return success(); 386 } 387 } 388 // TODO: Check all subviews requested are bound by a static constant. 389 // TODO: Check that the total footprint fits within a given size. 390 return failure(); 391 } 392 393 Optional<LinalgOp> mlir::linalg::promoteSubViews(OpBuilder &b, 394 LinalgOp linalgOp, 395 LinalgPromotionOptions options, 396 OperationFolder *folder) { 397 LinalgOpInstancePromotionOptions linalgOptions(linalgOp, options); 398 return ::promoteSubViews( 399 b, linalgOp, LinalgOpInstancePromotionOptions(linalgOp, options), folder); 400 } 401 402 namespace { 403 struct LinalgPromotionPass : public LinalgPromotionBase<LinalgPromotionPass> { 404 LinalgPromotionPass() = default; 405 LinalgPromotionPass(bool dynamicBuffers, bool useAlloca) { 406 this->dynamicBuffers = dynamicBuffers; 407 this->useAlloca = useAlloca; 408 } 409 410 void runOnFunction() override { 411 OperationFolder folder(&getContext()); 412 getFunction().walk([this, &folder](LinalgOp op) { 413 auto options = LinalgPromotionOptions() 414 .setDynamicBuffers(dynamicBuffers) 415 .setUseAlloca(useAlloca); 416 if (failed(promoteSubviewsPrecondition(op, options))) 417 return; 418 LLVM_DEBUG(llvm::dbgs() << "Promote: " << *(op.getOperation()) << "\n"); 419 OpBuilder b(op); 420 promoteSubViews(b, op, options, &folder); 421 }); 422 } 423 }; 424 } // namespace 425 426 // TODO: support more transformation options in the pass. 427 std::unique_ptr<OperationPass<FuncOp>> 428 mlir::createLinalgPromotionPass(bool dynamicBuffers, bool useAlloca) { 429 return std::make_unique<LinalgPromotionPass>(dynamicBuffers, useAlloca); 430 } 431 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgPromotionPass() { 432 return std::make_unique<LinalgPromotionPass>(); 433 } 434