[mlir][bufferize] Add a way for ops to fail the analysis

Add `BufferizableOpInterface::verifyAnalysis`. Ops can implement this method to check for expected invariants and limitations.

The purpose of this change is to introduce a modular way of checking assertions such as `assertScfForAliasingProperties`.

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

Author: matthias-springer

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

@@ -290,6 +290,23 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
         /*defaultImplementation=*/[{
           return false;
         }]
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Return `failure` if this op does not pass the analysis. This method
+          is run during One-Shot Bufferize (after all post-analysis steps). If
+          the op does not pass the analysis, bufferization is aborted.
+
+          This method can be used to check expected invariants and limitations
+          of the current bufferization implementation.
+        }],
+        /*retType=*/"LogicalResult",
+        /*methodName=*/"verifyAnalysis",
+        /*args=*/(ins "const BufferizationState &":$state),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          return success();
+        }]
       >
   ];
 

mlir/include/mlir/Dialect/SCF/BufferizableOpInterfaceImpl.h

@@ -20,16 +20,6 @@ class BufferizationAliasInfo;
 } // namespace bufferization
 
 namespace scf {
-/// Assert that yielded values of an scf.for op are aliasing their corresponding
-/// bbArgs. This is required because the i-th OpResult of an scf.for op is
-/// currently assumed to alias with the i-th iter_arg (in the absence of
-/// conflicts).
-LogicalResult
-assertScfForAliasingProperties(Operation *op,
-                               bufferization::BufferizationState &state,
-                               bufferization::BufferizationAliasInfo &aliasInfo,
-                               SmallVector<Operation *> &newOps);
-
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 } // namespace scf
 } // namespace mlir

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

@@ -778,6 +778,19 @@ LogicalResult bufferization::analyzeOp(Operation *op,
       return failure();
   }
 
+  // Analysis verification: After setting up alias/equivalence sets, each op
+  // can check for expected invariants/limitations and fail the analysis if
+  // necessary.
+  bool passedAnalysis = true;
+  op->walk([&](Operation *op) {
+    if (BufferizableOpInterface bufferizableOp =
+            options.dynCastBufferizableOp(op))
+      if (failed(bufferizableOp.verifyAnalysis(state)))
+        passedAnalysis = false;
+  });
+  if (!passedAnalysis)
+    return failure();
+
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly)
     annotateOpsWithBufferizationMarkers(op, aliasInfo, state);

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp

@@ -105,9 +105,6 @@ void LinalgComprehensiveModuleBufferize::runOnOperation() {
     opt = *options;
   }
 
-  // Only certain scf.for ops are supported by the analysis.
-  opt.addPostAnalysisStep(scf::assertScfForAliasingProperties);
-
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

@@ -385,6 +385,37 @@ struct ForOpInterface
 
     return success();
   }
+
+  /// Assert that yielded values of an scf.for op are aliasing with their
+  /// corresponding bbArgs. This is required because the i-th OpResult of an
+  /// scf.for op is currently assumed to alias with the i-th iter_arg (in the
+  /// absence of conflicts).
+  LogicalResult verifyAnalysis(Operation *op,
+                               const BufferizationState &state) const {
+    auto forOp = cast<scf::ForOp>(op);
+    auto yieldOp =
+        cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator());
+    for (OpOperand &operand : yieldOp->getOpOperands()) {
+      auto tensorType = operand.get().getType().dyn_cast<TensorType>();
+      if (!tensorType)
+        continue;
+
+      OpOperand &forOperand = forOp.getOpOperandForResult(
+          forOp->getResult(operand.getOperandNumber()));
+      auto bbArg = forOp.getRegionIterArgForOpOperand(forOperand);
+      // Note: This is overly strict. We should check for aliasing bufferized
+      // values. But we don't have a "must-alias" analysis yet.
+      if (!state.areEquivalentBufferizedValues(operand.get(), bbArg))
+        // TODO: this could get resolved with copies but it can also turn into
+        // swaps so we need to be careful about order of copies.
+        return yieldOp->emitError()
+               << "Yield operand #" << operand.getOperandNumber()
+               << " does not bufferize to a buffer that is aliasing the "
+                  "matching"
+               << " enclosing scf::for operand";
+    }
+    return success();
+  }
 };
 
 /// Bufferization of scf.yield. Bufferized as part of their enclosing ops, so
@@ -434,41 +465,6 @@ struct YieldOpInterface
 } // namespace scf
 } // namespace mlir
 
