lato.2_mrp_mlx/lato_mlx/sparse/conv.py

"""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)