[mlir] Remove unsafe tensor_to_memref -> load canonicalization pattern.

[mravishankar]

Summary

Folding tensor_to_memref -> load after bufferization results in
bufferization being undone. It can potentially violate dependencies,
if there was a store to the memref, followed by a load. This pattern
would return the value before the store. In general this pattern needs
to be applied selectively and adding it to canonicalizers would result
in subtle correctness errors.

[nicolasvasilache]

I did something similar for tensor_load(tensor_to_memref()).
The conservative way I came up with was to only fold when loadOp == tensorToMemRefOp->getNextNode(), this should also work here.
In my case it does provide useful canonicalizations.

It is also a good place to add a clear "// TODO: relax this once we have proper alias analysis support".

Can we just do the same ehre and avoid killing the pattern altogether?
This will also make it clear to whoever looking into this next that they should not reintroduce the pattern.

mlir/test/Dialect/Standard/canonicalize.mlir

90

this particular case is actually safe :)

[silvas]

LGTM. Retaining this in the "next node" case doesn't fix the core problem of undoing bufferization which I think was the issue you were hitting (and is my main objection to this pattern).

For context, the safety issue is technically arguable given the current upstream state (not saying this is great) -- upstream currently has a (not yet formally documented) quasi-SSA memref step during bufferization, which arises naturally from the opaque nature of tensor_to_memref (cannot assume writable) as it is inserted by the type conversion framework, which makes this pattern safe (but it still "undoes bufferization", so I still want to delete the pattern even given that assumption :) ). There has been an attempt to document the property D95522, but nothing committed.

It sounds like Nicolas is experimenting with a different approach, which doesn't have the same invariants, and I'm fine with that -- I still haven't seen a coherent solution to the inplace update problem with the upstream approach (the most is my hand waving that a sufficiently smart "regalloc" can do it), so any solution in this space is welcome, and I'm happy to supercede the current bufferization stuff by what Nicolas comes up with (though I suspect that the current upstream stuff will still be relevant to some flows, such as those that represent shapes as small tensors/memrefs -- will be "fun" to disentangle that). In chats with Nicolas he has managed to convince me that the "quasi-SSA" memref situation doesn't feel long-term, and I'm working with him to bring a new solution upstream.

[mravishankar]

Abandoning this for now (see https://llvm.discourse.group/t/tensor-to-memref-and-its-folding/3033)

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let extraClassDeclaration = [{
     void setMemRef(Value value) { setOperand(0, value); }
     MemRefType getMemRefType() {
       return getMemRef().getType().cast<MemRefType>();
     }
 
     operand_range getIndices() { return {operand_begin() + 1, operand_end()}; }
   }];
 
-  let hasCanonicalizer = 1;
   let hasFolder = 1;
 
   let assemblyFormat = "$memref `[` $indices `]` attr-dict `:` type($memref)";
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefCastOp
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 
 OpFoldResult LoadOp::fold(ArrayRef<Attribute> cstOperands) {
   /// load(memrefcast) -> load
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
   return OpFoldResult();
 }
 
-namespace {
-/// Fold a load on a tensor_to_memref operation into an tensor.extract on the
-/// corresponding tensor.
-struct LoadOfTensorToMemref : public OpRewritePattern<LoadOp> {
-  using OpRewritePattern<LoadOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(LoadOp load,
-                                PatternRewriter &rewriter) const override {
-    auto tensorToMemref = load.memref().getDefiningOp<TensorToMemrefOp>();
-    if (!tensorToMemref)
-      return failure();
-
-    rewriter.replaceOpWithNewOp<tensor::ExtractOp>(
-        load, tensorToMemref.tensor(), load.indices());
-    return success();
-  }
-};
-} // end anonymous namespace.
-
-void LoadOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
-                                         MLIRContext *context) {
-  results.insert<LoadOfTensorToMemref>(context);
-}
-
 //===----------------------------------------------------------------------===//
 // MemRefCastOp
 //===----------------------------------------------------------------------===//
 
 Value MemRefCastOp::getViewSource() { return source(); }
 
 bool MemRefCastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   if (inputs.size() != 1 || outputs.size() != 1)

mlir/test/Dialect/Standard/canonicalize.mlir

func @dim_of_tensor_load(%arg0: memref<?xf32>) -> index {
   %c0 = constant 0 : index
   %0 = tensor_load %arg0 : memref<?xf32>
   %1 = dim %0, %c0 : tensor<?xf32>
   return %1 : index
 }
 
 // -----
 
-// Test case: Folding of load(tensor_to_memref(%v, %idxs))
-//            -> tensor.extract(%v, %idx)
-// CHECK-LABEL: func @load_from_tensor_to_memref(
-//  CHECK-SAME:     %[[IDX0:[0-9a-z]+]]: index, %[[IDX1:[0-9a-z]+]]: index
-//  CHECK-SAME:     %[[TENSOR:[0-9a-z]+]]: tensor<?x?xf32>
-//       CHECK:   %[[RES:.*]] = tensor.extract %[[TENSOR]][%[[IDX0]], %[[IDX1]]]
-//   CHECK-NOT:   load
-//       CHECK:   return %[[RES]] : f32
-func @load_from_tensor_to_memref(%arg0: index, %arg1: index, %arg2: tensor<?x?xf32>) -> f32 {
-  %0 = tensor_to_memref %arg2 : memref<?x?xf32>
-  %1 = load %0[%arg0, %arg1] : memref<?x?xf32>
-  return %1 : f32
-}
-
-// -----
-
 // Test case: Folding of dim(tensor.generate %idx) -> %idx
 // CHECK-LABEL: func @dim_of_tensor.generate(
 //  CHECK-SAME:     %[[IDX0:[0-9a-z]+]]: index, %[[IDX1:[0-9a-z]+]]: index
 //   CHECK-NOT:   dim
 //       CHECK:   return %[[IDX1]] : index
 func @dim_of_tensor.generate(%arg0: index, %arg1: index) -> index {
   %c3 = constant 3 : index
   %0 = tensor.generate %arg0, %arg1 {