[mlir][bufferization] Improve memory leak detection

[springerm]

Summary

One-Shot Bufferize can currently not handle cases in which a buffer allocation is yielded from a block. This would cause a memory leak.

There is a check to detect such IR, but it is unnecessarily conservative. This change relaxes this restriction and lets us bufferize IR that we could not bufferize before.

Depends On: D143927

[mravishankar]

I am missing context for this change (sorry I am way behind on my reviews). Can you provide a few more details.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

public:
   /// bufferizable.
   AliasingOpOperandList getAliasingOpOperands(OpResult result) const;
 
   /// Determine which OpResult will alias with `opOperand` if the op is
   /// bufferized in place. Return all tensor OpResults if the op is not
   /// bufferizable.
   AliasingOpResultList getAliasingOpResults(OpOperand &opOperand) const;
 
+  /// Return `true` if `opResult` bufferizes to a memory allocation. Return
+  /// `true` if the op is not bufferizable.
+  bool bufferizesToAllocation(OpResult opResult) const;
+
   /// Return true if `opOperand` bufferizes to a memory read. Return `true` if
   /// the op is not bufferizable.
   bool bufferizesToMemoryRead(OpOperand &opOperand) const;
 
   /// Return true if `opOperand` bufferizes to a memory write. Return true` if
   /// the op is not bufferizable.
   bool bufferizesToMemoryWrite(OpOperand &opOperand) const;
 
public:
   /// %r = arith.select %cond, %a, %b : tensor<10xf32>
   /// findDefinitions(%r) = {%b}. %a is excluded because it does not define the
   /// contents of the tensor.
   ///
   /// Note: OpResults of unknown ops are handled conservatively and assumed to
   /// be definitions.
   SetVector<Value> findDefinitions(Value value) const;
 
+  /// Find the OpResults in the reverse use-def chain that are a buffer
+  /// allocation. Either because they bufferize to an allocation or because
+  /// at least one of their aliasing OpOperands is out-of-place.
+  SetVector<OpResult> findAllocations(Value value) const;
+
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   virtual bool isInPlace(OpOperand &opOperand) const;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   virtual bool areEquivalentBufferizedValues(Value v1, Value v2) const;
 
   /// Return true if `v1` and `v2` may bufferize to aliasing buffers.
   virtual bool areAliasingBufferizedValues(Value v1, Value v2) const;

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

AnalysisState::getAliasingOpResults(OpOperand &opOperand) const {
   if (auto bufferizableOp =
           getOptions().dynCastBufferizableOp(opOperand.getOwner()))
     return bufferizableOp.getAliasingOpResults(opOperand, *this);
 
   // The op is not bufferizable.
   return detail::unknownGetAliasingOpResults(opOperand);
 }
 
