From 5f6d5ca0f80c79813ab3821dc87ae01c86146398 Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Wed, 21 Dec 2022 12:39:19 +0100 Subject: [PATCH] [mlir][bufferize] Fix tensor copy insertion for dynamic tensors TensorCopyInsertion inserts bufferization.alloc_tensor ops in case of RaW conflicts. If such a tensor is dynamically shaped, tensor.dim ops are inserted. There is an optimization for ops such as tensor.extract_slice: A copy of the result is created instead of the operand. Afterwards, all uses of the result are updated. E.g.: ``` %0 = tensor.extract_slice ... : tensor to tensor %1 = tensor.dim %0, %c0 : tensor %2 = bufferization.alloc_tensor(%dim) : tensor ``` All uses of %0, except for tensor.dim and bufferization.alloc_tensor (if any), should be replaced. Before this change, the use in tensor.dim was also replaced, resulting in IR that had a dominance error. Note: There is no test case for this fix because the bug cannot be triggered with tensor.extract_slice, which implements an interface to reify result shapes. This bug appeared in an external project with a tensor.extract_slice-like op that does not implement that interface, in which case tensor.dim ops must be created. We do not have such an op in MLIR to trigger this bug. Differential Revision: https://reviews.llvm.org/D140471 --- mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp index f79db15..68990ea 100644 --- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp +++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp @@ -201,8 +201,13 @@ LogicalResult BufferizableOpInterface::resolveTensorOpOperandConflicts( opResult.getUses(), [](OpOperand &use) { return &use; })); for (OpOperand *use : uses) { // Do not update the alloc_tensor op that we just created. - if (use->getOwner() != copy->getDefiningOp()) - rewriter.updateRootInPlace(use->getOwner(), [&]() { use->set(*copy); }); + if (use->getOwner() == copy->getDefiningOp()) + continue; + // tensor.dim ops may have been created to be used as alloc_tensor op + // dynamic extents. Do not update these either. + if (isa(use->getOwner())) + continue; + rewriter.updateRootInPlace(use->getOwner(), [&]() { use->set(*copy); }); } } -- 2.7.4