Fix lock/flags array size validation in matmul_reduce_scatter and all_reduce

[Copilot]

FusedWorkspace.matches() does not check block sizes, so a workspace prepared with larger blocks (fewer tiles) can be silently reused with smaller blocks (more tiles), causing the kernel to write lock/flag entries past the end of the allocated array and corrupt adjacent symmetric heap objects.

Changes

iris/ops/matmul_reduce_scatter.py

• Preamble: raises ValueError before the allocation check if workspace.locks is set but smaller than total_tiles — catches direct preamble reuse with a differently-sized pre-allocation
• Main function: adds the same guard as a defensive post-preamble check (consistent with the fix pattern in matmul_all_reduce)

iris/ccl/all_reduce.py

• Spinlock variant: raises ValueError if workspace.locks.numel() < total_tiles — fires when workspace.matches() skips the preamble but block sizes shrank
• Ring variant: raises ValueError if workspace.flags.numel() < total_flags — same scenario for the flags array

Error message in all cases:

ValueError: Lock array too small: have 16 but need 64.
Pre-allocate workspace with the smallest block sizes you intend to use.

Tests

• test_matmul_reduce_scatter_lock_too_small: preamble with bm=128/bn=128 → preamble again with bm=64/bn=64, expects ValueError
• test_all_reduce_spinlock_lock_too_small: preamble with large blocks, all_reduce with small blocks (preamble skipped by matches()), expects ValueError
• test_all_reduce_ring_flags_too_small: same pattern for ring flags array

Original prompt

---

This section details on the original issue you should resolve

<issue_title>matmul_reduce_scatter and ccl.all_reduce: no validation that lock/flags arrays are large enough for current tile count</issue_title>
<issue_description>## Bug

Same class of bug as #463, but in two other operations:

1. iris/ops/matmul_reduce_scatter.py

The preamble allocates a locks array sized by total_tiles = num_pid_m * num_pid_n (lines 163-170). If the workspace is reused with smaller block sizes (more tiles), the kernel writes past the end of the lock array.

2. iris/ccl/all_reduce.py

• Spinlock variant (lines 118-124): Allocates locks array sized by total_tiles.
• Ring variant (lines 92-109): Allocates flags array sized by total_tiles * flags_per_tile.

Both can overflow if block sizes change between calls.

Root cause

FusedWorkspace.matches() in iris/ops/workspace.py (lines 47-75) checks operation, shape, dtype, world_size, and variant — but does not check block sizes. When block sizes change, the preamble is skipped and the undersized array is silently reused.

Impact

Silent symmetric heap corruption. No error raised. Wrong numerical results on subsequent kernel calls.

Fix

