[mlir][bufferize] Better error handling: Fail if ToMemrefOps are found

[springerm]

Summary

bufferization.to_memref ops are not supported in One-Shot Analysis. They often trigger a failed assertion that can be confusing. Instead, scan for to_memref ops before running the analysis and immediately abort with a proper error message.

[mamrami]

Hi!
I suspect this assertion is too strict.
I have here a piece of code that works without the assertion but fails when the assertion is there:

Command:
mlir-opt -allow-unregistered-dialect -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs=1 allow-unknown-ops=1 create-deallocs=0 function-boundary-type-conversion=identity-layout-map"

Input:

#map = affine_map<(d0, d1, d2, d3) -> (d0 * 6400 + d1 * 6400 + d2 * 80 + d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0 * 25 + d1 * 25 + d2 * 5 + d3)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0 * 92416 + d1 * 5776 + d2 * 76 + d3)>
    module {
      func.func @foo(%arg0: tensor<16xsi8>, %arg1: tensor<?x1x80x80xui8>, %arg2: tensor<16xsi8>, %arg3: tensor<16x1x5x5xsi8>) -> (tensor<?x16x76x76xsi8>) {
        %0 = "my.op"(%arg1, %arg3, %arg2, %arg0) : (tensor<?x1x80x80xui8>, tensor<16x1x5x5xsi8>, tensor<16xsi8>, tensor<16xsi8>) -> tensor<?x16x76x76xsi8>
        return %0 : tensor<?x16x76x76xsi8>
      }
      func.func @bar(%arg0: memref<16xsi8>, %arg1: memref<1x1x80x80xui8, #map>, %arg2: memref<16xsi8>, %arg3: memref<16x1x5x5xsi8, #map1>) -> (memref<1x16x76x76xsi8, #map2>) {
        %cast = memref.cast %arg1 : memref<1x1x80x80xui8, #map> to memref<?x1x80x80xui8, #map>
        %0 = bufferization.to_tensor %arg0 : memref<16xsi8>
        %1 = bufferization.to_tensor %cast : memref<?x1x80x80xui8, #map>
        %2 = bufferization.to_tensor %arg2 : memref<16xsi8>
        %3 = bufferization.to_tensor %arg3 : memref<16x1x5x5xsi8, #map1>
        %4 = call @foo(%0, %1, %2, %3) : (tensor<16xsi8>, tensor<?x1x80x80xui8>, tensor<16xsi8>, tensor<16x1x5x5xsi8>) -> tensor<?x16x76x76xsi8>
        %5 = bufferization.to_memref %4 : memref<?x16x76x76xsi8, #map2>
        %cast_0 = memref.cast %5 : memref<?x16x76x76xsi8, #map2> to memref<1x16x76x76xsi8, #map2>
        %cast_1 = memref.cast %cast_0 : memref<1x16x76x76xsi8, #map2> to memref<1x16x76x76xsi8, #map2>
        return %cast_1 : memref<1x16x76x76xsi8, #map2>
      }
    }

Output when ToMemref are not asserted:

#map = affine_map<(d0, d1, d2, d3) -> (d0 * 6400 + d1 * 6400 + d2 * 80 + d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0 * 25 + d1 * 25 + d2 * 5 + d3)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0 * 92416 + d1 * 5776 + d2 * 76 + d3)>
module {
  func.func @foo(%arg0: memref<16xsi8>, %arg1: memref<?x1x80x80xui8>, %arg2: memref<16xsi8>, %arg3: memref<16x1x5x5xsi8>) -> memref<?x16x76x76xsi8> {
    %0 = bufferization.to_tensor %arg3 : memref<16x1x5x5xsi8>
    %1 = bufferization.to_tensor %arg2 : memref<16xsi8>
    %2 = bufferization.to_tensor %arg1 : memref<?x1x80x80xui8>
    %3 = bufferization.to_tensor %arg0 : memref<16xsi8>
    %4 = "my.op"(%2, %0, %1, %3) : (tensor<?x1x80x80xui8>, tensor<16x1x5x5xsi8>, tensor<16xsi8>, tensor<16xsi8>) -> tensor<?x16x76x76xsi8>
    %5 = bufferization.to_memref %4 : memref<?x16x76x76xsi8>
    return %5 : memref<?x16x76x76xsi8>
  }
  func.func @bar(%arg0: memref<16xsi8>, %arg1: memref<1x1x80x80xui8, #map>, %arg2: memref<16xsi8>, %arg3: memref<16x1x5x5xsi8, #map1>) -> memref<1x16x76x76xsi8, #map2> {
    %cast = memref.cast %arg1 : memref<1x1x80x80xui8, #map> to memref<?x1x80x80xui8>
    %cast_0 = memref.cast %arg3 : memref<16x1x5x5xsi8, #map1> to memref<16x1x5x5xsi8>
    %0 = call @foo(%arg0, %cast, %arg2, %cast_0) : (memref<16xsi8>, memref<?x1x80x80xui8>, memref<16xsi8>, memref<16x1x5x5xsi8>) -> memref<?x16x76x76xsi8>
    %cast_1 = memref.cast %0 : memref<?x16x76x76xsi8> to memref<1x16x76x76xsi8, #map2>
    return %cast_1 : memref<1x16x76x76xsi8, #map2>
  }
}

