[mlir][linalg][bufferize] Allow unbufferizable ops in input

[springerm]

Summary

Allow ops that are not bufferizable in the input IR. (Deactivated by default.)

bufferization::ToMemrefOp and bufferization::ToTensorOp are generated at the bufferization boundaries.

Depends On D114661

[springerm]

update

[nicolasvasilache]

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

103

In one of the previous CLs, please write heavy documentation that to_memref represents the "future location in memory" to which the tensor will be bufferized.
This may connect nicely to the analysis we do in comprehensive bufferize which thinks in term of these "futures".
What do we know at a particular point in time re these futures for to_memref ?
It depends, often almost nothing and they may alias with the world.

115

please use m1 and print memref in the test so there is less ambiguity on what we expect.

118

Similar comment for this one but it seems easier:
this is a "load" of the whole tensor from whatever "future memory location" this will be bufferized to.
This also aliases with the world as a read until more information is available in the comprehensive bufferize analysis.

137

Can we add a comment here that this must alloc and copy?
This way if anyone ever tries to change the behavior they walk into a miniefield that they have to strongly justify.

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h

struct BufferizationOptions {
 
   /// Helper functions for allocation, deallocation, memory copying.
   std::unique_ptr<AllocationCallbacks> allocationFns;
 
   /// Specifies whether returning newly allocated memrefs should be allowed.
   /// Otherwise, a pass failure is triggered.
   bool allowReturnMemref = false;
 
+  /// Specifies whether not bufferizable ops are allowed in the input. If so,
+  /// bufferization.to_memref and bufferization.to_tensor ops are inserted at
+  /// the boundaries.
+  bool allowUnknownOps = false;
+
   /// Seed for the analysis fuzzer. If set to `0`, the fuzzer is deactivated.
   /// Should be used only with `testAnalysisOnly = true`.
   unsigned analysisFuzzerSeed = 0;
 
   /// If set to `true`, does not modify the IR apart from adding attributes (for
   /// checking the results of the analysis) and post analysis steps.
   bool testAnalysisOnly = false;
 
struct BufferizationState {
   /// Map a value to another value.
   void mapValue(Value from, Value to);
 
   /// Map a tensor value to a memref buffer.
   void mapBuffer(Value tensor, Value buffer);
 
   /// Lookup the memref buffer that is associated to the given tensor value.
   /// Asserts if no buffer is associated.
-  Value lookupBuffer(Value tensor) const;
+  Value lookupBuffer(Value tensor);
 
   /// Lookup the value that is associated to the given value. Asserts if no
   /// value is associated.
   Value lookupValue(Value value) const;
 
   /// Return `true` if the given value is mapped.
   bool isMapped(Value value) const;
 
struct AllocationHoistingBarrierOnly
 
   bool isWritable(Operation *op, Value value) const { return false; }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
     if (any_of(op->getOperandTypes(), isaTensor) ||
         any_of(op->getResultTypes(), isaTensor))
+      if (!state.options.allowUnknownOps)
         return op->emitError() << "unsupported op with tensors";
+
+    for (Region &region : op->getRegions())
+      if (failed(comprehensive_bufferize::bufferize(&region, state)))
+        return failure();
+
     return success();
   }
 
   bool isAllocationHoistingBarrier(Operation *op) const { return true; }
 };
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_

mlir/include/mlir/Dialect/Linalg/Passes.td

def LinalgComprehensiveModuleBufferize :
   }];
   let options = [
     Option<"testAnalysisOnly", "test-analysis-only", "bool",
             /*default=*/"false",
            "Only runs inplaceability analysis (for testing purposes only)">,
     Option<"allowReturnMemref", "allow-return-memref", "bool",
             /*default=*/"false",
            "Allows the return of memrefs (for testing purposes only)">,
+    Option<"allowUnknownOps", "allow-unknown-ops", "bool",
+           /*default=*/"false",
+           "Allows unknown (not bufferizable) ops in the input IR.">,
     Option<"useAlloca", "use-alloca", "bool",
            /*default=*/"false",
            "Use stack allocations for memrefs (for testing purposes only)">,
     Option<"analysisFuzzerSeed", "analysis-fuzzer-seed", "unsigned",
            /*default=*/"0",
            "Analyze ops in random order with a given seed (fuzzer)">
   ];
   let constructor = "mlir::createLinalgComprehensiveModuleBufferizePass()";

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp

BufferizationAliasInfo::getAliases(Value v) const {
   return BufferizationAliasInfo::EquivalenceClassRangeType(
       aliasInfo.member_begin(it), aliasInfo.member_end());
 }
 
 //===----------------------------------------------------------------------===//
 // Helper functions for BufferizableOpInterface
 //===----------------------------------------------------------------------===//
 
