50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
"""Unified 3D grid sampling — flex_gemm on Metal/CUDA, F.grid_sample fallback."""
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
try:
|
|
from flex_gemm.ops.grid_sample import grid_sample_3d as _flex_grid_sample
|
|
_HAS_FLEX_GEMM = True
|
|
except ImportError:
|
|
_HAS_FLEX_GEMM = False
|
|
|
|
|
|
def grid_sample_3d(feats, coords, shape, grid, mode='trilinear'):
|
|
"""Drop-in replacement for flex_gemm.ops.grid_sample.grid_sample_3d.
|
|
|
|
Args:
|
|
feats: [N, C] sparse voxel features
|
|
coords: [N, 4] voxel coordinates (batch_idx, x, y, z)
|
|
shape: torch.Size([B, C, D, H, W]) — sparse tensor shape
|
|
grid: [B, M, 3] query points in voxel space
|
|
mode: 'trilinear' (maps to F.grid_sample 'bilinear')
|
|
Returns:
|
|
[B*C, M] sampled features (matching flex_gemm output shape)
|
|
"""
|
|
if _HAS_FLEX_GEMM:
|
|
return _flex_grid_sample(feats, coords, shape, grid, mode=mode)
|
|
|
|
# Dense volume fallback
|
|
B, C = shape[0], shape[1]
|
|
D, H, W = shape[2], shape[3], shape[4]
|
|
device = feats.device
|
|
|
|
dense_vol = torch.zeros(B, C, D, H, W, dtype=feats.dtype, device=device)
|
|
batch_idx = coords[:, 0].long()
|
|
cx, cy, cz = coords[:, 1].long(), coords[:, 2].long(), coords[:, 3].long()
|
|
dense_vol[batch_idx, :, cx, cy, cz] = feats
|
|
|
|
# Normalize grid to [-1, 1] for F.grid_sample (expects z,y,x order)
|
|
grid_norm = torch.stack([
|
|
grid[..., 2] / (W - 1) * 2 - 1,
|
|
grid[..., 1] / (H - 1) * 2 - 1,
|
|
grid[..., 0] / (D - 1) * 2 - 1,
|
|
], dim=-1)
|
|
grid_norm = grid_norm.reshape(B, 1, 1, -1, 3)
|
|
|
|
sampled = F.grid_sample(dense_vol, grid_norm, mode='bilinear',
|
|
align_corners=True, padding_mode='border')
|
|
# sampled: [B, C, 1, 1, M] -> reshape to match flex_gemm output
|
|
M = grid.shape[1]
|
|
return sampled.reshape(B * C, M)
|