diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h --- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h +++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h @@ -315,12 +315,14 @@ /// tensor values. class AnalysisState { public: - /// 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 getAliasingOpOperand(OpResult result) const; + /// Determine which OpOperand* will alias with `opResult` if the op is + /// bufferized in place. Return all tensor OpOperand* if the op is not + /// bufferizable. + SmallVector getAliasingOpOperand(OpResult opResult) const; /// Determine which OpResult will alias with `opOperand` if the op is - /// bufferized in place. Return an empty vector if the op is not bufferizable. + /// bufferized in place. Return all tensor OpResults if the op is not + /// bufferizable. SmallVector getAliasingOpResult(OpOperand &opOperand) const; /// Return true if `opOperand` bufferizes to a memory read. Return `true` if diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp --- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp +++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp @@ -349,24 +349,39 @@ } } -/// Determine which OpOperand* will alias with `result` if the op is bufferized -/// in place. Return an empty vector if the op is not bufferizable. +/// Determine which OpOperand* will alias with `opResult` if the op is +/// bufferized in place. Return all tensor OpOperand* if the op is not +/// bufferizable. SmallVector -AnalysisState::getAliasingOpOperand(OpResult result) const { - if (Operation *op = result.getDefiningOp()) +AnalysisState::getAliasingOpOperand(OpResult opResult) const { + if (Operation *op = opResult.getDefiningOp()) if (auto bufferizableOp = getOptions().dynCastBufferizableOp(op)) - return bufferizableOp.getAliasingOpOperand(result, *this); - return {}; + return bufferizableOp.getAliasingOpOperand(opResult, *this); + + // The op is not bufferizable. Conservatively assume that everything is + // aliasing. + SmallVector r; + for (OpOperand &operand : opResult.getDefiningOp()->getOpOperands()) + if (operand.get().getType().isa()) + r.push_back(&operand); + return r; } /// Determine which OpResult will alias with `opOperand` if the op is bufferized -/// in place. Return an empty vector if the op is not bufferizable. +/// in place. Return all tensor OpResults if the op is not bufferizable. SmallVector AnalysisState::getAliasingOpResult(OpOperand &opOperand) const { if (auto bufferizableOp = getOptions().dynCastBufferizableOp(opOperand.getOwner())) return bufferizableOp.getAliasingOpResult(opOperand, *this); - return {}; + + // The op is not bufferizable. Conservatively assume that everything is + // aliasing. + SmallVector r; + for (OpResult result : opOperand.getOwner()->getOpResults()) + if (result.getType().isa()) + r.push_back(result); + return r; } /// Return true if `opOperand` bufferizes to a memory read. Return `true` if the diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp @@ -920,10 +920,9 @@ auto analyzeOp = [&](Operation *op) { for (OpOperand &opOperand : op->getOpOperands()) if (opOperand.get().getType().isa()) - if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op)) - if (failed(bufferizableInPlaceAnalysisImpl(opOperand, aliasInfo, - state, domInfo))) - return failure(); + if (failed(bufferizableInPlaceAnalysisImpl(opOperand, aliasInfo, state, + domInfo))) + return failure(); return success(); }; @@ -1049,11 +1048,10 @@ const BufferizationAliasInfo &aliasInfo, const BufferizationOptions &options) { // Add __inplace_operands_attr__. - op->walk([&](BufferizableOpInterface bufferizableOp) { - if (options.isOpAllowed(bufferizableOp.getOperation())) - for (OpOperand &opOperand : bufferizableOp->getOpOperands()) - if (opOperand.get().getType().isa()) - setInPlaceOpOperand(opOperand, aliasInfo.isInPlace(opOperand)); + op->walk([&](Operation *op) { + for (OpOperand &opOperand : op->getOpOperands()) + if (opOperand.get().getType().isa()) + setInPlaceOpOperand(opOperand, aliasInfo.isInPlace(opOperand)); }); } diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir @@ -0,0 +1,34 @@ +// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only" -allow-unregistered-dialect -split-input-file | FileCheck %s + +// CHECK-LABEL: func @unknown_op_aliasing( +func.func @unknown_op_aliasing(%f: f32, %f2: f32, %pos: index) -> f32 { + %0 = tensor.empty() : tensor<10xf32> + // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]} + %1 = linalg.fill ins(%f : f32) outs(%0 : tensor<10xf32>) -> tensor<10xf32> + + // Something must bufferize out-of-place because the op may return an alias + // of %1. + // CHECK: "dummy.dummy_op"(%{{.*}}) {__inplace_operands_attr__ = ["false"]} + %alias = "dummy.dummy_op"(%1) : (tensor<10xf32>) -> (tensor<10xf32>) + + // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]} + %2 = linalg.fill ins(%f2 : f32) outs(%1 : tensor<10xf32>) -> tensor<10xf32> + %3 = tensor.extract %alias[%pos] : tensor<10xf32> + return %3 : f32 +} + +// ----- + +// CHECK-LABEL: func @unknown_op_writing( +func.func @unknown_op_writing(%f: f32, %f2: f32, %pos: index) -> f32 { + %0 = tensor.empty() : tensor<10xf32> + // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]} + %1 = linalg.fill ins(%f : f32) outs(%0 : tensor<10xf32>) -> tensor<10xf32> + + // The op may bufferize to a memory write, so it must bufferize out-of-place. + // CHECK: "dummy.dummy_op"(%{{.*}}) {__inplace_operands_attr__ = ["false"]} + "dummy.dummy_op"(%1) : (tensor<10xf32>) -> () + + %3 = tensor.extract %1[%pos] : tensor<10xf32> + return %3 : f32 +} diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-partial.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-partial.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-partial.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-partial.mlir @@ -100,11 +100,14 @@ %t1: tensor, %o: index, %s: index) -> (tensor, tensor) { // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] // CHECK: %[[subview:.*]] = memref.subview %[[m1]] + // The op must alloc because "test.dummy" may bufferize to a memory write. + // CHECK: %[[alloc:.*]] = memref.alloc + // CHECK: memref.copy %[[subview]], %[[alloc]] %0 = tensor.extract_slice %t1[%o][%s][1] : tensor to tensor - // CHECK: %[[subview_tensor:.*]] = bufferization.to_tensor %[[subview]] - // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[subview_tensor]]) + // CHECK: %[[alloc_tensor:.*]] = bufferization.to_tensor %[[alloc]] + // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[alloc_tensor]]) %1 = "test.dummy_op"(%0) : (tensor) -> tensor - // CHECK: return %[[subview_tensor]], %[[dummy]] + // CHECK: return %[[alloc_tensor]], %[[dummy]] return %0, %1 : tensor, tensor } diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-pass-statistics.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-pass-statistics.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-pass-statistics.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-pass-statistics.mlir @@ -5,7 +5,7 @@ // CHECK: (S) 1 num-buffer-alloc // CHECK: (S) 1 num-buffer-dealloc // CHECK: (S) 1 num-tensor-in-place -// CHECK: (S) 1 num-tensor-out-of-place +// CHECK: (S) 2 num-tensor-out-of-place func.func @read_after_write_conflict(%cst : f32, %idx : index, %idx2 : index) -> (f32, f32) { %t = "test.dummy_op"() : () -> (tensor<10xf32>) diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir @@ -136,7 +136,7 @@ // CHECK-LABEL: func @select_different_tensors( // CHECK-SAME: %[[t:.*]]: tensor -func.func @select_different_tensors(%t: tensor, %sz: index, %c: i1) -> tensor { +func.func @select_different_tensors(%t: tensor, %sz: index, %pos: index, %c: i1) -> f32 { // CHECK-DAG: %[[m:.*]] = bufferization.to_memref %[[t]] : memref // CHECK-DAG: %[[alloc:.*]] = memref.alloc(%{{.*}}) {{.*}} : memref %0 = bufferization.alloc_tensor(%sz) : tensor @@ -145,7 +145,8 @@ // CHECK: %[[casted:.*]] = memref.cast %[[alloc]] : memref to memref // CHECK: arith.select %{{.*}}, %[[casted]], %[[m]] %1 = arith.select %c, %0, %t : tensor - return %1 : tensor + %2 = tensor.extract %1[%pos] : tensor + return %2 : f32 } // -----