Publish HOLYSHT kernel bundle (53048d01ccae)

.gitattributes

@@ -1,35 +1,11 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.so binary
+*.dylib binary
+*.pyd binary
+*.dll binary
+build/torch210-cxx11-cu126-x86_64-linux/_holysht_cuda_53048d0.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu128-x86_64-linux/_holysht_cuda_53048d0.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch210-cxx11-cu130-x86_64-linux/_holysht_cuda_53048d0.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch211-cxx11-cu126-x86_64-linux/_holysht_cuda_53048d0.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch211-cxx11-cu128-x86_64-linux/_holysht_cuda_53048d0.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch211-cxx11-cu130-x86_64-linux/_holysht_cuda_53048d0.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-cxx11-cu129-x86_64-linux/_holysht_cuda_53048d0.abi3.so filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Chris von Csefalvay
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,91 @@
+# HOLYSHT card
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+
+## Name
+
+**HOLYSHT**: highly optimised Legendre/Ylm/SHT
+
+## Summary
+
+HOLYSHT is a focused CUDA acceleration layer for spherical harmonic transforms
+in the `torch-harmonics` ecosystem. It accelerates the Legendre stage, vector
+SHT composition, and inverse-FFT preparation rather than replacing the whole
+upstream library.
+
+## Scope
+
+Included:
+
+- Scalar forward SHT
+- Scalar inverse SHT
+- Vector forward SHT
+- Vector inverse SHT
+- Explicit backward support for the custom scalar and vector kernels
+- BF16 forward paths backed by CUDA real-reduction kernels
+- Benchmark, profiling, and parity-test tooling
+
+Intentionally excluded:
+
+- The earlier DISCO fusion work
+- Experimental side-project code
+- Heavy multi-architecture kernel-builder fan-out in normal development
+
+## Benchmark summary
+
+Measured on 11 April 2026 with PyTorch `2.10.0+cu130`, CUDA `13.0`, batch size
+`4`, an allocation cap of `6 GiB`, and an NVIDIA GB10:
+
+- Scalar forward: up to `4.6x`
+- Scalar inverse: up to `2.0x`
+- Vector forward: up to `8.7x`
+- Vector inverse: up to `8.7x`
+- Scalar forward + backward: `2.3x`
+- Vector forward + backward: `4.3x`
+- BF16 scalar forward: `1.6x`
+- BF16 vector forward: `1.6x`
+
+All `20/20` executed correctness checks passed in that quick sweep.
+
+## Core techniques
+
+1. Architecture-aware launch selection between `tile_l=4` and `tile_l=8`.
+2. Shared-memory large-grid kernels for scalar, vector, and BF16 real paths.
+3. Dedicated vector forward and inverse kernels to remove the old Python-side
+   packing hot path.
+4. Real-valued BF16 Legendre reductions with float accumulation.
+5. Explicit autograd wrappers for scalar and vector CUDA paths.
+6. CUDA-side `irfft` preparation to avoid extra Python tensor passes.
+7. Local torch JIT build caching under `build/torch_extensions` for safer
+   day-to-day iteration on GB10.
+
+## Resource snapshot
+
+From `cuobjdump --dump-resource-usage` on the local `sm_120` build:
+
+- Scalar forward large kernel: `38` registers/thread, `3136` B shared/block,
+  `256` threads/block, `6` active blocks/SM.
+- Vector forward large kernel: `37` registers/thread, `5248` B shared/block,
+  `256` threads/block, `6` active blocks/SM.
+- BF16 forward large kernel: `34` registers/thread, `2080` B shared/block,
+  `256` threads/block, `6` active blocks/SM.
+- `prepare_irfft`: `19` registers/thread, no shared memory.
+
+## Practical constraints
+
+- HOLYSHT still depends on `torch-harmonics` for quadrature weight generation.
+- The default local build targets `12.0+PTX` on GB10-class systems.
+- `ncu` may require elevated GPU counter permissions; the supplied script falls
+  back to `cuobjdump` resource reporting when those counters are unavailable.
+- Very large inverse-heavy cases should still be run under an allocation cap on
+  unified-memory systems.
+
+## Intended audience
+
+- Researchers already using `torch-harmonics` who want faster SHT execution.
+- Neural operator workloads where SHT dominates forward or training cost.
+- CUDA developers who want a small, inspectable codebase with real profiling and
+  resource-report hooks rather than hand-wavy speed claims.

build.toml

@@ -0,0 +1,38 @@
+# HOLYSHT
+# Author: Chris von Csefalvay
+# Licence: MIT
+# Repository: https://github.com/chrisvoncsefalvay/holysht
+# Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+
+[general]
+name = "holysht"
+version = 1
+license = "mit"
+upstream = "https://github.com/NVIDIA/torch-harmonics"
+backends = ["cuda"]
+
+[general.hub]
+repo-id = "chrisvoncsefalvay/holysht"
+
+[torch]
+src = [
+    "torch-ext/torch_binding.cpp",
+    "torch-ext/torch_binding.h",
+    "torch-ext/registration_select.h",
+    "torch-ext/local_registration.h",
+]
+pyext = ["py"]
+
+[kernel.holysht_legendre]
+backend = "cuda"
+depends = ["torch"]
+src = ["cuda/fused_legendre.cu"]
+cuda-capabilities = ["12.0"]
+cuda-flags = ["-O3", "-lineinfo", "-Xptxas=-warn-spills", "--use_fast_math"]
+
+[kernel.holysht_sht]
+backend = "cuda"
+depends = ["torch"]
+src = ["cuda/fused_sht.cu"]
+cuda-capabilities = ["12.0"]
+cuda-flags = ["-O3", "-lineinfo", "-Xptxas=-warn-spills", "--use_fast_math"]

build/torch210-cxx11-cu126-x86_64-linux/__init__.py