Output when ToMemref are asserted:

<stdin>:16:14: error: to_memref ops not supported during One-Shot Analysis
        %5 = bufferization.to_memref %4 : memref<?x16x76x76xsi8, #map2>
             ^
<stdin>:16:14: note: see current operation: %6 = "bufferization.to_memref"(%5) : (tensor<?x16x76x76xsi8>) -> memref<?x16x76x76xsi8, affine_map<(d0, d1, d2, d3) -> (d0 * 92416 + d1 * 5776 + d2 * 76 + d3)>>

First we need to agree that this test is valid and should be bufferized.
If so, we may want to identify the supported cases.

Thanks in advance,
Maya

[springerm]

The input in your test case has mixed tensor/memref ops. How did you get to that state? Can you use -one-shot-bufferize="bufferize-function-boundaries" on the initial IR when everything is tensors?

We could indeed support a few cases with bufferization.to_memref and bufferization.to_tensor. But it would be pretty fragile. The main problem is that One-Shot Bufferize cannot analyze through memref code. Changing the IR just a little bit can make it so that One-Shot Bufferize has to bail. That's why we forbid to_memref in general -- to make it less surprising.

Btw, you should be able to run bufferize your IR with -one-shot-bufferize="copy-before-write". That will skip the analysis and to_memref/to_tensor should work. But it will introduce an extra copy for every op that writes to memory.

[mamrami]

The input in your test case has mixed tensor/memref ops. How did you get to that state? Can you use -one-shot-bufferize="bufferize-function-boundaries" on the initial IR when everything is tensors?

Actually the input to my all program uses memrefs (we created it). For example @bar from my reproducer.
Then I get tensor based functions (like @foo from my reproducer) and I need to bufferize them and connect them to the existing program.
So I don't have a state where everything is tensor based..

Btw, you should be able to run bufferize your IR with -one-shot-bufferize="copy-before-write". That will skip the analysis and to_memref/to_tensor should work. But it will introduce an extra copy for every op that writes to memory.

Interesting.. I just need to make sure I can avoid them later, because my program must run as fast as possible. Do you know some passes that can clean these copies?
Maybe I can change the code a bit so the "copy-before-write" will be decided per function? Because I don't mind if there will be copies in @bar. It's @foo that I want to keep clean.

I will be glad to contribute and I can also add a patch that supports more cases. I just don't know what other cases are legal.

Thanks a lot
Maya

[springerm]

Interesting.. I just need to make sure I can avoid them later, because my program must run as fast as possible. Do you know some passes that can clean these copies?

We don't have a pass for that. You would essentially have to reimplement the entire bufferization analysis. It may be possible to write a smaller removal pass that just considers simple cases.

Basically the current bufferization analyzes the IR to dectect conflict where copies are needed; i.e, insert copies when needed. One alternative that we considered was inserting copies everywhere, then having a pass that removes them. But we decided against that because analyses are generally easier on tensors than on memrefs.

Maybe I can change the code a bit so the "copy-before-write" will be decided per function? Because I don't mind if there will be copies in @bar. It's @foo that I want to keep clean.

