[compiler][flow] Move cast, reshape and bitcast after transfer op #18742
Conversation
sogartar
requested review from
hanhanW,
MaheshRavishankar,
IanWood1 and
benvanik
as code owners
We got incoming IR of the form
```mlir
%cast = tensor.cast %0 : tensor<2xf32> to tensor<?xf32>
%2 = flow.tensor.transfer %cast : tensor<?xf32>{%c2} to
#hal.device.affinity<@__device_0>
%cast_2 = tensor.cast %2 : tensor<?xf32> to tensor<2xf32>
```
We would like to allow for the 2 casts to get folded.
We would also like to reduce the dynamism of the transfer op. To
operate on a tensor with fewer dynamic dimensions.
Signed-off-by: Boian Petkantchin <boian.petkantchin@amd.com>
sogartar
force-pushed
the
cast-transfer-reorder
branch
from
f4c155a
to
d5794d3
Compare
|
I've been trying to fix a similar problem with If this is a recurring problem, maybe there is a more systematic way to address this? |
|
@IanWood1 This seems similar. I am not sure what is more appropriate, to fold the reshape or to move it. |
|
Folding is best whenever possible - the reshape/cast ops are only there to carry metadata and should be aggressively removed when the metadata they carry is not useful. For example, %0 = flow.tensor.reshape %arg : tensor<1x2xf32> -> tensor<2x1xf32>
%1 = flow.tensor.transfer %0 : tensor<2x1xf32> to "target"
%2 = flow.tensor.reshape %1 : tensor<2x1xf32> -> tensor<1x2xf32>should never exist - the canonical form would be: %2 = flow.tensor.transfer %arg : tensor<1x2xf32> to "target"If propagating the reshapes/casts/etc allows those to fold by having an intermediate state like: %0 = flow.tensor.transfer %arg : tensor<1x2xf32> to "target"
%1 = flow.tensor.reshape %0 : tensor<1x2xf32> -> tensor<2x1xf32>
%2 = flow.tensor.reshape %1 : tensor<2x1xf32> -> tensor<1x2xf32>and then the existing reshape-reshape canonicalizer kicks in that's a good thing. |
|
|
||
| // Unranked shapes are always considered to have more dynamic dimensions than | ||
| // ranked. | ||
| inline bool shapeHasLessOrEqualDynamicDimensions(ShapedType t1, ShapedType t2) { |
There was a problem hiding this comment.
YAGNI
on the 3rd or 4th time a 3 line block of code is used it's worth hoisting into a global shared util - but there's a bar for doing this - every shared util pollutes the project and adds a maintenance burden and has to be worth it. A single use in a single location is not worth it. Just inline this as static where this is used.
(this would also need to be static here - you can't put inline functions in header files)
| template <typename Op, unsigned homomorphismOpOperandIndex = 0, | ||
| unsigned homomorphismOpResultIndex = 0> | ||
| static void populateMoveOpAfterTransferPattern(RewritePatternSet &results) { | ||
|
|
| // before the transfer instead. We strive to reduce the dynamism of the | ||
| // transfer op. If there will be no strict dynamism improvement, we prefer the | ||
| // other op after the transfer. | ||
| // TODO: add the analogous move-befor-transfer pattern. |
There was a problem hiding this comment.
| // TODO: add the analogous move-befor-transfer pattern. | |
| // TODO: add the analogous move-before-transfer pattern. |
(a spell checker can help in your IDE)
There was a problem hiding this comment.
Done. I have one but it ignored the hyphened word.
| } // namespace | ||
|
|
||
| void TensorTransferOp::getCanonicalizationPatterns(RewritePatternSet &results, | ||
| MLIRContext *context) { | ||
| results.insert<ElideRedundantTransfer>(context); | ||
| populateMoveOpAfterTransferPattern<tensor::BitcastOp>(results); | ||
| populateMoveOpAfterTransferPattern<TensorBitCastOp>(results); |
There was a problem hiding this comment.
| populateMoveOpAfterTransferPattern<TensorBitCastOp>(results); | |
| populateMoveOpAfterTransferPattern<IREE::Flow::TensorBitCastOp>(results); |
use namespaced names where there may be ambiguity to make it clearer to readers
| populateMoveOpAfterTransferPattern<TensorBitCastOp>(results); | ||
| populateMoveOpAfterTransferPattern<tensor::CastOp>(results); | ||
| populateMoveOpAfterTransferPattern<tensor::ReshapeOp>(results); | ||
| populateMoveOpAfterTransferPattern<TensorReshapeOp>(results); |
There was a problem hiding this comment.
| populateMoveOpAfterTransferPattern<TensorReshapeOp>(results); | |
| populateMoveOpAfterTransferPattern<IREE::Flow::TensorReshapeOp>(results); |
We got incoming IR of the form
We would like to allow for the 2 casts to get folded. We would also like to reduce the dynamism of the transfer op. To operate on a tensor with fewer dynamic dimensions.