70cd436eb9
Extracted from lato.2_mrp_mlx. The same sparse module underlies LATO.2, Pixal3D and the rest of the TRELLIS.2 family, and all of them are blocked on Apple Silicon by the same single op - submanifold conv - so it belongs in one tested package rather than vendored per port. 13/13 tests: 7 for the conv against a hand-written reference (spconv is uninstallable here so there is no upstream oracle), 6 for the remaining layers against torch. SubMConv3d runs 13.8ms at 128^3/128ch on m3ultra, 5.9x faster than the obvious per-offset loop.
182 lines
7.3 KiB
Python
182 lines
7.3 KiB
Python
"""Submanifold sparse 3D convolution in pure MLX.
|
|
|
|
This is the ONLY genuinely CUDA-locked operation in LATO.2's inference path. Upstream
|
|
routes it to `spconv.SubMConv3d` (or torchsparse); neither has a Metal build, which is
|
|
what has kept LATO.2 — and TRELLIS before it — off Apple Silicon.
|
|
|
|
Every SparseConv3d in the LATO.2 model code is constructed with the defaults
|
|
`stride=1, padding=None`, which upstream dispatches to SubMConv3d. So only the
|
|
submanifold case is needed, at kernel sizes 3 and 1.
|
|
|
|
Submanifold semantics: the output occupies EXACTLY the input coordinates (no dilation of
|
|
the occupied set). For each output voxel c:
|
|
|
|
out[c] = bias + sum over kernel offsets d of W[d] @ feats[c + d] (c+d occupied)
|
|
|
|
Absent neighbours contribute nothing. Implemented by gathering into a feature matrix with
|
|
one appended zero row, so "missing" is index N and needs no masking in the hot loop.
|
|
|
|
Building the indice map is the expensive part and depends only on the coordinate set, so
|
|
it is cached on the SparseTensor under `indice_key` — the same trick spconv uses, and the
|
|
reason upstream threads `indice_key=f"res_{resolution}"` through the ResBlocks.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Optional
|
|
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
import numpy as np
|
|
|
|
from .tensor import SparseTensor
|
|
|
|
_MISSING = -1
|
|
|
|
|
|
def _kernel_offsets(k: int) -> np.ndarray:
|
|
"""Kernel offsets in C order over (dz, dy, dx), centred — matches spconv's ordering."""
|
|
r = np.arange(k) - (k // 2)
|
|
return np.stack(np.meshgrid(r, r, r, indexing="ij"), axis=-1).reshape(-1, 3)
|
|
|
|
|
|
def build_indice_map(coords: mx.array, kernel_size: int) -> np.ndarray:
|
|
"""[K^3, N] int32 — for each offset, the row of the neighbour, or -1 if unoccupied.
|
|
|
|
Uses a sorted-key binary search rather than a Python dict: at k=3 this is 27 lookups
|
|
per voxel, and a per-voxel dict lookup would dominate runtime for any real mesh.
|
|
"""
|
|
c = np.asarray(coords, dtype=np.int64)
|
|
n = c.shape[0]
|
|
if n == 0:
|
|
return np.full((kernel_size**3, 0), _MISSING, dtype=np.int32)
|
|
|
|
offsets = _kernel_offsets(kernel_size)
|
|
pad = kernel_size // 2
|
|
|
|
# Encode (batch,z,y,x) into one int64. Shift by `pad` so that neighbour coordinates
|
|
# of -1 stay non-negative and cannot alias onto a real cell at the opposite edge.
|
|
lo = c.min(axis=0) - pad
|
|
ext = (c.max(axis=0) + pad) - lo + 1
|
|
strides = np.array(
|
|
[ext[1] * ext[2] * ext[3], ext[2] * ext[3], ext[3], 1], dtype=np.int64
|
|
)
|
|
|
|
def encode(arr: np.ndarray) -> np.ndarray:
|
|
return ((arr - lo) * strides).sum(axis=1)
|
|
|
|
keys = encode(c)
|
|
order = np.argsort(keys, kind="stable")
|
|
sorted_keys = keys[order]
|
|
|
|
imap = np.empty((offsets.shape[0], n), dtype=np.int32)
|
|
for i, d in enumerate(offsets):
|
|
probe = c.copy()
|
|
probe[:, 1:] += d # batch index (column 0) never shifts
|
|
pk = encode(probe)
|
|
pos = np.searchsorted(sorted_keys, pk)
|
|
pos_clipped = np.clip(pos, 0, n - 1)
|
|
hit = sorted_keys[pos_clipped] == pk
|
|
imap[i] = np.where(hit, order[pos_clipped], _MISSING).astype(np.int32)
|
|
return imap
|
|
|
|
|
|
class SubMConv3d(nn.Module):
|
|
"""Submanifold sparse conv. Weight layout [K^3, in_channels, out_channels]."""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
kernel_size: int = 3,
|
|
bias: bool = True,
|
|
indice_key: Optional[str] = None,
|
|
):
|
|
super().__init__()
|
|
if kernel_size % 2 != 1:
|
|
raise ValueError(f"kernel_size must be odd, got {kernel_size}")
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
self.kernel_size = kernel_size
|
|
self.indice_key = indice_key
|
|
|
|
scale = (in_channels * kernel_size**3) ** -0.5
|
|
self.weight = mx.random.uniform(
|
|
-scale, scale, (kernel_size**3, in_channels, out_channels)
|
|
)
|
|
if bias:
|
|
self.bias = mx.zeros((out_channels,))
|
|
|
|
def _gather_index(self, x: SparseTensor, n: int) -> mx.array:
|
|
"""Cached [N, K^3] gather index, already on-device.
|
|
|
|
Caching only the raw indice map is not enough: rebuilding the "missing -> N"
|
|
substitution and re-uploading the index cost more than the convolution itself.
|
|
At 128^3/128ch that overhead was ~26ms against ~13ms of actual GPU work, i.e.
|
|
two thirds of the measured time was CPU-side bookkeeping repeated every layer.
|
|
The transposed, sentinel-substituted device array depends only on the
|
|
coordinate set, so it is cached whole.
|
|
"""
|
|
key = f"gidx_k{self.kernel_size}_{self.indice_key}"
|
|
cached = x.cache_get(key)
|
|
if cached is not None:
|
|
return cached
|
|
imap = build_indice_map(x.coords, self.kernel_size)
|
|
idx_t = mx.array(np.where(imap == _MISSING, n, imap).T) # [N, K^3]
|
|
x.cache_put(key, idx_t)
|
|
return idx_t
|
|
|
|
def __call__(self, x: SparseTensor) -> SparseTensor:
|
|
n = x.feats.shape[0]
|
|
|
|
# k=1 touches only the centre voxel, so it is exactly a per-voxel linear —
|
|
# skip the indice map entirely.
|
|
if self.kernel_size == 1:
|
|
out = x.feats @ self.weight[0]
|
|
if hasattr(self, "bias"):
|
|
out = out + self.bias
|
|
return x.replace(out)
|
|
|
|
# One appended zero row: absent neighbours index it and contribute nothing,
|
|
# which avoids a per-offset boolean mask.
|
|
feats_pad = mx.concatenate(
|
|
[x.feats, mx.zeros((1, self.in_channels), dtype=x.feats.dtype)], axis=0
|
|
)
|
|
idx_t = self._gather_index(x, n) # [N, K^3], cached on device
|
|
k3 = self.kernel_size**3
|
|
|
|
# Fuse the K^3 taps into ONE gather + ONE matmul.
|
|
#
|
|
# The obvious implementation loops over the K^3 offsets accumulating
|
|
# `gather(i) @ W[i]`, but that issues 2*K^3 tiny GPU dispatches with Python
|
|
# between them, and each is far too small to fill the machine. Fleet
|
|
# benchmarking made this unmistakable: the 80-core M3 Ultra came in SLOWER
|
|
# than a 38-core M2 Max (81.6ms vs 58.3ms at 128^3/128ch), i.e. throughput was
|
|
# anti-correlated with core count — the signature of launch-latency binding
|
|
# rather than compute binding.
|
|
#
|
|
# Concatenating neighbours along the channel axis turns the whole thing into a
|
|
# single [N, K^3*Cin] x [K^3*Cin, Cout] matmul, which is one dispatch big
|
|
# enough to actually occupy the GPU.
|
|
#
|
|
# That buffer is N*K^3*Cin floats, so it is chunked over rows to keep peak
|
|
# memory bounded (~256MB/chunk) — at 128^3 and 128ch the unchunked form alone
|
|
# would be ~2.9GB, which is fine on a Studio and not fine on an 8GB mini.
|
|
w_flat = self.weight.reshape(k3 * self.in_channels, self.out_channels).astype(
|
|
x.feats.dtype
|
|
)
|
|
bytes_per_row = k3 * self.in_channels * 4
|
|
chunk = max(1, min(n, (256 << 20) // max(bytes_per_row, 1)))
|
|
|
|
outs = []
|
|
for start in range(0, n, chunk):
|
|
stop = min(start + chunk, n)
|
|
g = mx.take(feats_pad, idx_t[start:stop].reshape(-1), axis=0)
|
|
g = g.reshape(stop - start, k3 * self.in_channels)
|
|
outs.append(g @ w_flat)
|
|
out = outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=0)
|
|
|
|
if hasattr(self, "bias"):
|
|
out = out + self.bias.astype(x.feats.dtype)
|
|
return x.replace(out)
|