1 //===- Bufferize.cpp - Bufferization utilities ----------------------------===// 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 #include "PassDetail.h" 10 11 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h" 12 #include "mlir/Dialect/Bufferization/IR/Bufferization.h" 13 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h" 14 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h" 15 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h" 16 #include "mlir/Dialect/Bufferization/Transforms/Passes.h" 17 #include "mlir/Dialect/Bufferization/Transforms/TensorCopyInsertion.h" 18 #include "mlir/Dialect/Func/IR/FuncOps.h" 19 #include "mlir/Dialect/MemRef/IR/MemRef.h" 20 #include "mlir/IR/Operation.h" 21 #include "mlir/Pass/PassManager.h" 22 #include "mlir/Transforms/GreedyPatternRewriteDriver.h" 23 #include "mlir/Transforms/Passes.h" 24 25 using namespace mlir; 26 using namespace mlir::bufferization; 27 28 //===----------------------------------------------------------------------===// 29 // BufferizeTypeConverter 30 //===----------------------------------------------------------------------===// 31 32 static Value materializeToTensor(OpBuilder &builder, TensorType type, 33 ValueRange inputs, Location loc) { 34 assert(inputs.size() == 1); 35 assert(inputs[0].getType().isa<BaseMemRefType>()); 36 return builder.create<bufferization::ToTensorOp>(loc, type, inputs[0]); 37 } 38 39 /// Registers conversions into BufferizeTypeConverter 40 BufferizeTypeConverter::BufferizeTypeConverter() { 41 // Keep all types unchanged. 42 addConversion([](Type type) { return type; }); 43 // Convert RankedTensorType to MemRefType. 44 addConversion([](RankedTensorType type) -> Type { 45 return MemRefType::get(type.getShape(), type.getElementType()); 46 }); 47 // Convert UnrankedTensorType to UnrankedMemRefType. 48 addConversion([](UnrankedTensorType type) -> Type { 49 return UnrankedMemRefType::get(type.getElementType(), 0); 50 }); 51 addArgumentMaterialization(materializeToTensor); 52 addSourceMaterialization(materializeToTensor); 53 addTargetMaterialization([](OpBuilder &builder, BaseMemRefType type, 54 ValueRange inputs, Location loc) -> Value { 55 assert(inputs.size() == 1 && "expected exactly one input"); 56 57 if (auto inputType = inputs[0].getType().dyn_cast<MemRefType>()) { 58 // MemRef to MemRef cast. 59 assert(inputType != type && "expected different types"); 60 // Unranked to ranked and ranked to unranked casts must be explicit. 61 auto rankedDestType = type.dyn_cast<MemRefType>(); 62 if (!rankedDestType) 63 return nullptr; 64 FailureOr<Value> replacement = 65 castOrReallocMemRefValue(builder, inputs[0], rankedDestType); 66 if (failed(replacement)) 67 return nullptr; 68 return *replacement; 69 } 70 71 if (inputs[0].getType().isa<TensorType>()) { 72 // Tensor to MemRef cast. 73 return builder.create<bufferization::ToMemrefOp>(loc, type, inputs[0]); 74 } 75 76 llvm_unreachable("only tensor/memref input types supported"); 77 }); 78 } 79 80 void mlir::bufferization::populateBufferizeMaterializationLegality( 81 ConversionTarget &target) { 82 target.addLegalOp<bufferization::ToTensorOp, bufferization::ToMemrefOp>(); 83 } 84 85 namespace { 86 // In a finalizing bufferize conversion, we know that all tensors have been 87 // converted to memrefs, thus, this op becomes an identity. 88 class BufferizeToTensorOp 89 : public OpConversionPattern<bufferization::ToTensorOp> { 90 public: 91 using OpConversionPattern::OpConversionPattern; 92 LogicalResult 93 matchAndRewrite(bufferization::ToTensorOp op, OpAdaptor adaptor, 94 ConversionPatternRewriter &rewriter) const override { 95 rewriter.replaceOp(op, adaptor.getMemref()); 96 return success(); 97 } 98 }; 99 } // namespace 100 101 namespace { 102 // In a finalizing bufferize conversion, we know that all tensors have been 103 // converted to memrefs, thus, this op becomes an identity. 104 class BufferizeToMemrefOp 105 : public OpConversionPattern<bufferization::ToMemrefOp> { 106 public: 107 using OpConversionPattern::OpConversionPattern; 108 LogicalResult 109 matchAndRewrite(bufferization::ToMemrefOp op, OpAdaptor adaptor, 110 ConversionPatternRewriter &rewriter) const override { 111 rewriter.replaceOp(op, adaptor.getTensor()); 112 return success(); 113 } 114 }; 115 } // namespace 116 117 void mlir::bufferization::populateEliminateBufferizeMaterializationsPatterns( 118 BufferizeTypeConverter &typeConverter, RewritePatternSet &patterns) { 119 patterns.add<BufferizeToTensorOp, BufferizeToMemrefOp>(typeConverter, 120 patterns.getContext()); 121 } 122 123 namespace { 124 struct FinalizingBufferizePass 125 : public FinalizingBufferizeBase<FinalizingBufferizePass> { 126 using FinalizingBufferizeBase< 127 FinalizingBufferizePass>::FinalizingBufferizeBase; 128 129 void runOnOperation() override { 130 auto func = getOperation(); 131 auto *context = &getContext(); 132 133 BufferizeTypeConverter typeConverter; 134 RewritePatternSet patterns(context); 135 ConversionTarget target(*context); 136 137 populateEliminateBufferizeMaterializationsPatterns(typeConverter, patterns); 138 139 // If all result types are legal, and all block arguments are legal (ensured 140 // by func conversion above), then all types in the program are legal. 141 // 142 // We also check that the operand types are legal to avoid creating invalid 143 // IR. For example, this prevents 144 // populateEliminateBufferizeMaterializationsPatterns from updating the 145 // types of the operands to a return op without updating the enclosing 146 // function. 147 target.markUnknownOpDynamicallyLegal( 148 [&](Operation *op) { return typeConverter.isLegal(op); }); 149 150 if (failed(applyFullConversion(func, target, std::move(patterns)))) 151 signalPassFailure(); 152 } 153 }; 154 155 static BufferizationOptions::LayoutMapOption 156 parseLayoutMapOption(const std::string &s) { 157 if (s == "fully-dynamic-layout-map") 158 return BufferizationOptions::LayoutMapOption::FullyDynamicLayoutMap; 159 if (s == "identity-layout-map") 160 return BufferizationOptions::LayoutMapOption::IdentityLayoutMap; 161 if (s == "infer-layout-map") 162 return BufferizationOptions::LayoutMapOption::InferLayoutMap; 163 llvm_unreachable("invalid layout map option"); 164 } 165 166 struct OneShotBufferizePass 167 : public OneShotBufferizeBase<OneShotBufferizePass> { 168 OneShotBufferizePass() : OneShotBufferizeBase<OneShotBufferizePass>() {} 169 170 explicit OneShotBufferizePass(const OneShotBufferizationOptions &options) 171 : options(options) {} 172 173 void getDependentDialects(DialectRegistry ®istry) const override { 174 registry 175 .insert<bufferization::BufferizationDialect, memref::MemRefDialect>(); 176 registerAllocationOpInterfaceExternalModels(registry); 177 } 178 179 void runOnOperation() override { 180 OneShotBufferizationOptions opt; 181 if (!options) { 182 // Make new bufferization options if none were provided when creating the 183 // pass. 184 opt.allowReturnAllocs = allowReturnAllocs; 185 opt.allowUnknownOps = allowUnknownOps; 186 opt.analysisFuzzerSeed = analysisFuzzerSeed; 187 opt.createDeallocs = createDeallocs; 188 opt.functionBoundaryTypeConversion = 189 parseLayoutMapOption(functionBoundaryTypeConversion); 190 if (mustInferMemorySpace) 191 opt.defaultMemorySpace = None; 192 opt.printConflicts = printConflicts; 193 opt.testAnalysisOnly = testAnalysisOnly; 194 opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries; 195 opt.unknownTypeConversion = parseLayoutMapOption(unknownTypeConversion); 196 197 OpFilter::Entry::FilterFn filterFn = 198 [&](Operation *op) { 199 // Filter may be specified via options. 200 if (this->dialectFilter.hasValue()) 201 return llvm::is_contained(this->dialectFilter, 202 op->getDialect()->getNamespace()); 203 // No filter specified: All other ops are allowed. 204 return true; 205 }; 206 opt.opFilter.allowOperation(filterFn); 207 } else { 208 opt = *options; 209 } 210 211 ModuleOp moduleOp = getOperation(); 212 if (opt.bufferizeFunctionBoundaries) { 213 if (failed(runOneShotModuleBufferize(moduleOp, opt))) { 214 signalPassFailure(); 215 return; 216 } 217 } else { 218 if (failed(runOneShotBufferize(moduleOp, opt))) { 219 signalPassFailure(); 220 return; 221 } 222 } 223 224 if (opt.testAnalysisOnly) 225 return; 226 227 OpPassManager cleanupPipeline("builtin.module"); 228 cleanupPipeline.addPass(createCanonicalizerPass()); 229 cleanupPipeline.addPass(createCSEPass()); 230 cleanupPipeline.addPass(createLoopInvariantCodeMotionPass()); 231 (void)runPipeline(cleanupPipeline, moduleOp); 232 } 233 234 private: 235 llvm::Optional<OneShotBufferizationOptions> options; 236 }; 237 } // namespace 238 239 namespace { 240 struct BufferizationBufferizePass 241 : public BufferizationBufferizeBase<BufferizationBufferizePass> { 242 void runOnOperation() override { 243 BufferizationOptions options = getPartialBufferizationOptions(); 244 options.opFilter.allowDialect<BufferizationDialect>(); 245 246 if (failed(bufferizeOp(getOperation(), options))) 247 signalPassFailure(); 248 } 249 250 void getDependentDialects(DialectRegistry ®istry) const override { 251 registry 252 .insert<bufferization::BufferizationDialect, memref::MemRefDialect>(); 253 } 254 }; 255 } // namespace 256 257 std::unique_ptr<Pass> mlir::bufferization::createBufferizationBufferizePass() { 258 return std::make_unique<BufferizationBufferizePass>(); 259 } 260 261 std::unique_ptr<Pass> mlir::bufferization::createOneShotBufferizePass() { 262 return std::make_unique<OneShotBufferizePass>(); 263 } 264 265 std::unique_ptr<Pass> mlir::bufferization::createOneShotBufferizePass( 266 const OneShotBufferizationOptions &options) { 267 return std::make_unique<OneShotBufferizePass>(options); 268 } 269 270 std::unique_ptr<OperationPass<func::FuncOp>> 271 mlir::bufferization::createFinalizingBufferizePass() { 272 return std::make_unique<FinalizingBufferizePass>(); 273 } 274 275 //===----------------------------------------------------------------------===// 276 // BufferizableOpInterface-based Bufferization 277 //===----------------------------------------------------------------------===// 278 279 static bool isaTensor(Type t) { return t.isa<TensorType>(); } 280 281 /// Return true if the given op has a tensor result or a tensor operand. 282 static bool hasTensorSemantics(Operation *op) { 283 if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) { 284 bool hasTensorArg = any_of(funcOp.getArgumentTypes(), isaTensor); 285 bool hasTensorResult = any_of(funcOp.getResultTypes(), isaTensor); 286 return hasTensorArg || hasTensorResult; 287 } 288 289 bool hasTensorResult = any_of(op->getResultTypes(), isaTensor); 290 bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor); 291 return hasTensorResult || hasTensorOperand; 292 } 293 294 namespace { 295 /// A rewriter that keeps track of extra information during bufferization. 296 class BufferizationRewriter : public IRRewriter { 297 public: 298 BufferizationRewriter(MLIRContext *ctx, DenseSet<Operation *> &erasedOps, 299 DenseSet<Operation *> &toMemrefOps, 300 SmallVector<Operation *> &worklist, 301 const BufferizationOptions &options, 302 const OpFilter *opFilter) 303 : IRRewriter(ctx), erasedOps(erasedOps), toMemrefOps(toMemrefOps), 304 worklist(worklist), analysisState(options), opFilter(opFilter) {} 305 306 protected: 307 void notifyOperationRemoved(Operation *op) override { 308 IRRewriter::notifyOperationRemoved(op); 309 erasedOps.insert(op); 310 // Erase if present. 311 toMemrefOps.erase(op); 312 } 313 314 void notifyOperationInserted(Operation *op) override { 315 IRRewriter::notifyOperationInserted(op); 316 erasedOps.erase(op); 317 318 // Keep track of to_memref ops. 319 if (isa<ToMemrefOp>(op)) { 320 toMemrefOps.insert(op); 321 return; 322 } 323 324 // Skip to_tensor ops. 325 if (isa<ToTensorOp>(op)) 326 return; 327 328 // Skip non-tensor ops. 329 if (!hasTensorSemantics(op)) 330 return; 331 332 // Skip ops that are not allowed to be bufferized. 333 auto const &options = analysisState.getOptions(); 334 if (!options.isOpAllowed(op) || (opFilter && !opFilter->isOpAllowed(op))) 335 return; 336 337 #ifndef NDEBUG 338 // Read-only tensor ops may be created during bufferization. Ops that are 339 // writing should not be created because such ops were never analyzed. 340 // Bufferizing such ops could introduce a RaW conflict. 341 for (OpOperand &operand : op->getOpOperands()) 342 if (operand.get().getType().isa<TensorType>()) 343 assert(!analysisState.bufferizesToMemoryWrite(operand) && 344 "creating tensor ops that bufferize to a memory write is not " 345 "allowed during bufferization"); 346 #endif // NDEBUG 347 348 // Add op to worklist. 349 worklist.push_back(op); 350 } 351 352 private: 353 /// A set of all erased ops. 354 DenseSet<Operation *> &erasedOps; 355 356 /// A set of all to_memref ops. 357 DenseSet<Operation *> &toMemrefOps; 358 359 /// The worklist of ops to be bufferized. 360 SmallVector<Operation *> &worklist; 361 362 /// The analysis state. Used for debug assertions and access to the 363 /// bufferization options. 364 const AnalysisState analysisState; 365 366 /// An extra op filter for bufferization. 367 const OpFilter *opFilter; 368 }; 369 } // namespace 370 371 LogicalResult bufferization::bufferizeOp(Operation *op, 372 const BufferizationOptions &options, 373 bool copyBeforeWrite, 374 const OpFilter *opFilter) { 375 assert(options.unknownTypeConversion != 376 BufferizationOptions::LayoutMapOption::InferLayoutMap && 377 "invalid layout map option"); 378 379 if (copyBeforeWrite) { 380 AnalysisState state(options); 381 if (failed(insertTensorCopies(op, state))) 382 return failure(); 383 } 384 385 // Keep track of to_memref ops. 386 DenseSet<Operation *> toMemrefOps; 387 op->walk([&](ToMemrefOp toMemrefOp) { toMemrefOps.insert(toMemrefOp); }); 388 389 // Gather all bufferizable ops in top-to-bottom order. 390 // 391 // We should ideally know the exact memref type of all operands when 392 // bufferizing an op. (This is the case when bufferizing top-to-bottom.) 393 // Otherwise, we have to use a memref type with a fully dynamic layout map to 394 // avoid copies. We are currently missing patterns for layout maps to 395 // canonicalize away (or canonicalize to more precise layouts). 396 // 397 // FuncOps must be bufferized before their bodies, so add them to the worklist 398 // first. 399 SmallVector<Operation *> worklist; 400 op->walk([&](func::FuncOp funcOp) { 401 if (hasTensorSemantics(funcOp)) 402 worklist.push_back(funcOp); 403 }); 404 op->walk<WalkOrder::PostOrder>([&](Operation *op) { 405 if (hasTensorSemantics(op) && !isa<func::FuncOp>(op)) 406 worklist.push_back(op); 407 }); 408 409 // Keep track of all erased ops. 410 DenseSet<Operation *> erasedOps; 411 412 // Bufferize all ops. 413 BufferizationRewriter rewriter(op->getContext(), erasedOps, toMemrefOps, 414 worklist, options, opFilter); 415 for (unsigned i = 0; i < worklist.size(); ++i) { 416 Operation *op = worklist[i]; 417 // Skip ops that were erased. 418 if (erasedOps.contains(op)) 419 continue; 420 // Skip ops that are not bufferizable or not allowed. 421 auto bufferizableOp = options.dynCastBufferizableOp(op); 422 if (!bufferizableOp) 423 continue; 424 if (opFilter && !opFilter->isOpAllowed(op)) 425 continue; 426 // Skip ops that no longer have tensor semantics. 427 if (!hasTensorSemantics(op)) 428 continue; 429 // Bufferize the op. 430 rewriter.setInsertionPoint(op); 431 if (failed(bufferizableOp.bufferize(rewriter, options))) 432 return op->emitError("failed to bufferize op"); 433 } 434 435 // Fold all to_memref(to_tensor(x)) pairs. 436 for (Operation *op : toMemrefOps) { 437 rewriter.setInsertionPoint(op); 438 (void)bufferization::foldToMemrefToTensorPair(rewriter, 439 cast<ToMemrefOp>(op)); 440 } 441 442 /// Check the result of bufferization. Return an error if an op was not 443 /// bufferized, unless partial bufferization is allowed. 444 if (options.allowUnknownOps) 445 return success(); 446 447 for (Operation *op : worklist) { 448 // Skip ops that are entirely gone. 449 if (erasedOps.contains(op)) 450 continue; 451 // Ops that no longer have tensor semantics (because they were updated 452 // in-place) are allowed. 453 if (!hasTensorSemantics(op)) 454 continue; 455 // Continue ops that are not allowed. 456 if (!options.isOpAllowed(op)) 457 continue; 458 if (opFilter && !opFilter->isOpAllowed(op)) 459 continue; 460 // Ops without any uses and no side effects will fold away. 461 if (op->getUses().empty() && MemoryEffectOpInterface::hasNoEffect(op)) 462 continue; 463 // ToTensorOps/ToMemrefOps are allowed in the output. 464 if (isa<ToTensorOp, ToMemrefOp>(op)) 465 continue; 466 return op->emitError("op was not bufferized"); 467 } 468 469 return success(); 470 } 471 472 BufferizationOptions bufferization::getPartialBufferizationOptions() { 473 BufferizationOptions options; 474 options.allowUnknownOps = true; 475 options.createDeallocs = false; 476 options.enforceAliasingInvariants = false; 477 options.unknownTypeConversion = 478 BufferizationOptions::LayoutMapOption::IdentityLayoutMap; 479 options.opFilter.allowDialect<BufferizationDialect>(); 480 return options; 481 } 482