Skip to content

[Mosaic GPU] Enable column slicing for TMEM layouts with varying row dimensions.#38129

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_925327243
Open

[Mosaic GPU] Enable column slicing for TMEM layouts with varying row dimensions.#38129
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_925327243

Conversation

@copybara-service
Copy link
Copy Markdown

[Mosaic GPU] Enable column slicing for TMEM layouts with varying row dimensions.

The TMEMRef slicing logic is updated to correctly handle column offsets when the base tile shape's row dimension differs from TMEM_ROWS. This involves scaling the column index by the row folding factor.

I dropped the "partitioned across lanes or warps" check because I believe it is not relevant when slicing across tiles (i.e. column offset and slice shape are multiple of tile column size).

…dimensions.

The TMEMRef slicing logic is updated to correctly handle column offsets when the base tile shape's row dimension differs from TMEM_ROWS. This involves scaling the column index by the row folding factor.

I dropped the "partitioned across lanes or warps" check because I believe it is not relevant when slicing across tiles (i.e. column offset and slice shape are multiple of tile column size).

PiperOrigin-RevId: 925327243
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant