"""Tests for the pytorch-compat Tensor API extensions. Covers the new ``torch.from_numpy`` factory and ``Tensor.numpy``, ``Tensor.copy_`` methods used by the unified ``ccl_allreduce`` bench. """ from __future__ import annotations import numpy as np import pytest from kernbench.policy.placement.dp import DPPolicy from kernbench.runtime_api.bench_runner import run_bench from kernbench.runtime_api.types import resolve_device from kernbench.sim_engine.engine import GraphEngine from kernbench.topology.builder import resolve_topology def _engine_factory(topology, device): return GraphEngine(getattr(topology, "topology_obj", topology), enable_data=True) def _run_with(bench_body): topo = resolve_topology("topology.yaml") return run_bench( topology=topo, bench_fn=bench_body, device=resolve_device("all"), engine_factory=_engine_factory, ) # ── from_numpy ────────────────────────────────────────────────────── def test_from_numpy_creates_host_tensor(): """torch.from_numpy returns a kernbench Tensor with the array stored in its host buffer (not deployed to any PE).""" def body(torch): arr = np.arange(8, dtype=np.float16).reshape(1, 8) h = torch.from_numpy(arr) # Host tensor has shape/dtype matching the array. assert h.shape == (1, 8) assert h.dtype == "f16" # numpy() round-trips the host buffer. assert np.array_equal(h.numpy(), arr) # No deploy → no real shards. assert h._handle is None # Submit a no-op so run_bench has at least one handle. torch.zeros((1, 8), dtype="f16", dp=DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1), name="dummy") _run_with(body) # ── single-PE replicated tensor ───────────────────────────────────── def test_copy_and_numpy_single_pe(): """copy_ from a numpy array, then numpy() round-trips correctly on a single-PE (no real sharding) tensor.""" def body(torch): dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1) t = torch.zeros((1, 16), dtype="f16", dp=dp, name="t") src = np.arange(16, dtype=np.float16).reshape(1, 16) t.copy_(torch.from_numpy(src)) gathered = t.numpy() assert gathered.shape == (1, 16) assert np.array_equal(gathered, src) _run_with(body) # ── multi-PE column-wise sharding (1 cube) ────────────────────────── def test_copy_and_numpy_multi_pe_column_wise(): """copy_ splits across 8 PEs in one cube, numpy() reassembles.""" def body(torch): n_pe = 8 dp = DPPolicy(cube="replicate", pe="column_wise", num_cubes=1, num_pes=n_pe) t = torch.zeros((1, n_pe * 4), dtype="f16", dp=dp, name="t") src = np.arange(n_pe * 4, dtype=np.float16).reshape(1, n_pe * 4) t.copy_(torch.from_numpy(src)) gathered = t.numpy() assert gathered.shape == (1, n_pe * 4) assert np.array_equal(gathered, src) # Sanity: there really were 8 shards. assert len(t._handle.shards) == n_pe _run_with(body) # ── multi-cube sharding ───────────────────────────────────────────── def test_copy_and_numpy_multi_cube(): """copy_ across 2 cubes (16 PEs total), numpy() reassembles.""" def body(torch): n_pe_per_cube = 8 n_cubes = 2 total = n_cubes * n_pe_per_cube # 16 dp = DPPolicy(cube="column_wise", pe="column_wise", num_cubes=n_cubes) t = torch.zeros((1, total * 4), dtype="f16", dp=dp, name="t") src = np.arange(total * 4, dtype=np.float16).reshape(1, total * 4) t.copy_(torch.from_numpy(src)) gathered = t.numpy() assert np.array_equal(gathered, src) assert len(t._handle.shards) == total _run_with(body) # ── shape mismatch raises ─────────────────────────────────────────── def test_copy_shape_mismatch_raises(): """copy_ with mismatched shapes raises ValueError.""" def body(torch): dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1) t = torch.zeros((1, 8), dtype="f16", dp=dp, name="t") src = np.zeros((1, 16), dtype=np.float16) with pytest.raises(ValueError, match="copy_ shape mismatch"): t.copy_(torch.from_numpy(src)) _run_with(body) # ── __setitem__ / __getitem__ (shard-aligned) ─────────────────────── def test_setitem_getitem_single_pe(): """Scalar and slice assignment on a single-PE tensor round-trips.""" def body(torch): dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1) t = torch.zeros((1, 8), dtype="f16", dp=dp, name="t") # Scalar broadcast t[0, 0:4] = 3.0 assert np.allclose(t[0, 0:4], 3.0) assert np.allclose(t[0, 4:8], 0.0) # untouched # Array assignment t[0, 4:8] = np.array([10, 20, 30, 40], dtype=np.float16) assert np.array_equal(t[0, 4:8], [10, 20, 30, 40]) # Full read-back via numpy() expected = np.array([[3, 3, 3, 3, 10, 20, 30, 40]], dtype=np.float16) assert np.array_equal(t.numpy(), expected) _run_with(body) def test_setitem_getitem_multi_pe_shard_aligned(): """Shard-aligned slice assignment on an 8-PE column-wise tensor.""" def body(torch): n_pe = 8 n_elem = 4 # per shard dp = DPPolicy(cube="replicate", pe="column_wise", num_cubes=1, num_pes=n_pe) t = torch.zeros((1, n_pe * n_elem), dtype="f16", dp=dp, name="t") # Write each shard with its rank value for r in range(n_pe): t[0, r * n_elem : (r + 1) * n_elem] = float(r + 1) # Read back each shard for r in range(n_pe): expected = float(r + 1) arr = t[0, r * n_elem : (r + 1) * n_elem] assert np.allclose(arr, expected), f"shard {r}: {arr} != {expected}" # Full gather full = t.numpy().reshape(-1) for r in range(n_pe): assert np.allclose(full[r * n_elem : (r + 1) * n_elem], float(r + 1)) _run_with(body) def test_setitem_cross_shard_raises(): """Slice spanning two shards raises NotImplementedError.""" def body(torch): n_pe = 4 n_elem = 4 dp = DPPolicy(cube="replicate", pe="column_wise", num_cubes=1, num_pes=n_pe) t = torch.zeros((1, n_pe * n_elem), dtype="f16", dp=dp, name="t") with pytest.raises(NotImplementedError, match="spans multiple shards"): t[0, 2:6] = 1.0 # crosses shard 0 (0:4) and shard 1 (4:8) _run_with(body)