BlockLayout2D.build: check capacity before generating regular origins

c397e08 moved the 64-bit capacity check after np.arange(nby*nbx), so an
over-limit regular layout allocated before rejecting. Check arithmetic
block count (regular) or validated selected count (sparse) first; mock
regression asserts no np.arange before the ValueError.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-10-16 03:27:13 +00:00
co-authored by Claude Opus 5.5
parent 68ecaf9e81
commit 93d1cddd27
2 changed files with 17 additions and 6 deletions
+10 -6
View File
@@ -140,13 +140,17 @@ class BlockLayout2D:
nby, nbx = ny // by, nx // bx
ly, lx = by + 2 * halo, bx + 2 * halo
cells = ly * lx
if origins is None:
jy, jx = np.divmod(np.arange(nby * nbx), nbx)
origins = np.stack([jy * by, jx * bx], axis=1)
selected = _selected_blocks(origins, ny, nx, by, bx)
origins = np.array(origins, dtype=np.int64).reshape(-1, 2) # own copy, never caller storage
if (len(origins) + 1) * NQ * cells > np.iinfo(np.int64).max // 2:
if origins is not None:
selected = _selected_blocks(origins, ny, nx, by, bx)
origins = np.array(origins, dtype=np.int64).reshape(-1, 2) # own copy, never caller storage
# Capacity from arithmetic before any grid-sized allocation; hole =
# nblocks * cells is below this bound too.
if (nby * nbx if origins is None else len(origins)) * NQ * cells > np.iinfo(np.int64).max // 2:
raise ValueError("blocked storage exceeds 64-bit indexing")
if origins is None:
selected = np.arange(nby * nbx)
jy, jx = np.divmod(selected, nbx)
origins = np.stack([jy * by, jx * bx], axis=1)
# Grid block -> stored block, or len(origins) for an absent block, so
# absent cells land exactly on hole = nblocks * cells.
+7
View File
@@ -1,6 +1,7 @@
"""CPU block layout contracts: ownership, halo pairs, packing and tagged pull."""
import sys
import unittest
from unittest import mock
from types import SimpleNamespace
import numpy as np
@@ -280,6 +281,12 @@ class BlockLayoutTests(unittest.TestCase):
lx = layout.local_shape[1]
self.assertEqual(int(mapped.write[0]), int(layout.slot(3 * lx + 0, 1)))
def test_capacity_rejected_before_allocation(self):
fail = AssertionError("grid-sized allocation before the capacity check")
with mock.patch.object(np, "arange", side_effect=fail):
with self.assertRaises(ValueError):
BlockLayout2D.build(2 ** 30, 2 ** 30, 1, 1)
def test_tables_are_readonly(self):
layout = BlockLayout2D.build(6, 8, 3, 4)
for array in (layout.origins, layout.owner, layout.halo_dst, layout.halo_src, layout.halo_kind):