[mlir][bufferize] Improve handling of to_memref ops during analysis

[springerm]

Summary

to_memref is usually not allowed in the input IR when running One-Shot Bufferize.

• Allow to_memref if there are no other bufferizable ops.
• Better error message when to_memref is found.

[nicolasvasilache]

Can you document the reason for this special casing?
Wouldn't this be better as a helper function that a client could call explicitly to know that bufferization won't do anything for them and they can avoid calling it in the first place ?
I am very concerned of anything that gives the impression to_memref and his twin are ok to have in this pass.

[springerm]

This is a pretty old revision. We don't need this anymore.

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

if (analysisFuzzerSeed) {
     // This is a fuzzer. For testing purposes only. Randomize the order in which
     // operations are analyzed. The bufferization quality is likely worse, but
     // we want to make sure that no assertions are triggered anywhere.
     std::mt19937 g(analysisFuzzerSeed);
     llvm::shuffle(ops.begin(), ops.end(), g);
   }
 
   // Walk ops in reverse for better interference analysis.
-  for (Operation *op : reverse(ops))
+  for (Operation *op : reverse(ops)) {
+    if (isa<ToMemrefOp>(op))
+      return op->emitError(
+          "analysis through bufferization.to_memref ops is not supported");
+
     for (OpOperand &opOperand : op->getOpOperands())
       if (opOperand.get().getType().isa<TensorType>())
         if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op))
           if (failed(bufferizableInPlaceAnalysisImpl(opOperand, aliasInfo,
                                                      state, domInfo)))
             return failure();
+  }
 
   return success();
 }
 
 /// Return true if the given op has a tensor result or a tensor operand.
 static bool hasTensorSemantics(Operation *op) {
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
   return hasTensorResult || hasTensorOperand;
 }
 
 /// Analyze all ops that are contained in `op`.
 static LogicalResult inPlaceAnalysis(Operation *op,
                                      BufferizationAliasInfo &aliasInfo,
                                      AnalysisState &state,
                                      const DominanceInfo &domInfo,
                                      unsigned analysisFuzzerSeed = 0) {
   // Collect ops so we can build our own reverse traversal.
   SmallVector<Operation *> ops;
   op->walk([&](Operation *op) {
     // No tensors => no buffers.
     if (!hasTensorSemantics(op))
       return;
+    // Skip ops that are not bufferizable;
+    if (!state.getOptions().dynCastBufferizableOp(op))
+      return;
     ops.push_back(op);
   });
 
+  // Shortcut: Nothing to do if there are only ToMemrefOps/ToTensorOps.
+  if (llvm::all_of(
+          ops, [&](Operation *op) { return isa<ToMemrefOp, ToTensorOp>(op); }))
+    return success();
+
   return inPlaceAnalysis(ops, aliasInfo, state, domInfo, analysisFuzzerSeed);
 }
 
 /// Analyze equivalence of tied OpResult/OpOperand pairs of the given ops.
 static void equivalenceAnalysis(SmallVector<Operation *> &ops,
                                 BufferizationAliasInfo &aliasInfo,
                                 AnalysisState &state) {
   for (Operation *op : ops)
static void equivalenceAnalysis(Operation *op,
                                 BufferizationAliasInfo &aliasInfo,
                                 AnalysisState &state) {
   // Traverse ops in PostOrder: Nested ops first, then enclosing ops.
   SmallVector<Operation *> ops;
   op->walk<WalkOrder::PostOrder>([&](Operation *op) {
     // No tensors => no buffers.
     if (none_of(op->getResultTypes(), isaTensor))
       return;
+    // Skip ops that are not bufferizable;
+    if (!state.getOptions().dynCastBufferizableOp(op))
+      return;
     ops.push_back(op);
   });
 
+  // Shortcut: Nothing to do if there are only ToMemrefOps/ToTensorOps.
+  if (llvm::all_of(
+          ops, [&](Operation *op) { return isa<ToMemrefOp, ToTensorOp>(op); }))
+    return;
+
   equivalenceAnalysis(ops, aliasInfo, state);
 }
 
 /// Assert that the current bufferization decisions are consistent.
 static LogicalResult
 checkAliasInfoConsistency(Operation *op, const DominanceInfo &domInfo,
                           AnalysisState &state,
                           const BufferizationAliasInfo &aliasInfo) {

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

// CHECK-SCF: } else {
     // CHECK-SCF: scf.yield %[[insert_memref]]
     scf.yield %1, %pos : tensor<?xf32>, index
   }
 
   // CHECK-SCF: %[[r_tensor:.*]] = bufferization.to_tensor %[[r]]
   // CHECK-SCF: return %[[r_tensor]], %[[pos]]
   return %r1, %r2 : tensor<?xf32>, index
 }
+
+// -----
+
+// to_memref ops are usually not supported by the analysis, but if there are no
+// other bufferizable ops, we allow them.
+
+// CHECK-LABEL: func @only_unbufferizable_ops() {
+//       CHECK:   test.dummy_op
+//       CHECK:   bufferization.to_memref
+//       CHECK:   test.another_dummy_op
+//       CHECK: }
+func.func @only_unbufferizable_ops() {
+  %0 = "test.dummy_op"() : () -> (tensor<?xf32>)
+  %1 = bufferization.to_memref %0 : memref<?xf32>
+  "test.another_dummy_op"(%1) : (memref<?xf32>) -> ()
+}

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

 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 {
   %0 = scf.execute_region -> tensor<?xf32> {
     %1 = linalg.init_tensor [%idx] : tensor<?xf32>
     // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}}
     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 {
   %1 = linalg.init_tensor [%idx] : tensor<?xf32>
 
   %0 = scf.execute_region -> tensor<?xf32> {
     // This YieldOp is in destination-passing style, 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
 }
+
+// -----
+
+func.func @to_memref_not_supported(%idx: index, %v: i32) -> memref<2x3xi32> {
+  %cst = arith.constant dense<0> : tensor<2x3xi32>
+  // expected-error @+1 {{analysis through bufferization.to_memref ops is not supported}}
+  %0 = bufferization.to_memref %cst : memref<2x3xi32>
+  memref.store %v, %0[%idx, %idx] : memref<2x3xi32>
+  return %0 : memref<2x3xi32>
+}
+