[mlir][bufferization] Factor out bufferization.dealloc lowering into separate pass

Moves the lowering of `bufferization.dealloc` to memref into a separate pass,
but still registers the pattern in the conversion pass.  This is helpful when
some tensor values (and thus `to_memref` or `to_tensor` operations) still
remain, e.g., when the function boundaries are not converted, or when constant
tensors are converted to memref.get_global at a later point.

However, it is still recommended to perform all bufferization before
deallocation to avoid memory leaks as all memref allocations inserted after the
deallocation pass was applied, have to be handled manually.

Note: The buffer deallocation pass assumes that memref values defined by
`bufferization.to_memref` don't return ownership and don't have to be
deallocated. `bufferization.to_tensor` operations are handled similarly to
`bufferization.clone` operations with the exception that the result value is
not handled because it's a tensor (not a memref).

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D159180

Author: maerhart

mlir/include/mlir/Conversion/Passes.td

@@ -195,12 +195,14 @@ def ConvertBufferizationToMemRef : Pass<"convert-bufferization-to-memref"> {
     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
-    is needed, since some clone operations could remain after applying several
-    transformation processes. Currently, only `canonicalize` transforms clone
-    operations or even eliminates them. This can lead to errors if any clone op
-    survived after all conversion passes (starting from the bufferization
-    dialect) are performed.
+    operation into `memref.alloc` and `memref.copy` operations and
+    `bufferization.dealloc` operations (the same way as the
+    `-bufferization-lower-deallocations` pass). The conversion of `clone`
+    operations is needed, since some clone operations could remain after
+    applying several transformation processes. Currently, only `canonicalize`
+    transforms clone operations or even eliminates them. This can lead to errors
+    if any clone op survived after all conversion passes (starting from the
+    bufferization dialect) are performed.
 
     See:
     https://llvm.discourse.group/t/bufferization-error-related-to-memref-clone/4665

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h

@@ -5,6 +5,9 @@
 
 namespace mlir {
 class ModuleOp;
+class RewritePatternSet;
+class OpBuilder;
+class SymbolTable;
 
 namespace func {
 class FuncOp;
@@ -29,6 +32,98 @@ std::unique_ptr<Pass> createBufferDeallocationPass();
 /// static alias analysis.
 std::unique_ptr<Pass> createBufferDeallocationSimplificationPass();
 
+/// Creates an instance of the LowerDeallocations pass to lower
+/// `bufferization.dealloc` operations to the `memref` dialect.
+std::unique_ptr<Pass> createLowerDeallocationsPass();
+
+/// Adds the conversion pattern of the `bufferization.dealloc` operation to the
+/// given pattern set for use in other transformation passes.
+void populateBufferizationDeallocLoweringPattern(
+    RewritePatternSet &patterns, func::FuncOp deallocLibraryFunc);
+
+/// Construct the library function needed for the fully generic
+/// `bufferization.dealloc` lowering implemented in the LowerDeallocations pass.
+/// The function can then be called at bufferization dealloc sites to determine
+/// aliasing and ownership.
+///
+/// The generated function takes two memrefs of indices and three memrefs of
+/// booleans as arguments:
+///   * The first argument A should contain the result of the
+///     extract_aligned_pointer_as_index operation applied to the memrefs to be
+///     deallocated
+///   * The second argument B should contain the result of the
+///     extract_aligned_pointer_as_index operation applied to the memrefs to be
+///     retained
+///   * The third argument C should contain the conditions as passed directly
+///     to the deallocation operation.
+///   * The fourth argument D is used to pass results to the caller. Those
+///     represent the condition under which the memref at the corresponding
+///     position in A should be deallocated.
+///   * The fifth argument E is used to pass results to the caller. It
+///     provides the ownership value corresponding the the memref at the same
+///     position in B
+///
+/// This helper function is supposed to be called once for each
+/// `bufferization.dealloc` operation to determine the deallocation need and new
+/// ownership indicator for the retained values, but does not perform the
+/// deallocation itself.
+///
+/// Generated code:
+/// ```
+/// func.func @dealloc_helper(
+///     %dyn_dealloc_base_pointer_list: memref<?xindex>,
+///     %dyn_retain_base_pointer_list: memref<?xindex>,
+///     %dyn_cond_list: memref<?xi1>,
+///     %dyn_dealloc_cond_out: memref<?xi1>,
+///     %dyn_ownership_out: memref<?xi1>) {
+///   %c0 = arith.constant 0 : index
+///   %c1 = arith.constant 1 : index
+///   %true = arith.constant true
+///   %false = arith.constant false
+///   %num_dealloc_memrefs = memref.dim %dyn_dealloc_base_pointer_list, %c0
+///   %num_retain_memrefs = memref.dim %dyn_retain_base_pointer_list, %c0
+///   // Zero initialize result buffer.
+///   scf.for %i = %c0 to %num_retain_memrefs step %c1 {
+///     memref.store %false, %dyn_ownership_out[%i] : memref<?xi1>
+///   }
+///   scf.for %i = %c0 to %num_dealloc_memrefs step %c1 {
+///     %dealloc_bp = memref.load %dyn_dealloc_base_pointer_list[%i]
+///     %cond = memref.load %dyn_cond_list[%i]
+///     // Check for aliasing with retained memrefs.
+///     %does_not_alias_retained = scf.for %j = %c0 to %num_retain_memrefs
+///         step %c1 iter_args(%does_not_alias_aggregated = %true) -> (i1) {
+///       %retain_bp = memref.load %dyn_retain_base_pointer_list[%j]
+///       %does_alias = arith.cmpi eq, %retain_bp, %dealloc_bp : index
+///       scf.if %does_alias {
+///         %curr_ownership = memref.load %dyn_ownership_out[%j]
+///         %updated_ownership = arith.ori %curr_ownership, %cond : i1
+///         memref.store %updated_ownership, %dyn_ownership_out[%j]
+///       }
+///       %does_not_alias = arith.cmpi ne, %retain_bp, %dealloc_bp : index
+///       %updated_aggregate = arith.andi %does_not_alias_aggregated,
+///                                       %does_not_alias : i1
+///       scf.yield %updated_aggregate : i1
+///     }
+///     // Check for aliasing with dealloc memrefs in the list before the
+///     // current one, i.e.,
+///     // `fix i, forall j < i: check_aliasing(%dyn_dealloc_base_pointer[j],
+///     // %dyn_dealloc_base_pointer[i])`
+///     %does_not_alias_any = scf.for %j = %c0 to %i step %c1
+///        iter_args(%does_not_alias_agg = %does_not_alias_retained) -> (i1) {
+///       %prev_dealloc_bp = memref.load %dyn_dealloc_base_pointer_list[%j]
+///       %does_not_alias = arith.cmpi ne, %prev_dealloc_bp, %dealloc_bp
+///       %updated_alias_agg = arith.andi %does_not_alias_agg, %does_not_alias
+///       scf.yield %updated_alias_agg : i1
+///     }
+///     %dealloc_cond = arith.andi %does_not_alias_any, %cond : i1
+///     memref.store %dealloc_cond, %dyn_dealloc_cond_out[%i] : memref<?xi1>
+///   }
+///   return
+/// }
+/// ```
+func::FuncOp buildDeallocationLibraryFunction(OpBuilder &builder, Location loc,
+                                              SymbolTable &symbolTable);
+
 /// Run buffer deallocation.
 LogicalResult deallocateBuffers(Operation *op);
 

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

@@ -108,6 +108,29 @@ def BufferDeallocationSimplification :
   ];
 }
 
+def LowerDeallocations : Pass<"bufferization-lower-deallocations"> {
+  let summary = "Lowers `bufferization.dealloc` operations to `memref.dealloc`"
+                "operations";
+  let description = [{
+    This pass lowers `bufferization.dealloc` operations to the `memref` dialect.
+    It can be applied to a `builtin.module` or operations implementing the
+    `FunctionOpInterface`. For the latter, only simple `dealloc` operations can
+    be lowered because the library function necessary for the fully generic
+    lowering cannot be inserted. In this case, an error will be emitted.
+    Next to `memref.dealloc` operations, it may also emit operations from the
+    `arith`, `scf`, and `func` dialects to build conditional deallocations and
+    library functions to avoid code-size blow-up.
+  }];
+
+  let constructor =
+    "mlir::bufferization::createLowerDeallocationsPass()";
+
+  let dependentDialects = [
+    "arith::ArithDialect", "memref::MemRefDialect", "scf::SCFDialect",
+    "func::FuncDialect"
+  ];
+}
+
 def BufferHoisting : Pass<"buffer-hoisting", "func::FuncOp"> {
   let summary = "Optimizes placement of allocation operations by moving them "
                 "into common dominators and out of nested regions";