Skip to content

Align the m dimension of pack_a outputs - #11136

Open
copybara-service[bot] wants to merge 1 commit into
masterfrom
test_976076613
Open

Align the m dimension of pack_a outputs#11136
copybara-service[bot] wants to merge 1 commit into
masterfrom
test_976076613

Conversation

@copybara-service

Copy link
Copy Markdown
Contributor

Align the m dimension of pack_a outputs

transpose_a dot kernels read multiples of tile_m rows at a time. This change adds a new dimension that correctly requires the padding, instead of relying on incidental allocation padding.

Additionally fixes two issues with transpose_a that broke multi-reduction and convolution targets:

  1. Corrects multi-reduction dimension stride ordering in define_transpose_a so that intermediate reduction dimensions (k2, k3) precede k1_outer.
  2. Disallows rewriting transpose_a(stencil_copy) to stencil_copy(transpose_a) when tile_m > 1, since stencil_copy cannot slide across tiled dimensions.

`transpose_a` dot kernels read multiples of `tile_m` rows at a time. This change adds a new dimension that correctly requires the padding, instead of relying on incidental allocation padding.

Additionally fixes two issues with `transpose_a` that broke multi-reduction and convolution targets:
1. Corrects multi-reduction dimension stride ordering in `define_transpose_a` so that intermediate reduction dimensions (k2, k3) precede `k1_outer`.
2. Disallows rewriting `transpose_a(stencil_copy)` to `stencil_copy(transpose_a)` when `tile_m > 1`, since `stencil_copy` cannot slide across tiled dimensions.

PiperOrigin-RevId: 976076613
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