Propagate layout permutations through split (#22594) - #22594
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22594
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit f200926 with merge base 978e584 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@mcremon-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D118743098. |
This PR needs a
|
Summary: Treat a `split_with_sizes_copy` and the `getitem` nodes that unpack it as one rank-preserving unit in `RemovePermutesAroundElementwiseOps`, remapping the split's `dim` under the region's permutation exactly as `cat`'s dim is already remapped. A permutation only reorders axes, so the extent of the split axis is unchanged and `split_sizes` stays valid. `visit` accepts a region only if every node in its transitive closure is permutable, so a single unrecognised node discards the whole subgraph. `getitem` is a plain `operator.getitem`: it carries no `_op`, so it is not tagged pointwise, it is not in the permutable set, and it is not a view. Any region reaching a split therefore failed, even though tuple indexing is layout-invariant and the split itself only needs its dim remapped. Guard the numel-one input check with an `isinstance` test. A multi-output op carries a list of fake tensors in `meta["val"]`, and calling `numel()` on that list raises. Differential Revision: D118743098
d229f31 to
f200926
Compare
Summary:
Treat a
split_with_sizes_copyand thegetitemnodes that unpack it as one rank-preserving unit inRemovePermutesAroundElementwiseOps, remapping the split'sdimunder the region's permutation exactly ascat's dim is already remapped. A permutation only reorders axes, so the extent of the split axis is unchanged andsplit_sizesstays valid.visitaccepts a region only if every node in its transitive closure is permutable, so a single unrecognised node discards the whole subgraph.getitemis a plainoperator.getitem: it carries no_op, so it is not tagged pointwise, it is not in the permutable set, and it is not a view. Any region reaching a split therefore failed, even though tuple indexing is layout-invariant and the split itself only needs its dim remapped.Guard the numel-one input check with an
isinstancetest. A multi-output op carries a list of fake tensors inmeta["val"], and callingnumel()on that list raises.Differential Revision: D118743098