[mlir][bufferize] Improve error message when returning allocs

[springerm]

Summary

The previous error message was confusing. Also improve code documentation and some minor code cleanups.

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

annotateOpsWithBufferizationMarkers(Operation *op,
   op->walk([&](BufferizableOpInterface bufferizableOp) {
     if (options.isOpAllowed(bufferizableOp.getOperation()))
       for (OpOperand &opOperand : bufferizableOp->getOpOperands())
         if (opOperand.get().getType().isa<TensorType>())
           setInPlaceOpOperand(opOperand, aliasInfo.isInPlace(opOperand));
   });
 }
 
-/// Assert that IR is in destination-passing style. I.e., every value that is
-/// returned or yielded from a block is:
-/// * aliasing a bbArg of that block or a parent block, or
-/// * aliasing an OpResult of a op in a parent block.
+/// Assert that every allocation can be deallocated in the same block. I.e.,
+/// every value that is returned or yielded from a block is:
+/// * guaranteed to be aliasing a bbArg of that block or a parent block, or
+/// * guaranteed to be aliasing an OpResult of a op in a parent block.
+///
+/// In that case, buffer deallocation is simple: Every allocated buffer can be
+/// deallocated in the same block. Otherwise, the buffer deallocation pass must
+/// be run.
+///
+/// Note: The current implementation checks for equivalent values instead of
+/// aliasing values, which is stricter than needed. We can currently not check
+/// for aliasing values because the analysis is a maybe-alias analysis and we
+/// need a must-alias analysis here.
 ///
 /// Example:
 /// ```
 /// %0 = "some_op" : tensor<?xf32>
 /// %1 = scf.if %c -> (tensor<?xf32>) {
 ///   scf.yield %0 : tensor<?xf32>
 /// } else {
 ///   %t = linalg.alloc_tensor : tensor<?xf32>
 ///   scf.yield %t : tensor<?xf32>
 /// }
 /// ```
-/// In the above example, the first scf.yield op satifies destination-passing
-/// style because the yielded value %0 is defined in the parent block. The
-/// second scf.yield op does not satisfy destination-passing style because the
-/// yielded value %t is defined in the same block as the scf.yield op.
-// TODO: The current implementation checks for equivalent values instead of
-// aliasing values, which is stricter than needed. We can currently not check
-// for aliasing values because the analysis is a maybe-alias analysis and we
-// need a must-alias analysis here.
-static LogicalResult
-assertDestinationPassingStyle(Operation *op, AnalysisState &state,
-                              BufferizationAliasInfo &aliasInfo,
-                              SmallVector<Operation *> &newOps) {
+///
+/// In the above example, the second scf.yield op is problematic because the
+/// yielded value %t is defined in the same block as the scf.yield op and
+/// and bufferizes to a new allocation.
+// TODO: Remove buffer deallocation from One-Shot Bufferize and fix the buffer
+// deallocation pass.
+static LogicalResult assertNoAllocsReturned(Operation *op,
+                                            const BufferizationOptions &options,
+                                            BufferizationAliasInfo &aliasInfo) {
   LogicalResult status = success();
   DominanceInfo domInfo(op);
   op->walk([&](Operation *returnOp) {
-    if (!isRegionReturnLike(returnOp) ||
-        !state.getOptions().isOpAllowed(returnOp))
+    if (!isRegionReturnLike(returnOp) || !options.isOpAllowed(returnOp))
       return WalkResult::advance();
 
     for (OpOperand &returnValOperand : returnOp->getOpOperands()) {
       Value returnVal = returnValOperand.get();
       // Skip non-tensor values.
       if (!returnVal.getType().isa<TensorType>())
         continue;
 
for (OpOperand &returnValOperand : returnOp->getOpOperands()) {
         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)
-        status =
-            returnOp->emitError()
+        status = returnOp->emitError()
                  << "operand #" << returnValOperand.getOperandNumber()
-            << " of ReturnLike op does not satisfy destination passing style";
+                 << " may return/yield a new buffer allocation";
     }
 
     return WalkResult::advance();
   });
 
   return status;
 }
 
 LogicalResult bufferization::analyzeOp(Operation *op,
                                        OneShotAnalysisState &state) {
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
   const OneShotBufferizationOptions &options = state.getOptions();
 
   if (failed(checkAliasInfoConsistency(op, domInfo, state, aliasInfo)))
     return failure();
 
   // If the analysis fails, just return.
   if (failed(inPlaceAnalysis(op, aliasInfo, state, domInfo,
                              options.analysisFuzzerSeed)))
     return failure();
   equivalenceAnalysis(op, aliasInfo, state);
 
   bool failedAnalysis = false;
-  if (!options.allowReturnAllocs) {
-    SmallVector<Operation *> newOps;
-    failedAnalysis |=
-        failed(assertDestinationPassingStyle(op, state, aliasInfo, newOps));
-  }
+  if (!options.allowReturnAllocs)
+    failedAnalysis |= failed(assertNoAllocsReturned(op, options, aliasInfo));
 
   // Gather some extra analysis data.
   state.gatherYieldedTensors(op);
   state.gatherUndefinedTensorUses(op);
 
   // Analysis verification: After setting up alias/equivalence sets, each op
   // can check for expected invariants/limitations and fail the analysis if
   // necessary.

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

 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 of ReturnLike op does not satisfy destination passing style}}
+    // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
     scf.yield %t2 : tensor<?xf32>
   }
-  // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}}
+  // 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>
-    // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}}
+    // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
     scf.yield %t2 : tensor<?xf32>
   }
   %f = tensor.extract %r[%idx] : tensor<?xf32>
   return %f : f32
 }
 
 // -----
 
func.func @extract_slice_fun(%A : tensor<?xf32> {bufferization.writable = true})
   // 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 of ReturnLike op does not satisfy destination passing style}}
+  // 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 of ReturnLike op does not satisfy destination passing style}}
+  // 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 of ReturnLike op does not satisfy destination passing style}}
+  // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
   return %r: tensor<4xf32>
 }
 
 // -----
 
 func.func @mini_test_case1() -> tensor<10x20xf32> {
   %f0 = arith.constant 0.0 : f32
   %t = bufferization.alloc_tensor() : tensor<10x20xf32>
   %r = linalg.fill ins(%f0 : f32) outs(%t : tensor<10x20xf32>) -> tensor<10x20xf32>
-  // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}}
+  // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
   return %r : tensor<10x20xf32>
 }
 
 // -----
 
 func.func @main() -> tensor<4xi32> {
   %r = scf.execute_region -> tensor<4xi32> {
     %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
-    // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}}
+    // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
     scf.yield %A: tensor<4xi32>
   }
 
-  // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}}
+  // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
   return %r: tensor<4xi32>
 }
 
 // -----
 
 func.func @to_memref_op_is_writing(
     %t1: tensor<?xf32> {bufferization.writable = true}, %idx1: index,
     %idx2: index, %idx3: index, %v1: vector<5xf32>) -> (vector<5xf32>, vector<5xf32>) {
func.func @call_to_unknown_tensor_returning_func(%t : tensor<?xf32>) {
   call @foo(%t) : (tensor<?xf32>) -> (f32, tensor<?xf32>, f32)
   return
 }
 
 // -----
 
 func.func @foo(%t : tensor<5xf32>) -> (tensor<5xf32>) {
   %0 = bufferization.alloc_tensor() : tensor<5xf32>
-  // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}}
+  // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
   return %0 : tensor<5xf32>
 }
 
 // Note: This function is not analyzed because there was an error in the
 // previous one.
 func.func @call_to_func_returning_non_equiv_tensor(%t : tensor<5xf32>) {
   call @foo(%t) : (tensor<5xf32>) -> (tensor<5xf32>)
   return
 }
 
 // -----
 
-func.func @destination_passing_style_dominance_test_1(%cst : f32, %idx : index,
-                                                      %idx2 : index) -> f32 {
+func.func @yield_alloc_dominance_test_1(%cst : f32, %idx : index,
   %0 = scf.execute_region -> tensor<?xf32> {
     %1 = bufferization.alloc_tensor(%idx) : tensor<?xf32>
-    // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}}
+    // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}}
     scf.yield %1 : tensor<?xf32>
   }
   %2 = tensor.insert %cst into %0[%idx] : tensor<?xf32>
   %r = tensor.extract %2[%idx2] : tensor<?xf32>
   return %r : f32
 }
 
 // -----
 
-func.func @destination_passing_style_dominance_test_2(%cst : f32, %idx : index,
-                                                      %idx2 : index) -> f32 {
+func.func @yield_alloc_dominance_test_2(%cst : f32, %idx : index,
   %1 = bufferization.alloc_tensor(%idx) : tensor<?xf32>
 
   %0 = scf.execute_region -> tensor<?xf32> {
-    // This YieldOp is in destination-passing style, thus no error.
+    // This YieldOp returns a value that is defined in a parent block, thus
+    // no error.
     scf.yield %1 : tensor<?xf32>
   }
   %2 = tensor.insert %cst into %0[%idx] : tensor<?xf32>
   %r = tensor.extract %2[%idx2] : tensor<?xf32>
   return %r : f32
 }