Repository navigation
refactor(dynamo): make _requires_output_allocator a pure predicate - #4369
Merged
Merged
Conversation
The _requires_output_allocator helper previously folded the fallback_data_dependent_ops setting into its return value, so it returned False for a node that genuinely required an output allocator whenever the setting was disabled. That made the name misleading at the call sites. Flip the semantics so the helper purely reports whether the converter selected for a node needs a TRT output allocator (data-dependent output shape), and move the fallback_data_dependent_ops gate to the two operator support call sites. The setting is still checked first, preserving the short-circuit that skips the registry lookup in the default (off) case. Behavior is unchanged end to end. Update the unit test to assert the pure predicate is setting-independent and to verify the setting-gated fallback decision through is_node_supported. Signed-off-by: cehongwang <wangcehong@gmail.com>
github-actions
Bot
added
component: tests
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Follow-up cleanup to #4355.
The
_requires_output_allocatorhelper in the partitioner previously folded thefallback_data_dependent_opssetting into its return value: it returnedFalsefor a node that genuinely required a TRT output allocator whenever the setting was disabled (the default). That made the name misleading at the call sites —_requires_output_allocator(node) == Falsedid not actually mean "this node does not require an output allocator," it usually meant "the setting is off."This flips the semantics so the helper is a pure predicate: it reports whether the converter selected for a node needs a TRT output allocator (data-dependent output shape, e.g.
nonzeroor booleanindex.Tensor), independent of any setting. Thefallback_data_dependent_opsgate moves to the two operator-support call sites (TorchTensorRTOperatorSupport.is_node_supportedandOpSupportTester.is_node_supported).The setting is still checked first at the call sites, preserving the short-circuit that skips the converter-registry lookup in the default (off) case. Behavior is unchanged end to end.
Details
_global_partitioner.py:_requires_output_allocatorno longer readscompilation_settings; the caller now guards withsettings is not None and settings.fallback_data_dependent_ops and self._requires_output_allocator(node)._adjacency_partitioner.py: same gate added at theOpSupportTestercall site.Tests
tests/py/dynamo/models/test_fallback_data_dependent_ops.py:test_requires_output_allocator_is_setting_independent: the pure predicate returnsTruefornonzeroregardless of the setting.test_output_allocator_node_falls_back_only_when_enabled: now drivesis_node_supported(...)and asserts the node stays on TensorRT (True) when the flag is off and falls back (False) when it is on — testing the real partitioner decision rather than the helper in isolation.Type of change
Checklist