Optional quantem-cuda kernel dispatch for K-Planes models and volume TV

src/quantem/core/config.py

@@ -34,6 +34,15 @@
     if "cuda" in str(e):
         NUM_DEVICES = 0
     _defaults["has_cupy"] = False
+try:
+    import quantem.cuda  # type: ignore  # noqa: F401
+
+    _defaults["has_quantem_cuda"] = True
+except ModuleNotFoundError:
+    _defaults["has_quantem_cuda"] = False
+except Exception:
+    # installed but unloadable (e.g. libcudart missing at runtime)
+    _defaults["has_quantem_cuda"] = False
 
 
 defaults: list[Mapping] = [_defaults]

src/quantem/core/ml/models/kplanes.py

@@ -10,6 +10,8 @@
 import torch.nn.functional as F
 from torch import nn
 
+from quantem.core import config
+
 from .model_base import PPLR, TensorDecompositionModel
 from .so3params import SO3ParamQuat, SO3ParamR9SVD
 
@@ -169,6 +171,7 @@ class KPlanes(PPLR, TensorDecompositionModel):
     """
     K-Planes model adapted from Fridovich-Keil et al., https://arxiv.org/abs/2301.10241
     """
+
     def __init__(
         self,
         # Grid parameters
@@ -312,10 +315,30 @@ def interpolate_ms_features_tilted(
     """
     Fully-vectorized multi-scale, multi-rotation K-Planes feature interpolation.
     Returns features of shape (B, C * T * num_scales).
+
+    When the optional ``quantem-cuda`` package is installed, the tensors live
+    on a CUDA device, and the ``use_cuda_kernels`` config option is true
+    (default), each scale dispatches to the fused CUDA kernel (rotate +
+    grid_sample + Hadamard product in one launch, analytic backward) —
+    identical semantics to the torch path below.
     """
     T = rotation_matrices.shape[0]
     B = pts.shape[0]
 
+    if (
+        pts.is_cuda
+        and pts.dtype == torch.float32
+        and rotation_matrices.dtype == torch.float32
+        and all(g.dtype == torch.float32 for g in ms_grids)
+        and config.get("has_quantem_cuda")
+        and config.get("use_cuda_kernels", default=True)
+    ):
+        from quantem.cuda.core.ml import kplanes_tilted_fuse
+
+        return torch.cat(
+            [kplanes_tilted_fuse(pts, rotation_matrices, g) for g in ms_grids], dim=-1
+        )
+
     # (T, B, 3)  — rotate all points by all rotations at once
     rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts)
 

src/quantem/core/quantem.yaml

@@ -14,6 +14,10 @@ dtype_complex: complex64
 # global verbosity, not currently used
 
 verbose: 1
+
+# Dispatch to the fused quantem-cuda kernels when installed (pip install quantem[cuda])
+# and the tensors are on a CUDA device. Set false to force the pure-torch paths.
+use_cuda_kernels: true
 cupy:
   # The size of the fft cache in MB used by cupy
   # https://docs.cupy.dev/en/stable/user_guide/fft.html#fft-plan-cache

src/quantem/diffractive_imaging/object_models.py

@@ -396,6 +396,30 @@ def get_tv_loss(
         return loss
 
     def _calc_tv_loss(self, array: torch.Tensor, weight: tuple[float, float]) -> torch.Tensor:
+        # Identical math via the fused quantem-cuda L1 kernel when available.
+        # Weighted size-1 axes fall through to torch so degenerate inputs
+        # behave exactly as before.
+        if (
+            array.ndim == 3
+            and array.is_cuda
+            and array.dtype == torch.float32
+            and config.get("has_quantem_cuda")
+            and config.get("use_cuda_kernels", default=True)
+            and not (weight[0] > 0 and array.shape[0] == 1)
+            and not (weight[1] > 0 and (array.shape[1] == 1 or array.shape[2] == 1))
+        ):
+            from quantem.cuda.core import tv_loss_l1_3d
+
+            s, h, w_ = array.shape
+            wts = array.new_tensor([weight[0], weight[1], weight[1]])
+            active = int((wts > 0).sum())
+            if active == 0:
+                return self._get_zero_loss_tensor()
+            counts = array.new_tensor(
+                [(s - 1) * h * w_, s * (h - 1) * w_, s * h * (w_ - 1)]
+            ).clamp(min=1)
+            return (wts * tv_loss_l1_3d(array) / counts).sum() / active
+
         loss = self._get_zero_loss_tensor()
         calc_dim = 0
         for dim in range(array.ndim):

src/quantem/tomography/object_models.py

@@ -18,6 +18,7 @@
 from quantem.core.utils.rng import RNGMixin
 from quantem.tomography.dataset_models import TomographyINRPretrainDataset
 from quantem.tomography.tomography_context import ReconstructionContext
+from quantem.tomography.utils import tv_loss_vol_sq
 
 
 class ObjConstraintParams:
@@ -439,11 +440,9 @@ def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor:
         # TV over the three trailing spatial dims, leaving any leading channel/batch axes
         # intact. Works for a 3-D volume, obj_view's [1, D, H, W], and a multimodal
         # [C, D, H, W] (channels = elemental compositions), matching the INR / tensor-decomp
-        # convention where the object carries a leading channel dimension.
-        tv_d = torch.pow(ctx.obj[..., 1:, :, :] - ctx.obj[..., :-1, :, :], 2).sum()
-        tv_h = torch.pow(ctx.obj[..., :, 1:, :] - ctx.obj[..., :, :-1, :], 2).sum()
-        tv_w = torch.pow(ctx.obj[..., :, :, 1:] - ctx.obj[..., :, :, :-1], 2).sum()
-        tv_loss = tv_d + tv_h + tv_w
+        # convention where the object carries a leading channel dimension. tv_loss_vol_sq
+        # dispatches to the fused quantem-cuda kernel when available.
+        tv_loss = tv_loss_vol_sq(ctx.obj)
 
         return tv_loss * self.constraints.tv_vol / ctx.obj.numel()
 
@@ -573,6 +572,10 @@ def apply_hard_constraints(self, pred: torch.Tensor) -> torch.Tensor:
 
         return pred
 
+    def sample_tv_tap_coords(self, coords: torch.Tensor) -> Optional[torch.Tensor]:
+        """Hook for the training loop: returns None (INR TV uses autograd, no tap merging)."""
+        return None
+
     # --- Define get_tv_loss ---
 
     def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor:
@@ -674,6 +677,35 @@ def forward(self, coords: Optional[torch.Tensor] = None) -> torch.Tensor:
 
         return all_densities
 
+    def forward_with_tv_taps(
+        self, coords: torch.Tensor, tap_coords: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Single model call covering the main batch and the volume-TV tap points.
+
+        The out-of-bounds mask and hard constraints apply to the main chunk
+        only; tap densities are returned raw (border-clamped), matching the
+        fallback path in ``get_volume_tv_loss``.
+        """
+        merged = self.model(torch.cat([coords, tap_coords], dim=0))
+        if isinstance(merged, tuple):
+            merged = merged[0]
+        main, taps = merged[: coords.shape[0]], merged[coords.shape[0] :]
+
+        if main.dim() > 1:
+            main = main.squeeze(-1)
+        valid_mask = (
+            (coords[:, 0] >= -1) & (coords[:, 0] <= 1) & (coords[:, 1] >= -1) & (coords[:, 1] <= 1)
+        ).float()
+        if main.dim() > 1:
+            valid_mask = valid_mask.unsqueeze(-1)
+        main = main * valid_mask
+        main = self.apply_hard_constraints(main)
+
+        if taps.dim() == 1:
+            taps = taps.unsqueeze(-1)
+        return main, taps
+
     # Pretrain Loop
 
     def pretrain(
@@ -911,13 +943,37 @@ def from_model(
         obj_model.to(device)
         return obj_model
 
+    def sample_tv_tap_coords(self, coords: torch.Tensor) -> Optional[torch.Tensor]:
+        """
+        Sample the finite-difference tap coordinates for the volume TV loss.
+
+        Returns a (4*n, 3) tensor [base; base+h*ex; base+h*ey; base+h*ez] for n
+        sampled base points, or None when tv_vol == 0.  The training loop
+        concatenates this to all_coords so the TV taps are evaluated in the
+        same model call as the main forward pass.
+        """
+        if self.constraints.tv_vol == 0:
+            return None
+        model = _unwrap(self.model)
+        h = 2.0 / min(model.resolution)
+        num_tv_samples = min(10_000, coords.shape[0])
+        tv_indices = torch.randperm(coords.shape[0], device=coords.device)[:num_tv_samples]
+        tv_coords = coords[tv_indices]  # (n, 3)
+        ex = torch.zeros(3, device=tv_coords.device)
+        ex[0] = h
+        ey = torch.zeros(3, device=tv_coords.device)
+        ey[1] = h
+        ez = torch.zeros(3, device=tv_coords.device)
+        ez[2] = h
+        return torch.cat([tv_coords, tv_coords + ex, tv_coords + ey, tv_coords + ez], dim=0)
+
     # --- Constraints ---
 
     def apply_soft_constraints(self, ctx: ReconstructionContext) -> torch.Tensor:
         soft_loss = torch.tensor(
             0.0, device=ctx.pred.device if ctx.pred is not None else self.device
         )
-        if self.constraints.tv_vol > 0:
+        if self.constraints.tv_plane > 0 or self.constraints.tv_vol > 0:
             assert ctx.coords is not None, "Coordinates must be provided for TV loss"
             assert ctx.pred is not None, "Prediction must be provided for TV loss"
             soft_loss += self.get_tv_loss(ctx)
@@ -943,8 +999,12 @@ def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor:
         assert ctx.coords is not None, "Coordinates must be provided for TV loss"
         assert ctx.pred is not None, "Prediction must be provided for TV loss"
         tv_loss = torch.tensor(0.0, device=ctx.pred.device)
-        tv_loss += self._get_plane_tv_loss()
-        tv_loss += self.get_volume_tv_loss(ctx.coords)
+        if self.constraints.tv_plane > 0:
+            tv_loss += self._get_plane_tv_loss()
+        if self.constraints.tv_vol > 0:
+            tv_loss += self.get_volume_tv_loss(
+                ctx.coords, precomputed_tap_densities=ctx.tv_tap_densities
+            )
         return tv_loss
 
     def _get_plane_tv_loss(self) -> torch.Tensor:
@@ -972,37 +1032,52 @@ def _get_plane_tv_loss(self) -> torch.Tensor:
 
         return self.constraints.tv_plane * torch.stack(per_level).sum()
 
-    def get_volume_tv_loss(self, coords: torch.Tensor) -> torch.Tensor:
+    def get_volume_tv_loss(
+        self,
+        coords: torch.Tensor,
+        precomputed_tap_densities: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         """
         Isotropic volume TV via finite differences. Same form as the autograd
         version (L1 of gradient L2-norm) but avoids double-backward, so it
         works for KPlanesTILTED, CPTilted, and anything else.
-        """
-        num_tv_samples = min(10_000, coords.shape[0])
-        tv_indices = torch.randperm(coords.shape[0], device=coords.device)[:num_tv_samples]
-        tv_coords = coords[tv_indices]  # (N, 3)
 
+        When *precomputed_tap_densities* is provided (a (4N, C) tensor from the
+        merged single-pass forward in the training loop), the model call is
+        skipped entirely and the supplied values are used directly.  When absent
+        the existing batched 4N-point fallback path runs unchanged.
+        """
         model = _unwrap(self.model)
         h = 2.0 / min(model.resolution)
 
-        pred = model(tv_coords)
-        if isinstance(pred, tuple):
-            pred = pred[0]
-        if pred.dim() == 1:
-            pred = pred.unsqueeze(-1)  # (N, 1)
-
-        grads = []
-        for axis in range(3):
-            offset = torch.zeros(3, device=tv_coords.device)
-            offset[axis] = h
-            shifted_pred = self.model(tv_coords + offset)
-            if isinstance(shifted_pred, tuple):
-                shifted_pred = shifted_pred[0]
-            if shifted_pred.dim() == 1:
-                shifted_pred = shifted_pred.unsqueeze(-1)
-            grads.append((shifted_pred - pred) / h)  # (N, 1)
-
-        grad_stack = torch.stack(grads, dim=-1)  # (N, C, 3)
+        if precomputed_tap_densities is not None:
+            batched_pred = precomputed_tap_densities
+        else:
+            num_tv_samples = min(10_000, coords.shape[0])
+            tv_indices = torch.randperm(coords.shape[0], device=coords.device)[:num_tv_samples]
+            tv_coords = coords[tv_indices]  # (N, 3)
+
+            ex = torch.zeros(3, device=tv_coords.device)
+            ex[0] = h
+            ey = torch.zeros(3, device=tv_coords.device)
+            ey[1] = h
+            ez = torch.zeros(3, device=tv_coords.device)
+            ez[2] = h
+
+            batched = torch.cat(
+                [tv_coords, tv_coords + ex, tv_coords + ey, tv_coords + ez], dim=0
+            )  # (4N, 3)
+            batched_pred = model(batched)
+            if isinstance(batched_pred, tuple):
+                batched_pred = batched_pred[0]
+
+        if batched_pred.dim() == 1:
+            batched_pred = batched_pred.unsqueeze(-1)  # (4N, C)
+
+        pred, px, py, pz = batched_pred.chunk(4, dim=0)  # each (N, C)
+        grad_stack = torch.stack(
+            [(px - pred) / h, (py - pred) / h, (pz - pred) / h], dim=-1
+        )  # (N, C, 3)
         grad_norm = torch.norm(grad_stack, dim=-1)  # (N, C)
 
         return self.constraints.tv_vol * grad_norm.mean()

src/quantem/tomography/tomography.py

@@ -210,7 +210,14 @@ def reconstruct(
                 ):
                     all_coords = self.dset.get_coords(batch, N, curr_num_samples_per_ray)
 
-                    all_densities = self.obj_model.forward(all_coords)
+                    tap_coords = self.obj_model.sample_tv_tap_coords(all_coords)
+                    if tap_coords is not None:
+                        all_densities, tv_tap_raw = self.obj_model.forward_with_tv_taps(
+                            all_coords, tap_coords
+                        )
+                    else:
+                        all_densities = self.obj_model.forward(all_coords)
+                        tv_tap_raw = None
 
                     integrated_densities = self.dset.integrate_rays(
                         all_densities,
@@ -225,6 +232,7 @@ def reconstruct(
                         coords=all_coords,
                         pred=pred,
                         all_densities=all_densities,
+                        tv_tap_densities=tv_tap_raw,
                     )
                 )
 

src/quantem/tomography/tomography_context.py

@@ -1,9 +1,10 @@
 from dataclasses import dataclass
 from typing import Optional
-from quantem.core.ml.constraints import BaseContext
 
 import torch
 
+from quantem.core.ml.constraints import BaseContext
+
 
 @dataclass
 class ReconstructionContext(BaseContext):
@@ -28,3 +29,4 @@ class ReconstructionContext(BaseContext):
     pred: Optional[torch.Tensor] = None
     all_densities: Optional[torch.Tensor] = None
     obj: Optional[torch.Tensor] = None
+    tv_tap_densities: Optional[torch.Tensor] = None

src/quantem/tomography/utils.py

@@ -1,6 +1,8 @@
 import torch
 import torch.nn.functional as F
 
+from quantem.core import config
+
 # --- Projection Operator Utils ---
 
 
@@ -69,6 +71,42 @@ def transform_slice(mag_slice):
     return rotated_mags.permute(1, 2, 3, 0)
 
 
+def tv_loss_vol_sq(obj: torch.Tensor) -> torch.Tensor:
+    """Squared-anisotropic volume TV: sum of squared forward differences.
+
+    Computes ``Σ (Δd)² + Σ (Δh)² + Σ (Δw)²`` over the three trailing
+    spatial dims, leaving any leading channel/batch axes intact (they are
+    included in the sum). This is the unnormalized ``tv_vol`` regularizer;
+    callers apply their own ``weight / numel`` scaling.
+
+    When the optional ``quantem-cuda`` package is installed
+    (``pip install quantem[cuda]``), the tensor is on a CUDA device, and the
+    ``use_cuda_kernels`` config option is true (default), this dispatches to
+    the fused CUDA forward/backward kernel — identical math, one kernel
+    launch instead of several large intermediates.
+
+    Args:
+        obj: Tensor of shape ``[..., D, H, W]`` (ndim >= 3).
+
+    Returns:
+        0-dim tensor on the same device as ``obj``; differentiable.
+    """
+    if (
+        obj.is_cuda
+        and obj.dtype == torch.float32
+        and config.get("has_quantem_cuda")
+        and config.get("use_cuda_kernels", default=True)
+    ):
+        from quantem.cuda.core import tv_loss_sq_3d
+
+        return tv_loss_sq_3d(obj)
+
+    tv_d = torch.pow(obj[..., 1:, :, :] - obj[..., :-1, :, :], 2).sum()
+    tv_h = torch.pow(obj[..., :, 1:, :] - obj[..., :, :-1, :], 2).sum()
+    tv_w = torch.pow(obj[..., :, :, 1:] - obj[..., :, :, :-1], 2).sum()
+    return tv_d + tv_h + tv_w
+
+
 def tv_loss_1d(x: torch.Tensor, reduction: str = "mean") -> torch.Tensor:
     """
     1D Total Variation Loss.