plugins/tensorrt/companion_sparse_trt.py

import math
import os
import sys
from typing import Tuple

import torch
import tensorrt as trt
import tensorrt.plugin as trtp

_NAMESPACE = "companionforge"
_APP_ROOT = os.environ.get("ANIGEN_APP_ROOT", "/home/user/app")
if _APP_ROOT not in sys.path:
    sys.path.insert(0, _APP_ROOT)

_CONV_CACHE = {}
_FLEXI_CACHE = {}


def _tt(x: trtp.Tensor) -> torch.Tensor:
    return torch.as_tensor(x, device="cuda")


def _stream_ctx(stream: int):
    return torch.cuda.stream(torch.cuda.ExternalStream(stream))


def _write_scalar(out: trtp.Tensor, value: int):
    t = _tt(out)
    t.fill_(int(value))


# -----------------------------------------------------------------------------
# SparseConv3D
# Explicit sparse representation: feats[N,C] + coords[N,4] (batch,x,y,z).
# Weight is stored in the native spconv layout, so exported checkpoints can be
# bound without transposition. The plugin intentionally delegates rulebook and
# GEMM selection to spconv.Native first; the ONNX ABI remains stable when this
# implementation is later replaced by an AOT CUDA kernel.
# -----------------------------------------------------------------------------
@trtp.register(f"{_NAMESPACE}::SparseConv3D")
def sparse_conv3d_desc(
    feats: trtp.TensorDesc,
    coords: trtp.TensorDesc,
    weight: trtp.TensorDesc,
    bias: trtp.TensorDesc,
    out_channels: int,
    kernel_size: int,
    stride: int,
    dilation: int,
    padding: int,
    subm: bool,
    spatial_x: int,
    spatial_y: int,
    spatial_z: int,
    batch_size: int,
) -> Tuple[trtp.TensorDesc, trtp.TensorDesc, trtp.TensorDesc]:
    # Production AniGen uses SubMConv3d here: cardinality is exactly N.
    # Keeping N as an ordinary dynamic dimension is essential because DDS
    # SizeTensor propagation through GroupNorm/reshape breaks TRT shape inference.
    n = feats.shape_expr[0]
    out_feats = trtp.from_shape_expr((n, int(out_channels)), dtype=feats.dtype)
    out_coords = trtp.from_shape_expr((n, 4), dtype=trt.int32)
    count = trtp.from_shape_expr((), dtype=trt.int32)
    return out_feats, out_coords, count


@trtp.impl(f"{_NAMESPACE}::SparseConv3D")
def sparse_conv3d_impl(
    feats: trtp.Tensor,
    coords: trtp.Tensor,
    weight: trtp.Tensor,
    bias: trtp.Tensor,
    out_channels: int,
    kernel_size: int,
    stride: int,
    dilation: int,
    padding: int,
    subm: bool,
    spatial_x: int,
    spatial_y: int,
    spatial_z: int,
    batch_size: int,
    outputs: Tuple[trtp.Tensor, trtp.Tensor, trtp.Tensor],
    stream: int,
) -> None:
    import spconv.pytorch as spconv

    with _stream_ctx(stream):
        f, c, w, b = _tt(feats), _tt(coords).to(torch.int32), _tt(weight), _tt(bias)
        # TensorRT's ONNX->Python-plugin bridge in 11.2 can corrupt scalar plugin
        # fields after serialization. For the production AniGen decoder all sparse
        # convolutions are SubMConv3d/stride=1; derive structural values from the
        # native spconv weight tensor instead of trusting serialized scalars.
        oc_runtime = int(w.shape[0])
        k_runtime = int(w.shape[1])
        subm_runtime = True
        stride_runtime = 1
        dilation_runtime = 1
        padding_runtime = 0
        key = (int(w.data_ptr()), int(b.data_ptr()), f.dtype, int(f.shape[1]), oc_runtime, k_runtime, subm_runtime)
        mod = _CONV_CACHE.get(key)
        if mod is None:
            algo = spconv.ConvAlgo.Native
            mod = spconv.SubMConv3d(
                int(f.shape[1]), oc_runtime, k_runtime,
                dilation=dilation_runtime, bias=(b.numel() != 0), algo=algo,
            )
            mod = mod.to(device=f.device, dtype=f.dtype).eval()
            with torch.no_grad():
                if tuple(mod.weight.shape) != tuple(w.shape):
                    raise RuntimeError(f"SparseConv3D weight layout mismatch: plugin={tuple(mod.weight.shape)} onnx={tuple(w.shape)}")
                mod.weight.copy_(w.to(dtype=mod.weight.dtype))
                if mod.bias is not None and b.numel():
                    mod.bias.copy_(b.to(dtype=mod.bias.dtype))
            _CONV_CACHE[key] = mod

        spatial = [max(1, int(c[:, i].max().item()) + 1) for i in (1, 2, 3)]
        bs = max(1, int(c[:, 0].max().item()) + 1)
        st = spconv.SparseConvTensor(f, c, spatial, bs)
        y = mod(st)
        m = int(y.features.shape[0])
        # SubMConv preserves the coordinate set but spconv may choose an internal
        # order that differs across independent modules. Normalize back to the
        # input coordinate order so geo/skin branches remain directly composable.
        yi = y.indices.to(torch.int32)
        yf = y.features
        if not torch.equal(yi, c):
            maxs=[max(int(c[:,j].max().item()),int(yi[:,j].max().item()))+1 for j in range(4)]
            mult=[maxs[1]*maxs[2]*maxs[3],maxs[2]*maxs[3],maxs[3],1]
            mul=torch.tensor(mult,device=c.device,dtype=torch.int64)
            ccode=(c.to(torch.int64)*mul).sum(-1); ycode=(yi.to(torch.int64)*mul).sum(-1)
            sy,perm=torch.sort(ycode); pos=torch.searchsorted(sy,ccode)
            if bool(torch.any(pos>=sy.numel()).item()) or not torch.equal(sy[pos],ccode):
                raise RuntimeError('SparseConv3D SubM output coordinate set differs from input')
            yf=yf[perm[pos]]
        of = _tt(outputs[0])
        oc = _tt(outputs[1])
        of.copy_(yf)
        oc.copy_(c)
        _write_scalar(outputs[2], m)


