Skip to content

Fix lock/flags array size validation in matmul_reduce_scatter and all_reduce #482

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
mawad-amd merged 2 commits into from
Mar 26, 2026
+169 −1
Merged

Fix lock/flags array size validation in matmul_reduce_scatter and all_reduce #482

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
868faed
Initial plan
Copilot Mar 25, 2026
16cfd0a
Fix: add validation that lock/flags arrays are large enough for curre…
Copilot Mar 25, 2026
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
19 changes: 19 additions & 0 deletions iris/ccl/all_reduce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -840,6 +840,15 @@ def all_reduce(
"Spinlock variant requires workspace preparation. Call all_reduce_preamble before all_reduce."
)

num_pid_m = (M + config.block_size_m - 1) // config.block_size_m
num_pid_n = (N + config.block_size_n - 1) // config.block_size_n
total_tiles = num_pid_m * num_pid_n
if 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."
)

persistent_all_reduce_spinlock[(config.comm_sms,)](
input_tensor,
output_tensor,
Expand Down Expand Up @@ -870,6 +879,16 @@ def all_reduce(
"Ring variant requires workspace preparation. Call all_reduce_preamble before all_reduce."
)

num_pid_m = (M + config.block_size_m - 1) // config.block_size_m
num_pid_n = (N + config.block_size_n - 1) // config.block_size_n
total_tiles = num_pid_m * num_pid_n
total_flags = total_tiles * workspace.flags_per_tile
if 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."
)

# Calculate next rank in the ring for group support
# next_rank must be a global rank for iris RMA operations
if group is None:
Expand Down
16 changes: 15 additions & 1 deletion iris/ops/matmul_reduce_scatter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,12 @@ def matmul_reduce_scatter_preamble(
num_pid_n = (N + config.block_size_n - 1) // config.block_size_n
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."
)

if workspace.locks is None or workspace.locks.numel() != total_tiles:
Comment on lines +167 to 173

Copilot AI Mar 26, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 uses AI. Check for mistakes.
workspace.locks = shmem.zeros((total_tiles,), dtype=torch.int32)
else:
Expand Down Expand Up @@ -224,7 +230,15 @@ def matmul_reduce_scatter(

num_pid_m = (M + config.block_size_m - 1) // config.block_size_m
num_pid_n = (N + config.block_size_n - 1) // config.block_size_n
grid = (num_pid_m * num_pid_n,)
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."
)

grid = (total_tiles,)

even_k = K % config.block_size_k == 0

