[Mlir-commits] [mlir] 5f6d5ca - [mlir][bufferize] Fix tensor copy insertion for dynamic tensors

[Matthias Springer]

Author: Matthias Springer
Date: 2022-12-21T12:42:02+01:00
New Revision: 5f6d5ca0f80c79813ab3821dc87ae01c86146398

URL: https://github.com/llvm/llvm-project/commit/5f6d5ca0f80c79813ab3821dc87ae01c86146398
DIFF: https://github.com/llvm/llvm-project/commit/5f6d5ca0f80c79813ab3821dc87ae01c86146398.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

Removed: 
    


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