Add runtime validation in each operation (matching the fix in #463 / PR #480):

iris/ops/matmul_reduce_scatter.py:

total_tiles = num_pid_m * num_pid_n
if workspace.locks is not None and workspace.locks.numel() < total_tiles:
    raise ValueError(
        f"Lock array too small: have {workspace.locks.numel()} but need {total_tiles}. "
        f"Pre-allocate workspace with the smallest block sizes you intend to use."
    )

iris/ccl/all_reduce.py (spinlock variant):

if workspace.locks is not None and workspace.locks.numel() < total_tiles:
    raise ValueError(
        f"Lock array too small: have {workspace.locks.numel()} but need {total_tiles}. "
        f"Pre-allocate workspace with the smallest block sizes you intend to use."
    )

iris/ccl/all_reduce.py (ring variant):

total_flags = total_tiles * flags_per_tile
if workspace.flags is not None and workspace.flags.numel() < total_flags:
    raise ValueError(
        f"Flags array too small: have {workspace.flags.numel()} but need {total_flags}. "
        f"Pre-allocate workspace with the smallest block sizes you intend to use."
    )

Add corresponding tests for each operation.

Components

• iris/ops/matmul_reduce_scatter.py
• iris/ccl/all_reduce.py

Related

• matmul_all_reduce: no validation that lock array is large enough for current tile count #463 (same bug in matmul_all_reduce, fixed by Fix lock array size validation in matmul_all_reduce #480)</issue_description>

Comments on the Issue (you are @copilot in this section)

• Fixes matmul_reduce_scatter and ccl.all_reduce: no validation that lock/flags arrays are large enough for current tile count #481

---

💡 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.

[Copilot]

Pull request overview

Adds runtime validation to prevent out-of-bounds writes when reusing workspaces whose lock/flag buffers were allocated for fewer tiles than a subsequent call requires (e.g., when block sizes shrink and workspace.matches() skips re-preparation).

Changes:

• Add “too small” guards for locks/flags in matmul_reduce_scatter and all_reduce (spinlock + ring).
• Add regression tests that exercise workspace reuse across differing block sizes and assert ValueError.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.

File	Description
iris/ops/matmul_reduce_scatter.py	Adds tile-count-based validation for the workspace lock array in both preamble and main entrypoint.
iris/ccl/all_reduce.py	Adds lock/flag size validation for spinlock and ring variants to prevent buffer overruns when preamble is skipped.
tests/ops/test_matmul_reduce_scatter.py	Adds a regression test for undersized lock reuse across different block sizes in preamble.
tests/ccl/test_all_reduce.py	Adds regression tests for undersized lock/flags reuse when workspace.matches() causes preamble to be skipped.

[Copilot]

The new error message instructs users to “pre-allocate workspace with the smallest block sizes”, which implies keeping a larger-than-needed lock array around for later reuse. However, the subsequent != total_tiles condition will shrink an oversized workspace.locks by allocating a new smaller array. In symmetric-heap scenarios this can also increase allocation churn and heap growth (old allocations may not be reclaimed), and it undermines the “pre-allocate small blocks once” guidance by potentially making the next small-block call fail. Recommendation (mandatory): only (re)allocate when workspace.locks is None or workspace.locks.numel() < total_tiles; when the array is larger, keep it and just zero it (optionally only the prefix used by total_tiles if supported).

Suggested change

	if workspace.locks is not None and workspace.locks.numel() < total_tiles:
	raise ValueError(
	f"Lock array too small: have {workspace.locks.numel()} but need {total_tiles}. "
	f"Pre-allocate workspace with the smallest block sizes you intend to use."
	)
	if workspace.locks is None or workspace.locks.numel() != total_tiles:
	if workspace.locks is None or workspace.locks.numel() < total_tiles:

[Copilot]

This logic can produce nonstandard block_size_n values (e.g., world_size=3 yields 126 and 63). If the all-reduce ring kernels/config validation expect specific block sizes (common in tiled kernels), this can make the test unexpectedly fail or behave inconsistently across environments. Recommendation (moderate): restrict the test to the intended “use 64 and 128” case by skipping unless world_size divides both 64 and 128 (e.g., 64 % world_size == 0 and 128 % world_size == 0), and then set block_size_n_large=128, block_size_n_small=64.

Suggested change

	# Choose block_size_n values divisible by world_size for both configs
	# Use 128 and 64 which are divisible by typical world sizes (1, 2, 4, 8)
	block_size_n_large = (128 // world_size) * world_size
	block_size_n_small = (64 // world_size) * world_size
	if block_size_n_large == 0 or block_size_n_small == 0 or block_size_n_large == block_size_n_small:
	del shmem
	pytest.skip(f"Cannot create two distinct block sizes divisible by world_size={world_size}")
	# Restrict to the intended case: use block_size_n values 128 and 64,
	# and only run the test when both are exactly divisible by world_size.
	if 128 % world_size != 0 or 64 % world_size != 0:
	del shmem
	pytest.skip(
	f"Skipping: world_size={world_size} does not divide both 128 and 64 for block_size_n"
	)
	block_size_n_large = 128
	block_size_n_small = 64

[Copilot]

The new tests add additional allocations of a very large symmetric heap (2**33) per test, which can significantly increase CI runtime and memory pressure (and may lead to OOM on smaller GPUs/nodes). Recommendation (optional): consider (1) reducing heap_size to the smallest value that still reproduces the condition (since these tests only need lock/flag arrays), and/or (2) using a shared fixture to reuse a single shmem across related tests when safe to do so.

Suggested change

	heap_size = 2**33
	heap_size = 2**26

[Copilot]

Similar to the all-reduce tests, this regression test allocates a very large heap for a scenario that primarily targets the lock-array sizing logic. Recommendation (optional): lower heap_size to the minimum required for the test or use a fixture/shared setup to reduce repeated heavy allocations and lower the likelihood of CI resource issues.

Suggested change

	heap_size = 2**33
	# Use a moderately sized heap sufficient for this test without stressing CI resources.
	heap_size = 2**24