Expand Down
93 changes: 93 additions & 0 deletions tests/ccl/test_all_reduce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,3 +169,96 @@ def test_all_reduce_two_shot_distribution(distribution, dtype=torch.float32, M=1
import gc

gc.collect()


def test_all_reduce_spinlock_lock_too_small():
"""Test that ValueError is raised when the spinlock lock array is too small for current tile count.

Scenario: workspace is prepared with larger block sizes (fewer tiles), then all_reduce
is called with smaller block sizes (more tiles). workspace.matches() skips the preamble,
and the undersized lock array is detected.
"""
if not dist.is_initialized():
pytest.skip("torch.distributed not initialized")

heap_size = 2**33

Copilot AI Mar 26, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 uses AI. Check for mistakes.
shmem = iris.iris(heap_size)

M, N = 512, 512

iris_input = shmem.zeros((M, N), dtype=torch.float32)
iris_output = shmem.zeros((M, N), dtype=torch.float32)

shmem.barrier()

# Step 1: run preamble with larger block sizes → allocates a smaller lock array
config_large = Config(all_reduce_variant="spinlock", block_size_m=128, block_size_n=128)
workspace = shmem.ccl.all_reduce_preamble(iris_output, iris_input, config=config_large)

# Step 2: call all_reduce with smaller block sizes that need more tiles —
# workspace.matches() returns True (same shape/dtype/variant), preamble is skipped,
# and the undersized lock array is detected.
config_small = Config(all_reduce_variant="spinlock", block_size_m=64, block_size_n=64)
with pytest.raises(ValueError, match="Lock array too small"):
shmem.ccl.all_reduce(iris_output, iris_input, config=config_small, workspace=workspace)

shmem.barrier()
del shmem
import gc

gc.collect()


def test_all_reduce_ring_flags_too_small():
"""Test that ValueError is raised when the ring flags array is too small for current tile count.

Scenario: workspace is prepared with larger block sizes (fewer tiles), then all_reduce
is called with smaller block sizes (more tiles). workspace.matches() skips the preamble,
and the undersized flags array is detected.
"""
if not dist.is_initialized():
pytest.skip("torch.distributed not initialized")

heap_size = 2**33
shmem = iris.iris(heap_size)
world_size = shmem.get_num_ranks()

M, N = 512, 512

# 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}")
Comment on lines +228 to +234

Copilot AI Mar 26, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 uses AI. Check for mistakes.

iris_input = shmem.zeros((M, N), dtype=torch.float32)
iris_output = shmem.zeros((M, N), dtype=torch.float32)

shmem.barrier()

# Step 1: run preamble with larger block sizes → allocates a smaller flags array
config_large = Config(
all_reduce_variant="ring",
block_size_m=128,
block_size_n=block_size_n_large,
)
workspace = shmem.ccl.all_reduce_preamble(iris_output, iris_input, config=config_large)

# Step 2: call all_reduce with smaller block sizes that need more tiles —
# workspace.matches() returns True (same shape/dtype/variant), preamble is skipped,
# and the undersized flags array is detected.
config_small = Config(
all_reduce_variant="ring",
block_size_m=64,
block_size_n=block_size_n_small,
)
with pytest.raises(ValueError, match="Flags array too small"):
shmem.ccl.all_reduce(iris_output, iris_input, config=config_small, workspace=workspace)

shmem.barrier()
del shmem
import gc

gc.collect()
42 changes: 42 additions & 0 deletions tests/ops/test_matmul_reduce_scatter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,3 +175,45 @@ def test_matmul_reduce_scatter_semantics(dtype, atol, rtol):
import gc

gc.collect()


def test_matmul_reduce_scatter_lock_too_small():
"""Test that ValueError is raised when the lock array is too small for current tile count.

Scenario: workspace is prepared with larger block sizes (fewer tiles), then the
preamble is called again with smaller block sizes (more tiles). The undersized
lock array is detected and a ValueError is raised.
"""
if not dist.is_initialized():
pytest.skip("torch.distributed not initialized")

heap_size = 2**33

Copilot AI Mar 26, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copilot uses AI. Check for mistakes.
shmem = iris.iris(heap_size)

from iris.ops.config import FusedConfig
from iris.ops.matmul_reduce_scatter import matmul_reduce_scatter_preamble

M, N, K = 512, 512, 64
dtype = torch.float16

iris_A = shmem.zeros((M, K), dtype=dtype)
iris_B = shmem.zeros((K, N), dtype=dtype)
iris_C = shmem.zeros((M, N), dtype=dtype)

shmem.barrier()

# Step 1: run preamble with larger block sizes → allocates a smaller lock array
config_large = FusedConfig(block_size_m=128, block_size_n=128, block_size_k=32)
workspace = matmul_reduce_scatter_preamble(shmem, iris_C, iris_A, iris_B, config=config_large)

# Step 2: call preamble again with smaller block sizes that need more tiles —
# the preamble detects that workspace.locks is too small and raises ValueError.
config_small = FusedConfig(block_size_m=64, block_size_n=64, block_size_k=32)
with pytest.raises(ValueError, match="Lock array too small"):
matmul_reduce_scatter_preamble(shmem, iris_C, iris_A, iris_B, config=config_small, workspace=workspace)

shmem.barrier()
del shmem
import gc

gc.collect()
Loading