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", ]