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

[springerm]

Summary

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.

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

let methods = [
         /*methodName=*/"isNotConflicting",
         /*args=*/(ins "OpOperand *":$uRead,
                       "OpOperand *":$uWrite,
                       "const BufferizationState &":$state),
         /*methodBody=*/"",
         /*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();
+        }]
       >
   ];
 
   let extraClassDeclaration = [{
     /// Return `true` if the given OpOperand creates an alias but does neither
     /// read nor write. This implies that `bufferizesToMemoryRead` and
     /// `bufferizesToMemoryWrite` must return `false`. This method will never
     /// be called on OpOperands that do not have a tensor type.

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

 class DialectRegistry;
 
 namespace bufferization {
 class BufferizationState;
 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
 
 #endif // MLIR_DIALECT_SCF_BUFFERIZABLEOPINTERFACEIMPL_H

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

LogicalResult bufferization::analyzeOp(Operation *op,
   }
 
   if (!options.allowReturnMemref) {
     SmallVector<Operation *> newOps;
     if (failed(assertDestinationPassingStyle(op, state, aliasInfo, newOps)))
       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);
 
   return success();
 }
 
 LogicalResult bufferization::runOneShotBufferize(
     Operation *op, std::unique_ptr<AnalysisBufferizationOptions> options) {
   AnalysisBufferizationState state(op, *options);
   if (failed(analyzeOp(op, state)))
     return failure();
   if (options->testAnalysisOnly)
     return success();
   return bufferizeOp(op, state);
 }

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

if (!options) {
     if (initTensorElimination) {
       opt.addPostAnalysisStep(
           linalg_ext::insertSliceAnchoredInitTensorEliminationStep);
     }
   } else {
     opt = *options;
   }
 
-  // Only certain scf.for ops are supported by the analysis.
-  opt.addPostAnalysisStep(scf::assertScfForAliasingProperties);
-
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 
   if (failed(runModuleBufferize(moduleOp, opt))) {
     signalPassFailure();
     return;
   }
 

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

SmallVector<Value> yieldValues =
         });
     yieldOp.getResultsMutable().assign(yieldValues);
 
     // Replace loop results.
     replaceOpWithBufferizedValues(rewriter, op, newForOp->getResults());
 
     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
 /// this is for analysis only.
 struct YieldOpInterface
     : public BufferizableOpInterface::ExternalModel<YieldOpInterface,
                                                     scf::YieldOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     return success();
   }
 };
 
 } // namespace
 } // 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>();
   registry.addOpInterface<ForOp, ForOpInterface>();
   registry.addOpInterface<IfOp, IfOpInterface>();
   registry.addOpInterface<YieldOp, YieldOpInterface>();
   registry
       .addOpInterface<ParallelOp, AllocationHoistingBarrierOnly<ParallelOp>>();
 }

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

 }
 
 // -----
 
 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>)
   {
     %ttA = tensor.insert_slice %C into %tA[0][4][1] : tensor<4xf32> into tensor<?xf32>
     %ttB = tensor.insert_slice %C into %tB[0][4][1] : tensor<4xf32> into tensor<?xf32>
 
     // Throw a wrench in the system by swapping yielded values: this result in a
     // ping-pong of values at each iteration on which we currently want to fail.
 
     // expected-error @+1 {{Yield operand #0 does not bufferize to a buffer that is aliasing}}
     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
 }
 
 // -----
 
 func private @fun_with_side_effects(%A: tensor<?xf32> {linalg.inplaceable = true})
 
 func @foo(%A: tensor<?xf32> {linalg.inplaceable = true}) -> (tensor<?xf32>) {
   call @fun_with_side_effects(%A) : (tensor<?xf32>) -> ()

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

func @simple_tensor_test(%t1 : tensor<?xf32>, %f : f32) -> tensor<?xf32> {
   // CHECK-TENSOR: memref.store %{{.*}}, %[[alloc]]
   %0 = tensor.insert %f into %t1[%c0] : tensor<?xf32>
   // CHECK-TENSOR: return %[[casted_tensor]]
   return %0 : 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)
     -> (tensor<?xf32>, index) {
   // CHECK-SCF: %[[r:.*]] = scf.if %[[c]] -> (memref<?xf32, #{{.*}}>) {
   %r1, %r2 = scf.if %c -> (tensor<?xf32>, index) {
     // CHECK-SCF: %[[t1_memref:.*]] = bufferization.to_memref %[[t1]]
     // CHECK-SCF: scf.yield %[[t1_memref]]

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

Option<bool> createDeallocs{
       *this, "create-deallocs",
       llvm::cl::desc("Specify if buffers should be deallocated"),
       llvm::cl::init(true)};
 };
 } // namespace
 
 void TestComprehensiveFunctionBufferize::runOnOperation() {
   auto options = std::make_unique<AnalysisBufferizationOptions>();
-
-  if (!allowReturnMemref)
-    options->addPostAnalysisStep(scf::assertScfForAliasingProperties);
-
   options->allowReturnMemref = allowReturnMemref;
   options->allowUnknownOps = allowUnknownOps;
   options->testAnalysisOnly = testAnalysisOnly;
   options->analysisFuzzerSeed = analysisFuzzerSeed;
   options->fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
   options->createDeallocs = createDeallocs;
 
   if (dialectFilter.hasValue()) {