[XLA] Use logical dimensions in dynamic CPU kernels - #61
Open
stevenvar wants to merge 2 commits into
Open
Conversation
Evaluate symbolic convolution and dot dimensions from the runtime root while retaining padded storage bounds. Preserve flattened expressions when 1x1 convolutions are rewritten to dots so optimized CPU execution does not process padded elements.
Match the packed runtime layout used by generated indexing and the existing Eigen matmul path. Test dynamic matrix dimensions in both layouts, including thunk reuse across changing sizes and untouched spare capacity.
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
1 participant
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.
Dynamic CPU execution reserves storage using padded shape bounds, but convolution and dot kernels can use those bounds as their computation sizes. This makes them process padded elements instead of stopping at the logical runtime dimensions.
This change evaluates symbolic dimensions from the runtime dynamic value. Convolution uses the logical batch size, while dot uses the logical matrix and batch dimensions for both computation and buffer strides. Allocation bounds still reserve enough capacity, while runtime matrices remain packed contiguously. For example, a runtime batch of 2 in storage sized for 4 computes two batches and leaves the remaining output storage untouched.
The convolution-to-dot rewrite also preserves flattened dimension expressions, so a 1x1 convolution keeps the same runtime dimensions after optimization.
Tests
ConvolutionThunkTest.UsesLogicalDynamicBatch: computes a logical batch of 2 within a bound of 4 and checks that padded output elements remain untouched.DotThunkTest.UsesLogicalDynamicDimension: checks the same behavior for a dot operation with a symbolic matrix dimension.DotThunkDynamicStrideTest.UsesLogicalMatrixStrides: checks dynamicm,n, andkin both row-major and column-major layouts. Each case reuses one thunk at runtime sizes 2, 4, then 1, verifies distinct results for both batches, and checks that spare output storage remains untouched.AlgebraicSimplifierTest.ConvertDynamicConvToMatmulPreservesFlattenedExpression: checks that converting a 1x1 convolution to a dot preserves the flattened expression and evaluates it correctly for a different runtime batch.Validation
All 24 tests in
@local_xla//xla/backends/cpu/runtime:dot_thunk_testpass locally on macOS with cached test results disabled. The six dynamic stride cases were also run before the stride fix and all failed, confirming that they catch the incorrect offsets.