perf(backend): trim redundant work in Variable value/dtype accessors#23297
Closed
pctablet505 wants to merge 10 commits into
Closed
perf(backend): trim redundant work in Variable value/dtype accessors#23297pctablet505 wants to merge 10 commits into
pctablet505 wants to merge 10 commits into
Conversation
…ribution() hot paths Part of keras-team#22561.
cast() had three sequential dtype-check-and-cast blocks: one for isinstance(x, torch.Tensor), one for isinstance(x, Variable), and one for is_tensor(x). Since is_tensor(x) is literally isinstance(x, torch.Tensor), the standalone torch.Tensor block was functionally identical to the is_tensor(x) block below it and unreachable in a meaningful way. Collapse to two blocks: unwrap Variable first, then let is_tensor(x) handle both raw tensors and unwrapped Variables.
…sor start indices
_to_static_index accepted 0-d integer tensors and coerced them via int(), specializing a data-dependent bound into a concrete Python int on the fast path; tensors now fall through to the tensor-native torch.narrow path instead. Also add an explicit length check for start_indices/shape: the fast path's zip() previously truncated silently on a mismatch instead of raising.
Rewrite the _to_static_index docstring and the slice() fast-path comment as flowing prose with single-backtick inline code, matching the surrounding torch backend, and rename the loop bindings for clarity. No behavior change.
The Python-primitive branch of convert_to_tensor mapped bool/int/float /complex to a torch.dtype through an intermediate string constant and to_torch_dtype(), which re-normalizes the string via standardize_dtype and does a dict lookup on every call. When dtype is None, that string round-trip is unnecessary — map directly to the torch.dtype constant instead, and only fall back to to_torch_dtype() when the caller passed an explicit dtype override. _tensor_on_device now spells out tensor_device/target instead of abbreviating, and checks torch.cuda.current_device()/torch.xpu. current_device() only once the tensor's device type already matches the target's, so a tensor can only reach that branch if the corresponding accelerator was available to create it in the first place — no availability guard needed. XPU is handled the same way as CUDA for index-less device targets. _convert_numpy_or_arraylike_to_tensor is defined ahead of convert_to_tensor, matching the caller precedes callee convention used elsewhere in this module.
Resolves the core_test.py conflict between P3's get_device/device_scope test and P4's convert_to_tensor/cast fast-path tests. Both sides' imports are deduped/reordered (isort, single-line, case-sensitive) and both sides' new test methods are kept: P4's 20 convert_to_tensor/cast tests plus P3's test_get_device_reflects_device_scope, appended at the class tail. core.py and distribution_lib.py auto-merged clean.
Resolves the core_test.py conflict between the P3/P4-derived tests (P4's 20 convert_to_tensor/cast fast-path tests plus P3's test_get_device_reflects_device_scope, carried in via P4's merge commit 89bb4d0) and P8's 7 new slice/slice_update SymInt and 0-d-tensor-start-index tests. Both sides' imports are deduped/reordered (isort, single-line, case-sensitive; get_device was added by both sides, kept once; slice_update is P8-only) and ALL test methods from both sides are kept, name-collision check confirmed disjoint: P8's 7 slice tests appended at the class tail after the P3/P4 block, 35 test methods total, no duplicates. core.py and distribution_lib.py auto-merged clean. Full file verified green post-merge: KERAS_BACKEND=torch pytest keras/src/backend/torch/core_test.py -n 1 -> 35 passed.
torch Variable.value: replace the per-access maybe_use_symbolic_tensor closure with a module-level _maybe_meta_stub helper (removes one closure allocation per access); drop the redundant str() wrapping around get_device() (it always returns a plain str -- see _parse_device_input, which raises on non-str input) and use value.device.type instead of str(value.device) for the meta-device comparison. common Variable._maybe_autocast: check self._autocast before reading the autocast-scope global, so autocast=False variables (e.g. LN/BN gamma/beta) skip the global-state read entirely. Same semantics, fewer global reads. common Variable.dtype: reorder the same way, and delete the trailing standardize_dtype() call -- both branches already hold standardized strings (self._dtype is standardized in __init__; AutocastScope.dtype is standardized in AutocastScope.__init__), so the call was a proven no-op on every path. No control-flow branch added or removed; only evaluation order, allocation removal, and deletion of a provably idempotent standardize_dtype() call.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces several performance optimizations and robustness improvements across the Keras backend, particularly for PyTorch and variable handling. Key changes include optimizing autocast checks in Variable to avoid redundant global state lookups and standardized dtype calls, refactoring Variable.value in the PyTorch backend to use a module-level helper instead of allocating a fresh closure on every access (preventing TorchDynamo graph breaks), and adding fast paths in convert_to_tensor for PyTorch tensors already on the target device. Additionally, the PyTorch slice operation is improved to safely handle SymInt and avoid specializing dynamic dimensions during torch.export. Comprehensive unit tests have been added to verify these optimizations, closure avoidance, and correctness under tracing. No review comments were provided, so there is no feedback to evaluate.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #23297 +/- ##
==========================================
- Coverage 84.70% 84.09% -0.62%
==========================================
Files 465 465
Lines 69335 69759 +424
Branches 11407 11518 +111
==========================================
- Hits 58729 58662 -67
- Misses 7672 8153 +481
- Partials 2934 2944 +10
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Collaborator
Author
|
Closing this one. On the models benchmarked the isolated effect sits within measurement noise — it's a deterministic micro-optimization (removes redundant per-read closure allocations and dtype re-derivations in the Variable accessors) with no measurable end-to-end speedup, and it doesn't affect torch.compile. Trimming the #22561 series to the changes with demonstrable macro impact so review stays focused. |
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
3 participants
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.
Closes #23293
Why this was closed
This change showed no measurable macro benefit on either axis. Its
isolated eager effect is within measurement noise on the models tested,
and it is compile-neutral: 0 graph-break delta, with a deterministic
dynamo decomposition (see
PER-PR-COMPILE-SCORECARD-2026-07-19.md).
It was backed only by a deterministic micro-benchmark (call-count
reduction), which is not sufficient justification for the change on its
own.
Summary
Variableaccessors were redoing work already fixed at construction timeon every read. This hoists a per-access closure in torch's
Variable.valueto module level, drops a redundantstr()call in itsmeta-device check, and in the common
Variable._maybe_autocast/dtypeproperties, checks the cheaper instance-local flag before the global
autocast-scope read and deletes a trailing
standardize_dtype()call thatwas a proven no-op (both possible dtype sources are already standardized
before this property sees them). No control flow changes — only order of
checks, one allocation removed, and one dead call deleted.
Why it's safe
The closure/
str()fix is confined tobackend/torch/core.py; theautocast/dtype fix is confined to
backend/common/variables.py, the baseevery backend inherits without overriding these two methods — no
backend-branching added. Hunks are disjoint from the other in-flight
#22561 PRs touching the same files (#23182, #23188, #23189, #23190),
verified via
git merge-treedry-runs.Benchmarks
Deterministic micro-proof: closure allocations on
Variable.valueaccess,500 → 0 over a 500-access loop; autocast/dtype global-state reads down
33% on a mixed autocast workload. At macro scale, stacked on its parent
PRs the full stack is a real ~1.20–1.22x over master on LLM
forward/generate, consistent with the series' cumulative curve — but this
PR's own marginal delta on top of that already-optimized tree is noise on
all three workloads, the expected outcome for a ~19-line accessor fix.
The micro-proof above is the load-bearing evidence for the mechanism.
Equivalence checked in both states: CNN forward exact, LLM generate-32
token ids identical, LLM forward max diff 1.7e-06 (pre-existing recompute
noise, unrelated to this change).
Tests
variables_test.pyandcore_test.py: new coverage for both accessorsreturning identical values regardless of check order, the removed
standardize_dtypecall, and the meta-device comparison; full torchbackend/suite clean.Contributor Agreement