Simplify aten_unbind when shape is static #2597
Conversation
Add static shape handling to aten_unbind function
justinchuby
added
the
module: torchlib
There was a problem hiding this comment.
Pull Request Overview
This PR simplifies the aten_unbind function by adding static shape handling to improve efficiency when the tensor dimension is known at compile time.
- Adds a static shape optimization path that uses
Split+Squeezeoperations instead ofSplitToSequence - Maintains backward compatibility by falling back to the original dynamic implementation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2597 +/- ##
=======================================
Coverage 69.95% 69.95%
=======================================
Files 222 222
Lines 26311 26314 +3
Branches 2604 2605 +1
=======================================
+ Hits 18406 18409 +3
Misses 6993 6993
Partials 912 912 ☔ View full report in Codecov by Sentry. |
Co-authored-by: Copilot <[email protected]>
| if isinstance(self.shape[dim], int): | ||
| # We can create a definitive split op if the input shape is static | ||
| outputs = op.Split(self, axis=dim, num_outputs=self.shape[dim]) | ||
| return [op.Squeeze(out, [dim]) for out in outputs] |
There was a problem hiding this comment.
I feel like it makes more sense to rewrite it? Although this PR probably works, it's adding another dimension on torchlib (covering both static and dynamic cases). Maybe let torchlib be as dynamic as possible, and we can optimize it after.
There was a problem hiding this comment.
How easy is it to create the optimization rules? I am fine either way
There was a problem hiding this comment.
Looks like SplitToSequence(self, axis=dim, keepdims=False) should be generically rewritable to the subgraph if the split axis is known. This can potentially cover more cases with other ops when encountered in the future
There was a problem hiding this comment.
@gramalingam for suggestions on the rewrite rule. Is this related to #2581 ?
There was a problem hiding this comment.
From a computation point of view, it is always better to generate the correct graph rather than producing a graph which needs to be rewritten. Matching a pattern takes time.
There was a problem hiding this comment.
But my understanding is that the original implementation is not incorrect. It's only the op is not preferred because of the backend implementation, which fits the category of rewritten rules. We surely can say it's more convenient to address this way (this PR), but I prefer an established/explicit rule to say when/what we should add support in torchlib, and under what condition we add rewrite rules/constat folding. Otherwise, it's just scattered around. And if we want it to be done in this way, do we consider upstream some other optimizations downstream that are optimized away because of static shapes as well?
There was a problem hiding this comment.
I kinda think we should do both. (1) we can measure the complexity of the torchlib implementation. I feel that the complexity of the current implementation is not high for the immediate benefits it brings. If the graph can be significantly simplified because we know some shapes are static, I think we should pursue that in torchlib. When looking at micro-optimizations like this I agree that we should be more careful and decide on a case by case basis. (2) we should still have a rule that will simplify this so that our tooling can handle SplitToSequence generally.
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
|
I updated tests and improved compatibility with torch<2.7 |
|
Merging for now. Please let me know if further changes are needed. |
Add static shape handling to aten_unbind function.
Fix #2596