-LogicalResult mlir::scf::assertScfForAliasingProperties(
-    Operation *op, BufferizationState &state, BufferizationAliasInfo &aliasInfo,
-    SmallVector<Operation *> &newOps) {
-  LogicalResult status = success();
-
-  op->walk([&](scf::ForOp forOp) {
-    auto yieldOp =
-        cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator());
-    for (OpOperand &operand : yieldOp->getOpOperands()) {
-      auto tensorType = operand.get().getType().dyn_cast<TensorType>();
-      if (!tensorType)
-        continue;
-
-      OpOperand &forOperand = forOp.getOpOperandForResult(
-          forOp->getResult(operand.getOperandNumber()));
-      auto bbArg = forOp.getRegionIterArgForOpOperand(forOperand);
-      // Note: This is overly strict. We should check for aliasing bufferized
-      // values. But we don't have a "must-alias" analysis yet.
-      if (!aliasInfo.areEquivalentBufferizedValues(operand.get(), bbArg)) {
-        // TODO: this could get resolved with copies but it can also turn into
-        // swaps so we need to be careful about order of copies.
-        status =
-            yieldOp->emitError()
-            << "Yield operand #" << operand.getOperandNumber()
-            << " does not bufferize to a buffer that is aliasing the matching"
-            << " enclosing scf::for operand";
-        return WalkResult::interrupt();
-      }
-    }
-    return WalkResult::advance();
-  });
-
-  return status;
-}
-
 void mlir::scf::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addOpInterface<ExecuteRegionOp, ExecuteRegionOpInterface>();

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir

@@ -87,7 +87,7 @@ func @scf_for(%A : tensor<?xf32>,
               %B : tensor<?xf32> {linalg.inplaceable = true},
               %C : tensor<4xf32>,
               %lb : index, %ub : index, %step : index)
-  -> (tensor<?xf32>, tensor<?xf32>)
+  -> (f32, f32)
 {
   %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)
       -> (tensor<?xf32>, tensor<?xf32>)
@@ -102,7 +102,9 @@ func @scf_for(%A : tensor<?xf32>,
     scf.yield %ttB, %ttA : tensor<?xf32>, tensor<?xf32>
   }
 
-  return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>
+  %f0 = tensor.extract %r0#0[%step] : tensor<?xf32>
+  %f1 = tensor.extract %r0#1[%step] : tensor<?xf32>
+  return %f0, %f1: f32, f32
 }
 
 // -----

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-partial.mlir

@@ -179,35 +179,6 @@ func @simple_tensor_test(%t1 : tensor<?xf32>, %f : f32) -> tensor<?xf32> {
 
 // -----
 
-// CHECK-SCF-LABEL: func @simple_scf_for(
-//  CHECK-SCF-SAME:     %[[t1:.*]]: tensor<?xf32>
-func @simple_scf_for(
-    %t1: tensor<?xf32>, %sz: index, %step: index, %f: f32) -> tensor<?xf32> {
-  %c0 = arith.constant 0 : index
-
-  // CHECK-SCF: %[[t1_memref:.*]] = bufferization.to_memref %[[t1]]
-  // CHECK-SCF: %[[alloc:.*]] = memref.alloc
-  // CHECK-SCF: %[[casted:.*]] = memref.cast %[[alloc]]
-  // CHECK-SCF: memref.copy %[[t1_memref]], %[[alloc]]
-  // CHECK-SCF: %[[scf_for:.*]] = scf.for %[[iv:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[arg0:.*]] = %[[casted]]) -> ({{.*}}) {
-  %0 = scf.for %iv = %c0 to %sz step %step iter_args(%arg0 = %t1) -> tensor<?xf32> {
-    // CHECK-SCF: %[[arg0_tensor:.*]] = bufferization.to_tensor %[[arg0]]
-    // CHECK-SCF: %[[insert:.*]] = tensor.insert %{{.*}} into %[[arg0_tensor]]
-    %1 = tensor.insert %f into %arg0[%iv] : tensor<?xf32>
-
-    // CHECK-SCF: %[[insert_memref:.*]] = bufferization.to_memref %[[insert]]
-    // CHECK-SCF: scf.yield %[[insert_memref]]
-    scf.yield %1 : tensor<?xf32>
-  }
-  // CHECK-SCF: }
-
-  // CHECK-SCF: %[[scf_for_tensor:.*]] = bufferization.to_tensor %[[scf_for]]
-  // CHECK-SCF: return %[[scf_for_tensor]]
-  return %0 : tensor<?xf32>
-}
-
-// -----
-
 // CHECK-SCF-LABEL: func @simple_scf_if(
 //  CHECK-SCF-SAME:     %[[t1:.*]]: tensor<?xf32> {linalg.inplaceable = true}, %[[c:.*]]: i1, %[[pos:.*]]: index
 func @simple_scf_if(%t1: tensor<?xf32> {linalg.inplaceable = true}, %c: i1, %pos: index, %f: f32)

mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp

@@ -102,10 +102,6 @@ struct TestComprehensiveFunctionBufferize
 
 void TestComprehensiveFunctionBufferize::runOnOperation() {
   auto options = std::make_unique<AnalysisBufferizationOptions>();
-
-  if (!allowReturnMemref)
-    options->addPostAnalysisStep(scf::assertScfForAliasingProperties);
-
   options->allowReturnMemref = allowReturnMemref;
   options->allowUnknownOps = allowUnknownOps;
   options->testAnalysisOnly = testAnalysisOnly;