| //===- LowerDeallocations.cpp - Bufferization Deallocs to MemRef pass -----===// | |
| // | |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. | |
| // See https://llvm.org/LICENSE.txt for license information. | |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |
| // | |
| //===----------------------------------------------------------------------===// | |
| // | |
| // This file implements patterns to convert `bufferization.dealloc` operations | |
| // to the MemRef dialect. | |
| // | |
| //===----------------------------------------------------------------------===// | |
| #include "mlir/Dialect/Arith/IR/Arith.h" | |
| #include "mlir/Dialect/Bufferization/IR/Bufferization.h" | |
| #include "mlir/Dialect/Bufferization/Transforms/Passes.h" | |
| #include "mlir/Dialect/Func/IR/FuncOps.h" | |
| #include "mlir/Dialect/MemRef/IR/MemRef.h" | |
| #include "mlir/Dialect/SCF/IR/SCF.h" | |
| #include "mlir/IR/BuiltinTypes.h" | |
| #include "mlir/Transforms/DialectConversion.h" | |
| namespace mlir { | |
| namespace bufferization { | |
| #define GEN_PASS_DEF_LOWERDEALLOCATIONSPASS | |
| #include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc" | |
| } // namespace bufferization | |
| } // namespace mlir | |
| using namespace mlir; | |
| namespace { | |
| /// The DeallocOpConversion transforms all bufferization dealloc operations into | |
| /// memref dealloc operations potentially guarded by scf if operations. | |
| /// Additionally, memref extract_aligned_pointer_as_index and arith operations | |
| /// are inserted to compute the guard conditions. We distinguish multiple cases | |
| /// to provide an overall more efficient lowering. In the general case, a helper | |
| /// func is created to avoid quadratic code size explosion (relative to the | |
| /// number of operands of the dealloc operation). For examples of each case, | |
| /// refer to the documentation of the member functions of this class. | |
| class DeallocOpConversion | |
| : public OpConversionPattern<bufferization::DeallocOp> { | |
| /// Lower a simple case without any retained values and a single memref to | |
| /// avoiding the helper function. Ideally, static analysis can provide enough | |
| /// aliasing information to split the dealloc operations up into this simple | |
| /// case as much as possible before running this pass. | |
| /// | |
| /// Example: | |
| /// ``` | |
| /// bufferization.dealloc (%arg0 : memref<2xf32>) if (%arg1) | |
| /// ``` | |
| /// is lowered to | |
| /// ``` | |
| /// scf.if %arg1 { | |
| /// memref.dealloc %arg0 : memref<2xf32> | |
| /// } | |
| /// ``` | |
| LogicalResult | |
| rewriteOneMemrefNoRetainCase(bufferization::DeallocOp op, OpAdaptor adaptor, | |
| ConversionPatternRewriter &rewriter) const { | |
| assert(adaptor.getMemrefs().size() == 1 && "expected only one memref"); | |
| assert(adaptor.getRetained().empty() && "expected no retained memrefs"); | |
| rewriter.replaceOpWithNewOp<scf::IfOp>( | |
| op, adaptor.getConditions()[0], [&](OpBuilder &builder, Location loc) { | |
| memref::DeallocOp::create(builder, loc, adaptor.getMemrefs()[0]); | |
| scf::YieldOp::create(builder, loc); | |
| }); | |
| return success(); | |
| } | |
| /// A special case lowering for the deallocation operation with exactly one | |
| /// memref, but arbitrary number of retained values. This avoids the helper | |
| /// function that the general case needs and thus also avoids storing indices | |
| /// to specifically allocated memrefs. The size of the code produced by this | |
| /// lowering is linear to the number of retained values. | |
| /// | |
| /// Example: | |
| /// ```mlir | |
| /// %0:2 = bufferization.dealloc (%m : memref<2xf32>) if (%cond) | |
| // retain (%r0, %r1 : memref<1xf32>, memref<2xf32>) | |
| /// return %0#0, %0#1 : i1, i1 | |
| /// ``` | |
| /// ```mlir | |
| /// %m_base_pointer = memref.extract_aligned_pointer_as_index %m | |
| /// %r0_base_pointer = memref.extract_aligned_pointer_as_index %r0 | |
| /// %r0_does_not_alias = arith.cmpi ne, %m_base_pointer, %r0_base_pointer | |
| /// %r1_base_pointer = memref.extract_aligned_pointer_as_index %r1 | |
| /// %r1_does_not_alias = arith.cmpi ne, %m_base_pointer, %r1_base_pointer | |
| /// %not_retained = arith.andi %r0_does_not_alias, %r1_does_not_alias : i1 | |
| /// %should_dealloc = arith.andi %not_retained, %cond : i1 | |
| /// scf.if %should_dealloc { | |
| /// memref.dealloc %m : memref<2xf32> | |
| /// } | |
| /// %true = arith.constant true | |
| /// %r0_does_alias = arith.xori %r0_does_not_alias, %true : i1 | |
| /// %r0_ownership = arith.andi %r0_does_alias, %cond : i1 | |
| /// %r1_does_alias = arith.xori %r1_does_not_alias, %true : i1 | |
| /// %r1_ownership = arith.andi %r1_does_alias, %cond : i1 | |
| /// return %r0_ownership, %r1_ownership : i1, i1 | |
| /// ``` | |
| LogicalResult rewriteOneMemrefMultipleRetainCase( | |
| bufferization::DeallocOp op, OpAdaptor adaptor, | |
| ConversionPatternRewriter &rewriter) const { | |
| assert(adaptor.getMemrefs().size() == 1 && "expected only one memref"); | |
| // Compute the base pointer indices, compare all retained indices to the | |
| // memref index to check if they alias. | |
| SmallVector<Value> doesNotAliasList; | |
| Value memrefAsIdx = memref::ExtractAlignedPointerAsIndexOp::create( | |
| rewriter, op->getLoc(), adaptor.getMemrefs()[0]); | |
| for (Value retained : adaptor.getRetained()) { | |
| Value retainedAsIdx = memref::ExtractAlignedPointerAsIndexOp::create( | |
| rewriter, op->getLoc(), retained); | |
| Value doesNotAlias = arith::CmpIOp::create(rewriter, op->getLoc(), | |
| arith::CmpIPredicate::ne, | |
| memrefAsIdx, retainedAsIdx); | |
| doesNotAliasList.push_back(doesNotAlias); | |
| } | |
| // AND-reduce the list of booleans from above. | |
| Value prev = doesNotAliasList.front(); | |
| for (Value doesNotAlias : ArrayRef(doesNotAliasList).drop_front()) | |
| prev = arith::AndIOp::create(rewriter, op->getLoc(), prev, doesNotAlias); | |
| // Also consider the condition given by the dealloc operation and perform a | |
| // conditional deallocation guarded by that value. | |
| Value shouldDealloc = arith::AndIOp::create(rewriter, op->getLoc(), prev, | |
| adaptor.getConditions()[0]); | |
| scf::IfOp::create(rewriter, op.getLoc(), shouldDealloc, | |
| [&](OpBuilder &builder, Location loc) { | |
| memref::DeallocOp::create(builder, loc, | |
| adaptor.getMemrefs()[0]); | |
| scf::YieldOp::create(builder, loc); | |
| }); | |
| // Compute the replacement values for the dealloc operation results. This | |
| // inserts an already canonicalized form of | |
| // `select(does_alias_with_memref(r), memref_cond, false)` for each retained | |
| // value r. | |
| SmallVector<Value> replacements; | |
| Value trueVal = arith::ConstantOp::create(rewriter, op->getLoc(), | |
| rewriter.getBoolAttr(true)); | |
| for (Value doesNotAlias : doesNotAliasList) { | |
| Value aliases = | |
| arith::XOrIOp::create(rewriter, op->getLoc(), doesNotAlias, trueVal); | |
| Value result = arith::AndIOp::create(rewriter, op->getLoc(), aliases, | |
| adaptor.getConditions()[0]); | |
| replacements.push_back(result); | |
| } | |
| rewriter.replaceOp(op, replacements); | |
| return success(); | |
| } | |
| /// Lowering that supports all features the dealloc operation has to offer. It | |
| /// computes the base pointer of each memref (as an index), stores it in a | |
| /// new memref helper structure and passes it to the helper function generated | |
| /// in 'buildDeallocationHelperFunction'. The results are stored in two lists | |
| /// (represented as memrefs) of booleans passed as arguments. The first list | |
| /// stores whether the corresponding condition should be deallocated, the | |
| /// second list stores the ownership of the retained values which can be used | |
| /// to replace the result values of the `bufferization.dealloc` operation. | |
| /// | |
| /// Example: | |
| /// ``` | |
| /// %0:2 = bufferization.dealloc (%m0, %m1 : memref<2xf32>, memref<5xf32>) | |
| /// if (%cond0, %cond1) | |
| /// retain (%r0, %r1 : memref<1xf32>, memref<2xf32>) | |
| /// ``` | |
| /// lowers to (simplified): | |
| /// ``` | |
| /// %c0 = arith.constant 0 : index | |
| /// %c1 = arith.constant 1 : index | |
| /// %dealloc_base_pointer_list = memref.alloc() : memref<2xindex> | |
| /// %cond_list = memref.alloc() : memref<2xi1> | |
| /// %retain_base_pointer_list = memref.alloc() : memref<2xindex> | |
| /// %m0_base_pointer = memref.extract_aligned_pointer_as_index %m0 | |
| /// memref.store %m0_base_pointer, %dealloc_base_pointer_list[%c0] | |
| /// %m1_base_pointer = memref.extract_aligned_pointer_as_index %m1 | |
| /// memref.store %m1_base_pointer, %dealloc_base_pointer_list[%c1] | |
| /// memref.store %cond0, %cond_list[%c0] | |
| /// memref.store %cond1, %cond_list[%c1] | |
| /// %r0_base_pointer = memref.extract_aligned_pointer_as_index %r0 | |
| /// memref.store %r0_base_pointer, %retain_base_pointer_list[%c0] | |
| /// %r1_base_pointer = memref.extract_aligned_pointer_as_index %r1 | |
| /// memref.store %r1_base_pointer, %retain_base_pointer_list[%c1] | |
| /// %dyn_dealloc_base_pointer_list = memref.cast %dealloc_base_pointer_list : | |
| /// memref<2xindex> to memref<?xindex> | |
| /// %dyn_cond_list = memref.cast %cond_list : memref<2xi1> to memref<?xi1> | |
| /// %dyn_retain_base_pointer_list = memref.cast %retain_base_pointer_list : | |
| /// memref<2xindex> to memref<?xindex> | |
| /// %dealloc_cond_out = memref.alloc() : memref<2xi1> | |
| /// %ownership_out = memref.alloc() : memref<2xi1> | |
| /// %dyn_dealloc_cond_out = memref.cast %dealloc_cond_out : | |
| /// memref<2xi1> to memref<?xi1> | |
| /// %dyn_ownership_out = memref.cast %ownership_out : | |
| /// memref<2xi1> to memref<?xi1> | |
| /// call @dealloc_helper(%dyn_dealloc_base_pointer_list, | |
| /// %dyn_retain_base_pointer_list, | |
| /// %dyn_cond_list, | |
| /// %dyn_dealloc_cond_out, | |
| /// %dyn_ownership_out) : (...) | |
| /// %m0_dealloc_cond = memref.load %dyn_dealloc_cond_out[%c0] : memref<2xi1> | |
| /// scf.if %m0_dealloc_cond { | |
| /// memref.dealloc %m0 : memref<2xf32> | |
| /// } | |
| /// %m1_dealloc_cond = memref.load %dyn_dealloc_cond_out[%c1] : memref<2xi1> | |
| /// scf.if %m1_dealloc_cond { | |
| /// memref.dealloc %m1 : memref<5xf32> | |
| /// } | |
| /// %r0_ownership = memref.load %dyn_ownership_out[%c0] : memref<2xi1> | |
| /// %r1_ownership = memref.load %dyn_ownership_out[%c1] : memref<2xi1> | |
| /// memref.dealloc %dealloc_base_pointer_list : memref<2xindex> | |
| /// memref.dealloc %retain_base_pointer_list : memref<2xindex> | |
| /// memref.dealloc %cond_list : memref<2xi1> | |
| /// memref.dealloc %dealloc_cond_out : memref<2xi1> | |
| /// memref.dealloc %ownership_out : memref<2xi1> | |
| /// // replace %0#0 with %r0_ownership | |
| /// // replace %0#1 with %r1_ownership | |
| /// ``` | |
| LogicalResult rewriteGeneralCase(bufferization::DeallocOp op, | |
| OpAdaptor adaptor, | |
| ConversionPatternRewriter &rewriter) const { | |
| // Allocate two memrefs holding the base pointer indices of the list of | |
| // memrefs to be deallocated and the ones to be retained. These can then be | |
| // passed to the helper function and the for-loops can iterate over them. | |
| // Without storing them to memrefs, we could not use for-loops but only a | |
| // completely unrolled version of it, potentially leading to code-size | |
| // blow-up. | |
| Value toDeallocMemref = memref::AllocOp::create( | |
| rewriter, op.getLoc(), | |
| MemRefType::get({(int64_t)adaptor.getMemrefs().size()}, | |
| rewriter.getIndexType())); | |
| Value conditionMemref = memref::AllocOp::create( | |
| rewriter, op.getLoc(), | |
| MemRefType::get({(int64_t)adaptor.getConditions().size()}, | |
| rewriter.getI1Type())); | |
| Value toRetainMemref = memref::AllocOp::create( | |
| rewriter, op.getLoc(), | |
| MemRefType::get({(int64_t)adaptor.getRetained().size()}, | |
| rewriter.getIndexType())); | |
| auto getConstValue = [&](uint64_t value) -> Value { | |
| return arith::ConstantOp::create(rewriter, op.getLoc(), | |
| rewriter.getIndexAttr(value)); | |
| }; | |
| // Extract the base pointers of the memrefs as indices to check for aliasing | |
| // at runtime. | |
| for (auto [i, toDealloc] : llvm::enumerate(adaptor.getMemrefs())) { | |
| Value memrefAsIdx = memref::ExtractAlignedPointerAsIndexOp::create( | |
| rewriter, op.getLoc(), toDealloc); | |
| memref::StoreOp::create(rewriter, op.getLoc(), memrefAsIdx, | |
| toDeallocMemref, getConstValue(i)); | |
| } | |
| for (auto [i, cond] : llvm::enumerate(adaptor.getConditions())) | |
| memref::StoreOp::create(rewriter, op.getLoc(), cond, conditionMemref, | |
| getConstValue(i)); | |
| for (auto [i, toRetain] : llvm::enumerate(adaptor.getRetained())) { | |
| Value memrefAsIdx = memref::ExtractAlignedPointerAsIndexOp::create( | |
| rewriter, op.getLoc(), toRetain); | |
| memref::StoreOp::create(rewriter, op.getLoc(), memrefAsIdx, | |
| toRetainMemref, getConstValue(i)); | |
| } | |
| // Cast the allocated memrefs to dynamic shape because we want only one | |
| // helper function no matter how many operands the bufferization.dealloc | |
| // has. | |
| Value castedDeallocMemref = memref::CastOp::create( | |
| rewriter, op->getLoc(), | |
| MemRefType::get({ShapedType::kDynamic}, rewriter.getIndexType()), | |
| toDeallocMemref); | |
| Value castedCondsMemref = memref::CastOp::create( | |
| rewriter, op->getLoc(), | |
| MemRefType::get({ShapedType::kDynamic}, rewriter.getI1Type()), | |
| conditionMemref); | |
| Value castedRetainMemref = memref::CastOp::create( | |
| rewriter, op->getLoc(), | |
| MemRefType::get({ShapedType::kDynamic}, rewriter.getIndexType()), | |
| toRetainMemref); | |
| Value deallocCondsMemref = memref::AllocOp::create( | |
| rewriter, op.getLoc(), | |
| MemRefType::get({(int64_t)adaptor.getMemrefs().size()}, | |
| rewriter.getI1Type())); | |
| Value retainCondsMemref = memref::AllocOp::create( | |
| rewriter, op.getLoc(), | |
| MemRefType::get({(int64_t)adaptor.getRetained().size()}, | |
| rewriter.getI1Type())); | |
| Value castedDeallocCondsMemref = memref::CastOp::create( | |
| rewriter, op->getLoc(), | |
| MemRefType::get({ShapedType::kDynamic}, rewriter.getI1Type()), | |
| deallocCondsMemref); | |
| Value castedRetainCondsMemref = memref::CastOp::create( | |
| rewriter, op->getLoc(), | |
| MemRefType::get({ShapedType::kDynamic}, rewriter.getI1Type()), | |
| retainCondsMemref); | |
| Operation *symtableOp = op->getParentWithTrait<OpTrait::SymbolTable>(); | |
| func::CallOp::create( | |
| rewriter, op.getLoc(), deallocHelperFuncMap.lookup(symtableOp), | |
| SmallVector<Value>{castedDeallocMemref, castedRetainMemref, | |
| castedCondsMemref, castedDeallocCondsMemref, | |
| castedRetainCondsMemref}); | |
| for (unsigned i = 0, e = adaptor.getMemrefs().size(); i < e; ++i) { | |
| Value idxValue = getConstValue(i); | |
| Value shouldDealloc = memref::LoadOp::create( | |
| rewriter, op.getLoc(), deallocCondsMemref, idxValue); | |
| scf::IfOp::create(rewriter, op.getLoc(), shouldDealloc, | |
| [&](OpBuilder &builder, Location loc) { | |
| memref::DeallocOp::create(builder, loc, | |
| adaptor.getMemrefs()[i]); | |
| scf::YieldOp::create(builder, loc); | |
| }); | |
| } | |
| SmallVector<Value> replacements; | |
| for (unsigned i = 0, e = adaptor.getRetained().size(); i < e; ++i) { | |
| Value idxValue = getConstValue(i); | |
| Value ownership = memref::LoadOp::create(rewriter, op.getLoc(), | |
| retainCondsMemref, idxValue); | |
| replacements.push_back(ownership); | |
| } | |
| // Deallocate above allocated memrefs again to avoid memory leaks. | |
| // Deallocation will not be run on code after this stage. | |
| memref::DeallocOp::create(rewriter, op.getLoc(), toDeallocMemref); | |
| memref::DeallocOp::create(rewriter, op.getLoc(), toRetainMemref); | |
| memref::DeallocOp::create(rewriter, op.getLoc(), conditionMemref); | |
| memref::DeallocOp::create(rewriter, op.getLoc(), deallocCondsMemref); | |
| memref::DeallocOp::create(rewriter, op.getLoc(), retainCondsMemref); | |
| rewriter.replaceOp(op, replacements); | |
| return success(); | |
| } | |
| public: | |
| DeallocOpConversion( | |
| MLIRContext *context, | |
| const bufferization::DeallocHelperMap &deallocHelperFuncMap) | |
| : OpConversionPattern<bufferization::DeallocOp>(context), | |
| deallocHelperFuncMap(deallocHelperFuncMap) {} | |
| LogicalResult | |
| matchAndRewrite(bufferization::DeallocOp op, OpAdaptor adaptor, | |
| ConversionPatternRewriter &rewriter) const override { | |
| // Lower the trivial case. | |
| if (adaptor.getMemrefs().empty()) { | |
| Value falseVal = arith::ConstantOp::create(rewriter, op.getLoc(), | |
| rewriter.getBoolAttr(false)); | |
| rewriter.replaceOp( | |
| op, SmallVector<Value>(adaptor.getRetained().size(), falseVal)); | |
| return success(); | |
| } | |
| if (adaptor.getMemrefs().size() == 1 && adaptor.getRetained().empty()) | |
| return rewriteOneMemrefNoRetainCase(op, adaptor, rewriter); | |
| if (adaptor.getMemrefs().size() == 1) | |
| return rewriteOneMemrefMultipleRetainCase(op, adaptor, rewriter); | |
| Operation *symtableOp = op->getParentWithTrait<OpTrait::SymbolTable>(); | |
| if (!deallocHelperFuncMap.contains(symtableOp)) | |
| return op->emitError( | |
| "library function required for generic lowering, but cannot be " | |
| "automatically inserted when operating on functions"); | |
| return rewriteGeneralCase(op, adaptor, rewriter); | |
| } | |
| private: | |
| const bufferization::DeallocHelperMap &deallocHelperFuncMap; | |
| }; | |
| } // namespace | |
| namespace { | |
| struct LowerDeallocationsPass | |
| : public bufferization::impl::LowerDeallocationsPassBase< | |
| LowerDeallocationsPass> { | |
| void runOnOperation() override { | |
| if (!isa<ModuleOp, FunctionOpInterface>(getOperation())) { | |
| emitError(getOperation()->getLoc(), | |
| "root operation must be a builtin.module or a function"); | |
| signalPassFailure(); | |
| return; | |
| } | |
| bufferization::DeallocHelperMap deallocHelperFuncMap; | |
| if (auto module = dyn_cast<ModuleOp>(getOperation())) { | |
| OpBuilder builder = OpBuilder::atBlockBegin(module.getBody()); | |
| // Build dealloc helper function if there are deallocs. | |
| getOperation()->walk([&](bufferization::DeallocOp deallocOp) { | |
| Operation *symtableOp = | |
| deallocOp->getParentWithTrait<OpTrait::SymbolTable>(); | |
| if (deallocOp.getMemrefs().size() > 1 && | |
| !deallocHelperFuncMap.contains(symtableOp)) { | |
| SymbolTable symbolTable(symtableOp); | |
| func::FuncOp helperFuncOp = | |
| bufferization::buildDeallocationLibraryFunction( | |
| builder, getOperation()->getLoc(), symbolTable); | |
| deallocHelperFuncMap[symtableOp] = helperFuncOp; | |
| } | |
| }); | |
| } | |
| RewritePatternSet patterns(&getContext()); | |
| bufferization::populateBufferizationDeallocLoweringPattern( | |
| patterns, deallocHelperFuncMap); | |
| ConversionTarget target(getContext()); | |
| target.addLegalDialect<memref::MemRefDialect, arith::ArithDialect, | |
| scf::SCFDialect, func::FuncDialect>(); | |
| target.addIllegalOp<bufferization::DeallocOp>(); | |
| if (failed(applyPartialConversion(getOperation(), target, | |
| std::move(patterns)))) | |
| signalPassFailure(); | |
| } | |
| }; | |
| } // namespace | |
| func::FuncOp mlir::bufferization::buildDeallocationLibraryFunction( | |
| OpBuilder &builder, Location loc, SymbolTable &symbolTable) { | |
| Type indexMemrefType = | |
| MemRefType::get({ShapedType::kDynamic}, builder.getIndexType()); | |
| Type boolMemrefType = | |
| MemRefType::get({ShapedType::kDynamic}, builder.getI1Type()); | |
| SmallVector<Type> argTypes{indexMemrefType, indexMemrefType, boolMemrefType, | |
| boolMemrefType, boolMemrefType}; | |
| builder.clearInsertionPoint(); | |
| // Generate the func operation itself. | |
| auto helperFuncOp = func::FuncOp::create( | |
| loc, "dealloc_helper", builder.getFunctionType(argTypes, {})); | |
| helperFuncOp.setVisibility(SymbolTable::Visibility::Private); | |
| symbolTable.insert(helperFuncOp); | |
| auto &block = helperFuncOp.getFunctionBody().emplaceBlock(); | |
| block.addArguments(argTypes, SmallVector<Location>(argTypes.size(), loc)); | |
| builder.setInsertionPointToStart(&block); | |
| Value toDeallocMemref = helperFuncOp.getArguments()[0]; | |
| Value toRetainMemref = helperFuncOp.getArguments()[1]; | |
| Value conditionMemref = helperFuncOp.getArguments()[2]; | |
| Value deallocCondsMemref = helperFuncOp.getArguments()[3]; | |
| Value retainCondsMemref = helperFuncOp.getArguments()[4]; | |
| // Insert some prerequisites. | |
| Value c0 = arith::ConstantOp::create(builder, loc, builder.getIndexAttr(0)); | |
| Value c1 = arith::ConstantOp::create(builder, loc, builder.getIndexAttr(1)); | |
| Value trueValue = | |
| arith::ConstantOp::create(builder, loc, builder.getBoolAttr(true)); | |
| Value falseValue = | |
| arith::ConstantOp::create(builder, loc, builder.getBoolAttr(false)); | |
| Value toDeallocSize = | |
| memref::DimOp::create(builder, loc, toDeallocMemref, c0); | |
| Value toRetainSize = memref::DimOp::create(builder, loc, toRetainMemref, c0); | |
| scf::ForOp::create( | |
| builder, loc, c0, toRetainSize, c1, ValueRange(), | |
| [&](OpBuilder &builder, Location loc, Value i, ValueRange iterArgs) { | |
| memref::StoreOp::create(builder, loc, falseValue, retainCondsMemref, i); | |
| scf::YieldOp::create(builder, loc); | |
| }); | |
| scf::ForOp::create( | |
| builder, loc, c0, toDeallocSize, c1, ValueRange(), | |
| [&](OpBuilder &builder, Location loc, Value outerIter, | |
| ValueRange iterArgs) { | |
| Value toDealloc = | |
| memref::LoadOp::create(builder, loc, toDeallocMemref, outerIter); | |
| Value cond = | |
| memref::LoadOp::create(builder, loc, conditionMemref, outerIter); | |
| // Build the first for loop that computes aliasing with retained | |
| // memrefs. | |
| Value | |
| noRetainAlias = | |
| scf::ForOp::create( | |
| builder, loc, c0, toRetainSize, c1, trueValue, | |
| [&](OpBuilder &builder, Location loc, Value i, | |
| ValueRange iterArgs) { | |
| Value retainValue = memref::LoadOp::create( | |
| builder, loc, toRetainMemref, i); | |
| Value doesAlias = arith::CmpIOp::create( | |
| builder, loc, arith::CmpIPredicate::eq, retainValue, | |
| toDealloc); | |
| scf::IfOp::create( | |
| builder, loc, doesAlias, | |
| [&](OpBuilder &builder, Location loc) { | |
| Value retainCondValue = memref::LoadOp::create( | |
| builder, loc, retainCondsMemref, i); | |
| Value aggregatedRetainCond = arith::OrIOp::create( | |
| builder, loc, retainCondValue, cond); | |
| memref::StoreOp::create(builder, loc, | |
| aggregatedRetainCond, | |
| retainCondsMemref, i); | |
| scf::YieldOp::create(builder, loc); | |
| }); | |
| Value doesntAlias = arith::CmpIOp::create( | |
| builder, loc, arith::CmpIPredicate::ne, retainValue, | |
| toDealloc); | |
| Value yieldValue = arith::AndIOp::create( | |
| builder, loc, iterArgs[0], doesntAlias); | |
| scf::YieldOp::create(builder, loc, yieldValue); | |
| }) | |
| .getResult(0); | |
| // Build the second for loop that adds aliasing with previously | |
| // deallocated memrefs. | |
| Value | |
| noAlias = | |
| scf::ForOp::create( | |
| builder, loc, c0, outerIter, c1, noRetainAlias, | |
| [&](OpBuilder &builder, Location loc, Value i, | |
| ValueRange iterArgs) { | |
| Value prevDeallocValue = memref::LoadOp::create( | |
| builder, loc, toDeallocMemref, i); | |
| Value doesntAlias = arith::CmpIOp::create( | |
| builder, loc, arith::CmpIPredicate::ne, | |
| prevDeallocValue, toDealloc); | |
| Value yieldValue = arith::AndIOp::create( | |
| builder, loc, iterArgs[0], doesntAlias); | |
| scf::YieldOp::create(builder, loc, yieldValue); | |
| }) | |
| .getResult(0); | |
| Value shouldDealoc = arith::AndIOp::create(builder, loc, noAlias, cond); | |
| memref::StoreOp::create(builder, loc, shouldDealoc, deallocCondsMemref, | |
| outerIter); | |
| scf::YieldOp::create(builder, loc); | |
| }); | |
| func::ReturnOp::create(builder, loc); | |
| return helperFuncOp; | |
| } | |
| void mlir::bufferization::populateBufferizationDeallocLoweringPattern( | |
| RewritePatternSet &patterns, | |
| const bufferization::DeallocHelperMap &deallocHelperFuncMap) { | |
| patterns.add<DeallocOpConversion>(patterns.getContext(), | |
| deallocHelperFuncMap); | |
| } |