# -----------------------------------------------------------------------------
# Sparse window self attention. Semantics match AniGen's shifted-window
# partition followed by flash_attn_varlen_qkvpacked_func.
# -----------------------------------------------------------------------------
@trtp.register(f"{_NAMESPACE}::SparseWindowAttention")
def sparse_window_attention_desc(
    qkv: trtp.TensorDesc,
    coords: trtp.TensorDesc,
    window_size: int,
    shift_x: int,
    shift_y: int,
    shift_z: int,
) -> trtp.TensorDesc:
    return trtp.from_shape_expr(
        (qkv.shape_expr[0], qkv.shape_expr[2], qkv.shape_expr[3]), dtype=qkv.dtype
    )


@trtp.impl(f"{_NAMESPACE}::SparseWindowAttention")
def sparse_window_attention_impl(
    qkv: trtp.Tensor,
    coords: trtp.Tensor,
    window_size: int,
    shift_x: int,
    shift_y: int,
    shift_z: int,
    outputs: Tuple[trtp.Tensor],
    stream: int,
) -> None:
    import flash_attn

    with _stream_ctx(stream):
        x, c = _tt(qkv), _tt(coords).to(torch.int32)
        ws = int(window_size)
        shifted = c.clone()
        shifted[:, 1:] += torch.tensor(
            [int(shift_x), int(shift_y), int(shift_z)], device=c.device, dtype=torch.int32
        )[None]
        max_coords = shifted[:, 1:].max(dim=0).values.tolist()
        nw = [math.ceil((int(v) + 1) / ws) for v in max_coords]
        offset = torch.cumprod(torch.tensor([1] + nw[::-1]), dim=0).tolist()[::-1]
        shifted[:, 1:] //= ws
        ids = (shifted * torch.tensor(offset, device=c.device, dtype=torch.int32)[None]).sum(dim=1)
        fwd = torch.argsort(ids)
        bwd = torch.empty_like(fwd)
        bwd[fwd] = torch.arange(fwd.shape[0], device=c.device)
        lens = torch.bincount(ids)
        lens = lens[lens != 0].to(torch.int32)
        sorted_qkv = x[fwd]
        cu = torch.cat(
            [torch.zeros(1, device=x.device, dtype=torch.int32), torch.cumsum(lens, 0, dtype=torch.int32)], 0
        )
        if x.dtype in (torch.float16, torch.bfloat16):
            y = flash_attn.flash_attn_varlen_qkvpacked_func(sorted_qkv, cu, int(lens.max().item()))
        else:
            # Debug/reference FP32 path. Production AniGen uses FP16.
            pieces = []
            start = 0
            for ln in lens.tolist():
                z = sorted_qkv[start:start + ln]
                q, k, v = z.unbind(1)
                q, k, v = [t.permute(1, 0, 2).unsqueeze(0) for t in (q, k, v)]
                p = torch.nn.functional.scaled_dot_product_attention(q, k, v)
                pieces.append(p.squeeze(0).permute(1, 0, 2))
                start += ln
            y = torch.cat(pieces, 0)
        _tt(outputs[0]).copy_(y[bwd])


