[mlir][bufferization] Add specialized lowering for deallocs with one memref but arbitrary retains

[maerhart]

Summary

It is often the case that many values in the memrefs operand list can be
split off to speparate dealloc operations by the
--buffer-deallocation-simplification pass, however, the retain list has to be
preserved initially. Further canonicalization can often trim it down
considerable, but some retains may remain. In those cases, the general lowering
would be chosen, but is very inefficient. This commit adds another lowering for
those cases which avoids allocation of auxillary memrefs and the helper
function while still producing code that is linear in the number of operands of
the dealloc operation.

[springerm]

mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp

175

I think you can write ArrayRef(doesNotAliasList).drop_front() here.

[maerhart]

Address comment

mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp

175

That's so much more elegant, thanks!

mlir/lib/Conversion/BufferizationToMemRef/BufferizationToMemRef.cpp

rewriteOneMemrefNoRetainCase(bufferization::DeallocOp op, OpAdaptor adaptor,
     rewriter.replaceOpWithNewOp<scf::IfOp>(
         op, adaptor.getConditions()[0], [&](OpBuilder &builder, Location loc) {
           builder.create<memref::DeallocOp>(loc, adaptor.getMemrefs()[0]);
           builder.create<scf::YieldOp>(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 = rewriter.create<memref::ExtractAlignedPointerAsIndexOp>(
+        op->getLoc(), adaptor.getMemrefs()[0]);
+    for (Value retained : adaptor.getRetained()) {
+      Value retainedAsIdx =
+          rewriter.create<memref::ExtractAlignedPointerAsIndexOp>(op->getLoc(),
+                                                                  retained);
+      Value doesNotAlias = rewriter.create<arith::CmpIOp>(
+          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 = rewriter.create<arith::AndIOp>(op->getLoc(), prev, doesNotAlias);
+
+    // Also consider the condition given by the dealloc operation and perform a
+    // conditional deallocation guarded by that value.
+    Value shouldDealloc = rewriter.create<arith::AndIOp>(
+        op->getLoc(), prev, adaptor.getConditions()[0]);
+
+    rewriter.create<scf::IfOp>(
+        op.getLoc(), shouldDealloc, [&](OpBuilder &builder, Location loc) {
+          builder.create<memref::DeallocOp>(loc, adaptor.getMemrefs()[0]);
+          builder.create<scf::YieldOp>(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 = rewriter.create<arith::ConstantOp>(
+        op->getLoc(), rewriter.getBoolAttr(true));
+    for (Value doesNotAlias : doesNotAliasList) {
+      Value aliases =
+          rewriter.create<arith::XOrIOp>(op->getLoc(), doesNotAlias, trueVal);
+      Value result = rewriter.create<arith::AndIOp>(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.
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().empty()) {
+      Value falseVal = rewriter.create<arith::ConstantOp>(
+          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);
+
     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 three memrefs of
   /// booleans as arguments:
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()) {
+      if (deallocOp.getMemrefs().size() > 1) {
         helperFuncOp = DeallocOpConversion::buildDeallocationHelperFunction(
             builder, getOperation()->getLoc(), symbolTable);
         return WalkResult::interrupt();
       }
       return WalkResult::advance();
     });
 
     RewritePatternSet patterns(&getContext());

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
+  // CHECK-NOT: bufferization.dealloc
   bufferization.dealloc
   return
 }
 
 // -----
 
+func.func @conversion_dealloc_empty_but_retains(%arg0: memref<2xi32>, %arg1: memref<2xi32>) -> (i1, i1) {
+  %0:2 = bufferization.dealloc retain (%arg0, %arg1 : memref<2xi32>, memref<2xi32>)
+  return %0#0, %0#1 : i1, i1
+}
+
+// CHECK-LABEL: func @conversion_dealloc_empty
+//  CHECK-NEXT: [[FALSE:%.+]] = arith.constant false
+//  CHECK-NEXT: return [[FALSE]], [[FALSE]] :
+
+// -----
+
 // 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) {
   bufferization.dealloc (%arg0 : memref<2xf32>) if (%arg1)
   return
 }
 
 //      CHECk: scf.if [[ARG1]] {
 // CHECk-NEXT:   memref.dealloc [[ARG0]] : memref<2xf32>
 // CHECk-NEXT: }
 // CHECk-NEXT: return
 
 // -----
 
+func.func @conversion_dealloc_one_memref_and_multiple_retained(%arg0: memref<2xf32>, %arg1: memref<1xf32>, %arg2: i1, %arg3: memref<2xf32>) -> (i1, i1) {
+  %0:2 = bufferization.dealloc (%arg0 : memref<2xf32>) if (%arg2) retain (%arg1, %arg3 : memref<1xf32>, memref<2xf32>)
+  return %0#0, %0#1 : i1, i1
+}
+
+// CHECK-LABEL: func @conversion_dealloc_one_memref_and_multiple_retained
+//  CHECK-SAME: ([[ARG0:%.+]]: memref<2xf32>, [[ARG1:%.+]]: memref<1xf32>, [[ARG2:%.+]]: i1, [[ARG3:%.+]]: memref<2xf32>)
+//   CHECK-DAG: [[M0:%.+]] = memref.extract_aligned_pointer_as_index [[ARG0]]
+//   CHECK-DAG: [[R0:%.+]] = memref.extract_aligned_pointer_as_index [[ARG1]]
+//   CHECK-DAG: [[R1:%.+]] = memref.extract_aligned_pointer_as_index [[ARG3]]
+//   CHECK-DAG: [[DOES_NOT_ALIAS_R0:%.+]] = arith.cmpi ne, [[M0]], [[R0]] : index
+//   CHECK-DAG: [[DOES_NOT_ALIAS_R1:%.+]] = arith.cmpi ne, [[M0]], [[R1]] : index
+//       CHECK: [[NOT_RETAINED:%.+]] = arith.andi [[DOES_NOT_ALIAS_R0]], [[DOES_NOT_ALIAS_R1]]
+//       CHECK: [[SHOULD_DEALLOC:%.+]] = arith.andi [[NOT_RETAINED]], [[ARG2]]
+//       CHECK: scf.if [[SHOULD_DEALLOC]]
+//       CHECK:   memref.dealloc [[ARG0]]
+//       CHECK: }
+//   CHECK-DAG: [[ALIASES_R0:%.+]] = arith.xori [[DOES_NOT_ALIAS_R0]], %true
+//   CHECK-DAG: [[ALIASES_R1:%.+]] = arith.xori [[DOES_NOT_ALIAS_R1]], %true
+//   CHECK-DAG: [[RES0:%.+]] = arith.andi [[ALIASES_R0]], [[ARG2]]
+//   CHECK-DAG: [[RES1:%.+]] = arith.andi [[ALIASES_R1]], [[ARG2]]
+//       CHECK: return [[RES0]], [[RES1]]
+
+// CHECK-NOT: func @dealloc_helper
+
+// -----
+
 func.func @conversion_dealloc_multiple_memrefs_and_retained(%arg0: memref<2xf32>, %arg1: memref<5xf32>, %arg2: memref<1xf32>, %arg3: i1, %arg4: i1, %arg5: memref<2xf32>) -> (i1, i1) {
   %0:2 = bufferization.dealloc (%arg0, %arg1 : memref<2xf32>, memref<5xf32>) if (%arg3, %arg4) retain (%arg2, %arg5 : memref<1xf32>, memref<2xf32>)
   return %0#0, %0#1 : i1, i1
 }
 
 // CHECK-LABEL: func @conversion_dealloc_multiple_memrefs_and_retained
 // CHECK-SAME: ([[ARG0:%.+]]: memref<2xf32>, [[ARG1:%.+]]: memref<5xf32>,
 // CHECK-SAME: [[ARG2:%.+]]: memref<1xf32>, [[ARG3:%.+]]: i1, [[ARG4:%.+]]: i1,