[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<?xf32> to tensor<?xf32>
%1 = tensor.dim %0, %c0 : tensor<?xf32>
%2 = bufferization.alloc_tensor(%dim) : tensor<?xf32>
```

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

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index f79db15..68990ea 100644 (file)
--- 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<tensor::DimOp>(use->getOwner()))
+        continue;
+      rewriter.updateRootInPlace(use->getOwner(), [&]() { use->set(*copy); });
     }
   }