@@ -0,0 +1,673 @@
+"""HOLYSHT: Highly Optimised Legendre/Ylm/SHT.
+
+CUDA-accelerated spherical harmonic transforms designed as a practical,
+production-oriented companion to torch-harmonics.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+import contextlib
+import os
+from typing import Optional
+import torch
+import torch.nn as nn
+
+__all__ = [
+    "RealSHT",
+    "InverseRealSHT",
+    "RealVectorSHT",
+    "InverseRealVectorSHT",
+    "legendre_forward",
+    "legendre_inverse",
+    "sht_forward",
+    "sht_inverse",
+]
+
+# Prefer kernel-builder's generated alias module on packaged builds, then fall
+# back to the local single-machine JIT loader for development.
+try:
+    from ._ops import ops as _ops
+    _HAS_CUDA_EXT = True
+except ModuleNotFoundError:
+    try:
+        from ._jit_ops import ops as _ops
+        _HAS_CUDA_EXT = True
+    except ImportError:
+        _HAS_CUDA_EXT = False
+except ImportError:
+    _HAS_CUDA_EXT = False
+
+
+def _can_use_cuda_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight_t.dtype == torch.float32
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_real_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype in (torch.float32, torch.bfloat16)
+        and weight_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_vector(input: torch.Tensor, weight0_t: Optional[torch.Tensor], weight1_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight0_t is not None
+        and weight1_t is not None
+        and input.is_cuda
+        and weight0_t.is_cuda
+        and weight1_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight0_t.dtype == torch.float32
+        and weight1_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight0_t.is_contiguous()
+        and weight1_t.is_contiguous()
+    )
+
+
+def _mul_i(x: torch.Tensor) -> torch.Tensor:
+    """Multiply a complex tensor by +i without promoting dtype."""
+    return torch.complex(-x.imag, x.real)
+
+
+@contextlib.contextmanager
+def _nvtx_range(name: str):
+    enabled = os.environ.get("HOLYSHT_ENABLE_NVTX", "0") == "1"
+    if enabled and torch.cuda.is_available():
+        torch.cuda.nvtx.range_push(name)
+        try:
+            yield
+        finally:
+            torch.cuda.nvtx.range_pop()
+    else:
+        yield
+
+
+def _prepare_irfft_input(x: torch.Tensor, nlon: int, active_mmax: Optional[int] = None) -> torch.Tensor:
+    """Pad/clean an rFFT-format complex tensor before irfft."""
+    active_mmax = x.size(-1) if active_mmax is None else active_mmax
+    full_mmax = nlon // 2 + 1
+
+    if x.size(-1) == full_mmax:
+        out = x.contiguous()
+    else:
+        out_shape = list(x.shape)
+        out_shape[-1] = full_mmax
+        out = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+        out[..., :x.size(-1)] = x
+
+    if _HAS_CUDA_EXT and out.is_cuda and out.dtype == torch.complex64:
+        orig_shape = out.shape
+        flat = out.reshape(-1, orig_shape[-2], orig_shape[-1]).contiguous()
+        _ops.sht_prepare_irfft(flat, active_mmax, nlon)
+        return flat.reshape(orig_shape)
+
+    out[..., 0] = out[..., 0].real.to(torch.complex64)
+    if nlon % 2 == 0:
+        nyquist_idx = nlon // 2
+        if nyquist_idx < out.size(-1):
+            out[..., nyquist_idx] = out[..., nyquist_idx].real.to(torch.complex64)
+    return out
+
+
+class _FusedLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(1), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(0), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreForwardRealFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_forward_real(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_inverse_real(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedVectorLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(0), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(1), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+class _FusedVectorLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(1), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(0), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+# ============================================================================
+# Fused Legendre Transform
+# ============================================================================
+
+def fused_legendre_forward(
+    input: torch.Tensor,       # [B, nlat, mmax] complex64
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32 (original torch-harmonics layout)
+    weight_t: Optional[torch.Tensor] = None,  # [lmax, nlat, mmax] pre-transposed
+) -> torch.Tensor:
+    """Fused forward Legendre transform operating on complex tensors.
+
+    Computes out[b,l,m] = Σ_k weights[m,l,k] · input[b,k,m] for complex input,
+    fusing the real and imaginary multiplications into a single pass.
+
+    Uses the custom CUDA kernels when the extension is available, otherwise
+    falls back to a stacked einsum.
+    """
+    B = input.size(0)
+    nlat = input.size(1)
+    mmax = input.size(2)
+    lmax = weights.size(1)
+
+    if weight_t is None:
+        weight_t = weights.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, weight_t):
+        # Adaptive CUDA kernel: small-grid direct path + large-grid tiled path.
+        return _FusedLegendreForwardFn.apply(input, weight_t)
+    else:
+        # Fallback: stacked einsum (1.9x speedup over reference 2x einsum)
+        x = torch.view_as_real(input)  # [B, nlat, mmax, 2]
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)  # [2B, nlat, mmax]
+        w = weights.to(x_stacked.dtype)
+        out_stacked = torch.einsum("bkm,mlk->blm", x_stacked, w)
+        out_re = out_stacked[:B]
+        out_im = out_stacked[B:]
+        return torch.complex(out_re, out_im)
+
+
+def fused_legendre_inverse(
+    input: torch.Tensor,       # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Fused inverse Legendre transform."""
+    B = input.size(0)
+    lmax = input.size(1)
+    mmax = input.size(2)
+    nlat = pct.size(2)
+
+    if pct_t is None:
+        pct_t = pct.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, pct_t):
+        return _FusedLegendreInverseFn.apply(input, pct_t)
+    else:
+        x = torch.view_as_real(input)
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)
+        p = pct.to(x_stacked.dtype)
+        out_stacked = torch.einsum("blm,mlk->bkm", x_stacked, p)
+        return torch.complex(out_stacked[:B], out_stacked[B:])
+
+
+def fused_legendre_forward_real(
+    input: torch.Tensor,       # [B, nlat, mmax] float32 or bfloat16
+    weight_t: torch.Tensor,    # [lmax, nlat, mmax] float32
+) -> torch.Tensor:
+    """Real-valued forward Legendre transform with float accumulation."""
+    if _can_use_cuda_real_legendre(input, weight_t):
+        return _FusedLegendreForwardRealFn.apply(input, weight_t)
+
+    return torch.einsum("bkm,lkm->blm", input.float(), weight_t)
+
+
+# ============================================================================
+# Fused SHT (complete pipeline)
+# ============================================================================
+
+def fused_sht_forward(
+    x: torch.Tensor,           # [B, nlat, nlon] float32
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32
+    mmax: int,
+    weight_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused forward SHT: rfft → fused Legendre → complex coefficients.
+
+    Replaces RealSHT.forward() with fewer intermediate allocations.
+    """
+    with _nvtx_range("holysht.scalar_forward"):
+        x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+        x_fft = x_fft[..., :mmax]
+        return fused_legendre_forward(x_fft, weights, weight_t)
+
+
+def fused_sht_inverse(
+    coeffs: torch.Tensor,      # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    nlon: int,
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused inverse SHT: fused Legendre → zero-pad → irfft."""
+    with _nvtx_range("holysht.scalar_inverse"):
+        x = fused_legendre_inverse(coeffs, pct, pct_t)
+        x = _prepare_irfft_input(x, nlon, coeffs.size(-1))
+        return torch.fft.irfft(x, n=nlon, dim=-1, norm="forward")
+
+
+# ============================================================================
+# nn.Module wrappers (drop-in replacements for torch-harmonics)
+# ============================================================================
+
+class RealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealSHT``.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealSHT
+        ref = RealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("weights", ref.weights.to(w_dtype))
+        self.register_buffer("weight_t", ref.weights.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._use_bf16:
+            with _nvtx_range("holysht.scalar_forward_bf16"):
+                x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+                x_fft = x_fft[..., :self.mmax]
+                xr = torch.view_as_real(x_fft)
+                if _HAS_CUDA_EXT and x.is_cuda and not x.requires_grad:
+                    xr_bf16 = xr.bfloat16().contiguous()
+                    out_re = fused_legendre_forward_real(xr_bf16[..., 0].contiguous(), self.weight_t)
+                    out_im = fused_legendre_forward_real(xr_bf16[..., 1].contiguous(), self.weight_t)
+                    return torch.complex(out_re, out_im)
+
+                B = x.size(0)
+                xs = torch.cat([xr[..., 0], xr[..., 1]], dim=0).bfloat16()
+                out = torch.einsum("bkm,mlk->blm", xs, self.weights).float()
+                return torch.complex(out[:B], out[B:])
+        return fused_sht_forward(x, self.weights, self.mmax, self.weight_t)
+
+
+class InverseRealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealSHT
+        ref = InverseRealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("pct", ref.pct.float())
+        self.register_buffer("pct_t", ref.pct.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return fused_sht_inverse(x, self.pct, self.nlon, self.pct_t)
+
+
+class RealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealVectorSHT``.
+
+    Reduces eight reference einsums to two composed Legendre passes on the
+    default FP32 CUDA path.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealVectorSHT
+        ref = RealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("w0", ref.weights[0].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w1", ref.weights[1].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w0_t", ref.weights[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("w1_t", ref.weights[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.nlat and x.shape[-1] == self.nlon
+
+        with _nvtx_range("holysht.vector_forward"):
+            x = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+            mmax = self.mmax
+            x = x[..., :mmax].contiguous()
+
+            if (not self._use_bf16) and _can_use_cuda_vector(x, self.w0_t, self.w1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.nlat, mmax).contiguous()
+                out = _FusedVectorLegendreForwardFn.apply(x_flat, self.w0_t, self.w1_t)
+                return out.reshape(B_shape + (2, self.lmax, mmax))
+
+            x = torch.view_as_real(x)  # [..., 2, nlat, mmax, 2]
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.nlat, mmax)
+            x01_flat = x01.reshape(-1, self.nlat, mmax)
+            x10_flat = x10.reshape(-1, self.nlat, mmax)
+            x11_flat = x11.reshape(-1, self.nlat, mmax)
+            B = x00_flat.shape[0]
+
+            if self._use_bf16 and _HAS_CUDA_EXT and x00_flat.is_cuda and not x.requires_grad:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0).bfloat16().contiguous()
+                out_w0 = fused_legendre_forward_real(stacked_w0, self.w0_t)
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0).bfloat16().contiguous()
+                out_w1 = fused_legendre_forward_real(stacked_w1, self.w1_t)
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+            else:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w0 = stacked_w0.bfloat16()
+                out_w0 = torch.einsum("bkm,mlk->blm", stacked_w0, self.w0)
+                if self._use_bf16:
+                    out_w0 = out_w0.float()
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w1 = stacked_w1.bfloat16()
+                out_w1 = torch.einsum("bkm,mlk->blm", stacked_w1, self.w1)
+                if self._use_bf16:
+                    out_w1 = out_w1.float()
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+
+            sph_re = r00 - s11
+            sph_im = r01 + s10
+            tor_re = -s01 - r10
+            tor_im = s00 - r11
+
+            out_shape = list(B_shape) + [2, self.lmax, mmax, 2]
+            xout = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+            xout[..., 0, :, :, 0] = sph_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 0, :, :, 1] = sph_im.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 0] = tor_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 1] = tor_im.reshape(B_shape + (self.lmax, mmax))
+            return torch.view_as_complex(xout)
+
+
+class InverseRealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealVectorSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealVectorSHT
+        ref = InverseRealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("d0", ref.dpct[0].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d1", ref.dpct[1].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d0_t", ref.dpct[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("d1_t", ref.dpct[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.lmax and x.shape[-1] == self.mmax
+
+        with _nvtx_range("holysht.vector_inverse"):
+            x = x.contiguous()
+            if _can_use_cuda_vector(x, self.d0_t, self.d1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.lmax, self.mmax).contiguous()
+                x_out = _FusedVectorLegendreInverseFn.apply(x_flat, self.d0_t, self.d1_t)
+                x_out = x_out.reshape(B_shape + (2, self.nlat, self.mmax))
+                x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+                return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+            x = torch.view_as_real(x)  # [..., 2, lmax, mmax, 2]
+            mmax = self.mmax
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.lmax, mmax)
+            x01_flat = x01.reshape(-1, self.lmax, mmax)
+            x10_flat = x10.reshape(-1, self.lmax, mmax)
+            x11_flat = x11.reshape(-1, self.lmax, mmax)
+            B = x00_flat.shape[0]
+
+            stacked_d0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+            out_d0 = torch.einsum("blm,mlk->bkm", stacked_d0, self.d0)
+            r00, r01, r10, r11 = out_d0[:B], out_d0[B:2 * B], out_d0[2 * B:3 * B], out_d0[3 * B:]
+
+            stacked_d1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+            out_d1 = torch.einsum("blm,mlk->bkm", stacked_d1, self.d1)
+            s11, s10, s01, s00 = out_d1[:B], out_d1[B:2 * B], out_d1[2 * B:3 * B], out_d1[3 * B:]
+
+            srl = r00 - s11
+            sim = r01 + s10
+            trl = -s01 - r10
+            tim = s00 - r11
+
+            out_k = self.nlat
+            srl = srl.reshape(B_shape + (out_k, mmax))
+            sim = sim.reshape(B_shape + (out_k, mmax))
+            trl = trl.reshape(B_shape + (out_k, mmax))
+            tim = tim.reshape(B_shape + (out_k, mmax))
+
+            s = torch.stack((srl, sim), -1)
+            t = torch.stack((trl, tim), -1)
+            xs = torch.stack((s, t), -4)
+            x_out = torch.view_as_complex(xs)
+            x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+            return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+
+legendre_forward = fused_legendre_forward
+legendre_inverse = fused_legendre_inverse
+sht_forward = fused_sht_forward
+sht_inverse = fused_sht_inverse
+
+# Backwards-compatible aliases from the research prototype.
+FusedRealSHT = RealSHT
+FusedInverseRealSHT = InverseRealSHT
+FusedRealVectorSHT = RealVectorSHT
+FusedInverseRealVectorSHT = InverseRealVectorSHT

build/torch210-cxx11-cu126-x86_64-linux/_jit_ops.py

@@ -0,0 +1,72 @@
+"""Local CUDA extension loader for HOLYSHT.
+
+This keeps development off the heavyweight `kernel-builder` path by compiling a
+small torch extension for the current machine only. The compiled artefacts are
+cached under `build/torch_extensions`.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+from torch.utils.cpp_extension import load
+
+
+_ROOT = Path(__file__).resolve().parents[2]
+_BUILD_DIR = _ROOT / "build" / "torch_extensions"
+_EXTENSION_NAME = "holysht_ops"
+
+
+def _default_arch_list() -> str:
+    if not torch.cuda.is_available():
+        return ""
+
+    major, minor = torch.cuda.get_device_capability(0)
+    if (major, minor) >= (12, 1):
+        return "12.0+PTX"
+    return f"{major}.{minor}"
+
+
+def _cuda_flags() -> list[str]:
+    flags = [
+        "-O3",
+        "-lineinfo",
+        "-Xptxas=-warn-spills",
+        "--expt-relaxed-constexpr",
+    ]
+    if os.environ.get("HOLYSHT_USE_FAST_MATH", "1") != "0":
+        flags.append("--use_fast_math")
+    return flags
+
+
+def _load_ops():
+    os.environ.setdefault("MAX_JOBS", "1")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _default_arch_list())
+    _BUILD_DIR.mkdir(parents=True, exist_ok=True)
+
+    load(
+        name=_EXTENSION_NAME,
+        sources=[
+            str(_ROOT / "torch-ext" / "torch_binding.cpp"),
+            str(_ROOT / "cuda" / "fused_legendre.cu"),
+            str(_ROOT / "cuda" / "fused_sht.cu"),
+        ],
+        extra_include_paths=[str(_ROOT / "torch-ext")],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=_cuda_flags(),
+        build_directory=str(_BUILD_DIR),
+        verbose=os.environ.get("HOLYSHT_VERBOSE_BUILD", "0") == "1",
+        with_cuda=True,
+        is_python_module=False,
+    )
+    return getattr(torch.ops, _EXTENSION_NAME)
+
+
+ops = _load_ops()

build/torch210-cxx11-cu126-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _holysht_cuda_53048d0
+ops = torch.ops._holysht_cuda_53048d0
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_holysht_cuda_53048d0::{op_name}"

build/torch210-cxx11-cu126-x86_64-linux/holysht/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch210-cxx11-cu126-x86_64-linux/metadata.json

@@ -0,0 +1,9 @@
+{
+  "version": 1,
+  "license": "mit",
+  "upstream": "https://github.com/NVIDIA/torch-harmonics",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda"
+  }
+}

build/torch210-cxx11-cu128-x86_64-linux/__init__.py

@@ -0,0 +1,673 @@
+"""HOLYSHT: Highly Optimised Legendre/Ylm/SHT.
+
+CUDA-accelerated spherical harmonic transforms designed as a practical,
+production-oriented companion to torch-harmonics.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+import contextlib
+import os
+from typing import Optional
+import torch
+import torch.nn as nn
+
+__all__ = [
+    "RealSHT",
+    "InverseRealSHT",
+    "RealVectorSHT",
+    "InverseRealVectorSHT",
+    "legendre_forward",
+    "legendre_inverse",
+    "sht_forward",
+    "sht_inverse",
+]
+
+# Prefer kernel-builder's generated alias module on packaged builds, then fall
+# back to the local single-machine JIT loader for development.
+try:
+    from ._ops import ops as _ops
+    _HAS_CUDA_EXT = True
+except ModuleNotFoundError:
+    try:
+        from ._jit_ops import ops as _ops
+        _HAS_CUDA_EXT = True
+    except ImportError:
+        _HAS_CUDA_EXT = False
+except ImportError:
+    _HAS_CUDA_EXT = False
+
+
+def _can_use_cuda_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight_t.dtype == torch.float32
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_real_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype in (torch.float32, torch.bfloat16)
+        and weight_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_vector(input: torch.Tensor, weight0_t: Optional[torch.Tensor], weight1_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight0_t is not None
+        and weight1_t is not None
+        and input.is_cuda
+        and weight0_t.is_cuda
+        and weight1_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight0_t.dtype == torch.float32
+        and weight1_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight0_t.is_contiguous()
+        and weight1_t.is_contiguous()
+    )
+
+
+def _mul_i(x: torch.Tensor) -> torch.Tensor:
+    """Multiply a complex tensor by +i without promoting dtype."""
+    return torch.complex(-x.imag, x.real)
+
+
+@contextlib.contextmanager
+def _nvtx_range(name: str):
+    enabled = os.environ.get("HOLYSHT_ENABLE_NVTX", "0") == "1"
+    if enabled and torch.cuda.is_available():
+        torch.cuda.nvtx.range_push(name)
+        try:
+            yield
+        finally:
+            torch.cuda.nvtx.range_pop()
+    else:
+        yield
+
+
+def _prepare_irfft_input(x: torch.Tensor, nlon: int, active_mmax: Optional[int] = None) -> torch.Tensor:
+    """Pad/clean an rFFT-format complex tensor before irfft."""
+    active_mmax = x.size(-1) if active_mmax is None else active_mmax
+    full_mmax = nlon // 2 + 1
+
+    if x.size(-1) == full_mmax:
+        out = x.contiguous()
+    else:
+        out_shape = list(x.shape)
+        out_shape[-1] = full_mmax
+        out = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+        out[..., :x.size(-1)] = x
+
+    if _HAS_CUDA_EXT and out.is_cuda and out.dtype == torch.complex64:
+        orig_shape = out.shape
+        flat = out.reshape(-1, orig_shape[-2], orig_shape[-1]).contiguous()
+        _ops.sht_prepare_irfft(flat, active_mmax, nlon)
+        return flat.reshape(orig_shape)
+
+    out[..., 0] = out[..., 0].real.to(torch.complex64)
+    if nlon % 2 == 0:
+        nyquist_idx = nlon // 2
+        if nyquist_idx < out.size(-1):
+            out[..., nyquist_idx] = out[..., nyquist_idx].real.to(torch.complex64)
+    return out
+
+
+class _FusedLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(1), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(0), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreForwardRealFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_forward_real(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_inverse_real(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedVectorLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(0), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(1), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+class _FusedVectorLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(1), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(0), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+# ============================================================================
+# Fused Legendre Transform
+# ============================================================================
+
+def fused_legendre_forward(
+    input: torch.Tensor,       # [B, nlat, mmax] complex64
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32 (original torch-harmonics layout)
+    weight_t: Optional[torch.Tensor] = None,  # [lmax, nlat, mmax] pre-transposed
+) -> torch.Tensor:
+    """Fused forward Legendre transform operating on complex tensors.
+
+    Computes out[b,l,m] = Σ_k weights[m,l,k] · input[b,k,m] for complex input,
+    fusing the real and imaginary multiplications into a single pass.
+
+    Uses the custom CUDA kernels when the extension is available, otherwise
+    falls back to a stacked einsum.
+    """
+    B = input.size(0)
+    nlat = input.size(1)
+    mmax = input.size(2)
+    lmax = weights.size(1)
+
+    if weight_t is None:
+        weight_t = weights.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, weight_t):
+        # Adaptive CUDA kernel: small-grid direct path + large-grid tiled path.
+        return _FusedLegendreForwardFn.apply(input, weight_t)
+    else:
+        # Fallback: stacked einsum (1.9x speedup over reference 2x einsum)
+        x = torch.view_as_real(input)  # [B, nlat, mmax, 2]
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)  # [2B, nlat, mmax]
+        w = weights.to(x_stacked.dtype)
+        out_stacked = torch.einsum("bkm,mlk->blm", x_stacked, w)
+        out_re = out_stacked[:B]
+        out_im = out_stacked[B:]
+        return torch.complex(out_re, out_im)
+
+
+def fused_legendre_inverse(
+    input: torch.Tensor,       # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Fused inverse Legendre transform."""
+    B = input.size(0)
+    lmax = input.size(1)
+    mmax = input.size(2)
+    nlat = pct.size(2)
+
+    if pct_t is None:
+        pct_t = pct.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, pct_t):
+        return _FusedLegendreInverseFn.apply(input, pct_t)
+    else:
+        x = torch.view_as_real(input)
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)
+        p = pct.to(x_stacked.dtype)
+        out_stacked = torch.einsum("blm,mlk->bkm", x_stacked, p)
+        return torch.complex(out_stacked[:B], out_stacked[B:])
+
+
+def fused_legendre_forward_real(
+    input: torch.Tensor,       # [B, nlat, mmax] float32 or bfloat16
+    weight_t: torch.Tensor,    # [lmax, nlat, mmax] float32
+) -> torch.Tensor:
+    """Real-valued forward Legendre transform with float accumulation."""
+    if _can_use_cuda_real_legendre(input, weight_t):
+        return _FusedLegendreForwardRealFn.apply(input, weight_t)
+
+    return torch.einsum("bkm,lkm->blm", input.float(), weight_t)
+
+
+# ============================================================================
+# Fused SHT (complete pipeline)
+# ============================================================================
+
+def fused_sht_forward(
+    x: torch.Tensor,           # [B, nlat, nlon] float32
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32
+    mmax: int,
+    weight_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused forward SHT: rfft → fused Legendre → complex coefficients.
+
+    Replaces RealSHT.forward() with fewer intermediate allocations.
+    """
+    with _nvtx_range("holysht.scalar_forward"):
+        x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+        x_fft = x_fft[..., :mmax]
+        return fused_legendre_forward(x_fft, weights, weight_t)
+
+
+def fused_sht_inverse(
+    coeffs: torch.Tensor,      # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    nlon: int,
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused inverse SHT: fused Legendre → zero-pad → irfft."""
+    with _nvtx_range("holysht.scalar_inverse"):
+        x = fused_legendre_inverse(coeffs, pct, pct_t)
+        x = _prepare_irfft_input(x, nlon, coeffs.size(-1))
+        return torch.fft.irfft(x, n=nlon, dim=-1, norm="forward")
+
+
+# ============================================================================
+# nn.Module wrappers (drop-in replacements for torch-harmonics)
+# ============================================================================
+
+class RealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealSHT``.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealSHT
+        ref = RealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("weights", ref.weights.to(w_dtype))
+        self.register_buffer("weight_t", ref.weights.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._use_bf16:
+            with _nvtx_range("holysht.scalar_forward_bf16"):
+                x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+                x_fft = x_fft[..., :self.mmax]
+                xr = torch.view_as_real(x_fft)
+                if _HAS_CUDA_EXT and x.is_cuda and not x.requires_grad:
+                    xr_bf16 = xr.bfloat16().contiguous()
+                    out_re = fused_legendre_forward_real(xr_bf16[..., 0].contiguous(), self.weight_t)
+                    out_im = fused_legendre_forward_real(xr_bf16[..., 1].contiguous(), self.weight_t)
+                    return torch.complex(out_re, out_im)
+
+                B = x.size(0)
+                xs = torch.cat([xr[..., 0], xr[..., 1]], dim=0).bfloat16()
+                out = torch.einsum("bkm,mlk->blm", xs, self.weights).float()
+                return torch.complex(out[:B], out[B:])
+        return fused_sht_forward(x, self.weights, self.mmax, self.weight_t)
+
+
+class InverseRealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealSHT
+        ref = InverseRealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("pct", ref.pct.float())
+        self.register_buffer("pct_t", ref.pct.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return fused_sht_inverse(x, self.pct, self.nlon, self.pct_t)
+
+
+class RealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealVectorSHT``.
+
+    Reduces eight reference einsums to two composed Legendre passes on the
+    default FP32 CUDA path.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealVectorSHT
+        ref = RealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("w0", ref.weights[0].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w1", ref.weights[1].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w0_t", ref.weights[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("w1_t", ref.weights[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.nlat and x.shape[-1] == self.nlon
+
+        with _nvtx_range("holysht.vector_forward"):
+            x = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+            mmax = self.mmax
+            x = x[..., :mmax].contiguous()
+
+            if (not self._use_bf16) and _can_use_cuda_vector(x, self.w0_t, self.w1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.nlat, mmax).contiguous()
+                out = _FusedVectorLegendreForwardFn.apply(x_flat, self.w0_t, self.w1_t)
+                return out.reshape(B_shape + (2, self.lmax, mmax))
+
+            x = torch.view_as_real(x)  # [..., 2, nlat, mmax, 2]
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.nlat, mmax)
+            x01_flat = x01.reshape(-1, self.nlat, mmax)
+            x10_flat = x10.reshape(-1, self.nlat, mmax)
+            x11_flat = x11.reshape(-1, self.nlat, mmax)
+            B = x00_flat.shape[0]
+
+            if self._use_bf16 and _HAS_CUDA_EXT and x00_flat.is_cuda and not x.requires_grad:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0).bfloat16().contiguous()
+                out_w0 = fused_legendre_forward_real(stacked_w0, self.w0_t)
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0).bfloat16().contiguous()
+                out_w1 = fused_legendre_forward_real(stacked_w1, self.w1_t)
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+            else:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w0 = stacked_w0.bfloat16()
+                out_w0 = torch.einsum("bkm,mlk->blm", stacked_w0, self.w0)
+                if self._use_bf16:
+                    out_w0 = out_w0.float()
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w1 = stacked_w1.bfloat16()
+                out_w1 = torch.einsum("bkm,mlk->blm", stacked_w1, self.w1)
+                if self._use_bf16:
+                    out_w1 = out_w1.float()
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+
+            sph_re = r00 - s11
+            sph_im = r01 + s10
+            tor_re = -s01 - r10
+            tor_im = s00 - r11
+
+            out_shape = list(B_shape) + [2, self.lmax, mmax, 2]
+            xout = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+            xout[..., 0, :, :, 0] = sph_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 0, :, :, 1] = sph_im.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 0] = tor_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 1] = tor_im.reshape(B_shape + (self.lmax, mmax))
+            return torch.view_as_complex(xout)
+
+
+class InverseRealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealVectorSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealVectorSHT
+        ref = InverseRealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("d0", ref.dpct[0].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d1", ref.dpct[1].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d0_t", ref.dpct[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("d1_t", ref.dpct[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.lmax and x.shape[-1] == self.mmax
+
+        with _nvtx_range("holysht.vector_inverse"):
+            x = x.contiguous()
+            if _can_use_cuda_vector(x, self.d0_t, self.d1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.lmax, self.mmax).contiguous()
+                x_out = _FusedVectorLegendreInverseFn.apply(x_flat, self.d0_t, self.d1_t)
+                x_out = x_out.reshape(B_shape + (2, self.nlat, self.mmax))
+                x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+                return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+            x = torch.view_as_real(x)  # [..., 2, lmax, mmax, 2]
+            mmax = self.mmax
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.lmax, mmax)
+            x01_flat = x01.reshape(-1, self.lmax, mmax)
+            x10_flat = x10.reshape(-1, self.lmax, mmax)
+            x11_flat = x11.reshape(-1, self.lmax, mmax)
+            B = x00_flat.shape[0]
+
+            stacked_d0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+            out_d0 = torch.einsum("blm,mlk->bkm", stacked_d0, self.d0)
+            r00, r01, r10, r11 = out_d0[:B], out_d0[B:2 * B], out_d0[2 * B:3 * B], out_d0[3 * B:]
+
+            stacked_d1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+            out_d1 = torch.einsum("blm,mlk->bkm", stacked_d1, self.d1)
+            s11, s10, s01, s00 = out_d1[:B], out_d1[B:2 * B], out_d1[2 * B:3 * B], out_d1[3 * B:]
+
+            srl = r00 - s11
+            sim = r01 + s10
+            trl = -s01 - r10
+            tim = s00 - r11
+
+            out_k = self.nlat
+            srl = srl.reshape(B_shape + (out_k, mmax))
+            sim = sim.reshape(B_shape + (out_k, mmax))
+            trl = trl.reshape(B_shape + (out_k, mmax))
+            tim = tim.reshape(B_shape + (out_k, mmax))
+
+            s = torch.stack((srl, sim), -1)
+            t = torch.stack((trl, tim), -1)
+            xs = torch.stack((s, t), -4)
+            x_out = torch.view_as_complex(xs)
+            x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+            return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+
+legendre_forward = fused_legendre_forward
+legendre_inverse = fused_legendre_inverse
+sht_forward = fused_sht_forward
+sht_inverse = fused_sht_inverse
+
+# Backwards-compatible aliases from the research prototype.
+FusedRealSHT = RealSHT
+FusedInverseRealSHT = InverseRealSHT
+FusedRealVectorSHT = RealVectorSHT
+FusedInverseRealVectorSHT = InverseRealVectorSHT

build/torch210-cxx11-cu128-x86_64-linux/_jit_ops.py

@@ -0,0 +1,72 @@
+"""Local CUDA extension loader for HOLYSHT.
+
+This keeps development off the heavyweight `kernel-builder` path by compiling a
+small torch extension for the current machine only. The compiled artefacts are
+cached under `build/torch_extensions`.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+from torch.utils.cpp_extension import load
+
+
+_ROOT = Path(__file__).resolve().parents[2]
+_BUILD_DIR = _ROOT / "build" / "torch_extensions"
+_EXTENSION_NAME = "holysht_ops"
+
+
+def _default_arch_list() -> str:
+    if not torch.cuda.is_available():
+        return ""
+
+    major, minor = torch.cuda.get_device_capability(0)
+    if (major, minor) >= (12, 1):
+        return "12.0+PTX"
+    return f"{major}.{minor}"
+
+
+def _cuda_flags() -> list[str]:
+    flags = [
+        "-O3",
+        "-lineinfo",
+        "-Xptxas=-warn-spills",
+        "--expt-relaxed-constexpr",
+    ]
+    if os.environ.get("HOLYSHT_USE_FAST_MATH", "1") != "0":
+        flags.append("--use_fast_math")
+    return flags
+
+
+def _load_ops():
+    os.environ.setdefault("MAX_JOBS", "1")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _default_arch_list())
+    _BUILD_DIR.mkdir(parents=True, exist_ok=True)
+
+    load(
+        name=_EXTENSION_NAME,
+        sources=[
+            str(_ROOT / "torch-ext" / "torch_binding.cpp"),
+            str(_ROOT / "cuda" / "fused_legendre.cu"),
+            str(_ROOT / "cuda" / "fused_sht.cu"),
+        ],
+        extra_include_paths=[str(_ROOT / "torch-ext")],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=_cuda_flags(),
+        build_directory=str(_BUILD_DIR),
+        verbose=os.environ.get("HOLYSHT_VERBOSE_BUILD", "0") == "1",
+        with_cuda=True,
+        is_python_module=False,
+    )
+    return getattr(torch.ops, _EXTENSION_NAME)
+
+
+ops = _load_ops()

build/torch210-cxx11-cu128-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _holysht_cuda_53048d0
+ops = torch.ops._holysht_cuda_53048d0
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_holysht_cuda_53048d0::{op_name}"

build/torch210-cxx11-cu128-x86_64-linux/holysht/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch210-cxx11-cu128-x86_64-linux/metadata.json

@@ -0,0 +1,12 @@
+{
+  "version": 1,
+  "license": "mit",
+  "upstream": "https://github.com/NVIDIA/torch-harmonics",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  }
+}

build/torch210-cxx11-cu130-x86_64-linux/__init__.py

@@ -0,0 +1,673 @@
+"""HOLYSHT: Highly Optimised Legendre/Ylm/SHT.
+
+CUDA-accelerated spherical harmonic transforms designed as a practical,
+production-oriented companion to torch-harmonics.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+import contextlib
+import os
+from typing import Optional
+import torch
+import torch.nn as nn
+
+__all__ = [
+    "RealSHT",
+    "InverseRealSHT",
+    "RealVectorSHT",
+    "InverseRealVectorSHT",
+    "legendre_forward",
+    "legendre_inverse",
+    "sht_forward",
+    "sht_inverse",
+]
+
+# Prefer kernel-builder's generated alias module on packaged builds, then fall
+# back to the local single-machine JIT loader for development.
+try:
+    from ._ops import ops as _ops
+    _HAS_CUDA_EXT = True
+except ModuleNotFoundError:
+    try:
+        from ._jit_ops import ops as _ops
+        _HAS_CUDA_EXT = True
+    except ImportError:
+        _HAS_CUDA_EXT = False
+except ImportError:
+    _HAS_CUDA_EXT = False
+
+
+def _can_use_cuda_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight_t.dtype == torch.float32
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_real_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype in (torch.float32, torch.bfloat16)
+        and weight_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_vector(input: torch.Tensor, weight0_t: Optional[torch.Tensor], weight1_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight0_t is not None
+        and weight1_t is not None
+        and input.is_cuda
+        and weight0_t.is_cuda
+        and weight1_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight0_t.dtype == torch.float32
+        and weight1_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight0_t.is_contiguous()
+        and weight1_t.is_contiguous()
+    )
+
+
+def _mul_i(x: torch.Tensor) -> torch.Tensor:
+    """Multiply a complex tensor by +i without promoting dtype."""
+    return torch.complex(-x.imag, x.real)
+
+
+@contextlib.contextmanager
+def _nvtx_range(name: str):
+    enabled = os.environ.get("HOLYSHT_ENABLE_NVTX", "0") == "1"
+    if enabled and torch.cuda.is_available():
+        torch.cuda.nvtx.range_push(name)
+        try:
+            yield
+        finally:
+            torch.cuda.nvtx.range_pop()
+    else:
+        yield
+
+
+def _prepare_irfft_input(x: torch.Tensor, nlon: int, active_mmax: Optional[int] = None) -> torch.Tensor:
+    """Pad/clean an rFFT-format complex tensor before irfft."""
+    active_mmax = x.size(-1) if active_mmax is None else active_mmax
+    full_mmax = nlon // 2 + 1
+
+    if x.size(-1) == full_mmax:
+        out = x.contiguous()
+    else:
+        out_shape = list(x.shape)
+        out_shape[-1] = full_mmax
+        out = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+        out[..., :x.size(-1)] = x
+
+    if _HAS_CUDA_EXT and out.is_cuda and out.dtype == torch.complex64:
+        orig_shape = out.shape
+        flat = out.reshape(-1, orig_shape[-2], orig_shape[-1]).contiguous()
+        _ops.sht_prepare_irfft(flat, active_mmax, nlon)
+        return flat.reshape(orig_shape)
+
+    out[..., 0] = out[..., 0].real.to(torch.complex64)
+    if nlon % 2 == 0:
+        nyquist_idx = nlon // 2
+        if nyquist_idx < out.size(-1):
+            out[..., nyquist_idx] = out[..., nyquist_idx].real.to(torch.complex64)
+    return out
+
+
+class _FusedLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(1), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(0), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreForwardRealFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_forward_real(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_inverse_real(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedVectorLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(0), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(1), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+class _FusedVectorLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(1), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(0), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+# ============================================================================
+# Fused Legendre Transform
+# ============================================================================
+
+def fused_legendre_forward(
+    input: torch.Tensor,       # [B, nlat, mmax] complex64
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32 (original torch-harmonics layout)
+    weight_t: Optional[torch.Tensor] = None,  # [lmax, nlat, mmax] pre-transposed
+) -> torch.Tensor:
+    """Fused forward Legendre transform operating on complex tensors.
+
+    Computes out[b,l,m] = Σ_k weights[m,l,k] · input[b,k,m] for complex input,
+    fusing the real and imaginary multiplications into a single pass.
+
+    Uses the custom CUDA kernels when the extension is available, otherwise
+    falls back to a stacked einsum.
+    """
+    B = input.size(0)
+    nlat = input.size(1)
+    mmax = input.size(2)
+    lmax = weights.size(1)
+
+    if weight_t is None:
+        weight_t = weights.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, weight_t):
+        # Adaptive CUDA kernel: small-grid direct path + large-grid tiled path.
+        return _FusedLegendreForwardFn.apply(input, weight_t)
+    else:
+        # Fallback: stacked einsum (1.9x speedup over reference 2x einsum)
+        x = torch.view_as_real(input)  # [B, nlat, mmax, 2]
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)  # [2B, nlat, mmax]
+        w = weights.to(x_stacked.dtype)
+        out_stacked = torch.einsum("bkm,mlk->blm", x_stacked, w)
+        out_re = out_stacked[:B]
+        out_im = out_stacked[B:]
+        return torch.complex(out_re, out_im)
+
+
+def fused_legendre_inverse(
+    input: torch.Tensor,       # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Fused inverse Legendre transform."""
+    B = input.size(0)
+    lmax = input.size(1)
+    mmax = input.size(2)
+    nlat = pct.size(2)
+
+    if pct_t is None:
+        pct_t = pct.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, pct_t):
+        return _FusedLegendreInverseFn.apply(input, pct_t)
+    else:
+        x = torch.view_as_real(input)
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)
+        p = pct.to(x_stacked.dtype)
+        out_stacked = torch.einsum("blm,mlk->bkm", x_stacked, p)
+        return torch.complex(out_stacked[:B], out_stacked[B:])
+
+
+def fused_legendre_forward_real(
+    input: torch.Tensor,       # [B, nlat, mmax] float32 or bfloat16
+    weight_t: torch.Tensor,    # [lmax, nlat, mmax] float32
+) -> torch.Tensor:
+    """Real-valued forward Legendre transform with float accumulation."""
+    if _can_use_cuda_real_legendre(input, weight_t):
+        return _FusedLegendreForwardRealFn.apply(input, weight_t)
+
+    return torch.einsum("bkm,lkm->blm", input.float(), weight_t)
+
+
+# ============================================================================
+# Fused SHT (complete pipeline)
+# ============================================================================
+
+def fused_sht_forward(
+    x: torch.Tensor,           # [B, nlat, nlon] float32
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32
+    mmax: int,
+    weight_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused forward SHT: rfft → fused Legendre → complex coefficients.
+
+    Replaces RealSHT.forward() with fewer intermediate allocations.
+    """
+    with _nvtx_range("holysht.scalar_forward"):
+        x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+        x_fft = x_fft[..., :mmax]
+        return fused_legendre_forward(x_fft, weights, weight_t)
+
+
+def fused_sht_inverse(
+    coeffs: torch.Tensor,      # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    nlon: int,
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused inverse SHT: fused Legendre → zero-pad → irfft."""
+    with _nvtx_range("holysht.scalar_inverse"):
+        x = fused_legendre_inverse(coeffs, pct, pct_t)
+        x = _prepare_irfft_input(x, nlon, coeffs.size(-1))
+        return torch.fft.irfft(x, n=nlon, dim=-1, norm="forward")
+
+
+# ============================================================================
+# nn.Module wrappers (drop-in replacements for torch-harmonics)
+# ============================================================================
+
+class RealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealSHT``.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealSHT
+        ref = RealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("weights", ref.weights.to(w_dtype))
+        self.register_buffer("weight_t", ref.weights.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._use_bf16:
+            with _nvtx_range("holysht.scalar_forward_bf16"):
+                x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+                x_fft = x_fft[..., :self.mmax]
+                xr = torch.view_as_real(x_fft)
+                if _HAS_CUDA_EXT and x.is_cuda and not x.requires_grad:
+                    xr_bf16 = xr.bfloat16().contiguous()
+                    out_re = fused_legendre_forward_real(xr_bf16[..., 0].contiguous(), self.weight_t)
+                    out_im = fused_legendre_forward_real(xr_bf16[..., 1].contiguous(), self.weight_t)
+                    return torch.complex(out_re, out_im)
+
+                B = x.size(0)
+                xs = torch.cat([xr[..., 0], xr[..., 1]], dim=0).bfloat16()
+                out = torch.einsum("bkm,mlk->blm", xs, self.weights).float()
+                return torch.complex(out[:B], out[B:])
+        return fused_sht_forward(x, self.weights, self.mmax, self.weight_t)
+
+
+class InverseRealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealSHT
+        ref = InverseRealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("pct", ref.pct.float())
+        self.register_buffer("pct_t", ref.pct.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return fused_sht_inverse(x, self.pct, self.nlon, self.pct_t)
+
+
+class RealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealVectorSHT``.
+
+    Reduces eight reference einsums to two composed Legendre passes on the
+    default FP32 CUDA path.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealVectorSHT
+        ref = RealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("w0", ref.weights[0].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w1", ref.weights[1].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w0_t", ref.weights[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("w1_t", ref.weights[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.nlat and x.shape[-1] == self.nlon
+
+        with _nvtx_range("holysht.vector_forward"):
+            x = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+            mmax = self.mmax
+            x = x[..., :mmax].contiguous()
+
+            if (not self._use_bf16) and _can_use_cuda_vector(x, self.w0_t, self.w1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.nlat, mmax).contiguous()
+                out = _FusedVectorLegendreForwardFn.apply(x_flat, self.w0_t, self.w1_t)
+                return out.reshape(B_shape + (2, self.lmax, mmax))
+
+            x = torch.view_as_real(x)  # [..., 2, nlat, mmax, 2]
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.nlat, mmax)
+            x01_flat = x01.reshape(-1, self.nlat, mmax)
+            x10_flat = x10.reshape(-1, self.nlat, mmax)
+            x11_flat = x11.reshape(-1, self.nlat, mmax)
+            B = x00_flat.shape[0]
+
+            if self._use_bf16 and _HAS_CUDA_EXT and x00_flat.is_cuda and not x.requires_grad:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0).bfloat16().contiguous()
+                out_w0 = fused_legendre_forward_real(stacked_w0, self.w0_t)
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0).bfloat16().contiguous()
+                out_w1 = fused_legendre_forward_real(stacked_w1, self.w1_t)
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+            else:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w0 = stacked_w0.bfloat16()
+                out_w0 = torch.einsum("bkm,mlk->blm", stacked_w0, self.w0)
+                if self._use_bf16:
+                    out_w0 = out_w0.float()
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w1 = stacked_w1.bfloat16()
+                out_w1 = torch.einsum("bkm,mlk->blm", stacked_w1, self.w1)
+                if self._use_bf16:
+                    out_w1 = out_w1.float()
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+
+            sph_re = r00 - s11
+            sph_im = r01 + s10
+            tor_re = -s01 - r10
+            tor_im = s00 - r11
+
+            out_shape = list(B_shape) + [2, self.lmax, mmax, 2]
+            xout = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+            xout[..., 0, :, :, 0] = sph_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 0, :, :, 1] = sph_im.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 0] = tor_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 1] = tor_im.reshape(B_shape + (self.lmax, mmax))
+            return torch.view_as_complex(xout)
+
+
+class InverseRealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealVectorSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealVectorSHT
+        ref = InverseRealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("d0", ref.dpct[0].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d1", ref.dpct[1].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d0_t", ref.dpct[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("d1_t", ref.dpct[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.lmax and x.shape[-1] == self.mmax
+
+        with _nvtx_range("holysht.vector_inverse"):
+            x = x.contiguous()
+            if _can_use_cuda_vector(x, self.d0_t, self.d1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.lmax, self.mmax).contiguous()
+                x_out = _FusedVectorLegendreInverseFn.apply(x_flat, self.d0_t, self.d1_t)
+                x_out = x_out.reshape(B_shape + (2, self.nlat, self.mmax))
+                x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+                return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+            x = torch.view_as_real(x)  # [..., 2, lmax, mmax, 2]
+            mmax = self.mmax
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.lmax, mmax)
+            x01_flat = x01.reshape(-1, self.lmax, mmax)
+            x10_flat = x10.reshape(-1, self.lmax, mmax)
+            x11_flat = x11.reshape(-1, self.lmax, mmax)
+            B = x00_flat.shape[0]
+
+            stacked_d0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+            out_d0 = torch.einsum("blm,mlk->bkm", stacked_d0, self.d0)
+            r00, r01, r10, r11 = out_d0[:B], out_d0[B:2 * B], out_d0[2 * B:3 * B], out_d0[3 * B:]
+
+            stacked_d1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+            out_d1 = torch.einsum("blm,mlk->bkm", stacked_d1, self.d1)
+            s11, s10, s01, s00 = out_d1[:B], out_d1[B:2 * B], out_d1[2 * B:3 * B], out_d1[3 * B:]
+
+            srl = r00 - s11
+            sim = r01 + s10
+            trl = -s01 - r10
+            tim = s00 - r11
+
+            out_k = self.nlat
+            srl = srl.reshape(B_shape + (out_k, mmax))
+            sim = sim.reshape(B_shape + (out_k, mmax))
+            trl = trl.reshape(B_shape + (out_k, mmax))
+            tim = tim.reshape(B_shape + (out_k, mmax))
+
+            s = torch.stack((srl, sim), -1)
+            t = torch.stack((trl, tim), -1)
+            xs = torch.stack((s, t), -4)
+            x_out = torch.view_as_complex(xs)
+            x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+            return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+
+legendre_forward = fused_legendre_forward
+legendre_inverse = fused_legendre_inverse
+sht_forward = fused_sht_forward
+sht_inverse = fused_sht_inverse
+
+# Backwards-compatible aliases from the research prototype.
+FusedRealSHT = RealSHT
+FusedInverseRealSHT = InverseRealSHT
+FusedRealVectorSHT = RealVectorSHT
+FusedInverseRealVectorSHT = InverseRealVectorSHT

build/torch210-cxx11-cu130-x86_64-linux/_jit_ops.py

@@ -0,0 +1,72 @@
+"""Local CUDA extension loader for HOLYSHT.
+
+This keeps development off the heavyweight `kernel-builder` path by compiling a
+small torch extension for the current machine only. The compiled artefacts are
+cached under `build/torch_extensions`.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+from torch.utils.cpp_extension import load
+
+
+_ROOT = Path(__file__).resolve().parents[2]
+_BUILD_DIR = _ROOT / "build" / "torch_extensions"
+_EXTENSION_NAME = "holysht_ops"
+
+
+def _default_arch_list() -> str:
+    if not torch.cuda.is_available():
+        return ""
+
+    major, minor = torch.cuda.get_device_capability(0)
+    if (major, minor) >= (12, 1):
+        return "12.0+PTX"
+    return f"{major}.{minor}"
+
+
+def _cuda_flags() -> list[str]:
+    flags = [
+        "-O3",
+        "-lineinfo",
+        "-Xptxas=-warn-spills",
+        "--expt-relaxed-constexpr",
+    ]
+    if os.environ.get("HOLYSHT_USE_FAST_MATH", "1") != "0":
+        flags.append("--use_fast_math")
+    return flags
+
+
+def _load_ops():
+    os.environ.setdefault("MAX_JOBS", "1")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _default_arch_list())
+    _BUILD_DIR.mkdir(parents=True, exist_ok=True)
+
+    load(
+        name=_EXTENSION_NAME,
+        sources=[
+            str(_ROOT / "torch-ext" / "torch_binding.cpp"),
+            str(_ROOT / "cuda" / "fused_legendre.cu"),
+            str(_ROOT / "cuda" / "fused_sht.cu"),
+        ],
+        extra_include_paths=[str(_ROOT / "torch-ext")],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=_cuda_flags(),
+        build_directory=str(_BUILD_DIR),
+        verbose=os.environ.get("HOLYSHT_VERBOSE_BUILD", "0") == "1",
+        with_cuda=True,
+        is_python_module=False,
+    )
+    return getattr(torch.ops, _EXTENSION_NAME)
+
+
+ops = _load_ops()

build/torch210-cxx11-cu130-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _holysht_cuda_53048d0
+ops = torch.ops._holysht_cuda_53048d0
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_holysht_cuda_53048d0::{op_name}"

build/torch210-cxx11-cu130-x86_64-linux/holysht/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch210-cxx11-cu130-x86_64-linux/metadata.json

@@ -0,0 +1,12 @@
+{
+  "version": 1,
+  "license": "mit",
+  "upstream": "https://github.com/NVIDIA/torch-harmonics",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  }
+}

build/torch211-cxx11-cu126-x86_64-linux/__init__.py

@@ -0,0 +1,673 @@
+"""HOLYSHT: Highly Optimised Legendre/Ylm/SHT.
+
+CUDA-accelerated spherical harmonic transforms designed as a practical,
+production-oriented companion to torch-harmonics.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+import contextlib
+import os
+from typing import Optional
+import torch
+import torch.nn as nn
+
+__all__ = [
+    "RealSHT",
+    "InverseRealSHT",
+    "RealVectorSHT",
+    "InverseRealVectorSHT",
+    "legendre_forward",
+    "legendre_inverse",
+    "sht_forward",
+    "sht_inverse",
+]
+
+# Prefer kernel-builder's generated alias module on packaged builds, then fall
+# back to the local single-machine JIT loader for development.
+try:
+    from ._ops import ops as _ops
+    _HAS_CUDA_EXT = True
+except ModuleNotFoundError:
+    try:
+        from ._jit_ops import ops as _ops
+        _HAS_CUDA_EXT = True
+    except ImportError:
+        _HAS_CUDA_EXT = False
+except ImportError:
+    _HAS_CUDA_EXT = False
+
+
+def _can_use_cuda_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight_t.dtype == torch.float32
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_real_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype in (torch.float32, torch.bfloat16)
+        and weight_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_vector(input: torch.Tensor, weight0_t: Optional[torch.Tensor], weight1_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight0_t is not None
+        and weight1_t is not None
+        and input.is_cuda
+        and weight0_t.is_cuda
+        and weight1_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight0_t.dtype == torch.float32
+        and weight1_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight0_t.is_contiguous()
+        and weight1_t.is_contiguous()
+    )
+
+
+def _mul_i(x: torch.Tensor) -> torch.Tensor:
+    """Multiply a complex tensor by +i without promoting dtype."""
+    return torch.complex(-x.imag, x.real)
+
+
+@contextlib.contextmanager
+def _nvtx_range(name: str):
+    enabled = os.environ.get("HOLYSHT_ENABLE_NVTX", "0") == "1"
+    if enabled and torch.cuda.is_available():
+        torch.cuda.nvtx.range_push(name)
+        try:
+            yield
+        finally:
+            torch.cuda.nvtx.range_pop()
+    else:
+        yield
+
+
+def _prepare_irfft_input(x: torch.Tensor, nlon: int, active_mmax: Optional[int] = None) -> torch.Tensor:
+    """Pad/clean an rFFT-format complex tensor before irfft."""
+    active_mmax = x.size(-1) if active_mmax is None else active_mmax
+    full_mmax = nlon // 2 + 1
+
+    if x.size(-1) == full_mmax:
+        out = x.contiguous()
+    else:
+        out_shape = list(x.shape)
+        out_shape[-1] = full_mmax
+        out = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+        out[..., :x.size(-1)] = x
+
+    if _HAS_CUDA_EXT and out.is_cuda and out.dtype == torch.complex64:
+        orig_shape = out.shape
+        flat = out.reshape(-1, orig_shape[-2], orig_shape[-1]).contiguous()
+        _ops.sht_prepare_irfft(flat, active_mmax, nlon)
+        return flat.reshape(orig_shape)
+
+    out[..., 0] = out[..., 0].real.to(torch.complex64)
+    if nlon % 2 == 0:
+        nyquist_idx = nlon // 2
+        if nyquist_idx < out.size(-1):
+            out[..., nyquist_idx] = out[..., nyquist_idx].real.to(torch.complex64)
+    return out
+
+
+class _FusedLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(1), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(0), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreForwardRealFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_forward_real(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_inverse_real(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedVectorLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(0), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(1), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+class _FusedVectorLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(1), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(0), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+# ============================================================================
+# Fused Legendre Transform
+# ============================================================================
+
+def fused_legendre_forward(
+    input: torch.Tensor,       # [B, nlat, mmax] complex64
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32 (original torch-harmonics layout)
+    weight_t: Optional[torch.Tensor] = None,  # [lmax, nlat, mmax] pre-transposed
+) -> torch.Tensor:
+    """Fused forward Legendre transform operating on complex tensors.
+
+    Computes out[b,l,m] = Σ_k weights[m,l,k] · input[b,k,m] for complex input,
+    fusing the real and imaginary multiplications into a single pass.
+
+    Uses the custom CUDA kernels when the extension is available, otherwise
+    falls back to a stacked einsum.
+    """
+    B = input.size(0)
+    nlat = input.size(1)
+    mmax = input.size(2)
+    lmax = weights.size(1)
+
+    if weight_t is None:
+        weight_t = weights.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, weight_t):
+        # Adaptive CUDA kernel: small-grid direct path + large-grid tiled path.
+        return _FusedLegendreForwardFn.apply(input, weight_t)
+    else:
+        # Fallback: stacked einsum (1.9x speedup over reference 2x einsum)
+        x = torch.view_as_real(input)  # [B, nlat, mmax, 2]
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)  # [2B, nlat, mmax]
+        w = weights.to(x_stacked.dtype)
+        out_stacked = torch.einsum("bkm,mlk->blm", x_stacked, w)
+        out_re = out_stacked[:B]
+        out_im = out_stacked[B:]
+        return torch.complex(out_re, out_im)
+
+
+def fused_legendre_inverse(
+    input: torch.Tensor,       # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Fused inverse Legendre transform."""
+    B = input.size(0)
+    lmax = input.size(1)
+    mmax = input.size(2)
+    nlat = pct.size(2)
+
+    if pct_t is None:
+        pct_t = pct.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, pct_t):
+        return _FusedLegendreInverseFn.apply(input, pct_t)
+    else:
+        x = torch.view_as_real(input)
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)
+        p = pct.to(x_stacked.dtype)
+        out_stacked = torch.einsum("blm,mlk->bkm", x_stacked, p)
+        return torch.complex(out_stacked[:B], out_stacked[B:])
+
+
+def fused_legendre_forward_real(
+    input: torch.Tensor,       # [B, nlat, mmax] float32 or bfloat16
+    weight_t: torch.Tensor,    # [lmax, nlat, mmax] float32
+) -> torch.Tensor:
+    """Real-valued forward Legendre transform with float accumulation."""
+    if _can_use_cuda_real_legendre(input, weight_t):
+        return _FusedLegendreForwardRealFn.apply(input, weight_t)
+
+    return torch.einsum("bkm,lkm->blm", input.float(), weight_t)
+
+
+# ============================================================================
+# Fused SHT (complete pipeline)
+# ============================================================================
+
+def fused_sht_forward(
+    x: torch.Tensor,           # [B, nlat, nlon] float32
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32
+    mmax: int,
+    weight_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused forward SHT: rfft → fused Legendre → complex coefficients.
+
+    Replaces RealSHT.forward() with fewer intermediate allocations.
+    """
+    with _nvtx_range("holysht.scalar_forward"):
+        x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+        x_fft = x_fft[..., :mmax]
+        return fused_legendre_forward(x_fft, weights, weight_t)
+
+
+def fused_sht_inverse(
+    coeffs: torch.Tensor,      # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    nlon: int,
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused inverse SHT: fused Legendre → zero-pad → irfft."""
+    with _nvtx_range("holysht.scalar_inverse"):
+        x = fused_legendre_inverse(coeffs, pct, pct_t)
+        x = _prepare_irfft_input(x, nlon, coeffs.size(-1))
+        return torch.fft.irfft(x, n=nlon, dim=-1, norm="forward")
+
+
+# ============================================================================
+# nn.Module wrappers (drop-in replacements for torch-harmonics)
+# ============================================================================
+
+class RealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealSHT``.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealSHT
+        ref = RealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("weights", ref.weights.to(w_dtype))
+        self.register_buffer("weight_t", ref.weights.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._use_bf16:
+            with _nvtx_range("holysht.scalar_forward_bf16"):
+                x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+                x_fft = x_fft[..., :self.mmax]
+                xr = torch.view_as_real(x_fft)
+                if _HAS_CUDA_EXT and x.is_cuda and not x.requires_grad:
+                    xr_bf16 = xr.bfloat16().contiguous()
+                    out_re = fused_legendre_forward_real(xr_bf16[..., 0].contiguous(), self.weight_t)
+                    out_im = fused_legendre_forward_real(xr_bf16[..., 1].contiguous(), self.weight_t)
+                    return torch.complex(out_re, out_im)
+
+                B = x.size(0)
+                xs = torch.cat([xr[..., 0], xr[..., 1]], dim=0).bfloat16()
+                out = torch.einsum("bkm,mlk->blm", xs, self.weights).float()
+                return torch.complex(out[:B], out[B:])
+        return fused_sht_forward(x, self.weights, self.mmax, self.weight_t)
+
+
+class InverseRealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealSHT
+        ref = InverseRealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("pct", ref.pct.float())
+        self.register_buffer("pct_t", ref.pct.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return fused_sht_inverse(x, self.pct, self.nlon, self.pct_t)
+
+
+class RealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealVectorSHT``.
+
+    Reduces eight reference einsums to two composed Legendre passes on the
+    default FP32 CUDA path.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealVectorSHT
+        ref = RealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("w0", ref.weights[0].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w1", ref.weights[1].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w0_t", ref.weights[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("w1_t", ref.weights[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.nlat and x.shape[-1] == self.nlon
+
+        with _nvtx_range("holysht.vector_forward"):
+            x = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+            mmax = self.mmax
+            x = x[..., :mmax].contiguous()
+
+            if (not self._use_bf16) and _can_use_cuda_vector(x, self.w0_t, self.w1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.nlat, mmax).contiguous()
+                out = _FusedVectorLegendreForwardFn.apply(x_flat, self.w0_t, self.w1_t)
+                return out.reshape(B_shape + (2, self.lmax, mmax))
+
+            x = torch.view_as_real(x)  # [..., 2, nlat, mmax, 2]
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.nlat, mmax)
+            x01_flat = x01.reshape(-1, self.nlat, mmax)
+            x10_flat = x10.reshape(-1, self.nlat, mmax)
+            x11_flat = x11.reshape(-1, self.nlat, mmax)
+            B = x00_flat.shape[0]
+
+            if self._use_bf16 and _HAS_CUDA_EXT and x00_flat.is_cuda and not x.requires_grad:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0).bfloat16().contiguous()
+                out_w0 = fused_legendre_forward_real(stacked_w0, self.w0_t)
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0).bfloat16().contiguous()
+                out_w1 = fused_legendre_forward_real(stacked_w1, self.w1_t)
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+            else:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w0 = stacked_w0.bfloat16()
+                out_w0 = torch.einsum("bkm,mlk->blm", stacked_w0, self.w0)
+                if self._use_bf16:
+                    out_w0 = out_w0.float()
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w1 = stacked_w1.bfloat16()
+                out_w1 = torch.einsum("bkm,mlk->blm", stacked_w1, self.w1)
+                if self._use_bf16:
+                    out_w1 = out_w1.float()
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+
+            sph_re = r00 - s11
+            sph_im = r01 + s10
+            tor_re = -s01 - r10
+            tor_im = s00 - r11
+
+            out_shape = list(B_shape) + [2, self.lmax, mmax, 2]
+            xout = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+            xout[..., 0, :, :, 0] = sph_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 0, :, :, 1] = sph_im.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 0] = tor_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 1] = tor_im.reshape(B_shape + (self.lmax, mmax))
+            return torch.view_as_complex(xout)
+
+
+class InverseRealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealVectorSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealVectorSHT
+        ref = InverseRealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("d0", ref.dpct[0].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d1", ref.dpct[1].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d0_t", ref.dpct[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("d1_t", ref.dpct[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.lmax and x.shape[-1] == self.mmax
+
+        with _nvtx_range("holysht.vector_inverse"):
+            x = x.contiguous()
+            if _can_use_cuda_vector(x, self.d0_t, self.d1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.lmax, self.mmax).contiguous()
+                x_out = _FusedVectorLegendreInverseFn.apply(x_flat, self.d0_t, self.d1_t)
+                x_out = x_out.reshape(B_shape + (2, self.nlat, self.mmax))
+                x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+                return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+            x = torch.view_as_real(x)  # [..., 2, lmax, mmax, 2]
+            mmax = self.mmax
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.lmax, mmax)
+            x01_flat = x01.reshape(-1, self.lmax, mmax)
+            x10_flat = x10.reshape(-1, self.lmax, mmax)
+            x11_flat = x11.reshape(-1, self.lmax, mmax)
+            B = x00_flat.shape[0]
+
+            stacked_d0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+            out_d0 = torch.einsum("blm,mlk->bkm", stacked_d0, self.d0)
+            r00, r01, r10, r11 = out_d0[:B], out_d0[B:2 * B], out_d0[2 * B:3 * B], out_d0[3 * B:]
+
+            stacked_d1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+            out_d1 = torch.einsum("blm,mlk->bkm", stacked_d1, self.d1)
+            s11, s10, s01, s00 = out_d1[:B], out_d1[B:2 * B], out_d1[2 * B:3 * B], out_d1[3 * B:]
+
+            srl = r00 - s11
+            sim = r01 + s10
+            trl = -s01 - r10
+            tim = s00 - r11
+
+            out_k = self.nlat
+            srl = srl.reshape(B_shape + (out_k, mmax))
+            sim = sim.reshape(B_shape + (out_k, mmax))
+            trl = trl.reshape(B_shape + (out_k, mmax))
+            tim = tim.reshape(B_shape + (out_k, mmax))
+
+            s = torch.stack((srl, sim), -1)
+            t = torch.stack((trl, tim), -1)
+            xs = torch.stack((s, t), -4)
+            x_out = torch.view_as_complex(xs)
+            x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+            return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+
+legendre_forward = fused_legendre_forward
+legendre_inverse = fused_legendre_inverse
+sht_forward = fused_sht_forward
+sht_inverse = fused_sht_inverse
+
+# Backwards-compatible aliases from the research prototype.
+FusedRealSHT = RealSHT
+FusedInverseRealSHT = InverseRealSHT
+FusedRealVectorSHT = RealVectorSHT
+FusedInverseRealVectorSHT = InverseRealVectorSHT

build/torch211-cxx11-cu126-x86_64-linux/_jit_ops.py

@@ -0,0 +1,72 @@
+"""Local CUDA extension loader for HOLYSHT.
+
+This keeps development off the heavyweight `kernel-builder` path by compiling a
+small torch extension for the current machine only. The compiled artefacts are
+cached under `build/torch_extensions`.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+from torch.utils.cpp_extension import load
+
+
+_ROOT = Path(__file__).resolve().parents[2]
+_BUILD_DIR = _ROOT / "build" / "torch_extensions"
+_EXTENSION_NAME = "holysht_ops"
+
+
+def _default_arch_list() -> str:
+    if not torch.cuda.is_available():
+        return ""
+
+    major, minor = torch.cuda.get_device_capability(0)
+    if (major, minor) >= (12, 1):
+        return "12.0+PTX"
+    return f"{major}.{minor}"
+
+
+def _cuda_flags() -> list[str]:
+    flags = [
+        "-O3",
+        "-lineinfo",
+        "-Xptxas=-warn-spills",
+        "--expt-relaxed-constexpr",
+    ]
+    if os.environ.get("HOLYSHT_USE_FAST_MATH", "1") != "0":
+        flags.append("--use_fast_math")
+    return flags
+
+
+def _load_ops():
+    os.environ.setdefault("MAX_JOBS", "1")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _default_arch_list())
+    _BUILD_DIR.mkdir(parents=True, exist_ok=True)
+
+    load(
+        name=_EXTENSION_NAME,
+        sources=[
+            str(_ROOT / "torch-ext" / "torch_binding.cpp"),
+            str(_ROOT / "cuda" / "fused_legendre.cu"),
+            str(_ROOT / "cuda" / "fused_sht.cu"),
+        ],
+        extra_include_paths=[str(_ROOT / "torch-ext")],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=_cuda_flags(),
+        build_directory=str(_BUILD_DIR),
+        verbose=os.environ.get("HOLYSHT_VERBOSE_BUILD", "0") == "1",
+        with_cuda=True,
+        is_python_module=False,
+    )
+    return getattr(torch.ops, _EXTENSION_NAME)
+
+
+ops = _load_ops()

build/torch211-cxx11-cu126-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _holysht_cuda_53048d0
+ops = torch.ops._holysht_cuda_53048d0
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_holysht_cuda_53048d0::{op_name}"

build/torch211-cxx11-cu126-x86_64-linux/holysht/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch211-cxx11-cu126-x86_64-linux/metadata.json

@@ -0,0 +1,9 @@
+{
+  "version": 1,
+  "license": "mit",
+  "upstream": "https://github.com/NVIDIA/torch-harmonics",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda"
+  }
+}

build/torch211-cxx11-cu128-x86_64-linux/__init__.py

@@ -0,0 +1,673 @@
+"""HOLYSHT: Highly Optimised Legendre/Ylm/SHT.
+
+CUDA-accelerated spherical harmonic transforms designed as a practical,
+production-oriented companion to torch-harmonics.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+import contextlib
+import os
+from typing import Optional
+import torch
+import torch.nn as nn
+
+__all__ = [
+    "RealSHT",
+    "InverseRealSHT",
+    "RealVectorSHT",
+    "InverseRealVectorSHT",
+    "legendre_forward",
+    "legendre_inverse",
+    "sht_forward",
+    "sht_inverse",
+]
+
+# Prefer kernel-builder's generated alias module on packaged builds, then fall
+# back to the local single-machine JIT loader for development.
+try:
+    from ._ops import ops as _ops
+    _HAS_CUDA_EXT = True
+except ModuleNotFoundError:
+    try:
+        from ._jit_ops import ops as _ops
+        _HAS_CUDA_EXT = True
+    except ImportError:
+        _HAS_CUDA_EXT = False
+except ImportError:
+    _HAS_CUDA_EXT = False
+
+
+def _can_use_cuda_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight_t.dtype == torch.float32
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_real_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype in (torch.float32, torch.bfloat16)
+        and weight_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_vector(input: torch.Tensor, weight0_t: Optional[torch.Tensor], weight1_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight0_t is not None
+        and weight1_t is not None
+        and input.is_cuda
+        and weight0_t.is_cuda
+        and weight1_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight0_t.dtype == torch.float32
+        and weight1_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight0_t.is_contiguous()
+        and weight1_t.is_contiguous()
+    )
+
+
+def _mul_i(x: torch.Tensor) -> torch.Tensor:
+    """Multiply a complex tensor by +i without promoting dtype."""
+    return torch.complex(-x.imag, x.real)
+
+
+@contextlib.contextmanager
+def _nvtx_range(name: str):
+    enabled = os.environ.get("HOLYSHT_ENABLE_NVTX", "0") == "1"
+    if enabled and torch.cuda.is_available():
+        torch.cuda.nvtx.range_push(name)
+        try:
+            yield
+        finally:
+            torch.cuda.nvtx.range_pop()
+    else:
+        yield
+
+
+def _prepare_irfft_input(x: torch.Tensor, nlon: int, active_mmax: Optional[int] = None) -> torch.Tensor:
+    """Pad/clean an rFFT-format complex tensor before irfft."""
+    active_mmax = x.size(-1) if active_mmax is None else active_mmax
+    full_mmax = nlon // 2 + 1
+
+    if x.size(-1) == full_mmax:
+        out = x.contiguous()
+    else:
+        out_shape = list(x.shape)
+        out_shape[-1] = full_mmax
+        out = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+        out[..., :x.size(-1)] = x
+
+    if _HAS_CUDA_EXT and out.is_cuda and out.dtype == torch.complex64:
+        orig_shape = out.shape
+        flat = out.reshape(-1, orig_shape[-2], orig_shape[-1]).contiguous()
+        _ops.sht_prepare_irfft(flat, active_mmax, nlon)
+        return flat.reshape(orig_shape)
+
+    out[..., 0] = out[..., 0].real.to(torch.complex64)
+    if nlon % 2 == 0:
+        nyquist_idx = nlon // 2
+        if nyquist_idx < out.size(-1):
+            out[..., nyquist_idx] = out[..., nyquist_idx].real.to(torch.complex64)
+    return out
+
+
+class _FusedLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(1), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(0), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreForwardRealFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_forward_real(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_inverse_real(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedVectorLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(0), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(1), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+class _FusedVectorLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(1), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(0), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+# ============================================================================
+# Fused Legendre Transform
+# ============================================================================
+
+def fused_legendre_forward(
+    input: torch.Tensor,       # [B, nlat, mmax] complex64
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32 (original torch-harmonics layout)
+    weight_t: Optional[torch.Tensor] = None,  # [lmax, nlat, mmax] pre-transposed
+) -> torch.Tensor:
+    """Fused forward Legendre transform operating on complex tensors.
+
+    Computes out[b,l,m] = Σ_k weights[m,l,k] · input[b,k,m] for complex input,
+    fusing the real and imaginary multiplications into a single pass.
+
+    Uses the custom CUDA kernels when the extension is available, otherwise
+    falls back to a stacked einsum.
+    """
+    B = input.size(0)
+    nlat = input.size(1)
+    mmax = input.size(2)
+    lmax = weights.size(1)
+
+    if weight_t is None:
+        weight_t = weights.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, weight_t):
+        # Adaptive CUDA kernel: small-grid direct path + large-grid tiled path.
+        return _FusedLegendreForwardFn.apply(input, weight_t)
+    else:
+        # Fallback: stacked einsum (1.9x speedup over reference 2x einsum)
+        x = torch.view_as_real(input)  # [B, nlat, mmax, 2]
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)  # [2B, nlat, mmax]
+        w = weights.to(x_stacked.dtype)
+        out_stacked = torch.einsum("bkm,mlk->blm", x_stacked, w)
+        out_re = out_stacked[:B]
+        out_im = out_stacked[B:]
+        return torch.complex(out_re, out_im)
+
+
+def fused_legendre_inverse(
+    input: torch.Tensor,       # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Fused inverse Legendre transform."""
+    B = input.size(0)
+    lmax = input.size(1)
+    mmax = input.size(2)
+    nlat = pct.size(2)
+
+    if pct_t is None:
+        pct_t = pct.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, pct_t):
+        return _FusedLegendreInverseFn.apply(input, pct_t)
+    else:
+        x = torch.view_as_real(input)
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)
+        p = pct.to(x_stacked.dtype)
+        out_stacked = torch.einsum("blm,mlk->bkm", x_stacked, p)
+        return torch.complex(out_stacked[:B], out_stacked[B:])
+
+
+def fused_legendre_forward_real(
+    input: torch.Tensor,       # [B, nlat, mmax] float32 or bfloat16
+    weight_t: torch.Tensor,    # [lmax, nlat, mmax] float32
+) -> torch.Tensor:
+    """Real-valued forward Legendre transform with float accumulation."""
+    if _can_use_cuda_real_legendre(input, weight_t):
+        return _FusedLegendreForwardRealFn.apply(input, weight_t)
+
+    return torch.einsum("bkm,lkm->blm", input.float(), weight_t)
+
+
+# ============================================================================
+# Fused SHT (complete pipeline)
+# ============================================================================
+
+def fused_sht_forward(
+    x: torch.Tensor,           # [B, nlat, nlon] float32
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32
+    mmax: int,
+    weight_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused forward SHT: rfft → fused Legendre → complex coefficients.
+
+    Replaces RealSHT.forward() with fewer intermediate allocations.
+    """
+    with _nvtx_range("holysht.scalar_forward"):
+        x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+        x_fft = x_fft[..., :mmax]
+        return fused_legendre_forward(x_fft, weights, weight_t)
+
+
+def fused_sht_inverse(
+    coeffs: torch.Tensor,      # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    nlon: int,
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused inverse SHT: fused Legendre → zero-pad → irfft."""
+    with _nvtx_range("holysht.scalar_inverse"):
+        x = fused_legendre_inverse(coeffs, pct, pct_t)
+        x = _prepare_irfft_input(x, nlon, coeffs.size(-1))
+        return torch.fft.irfft(x, n=nlon, dim=-1, norm="forward")
+
+
+# ============================================================================
+# nn.Module wrappers (drop-in replacements for torch-harmonics)
+# ============================================================================
+
+class RealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealSHT``.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealSHT
+        ref = RealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("weights", ref.weights.to(w_dtype))
+        self.register_buffer("weight_t", ref.weights.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._use_bf16:
+            with _nvtx_range("holysht.scalar_forward_bf16"):
+                x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+                x_fft = x_fft[..., :self.mmax]
+                xr = torch.view_as_real(x_fft)
+                if _HAS_CUDA_EXT and x.is_cuda and not x.requires_grad:
+                    xr_bf16 = xr.bfloat16().contiguous()
+                    out_re = fused_legendre_forward_real(xr_bf16[..., 0].contiguous(), self.weight_t)
+                    out_im = fused_legendre_forward_real(xr_bf16[..., 1].contiguous(), self.weight_t)
+                    return torch.complex(out_re, out_im)
+
+                B = x.size(0)
+                xs = torch.cat([xr[..., 0], xr[..., 1]], dim=0).bfloat16()
+                out = torch.einsum("bkm,mlk->blm", xs, self.weights).float()
+                return torch.complex(out[:B], out[B:])
+        return fused_sht_forward(x, self.weights, self.mmax, self.weight_t)
+
+
+class InverseRealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealSHT
+        ref = InverseRealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("pct", ref.pct.float())
+        self.register_buffer("pct_t", ref.pct.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return fused_sht_inverse(x, self.pct, self.nlon, self.pct_t)
+
+
+class RealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealVectorSHT``.
+
+    Reduces eight reference einsums to two composed Legendre passes on the
+    default FP32 CUDA path.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealVectorSHT
+        ref = RealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("w0", ref.weights[0].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w1", ref.weights[1].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w0_t", ref.weights[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("w1_t", ref.weights[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.nlat and x.shape[-1] == self.nlon
+
+        with _nvtx_range("holysht.vector_forward"):
+            x = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+            mmax = self.mmax
+            x = x[..., :mmax].contiguous()
+
+            if (not self._use_bf16) and _can_use_cuda_vector(x, self.w0_t, self.w1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.nlat, mmax).contiguous()
+                out = _FusedVectorLegendreForwardFn.apply(x_flat, self.w0_t, self.w1_t)
+                return out.reshape(B_shape + (2, self.lmax, mmax))
+
+            x = torch.view_as_real(x)  # [..., 2, nlat, mmax, 2]
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.nlat, mmax)
+            x01_flat = x01.reshape(-1, self.nlat, mmax)
+            x10_flat = x10.reshape(-1, self.nlat, mmax)
+            x11_flat = x11.reshape(-1, self.nlat, mmax)
+            B = x00_flat.shape[0]
+
+            if self._use_bf16 and _HAS_CUDA_EXT and x00_flat.is_cuda and not x.requires_grad:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0).bfloat16().contiguous()
+                out_w0 = fused_legendre_forward_real(stacked_w0, self.w0_t)
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0).bfloat16().contiguous()
+                out_w1 = fused_legendre_forward_real(stacked_w1, self.w1_t)
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+            else:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w0 = stacked_w0.bfloat16()
+                out_w0 = torch.einsum("bkm,mlk->blm", stacked_w0, self.w0)
+                if self._use_bf16:
+                    out_w0 = out_w0.float()
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w1 = stacked_w1.bfloat16()
+                out_w1 = torch.einsum("bkm,mlk->blm", stacked_w1, self.w1)
+                if self._use_bf16:
+                    out_w1 = out_w1.float()
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+
+            sph_re = r00 - s11
+            sph_im = r01 + s10
+            tor_re = -s01 - r10
+            tor_im = s00 - r11
+
+            out_shape = list(B_shape) + [2, self.lmax, mmax, 2]
+            xout = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+            xout[..., 0, :, :, 0] = sph_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 0, :, :, 1] = sph_im.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 0] = tor_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 1] = tor_im.reshape(B_shape + (self.lmax, mmax))
+            return torch.view_as_complex(xout)
+
+
+class InverseRealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealVectorSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealVectorSHT
+        ref = InverseRealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("d0", ref.dpct[0].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d1", ref.dpct[1].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d0_t", ref.dpct[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("d1_t", ref.dpct[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.lmax and x.shape[-1] == self.mmax
+
+        with _nvtx_range("holysht.vector_inverse"):
+            x = x.contiguous()
+            if _can_use_cuda_vector(x, self.d0_t, self.d1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.lmax, self.mmax).contiguous()
+                x_out = _FusedVectorLegendreInverseFn.apply(x_flat, self.d0_t, self.d1_t)
+                x_out = x_out.reshape(B_shape + (2, self.nlat, self.mmax))
+                x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+                return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+            x = torch.view_as_real(x)  # [..., 2, lmax, mmax, 2]
+            mmax = self.mmax
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.lmax, mmax)
+            x01_flat = x01.reshape(-1, self.lmax, mmax)
+            x10_flat = x10.reshape(-1, self.lmax, mmax)
+            x11_flat = x11.reshape(-1, self.lmax, mmax)
+            B = x00_flat.shape[0]
+
+            stacked_d0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+            out_d0 = torch.einsum("blm,mlk->bkm", stacked_d0, self.d0)
+            r00, r01, r10, r11 = out_d0[:B], out_d0[B:2 * B], out_d0[2 * B:3 * B], out_d0[3 * B:]
+
+            stacked_d1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+            out_d1 = torch.einsum("blm,mlk->bkm", stacked_d1, self.d1)
+            s11, s10, s01, s00 = out_d1[:B], out_d1[B:2 * B], out_d1[2 * B:3 * B], out_d1[3 * B:]
+
+            srl = r00 - s11
+            sim = r01 + s10
+            trl = -s01 - r10
+            tim = s00 - r11
+
+            out_k = self.nlat
+            srl = srl.reshape(B_shape + (out_k, mmax))
+            sim = sim.reshape(B_shape + (out_k, mmax))
+            trl = trl.reshape(B_shape + (out_k, mmax))
+            tim = tim.reshape(B_shape + (out_k, mmax))
+
+            s = torch.stack((srl, sim), -1)
+            t = torch.stack((trl, tim), -1)
+            xs = torch.stack((s, t), -4)
+            x_out = torch.view_as_complex(xs)
+            x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+            return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+
+legendre_forward = fused_legendre_forward
+legendre_inverse = fused_legendre_inverse
+sht_forward = fused_sht_forward
+sht_inverse = fused_sht_inverse
+
+# Backwards-compatible aliases from the research prototype.
+FusedRealSHT = RealSHT
+FusedInverseRealSHT = InverseRealSHT
+FusedRealVectorSHT = RealVectorSHT
+FusedInverseRealVectorSHT = InverseRealVectorSHT

build/torch211-cxx11-cu128-x86_64-linux/_jit_ops.py

@@ -0,0 +1,72 @@
+"""Local CUDA extension loader for HOLYSHT.
+
+This keeps development off the heavyweight `kernel-builder` path by compiling a
+small torch extension for the current machine only. The compiled artefacts are
+cached under `build/torch_extensions`.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+from torch.utils.cpp_extension import load
+
+
+_ROOT = Path(__file__).resolve().parents[2]
+_BUILD_DIR = _ROOT / "build" / "torch_extensions"
+_EXTENSION_NAME = "holysht_ops"
+
+
+def _default_arch_list() -> str:
+    if not torch.cuda.is_available():
+        return ""
+
+    major, minor = torch.cuda.get_device_capability(0)
+    if (major, minor) >= (12, 1):
+        return "12.0+PTX"
+    return f"{major}.{minor}"
+
+
+def _cuda_flags() -> list[str]:
+    flags = [
+        "-O3",
+        "-lineinfo",
+        "-Xptxas=-warn-spills",
+        "--expt-relaxed-constexpr",
+    ]
+    if os.environ.get("HOLYSHT_USE_FAST_MATH", "1") != "0":
+        flags.append("--use_fast_math")
+    return flags
+
+
+def _load_ops():
+    os.environ.setdefault("MAX_JOBS", "1")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _default_arch_list())
+    _BUILD_DIR.mkdir(parents=True, exist_ok=True)
+
+    load(
+        name=_EXTENSION_NAME,
+        sources=[
+            str(_ROOT / "torch-ext" / "torch_binding.cpp"),
+            str(_ROOT / "cuda" / "fused_legendre.cu"),
+            str(_ROOT / "cuda" / "fused_sht.cu"),
+        ],
+        extra_include_paths=[str(_ROOT / "torch-ext")],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=_cuda_flags(),
+        build_directory=str(_BUILD_DIR),
+        verbose=os.environ.get("HOLYSHT_VERBOSE_BUILD", "0") == "1",
+        with_cuda=True,
+        is_python_module=False,
+    )
+    return getattr(torch.ops, _EXTENSION_NAME)
+
+
+ops = _load_ops()

build/torch211-cxx11-cu128-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _holysht_cuda_53048d0
+ops = torch.ops._holysht_cuda_53048d0
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_holysht_cuda_53048d0::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/holysht/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch211-cxx11-cu128-x86_64-linux/metadata.json

@@ -0,0 +1,12 @@
+{
+  "version": 1,
+  "license": "mit",
+  "upstream": "https://github.com/NVIDIA/torch-harmonics",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  }
+}

build/torch211-cxx11-cu130-x86_64-linux/__init__.py

@@ -0,0 +1,673 @@
+"""HOLYSHT: Highly Optimised Legendre/Ylm/SHT.
+
+CUDA-accelerated spherical harmonic transforms designed as a practical,
+production-oriented companion to torch-harmonics.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+import contextlib
+import os
+from typing import Optional
+import torch
+import torch.nn as nn
+
+__all__ = [
+    "RealSHT",
+    "InverseRealSHT",
+    "RealVectorSHT",
+    "InverseRealVectorSHT",
+    "legendre_forward",
+    "legendre_inverse",
+    "sht_forward",
+    "sht_inverse",
+]
+
+# Prefer kernel-builder's generated alias module on packaged builds, then fall
+# back to the local single-machine JIT loader for development.
+try:
+    from ._ops import ops as _ops
+    _HAS_CUDA_EXT = True
+except ModuleNotFoundError:
+    try:
+        from ._jit_ops import ops as _ops
+        _HAS_CUDA_EXT = True
+    except ImportError:
+        _HAS_CUDA_EXT = False
+except ImportError:
+    _HAS_CUDA_EXT = False
+
+
+def _can_use_cuda_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight_t.dtype == torch.float32
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_real_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype in (torch.float32, torch.bfloat16)
+        and weight_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_vector(input: torch.Tensor, weight0_t: Optional[torch.Tensor], weight1_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight0_t is not None
+        and weight1_t is not None
+        and input.is_cuda
+        and weight0_t.is_cuda
+        and weight1_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight0_t.dtype == torch.float32
+        and weight1_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight0_t.is_contiguous()
+        and weight1_t.is_contiguous()
+    )
+
+
+def _mul_i(x: torch.Tensor) -> torch.Tensor:
+    """Multiply a complex tensor by +i without promoting dtype."""
+    return torch.complex(-x.imag, x.real)
+
+
+@contextlib.contextmanager
+def _nvtx_range(name: str):
+    enabled = os.environ.get("HOLYSHT_ENABLE_NVTX", "0") == "1"
+    if enabled and torch.cuda.is_available():
+        torch.cuda.nvtx.range_push(name)
+        try:
+            yield
+        finally:
+            torch.cuda.nvtx.range_pop()
+    else:
+        yield
+
+
+def _prepare_irfft_input(x: torch.Tensor, nlon: int, active_mmax: Optional[int] = None) -> torch.Tensor:
+    """Pad/clean an rFFT-format complex tensor before irfft."""
+    active_mmax = x.size(-1) if active_mmax is None else active_mmax
+    full_mmax = nlon // 2 + 1
+
+    if x.size(-1) == full_mmax:
+        out = x.contiguous()
+    else:
+        out_shape = list(x.shape)
+        out_shape[-1] = full_mmax
+        out = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+        out[..., :x.size(-1)] = x
+
+    if _HAS_CUDA_EXT and out.is_cuda and out.dtype == torch.complex64:
+        orig_shape = out.shape
+        flat = out.reshape(-1, orig_shape[-2], orig_shape[-1]).contiguous()
+        _ops.sht_prepare_irfft(flat, active_mmax, nlon)
+        return flat.reshape(orig_shape)
+
+    out[..., 0] = out[..., 0].real.to(torch.complex64)
+    if nlon % 2 == 0:
+        nyquist_idx = nlon // 2
+        if nyquist_idx < out.size(-1):
+            out[..., nyquist_idx] = out[..., nyquist_idx].real.to(torch.complex64)
+    return out
+
+
+class _FusedLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(1), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(0), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreForwardRealFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_forward_real(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_inverse_real(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedVectorLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(0), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(1), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+class _FusedVectorLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(1), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(0), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+# ============================================================================
+# Fused Legendre Transform
+# ============================================================================
+
+def fused_legendre_forward(
+    input: torch.Tensor,       # [B, nlat, mmax] complex64
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32 (original torch-harmonics layout)
+    weight_t: Optional[torch.Tensor] = None,  # [lmax, nlat, mmax] pre-transposed
+) -> torch.Tensor:
+    """Fused forward Legendre transform operating on complex tensors.
+
+    Computes out[b,l,m] = Σ_k weights[m,l,k] · input[b,k,m] for complex input,
+    fusing the real and imaginary multiplications into a single pass.
+
+    Uses the custom CUDA kernels when the extension is available, otherwise
+    falls back to a stacked einsum.
+    """
+    B = input.size(0)
+    nlat = input.size(1)
+    mmax = input.size(2)
+    lmax = weights.size(1)
+
+    if weight_t is None:
+        weight_t = weights.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, weight_t):
+        # Adaptive CUDA kernel: small-grid direct path + large-grid tiled path.
+        return _FusedLegendreForwardFn.apply(input, weight_t)
+    else:
+        # Fallback: stacked einsum (1.9x speedup over reference 2x einsum)
+        x = torch.view_as_real(input)  # [B, nlat, mmax, 2]
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)  # [2B, nlat, mmax]
+        w = weights.to(x_stacked.dtype)
+        out_stacked = torch.einsum("bkm,mlk->blm", x_stacked, w)
+        out_re = out_stacked[:B]
+        out_im = out_stacked[B:]
+        return torch.complex(out_re, out_im)
+
+
+def fused_legendre_inverse(
+    input: torch.Tensor,       # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Fused inverse Legendre transform."""
+    B = input.size(0)
+    lmax = input.size(1)
+    mmax = input.size(2)
+    nlat = pct.size(2)
+
+    if pct_t is None:
+        pct_t = pct.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, pct_t):
+        return _FusedLegendreInverseFn.apply(input, pct_t)
+    else:
+        x = torch.view_as_real(input)
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)
+        p = pct.to(x_stacked.dtype)
+        out_stacked = torch.einsum("blm,mlk->bkm", x_stacked, p)
+        return torch.complex(out_stacked[:B], out_stacked[B:])
+
+
+def fused_legendre_forward_real(
+    input: torch.Tensor,       # [B, nlat, mmax] float32 or bfloat16
+    weight_t: torch.Tensor,    # [lmax, nlat, mmax] float32
+) -> torch.Tensor:
+    """Real-valued forward Legendre transform with float accumulation."""
+    if _can_use_cuda_real_legendre(input, weight_t):
+        return _FusedLegendreForwardRealFn.apply(input, weight_t)
+
+    return torch.einsum("bkm,lkm->blm", input.float(), weight_t)
+
+
+# ============================================================================
+# Fused SHT (complete pipeline)
+# ============================================================================
+
+def fused_sht_forward(
+    x: torch.Tensor,           # [B, nlat, nlon] float32
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32
+    mmax: int,
+    weight_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused forward SHT: rfft → fused Legendre → complex coefficients.
+
+    Replaces RealSHT.forward() with fewer intermediate allocations.
+    """
+    with _nvtx_range("holysht.scalar_forward"):
+        x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+        x_fft = x_fft[..., :mmax]
+        return fused_legendre_forward(x_fft, weights, weight_t)
+
+
+def fused_sht_inverse(
+    coeffs: torch.Tensor,      # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    nlon: int,
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused inverse SHT: fused Legendre → zero-pad → irfft."""
+    with _nvtx_range("holysht.scalar_inverse"):
+        x = fused_legendre_inverse(coeffs, pct, pct_t)
+        x = _prepare_irfft_input(x, nlon, coeffs.size(-1))
+        return torch.fft.irfft(x, n=nlon, dim=-1, norm="forward")
+
+
+# ============================================================================
+# nn.Module wrappers (drop-in replacements for torch-harmonics)
+# ============================================================================
+
+class RealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealSHT``.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealSHT
+        ref = RealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("weights", ref.weights.to(w_dtype))
+        self.register_buffer("weight_t", ref.weights.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._use_bf16:
+            with _nvtx_range("holysht.scalar_forward_bf16"):
+                x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+                x_fft = x_fft[..., :self.mmax]
+                xr = torch.view_as_real(x_fft)
+                if _HAS_CUDA_EXT and x.is_cuda and not x.requires_grad:
+                    xr_bf16 = xr.bfloat16().contiguous()
+                    out_re = fused_legendre_forward_real(xr_bf16[..., 0].contiguous(), self.weight_t)
+                    out_im = fused_legendre_forward_real(xr_bf16[..., 1].contiguous(), self.weight_t)
+                    return torch.complex(out_re, out_im)
+
+                B = x.size(0)
+                xs = torch.cat([xr[..., 0], xr[..., 1]], dim=0).bfloat16()
+                out = torch.einsum("bkm,mlk->blm", xs, self.weights).float()
+                return torch.complex(out[:B], out[B:])
+        return fused_sht_forward(x, self.weights, self.mmax, self.weight_t)
+
+
+class InverseRealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealSHT
+        ref = InverseRealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("pct", ref.pct.float())
+        self.register_buffer("pct_t", ref.pct.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return fused_sht_inverse(x, self.pct, self.nlon, self.pct_t)
+
+
+class RealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealVectorSHT``.
+
+    Reduces eight reference einsums to two composed Legendre passes on the
+    default FP32 CUDA path.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealVectorSHT
+        ref = RealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("w0", ref.weights[0].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w1", ref.weights[1].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w0_t", ref.weights[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("w1_t", ref.weights[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.nlat and x.shape[-1] == self.nlon
+
+        with _nvtx_range("holysht.vector_forward"):
+            x = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+            mmax = self.mmax
+            x = x[..., :mmax].contiguous()
+
+            if (not self._use_bf16) and _can_use_cuda_vector(x, self.w0_t, self.w1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.nlat, mmax).contiguous()
+                out = _FusedVectorLegendreForwardFn.apply(x_flat, self.w0_t, self.w1_t)
+                return out.reshape(B_shape + (2, self.lmax, mmax))
+
+            x = torch.view_as_real(x)  # [..., 2, nlat, mmax, 2]
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.nlat, mmax)
+            x01_flat = x01.reshape(-1, self.nlat, mmax)
+            x10_flat = x10.reshape(-1, self.nlat, mmax)
+            x11_flat = x11.reshape(-1, self.nlat, mmax)
+            B = x00_flat.shape[0]
+
+            if self._use_bf16 and _HAS_CUDA_EXT and x00_flat.is_cuda and not x.requires_grad:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0).bfloat16().contiguous()
+                out_w0 = fused_legendre_forward_real(stacked_w0, self.w0_t)
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0).bfloat16().contiguous()
+                out_w1 = fused_legendre_forward_real(stacked_w1, self.w1_t)
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+            else:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w0 = stacked_w0.bfloat16()
+                out_w0 = torch.einsum("bkm,mlk->blm", stacked_w0, self.w0)
+                if self._use_bf16:
+                    out_w0 = out_w0.float()
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w1 = stacked_w1.bfloat16()
+                out_w1 = torch.einsum("bkm,mlk->blm", stacked_w1, self.w1)
+                if self._use_bf16:
+                    out_w1 = out_w1.float()
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+
+            sph_re = r00 - s11
+            sph_im = r01 + s10
+            tor_re = -s01 - r10
+            tor_im = s00 - r11
+
+            out_shape = list(B_shape) + [2, self.lmax, mmax, 2]
+            xout = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+            xout[..., 0, :, :, 0] = sph_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 0, :, :, 1] = sph_im.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 0] = tor_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 1] = tor_im.reshape(B_shape + (self.lmax, mmax))
+            return torch.view_as_complex(xout)
+
+
+class InverseRealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealVectorSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealVectorSHT
+        ref = InverseRealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("d0", ref.dpct[0].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d1", ref.dpct[1].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d0_t", ref.dpct[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("d1_t", ref.dpct[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.lmax and x.shape[-1] == self.mmax
+
+        with _nvtx_range("holysht.vector_inverse"):
+            x = x.contiguous()
+            if _can_use_cuda_vector(x, self.d0_t, self.d1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.lmax, self.mmax).contiguous()
+                x_out = _FusedVectorLegendreInverseFn.apply(x_flat, self.d0_t, self.d1_t)
+                x_out = x_out.reshape(B_shape + (2, self.nlat, self.mmax))
+                x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+                return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+            x = torch.view_as_real(x)  # [..., 2, lmax, mmax, 2]
+            mmax = self.mmax
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.lmax, mmax)
+            x01_flat = x01.reshape(-1, self.lmax, mmax)
+            x10_flat = x10.reshape(-1, self.lmax, mmax)
+            x11_flat = x11.reshape(-1, self.lmax, mmax)
+            B = x00_flat.shape[0]
+
+            stacked_d0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+            out_d0 = torch.einsum("blm,mlk->bkm", stacked_d0, self.d0)
+            r00, r01, r10, r11 = out_d0[:B], out_d0[B:2 * B], out_d0[2 * B:3 * B], out_d0[3 * B:]
+
+            stacked_d1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+            out_d1 = torch.einsum("blm,mlk->bkm", stacked_d1, self.d1)
+            s11, s10, s01, s00 = out_d1[:B], out_d1[B:2 * B], out_d1[2 * B:3 * B], out_d1[3 * B:]
+
+            srl = r00 - s11
+            sim = r01 + s10
+            trl = -s01 - r10
+            tim = s00 - r11
+
+            out_k = self.nlat
+            srl = srl.reshape(B_shape + (out_k, mmax))
+            sim = sim.reshape(B_shape + (out_k, mmax))
+            trl = trl.reshape(B_shape + (out_k, mmax))
+            tim = tim.reshape(B_shape + (out_k, mmax))
+
+            s = torch.stack((srl, sim), -1)
+            t = torch.stack((trl, tim), -1)
+            xs = torch.stack((s, t), -4)
+            x_out = torch.view_as_complex(xs)
+            x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+            return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+
+legendre_forward = fused_legendre_forward
+legendre_inverse = fused_legendre_inverse
+sht_forward = fused_sht_forward
+sht_inverse = fused_sht_inverse
+
+# Backwards-compatible aliases from the research prototype.
+FusedRealSHT = RealSHT
+FusedInverseRealSHT = InverseRealSHT
+FusedRealVectorSHT = RealVectorSHT
+FusedInverseRealVectorSHT = InverseRealVectorSHT

build/torch211-cxx11-cu130-x86_64-linux/_jit_ops.py

@@ -0,0 +1,72 @@
+"""Local CUDA extension loader for HOLYSHT.
+
+This keeps development off the heavyweight `kernel-builder` path by compiling a
+small torch extension for the current machine only. The compiled artefacts are
+cached under `build/torch_extensions`.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+from torch.utils.cpp_extension import load
+
+
+_ROOT = Path(__file__).resolve().parents[2]
+_BUILD_DIR = _ROOT / "build" / "torch_extensions"
+_EXTENSION_NAME = "holysht_ops"
+
+
+def _default_arch_list() -> str:
+    if not torch.cuda.is_available():
+        return ""
+
+    major, minor = torch.cuda.get_device_capability(0)
+    if (major, minor) >= (12, 1):
+        return "12.0+PTX"
+    return f"{major}.{minor}"
+
+
+def _cuda_flags() -> list[str]:
+    flags = [
+        "-O3",
+        "-lineinfo",
+        "-Xptxas=-warn-spills",
+        "--expt-relaxed-constexpr",
+    ]
+    if os.environ.get("HOLYSHT_USE_FAST_MATH", "1") != "0":
+        flags.append("--use_fast_math")
+    return flags
+
+
+def _load_ops():
+    os.environ.setdefault("MAX_JOBS", "1")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _default_arch_list())
+    _BUILD_DIR.mkdir(parents=True, exist_ok=True)
+
+    load(
+        name=_EXTENSION_NAME,
+        sources=[
+            str(_ROOT / "torch-ext" / "torch_binding.cpp"),
+            str(_ROOT / "cuda" / "fused_legendre.cu"),
+            str(_ROOT / "cuda" / "fused_sht.cu"),
+        ],
+        extra_include_paths=[str(_ROOT / "torch-ext")],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=_cuda_flags(),
+        build_directory=str(_BUILD_DIR),
+        verbose=os.environ.get("HOLYSHT_VERBOSE_BUILD", "0") == "1",
+        with_cuda=True,
+        is_python_module=False,
+    )
+    return getattr(torch.ops, _EXTENSION_NAME)
+
+
+ops = _load_ops()

build/torch211-cxx11-cu130-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _holysht_cuda_53048d0
+ops = torch.ops._holysht_cuda_53048d0
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_holysht_cuda_53048d0::{op_name}"

build/torch211-cxx11-cu130-x86_64-linux/holysht/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch211-cxx11-cu130-x86_64-linux/metadata.json

@@ -0,0 +1,12 @@
+{
+  "version": 1,
+  "license": "mit",
+  "upstream": "https://github.com/NVIDIA/torch-harmonics",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  }
+}

build/torch29-cxx11-cu129-x86_64-linux/__init__.py

@@ -0,0 +1,673 @@
+"""HOLYSHT: Highly Optimised Legendre/Ylm/SHT.
+
+CUDA-accelerated spherical harmonic transforms designed as a practical,
+production-oriented companion to torch-harmonics.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+import contextlib
+import os
+from typing import Optional
+import torch
+import torch.nn as nn
+
+__all__ = [
+    "RealSHT",
+    "InverseRealSHT",
+    "RealVectorSHT",
+    "InverseRealVectorSHT",
+    "legendre_forward",
+    "legendre_inverse",
+    "sht_forward",
+    "sht_inverse",
+]
+
+# Prefer kernel-builder's generated alias module on packaged builds, then fall
+# back to the local single-machine JIT loader for development.
+try:
+    from ._ops import ops as _ops
+    _HAS_CUDA_EXT = True
+except ModuleNotFoundError:
+    try:
+        from ._jit_ops import ops as _ops
+        _HAS_CUDA_EXT = True
+    except ImportError:
+        _HAS_CUDA_EXT = False
+except ImportError:
+    _HAS_CUDA_EXT = False
+
+
+def _can_use_cuda_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight_t.dtype == torch.float32
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_real_legendre(input: torch.Tensor, weight_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight_t is not None
+        and input.is_cuda
+        and weight_t.is_cuda
+        and input.dtype in (torch.float32, torch.bfloat16)
+        and weight_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight_t.is_contiguous()
+    )
+
+
+def _can_use_cuda_vector(input: torch.Tensor, weight0_t: Optional[torch.Tensor], weight1_t: Optional[torch.Tensor]) -> bool:
+    return (
+        _HAS_CUDA_EXT
+        and weight0_t is not None
+        and weight1_t is not None
+        and input.is_cuda
+        and weight0_t.is_cuda
+        and weight1_t.is_cuda
+        and input.dtype == torch.complex64
+        and weight0_t.dtype == torch.float32
+        and weight1_t.dtype == torch.float32
+        and input.is_contiguous()
+        and weight0_t.is_contiguous()
+        and weight1_t.is_contiguous()
+    )
+
+
+def _mul_i(x: torch.Tensor) -> torch.Tensor:
+    """Multiply a complex tensor by +i without promoting dtype."""
+    return torch.complex(-x.imag, x.real)
+
+
+@contextlib.contextmanager
+def _nvtx_range(name: str):
+    enabled = os.environ.get("HOLYSHT_ENABLE_NVTX", "0") == "1"
+    if enabled and torch.cuda.is_available():
+        torch.cuda.nvtx.range_push(name)
+        try:
+            yield
+        finally:
+            torch.cuda.nvtx.range_pop()
+    else:
+        yield
+
+
+def _prepare_irfft_input(x: torch.Tensor, nlon: int, active_mmax: Optional[int] = None) -> torch.Tensor:
+    """Pad/clean an rFFT-format complex tensor before irfft."""
+    active_mmax = x.size(-1) if active_mmax is None else active_mmax
+    full_mmax = nlon // 2 + 1
+
+    if x.size(-1) == full_mmax:
+        out = x.contiguous()
+    else:
+        out_shape = list(x.shape)
+        out_shape[-1] = full_mmax
+        out = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+        out[..., :x.size(-1)] = x
+
+    if _HAS_CUDA_EXT and out.is_cuda and out.dtype == torch.complex64:
+        orig_shape = out.shape
+        flat = out.reshape(-1, orig_shape[-2], orig_shape[-1]).contiguous()
+        _ops.sht_prepare_irfft(flat, active_mmax, nlon)
+        return flat.reshape(orig_shape)
+
+    out[..., 0] = out[..., 0].real.to(torch.complex64)
+    if nlon % 2 == 0:
+        nyquist_idx = nlon // 2
+        if nyquist_idx < out.size(-1):
+            out[..., nyquist_idx] = out[..., nyquist_idx].real.to(torch.complex64)
+    return out
+
+
+class _FusedLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(1), input_c.size(2),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_inverse(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(0), grad_output.size(2),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_legendre_forward(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedLegendreForwardRealFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), weight_t.size(0), input_c.size(2),
+            device=input_c.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_forward_real(output, input_c, weight_t)
+        ctx.save_for_backward(weight_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (weight_t,) = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), weight_t.size(1), grad_output.size(2),
+            device=grad_output.device, dtype=torch.float32
+        )
+        _ops.fused_legendre_inverse_real(grad_input, grad_output, weight_t)
+        return grad_input, None
+
+
+class _FusedVectorLegendreForwardFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(0), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(1), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+class _FusedVectorLegendreInverseFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, weight0_t: torch.Tensor, weight1_t: torch.Tensor) -> torch.Tensor:
+        input_c = input.contiguous()
+        output = torch.empty(
+            input_c.size(0), 2, weight0_t.size(1), input_c.size(3),
+            device=input_c.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_inverse(output, input_c, weight0_t, weight1_t)
+        ctx.save_for_backward(weight0_t, weight1_t)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        weight0_t, weight1_t = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+        grad_input = torch.empty(
+            grad_output.size(0), 2, weight0_t.size(0), grad_output.size(3),
+            device=grad_output.device, dtype=torch.complex64
+        )
+        _ops.fused_vector_legendre_forward(grad_input, grad_output, weight0_t, weight1_t)
+        return grad_input, None, None
+
+
+# ============================================================================
+# Fused Legendre Transform
+# ============================================================================
+
+def fused_legendre_forward(
+    input: torch.Tensor,       # [B, nlat, mmax] complex64
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32 (original torch-harmonics layout)
+    weight_t: Optional[torch.Tensor] = None,  # [lmax, nlat, mmax] pre-transposed
+) -> torch.Tensor:
+    """Fused forward Legendre transform operating on complex tensors.
+
+    Computes out[b,l,m] = Σ_k weights[m,l,k] · input[b,k,m] for complex input,
+    fusing the real and imaginary multiplications into a single pass.
+
+    Uses the custom CUDA kernels when the extension is available, otherwise
+    falls back to a stacked einsum.
+    """
+    B = input.size(0)
+    nlat = input.size(1)
+    mmax = input.size(2)
+    lmax = weights.size(1)
+
+    if weight_t is None:
+        weight_t = weights.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, weight_t):
+        # Adaptive CUDA kernel: small-grid direct path + large-grid tiled path.
+        return _FusedLegendreForwardFn.apply(input, weight_t)
+    else:
+        # Fallback: stacked einsum (1.9x speedup over reference 2x einsum)
+        x = torch.view_as_real(input)  # [B, nlat, mmax, 2]
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)  # [2B, nlat, mmax]
+        w = weights.to(x_stacked.dtype)
+        out_stacked = torch.einsum("bkm,mlk->blm", x_stacked, w)
+        out_re = out_stacked[:B]
+        out_im = out_stacked[B:]
+        return torch.complex(out_re, out_im)
+
+
+def fused_legendre_inverse(
+    input: torch.Tensor,       # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Fused inverse Legendre transform."""
+    B = input.size(0)
+    lmax = input.size(1)
+    mmax = input.size(2)
+    nlat = pct.size(2)
+
+    if pct_t is None:
+        pct_t = pct.float().permute(1, 2, 0).contiguous()
+
+    if _can_use_cuda_legendre(input, pct_t):
+        return _FusedLegendreInverseFn.apply(input, pct_t)
+    else:
+        x = torch.view_as_real(input)
+        x_stacked = torch.cat([x[..., 0], x[..., 1]], dim=0)
+        p = pct.to(x_stacked.dtype)
+        out_stacked = torch.einsum("blm,mlk->bkm", x_stacked, p)
+        return torch.complex(out_stacked[:B], out_stacked[B:])
+
+
+def fused_legendre_forward_real(
+    input: torch.Tensor,       # [B, nlat, mmax] float32 or bfloat16
+    weight_t: torch.Tensor,    # [lmax, nlat, mmax] float32
+) -> torch.Tensor:
+    """Real-valued forward Legendre transform with float accumulation."""
+    if _can_use_cuda_real_legendre(input, weight_t):
+        return _FusedLegendreForwardRealFn.apply(input, weight_t)
+
+    return torch.einsum("bkm,lkm->blm", input.float(), weight_t)
+
+
+# ============================================================================
+# Fused SHT (complete pipeline)
+# ============================================================================
+
+def fused_sht_forward(
+    x: torch.Tensor,           # [B, nlat, nlon] float32
+    weights: torch.Tensor,     # [mmax, lmax, nlat] float32
+    mmax: int,
+    weight_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused forward SHT: rfft → fused Legendre → complex coefficients.
+
+    Replaces RealSHT.forward() with fewer intermediate allocations.
+    """
+    with _nvtx_range("holysht.scalar_forward"):
+        x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+        x_fft = x_fft[..., :mmax]
+        return fused_legendre_forward(x_fft, weights, weight_t)
+
+
+def fused_sht_inverse(
+    coeffs: torch.Tensor,      # [B, lmax, mmax] complex64
+    pct: torch.Tensor,         # [mmax, lmax, nlat] float32
+    nlon: int,
+    pct_t: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Complete fused inverse SHT: fused Legendre → zero-pad → irfft."""
+    with _nvtx_range("holysht.scalar_inverse"):
+        x = fused_legendre_inverse(coeffs, pct, pct_t)
+        x = _prepare_irfft_input(x, nlon, coeffs.size(-1))
+        return torch.fft.irfft(x, n=nlon, dim=-1, norm="forward")
+
+
+# ============================================================================
+# nn.Module wrappers (drop-in replacements for torch-harmonics)
+# ============================================================================
+
+class RealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealSHT``.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealSHT
+        ref = RealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("weights", ref.weights.to(w_dtype))
+        self.register_buffer("weight_t", ref.weights.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._use_bf16:
+            with _nvtx_range("holysht.scalar_forward_bf16"):
+                x_fft = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+                x_fft = x_fft[..., :self.mmax]
+                xr = torch.view_as_real(x_fft)
+                if _HAS_CUDA_EXT and x.is_cuda and not x.requires_grad:
+                    xr_bf16 = xr.bfloat16().contiguous()
+                    out_re = fused_legendre_forward_real(xr_bf16[..., 0].contiguous(), self.weight_t)
+                    out_im = fused_legendre_forward_real(xr_bf16[..., 1].contiguous(), self.weight_t)
+                    return torch.complex(out_re, out_im)
+
+                B = x.size(0)
+                xs = torch.cat([xr[..., 0], xr[..., 1]], dim=0).bfloat16()
+                out = torch.einsum("bkm,mlk->blm", xs, self.weights).float()
+                return torch.complex(out[:B], out[B:])
+        return fused_sht_forward(x, self.weights, self.mmax, self.weight_t)
+
+
+class InverseRealSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealSHT
+        ref = InverseRealSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("pct", ref.pct.float())
+        self.register_buffer("pct_t", ref.pct.float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return fused_sht_inverse(x, self.pct, self.nlon, self.pct_t)
+
+
+class RealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.RealVectorSHT``.
+
+    Reduces eight reference einsums to two composed Legendre passes on the
+    default FP32 CUDA path.
+
+    Args:
+        dtype: Weight precision. ``"fp32"`` (default) or ``"bf16"``.
+    """
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+        dtype: str = "fp32",
+    ):
+        super().__init__()
+        from torch_harmonics import RealVectorSHT
+        ref = RealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self._use_bf16 = (dtype == "bf16")
+        w_dtype = torch.bfloat16 if self._use_bf16 else torch.float32
+        self.register_buffer("w0", ref.weights[0].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w1", ref.weights[1].to(w_dtype))  # [mmax, lmax, nlat]
+        self.register_buffer("w0_t", ref.weights[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("w1_t", ref.weights[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.nlat and x.shape[-1] == self.nlon
+
+        with _nvtx_range("holysht.vector_forward"):
+            x = 2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+            mmax = self.mmax
+            x = x[..., :mmax].contiguous()
+
+            if (not self._use_bf16) and _can_use_cuda_vector(x, self.w0_t, self.w1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.nlat, mmax).contiguous()
+                out = _FusedVectorLegendreForwardFn.apply(x_flat, self.w0_t, self.w1_t)
+                return out.reshape(B_shape + (2, self.lmax, mmax))
+
+            x = torch.view_as_real(x)  # [..., 2, nlat, mmax, 2]
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.nlat, mmax)
+            x01_flat = x01.reshape(-1, self.nlat, mmax)
+            x10_flat = x10.reshape(-1, self.nlat, mmax)
+            x11_flat = x11.reshape(-1, self.nlat, mmax)
+            B = x00_flat.shape[0]
+
+            if self._use_bf16 and _HAS_CUDA_EXT and x00_flat.is_cuda and not x.requires_grad:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0).bfloat16().contiguous()
+                out_w0 = fused_legendre_forward_real(stacked_w0, self.w0_t)
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0).bfloat16().contiguous()
+                out_w1 = fused_legendre_forward_real(stacked_w1, self.w1_t)
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+            else:
+                stacked_w0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w0 = stacked_w0.bfloat16()
+                out_w0 = torch.einsum("bkm,mlk->blm", stacked_w0, self.w0)
+                if self._use_bf16:
+                    out_w0 = out_w0.float()
+                r00, r01, r10, r11 = out_w0[:B], out_w0[B:2 * B], out_w0[2 * B:3 * B], out_w0[3 * B:]
+
+                stacked_w1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+                if self._use_bf16:
+                    stacked_w1 = stacked_w1.bfloat16()
+                out_w1 = torch.einsum("bkm,mlk->blm", stacked_w1, self.w1)
+                if self._use_bf16:
+                    out_w1 = out_w1.float()
+                s11, s10, s01, s00 = out_w1[:B], out_w1[B:2 * B], out_w1[2 * B:3 * B], out_w1[3 * B:]
+
+            sph_re = r00 - s11
+            sph_im = r01 + s10
+            tor_re = -s01 - r10
+            tor_im = s00 - r11
+
+            out_shape = list(B_shape) + [2, self.lmax, mmax, 2]
+            xout = torch.zeros(out_shape, dtype=x.dtype, device=x.device)
+            xout[..., 0, :, :, 0] = sph_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 0, :, :, 1] = sph_im.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 0] = tor_re.reshape(B_shape + (self.lmax, mmax))
+            xout[..., 1, :, :, 1] = tor_im.reshape(B_shape + (self.lmax, mmax))
+            return torch.view_as_complex(xout)
+
+
+class InverseRealVectorSHT(nn.Module):
+    """Optimised drop-in replacement for ``torch_harmonics.InverseRealVectorSHT``."""
+
+    def __init__(
+        self,
+        nlat: int,
+        nlon: int,
+        lmax: Optional[int] = None,
+        mmax: Optional[int] = None,
+        grid: str = "equiangular",
+        norm: str = "ortho",
+        csphase: bool = True,
+    ):
+        super().__init__()
+        from torch_harmonics import InverseRealVectorSHT
+        ref = InverseRealVectorSHT(
+            nlat,
+            nlon,
+            lmax=lmax,
+            mmax=mmax,
+            grid=grid,
+            norm=norm,
+            csphase=csphase,
+        )
+        self.nlat = ref.nlat
+        self.nlon = ref.nlon
+        self.lmax = ref.lmax
+        self.mmax = ref.mmax
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+        self.register_buffer("d0", ref.dpct[0].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d1", ref.dpct[1].float())  # [mmax, lmax, nlat]
+        self.register_buffer("d0_t", ref.dpct[0].float().permute(1, 2, 0).contiguous())
+        self.register_buffer("d1_t", ref.dpct[1].float().permute(1, 2, 0).contiguous())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.shape[-2] == self.lmax and x.shape[-1] == self.mmax
+
+        with _nvtx_range("holysht.vector_inverse"):
+            x = x.contiguous()
+            if _can_use_cuda_vector(x, self.d0_t, self.d1_t):
+                B_shape = x.shape[:-3]
+                x_flat = x.reshape(-1, 2, self.lmax, self.mmax).contiguous()
+                x_out = _FusedVectorLegendreInverseFn.apply(x_flat, self.d0_t, self.d1_t)
+                x_out = x_out.reshape(B_shape + (2, self.nlat, self.mmax))
+                x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+                return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+            x = torch.view_as_real(x)  # [..., 2, lmax, mmax, 2]
+            mmax = self.mmax
+
+            x00 = x[..., 0, :, :, 0]
+            x01 = x[..., 0, :, :, 1]
+            x10 = x[..., 1, :, :, 0]
+            x11 = x[..., 1, :, :, 1]
+
+            B_shape = x00.shape[:-2]
+            x00_flat = x00.reshape(-1, self.lmax, mmax)
+            x01_flat = x01.reshape(-1, self.lmax, mmax)
+            x10_flat = x10.reshape(-1, self.lmax, mmax)
+            x11_flat = x11.reshape(-1, self.lmax, mmax)
+            B = x00_flat.shape[0]
+
+            stacked_d0 = torch.cat([x00_flat, x01_flat, x10_flat, x11_flat], dim=0)
+            out_d0 = torch.einsum("blm,mlk->bkm", stacked_d0, self.d0)
+            r00, r01, r10, r11 = out_d0[:B], out_d0[B:2 * B], out_d0[2 * B:3 * B], out_d0[3 * B:]
+
+            stacked_d1 = torch.cat([x11_flat, x10_flat, x01_flat, x00_flat], dim=0)
+            out_d1 = torch.einsum("blm,mlk->bkm", stacked_d1, self.d1)
+            s11, s10, s01, s00 = out_d1[:B], out_d1[B:2 * B], out_d1[2 * B:3 * B], out_d1[3 * B:]
+
+            srl = r00 - s11
+            sim = r01 + s10
+            trl = -s01 - r10
+            tim = s00 - r11
+
+            out_k = self.nlat
+            srl = srl.reshape(B_shape + (out_k, mmax))
+            sim = sim.reshape(B_shape + (out_k, mmax))
+            trl = trl.reshape(B_shape + (out_k, mmax))
+            tim = tim.reshape(B_shape + (out_k, mmax))
+
+            s = torch.stack((srl, sim), -1)
+            t = torch.stack((trl, tim), -1)
+            xs = torch.stack((s, t), -4)
+            x_out = torch.view_as_complex(xs)
+            x_out = _prepare_irfft_input(x_out, self.nlon, self.mmax)
+            return torch.fft.irfft(x_out, n=self.nlon, dim=-1, norm="forward")
+
+
+legendre_forward = fused_legendre_forward
+legendre_inverse = fused_legendre_inverse
+sht_forward = fused_sht_forward
+sht_inverse = fused_sht_inverse
+
+# Backwards-compatible aliases from the research prototype.
+FusedRealSHT = RealSHT
+FusedInverseRealSHT = InverseRealSHT
+FusedRealVectorSHT = RealVectorSHT
+FusedInverseRealVectorSHT = InverseRealVectorSHT

build/torch29-cxx11-cu129-x86_64-linux/_jit_ops.py

@@ -0,0 +1,72 @@
+"""Local CUDA extension loader for HOLYSHT.
+
+This keeps development off the heavyweight `kernel-builder` path by compiling a
+small torch extension for the current machine only. The compiled artefacts are
+cached under `build/torch_extensions`.
+
+Author: Chris von Csefalvay
+Licence: MIT
+Repository: https://github.com/chrisvoncsefalvay/holysht
+Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+from torch.utils.cpp_extension import load
+
+
+_ROOT = Path(__file__).resolve().parents[2]
+_BUILD_DIR = _ROOT / "build" / "torch_extensions"
+_EXTENSION_NAME = "holysht_ops"
+
+
+def _default_arch_list() -> str:
+    if not torch.cuda.is_available():
+        return ""
+
+    major, minor = torch.cuda.get_device_capability(0)
+    if (major, minor) >= (12, 1):
+        return "12.0+PTX"
+    return f"{major}.{minor}"
+
+
+def _cuda_flags() -> list[str]:
+    flags = [
+        "-O3",
+        "-lineinfo",
+        "-Xptxas=-warn-spills",
+        "--expt-relaxed-constexpr",
+    ]
+    if os.environ.get("HOLYSHT_USE_FAST_MATH", "1") != "0":
+        flags.append("--use_fast_math")
+    return flags
+
+
+def _load_ops():
+    os.environ.setdefault("MAX_JOBS", "1")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _default_arch_list())
+    _BUILD_DIR.mkdir(parents=True, exist_ok=True)
+
+    load(
+        name=_EXTENSION_NAME,
+        sources=[
+            str(_ROOT / "torch-ext" / "torch_binding.cpp"),
+            str(_ROOT / "cuda" / "fused_legendre.cu"),
+            str(_ROOT / "cuda" / "fused_sht.cu"),
+        ],
+        extra_include_paths=[str(_ROOT / "torch-ext")],
+        extra_cflags=["-O3", "-std=c++17"],
+        extra_cuda_cflags=_cuda_flags(),
+        build_directory=str(_BUILD_DIR),
+        verbose=os.environ.get("HOLYSHT_VERBOSE_BUILD", "0") == "1",
+        with_cuda=True,
+        is_python_module=False,
+    )
+    return getattr(torch.ops, _EXTENSION_NAME)
+
+
+ops = _load_ops()

build/torch29-cxx11-cu129-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _holysht_cuda_53048d0
+ops = torch.ops._holysht_cuda_53048d0
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_holysht_cuda_53048d0::{op_name}"

build/torch29-cxx11-cu129-x86_64-linux/holysht/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch29-cxx11-cu129-x86_64-linux/metadata.json

@@ -0,0 +1,12 @@
+{
+  "version": 1,
+  "license": "mit",
+  "upstream": "https://github.com/NVIDIA/torch-harmonics",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "12.0"
+    ]
+  }
+}

flake.lock

@@ -0,0 +1,117 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1767039857,
+        "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs",
+        "rust-overlay": "rust-overlay"
+      },
+      "locked": {
+        "lastModified": 1775826501,
+        "narHash": "sha256-Do1JnE2aRIPW5Y2dbgx/Wj77/0wbvuROtCFGYnDUaKY=",
+        "owner": "huggingface",
+        "repo": "kernels",
+        "rev": "2fe18de4da9a6769e9cc3b615aff27db5b3d21dd",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernels",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1774935083,
+        "narHash": "sha256-Mh6bLcYAcENBAZk3RoMPMFCGGMZmfaGMERE4siZOgP4=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "2f4fd5e1abf9bac8c1d22750c701a7a5e6b524c6",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "ref": "nixos-unstable-small",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "rust-overlay": {
+      "inputs": {
+        "nixpkgs": [
+          "kernel-builder",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1774926780,
+        "narHash": "sha256-JMdDYn0F+swYBILlpCeHDbCSyzqkeSGNxZ/Q5J584jM=",
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "rev": "962a0934d0e32f42d1b5e49186f9595f9b178d2d",
+        "type": "github"
+      },
+      "original": {
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "type": "github"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -0,0 +1,17 @@
+# HOLYSHT
+# Author: Chris von Csefalvay
+# Licence: MIT
+# Repository: https://github.com/chrisvoncsefalvay/holysht
+# Hugging Face kernel: https://hf.co/chrisvoncsefalvay/holysht
+
+{
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernels";
+  };
+  outputs =
+    { self, kernel-builder, ... }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}