[mlir][Bufferization] Extend the folding of to_tensor(to_memref)

[mamrami]

Summary

The folding was allowed only if to_tensor appeared right after to_memref.
This change allows the folding as long as there are no interleaved users
of the result memref.

[mamrami]

Change operation in lit test.

[mamrami]

lit test change

[springerm]

This would fold in cases such as this one right? That's not safe.

%alias = memref.subview %arg0
%0 = bufferization.to_memref %arg0 : memref<?xf32>
%1 = "use"(%alias) : (memref<?xf32>) -> memref<?xf32>
%2 = bufferization.to_tensor %0 : memref<?xf32>

My suggestion would be to remove to_tensor and to_memref entirely and using unrealized_conversion_cast instead. We can discuss on Discourse.

[mehdi_amini]

In D142195#4068203, @springerm wrote:

This would fold in cases such as this one right? That's not safe.

%alias = memref.subview %arg0
%0 = bufferization.to_memref %arg0 : memref<?xf32>
%1 = "use"(%alias) : (memref<?xf32>) -> memref<?xf32>
%2 = bufferization.to_tensor %0 : memref<?xf32>

Isn't this covered by the isBeforeInBlock condition?

[mehdi_amini]

mlir/test/Dialect/Bufferization/canonicalize.mlir

48

We should also have the positive test case for the ordering:

  %0 = bufferization.to_memref %arg0 : memref<?xf32>
  %2 = bufferization.to_tensor %0 : memref<?xf32>
  %1 = "use"(%0) : (memref<?xf32>) -> memref<?xf32>
  return %2 : tensor<?xf32>
}

This should fold right?

[springerm]

In D142195#4068240, @mehdi_amini wrote:

In D142195#4068203, @springerm wrote:

This would fold in cases such as this one right? That's not safe.

%alias = memref.subview %arg0
%0 = bufferization.to_memref %arg0 : memref<?xf32>
%1 = "use"(%alias) : (memref<?xf32>) -> memref<?xf32>
%2 = bufferization.to_tensor %0 : memref<?xf32>

Isn't this covered by the isBeforeInBlock condition?

Ahh yes, the IR that I wrote is not even valid. The case that I was thinking of cannot happen. (Was thinking of to_memref(to_tensor) folding....)

[nicolasvasilache]

