[XLA] Reorganize dynamic-size support into feature-based commits - #53
Open
stevenvar wants to merge 20 commits into
Open
[XLA] Reorganize dynamic-size support into feature-based commits#53stevenvar wants to merge 20 commits into
stevenvar wants to merge 20 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
14 times, most recently
from
August 18, 2026 12:49
1edf990 to
54a0c74
Compare
Use the same owned expression representation for TensorFlow TensorShape dimensions and XLA Shape dimensions. Add serialization, simplification, substitution, solving, and focused expression tests.
Extend TensorFlow shape inference to preserve symbolic dimension identities and shape-derived tensor values. Keep unknown physical sizes separate from their symbolic expressions.
Carry dynamic expressions through XLA shape construction, HLO instructions, and shape inference. This keeps symbolic dimensions available while HLO graphs are built and transformed.
stevenvar
force-pushed
the
steven.for-serving-2.20-clean-history
branch
from
August 18, 2026 14:37
54a0c74 to
c3d2540
Compare
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.
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.
stevenvar
force-pushed
the
steven.for-serving-2.20-clean-history
branch
from
August 19, 2026 15:07
c3d2540 to
edd2ef7
Compare
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.
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.
stevenvar
force-pushed
the
steven.for-serving-2.20-clean-history
branch
from
August 20, 2026 09:26
edd2ef7 to
e8b6182
Compare
Previously, we used a padding policy of powers of 2, which would result in the loss of the multiple relationship. This commit changes the padding policy to a multiple padding with k value. All padding shapes are recorded for each function, if there exists a upper_bound for k*dim, select the upper_bound shape as padding shape
TF_XLA_BATCH_SMALL_FACTOR is padding multiplier for size less than 10; TF_XLA_BATCH_LARGE_FACTOR is padding multiplier for size not less than 10.
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.
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
2 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.