Skip to content

Fix iris.store hint: use (1, BLOCK_N) instead of (BLOCK_M, BLOCK_N) across all CCL/X collectives#437

Merged
mawad-amd merged 2 commits intomuhaawad/hintsfrom
copilot/sub-pr-432-again
Mar 7, 2026
Merged

Fix iris.store hint: use (1, BLOCK_N) instead of (BLOCK_M, BLOCK_N) across all CCL/X collectives#437
mawad-amd merged 2 commits intomuhaawad/hintsfrom
copilot/sub-pr-432-again

Conversation

Copy link
Contributor

Copilot AI commented Mar 7, 2026

hint=(BLOCK_SIZE_M, BLOCK_SIZE_N) asserts 2D cross-row contiguity, which is only valid when BLOCK_N == N. When BLOCK_N < N, stride_m > BLOCK_N so rows aren't contiguous — getOrderFromContiguity picks dim-0 for vectorization and emits scalar writes to wrong addresses, leaving zeros in the output (e.g., 256×128 tensor with block 32×16).

The fix from PR #434 (x.all_gather) was not applied consistently; the same bug existed in all other collective iris.store call sites.

Changes

  • iris/ccl/all_gather.py (2 sites) — hint=(BLOCK_SIZE_M, BLOCK_SIZE_N)hint=(1, BLOCK_SIZE_N)
  • iris/ccl/all_reduce.py (4 sites) — hint=(BLOCK_SIZE_M, BLOCK_SIZE_N)hint=(1, BLOCK_SIZE_N)
  • iris/ccl/all_to_all.py (2 sites) — hint=(BLOCK_SIZE_M, BLOCK_SIZE_N)hint=(1, BLOCK_SIZE_N)
  • iris/x/all_reduce.py (1 site) — hint=(tile.block_m, tile.block_n)hint=(1, tile.block_n)

hint=(1, BLOCK_N) asserts only per-row contiguity (BLOCK_N consecutive elements within a row), which holds unconditionally regardless of tile shape.


💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.

… all store calls

Co-authored-by: mawad-amd <112003944+mawad-amd@users.noreply.github.com>
Copilot AI changed the title [WIP] Add hints in CCL and X implementations Fix iris.store hint: use (1, BLOCK_N) instead of (BLOCK_M, BLOCK_N) across all CCL/X collectives Mar 7, 2026
@mawad-amd mawad-amd marked this pull request as ready for review March 7, 2026 20:33
@mawad-amd mawad-amd requested review from BKP and neoblizz as code owners March 7, 2026 20:33
@mawad-amd mawad-amd merged commit 1c55396 into muhaawad/hints Mar 7, 2026
@mawad-amd mawad-amd deleted the copilot/sub-pr-432-again branch March 7, 2026 20:33
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.

2 participants