Skip to content

TileOP: fix RedRows per-PE semantics for Local-A/Shared-B groupM #86

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
VV0003 merged 1 commit into from
Sep 8, 2026
Merged

TileOP: fix RedRows per-PE semantics for Local-A/Shared-B groupM #86

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 21 additions & 3 deletions include/jcore/template_asm.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6335,8 +6335,14 @@ PTO_SHARED_INLINE void TMATMUL_ACC(tile_shape_d &d, tile_shape_c &c, tile_shape_
// Reduction outputs (RowMax/GroupMax) reduce the per-PE D rows: ASL
// MatrixRowMaxResult iterates input.valid_rows, which for a cooperative
// TMATMUL is the per-PE clamp of group_M, not the core-total group_M.
// For cooperative forms the reduction/output rows are the per-PE M.
// Shared-A: per-PE block = rows_per_pe(group_M) derived from the core
// total. Local-A/Shared-B: A::ValidRow already IS the per-PE shard size
// (M_per_PE itself, per ADR-0100), so it must be used directly — feeding
// it back through rows_per_pe() would map M_per_PE=32 to 16 and break
// CubeM32 group_M>64 configurations.
constexpr int RedRows =
(is_shared_tile_v<tile_shape_a> || is_shared_tile_v<tile_shape_b>)
(is_shared_tile_v<tile_shape_a>)
? pto_matmul_detail::cooperative_group_m_rows_per_pe(EffectiveM)
: EffectiveM;

Expand Down Expand Up @@ -6443,8 +6449,14 @@ PTO_SHARED_INLINE void TMATMUL(tile_shape_d &d, tile_shape_a &a,
? tile_shape_a::ValidCol : tile_shape_a::ValidRow;
// Reduction outputs still use the per-PE clamp of group_M for cooperative
// Local-A/Shared-B; the explicit groupM only changes the encoded LB0 M.
// For cooperative forms the reduction/output rows are the per-PE M.
// Shared-A: per-PE block = rows_per_pe(group_M) derived from the core
// total. Local-A/Shared-B: A::ValidRow already IS the per-PE shard size
// (M_per_PE itself, per ADR-0100), so it must be used directly — feeding
// it back through rows_per_pe() would map M_per_PE=32 to 16 and break
// CubeM32 group_M>64 configurations.
constexpr int RedRows =
(is_shared_tile_v<tile_shape_a> || is_shared_tile_v<tile_shape_b>)
(is_shared_tile_v<tile_shape_a>)
? pto_matmul_detail::cooperative_group_m_rows_per_pe(EffectiveM)
: EffectiveM;

Expand Down Expand Up @@ -6564,8 +6576,14 @@ TMATMUL(tile_shape_d &d, tile_shape_a &a,
// Reduction outputs (RowMax/GroupMax) reduce the per-PE D rows: ASL
// MatrixRowMaxResult iterates input.valid_rows, which for a cooperative
// TMATMUL is the per-PE clamp of group_M, not the core-total group_M.
// For cooperative forms the reduction/output rows are the per-PE M.
// Shared-A: per-PE block = rows_per_pe(group_M) derived from the core
// total. Local-A/Shared-B: A::ValidRow already IS the per-PE shard size
// (M_per_PE itself, per ADR-0100), so it must be used directly — feeding
// it back through rows_per_pe() would map M_per_PE=32 to 16 and break
// CubeM32 group_M>64 configurations.
constexpr int RedRows =
(is_shared_tile_v<tile_shape_a> || is_shared_tile_v<tile_shape_b>)
(is_shared_tile_v<tile_shape_a>)
? pto_matmul_detail::cooperative_group_m_rows_per_pe(EffectiveM)
: EffectiveM;

Expand Down
Loading