+bool AnalysisState::bufferizesToAllocation(OpResult opResult) const {
+  if (Operation *op = opResult.getDefiningOp())
+    if (auto bufferizableOp = getOptions().dynCastBufferizableOp(op))
+      return bufferizableOp.bufferizesToAllocation(opResult);
+
+  // The op is not bufferizable. Conservatively return true.
+  return true;
+}
+
 /// Return true if `opOperand` bufferizes to a memory read. Return `true` if the
 /// op is not bufferizable.
 bool AnalysisState::bufferizesToMemoryRead(OpOperand &opOperand) const {
   if (auto bufferizableOp =
           getOptions().dynCastBufferizableOp(opOperand.getOwner()))
     return bufferizableOp.bufferizesToMemoryRead(opOperand, *this);
 
   // Unknown op that returns a tensor. The inplace analysis does not support it.

 // Find the values that define the contents of the given value.
 llvm::SetVector<Value> AnalysisState::findDefinitions(Value value) const {
   TraversalConfig config;
   config.alwaysIncludeLeaves = false;
   return findValueInReverseUseDefChain(
       value, [&](Value v) { return this->bufferizesToMemoryWrite(v); }, config);
 }
 
+llvm::SetVector<OpResult> AnalysisState::findAllocations(Value value) const {
+  TraversalConfig config;
+  config.alwaysIncludeLeaves = true;
+  config.followInPlaceOnly = true;
+  auto bufferizesToAlloc = [&](Value v) {
+    auto opResult = v.dyn_cast<OpResult>();
+    if (!opResult)
+      return true;
+    return bufferizesToAllocation(opResult);
+  };
+  SetVector<Value> potentialAllocs =
+      findValueInReverseUseDefChain(value, bufferizesToAlloc, config);
+  SetVector<OpResult> result;
+  for (Value v : potentialAllocs) {
+    auto opResult = v.dyn_cast<OpResult>();
+    if (!opResult)
+      continue;
+    if (bufferizesToAllocation(opResult))
+      result.insert(opResult);
+    if (llvm::any_of(getAliasingOpOperands(opResult),
+                     [&](AliasingOpOperand alias) {
+                       return !isInPlace(*alias.opOperand);
+                     }))
+      result.insert(opResult);
+  }
+  return result;
+}
+
 AnalysisState::AnalysisState(const BufferizationOptions &options)
     : AnalysisState(options, TypeID::get<AnalysisState>()) {}
 
 AnalysisState::AnalysisState(const BufferizationOptions &options, TypeID type)
     : options(options), type(type) {
   for (const BufferizationOptions::AnalysisStateInitFn &fn :
        options.stateInitializers)
     fn(*this);

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

static LogicalResult assertNoAllocsReturned(Operation *op,
                                             const OneShotAnalysisState &state) {
   LogicalResult status = success();
   DominanceInfo domInfo(op);
   op->walk([&](Operation *returnOp) {
     if (!isRegionReturnLike(returnOp) ||
         !state.getOptions().isOpAllowed(returnOp))
       return WalkResult::advance();
 
+    // The block from which the op is yielding.
+    Block *block = returnOp->getBlock();
+
     for (OpOperand &returnValOperand : returnOp->getOpOperands()) {
       Value returnVal = returnValOperand.get();
+
       // Skip non-tensor values.
       if (!returnVal.getType().isa<TensorType>())
         continue;
 
-      bool foundEquivValue = false;
-      state.applyOnEquivalenceClass(returnVal, [&](Value equivVal) {
-        if (auto bbArg = equivVal.dyn_cast<BlockArgument>()) {
-          Operation *definingOp = bbArg.getOwner()->getParentOp();
-          if (definingOp->isProperAncestor(returnOp))
-            foundEquivValue = true;
-          return;
-        }
-
-        Operation *definingOp = equivVal.getDefiningOp();
-        if (definingOp->getBlock()->findAncestorOpInBlock(
-                *returnOp->getParentOp()))
-          // Skip ops that happen after `returnOp` and parent ops.
-          if (happensBefore(definingOp, returnOp, domInfo))
-            foundEquivValue = true;
-      });
-
-      // Note: Returning/yielding buffer allocations is allowed only if
-      // `allowReturnAllocs` is set.
-      if (!foundEquivValue)
+      for (OpResult alloc : state.findAllocations(returnVal)) {
+        if (block->findAncestorOpInBlock(*alloc.getDefiningOp()))
           status = returnOp->emitError()
                    << "operand #" << returnValOperand.getOperandNumber()
                    << " may return/yield a new buffer allocation";
       }
+    }
 
     return WalkResult::advance();
   });
 
   return status;
 }
 
 LogicalResult bufferization::analyzeOp(Operation *op,

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir

 func.func @read_of_alloc_tensor_is_not_a_conflict(%f: f32, %idx: index) -> f32 {
   %0 = bufferization.alloc_tensor() : tensor<10xf32>
   // This can be in-place because the read below does reads undefined data.
   // CHECK: tensor.insert {{.*}} {__inplace_operands_attr__ = ["none", "true", "none"]}
   %1 = tensor.insert %f into %0[%idx] : tensor<10xf32>
   %2 = tensor.extract %0[%idx] : tensor<10xf32>
   return %2 : f32
 }
+
+// -----
+
+// CHECK-LABEL: func @scf_if_definitely_aliasing(
+func.func @scf_if_definitely_aliasing(
+    %cond: i1, %t1: tensor<?xf32> {bufferization.writable = true},
+    %idx: index) -> tensor<?xf32> {
+  %r = scf.if %cond -> (tensor<?xf32>) {
+    scf.yield %t1 : tensor<?xf32>
+  } else {
+    // This buffer definitely aliases, so it can be yielded from the block.
+    // CHECK: tensor.extract_slice
+    // CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none"]}
+    %t2 = tensor.extract_slice %t1 [%idx] [%idx] [1] : tensor<?xf32> to tensor<?xf32>
+    scf.yield %t2 : tensor<?xf32>
+  }
+  return %r : tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @scf_yield_no_allocs(
+func.func @scf_yield_no_allocs(
+    %c1: i1, %c2: i1, %A: tensor<4xf32>, %B: tensor<4xf32>, %C: tensor<4xf32>)
+  -> tensor<4xf32>
+{
+  %r = scf.if %c1 -> (tensor<4xf32>) {
+    // CHECK: arith.select
+    // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "true"]}
+    %0 = arith.select %c2, %A, %B : tensor<4xf32>
+    scf.yield %0 : tensor<4xf32>
+  } else {
+    scf.yield %C : tensor<4xf32>
+  }
+  return %r: tensor<4xf32>
+}
+
+// -----
+
+// Just make sure that bufferization succeeds. The arith.constant ops do not
+// allocate.
+
+// CHECK-LABEL: func @block_yield_constant(
+func.func @block_yield_constant(%c: i1) -> (tensor<3x4xf32>) {
+  %r = scf.if %c -> (tensor<3x4xf32>) {
+    %0 = arith.constant dense<7.0> : tensor<3x4xf32>
+    scf.yield %0 : tensor<3x4xf32>
+  } else {
+    %0 = arith.constant dense<8.0> : tensor<3x4xf32>
+    scf.yield %0 : tensor<3x4xf32>
+  }
+  return %r : tensor<3x4xf32>
+}

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir

^bb1:
     }
     return %T#0, %T#1 : tensor<f32>, tensor<f32>
   ^bb2:
     return %t2, %t1 : tensor<f32>, tensor<f32>
 }
 
 // -----
 
-func.func @scf_if_not_equivalent(
-    %cond: i1, %t1: tensor<?xf32> {bufferization.writable = true},
-    %idx: index) -> tensor<?xf32> {
-  %r = scf.if %cond -> (tensor<?xf32>) {
-    scf.yield %t1 : tensor<?xf32>
-  } else {
-    // This buffer aliases, but it is not equivalent.
-    %t2 = tensor.extract_slice %t1 [%idx] [%idx] [1] : tensor<?xf32> to tensor<?xf32>
-    // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
-    scf.yield %t2 : tensor<?xf32>
-  }
-  // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
-  return %r : tensor<?xf32>
-}
-
-// -----
-
 func.func @scf_if_not_aliasing(
     %cond: i1, %t1: tensor<?xf32> {bufferization.writable = true},
     %idx: index) -> f32 {
   %r = scf.if %cond -> (tensor<?xf32>) {
     scf.yield %t1 : tensor<?xf32>
   } else {
     // This buffer aliases.
     %t2 = bufferization.alloc_tensor(%idx) : tensor<?xf32>
%res = scf.for %arg0 = %c0 to %iters step %c1 iter_args(%bbarg = %A) -> (tensor<?xf32>) {
     scf.yield %r : tensor<?xf32>
   }
   call @fun_with_side_effects(%res) : (tensor<?xf32>) -> ()
   return
 }
 
 // -----
 
-func.func @extract_slice_fun(%A : tensor<?xf32> {bufferization.writable = true})
-  ->  tensor<4xf32>
-{
-  // This bufferizes to a pattern that the cross-function boundary pass needs to
-  // convert into a new memref argument at all call site; this may be either:
-  //   - an externally created aliasing subview (if we want to allow aliasing
-  //     function arguments).
-  //   - a new alloc + copy (more expensive but does not create new function
-  //     argument aliasing).
-  %r0 = tensor.extract_slice %A[0][4][1] : tensor<?xf32> to tensor<4xf32>
-
-  // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
-  return %r0: tensor<4xf32>
-}
-
-// -----
-
-func.func @scf_yield(%b : i1, %A : tensor<4xf32>, %B : tensor<4xf32>) -> tensor<4xf32>
-{
-  %r = scf.if %b -> (tensor<4xf32>) {
-    scf.yield %A : tensor<4xf32>
-  } else {
-    scf.yield %B : tensor<4xf32>
-  }
-  // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
-  return %r: tensor<4xf32>
-}
-
-// -----
-
 func.func @unknown_op(%A : tensor<4xf32>) -> tensor<4xf32>
 {
   // expected-error: @+1 {{op was not bufferized}}
   %r = "marklar"(%A) : (tensor<4xf32>) -> (tensor<4xf32>)
   // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
   return %r: tensor<4xf32>
 }