That could work. The entry point is bufferization::bufferizeModuleOp. This function calls bufferizeOp, which uses options.copyBeforeWrite. You could try toggling that flag based on the function. Probably some smaller changes to the analysis are needed to set up the state in FuncBufferizableOpInterfaceImpl.cpp. E.g., when copy-before-write is on, aliasingFuncArgs and aliasingReturnVals should be set such that every funcBbArg is aliasing every func result. Similarly, readBbArgs should contain all bbArgs. These fields are usually filled by OneShotAnalysis, but if you run with copu-before-write=true, the analysis is skipped. But the caller of a function expects the analysis results of the callee, so you have to set it up manually.

Another problem is memory deallocation. It looks like your operation returns a new allocation after bufferization (tensor<?x16x76x76xsi8> is different from all other operand types). This is generally not supported in the current bufferization because it would require reference counting or a similar technique. What you can do is adding a "destination" tensor to your op and passing that tensor to the function. This is similar to how linalg.generic has an outs operand.

[mamrami]

Thanks @springerm
I tried it, added this patch: https://reviews.llvm.org/D142631

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

 }
 
 /// Assert that the current bufferization decisions are consistent.
 static LogicalResult
 checkAliasInfoConsistency(Operation *op, const DominanceInfo &domInfo,
                           AnalysisState &state,
                           const BufferizationAliasInfo &aliasInfo) {
   const BufferizationOptions &options = state.getOptions();
-  Operation *inconsistentOp = nullptr;
-  WalkResult walkResult = op->walk([&](Operation *op) {
-    if (auto bufferizableOp = options.dynCastBufferizableOp(op))
-      for (OpOperand &opOperand : op->getOpOperands())
-        if (opOperand.get().getType().isa<TensorType>()) {
-          if (wouldCreateReadAfterWriteInterference(
-                  opOperand, domInfo, state, aliasInfo,
-                  /*checkConsistencyOnly=*/true)) {
-            // This error can happen if certain "mustBufferizeInPlace" interface
-            // methods are implemented incorrectly, such that the IR already has
-            // a RaW conflict before making any bufferization decisions.
-            inconsistentOp = op;
-            return WalkResult::interrupt();
-          }
-        }
+
+  WalkResult walkResult = op->walk([&](BufferizableOpInterface op) {
+    // Skip ops that are not in the filter.
+    if (!options.isOpAllowed(op.getOperation()))
+      return WalkResult::advance();
+
+    // Input IR may not contain any ToMemrefOps. These are not supported because
+    // the analysis cannot follow the data flow through memrefs.
+    if (isa<ToMemrefOp>(op.getOperation())) {
+      op->emitError("to_memref ops not supported during One-Shot Analysis");
+      return WalkResult::interrupt();
+    }
+
+    for (OpOperand &opOperand : op->getOpOperands()) {
+          op->emitError("input IR has RaW conflict");
+    }
+
     return WalkResult::advance();
   });
 
-  if (walkResult.wasInterrupted())
-    return inconsistentOp->emitError("input IR has RaW conflict");
-  return success();
+  return success(!walkResult.wasInterrupted());
 }
 
 /// Annotate the IR with the result of the analysis. For testing/debugging only.
 static void
 annotateOpsWithBufferizationMarkers(Operation *op,
                                     const BufferizationAliasInfo &aliasInfo,
                                     AnalysisState &state) {
   op->walk([&](Operation *op) {

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

func.func @to_tensor_op_not_writable(%m: memref<?xf32>, %v:  vector<5xf32>,
   // Read from the tensor and return result.
   %cst = arith.constant 0.0 : f32
   %r = vector.transfer_read %w[%idx2], %cst : tensor<?xf32>, vector<10xf32>
   return %r : vector<10xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @to_memref_op_is_reading
-func.func @to_memref_op_is_reading(%t1: tensor<?xf32> {bufferization.writable = true},
-                                   %idx1: index, %idx2: index, %idx3: index,
-                                   %v1: vector<5xf32>)
-    -> (vector<5xf32>, vector<5xf32>) {
-  // Write + read to/from tensor.
-  //      CHECK: vector.transfer_write
-  // CHECK-SAME: {__inplace_operands_attr__ = ["none", "false", "none"]
-  %1 = vector.transfer_write %v1, %t1[%idx2] : vector<5xf32>, tensor<?xf32>
-  %cst = arith.constant 0.0 : f32
-  %r1 = vector.transfer_read %1[%idx3], %cst : tensor<?xf32>, vector<5xf32>
-
-  // Write + read to/from same memref.
-  %0 = bufferization.to_memref %t1 : memref<?xf32>
-  vector.transfer_write %v1, %0[%idx1] : vector<5xf32>, memref<?xf32>
-  %r2 = vector.transfer_read %0[%idx3], %cst : memref<?xf32>, vector<5xf32>
-
-  return %r1, %r2 : vector<5xf32>, vector<5xf32>
-}
-
-// -----
-
 // CHECK-LABEL: func @inner_func
 func.func @inner_func(%t: tensor<?xf32>) -> tensor<?xf32> {
   //      CHECK: return
   // CHECK-SAME: __equivalent_func_args__ = [0]
   return %t : tensor<?xf32>
 }
 
 func.func @equivalent_func_arg(%c0: index, %c10: index, %c1: index, %t0: tensor<?xf32>) -> tensor<?xf32> {

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

 
 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>) {
   // This is a RaW conflict because to_memref is an inplace write and %t1 is
   // read further down. This will likely have to change with partial
   // bufferization.
 
-  // expected-error @+1 {{input IR has RaW conflict}}
+  // expected-error @+1 {{to_memref ops not supported during One-Shot Analysis}}
   %0 = bufferization.to_memref %t1 : memref<?xf32>
 
   // Read from both.
   %cst = arith.constant 0.0 : f32
   %r1 = vector.transfer_read %t1[%idx3], %cst : tensor<?xf32>, vector<5xf32>
   %r2 = vector.transfer_read %0[%idx3], %cst : memref<?xf32>, vector<5xf32>
 
   return %r1, %r2 : vector<5xf32>, vector<5xf32>

 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 = bufferization.alloc_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 = bufferization.alloc_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
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/concatenate.mlir

func.func @concat_mix_sparse_dyn(%arg0: tensor<4x2xf64>, %arg1: tensor<4x3xf64, #MAT_C_D>, %arg2: tensor<4x4xf64, #MAT_D_C>) -> tensor<?x?xf64, #MAT_C_C> {
     return %0 : tensor<?x?xf64, #MAT_C_C>
   }
 
   func.func @dump_mat_9x4(%A: tensor<9x4xf64, #MAT_C_C>) {
     %c0 = arith.constant 0 : index
     %du = arith.constant -1.0 : f64
 
     %c = sparse_tensor.convert %A : tensor<9x4xf64, #MAT_C_C> to tensor<9x4xf64>
-    %m = bufferization.to_memref %c : memref<9x4xf64>
-    %v = vector.transfer_read %m[%c0, %c0], %du: memref<9x4xf64>, vector<9x4xf64>
+    %v = vector.transfer_read %c[%c0, %c0], %du: tensor<9x4xf64>, vector<9x4xf64>
     vector.print %v : vector<9x4xf64>
 
     %1 = sparse_tensor.values %A : tensor<9x4xf64, #MAT_C_C> to memref<?xf64>
     %2 = vector.transfer_read %1[%c0], %du: memref<?xf64>, vector<36xf64>
     vector.print %2 : vector<36xf64>
 
     return
   }
 
   func.func @dump_mat_perm_9x4(%A: tensor<9x4xf64, #MAT_C_C_P>) {
     %c0 = arith.constant 0 : index
     %du = arith.constant -1.0 : f64
 
     %c = sparse_tensor.convert %A : tensor<9x4xf64, #MAT_C_C_P> to tensor<9x4xf64>
-    %m = bufferization.to_memref %c : memref<9x4xf64>
-    %v = vector.transfer_read %m[%c0, %c0], %du: memref<9x4xf64>, vector<9x4xf64>
+    %v = vector.transfer_read %c[%c0, %c0], %du: tensor<9x4xf64>, vector<9x4xf64>
     vector.print %v : vector<9x4xf64>
 
     %1 = sparse_tensor.values %A : tensor<9x4xf64, #MAT_C_C_P> to memref<?xf64>
     %2 = vector.transfer_read %1[%c0], %du: memref<?xf64>, vector<36xf64>
     vector.print %2 : vector<36xf64>
 
     return
   }
 
   func.func @dump_mat_dense_9x4(%A: tensor<9x4xf64>) {
     %c0 = arith.constant 0 : index
     %du = arith.constant -1.0 : f64
 
-    %m = bufferization.to_memref %A : memref<9x4xf64>
-    %v = vector.transfer_read %m[%c0, %c0], %du: memref<9x4xf64>, vector<9x4xf64>
+    %v = vector.transfer_read %A[%c0, %c0], %du: tensor<9x4xf64>, vector<9x4xf64>
     vector.print %v : vector<9x4xf64>
 
     return
   }
 
   func.func @dump_mat_4x9(%A: tensor<4x9xf64, #MAT_C_C>) {
     %c0 = arith.constant 0 : index
     %du = arith.constant -1.0 : f64
 
     %c = sparse_tensor.convert %A : tensor<4x9xf64, #MAT_C_C> to tensor<4x9xf64>
-    %m = bufferization.to_memref %c : memref<4x9xf64>
-    %v = vector.transfer_read %m[%c0, %c0], %du: memref<4x9xf64>, vector<4x9xf64>
+    %v = vector.transfer_read %c[%c0, %c0], %du: tensor<4x9xf64>, vector<4x9xf64>
     vector.print %v : vector<4x9xf64>
 
     %1 = sparse_tensor.values %A : tensor<4x9xf64, #MAT_C_C> to memref<?xf64>
     %2 = vector.transfer_read %1[%c0], %du: memref<?xf64>, vector<36xf64>
     vector.print %2 : vector<36xf64>
 
     return
   }
 
   func.func @dump_mat_dyn(%A: tensor<?x?xf64, #MAT_C_C>) {
     %c0 = arith.constant 0 : index
     %du = arith.constant -1.0 : f64
 
     %c = sparse_tensor.convert %A : tensor<?x?xf64, #MAT_C_C> to tensor<?x?xf64>
-    %m = bufferization.to_memref %c : memref<?x?xf64>
-    %v = vector.transfer_read %m[%c0, %c0], %du: memref<?x?xf64>, vector<4x9xf64>
+    %v = vector.transfer_read %c[%c0, %c0], %du: tensor<?x?xf64>, vector<4x9xf64>
     vector.print %v : vector<4x9xf64>
 
     %1 = sparse_tensor.values %A : tensor<?x?xf64, #MAT_C_C> to memref<?xf64>
     %2 = vector.transfer_read %1[%c0], %du: memref<?xf64>, vector<36xf64>
     vector.print %2 : vector<36xf64>
 
     return
   }
 
   func.func @dump_mat_perm_4x9(%A: tensor<4x9xf64, #MAT_C_C_P>) {
     %c0 = arith.constant 0 : index
     %du = arith.constant -1.0 : f64
 
     %c = sparse_tensor.convert %A : tensor<4x9xf64, #MAT_C_C_P> to tensor<4x9xf64>
-    %m = bufferization.to_memref %c : memref<4x9xf64>
-    %v = vector.transfer_read %m[%c0, %c0], %du: memref<4x9xf64>, vector<4x9xf64>
+    %v = vector.transfer_read %c[%c0, %c0], %du: tensor<4x9xf64>, vector<4x9xf64>
     vector.print %v : vector<4x9xf64>
 
     %1 = sparse_tensor.values %A : tensor<4x9xf64, #MAT_C_C_P> to memref<?xf64>
     %2 = vector.transfer_read %1[%c0], %du: memref<?xf64>, vector<36xf64>
     vector.print %2 : vector<36xf64>
 
     return
   }
 
   func.func @dump_mat_dense_4x9(%A: tensor<4x9xf64>) {
     %c0 = arith.constant 0 : index
     %du = arith.constant -1.0 : f64
 
-    %m = bufferization.to_memref %A : memref<4x9xf64>
-    %v = vector.transfer_read %m[%c0, %c0], %du: memref<4x9xf64>, vector<4x9xf64>
+    %v = vector.transfer_read %A[%c0, %c0], %du: tensor<4x9xf64>, vector<4x9xf64>
     vector.print %v : vector<4x9xf64>
 
     return
   }
 
   // Driver method to call and verify kernels.
   func.func @entry() {
     %m42 = arith.constant dense<