to_tensor and to_memref are fundamentally broken abstractions with tricky constraints specified in doc text that only amount to wishful thinking.
They are part of the difficulties I had with using bufferization in a sane way a while back (https://discourse.llvm.org/t/properly-using-bufferization-related-passes/2913) until I decided things were too broken and started a new bufferization effort (https://discourse.llvm.org/t/rfc-linalg-on-tensors-update-and-comprehensive-bufferization-rfc/3373).

They have survived until now for purposes of "compatibility" with some downstream uses and "composability".
I believe it is time to cut the cord.

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

569

you need to consider all possible aliases of the memref, this requires a much more serious analysis

[springerm]

In D142195#4068251, @nicolasvasilache wrote:

I believe it is time to cut the cord.

It would indeed be better to remove them entirely (as a cleanup), as we no longer have a need for them. We moved to a bufferization that bufferizes to entire IR in one go. @mamrami Can you describe your use case in Discourse, so we can discuss alternatives?

[mamrami]

@springerm - my case is a bit different/special. I have an IR with functions that have memref based signature.
Some of them have implementation - and it is a tensor based implementation, meaning I have to_tensor/to_memrefs on args/results.
I had to inline some of them to one function that will be pure tensor based function. On this function I run OneShotBufferize.
After inlining I had the intermediate to_tensor(to_memref), and I expected them to fold.
That's how I got to all the wondering about the folding.
I understand now that the general case is much more complicated than my case.
So I decided to manually move the to_tensor right after their to_memref - because in my case I know it is not changing the meaning of the IR.

[springerm]

In D142195#4072693, @mamrami wrote:

my case is a bit different/special. I have an IR with functions that have memref based signature.
Some of them have implementation - and it is a tensor based implementation, meaning I have to_tensor/to_memrefs on args/results.

Is it possible to give them a tensor signature? Then you could bufferize the entire thing with One-Shot Bufferize.

[mamrami]

In D142195#4073141, @springerm wrote:

Is it possible to give them a tensor signature? Then you could bufferize the entire thing with One-Shot Bufferize.

It will require redesigning our whole flow 😐

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

LogicalResult DeallocTensorOp::bufferize(RewriterBase &rewriter,
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ToTensorOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ToTensorOp::fold(FoldAdaptor) {
-  if (auto toMemref = getMemref().getDefiningOp<ToMemrefOp>())
-    // Approximate alias analysis by conservatively folding only when no there
-    // is no interleaved operation.
+  if (auto toMemref = getMemref().getDefiningOp<ToMemrefOp>()) {
+    // Approximate alias analysis by conservatively folding only when know there
+    // is no interleaved user of the result memref.
+    bool canResultMemrefMutate =
+        llvm::any_of(toMemref->getUsers(), [&](Operation *op) {
+          if (op->getBlock() == this->getOperation()->getBlock())
+            return op->isBeforeInBlock(*this);
+          return true;
+        });
     if (toMemref->getBlock() == this->getOperation()->getBlock() &&
-        toMemref->getNextNode() == this->getOperation())
+        !canResultMemrefMutate)
       return toMemref.getTensor();
+  }
   return {};
 }
 
 namespace {
 struct DimOfToTensorFolder : public OpRewritePattern<tensor::DimOp> {
   using OpRewritePattern<tensor::DimOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tensor::DimOp dimOp,

mlir/test/Dialect/Bufferization/canonicalize.mlir

func.func @tensor_load_of_buffer_cast(%arg0: tensor<?xf32>) -> tensor<?xf32> {
   %1 = bufferization.to_tensor %0 : memref<?xf32>
   return %1 : tensor<?xf32>
 }
 // CHECK-SAME:   %[[TENSOR:.*]]: tensor<?xf32>) -> tensor<?xf32> {
 // CHECK: return %[[TENSOR]]
 
 // -----
 
+// Folding of to_tensor(to_memref(t)) -> t is allowed when the buffer is not mutated before the to_tensor.
+// CHECK-LABEL:   func.func @tensor_load_of_unmutated_buffer(
+func.func @tensor_load_of_unmutated_buffer(%arg0: tensor<?xf32>) -> (tensor<?xf32>, tensor<?xf32>) {
+  %0 = bufferization.to_memref %arg0 : memref<?xf32>
+  %1 = "use"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
+  %2 = bufferization.to_tensor %0 : memref<?xf32>
+  return %1, %2 : tensor<?xf32>, tensor<?xf32>
+}
+
+// CHECK-SAME:   %[[arg0:.*]]: tensor<?xf32>) -> (tensor<?xf32>, tensor<?xf32>) {
+// CHECK: %[[VAL_0:.*]] = "use"(%[[arg0]]) : (tensor<?xf32>) -> tensor<?xf32>
+// CHECK: return %[[VAL_0]], %[[arg0]] : tensor<?xf32>, tensor<?xf32>
+
+// -----
+
+// Folding of to_tensor(to_memref(t)) -> t is not allowed since t may mutate before the to_tensor.
+// CHECK-LABEL:   func.func @tensor_load_of_mutated_buffer(
+func.func @tensor_load_of_mutated_buffer(%arg0: tensor<?xf32>) -> (tensor<?xf32>) {
+  %0 = bufferization.to_memref %arg0 : memref<?xf32>
+  %1 = "use"(%0) : (memref<?xf32>) -> memref<?xf32>
+  %2 = bufferization.to_tensor %0 : memref<?xf32>
+  return %2 : tensor<?xf32>
+}
+
+// CHECK-SAME:  %[[VAL_0:.*]]: tensor<?xf32>) -> tensor<?xf32> {
+// CHECK: %[[VAL_1:.*]] = bufferization.to_memref %[[VAL_0]] : memref<?xf32>
+// CHECK: %[[VAL_2:.*]] = "use"(%[[VAL_1]]) : (memref<?xf32>) -> memref<?xf32>
+// CHECK: %[[VAL_3:.*]] = bufferization.to_tensor %[[VAL_1]] : memref<?xf32>
+// CHECK: return %[[VAL_3]] : tensor<?xf32>
+
+// -----
+
 // Basic folding of to_memref(to_tensor(m)) -> m
 // CHECK-LABEL: func @buffer_cast_of_tensor_load(
 func.func @buffer_cast_of_tensor_load(%arg0: memref<?xf32>) -> memref<?xf32> {
   %0 = bufferization.to_tensor %arg0 : memref<?xf32>
   %1 = bufferization.to_memref %0 : memref<?xf32>
   return %1 : memref<?xf32>
 }
 // CHECK-SAME:   %[[MEMREF:.*]]: memref<?xf32>) -> memref<?xf32> {