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:
co-authored by
Claude Opus 5.5
parent
68ecaf9e81
commit
93d1cddd27
@@ -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.
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user