97dcdfb54a
The blocker for LATO.2 on Apple Silicon is one op, not the whole setup.sh --all CUDA stack. Measured: 5 of 7 checkpoints are fully dense, and every SparseConv3d in the model is constructed stride=1/padding=None, which upstream dispatches to spconv's SubMConv3d. No strided or inverse sparse conv is ever instantiated. - SubMConv3d in pure MLX via a sorted-key indice map (27 lookups/voxel vectorised, cached per coordinate set the way spconv uses indice_key) - SparseTensor container + subdivide upsampling - Converter handles the 5-D layout collision: spconv KRSC [O,kz,ky,kx,I] vs torch Conv3d [O,I,kz,ky,kx]. Rank alone is ambiguous; misreading it silently mangles the voxel encoder. - 7/7 tests pass vs an independent naive reference, max err 3e-7. spconv has no Metal build so there is no upstream oracle; the reference shares no indexing code. Kernel orientation (feats[c+d] vs feats[c-d]) remains unverified and is silent when wrong; --flip-kernel builds the mirror for an end-to-end A/B.
163 lines
6.4 KiB
Python
163 lines
6.4 KiB
Python
"""Correctness tests for the MLX sparse core.
|
|
|
|
spconv cannot be installed on this machine — that is the entire reason this port exists —
|
|
so there is no way to diff against upstream numerically here. Instead the vectorised MLX
|
|
implementation is checked against a deliberately naive, obviously-correct reference
|
|
written straight from the definition of submanifold convolution (dict lookup, per-voxel
|
|
Python loop). The two share no indexing code, so an off-by-one in the fast path cannot
|
|
hide in both.
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from lato_mlx.sparse.conv import SubMConv3d, build_indice_map, _kernel_offsets
|
|
from lato_mlx.sparse.tensor import SparseTensor, subdivide
|
|
|
|
|
|
def reference_subm_conv(coords, feats, weight, bias, k):
|
|
"""Definition of submanifold conv, written for obviousness, not speed."""
|
|
occupied = {tuple(c): i for i, c in enumerate(coords.tolist())}
|
|
offsets = _kernel_offsets(k)
|
|
n, out_c = coords.shape[0], weight.shape[2]
|
|
out = np.zeros((n, out_c), dtype=np.float64)
|
|
for i, c in enumerate(coords.tolist()):
|
|
for oi, d in enumerate(offsets):
|
|
nb = (c[0], c[1] + d[0], c[2] + d[1], c[3] + d[2])
|
|
j = occupied.get(nb)
|
|
if j is not None:
|
|
out[i] += feats[j].astype(np.float64) @ weight[oi].astype(np.float64)
|
|
if bias is not None:
|
|
out += bias.astype(np.float64)
|
|
return out
|
|
|
|
|
|
def random_sparse(n_vox, channels, batch=2, res=8, seed=0):
|
|
rng = np.random.default_rng(seed)
|
|
seen, coords = set(), []
|
|
while len(coords) < n_vox:
|
|
b = int(rng.integers(0, batch))
|
|
z, y, x = (int(v) for v in rng.integers(0, res, 3))
|
|
if (b, z, y, x) in seen:
|
|
continue
|
|
seen.add((b, z, y, x))
|
|
coords.append((b, z, y, x))
|
|
coords.sort() # upstream requires batch-contiguous rows
|
|
coords = np.array(coords, dtype=np.int32)
|
|
feats = rng.standard_normal((n_vox, channels)).astype(np.float32)
|
|
return coords, feats
|
|
|
|
|
|
def test_subm_conv_matches_reference(k=3, n=180, cin=12, cout=7):
|
|
coords, feats = random_sparse(n, cin, seed=1)
|
|
rng = np.random.default_rng(2)
|
|
w = rng.standard_normal((k**3, cin, cout)).astype(np.float32) * 0.1
|
|
b = rng.standard_normal((cout,)).astype(np.float32)
|
|
|
|
conv = SubMConv3d(cin, cout, k, bias=True, indice_key="t")
|
|
conv.weight = mx.array(w)
|
|
conv.bias = mx.array(b)
|
|
|
|
got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats)
|
|
want = reference_subm_conv(coords, feats, w, b, k)
|
|
err = np.abs(got - want).max()
|
|
assert err < 2e-4, f"k={k} max abs err {err:.3g}"
|
|
return err
|
|
|
|
|
|
def test_kernel1_is_pointwise(n=64, cin=8, cout=5):
|
|
coords, feats = random_sparse(n, cin, seed=3)
|
|
rng = np.random.default_rng(4)
|
|
w = rng.standard_normal((1, cin, cout)).astype(np.float32)
|
|
conv = SubMConv3d(cin, cout, 1, bias=False, indice_key="p")
|
|
conv.weight = mx.array(w)
|
|
got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats)
|
|
err = np.abs(got - feats @ w[0]).max()
|
|
assert err < 1e-4, f"k=1 err {err:.3g}"
|
|
return err
|
|
|
|
|
|
def test_isolated_voxel_sees_only_itself():
|
|
"""A voxel with no occupied neighbours must reduce to the centre tap alone."""
|
|
coords = np.array([[0, 0, 0, 0], [0, 50, 50, 50]], dtype=np.int32)
|
|
feats = np.ones((2, 3), dtype=np.float32)
|
|
w = np.zeros((27, 3, 3), dtype=np.float32)
|
|
centre = 13 # index of (0,0,0) in centred C-order offsets
|
|
assert tuple(_kernel_offsets(3)[centre]) == (0, 0, 0)
|
|
w[centre] = np.eye(3)
|
|
conv = SubMConv3d(3, 3, 3, bias=False, indice_key="iso")
|
|
conv.weight = mx.array(w)
|
|
got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats)
|
|
assert np.abs(got - feats).max() < 1e-6, got
|
|
return 0.0
|
|
|
|
|
|
def test_batches_do_not_leak():
|
|
"""Same spatial cell in two batch items must not become neighbours."""
|
|
coords = np.array([[0, 1, 1, 1], [1, 1, 1, 1]], dtype=np.int32)
|
|
feats = np.array([[1.0], [100.0]], dtype=np.float32)
|
|
w = np.ones((27, 1, 1), dtype=np.float32) # sum every occupied neighbour
|
|
conv = SubMConv3d(1, 1, 3, bias=False, indice_key="b")
|
|
conv.weight = mx.array(w)
|
|
got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats)
|
|
assert np.allclose(got, [[1.0], [100.0]]), f"batch leak: {got}"
|
|
return 0.0
|
|
|
|
|
|
def test_indice_map_hit_rate():
|
|
"""A fully dense block: interior voxels must find all 27 neighbours."""
|
|
coords = np.array(
|
|
[[0, z, y, x] for z in range(4) for y in range(4) for x in range(4)],
|
|
dtype=np.int32,
|
|
)
|
|
imap = build_indice_map(mx.array(coords), 3)
|
|
lin = {tuple(c): i for i, c in enumerate(coords.tolist())}
|
|
interior = [lin[(0, z, y, x)] for z in (1, 2) for y in (1, 2) for x in (1, 2)]
|
|
assert (imap[:, interior] != -1).all(), "interior voxel missing a neighbour"
|
|
corner = lin[(0, 0, 0, 0)]
|
|
assert (imap[:, corner] != -1).sum() == 8, "corner should see exactly 8 of 27"
|
|
return 0.0
|
|
|
|
|
|
def test_subdivide():
|
|
coords = np.array([[0, 1, 2, 3]], dtype=np.int32)
|
|
feats = np.array([[5.0, 6.0]], dtype=np.float32)
|
|
out = subdivide(SparseTensor(mx.array(feats), mx.array(coords)))
|
|
c = np.asarray(out.coords)
|
|
assert c.shape == (8, 4) and out.feats.shape == (8, 2)
|
|
assert set(map(tuple, c[:, 1:].tolist())) == {
|
|
(2 + a, 4 + b, 6 + d) for a in (0, 1) for b in (0, 1) for d in (0, 1)
|
|
}
|
|
assert np.abs(np.asarray(out.feats) - feats).max() < 1e-6
|
|
return 0.0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tests = [
|
|
("subm k=3 vs reference", lambda: test_subm_conv_matches_reference(3)),
|
|
("subm k=5 vs reference", lambda: test_subm_conv_matches_reference(5, n=140)),
|
|
("k=1 is pointwise", test_kernel1_is_pointwise),
|
|
("isolated voxel", test_isolated_voxel_sees_only_itself),
|
|
("batch isolation", test_batches_do_not_leak),
|
|
("indice map hit rate", test_indice_map_hit_rate),
|
|
("subdivide", test_subdivide),
|
|
]
|
|
failed = 0
|
|
for name, fn in tests:
|
|
try:
|
|
err = fn()
|
|
print(f" PASS {name:28s} (max err {err:.2e})")
|
|
except AssertionError as e:
|
|
print(f" FAIL {name:28s} {e}")
|
|
failed += 1
|
|
except Exception as e: # noqa: BLE001
|
|
print(f" ERROR {name:28s} {type(e).__name__}: {e}")
|
|
failed += 1
|
|
print(f"\n{len(tests)-failed}/{len(tests)} passed")
|
|
sys.exit(1 if failed else 0)
|