[XLA] Reorganize dynamic-size support into feature-based commits - #53
Open
stevenvar wants to merge 21 commits into
Open
[XLA] Reorganize dynamic-size support into feature-based commits#53stevenvar wants to merge 21 commits into
stevenvar wants to merge 21 commits into
Conversation
stevenvar
force-pushed
the
steven.for-serving-2.20-clean-history
branch
4 times, most recently
from
August 17, 2026 09:25
d53f0c2 to
d2858fa
Compare
stevenvar
changed the title
stevenvar
force-pushed
the
steven.for-serving-2.20-clean-history
branch
17 times, most recently
from
August 20, 2026 09:26
edd2ef7 to
e8b6182
Compare
Rewrite eligible scalar broadcast divisions to multiply by one shared reciprocal when fast math permits it. Preserve the original divide shape and cover both the simplifier and CPU compiler pipeline.
Use the same owned expression representation for TensorFlow TensorShape dimensions and XLA Shape dimensions. Add serialization, simplification, substitution, solving, and focused expression tests.
stevenvar
force-pushed
the
steven.for-serving-2.20-clean-history
branch
from
August 24, 2026 10:48
199fc4f to
a2f4dc3
Compare
Extend TensorFlow shape inference to preserve symbolic dimension identities and shape-derived tensor values. Keep unknown physical sizes separate from their symbolic expressions.
Assign one configurable symbolic variable to dynamic dimensions of top-level graph inputs while keeping inferred internal dimensions distinct. Add GraphProperties coverage for shared input symbols and non-input dimensions.
Carry dynamic expressions through XLA shape construction, HLO instructions, and shape inference. This keeps symbolic dimensions available while HLO graphs are built and transformed.
Add the flags, shape conversion helpers, argument metadata, and XlaExpression APIs needed to carry symbolic dimensions and shape-derived contents through tf2xla. Keep concrete kernel propagation separate so this commit defines the common contract first.
stevenvar
force-pushed
the
steven.for-serving-2.20-clean-history
branch
from
August 25, 2026 11:56
a2f4dc3 to
894136d
Compare
Update individual tf2xla kernel lowerings to derive and preserve symbolic shape and value expressions. Cover shape-producing, reshaping, slicing, sequence, partitioning, and related kernels using the shared expression plumbing.
Derive binary-op output expressions from the operand XLA shapes during broadcasting. Preserve the non-singleton expression when an operand expands from a singleton dimension.
Retain symbolic contents when shape-derived operations are rewritten or folded to constants. Add a dedicated rewrite path so optimization does not erase runtime shape dependencies.
Normalize compatible dynamic expressions around their smallest shared variable-covering subexpression before clustering. Reject cluster merges whose dynamic inputs and outputs cannot use one consistent symbolic core.
stevenvar
force-pushed
the
steven.for-serving-2.20-clean-history
branch
from
August 25, 2026 12:37
894136d to
55b5c4d
Compare
Map compiled XLA parameters back to their runtime TensorFlow inputs and solve the shared dynamic value from concrete dimensions. Filter broadcast-singleton evidence, pass the solved value through execution, and substitute it into dynamic output shapes after execution.
Normalize symbolic input metadata, pad solved dynamic values to configured compilation buckets, and include the resulting constants in executable cache signatures. Retry inconsistent dynamic requests with concrete static arguments instead of reusing an incompatible executable.
Prevent MLIR XLA kernels from silently consuming dynamic expressions they cannot preserve. Keep the existing MLIR path for operations without dynamic expression metadata.
Store the solved dynamic batch value in executable run options and expose it to CPU execution paths. Keep thunk and non-thunk execution consistent at the cluster boundary.
Generate LLVM loop bounds from the runtime batch dimension when an HLO dimension carries a dynamic expression. Add focused coverage for runtime-dependent bounds.
Teach CPU emitters to use runtime dynamic dimensions for affected operations instead of fixed padded extents. Preserve the padded allocation while limiting computation to the real shape.
Make CPU HLO profiling tolerate dynamic execution metadata and missing profile counters. Avoid invalid accesses while retaining profiling for supported instructions.
Add end-to-end coverage for symbolic shape propagation, clustering, compilation, and CPU execution. Exercise representative dynamic models across the complete TensorFlow-to-XLA path.
Replace power-of-two padding with configurable factor-preserving batch buckets. Use separate environment-controlled multipliers for small and large batches while retaining discovered bounds per cluster.
Avoid representing ceil-aligned reshape padding with ordinary integer division, which can simplify to an incorrect symbolic bound.
Solve the dynamic branch when the observed result is strictly greater than the constant bound, while keeping equality ambiguous.
stevenvar
force-pushed
the
steven.for-serving-2.20-clean-history
branch
from
August 25, 2026 14:17
55b5c4d to
28f7813
Compare
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.
Summary
This PR reconstructs the current
for-serving-2.20source tree from commit92148e0b108c3547ca73c2b08b2149b6bf349e75using a dependency-ordered, feature-based history.The new history separates:
The final source tree is byte-for-byte identical to the current
for-serving-2.20branch. This PR is intended for reviewing the proposed history rather than introducing new source changes.Verification
for-serving-2.20have the same Git tree hash:e92c08b099aa105e9e75fb7b26b862a071856983.