+static void setInsertionPointAfter(OpBuilder &b, Value value) {
+  if (auto bbArg = value.dyn_cast<BlockArgument>()) {
+    b.setInsertionPointToStart(bbArg.getOwner());
+  } else {
+    b.setInsertionPointAfter(value.getDefiningOp());
+  }
+}
+
 /// Determine which OpOperand* will alias with `result` if the op is bufferized
 /// in place. Return an empty vector if the op is not bufferizable.
 SmallVector<OpOperand *>
 mlir::linalg::comprehensive_bufferize::getAliasingOpOperand(OpResult result) {
   if (Operation *op = result.getDefiningOp())
     if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
       return bufferizableOp.getAliasingOpOperand(result);
   return {};
Value mlir::linalg::comprehensive_bufferize::getResultBuffer(
   OpOperand *opOperand = aliasingOperands.front();
   Value operand = opOperand->get();
   Value operandBuffer = state.lookupBuffer(operand);
   // Make sure that all OpOperands are the same buffer. If this is not the case,
   // we would have to materialize a memref value.
   // TODO: Should be looking for checking for "equivalent buffers" instead of
   // operator== here, but equivalent buffers for scf.if yield values are not
   // set up yet.
-  if (!llvm::all_of(aliasingOperands, [&](OpOperand *o) {
+  if (aliasingOperands.size() > 1 &&
+      !llvm::all_of(aliasingOperands, [&](OpOperand *o) {
         return state.lookupBuffer(o->get()) == operandBuffer;
       })) {
     op->emitError("result buffer is ambiguous");
     return Value();
   }
 
   // If bufferizing out-of-place, allocate a new buffer.
   if (!state.aliasInfo.isInPlace(result)) {
     // Ops with multiple aliasing operands can currently not bufferize
     // out-of-place.
     assert(
         aliasingOperands.size() == 1 &&
         "ops with multiple aliasing OpOperands cannot bufferize out-of-place");
     Location loc = op->getLoc();
     // Move insertion point right after `operandBuffer`. That is where the
     // allocation should be inserted (in the absence of allocation hoisting).
-    if (auto bbArg = operandBuffer.dyn_cast<BlockArgument>()) {
-      b.setInsertionPointToStart(bbArg.getOwner());
-    } else {
-      b.setInsertionPointAfter(operandBuffer.getDefiningOp());
-    }
+    setInsertionPointAfter(b, operandBuffer);
     // Allocate the result buffer.
     Value resultBuffer = state.createAllocDeallocFn(b, loc, operandBuffer);
     bool skipCopy = false;
     // Do not copy if the last preceding write of `operand` is an op that does
     // not write (skipping ops that merely create aliases). E.g., InitTensorOp.
     // Note: If `findLastPrecedingWrite` reaches the end of the reverse SSA
     // use-def chain, it returns that value, regardless of whether it is a
     // memory write or not.
if (!hasTensorResult && !hasTensorOperand) {
     for (Region &region : op->getRegions())
       if (failed(bufferize(&region, state)))
         return failure();
     return success();
   }
 
   // Bufferize using `BufferizableOpInterface`. Interface implementations are
   // responsible for bufferizing nested ops.
+  if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op)) {
     b.setInsertionPoint(op);
-  if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
     return bufferizableOp.bufferize(b, state);
+  }
 
-  // Emit error if tensor op is not bufferizable.
+  // `op` is an unbufferizable tensor op.
+  if (!state.options.allowUnknownOps)
     return op->emitError() << "unsupported op with tensors";
+
+  // Replace all OpOperands with "to-tensor casted" bufferized values.
+  for (OpOperand &operand : op->getOpOperands()) {
+    if (operand.get().getType().isa<TensorType>() &&
+        state.isMapped(operand.get())) {
+      b.setInsertionPoint(op);
+      Value toTensorOp = b.create<bufferization::ToTensorOp>(
+          op->getLoc(), state.lookupBuffer(operand.get()));
+      operand.set(toTensorOp);
+    }
+  }
+
+  for (Region &region : op->getRegions())
+    if (failed(bufferize(&region, state)))
+      return failure();
+
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific scoped alloc/dealloc insertion support.
 //===----------------------------------------------------------------------===//
 
 /// Move the insertion point of the given builder to the beginning of a
 /// surrounding block as much as possible, while not crossing any allocation

 void mlir::linalg::comprehensive_bufferize::BufferizationState::mapValue(
     Value from, Value to) {
   assert(from && "unexpected empty value");
   return mapping.map(from, to);
 }
 
 /// Wrapper for better debugging.
 Value mlir::linalg::comprehensive_bufferize::BufferizationState::lookupBuffer(
-    Value tensor) const {
+    Value tensor) {
   // TODO: if key comes from bbArg, forward.
   assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
-  Value v = mapping.lookupOrNull(tensor);
+  Value buffer = mapping.lookupOrNull(tensor);
+
+  if (!buffer) {
+    if (options.allowUnknownOps) {
+      // `tensor` was not bufferized yet. This should never happen with
+      // bufferizable ops.
+      assert(!tensor.getDefiningOp<BufferizableOpInterface>() &&
+             "tensor is not mapped");
+      // Insert to_memref op.
+      OpBuilder b(tensor.getContext());
+      setInsertionPointAfter(b, tensor);
+      return b.create<bufferization::ToMemrefOp>(
+          tensor.getLoc(),
+          getDynamicMemRefType(tensor.getType().cast<RankedTensorType>()),
+          tensor);
+    }
 
-  if (!v) {
     // Dump tensor for easier debugging.
     tensor.dump();
     llvm_unreachable("tensor is not mapped");
     return Value();
   }
 
-  assert((v.getType().isa<MemRefType>() ||
-          v.getType().isa<UnrankedMemRefType>()) &&
+  assert((buffer.getType().isa<MemRefType>() ||
+          buffer.getType().isa<UnrankedMemRefType>()) &&
          "expected that tensor is mapped to memref");
-  return v;
+  return buffer;
 }
 
 Value mlir::linalg::comprehensive_bufferize::BufferizationState::lookupValue(
     Value value) const {
   Value v = mapping.lookupOrNull(value);
   if (!v) {
     llvm_unreachable("tensor is not mapped");
     return Value();

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

void LinalgComprehensiveModuleBufferize::runOnOperation() {
   }
   // TODO: Change to memref::CopyOp (default memCpyFn).
   options.allocationFns->memCpyFn = [](OpBuilder &b, Location loc, Value from,
                                        Value to) {
     b.create<linalg::CopyOp>(loc, from, to);
   };
 
   options.allowReturnMemref = allowReturnMemref;
+  options.allowUnknownOps = allowUnknownOps;
   options.analysisFuzzerSeed = analysisFuzzerSeed;
   options.testAnalysisOnly = testAnalysisOnly;
 
   // Enable InitTensorOp elimination.
   options.addPostAnalysisStep<
       linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
   // TODO: Find a way to enable this step automatically when bufferizing tensor
   // dialect ops.

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

+// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref allow-unknown-ops" -split-input-file | FileCheck %s
+
+// TODO: Bufferize result IR of bufferization.
+// TODO: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref allow-unknown-ops" -linalg-comprehensive-module-bufferize="allow-return-memref allow-unknown-ops" -split-input-file | FileCheck %s
+
+// Run fuzzer with different seeds.
+// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
+
+// CHECK-LABEL: func @use_of_unknown_op_1(
+//  CHECK-SAME:     %[[t1:.*]]: memref<?xf32
+func @use_of_unknown_op_1(%t1: tensor<?xf32> {linalg.inplaceable = true})
+    -> vector<5xf32> {
+  // CHECK: %[[t1_tensor:.*]] = bufferization.to_tensor %[[t1]]
+  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[t1_tensor]])
+  %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
+
+  %idx = arith.constant 0 : index
+  %cst = arith.constant 0.0 : f32
+  // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
+  // CHECK: vector.transfer_read %[[dummy_memref]]
+  %1 = vector.transfer_read %0[%idx], %cst : tensor<?xf32>, vector<5xf32>
+  return %1 : vector<5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @use_of_unknown_op_2(
+//  CHECK-SAME:     %[[t1:.*]]: memref<?xf32
+func @use_of_unknown_op_2(%t1: tensor<?xf32> {linalg.inplaceable = true})
+    -> tensor<?xf32> {
+  // CHECK: %[[t1_tensor:.*]] = bufferization.to_tensor %[[t1]]
+
+  // CHECK: %[[dummy1:.*]] = "test.dummy_op"(%[[t1_tensor]])
+  %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
+  // CHECK: %[[dummy2:.*]] = "test.another_dummy_op"(%[[dummy1]])
+  %1 = "test.another_dummy_op"(%0) : (tensor<?xf32>) -> tensor<?xf32>
+
+  // CHECK: %[[dummy2_memref:.*]] = bufferization.to_memref %[[dummy2]]
+  // CHECK: return %[[dummy2_memref]]
+  return %1 : tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @use_of_unknown_op_3(
+//  CHECK-SAME:     %[[t1:.*]]: memref<?xf32
+func @use_of_unknown_op_3(%t1: tensor<?xf32> {linalg.inplaceable = true})
+    -> (vector<5xf32>, vector<5xf32>) {
+  %idx = arith.constant 0 : index
+  %cst = arith.constant 0.0 : f32
+  // CHECK: %[[v1:.*]] = vector.transfer_read %[[t1]]
+  %1 = vector.transfer_read %t1[%idx], %cst : tensor<?xf32>, vector<5xf32>
+
+  // CHECK: %[[t1_tensor:.*]] = bufferization.to_tensor %[[t1]]
+  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[t1_tensor]])
+  %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
+  // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
+  // CHECK: %[[v2:.*]] = vector.transfer_read %[[dummy_memref]]
+  %2 = vector.transfer_read %0[%idx], %cst : tensor<?xf32>, vector<5xf32>
+
+  // CHECK: return %[[v1]], %[[v2]]
+  return %1, %2 : vector<5xf32>, vector<5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @use_of_unknown_op_4(
+//  CHECK-SAME:     %[[t1:.*]]: memref<?xf32
+func @use_of_unknown_op_4(%t1: tensor<?xf32> {linalg.inplaceable = true})
+    -> (vector<5xf32>, tensor<?xf32>) {
+  %idx = arith.constant 0 : index
+  %cst = arith.constant 0.0 : f32
+
+  // CHECK: %[[t1_tensor:.*]] = bufferization.to_tensor %[[t1]]
+  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[t1_tensor]])
+  %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
+
+  // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
+  // CHECK: %[[v1:.*]] = vector.transfer_read %[[dummy_memref]]
+  %1 = vector.transfer_read %0[%idx], %cst : tensor<?xf32>, vector<5xf32>
+
+  // CHECK: %[[another_dummy:.*]] = "test.another_dummy_op"(%[[dummy]])
+  %2 = "test.another_dummy_op"(%0) : (tensor<?xf32>) -> tensor<?xf32>
+
+  // CHECK: %[[another_dummy_memref:.*]] = bufferization.to_memref %[[another_dummy]]
+  // CHECK: return %[[v1]], %[[another_dummy_memref]]
+  return %1, %2 : vector<5xf32>, tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @use_of_bufferizable_op_in_unbufferizable_op
+//  CHECK-SAME:     %[[t1:.*]]: memref<?xf32
+func @use_of_bufferizable_op_in_unbufferizable_op(
+    %t1: tensor<?xf32>, %o: index, %s: index) -> (tensor<?xf32>, tensor<?xf32>) {
+  // CHECK: %[[subview:.*]] = memref.subview %[[t1]]
+  %0 = tensor.extract_slice %t1[%o][%s][1] : tensor<?xf32> to tensor<?xf32>
+  // CHECK: %[[subview_tensor:.*]] = bufferization.to_tensor %[[subview]]
+  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[subview_tensor]])
+  %1 = "test.dummy_op"(%0) : (tensor<?xf32>) -> tensor<?xf32>
+  // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
+  // CHECK: return %[[subview]], %[[dummy_memref]]
+  return %0, %1 : tensor<?xf32>, tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @unused_unknown_op(
+//  CHECK-SAME:     %[[t1:.*]]: memref<?xf32
+func @unused_unknown_op(%t1 : tensor<?xf32>) -> vector<5xf32> {
+  %idx = arith.constant 0 : index
+  %cst = arith.constant 0.0 : f32
+  // CHECK: vector.transfer_read %[[t1]]
+  %1 = vector.transfer_read %t1[%idx], %cst : tensor<?xf32>, vector<5xf32>
+
+  // CHECK: %[[t1_tensor:.*]] = bufferization.to_tensor %[[t1]]
+  // CHECK: "test.dummy_op"(%[[t1_tensor]])
+  "test.dummy_op"(%t1) : (tensor<?xf32>) -> ()
+
+  return %1 : vector<5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @unknown_op_not_writable
+//  CHECK-SAME:     %[[t1:.*]]: memref<?xf32
+func @unknown_op_not_writable(
+    %t1 : tensor<?xf32>, %v :  vector<5xf32>, %idx : index) -> tensor<?xf32> {
+  // CHECK: %[[t1_tensor:.*]] = bufferization.to_tensor %[[t1]]
+  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[t1_tensor]])
+  // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
+  %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> (tensor<?xf32>)
+
+  // CHECK: %[[dim:.*]] = tensor.dim %[[dummy]]
+  // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
+  // CHECK: linalg.copy(%[[dummy_memref]], %[[alloc]])
+  // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
+  %1 = vector.transfer_write %v, %0[%idx] : vector<5xf32>, tensor<?xf32>
+
+  // CHECK: return %[[alloc]]
+  return %1 : tensor<?xf32>
+}