-
Couldn't load subscription status.
- Fork 608
Description
From discussion with @zjgarvey, we have the following findings (mostly Zach, huge thanks):
The torch.chunk operation currently relies on aten.as_strided which is overkill for what the op actually does, and it can use .slice instead. However, we also note that aten.as_strided relies on accurate stride input layout, but in torch-mlir, tensors are materialized as contiguous (torch can have non-contiguous tensors, we don't. For example, transpose the input.). Therefore the static accessings args from as_strided refer to data that is no longer in the same format.
aten.as_strided is only generated in decompositions of split.Tensor. We propose that
- Remove
aten.as_stridedfrom all decomposition lists in torch-mlir (all paths reachable post-import) fix: remove split.Tensor and split_with_sizes from decomp tables #4340 - On the IREE side, we can add the
RecomposeComplexOpspass from torch-mlir that will decompose.chunkinto.sliceinstead ofas_strided[Codegen] Add RecomposeComplexOps pass in Torch/InputConversion/Passes iree-org/iree#22276 - Remove
split.Tensorandsplit_with_sizesfrom IREE-Turbine's decomposition table http://github.com/iree-org/iree-turbine/pull/1175
Eventually we should probably deprecate theaten.as_stridedsupport and throw an error on input when we see aaten.as_stridedop, since we cannot guarantee the correctness.
Related issues: