// 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 *>
// 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");
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;
// Bufferize using `BufferizableOpInterface`. Interface implementations are
// responsible for bufferizing nested ops.
- b.setInsertionPoint(op);
- if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
+ if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op)) {
+ b.setInsertionPoint(op);
return bufferizableOp.bufferize(b, state);
+ }
+
+ // `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 ®ion : op->getRegions())
+ if (failed(bufferize(®ion, state)))
+ return failure();
- // Emit error if tensor op is not bufferizable.
- return op->emitError() << "unsupported op with tensors";
+ return success();
}
//===----------------------------------------------------------------------===//
/// 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(
--- /dev/null
+// 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: %[[m1:.*]]: memref<?xf32
+func @use_of_unknown_op_1(%t1: tensor<?xf32> {linalg.inplaceable = true})
+ -> vector<5xf32> {
+ // ToTensorOp is generated because the function is bufferized and has a
+ // memref block argument.
+ // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
+ // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[m1_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: %[[m1:.*]]: memref<?xf32
+func @use_of_unknown_op_2(%t1: tensor<?xf32> {linalg.inplaceable = true})
+ -> tensor<?xf32> {
+ // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
+
+ // CHECK: %[[dummy1:.*]] = "test.dummy_op"(%[[m1_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: %[[m1:.*]]: 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 %[[m1]]
+ %1 = vector.transfer_read %t1[%idx], %cst : tensor<?xf32>, vector<5xf32>
+
+ // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
+ // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[m1_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: %[[m1:.*]]: 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: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
+ // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[m1_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: %[[m1:.*]]: 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 %[[m1]]
+ %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: %[[m1:.*]]: 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 %[[m1]]
+ %1 = vector.transfer_read %t1[%idx], %cst : tensor<?xf32>, vector<5xf32>
+
+ // ToTensorOp is inserted to pass in the result of the above bufferized op.
+ // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
+ // CHECK: "test.dummy_op"(%[[m1_tensor]])
+ "test.dummy_op"(%t1) : (tensor<?xf32>) -> ()
+
+ return %1 : vector<5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @unknown_op_not_writable
+// CHECK-SAME: %[[m1:.*]]: memref<?xf32
+func @unknown_op_not_writable(
+ %t1 : tensor<?xf32>, %v : vector<5xf32>, %idx : index) -> tensor<?xf32> {
+ // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
+ // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[m1_tensor]])
+ // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
+ %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> (tensor<?xf32>)
+
+ // The result of an unknown op is not writable. Always generate a copy.
+ // Note: This copy is essential for partial bufferization. Otherwise, we could
+ // introducing a RaW conflict.
+ // 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>
+}