# -----------------------------------------------------------------------------
# SparseDownsample: AniGen average-pooling semantics + inverse map used by the
# matching SparseUpsample.
# -----------------------------------------------------------------------------
@trtp.register(f"{_NAMESPACE}::SparseDownsample")
def sparse_downsample_desc(
    feats: trtp.TensorDesc,
    coords: trtp.TensorDesc,
    factor_x: int,
    factor_y: int,
    factor_z: int,
) -> Tuple[trtp.TensorDesc, trtp.TensorDesc, trtp.TensorDesc, trtp.TensorDesc]:
    n, ch = feats.shape_expr[0], feats.shape_expr[1]
    st = trtp.size_tensor(n // 2, n)
    return (
        trtp.from_shape_expr((st.expr(), ch), dtype=feats.dtype),
        trtp.from_shape_expr((st.expr(), 4), dtype=trt.int32),
        trtp.from_shape_expr((n,), dtype=trt.int32),
        st,
    )


@trtp.impl(f"{_NAMESPACE}::SparseDownsample")
def sparse_downsample_impl(
    feats: trtp.Tensor,
    coords: trtp.Tensor,
    factor_x: int,
    factor_y: int,
    factor_z: int,
    outputs: Tuple[trtp.Tensor, trtp.Tensor, trtp.Tensor, trtp.Tensor],
    stream: int,
) -> None:
    with _stream_ctx(stream):
        f, c = _tt(feats), _tt(coords).to(torch.int32)
        factor = (int(factor_x), int(factor_y), int(factor_z))
        parts = list(c.unbind(-1))
        for i, fac in enumerate(factor):
            parts[i + 1] = parts[i + 1] // fac
        maxs = [int(parts[i + 1].max().item()) + 1 for i in range(3)]
        off = torch.cumprod(torch.tensor(maxs[::-1], dtype=torch.int64), 0).tolist()[::-1] + [1]
        code = sum(x.to(torch.int64) * int(o) for x, o in zip(parts, off))
        u, inv = code.unique(return_inverse=True)
        m = int(u.shape[0])
        y = torch.zeros((m, f.shape[1]), device=f.device, dtype=f.dtype)
        y = torch.scatter_reduce(y, 0, inv[:, None].expand(-1, f.shape[1]), f, reduce="mean")
        yc = torch.stack(
            [u // off[0]] + [(u // off[i + 1]) % maxs[i] for i in range(3)], -1
        ).to(torch.int32)
        out0 = _tt(outputs[0].aliased((int(f.shape[0]), int(f.shape[1]))))
        out1 = _tt(outputs[1].aliased((int(c.shape[0]), 4)))
        out0[:m].copy_(y)
        out1[:m].copy_(yc)
        _tt(outputs[2]).copy_(inv.to(torch.int32))
        _write_scalar(outputs[3], m)


@trtp.register(f"{_NAMESPACE}::SparseUpsample")
def sparse_upsample_desc(
    feats: trtp.TensorDesc,
    target_coords: trtp.TensorDesc,
    inverse: trtp.TensorDesc,
) -> Tuple[trtp.TensorDesc, trtp.TensorDesc]:
    n = target_coords.shape_expr[0]
    return (
        trtp.from_shape_expr((n, feats.shape_expr[1]), dtype=feats.dtype),
        target_coords.like(),
    )


@trtp.impl(f"{_NAMESPACE}::SparseUpsample")
def sparse_upsample_impl(
    feats: trtp.Tensor,
    target_coords: trtp.Tensor,
    inverse: trtp.Tensor,
    outputs: Tuple[trtp.Tensor, trtp.Tensor],
    stream: int,
) -> None:
    with _stream_ctx(stream):
        f, tc, inv = _tt(feats), _tt(target_coords), _tt(inverse).to(torch.long)
        _tt(outputs[0]).copy_(f[inv])
        _tt(outputs[1]).copy_(tc)


@trtp.register(f"{_NAMESPACE}::SparseSubdivide")
def sparse_subdivide_desc(
    feats: trtp.TensorDesc,
    coords: trtp.TensorDesc,
) -> Tuple[trtp.TensorDesc, trtp.TensorDesc]:
    n8 = feats.shape_expr[0] * 8
    return (
        trtp.from_shape_expr((n8, feats.shape_expr[1]), dtype=feats.dtype),
        trtp.from_shape_expr((n8, 4), dtype=trt.int32),
    )


@trtp.impl(f"{_NAMESPACE}::SparseSubdivide")
def sparse_subdivide_impl(
    feats: trtp.Tensor,
    coords: trtp.Tensor,
    outputs: Tuple[trtp.Tensor, trtp.Tensor],
    stream: int,
) -> None:
    with _stream_ctx(stream):
        f, c = _tt(feats), _tt(coords).to(torch.int32)
        offsets = torch.tensor(
            [[0, x, y, z] for x in (0, 1) for y in (0, 1) for z in (0, 1)],
            device=c.device,
            dtype=torch.int32,
        )
        oc = c.clone()
        oc[:, 1:] *= 2
        oc = (oc[:, None, :] + offsets[None, :, :]).flatten(0, 1)
        of = f[:, None, :].expand(f.shape[0], 8, f.shape[1]).flatten(0, 1)
        _tt(outputs[0]).copy_(of)
        _tt(outputs[1]).copy_(oc)


# -----------------------------------------------------------------------------
# MeshTopologyExtract: inference-only FlexiCubes topology extraction. Dynamic
# vertices/faces are exposed using two TensorRT size tensors.
# -----------------------------------------------------------------------------
@trtp.register(f"{_NAMESPACE}::MeshTopologyExtract")
def mesh_topology_desc(
    voxelgrid_vertices: trtp.TensorDesc,
    scalar_field: trtp.TensorDesc,
    cube_idx: trtp.TensorDesc,
    beta: trtp.TensorDesc,
    alpha: trtp.TensorDesc,
    gamma_f: trtp.TensorDesc,
    voxelgrid_colors: trtp.TensorDesc,
    resolution: int,
    no_sigmoid: bool,
) -> Tuple[trtp.TensorDesc, trtp.TensorDesc, trtp.TensorDesc, trtp.TensorDesc, trtp.TensorDesc]:
    nc = cube_idx.shape_expr[0]
    # Conservative FlexiCubes bounds; actual extents are communicated by DDS.
    vst = trtp.size_tensor(nc * 2, nc * 32)
    fst = trtp.size_tensor(nc * 4, nc * 64)
    return (
        trtp.from_shape_expr((vst.expr(), 3), dtype=voxelgrid_vertices.dtype),
        trtp.from_shape_expr((fst.expr(), 3), dtype=trt.int32),
        trtp.from_shape_expr((vst.expr(), voxelgrid_colors.shape_expr[1]), dtype=voxelgrid_colors.dtype),
        vst,
        fst,
    )


@trtp.impl(f"{_NAMESPACE}::MeshTopologyExtract")
def mesh_topology_impl(
    voxelgrid_vertices: trtp.Tensor,
    scalar_field: trtp.Tensor,
    cube_idx: trtp.Tensor,
    beta: trtp.Tensor,
    alpha: trtp.Tensor,
    gamma_f: trtp.Tensor,
    voxelgrid_colors: trtp.Tensor,
    resolution: int,
    no_sigmoid: bool,
    outputs: Tuple[trtp.Tensor, trtp.Tensor, trtp.Tensor, trtp.Tensor, trtp.Tensor],
    stream: int,
) -> None:
    from anigen.representations.mesh.flexicubes.flexicubes import FlexiCubes

    with _stream_ctx(stream):
        v = _tt(voxelgrid_vertices)
        s = _tt(scalar_field)
        cubes = _tt(cube_idx).to(torch.long)
        be, al, ga = _tt(beta), _tt(alpha), _tt(gamma_f)
        col = _tt(voxelgrid_colors)
        key = (int(v.device.index or 0), bool(col.shape[1] > 0))
        fc = _FLEXI_CACHE.get(key)
        if fc is None:
            fc = FlexiCubes(device=str(v.device), use_color=bool(col.shape[1] > 0))
            _FLEXI_CACHE[key] = fc
        verts, faces, _ldev, colors = fc(
            voxelgrid_vertices=v,
            scalar_field=s,
            cube_idx=cubes,
            resolution=int(resolution),
            beta=be,
            alpha=al,
            gamma_f=ga,
            voxelgrid_colors=col,
            training=False,
            no_sigmoid=bool(no_sigmoid),
        )
        nv, nf = int(verts.shape[0]), int(faces.shape[0])
        nc = int(cubes.shape[0]); vcap = nc * 32; fcap = nc * 64
        outv = _tt(outputs[0].aliased((vcap, 3)))
        outf = _tt(outputs[1].aliased((fcap, 3)))
        outc = _tt(outputs[2].aliased((vcap, int(col.shape[1]))))
        outv[:nv].copy_(verts)
        outf[:nf].copy_(faces.to(torch.int32))
        if colors is not None and outc.shape[1] > 0:
            outc[:nv].copy_(colors)
        _write_scalar(outputs[3], nv)
        _write_scalar(outputs[4], nf)



# -----------------------------------------------------------------------------
# Production high-level sparse mesh extractor. This consumes the actual SLat
# DAE head output [cube_feats, cube_coords], encapsulating sparse_cube2verts,
# dense attribute staging, FlexiCubes topology, RGB/normal extraction and
# vertex skin-feature extraction behind one TensorRT DDS node.
# -----------------------------------------------------------------------------
@trtp.register(f"{_NAMESPACE}::SparseMeshTopologyExtract")
def sparse_mesh_topology_desc(
    cube_feats: trtp.TensorDesc,
    cube_coords: trtp.TensorDesc,
    resolution: int,
) -> Tuple[trtp.TensorDesc, trtp.TensorDesc, trtp.TensorDesc, trtp.TensorDesc, trtp.TensorDesc, trtp.TensorDesc]:
    n = cube_feats.shape_expr[0]
    # Conservative bounds from sparse surface cubes; unlike the dense 256^3
    # grid these scale only with generated sparse cubes.
    vst = trtp.size_tensor(n, n * 16)
    fst = trtp.size_tensor(n * 2, n * 32)
    return (
        trtp.from_shape_expr((vst.expr(), 3), dtype=trt.float32),
        trtp.from_shape_expr((fst.expr(), 3), dtype=trt.int32),
        trtp.from_shape_expr((vst.expr(), 6), dtype=trt.float32),
        trtp.from_shape_expr((vst.expr(), 4), dtype=cube_feats.dtype),
        vst,
        fst,
    )


@trtp.impl(f"{_NAMESPACE}::SparseMeshTopologyExtract")
def sparse_mesh_topology_impl(
    cube_feats: trtp.Tensor,
    cube_coords: trtp.Tensor,
    resolution: int,
    outputs: Tuple[trtp.Tensor, trtp.Tensor, trtp.Tensor, trtp.Tensor, trtp.Tensor, trtp.Tensor],
    stream: int,
) -> None:
    from anigen.modules.sparse import SparseTensor
    from anigen.representations.mesh.cube2mesh_skeleton import AniGenSparseFeatures2Mesh
    with _stream_ctx(stream):
        f=_tt(cube_feats); c=_tt(cube_coords).to(torch.int32)
        # Production AniGen slat_dae config is resolution=64 and extracts mesh at x4.
        res_runtime=256
        key=('sparse-mesh',res_runtime,int(f.shape[1]))
        ext=_FLEXI_CACHE.get(key)
        if ext is None:
            ext=AniGenSparseFeatures2Mesh(res=res_runtime,use_color=True,skin_feat_channels=4,predict_skin=True,device='cuda')
            _FLEXI_CACHE[key]=ext
        st=SparseTensor(feats=f,coords=c)
        mesh=ext(st,training=False)
        verts=mesh.vertices.float(); faces=mesh.faces.to(torch.int32)
        attrs=mesh.vertex_attrs.float() if mesh.vertex_attrs is not None else torch.zeros((verts.shape[0],6),device=verts.device,dtype=torch.float32)
        skin=mesh.vertex_skin_feats if mesh.vertex_skin_feats is not None else torch.zeros((verts.shape[0],4),device=verts.device,dtype=f.dtype)
        nv,nf=int(verts.shape[0]),int(faces.shape[0]);vcap=int(f.shape[0])*16;fcap=int(f.shape[0])*32
        ov=_tt(outputs[0].aliased((vcap,3)));of=_tt(outputs[1].aliased((fcap,3)));oa=_tt(outputs[2].aliased((vcap,6)));os=_tt(outputs[3].aliased((vcap,4)))
        ov[:nv].copy_(verts);of[:nf].copy_(faces);oa[:nv].copy_(attrs);os[:nv].copy_(skin.to(os.dtype))
        _write_scalar(outputs[4],nv);_write_scalar(outputs[5],nf)

def registered_ops():
    return [
        "SparseConv3D",
        "SparseWindowAttention",
        "SparseDownsample",
        "SparseUpsample",
        "SparseSubdivide",
        "MeshTopologyExtract",
        "SparseMeshTopologyExtract",
    ]