[mlir][bufferization] Add lowering of bufferization.dealloc to memref.dealloc

[maerhart]

Summary

Adds a generic lowering that suppors all cases of bufferization.dealloc
and one specialized, more efficient lowering for the simple case. Using
a helper function with for loops in the general case enables
O(|num_dealloc_memrefs|+|num_retain_memrefs|) size of the lowered code.

Depends on D155467

[maerhart]

• Update to new assembly format
• Fix mixed signed/unsigned comparison

[maerhart]

mlir/test/Conversion/BufferizationToMemRef/bufferization-to-memref.mlir

127

This should be TO_RETAIN_MR instead of TO_DEALLOC_MR

[maerhart]

mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp

158

Need to pass toRetainMemref here

[springerm]

mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp

88

Can we put an example here.

93–105

You don't need to cache these, we generate duplicate constants all the time and they get DCE'd. (memref.store should take an OpFoldResult instead of Value, but that's a different story...)

128

"base"?

131–132

grammar

220

Let's add a comment here that describes what the function is doing and interface (e.g., memrefs are being dellocated by the function).

226–227

You can insert the function via the module's symbol table, which will unique the name automatically when necessary. To do this, you have to clear the insertion point before creating the FuncOp, then use SymbolTable::insert.
E.g.: https://github.com/matthias-springer/iree/commit/ea9032e08fa2dbba7b93763964ea9ada8f3d8e3e

237–238

These are not needed. (?) You can call memref.dim in the helper function, then the function signature will be a bit simpler.

403

Add comment before this line: Build dealloc helper function if there are deallocs.

mlir/test/Conversion/BufferizationToMemRef/bufferization-to-memref.mlir

88

CHECK

101

We should capture all SSA names, e.g. %[[arg0:.*]]. You need a CHECK-SAME in a new line to do that.

127

TO_RETAIN_MR?

[maerhart]

Address review comments

mlir/include/mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h

 //===- BufferizationToMemRef.h - Bufferization to MemRef conversion -------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_BUFFERIZATIONTOMEMREF_BUFFERIZATIONTOMEMREF_H
 #define MLIR_CONVERSION_BUFFERIZATIONTOMEMREF_BUFFERIZATIONTOMEMREF_H
 
+#include "mlir/Pass/Pass.h"
 #include <memory>
 
 namespace mlir {
-class Pass;
-class RewritePatternSet;
+class ModuleOp;
 
 #define GEN_PASS_DECL_CONVERTBUFFERIZATIONTOMEMREF
 #include "mlir/Conversion/Passes.h.inc"
 
-/// Collect a set of patterns to convert memory-related operations from the
-/// Bufferization dialect to the MemRef dialect.
-void populateBufferizationToMemRefConversionPatterns(
-    RewritePatternSet &patterns);
-
-std::unique_ptr<Pass> createBufferizationToMemRefPass();
+std::unique_ptr<OperationPass<ModuleOp>> createBufferizationToMemRefPass();
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_BUFFERIZATIONTOMEMREF_BUFFERIZATIONTOMEMREF_H

mlir/include/mlir/Conversion/Passes.td

Option<"useOpaquePointers", "use-opaque-pointers", "bool",
            "instead of typed pointers">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // BufferizationToMemRef
 //===----------------------------------------------------------------------===//
 
-def ConvertBufferizationToMemRef : Pass<"convert-bufferization-to-memref"> {
+def ConvertBufferizationToMemRef : Pass<"convert-bufferization-to-memref",
+                                        "mlir::ModuleOp"> {
   let summary = "Convert operations from the Bufferization dialect to the "
                 "MemRef dialect";
   let description = [{
 
     This pass converts bufferization operations into memref operations.
 
     In the current state, this pass only transforms a `bufferization.clone`
     operation into `memref.alloc` and `memref.copy` operations. This conversion
let description = [{
     e.g.).
 
     Note that this pass only transforms the operation without any further
     analyses. This pass does not consider any memory analysis or optimization
     and hence does not resolve any memory leaks.
 
   }];
   let constructor = "mlir::createBufferizationToMemRefPass()";
-  let dependentDialects = ["arith::ArithDialect", "memref::MemRefDialect"];
+  let dependentDialects = [
+    "arith::ArithDialect", "memref::MemRefDialect", "scf::SCFDialect",
+    "func::FuncDialect"
+  ];
 }
 
 //===----------------------------------------------------------------------===//
 // ComplexToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertComplexToLLVMPass : Pass<"convert-complex-to-llvm"> {
   let summary = "Convert Complex dialect to LLVM dialect";

mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp

 // dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.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/Pass/Pass.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_CONVERTBUFFERIZATIONTOMEMREF
 #include "mlir/Conversion/Passes.h.inc"
matchAndRewrite(bufferization::CloneOp op, OpAdaptor adaptor,
     // Cast the allocation to the specified type if needed.
     if (memrefType != allocType)
       alloc = rewriter.create<memref::CastOp>(op->getLoc(), memrefType, alloc);
     rewriter.replaceOp(op, alloc);
     rewriter.create<memref::CopyOp>(loc, op.getInput(), alloc);
     return success();
   }
 };
-} // namespace
 
-void mlir::populateBufferizationToMemRefConversionPatterns(
-    RewritePatternSet &patterns) {
-  patterns.add<CloneOpConversion>(patterns.getContext());
+/// 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 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:
+  /// ```
+  /// %0 = bufferization.dealloc (%arg0 : memref<2xf32>) if (%arg1)
+  /// ```
+  /// is lowered to
+  /// ```
+  /// scf.if %arg1 {
+  ///   memref.dealloc %arg0 : memref<2xf32>
+  /// }
+  /// %0 = arith.constant false
+  /// ```
+  LogicalResult
+  rewriteOneMemrefNoRetainCase(bufferization::DeallocOp op, OpAdaptor adaptor,
+                               ConversionPatternRewriter &rewriter) const {
+    rewriter.create<scf::IfOp>(op.getLoc(), adaptor.getConditions()[0],
+                               [&](OpBuilder &builder, Location loc) {
+                                 builder.create<memref::DeallocOp>(
+                                     loc, adaptor.getMemrefs()[0]);
+                                 builder.create<scf::YieldOp>(loc);
+                               });
+    rewriter.replaceOpWithNewOp<arith::ConstantOp>(op,
+                                                   rewriter.getBoolAttr(false));
+    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 them in a
+  /// new memref and passes it to the helper function generated in
+  /// 'buildDeallocationHelperFunction'. The two return values are used as
+  /// condition for the scf if operation containing the memref deallocate and as
+  /// replacement for the original bufferization dealloc respectively.
+  ///
+  /// Example:
+  /// ```
+  /// %0:2 = bufferization.dealloc (%arg0, %arg1 : memref<2xf32>, memref<5xf32>)
+  ///                           if (%arg3, %arg4) retain (%arg2 : memref<1xf32>)
+  /// ```
+  /// lowers to (simplified):
+  /// ```
+  /// %c0 = arith.constant 0 : index
+  /// %c1 = arith.constant 1 : index
+  /// %alloc = memref.alloc() : memref<2xindex>
+  /// %alloc_0 = memref.alloc() : memref<1xindex>
+  /// %intptr = memref.extract_aligned_pointer_as_index %arg0
+  /// memref.store %intptr, %alloc[%c0] : memref<2xindex>
+  /// %intptr_1 = memref.extract_aligned_pointer_as_index %arg1
+  /// memref.store %intptr_1, %alloc[%c1] : memref<2xindex>
+  /// %intptr_2 = memref.extract_aligned_pointer_as_index %arg2
+  /// memref.store %intptr_2, %alloc_0[%c0] : memref<1xindex>
+  /// %cast = memref.cast %alloc : memref<2xindex> to memref<?xindex>
+  /// %cast_4 = memref.cast %alloc_0 : memref<1xindex> to memref<?xindex>
+  /// %0:2 = call @dealloc_helper(%cast, %cast_4, %c0)
+  /// %1 = arith.andi %0#0, %arg3 : i1
+  /// %2 = arith.andi %0#1, %arg3 : i1
+  /// scf.if %1 {
+  ///   memref.dealloc %arg0 : memref<2xf32>
+  /// }
+  /// %3:2 = call @dealloc_helper(%cast, %cast_4, %c1)
+  /// %4 = arith.andi %3#0, %arg4 : i1
+  /// %5 = arith.andi %3#1, %arg4 : i1
+  /// scf.if %4 {
+  ///   memref.dealloc %arg1 : memref<5xf32>
+  /// }
+  /// memref.dealloc %alloc : memref<2xindex>
+  /// memref.dealloc %alloc_0 : memref<1xindex>
+  /// // replace %0#0 with %2
+  /// // replace %0#1 with %5
+  /// ```
+  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 = rewriter.create<memref::AllocOp>(
+        op.getLoc(), MemRefType::get({(int64_t)adaptor.getMemrefs().size()},
+                                     rewriter.getIndexType()));
+    Value toRetainMemref = rewriter.create<memref::AllocOp>(
+        op.getLoc(), MemRefType::get({(int64_t)adaptor.getRetained().size()},
+                                     rewriter.getIndexType()));
+
+    auto getConstValue = [&](uint64_t value) -> Value {
+      return rewriter.create<arith::ConstantOp>(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 =
+          rewriter.create<memref::ExtractAlignedPointerAsIndexOp>(op.getLoc(),
+                                                                  toDealloc);
+      rewriter.create<memref::StoreOp>(op.getLoc(), memrefAsIdx,
+                                       toDeallocMemref, getConstValue(i));
     }
+    for (auto [i, toRetain] : llvm::enumerate(adaptor.getRetained())) {
+      Value memrefAsIdx =
+          rewriter.create<memref::ExtractAlignedPointerAsIndexOp>(op.getLoc(),
+                                                                  toRetain);
+      rewriter.create<memref::StoreOp>(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 = rewriter.create<memref::CastOp>(
+        op->getLoc(),
+        MemRefType::get({ShapedType::kDynamic}, rewriter.getIndexType()),
+        toDeallocMemref);
+    Value castedRetainMemref = rewriter.create<memref::CastOp>(
+        op->getLoc(),
+        MemRefType::get({ShapedType::kDynamic}, rewriter.getIndexType()),
+        toRetainMemref);
+
+    SmallVector<Value> replacements;
+    for (unsigned i = 0, e = adaptor.getMemrefs().size(); i < e; ++i) {
+      auto callOp = rewriter.create<func::CallOp>(
+          op.getLoc(), deallocHelperFunc,
+          SmallVector<Value>{castedDeallocMemref, castedRetainMemref,
+                             getConstValue(i)});
+      Value shouldDealloc = rewriter.create<arith::AndIOp>(
+          op.getLoc(), callOp.getResult(0), adaptor.getConditions()[i]);
+      Value ownership = rewriter.create<arith::AndIOp>(
+          op.getLoc(), callOp.getResult(1), adaptor.getConditions()[i]);
+      replacements.push_back(ownership);
+      rewriter.create<scf::IfOp>(
+          op.getLoc(), shouldDealloc, [&](OpBuilder &builder, Location loc) {
+            builder.create<memref::DeallocOp>(loc, adaptor.getMemrefs()[i]);
+            builder.create<scf::YieldOp>(loc);
+          });
+    }
+
+    // Deallocate above allocated memrefs again to avoid memory leaks.
+    // Deallocation will not be run on code after this stage.
+    rewriter.create<memref::DeallocOp>(op.getLoc(), toDeallocMemref);
+    rewriter.create<memref::DeallocOp>(op.getLoc(), toRetainMemref);
+
+    rewriter.replaceOp(op, replacements);
+    return success();
+  }
+
+public:
+  DeallocOpConversion(MLIRContext *context, func::FuncOp deallocHelperFunc)
+      : OpConversionPattern<bufferization::DeallocOp>(context),
+        deallocHelperFunc(deallocHelperFunc) {}
+
+  LogicalResult
+  matchAndRewrite(bufferization::DeallocOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Lower the trivial case.
+    if (adaptor.getMemrefs().empty())
+      return rewriter.eraseOp(op), success();
+
+    if (adaptor.getMemrefs().size() == 1 && adaptor.getRetained().empty())
+      return rewriteOneMemrefNoRetainCase(op, adaptor, rewriter);
+
+    return rewriteGeneralCase(op, adaptor, rewriter);
+  }
+
+  /// Build a helper function per compilation unit that can be called at
+  /// bufferization dealloc sites to determine aliasing and ownership.
+  ///
+  /// The generated function takes two memrefs of indices and one index value as
+  /// arguments and returns two boolean values:
+  ///   * The first memref argument A should contain the result of the
+  ///   extract_aligned_pointer_as_index operation applied to the memrefs to be
+  ///   deallocated
+  ///   * The second memref argument B should contain the result of the
+  ///   extract_aligned_pointer_as_index operation applied to the memrefs to be
+  ///   retained
+  ///   * The index argument I represents the currently processed index of
+  ///   memref A and is needed because aliasing with all previously deallocated
+  ///   memrefs has to be checked to avoid double deallocation
+  ///   * The first result indicates whether the memref at position I should be
+  ///   deallocated
+  ///   * The second result provides the updated ownership value corresponding
+  ///   the the memref at position I
+  ///
+  /// This helper function is supposed to be called for each element in the list
+  /// of memrefs to be deallocated to determine the deallocation need and new
+  /// ownership indicator, but does not perform the deallocation itself.
+  ///
+  /// The first scf for loop in the body computes whether the memref at index I
+  /// aliases with any memref in the list of retained memrefs.
+  /// The second loop additionally checks whether one of the previously
+  /// deallocated memrefs aliases with the currently processed one.
+  ///
+  /// Generated code:
+  /// ```
+  /// func.func @dealloc_helper(%arg0: memref<?xindex>,
+  ///                           %arg1: memref<?xindex>,
+  ///                           %arg2: index) -> (i1, i1) {
+  ///   %c0 = arith.constant 0 : index
+  ///   %c1 = arith.constant 1 : index
+  ///   %true = arith.constant true
+  ///   %dim = memref.dim %arg1, %c0 : memref<?xindex>
+  ///   %0 = memref.load %arg0[%arg2] : memref<?xindex>
+  ///   %1 = scf.for %i = %c0 to %dim step %c1 iter_args(%arg4 = %true) -> (i1){
+  ///     %4 = memref.load %arg1[%i] : memref<?xindex>
+  ///     %5 = arith.cmpi ne, %4, %0 : index
+  ///     %6 = arith.andi %arg4, %5 : i1
+  ///     scf.yield %6 : i1
+  ///   }
+  ///   %2 = scf.for %i = %c0 to %arg2 step %c1 iter_args(%arg4 = %1) -> (i1) {
+  ///     %4 = memref.load %arg0[%i] : memref<?xindex>
+  ///     %5 = arith.cmpi ne, %4, %0 : index
+  ///     %6 = arith.andi %arg4, %5 : i1
+  ///     scf.yield %6 : i1
+  ///   }
+  ///   %3 = arith.xori %1, %true : i1
+  ///   return %2, %3 : i1, i1
+  /// }
+  /// ```
+  static func::FuncOp
+  buildDeallocationHelperFunction(OpBuilder &builder, Location loc,
+                                  SymbolTable &symbolTable) {
+    Type idxType = builder.getIndexType();
+    Type memrefArgType = MemRefType::get({ShapedType::kDynamic}, idxType);
+    SmallVector<Type> argTypes{memrefArgType, memrefArgType, idxType};
+    builder.clearInsertionPoint();
+
+    // Generate the func operation itself.
+    auto helperFuncOp = func::FuncOp::create(
+        loc, "dealloc_helper",
+        builder.getFunctionType(argTypes,
+                                {builder.getI1Type(), builder.getI1Type()}));
+    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 idxArg = helperFuncOp.getArguments()[2];
+
+    // Insert some prerequisites.
+    Value c0 = builder.create<arith::ConstantOp>(loc, builder.getIndexAttr(0));
+    Value c1 = builder.create<arith::ConstantOp>(loc, builder.getIndexAttr(1));
+    Value trueValue =
+        builder.create<arith::ConstantOp>(loc, builder.getBoolAttr(true));
+    Value toRetainSize = builder.create<memref::DimOp>(loc, toRetainMemref, c0);
+    Value toDealloc =
+        builder.create<memref::LoadOp>(loc, toDeallocMemref, idxArg);
+
+    // Build the first for loop that computes aliasing with retained memrefs.
+    Value noRetainAlias =
+        builder
+            .create<scf::ForOp>(
+                loc, c0, toRetainSize, c1, trueValue,
+                [&](OpBuilder &builder, Location loc, Value i,
+                    ValueRange iterArgs) {
+                  Value retainValue =
+                      builder.create<memref::LoadOp>(loc, toRetainMemref, i);
+                  Value doesntAlias = builder.create<arith::CmpIOp>(
+                      loc, arith::CmpIPredicate::ne, retainValue, toDealloc);
+                  Value yieldValue = builder.create<arith::AndIOp>(
+                      loc, iterArgs[0], doesntAlias);
+                  builder.create<scf::YieldOp>(loc, yieldValue);
+                })
+            .getResult(0);
+
+    // Build the second for loop that adds aliasing with previously deallocated
+    // memrefs.
+    Value noAlias =
+        builder
+            .create<scf::ForOp>(
+                loc, c0, idxArg, c1, noRetainAlias,
+                [&](OpBuilder &builder, Location loc, Value i,
+                    ValueRange iterArgs) {
+                  Value prevDeallocValue =
+                      builder.create<memref::LoadOp>(loc, toDeallocMemref, i);
+                  Value doesntAlias = builder.create<arith::CmpIOp>(
+                      loc, arith::CmpIPredicate::ne, prevDeallocValue,
+                      toDealloc);
+                  Value yieldValue = builder.create<arith::AndIOp>(
+                      loc, iterArgs[0], doesntAlias);
+                  builder.create<scf::YieldOp>(loc, yieldValue);
+                })
+            .getResult(0);
+
+    Value ownership =
+        builder.create<arith::XOrIOp>(loc, noRetainAlias, trueValue);
+    builder.create<func::ReturnOp>(loc, SmallVector<Value>{noAlias, ownership});
+
+    return helperFuncOp;
+  }
+
+private:
+  func::FuncOp deallocHelperFunc;
+};
+} // namespace
 
 namespace {
 struct BufferizationToMemRefPass
     : public impl::ConvertBufferizationToMemRefBase<BufferizationToMemRefPass> {
   BufferizationToMemRefPass() = default;
 
   void runOnOperation() override {
+    ModuleOp module = cast<ModuleOp>(getOperation());
+    OpBuilder builder =
+        OpBuilder::atBlockBegin(&module.getBodyRegion().front());
+    SymbolTable symbolTable(module);
+
+    // Build dealloc helper function if there are deallocs.
+    func::FuncOp helperFuncOp;
+    getOperation()->walk([&](bufferization::DeallocOp deallocOp) {
+      if (deallocOp.getMemrefs().size() > 1 ||
+          !deallocOp.getRetained().empty()) {
+        helperFuncOp = DeallocOpConversion::buildDeallocationHelperFunction(
+            builder, getOperation()->getLoc(), symbolTable);
+        return WalkResult::interrupt();
+      }
+      return WalkResult::advance();
+    });
+
     RewritePatternSet patterns(&getContext());
-    populateBufferizationToMemRefConversionPatterns(patterns);
+    patterns.add<CloneOpConversion>(patterns.getContext());
+    patterns.add<DeallocOpConversion>(patterns.getContext(), helperFuncOp);
 
     ConversionTarget target(getContext());
-    target.addLegalDialect<memref::MemRefDialect>();
-    target.addLegalOp<arith::ConstantOp>();
+    target.addLegalDialect<memref::MemRefDialect, arith::ArithDialect,
+                           scf::SCFDialect, func::FuncDialect>();
     target.addIllegalDialect<bufferization::BufferizationDialect>();
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
-std::unique_ptr<Pass> mlir::createBufferizationToMemRefPass() {
+std::unique_ptr<OperationPass<ModuleOp>>
+mlir::createBufferizationToMemRefPass() {
   return std::make_unique<BufferizationToMemRefPass>();
 }

mlir/lib/Conversion/BufferizationToMemRef/CMakeLists.txt

 add_mlir_conversion_library(MLIRBufferizationToMemRef
   BufferizationToMemRef.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/BufferizationToMemRef
 
   DEPENDS
   MLIRConversionPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRBufferizationDialect
+  MLIRSCFDialect
+  MLIRFuncDialect
+  MLIRArithDialect
+  MLIRMemRefDialect
   MLIRPass
   MLIRTransforms
   )

mlir/test/Conversion/BufferizationToMemRef/bufferization-to-memref.mlir

 
 func.func @conversion_with_invalid_layout_map(%arg0 : memref<?xf32, strided<[10], offset: ?>>)
     -> memref<?xf32, strided<[10], offset: ?>> {
 // expected-error@+1 {{failed to legalize operation 'bufferization.clone' that was explicitly marked illegal}}
   %1 = bufferization.clone %arg0 : memref<?xf32, strided<[10], offset: ?>> to memref<?xf32, strided<[10], offset: ?>>
   memref.dealloc %arg0 : memref<?xf32, strided<[10], offset: ?>>
   return %1 : memref<?xf32, strided<[10], offset: ?>>
 }
+// -----
+
+// CHECK-LABEL: func @conversion_dealloc_empty
+func.func @conversion_dealloc_empty() {
+  // CHECK-NEXT: return
+  bufferization.dealloc
+  return
+}
+
+// -----
+
+// CHECK-NOT: func @deallocHelper
+// CHECK-LABEL: func @conversion_dealloc_simple
+// CHECK-SAME: [[ARG0:%.+]]: memref<2xf32>
+// CHECK-SAME: [[ARG1:%.+]]: i1
+func.func @conversion_dealloc_simple(%arg0: memref<2xf32>, %arg1: i1) -> i1 {
+  %0 = bufferization.dealloc (%arg0 : memref<2xf32>) if (%arg1)
+  return %0 : i1
+}
+
+//      CHECk: scf.if [[ARG1]] {
+// CHECk-NEXT:   memref.dealloc [[ARG0]] : memref<2xf32>
+// CHECk-NEXT: }
+// CHECk-NEXT: [[FALSE:%.+]] = arith.constant false
+// CHECk-NEXT: return [[FALSE]] : i1
+
+// -----
+
+func.func @conversion_dealloc_multiple_memrefs_and_retained(%arg0: memref<2xf32>, %arg1: memref<5xf32>, %arg2: memref<1xf32>, %arg3: i1, %arg4: i1) -> (i1, i1) {
+  %0:2 = bufferization.dealloc (%arg0, %arg1 : memref<2xf32>, memref<5xf32>) if (%arg3, %arg4) retain (%arg2 : memref<1xf32>)
+  return %0#0, %0#1 : i1, i1
+}
+
+// CHECK-LABEL: func @conversion_dealloc_multiple_memrefs_and_retained
+// CHECK-SAME: [[ARG0:%.+]]: memref<2xf32>,
+// CHECK-SAME: [[ARG1:%.+]]: memref<5xf32>,
+// CHECK-SAME: [[ARG2:%.+]]: memref<1xf32>,
+// CHECK-SAME: [[ARG3:%.+]]: i1,
+// CHECK-SAME: [[ARG4:%.+]]: i1
+//      CHECK: [[TO_DEALLOC_MR:%.+]] = memref.alloc() : memref<2xindex>
+//      CHECK: [[TO_RETAIN_MR:%.+]] = memref.alloc() : memref<1xindex>
+//  CHECK-DAG: [[V0:%.+]] = memref.extract_aligned_pointer_as_index [[ARG0]]
+//  CHECK-DAG: [[C0:%.+]] = arith.constant 0 : index
+//  CHECK-DAG: memref.store [[V0]], [[TO_DEALLOC_MR]][[[C0]]]
+//  CHECK-DAG: [[V1:%.+]] = memref.extract_aligned_pointer_as_index [[ARG1]]
+//  CHECK-DAG: [[C1:%.+]] = arith.constant 1 : index
+//  CHECK-DAG: memref.store [[V1]], [[TO_DEALLOC_MR]][[[C1]]]
+//  CHECK-DAG: [[V2:%.+]] = memref.extract_aligned_pointer_as_index [[ARG2]]
+//  CHECK-DAG: [[C0:%.+]] = arith.constant 0 : index
+//  CHECK-DAG: memref.store [[V2]], [[TO_RETAIN_MR]][[[C0]]]
+//  CHECK-DAG: [[CAST_DEALLOC:%.+]] = memref.cast [[TO_DEALLOC_MR]] : memref<2xindex> to memref<?xindex>
+//  CHECK-DAG: [[CAST_RETAIN:%.+]] = memref.cast [[TO_RETAIN_MR]] : memref<1xindex> to memref<?xindex>
+//  CHECK-DAG: [[C0:%.+]] = arith.constant 0 : index
+//      CHECK: [[RES0:%.+]]:2 = call @dealloc_helper([[CAST_DEALLOC]], [[CAST_RETAIN]], [[C0]])
+//      CHECK: [[SHOULD_DEALLOC_0:%.+]] = arith.andi [[RES0]]#0, [[ARG3]]
+//      CHECK: [[OWNERSHIP0:%.+]] = arith.andi [[RES0]]#1, [[ARG3]]
+//      CHECK: scf.if [[SHOULD_DEALLOC_0]] {
+//      CHECK:   memref.dealloc %arg0
+//      CHECK: }
+//      CHECK: [[C1:%.+]] = arith.constant 1 : index
+//      CHECK: [[RES1:%.+]]:2 = call @dealloc_helper([[CAST_DEALLOC]], [[CAST_RETAIN]], [[C1]])
+//      CHECK: [[SHOULD_DEALLOC_1:%.+]] = arith.andi [[RES1:%.+]]#0, [[ARG4]]
+//      CHECK: [[OWNERSHIP1:%.+]] = arith.andi [[RES1:%.+]]#1, [[ARG4]]
+//      CHECK: scf.if [[SHOULD_DEALLOC_1]]
+//      CHECK:   memref.dealloc [[ARG1]]
+//      CHECK: }
+//      CHECK: memref.dealloc [[TO_DEALLOC_MR]]
+//      CHECK: memref.dealloc [[TO_RETAIN_MR]]
+//      CHECK: return [[OWNERSHIP0]], [[OWNERSHIP1]]
+
+//      CHECK: func @dealloc_helper
+// CHECK-SAME: [[ARG0:%.+]]: memref<?xindex>, [[ARG1:%.+]]: memref<?xindex>
+// CHECK-SAME: [[ARG2:%.+]]: index
+// CHECK-SAME:  -> (i1, i1)
+//      CHECK:   [[TO_RETAIN_SIZE:%.+]] = memref.dim [[ARG1]], %c0
+//      CHECK:   [[TO_DEALLOC:%.+]] = memref.load [[ARG0]][[[ARG2]]] : memref<?xindex>
+// CHECK-NEXT:   [[NO_RETAIN_ALIAS:%.+]] = scf.for [[ITER:%.+]] = %c0 to [[TO_RETAIN_SIZE]] step %c1 iter_args([[ITER_ARG:%.+]] = %true) -> (i1) {
+// CHECK-NEXT:     [[RETAIN_VAL:%.+]] = memref.load [[ARG1]][[[ITER]]] : memref<?xindex>
+// CHECK-NEXT:     [[DOES_ALIAS:%.+]] = arith.cmpi ne, [[RETAIN_VAL]], [[TO_DEALLOC]] : index
+// CHECK-NEXT:     [[AGG_DOES_ALIAS:%.+]] = arith.andi [[ITER_ARG]], [[DOES_ALIAS]] : i1
+// CHECK-NEXT:     scf.yield [[AGG_DOES_ALIAS]] : i1
+// CHECK-NEXT:   }
+// CHECK-NEXT:   [[SHOULD_DEALLOC:%.+]] = scf.for [[ITER:%.+]] = %c0 to [[ARG2]] step %c1 iter_args([[ITER_ARG:%.+]] = [[NO_RETAIN_ALIAS]]) -> (i1) {
+// CHECK-NEXT:     [[OTHER_DEALLOC_VAL:%.+]] = memref.load [[ARG0]][[[ITER]]] : memref<?xindex>
+// CHECK-NEXT:     [[DOES_ALIAS:%.+]] = arith.cmpi ne, [[OTHER_DEALLOC_VAL]], [[TO_DEALLOC]] : index
+// CHECK-NEXT:     [[AGG_DOES_ALIAS:%.+]] = arith.andi [[ITER_ARG]], [[DOES_ALIAS]] : i1
+// CHECK-NEXT:     scf.yield [[AGG_DOES_ALIAS]] : i1
+// CHECK-NEXT:   }
+// CHECK-NEXT:   [[OWNERSHIP:%.+]] = arith.xori [[NO_RETAIN_ALIAS]], %true : i1
+// CHECK-NEXT:   return [[SHOULD_DEALLOC]], [[OWNERSHIP]] : i1, i1

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     deps = [
         ":ArithDialect",
         ":BufferizationDialect",
         ":ConversionPassIncGen",
         ":FuncDialect",
         ":IR",
         ":MemRefDialect",
         ":Pass",
+        ":SCFDialect",
         ":Support",
         ":Transforms",
     ],
 )
 
 td_library(
     name = "DLTIDialectTdFiles",
     srcs = [