7860148eb2
- sparse_conv_metal: fused gather-GEMM kernel, parity 8e-4 vs stock; honest verdict: stock chunking already bounds memory at decoder scale (1.97GB synthetic peak) and beats the scalar kernel on speed — kept as TRELLIS2_METAL_SPCONV=1 opt-in + upstream reference, NOT default. - _attempt_schedule: preferred=vertex now runs the vertex baker (was silently falling through to the 5h+ pure-python kdtree grind). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
86 lines
3.0 KiB
Python
86 lines
3.0 KiB
Python
"""Parity + memory gate: fused Metal sparse conv vs the stock chunked path.
|
|
Small case first (exact check vs naive reference), then a decoder-scale
|
|
case for peak-memory + timing comparison.
|
|
"""
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from mlx_backend.sparse_conv import MlxSparseConv3d, _build_neighbor_map
|
|
from mlx_backend.sparse_conv_metal import sparse_conv_metal
|
|
from mlx_backend.sparse_tensor import MlxSparseTensor
|
|
|
|
mx.random.seed(0)
|
|
|
|
|
|
def make_case(n_vox, ci, co, res):
|
|
# random unique voxel coords in a res^3 grid, batch 0
|
|
rng = np.random.default_rng(0)
|
|
flat = rng.choice(res ** 3, size=n_vox, replace=False)
|
|
x, rem = flat // (res * res), flat % (res * res)
|
|
y, z = rem // res, rem % res
|
|
coords = np.stack([np.zeros_like(x), x, y, z], 1).astype(np.int32)
|
|
st = MlxSparseTensor(
|
|
feats=mx.random.normal((n_vox, ci)).astype(mx.float16),
|
|
coords=mx.array(coords))
|
|
conv = MlxSparseConv3d(ci, co, 3)
|
|
conv.weight = mx.random.normal((co, 3, 3, 3, ci)).astype(mx.float16) * 0.05
|
|
conv.bias = mx.random.normal((co,)).astype(mx.float16) * 0.1
|
|
return st, conv
|
|
|
|
|
|
def run_metal(st, conv):
|
|
nmap = _build_neighbor_map(st.coords, st.shape[0], st.spatial_shape,
|
|
conv.kernel_size, conv.dilation)
|
|
mx.eval(nmap)
|
|
K = nmap.shape[0]
|
|
Co, Ci = conv.out_channels, conv.in_channels
|
|
w = conv.weight.reshape(Co, K, Ci).transpose(1, 2, 0) # (K, Ci, Co)
|
|
feats_padded = mx.concatenate(
|
|
[st.feats, mx.zeros((1, Ci), dtype=st.feats.dtype)], axis=0)
|
|
return sparse_conv_metal(feats_padded, nmap, w, conv.bias)
|
|
|
|
|
|
# ---- small exact-parity case ----
|
|
st, conv = make_case(5000, 64, 96, 64)
|
|
ref = conv(st).feats
|
|
got = run_metal(st, conv)
|
|
mx.eval(ref, got)
|
|
d = mx.max(mx.abs(ref.astype(mx.float32) - got.astype(mx.float32))).item()
|
|
rel = d / max(1e-6, mx.max(mx.abs(ref.astype(mx.float32))).item())
|
|
print(f"small: max|diff|={d:.3e} rel={rel:.3e}", "PASS" if rel < 2e-2 else "FAIL")
|
|
ok_small = rel < 2e-2
|
|
|
|
# ---- decoder-scale case: memory + time ----
|
|
st2, conv2 = make_case(400_000, 512, 512, 512)
|
|
mx.eval(st2.feats, conv2.weight)
|
|
|
|
mx.reset_peak_memory()
|
|
t0 = time.perf_counter()
|
|
ref2 = conv2(st2).feats
|
|
mx.eval(ref2)
|
|
t_stock = time.perf_counter() - t0
|
|
m_stock = mx.get_peak_memory() / 2**30
|
|
|
|
st2._spatial_cache = {} # drop cached neighbor map so both build it
|
|
mx.reset_peak_memory()
|
|
t0 = time.perf_counter()
|
|
got2 = run_metal(st2, conv2)
|
|
mx.eval(got2)
|
|
t_metal = time.perf_counter() - t0
|
|
m_metal = mx.get_peak_memory() / 2**30
|
|
|
|
d2 = mx.max(mx.abs(ref2.astype(mx.float32) - got2.astype(mx.float32))).item()
|
|
r2 = d2 / max(1e-6, mx.max(mx.abs(ref2.astype(mx.float32))).item())
|
|
print(f"large: stock {t_stock:.2f}s peak {m_stock:.2f}GB | "
|
|
f"metal {t_metal:.2f}s peak {m_metal:.2f}GB | rel diff {r2:.3e}")
|
|
ok_large = r2 < 2e-2
|
|
print("SPCONV_GATE:", "PASS" if (ok_small and ok_large) else "FAIL")
|
|
sys.exit(0 if (ok_small and ok_large) else 1)
|