[Pallas] Fix scalar .begin index not collapsing tensor dimensions#1972
Draft
[Pallas] Fix scalar .begin index not collapsing tensor dimensions#1972
Conversation
meta-cla
bot
added
the
CLA Signed
norx1991
force-pushed
the
yifeixu/pallas-scalar-index-test-only
branch
from
e419f95 to
643a8d4
Compare
norx1991
force-pushed
the
yifeixu/pallas-scalar-index-fix
branch
from
a52bee5 to
2180150
Compare
norx1991
force-pushed
the
yifeixu/pallas-scalar-index-test-only
branch
from
643a8d4 to
8a58474
Compare
norx1991
force-pushed
the
yifeixu/pallas-scalar-index-fix
branch
from
2180150 to
7eff4ec
Compare
norx1991
force-pushed
the
yifeixu/pallas-scalar-index-test-only
branch
from
8a58474 to
c215b44
Compare
norx1991
force-pushed
the
yifeixu/pallas-scalar-index-fix
branch
from
7eff4ec to
47aff58
Compare
norx1991
force-pushed
the
yifeixu/pallas-scalar-index-test-only
branch
from
c215b44 to
8d9cb31
Compare
norx1991
force-pushed
the
yifeixu/pallas-scalar-index-fix
branch
from
47aff58 to
042d204
Compare
norx1991
force-pushed
the
yifeixu/pallas-scalar-index-test-only
branch
from
8d9cb31 to
a2ebcf1
Compare
norx1991
force-pushed
the
yifeixu/pallas-scalar-index-fix
branch
from
042d204 to
10b4184
Compare
norx1991
changed the base branch from
yifeixu/pallas-scalar-index-test-only
to
main
norx1991
force-pushed
the
yifeixu/pallas-scalar-index-fix
branch
from
10b4184 to
32711fc
Compare
When a kernel uses tile.begin as a scalar index, the Pallas codegen emitted ':' keeping the dimension alive. Operations like .T then failed with a rank mismatch (e.g. 2D permutation on a 3D tensor). Fixes: 1. compile_environment.py: Add TileBeginOrigin to get_block_id's accepted types so .begin SymInts get a block_id. 2. memory_ops.py: Add _is_scalar_tile_offset helper and modify _pallas_index_str to emit literal 0 for .begin indices (collapsing the dim), mirroring Triton's scalar SymInt handling. Record in dim_map so BlockSpec tiles the dim correctly. Skip out_pos increment for collapsed dims so None positions stay correct. 3. backend.py: Skip alignment constraints when no tensor dim mapping is found, fixing a fallback that incorrectly assumed spec ordering matches tensor dim ordering.
norx1991
force-pushed
the
yifeixu/pallas-scalar-index-fix
branch
from
32711fc to
441d03e
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
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.
Summary
TileBeginOrigintoget_block_id's accepted types so.beginSymInts get a block_id0for.beginindices in_pallas_index_strto collapse the dim, mirroring Triton's scalar SymInt handlingadjust_block_size_constraintsfallback that incorrectly assumed spec ordering matches tensor dim orderingout_posincrement for collapsed dims soNonepositions stay correct