[XLA] Fix dynamic expression edge cases and batch bucketing by stevenvar · Pull Request #54 · stevenvar/tensorflow

stevenvar · 2026-08-17T15:23:01Z

Summary

treat unknown dynamic expressions as neither constants nor valid dynamic expressions, and avoid emitting them as the integer value -1
preserve division expressions with a zero divisor instead of simplifying {0 / 0} to {1}
keep configured XLA batch buckets immutable so batch selection is deterministic and safe for concurrent compilation requests
remove stale debug-only output and clarify the StridedSlice expression fallback

Tests

//tensorflow/compiler/jit:xla_batch_matcher_test checks request-order independence and immutable configured buckets
@local_xla//xla:shape_test checks unknown-expression classification and zero-divisor simplification
@local_xla//xla/service/llvm_ir:llvm_util_test checks that unknown expressions are not emitted as constants
related XLA compiler, runtime, shape-inference, GraphProperties, and constant-folding tests pass without cached test results

Update tf2xla kernels to derive and forward symbolic shape and value expressions. Gate symbolic content propagation behind its dedicated configuration flag.

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.

Map runtime arguments to symbolic cluster inputs, solve their dynamic values, and build padded compilation arguments for cache lookup. Add batch bucketing and runtime boundary tests for dynamic cluster compilation.

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.fix-clean-history-review-findings branch from 18b123b to 9ad0149 Compare August 18, 2026 08:21

stevenvar force-pushed the steven.for-serving-2.20-clean-history branch from 4b41aa2 to 51b14f8 Compare August 18, 2026 09:52

stevenvar force-pushed the steven.fix-clean-history-review-findings branch from 9ad0149 to 749308d Compare August 18, 2026 09:52

stevenvar force-pushed the steven.for-serving-2.20-clean-history branch from 51b14f8 to 4f36467 Compare August 18, 2026 09:55

stevenvar force-pushed the steven.fix-clean-history-review-findings branch from 749308d to 575bc44 Compare August 18, 2026 09:55

stevenvar force-pushed the steven.for-serving-2.20-clean-history branch from 4f36467 to 1edf990 Compare August 18, 2026 11:45

stevenvar force-pushed the steven.fix-clean-history-review-findings branch from 575bc44 to 3d05104 Compare August 18, 2026 11:45

stevenvar force-pushed the steven.for-serving-2.20-clean-history branch from 1edf990 to 54a0c74 Compare August 18, 2026 12:49

stevenvar force-pushed the steven.fix-clean-history-review-findings branch from 3d05104 to 41e8ef0 Compare August 18, 2026 12:52

stevenvar added 11 commits August 18, 2026 15:27

Propagate dynamic expressions through tf2xla kernels

15f1c18

Update tf2xla kernels to derive and forward symbolic shape and value expressions. Gate symbolic content propagation behind its dedicated configuration flag.

Preserve dynamic expressions through binary broadcasting

797aa93

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.

Preserve shape-derived expressions through constant folding

1b6dcea

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.

Cluster graphs using normalized dynamic subexpressions

a962070

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.

Compile and execute clusters with dynamic expressions

3d1a28c

Map runtime arguments to symbolic cluster inputs, solve their dynamic values, and build padded compilation arguments for cache lookup. Add batch bucketing and runtime boundary tests for dynamic cluster compilation.

Guard MLIR lowering for dynamic shape expressions

9a26c5d

Prevent MLIR XLA kernels from silently consuming dynamic expressions they cannot preserve. Keep the existing MLIR path for operations without dynamic expression metadata.

Pass dynamic batch dimensions through CPU execution

676bfdf

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 runtime-dependent CPU loop bounds

3ac507f

Generate LLVM loop bounds from the runtime batch dimension when an HLO dimension carries a dynamic expression. Add focused coverage for runtime-dependent bounds.

Support dynamic dimensions in CPU operation emitters

6eb28c5

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.

Prevent CPU HLO profiling crashes

ce28e7a

Make CPU HLO profiling tolerate dynamic execution metadata and missing profile counters. Avoid invalid accesses while retaining profiling for supported instructions.

Add dynamic expression integration tests

c3d2540

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 54a0c74 to c3d2540 Compare August 18, 2026 14:37

Fix dynamic expression edge cases and batch bucketing

85612da

stevenvar force-pushed the steven.fix-clean-history-review-findings branch from 41e8ef0 to 85612da Compare August 18, 2026 14:38

stevenvar force-pushed the steven.for-serving-2.20-clean-history branch 2 times, most recently from edd2ef7 to e8b6182 Compare August 20, 2026 09:26

Navigation Menu

[XLA] Fix dynamic expression edge cases and batch bucketing - #54

[XLA] Fix dynamic expression edge cases and batch bucketing#54
stevenvar wants to merge 12 commits into
steven.for-serving-2.20-clean-historyfrom
steven.fix-clean-history-review-findings

stevenvar commented Aug 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

stevenvar commented Aug 17, 2026

Summary

Tests

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant