diff --git a/CHANGELOG.md b/CHANGELOG.md index 7e951ccd..5caf8d3d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,55 @@ All notable changes to nvSubquadratic are documented here. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## \[Unreleased\] + +### Added + +- **`fft_backend="subq_ops_fused"` on `CKConvND`**, backed by + `subquadratic_ops_torch.fused_fft_conv2d`. It runs the whole + rfft2 → multiply → irfft2 pipeline in a single cuFFTDx launch and, unlike + every other FFT path, **natively in fp32/fp16/bf16** instead of upcasting to + fp32. Roughly 3–4× faster than `torch_fft` and 1.2–2.4× faster than + `subq_ops` on forward+backward in bf16. + + Restricted to `data_dim=2`, `is_causal=False`, `fft_padding="zero"`, and + spatial extents of at most **64 per axis** — the kernel's largest FFT tile is + 128 and it requires `max(X, Y) <= fft_size // 2`. The spatial cap is enforced + on the first forward pass rather than at construction, because the input size + is not known when the module is built. Per-sample (FiLM) kernels are + supported. Requires `subquadratic-ops-torch >= 0.2.2`. + + The upstream kernel crops the 'same' window at `fft_size // 2` whereas + `fftconv.py` crops at `K // 2`; the wrapper pre-pads the filter's top/left by + the difference so results are interchangeable with the other backends + (verified to ~3e-7 normwise in fp32). Without that pre-pad the output is + shifted by `fft_size // 2 - K // 2` pixels. + +- **`nvsubquadratic.ops.fftconv_lowering`** — a `torch.compile` pre-grad pass + that detects `fftconv.py`'s 2D FFT-conv chain and rewrites it onto the fused + kernel, so a model already on `fft_backend="torch_fft"` picks it up without a + config change. This matters because inductor cannot generate code for complex + operators and otherwise falls back to eager cuFFT for the entire chain. + + Enable it per-callable with + `torch.compile(model, options=fused_fftconv2d_options())`, or globally for a + scope with the `fused_fftconv2d_lowering()` context manager when a framework + owns the `torch.compile` call. + + The pass only fires on an exact match of the reference recipe (padding rule, + crop offset, shape limits, CUDA device, and a compute capability that + supports the required FFT tile — the 128 tile needs SM90+). `lowering_stats()` + reports rewrite and per-reason skip counts, since a silent pass is otherwise + hard to tell apart from one that never fired. + +### Fixed + +- `tests/conftest.py` only queried the `subquadratic-ops-torch-cu12` + distribution when resolving the installed kernel version. On a + `-cu13` install (what `pyproject.toml` pins) it resolved to `(0, 0, 0)`, + which silently turned `requires_subq_ops_v2` into a blanket `xfail` and hid + every `subq_ops` test result. It now checks both distributions. + ## \[0.1.1\] ### Changed diff --git a/Dockerfile b/Dockerfile index 4297d96a..6505475f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,7 +1,10 @@ -# Development Dockerfile for nvSubquadratic +# Development Dockerfile for nvSubquadratic (CUDA 13.2 / PyTorch cu132) # # Build instructions: -# docker build -t nvsubquadratic:dev . +# DOCKER_BUILDKIT=1 docker build -t nvsubquadratic:dev . +# +# Requires an NVIDIA driver >= 580 on the host (CUDA 13.x minimum), or the +# forward-compatibility package. Volta (sm_70) GPUs are no longer supported. # # Layer order is intentional for CI cache efficiency: # 1. Base image + conda + torch + DALI (never changes) @@ -9,11 +12,20 @@ # 3. requirements-dev.txt (changes when dev deps change) # 4. COPY . . + pip install (changes on every code push — fast) -FROM nvcr.io/nvidia/cuda:12.9.0-devel-ubuntu22.04 +FROM nvcr.io/nvidia/cuda:13.2.1-devel-ubuntu22.04 ARG MINIFORGE_NAME=Miniforge3 ARG MINIFORGE_VERSION=25.3.0-3 +# ── CUDA 13.2 toolchain pins ───────────────────────────────────────────────── +# CUDA 13 wheels live under the cu132 index. The first PyTorch release built +# against CUDA 13.2 is 2.12.0, so the previous 2.10.0 pin cannot be kept here. +ARG TORCH_VERSION=2.12.1 +ARG TORCHVISION_VERSION=0.27.1 +ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cu132 +# DALI ships one build per CUDA major version; cuda130 is the CUDA 13.x build. +ARG DALI_PACKAGE=nvidia-dali-cuda130 + ENV CONDA_DIR=/opt/conda ENV LANG=C.UTF-8 LC_ALL=C.UTF-8 ENV PATH=${CONDA_DIR}/bin:${PATH} @@ -40,11 +52,14 @@ RUN conda install --yes \ python=3.12 \ && conda clean --all --yes +# The torch install is repeated after DALI on purpose: DALI pulls its own +# nvidia-* CUDA runtime wheels and can shuffle the ones torch depends on, so we +# re-assert the pinned cu132 build afterwards. RUN pip install --no-cache-dir \ - torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu129 \ - && pip install --no-cache-dir nvidia-dali-cuda120 \ + torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} --index-url ${TORCH_INDEX_URL} \ + && pip install --no-cache-dir ${DALI_PACKAGE} \ && pip install --no-cache-dir \ - torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu129 \ + torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} --index-url ${TORCH_INDEX_URL} \ && conda clean --all --yes # Create ubuntu user with sudo privileges @@ -66,7 +81,11 @@ WORKDIR /workspaces/nvSubquadratic # ── Heavy build: Apex from source (cached until apex commit changes) ────────── # This layer is intentionally placed before COPY so code changes do not # trigger a rebuild. Apex does not depend on the project source. -ARG TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6;8.9;9.0;10.0;12.0" +# CUDA 13 dropped offline compilation for Maxwell/Pascal/Volta — `nvcc +# --list-gpu-arch` in 13.2 starts at compute_75, so 7.0 must be removed or the +# apex build fails with "Unsupported gpu architecture 'compute_70'". +# Newly available in 13.2 if you need them: 8.7, 8.8, 10.3 (B300), 11.0, 12.1. +ARG TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;12.0" ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} ARG MAX_JOBS="" RUN MAX_JOBS="${MAX_JOBS}" pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \ @@ -90,7 +109,7 @@ RUN git config --global --add safe.directory /workspaces/nvSubquadratic # pulls them — [all] restores the complete pre-restructure dependency set. RUN pip install --no-cache-dir wheel-stub \ && pip install --no-cache-dir --no-build-isolation ".[all]" \ - --extra-index-url https://download.pytorch.org/whl/cu129 + --extra-index-url ${TORCH_INDEX_URL} # Set up ubuntu user's home directory and permissions RUN chown -R ubuntu:ubuntu /workspaces && \ diff --git a/README.md b/README.md index 4a300031..89a1e8fc 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ nvSubquadratic consolidates efforts from across NVIDIA Research teams (nvResearc - **B2B CausalConv1d**: Back-to-back causal convolutions for striped Hyena architectures - **CausalConv1d**: Standard causal convolutions with various kernel sizes (2-256) - **FFT CausalConv1d**: FFT-based causal convolutions for large kernel sizes (up to 8K-16M) +- **Fused FFT Conv2d**: single-launch 2D FFT convolution running natively in fp32/fp16/bf16 (spatial dims up to 64 per axis); requires `subquadratic-ops-torch >= 0.2.2` **Requirements**: @@ -60,8 +61,15 @@ pip install "nvsubquadratic[all]" # all of the above The accelerated CUDA kernels (`[cuda]`) are a source build that requires `nvcc`, so they are kept out of the core install — this is what lets `pip install nvsubquadratic` succeed in environments without the CUDA toolkit (e.g. a downstream project's CPU CI). The operators default to the portable `torch.fft` -backend; selecting `fft_backend="subq_ops"` without `[cuda]` installed raises a -clear `ImportError` pointing you to the extra. +backend; selecting `fft_backend="subq_ops"` (or `"subq_ops_fused"`) without +`[cuda]` installed raises a clear `ImportError` pointing you to the extra. + +On 2D problems with spatial dims of at most 64 per axis, `fft_backend="subq_ops_fused"` +is the fastest option: it fuses the whole FFT-conv pipeline into one launch and +runs it natively in bf16/fp16 instead of upcasting to fp32, for roughly a 3-4x +speedup over `torch_fft`. Models already written against `torch_fft` can pick up +the same kernel under `torch.compile` without a config change — see the +[torch.compile lowering](docs/ops/README.md#torchcompile-lowering). ### Package Manager diff --git a/docs/api_reference/ops.rst b/docs/api_reference/ops.rst index bd0a90d2..75396606 100644 --- a/docs/api_reference/ops.rst +++ b/docs/api_reference/ops.rst @@ -49,6 +49,47 @@ fp32 reference ops above. ~ops.fftconv_custom.causal_fftconv1d_bhl ~ops.fftconv_custom.causal_fftconv1d_bhl_w_reshape +.. _ops-fftconv-fused: + +Fused 2D FFT convolutions (native dtype) +----------------------------------------- + +Wrappers around ``subquadratic_ops_torch.fused_fft_conv2d``, which fuses the +whole rfft2/multiply/irfft2 pipeline into one launch and runs it in the input +dtype (fp32/fp16/bf16) rather than upcasting to fp32. Spatial dims are capped +at 64 per axis. Selected on ``CKConvND`` via ``fft_backend="subq_ops_fused"``. + +.. autosummary:: + :toctree: generated/ + :template: function_template.rst + + ~ops.fftconv_custom.fused_fftconv2d_blh + ~ops.fftconv_custom.fused_fftconv2d_bhl + ~ops.fftconv_custom.fused_fftconv2d_bhl_w_reshape + ~ops.fftconv_custom.fused_fftconv2d_bhl_chunked + ~ops.fftconv_custom.resolve_fused_fft_size + ~ops.fftconv_custom.fused_fftconv2d_supported + ~ops.fftconv_custom.fused_fftconv2d_max_spatial + ~ops.fftconv_custom.load_fused_fft_conv2d + +.. _ops-fftconv-lowering: + +torch.compile lowering +----------------------- + +Inductor pre-grad pass that rewrites the reference 2D FFT-conv chain onto the +fused CUDA kernel, so an existing ``fft_backend="torch_fft"`` model picks it up +without a config change. + +.. autosummary:: + :toctree: generated/ + :template: function_template.rst + + ~ops.fftconv_lowering.fused_fftconv2d_options + ~ops.fftconv_lowering.fused_fftconv2d_lowering + ~ops.fftconv_lowering.lowering_stats + ~ops.fftconv_lowering.reset_lowering_stats + .. _ops-causal-conv1d: Direct 1D causal convolutions (CUDA-accelerated) diff --git a/docs/getting_started.md b/docs/getting_started.md index f2abbbff..45c5a4d8 100644 --- a/docs/getting_started.md +++ b/docs/getting_started.md @@ -100,7 +100,10 @@ hyena_cfg = LazyConfig(Hyena)( mask_cfg=LazyConfig(torch.nn.Identity)(), grid_type="double", # linear (non-circular) convolution fft_padding="zero", - fft_backend="torch_fft", # portable; "subq_ops" uses the fused 2D CUDA kernel + # Portable default. "subq_ops" uses the fused 2D CUDA kernel; + # "subq_ops_fused" is faster still and runs natively in bf16/fp16, + # but caps spatial dims at 64 per axis. + fft_backend="torch_fft", is_causal=False, ), # Depthwise short conv on the concatenated [Q; K; V] (3 * H channels). diff --git a/docs/ops/README.md b/docs/ops/README.md index b7bbc934..bbb03b8f 100644 --- a/docs/ops/README.md +++ b/docs/ops/README.md @@ -41,7 +41,8 @@ ______________________________________________________________________ | {ref}`circular_fftconv.py ` | fp32 | circular | depthwise | Periodic boundaries (e.g. PDEs, ARC grids), or when `K = N` so padding is wasteful. | | {ref}`mixed_fftconv.py ` | fp32 | per-axis BC | depthwise | **Mixed boundaries**: periodic on some spatial axes, zero-padded on others (e.g. Well's `rayleigh_benard`, `viscoelastic_instability`, `turbulent_radiative_layer`). Routes to the existing linear/circular ops in the all-False/all-True cases. | | {ref}`fftconv_chunked.py ` | fp32 | linear | depthwise | Memory-constrained training; processes channels in chunks. Has a global flag so models can opt in transparently. | -| {ref}`fftconv_custom.py ` | fp32 | linear | depthwise | Wraps optional fused CUDA kernels (`subquadratic_ops_torch.fft_conv2d` for 2D non-causal, `fft_causal_conv1d` for 1D causal) behind the same API as `fftconv.py`. | +| {ref}`fftconv_custom.py ` | fp32 / native | linear | depthwise | Wraps optional fused CUDA kernels (`subquadratic_ops_torch.fft_conv2d` for 2D non-causal, `fft_causal_conv1d` for 1D causal) behind the same API as `fftconv.py`. Also hosts the `fused_fftconv2d_*` wrappers around `fused_fft_conv2d`, which run natively in fp32/fp16/bf16 but cap spatial dims at 64 per axis. | +| {ref}`fftconv_lowering.py ` | native | linear | depthwise | `torch.compile` pre-grad pass that rewrites `fftconv.py`'s 2D chain onto `fused_fft_conv2d`. Use when you want the fused kernel without changing `fft_backend` on an existing model. | | {ref}`causal_conv1d_custom.py ` | fp32 | direct causal | depthwise | Non-FFT 1D causal kernels (`causal_conv1d` short conv, `b2b_causal_conv1d` fused proj-gate-mixer-gate). Use for kernels short enough that FFT overhead dominates, or as a fused-Hyena building block. | The [FP16 circular FFT convolution report](https://github.com/NVIDIA-BioNeMo/nvSubquadratic/blob/main/reports/fp16_fft_convolution/REPORT.md) contains the full derivation of the numerically stable fp16 circular conv (dual mean-centering + inclusion-exclusion geometric correction). Read it if you are touching the fp16 path or want to understand the math behind those `T1, T2, T3, T4` terms in the code. @@ -123,14 +124,74 @@ ______________________________________________________________________ 1. **Is there a fused CUDA kernel for my shape?** - 2D non-causal or 1D causal long-conv → `fftconv_custom.py` exposes the upstream fused FFT kernels (`fft_conv2d`, `fft_causal_conv1d`) through the same API. Wire in via the `fft_backend="subq_ops"` flag on `CKConvND`. The 1D path requires `data_dim=1, is_causal=True`; the 2D path requires `data_dim=2, is_causal=False`. + - 2D non-causal with spatial dims **≤ 64 per axis** → prefer `fft_backend="subq_ops_fused"`, which uses the newer `fused_fft_conv2d` kernel. It runs the whole rfft2/multiply/irfft2 pipeline in one launch **natively in fp32/fp16/bf16** instead of upcasting to fp32 — roughly 3–4× faster than `torch_fft` end-to-end in bf16 and ~1.2–2.4× faster than `subq_ops`. Requires `subquadratic-ops-torch >= 0.2.2`. The 64×64 cap comes from the kernel's largest FFT tile (128) combined with its `max(X, Y) ≤ fft_size // 2` requirement, and is enforced on the first forward pass rather than at construction. Beyond that, fall back to `subq_ops` or `torch_fft`. + - Already have a model on `fft_backend="torch_fft"` and want the fused kernel without a config change → `fftconv_lowering.py` provides a `torch.compile` pre-grad pass that detects the FFT-conv chain and rewrites it. See [the lowering section below](#torchcompile-lowering). - 1D causal short conv (typical short_conv slot in a Hyena block) → `causal_conv1d_custom.py` exposes `causal_conv1d` directly, and `nvsubquadratic.modules.subq_ops_causal_conv1d.SubqOpsCausalConv1d` wraps it as a depthwise `nn.Conv1d`-compatible module. - 1D causal fused proj+gate+mixer+gate block → `b2b_causal_conv1d` in `causal_conv1d_custom.py`. Not yet wired into a `Hyena` variant; exposed as a building block. ______________________________________________________________________ +## torch.compile lowering + +Inductor cannot generate code for complex operators — it warns +`Torchinductor does not support code generation for complex operators` and falls +back to eager cuFFT for the whole `rfft2 → multiply → irfft2` chain. So +`torch.compile` on its own buys essentially nothing for `fftconv.py`. + +`fftconv_lowering.py` closes that gap. It registers an inductor **pre-grad** +pass (before AOTAutograd, so autograd comes from the custom op's registered +backward) that detects the chain and swaps in the fused kernel: + +```python +import torch +from nvsubquadratic.ops.fftconv_lowering import fused_fftconv2d_options, lowering_stats + +compiled = torch.compile(model, options=fused_fftconv2d_options()) +out = compiled(x) + +print(lowering_stats()) # {'rewritten': 1} +``` + +`fused_fftconv2d_options()` is scoped to that one compiled callable and leaves +global inductor config untouched. When a trainer or framework owns the +`torch.compile` call and you cannot pass `options`, the +`fused_fftconv2d_lowering()` context manager does the same thing by patching +global config for its duration: + +```python +from nvsubquadratic.ops.fftconv_lowering import fused_fftconv2d_lowering + +with fused_fftconv2d_lowering(): + out = trainer.fit(model) # whatever compiles internally +``` + +The pass only fires when the graph is *exactly* `fftconv.py`'s recipe — the +`min(N + (K+1)//2, 2N)` padding, the `K // 2` crop, shapes within the fused +kernel's limits, CUDA device, and a compute capability that supports the +required FFT tile (the 128 tile needs SM90+; SM80/SM86 lack the shared memory). +Anything else is left on the eager path. + +Two things to know: + +- **Setting `fft_backend="subq_ops_fused"` is the more predictable route** and + needs no pass at all. Reach for the lowering when you want the kernel under an + existing `torch_fft` model without touching its config. +- **`lowering_stats()` only counts passes that actually ran.** Inductor caches + compiled artifacts on disk, and a cache hit skips every pre-grad pass, so + empty counters can mean "served from cache" rather than "did not fire". Set + `torch._inductor.config.force_disable_caches = True` when verifying. + +By default the pass also rewrites fp16/bf16 graphs, which moves the convolution +from fp32-internal to native-dtype — that is where most of the speedup comes +from, and it is a real numerics change (~2e-3 normwise in bf16). Pass +`fused_fftconv2d_lowering(allow_reduced_precision=False)` to restrict it to +fp32, where the rewrite is numerically neutral. + +______________________________________________________________________ + ## Numerical notes -- All operators **accept any input dtype** but cast to the internal compute precision (fp32 or fp16) before the FFT. The output is returned in the **original dtype of `x`**, so no manual cast is needed on the caller side. +- All operators **accept any input dtype** but cast to the internal compute precision (fp32 or fp16) before the FFT. The output is returned in the **original dtype of `x`**, so no manual cast is needed on the caller side. The one exception is the `fused_fftconv2d_*` family, which runs in `x.dtype` end to end — that is the point of it. Measured against the fp32 reference, its normwise relative error is ~3e-7 in fp32, ~3e-4 in fp16, and ~3e-3 in bf16. - The fp32 ops are correct for any input range. The fp16 ops impose two constraints: spatial dims must be powers of two (cuFFT), and the dynamic range is handled by mean-centering both `x` and `k` (see derivation doc). - The non-causal linear ops match a standard `torch.nn.ConvNd(padding='same')` up to floating-point rounding. The circular ops match `torch.nn.functional.conv*d` after a circular pad. Both are exercised in `tests/`. diff --git a/nvsubquadratic/modules/ckconv_nd.py b/nvsubquadratic/modules/ckconv_nd.py index 55d5ef59..f8c269f6 100644 --- a/nvsubquadratic/modules/ckconv_nd.py +++ b/nvsubquadratic/modules/ckconv_nd.py @@ -497,8 +497,8 @@ class CKConvND(torch.nn.Module): convolution. Only valid when ``data_dim=1``. use_chunked_fftconv (bool): Whether to process channels in chunks to reduce peak GPU memory. - fft_backend (str): FFT backend identifier, ``"torch_fft"`` or - ``"subq_ops"``. + fft_backend (str): FFT backend identifier, ``"torch_fft"``, + ``"subq_ops"``, or ``"subq_ops_fused"``. grid_type (str or None): Kernel grid size mode (``"single"``, ``"double"``, or ``None`` for per-axis auto-derivation). kernel (nn.Module): Implicit kernel generator (produces @@ -531,7 +531,7 @@ def __init__( fft_padding: "Literal['zero', 'circular'] | str | Sequence[str]", is_causal: bool = False, use_chunked_fftconv: bool = False, - fft_backend: Literal["torch_fft", "subq_ops"] = "torch_fft", + fft_backend: Literal["torch_fft", "subq_ops", "subq_ops_fused"] = "torch_fft", ): """Construct a CKConvND operator. @@ -614,6 +614,19 @@ def __init__( Does not support fp16 FFT, per-axis ``fft_padding``, or ``data_dim=3``. + * ``"subq_ops_fused"``: the fused ``fused_fft_conv2d`` CUDA + kernel from ``subquadratic_ops_torch``, which runs the whole + rfft2 → multiply → irfft2 pipeline in a single launch and + **natively in the input dtype** (fp32/fp16/bf16) instead of + upcasting to fp32. Roughly 3-4x faster than ``"torch_fft"`` + end-to-end in bf16. Restricted to ``data_dim=2``, + ``is_causal=False``, ``fft_padding="zero"``, and spatial + extents of at most 64 per axis (the kernel's largest FFT tile + is 128 and it requires ``max(X, Y) <= fft_size // 2``). The + spatial cap is enforced on the first forward pass, not at + construction, since the input size is not known here. + Per-sample (FiLM) batched kernel weights are supported. + Raises: AssertionError: If ``fft_backend`` is not one of the recognised values, or if a constraint between ``grid_type``, @@ -625,8 +638,8 @@ def __init__( or if the resolved ``(fft_padding, data_dim)`` combination has no registered FFT function. """ - assert fft_backend in ["torch_fft", "subq_ops"], ( - f"Invalid fft_backend: {fft_backend!r}. Must be 'torch_fft' or 'subq_ops'." + assert fft_backend in ["torch_fft", "subq_ops", "subq_ops_fused"], ( + f"Invalid fft_backend: {fft_backend!r}. Must be 'torch_fft', 'subq_ops', or 'subq_ops_fused'." ) # ---- Normalise fft_padding & grid_type -------------------------------- @@ -698,13 +711,34 @@ def __init__( ) # subq_ops backend constraints - if fft_backend == "subq_ops": + if fft_backend in ("subq_ops", "subq_ops_fused"): if _is_tuple_mode: raise ValueError( - "fft_backend='subq_ops' does not support a per-axis fft_padding. " + f"fft_backend='{fft_backend}' does not support a per-axis fft_padding. " "The CUDA kernel implements zero-padded conv only. " "Use fft_backend='torch_fft' for mixed boundary conditions." ) + + # The fused kernel is 2D-only: there is no 1D or 3D fused variant. + # Its spatial cap (<= 64 per axis) depends on the input size, which is + # unknown until forward(), so it is enforced by resolve_fused_fft_size + # inside the op wrapper. + if fft_backend == "subq_ops_fused": + if data_dim != 2: + raise AssertionError( + f"fft_backend='subq_ops_fused' only supports data_dim=2 " + f"(no fused 1D/3D kernel exists). Got data_dim={data_dim}. " + "Use fft_backend='subq_ops' for 1D causal." + ) + assert not is_causal, ( + "fft_backend='subq_ops_fused' does not support causal convolutions (causal is 1D only)." + ) + assert fft_padding == "zero", ( + "fft_backend='subq_ops_fused' only supports zero-padded convolutions. " + f"Got fft_padding='{fft_padding}'." + ) + + if fft_backend == "subq_ops": if data_dim == 1: assert is_causal, ( "fft_backend='subq_ops' on 1D requires is_causal=True " @@ -829,6 +863,22 @@ def __init__( f"fft_backend='subq_ops' dispatch reached unexpected data_dim={data_dim}; " "the constraint block above should have rejected this." ) + elif fft_backend == "subq_ops_fused": + # 2D-only (enforced above). The fused op keeps the input dtype end + # to end, so unlike the "subq_ops" path there is no fp32 round-trip. + from nvsubquadratic.ops.fftconv_custom import ( + fused_fftconv2d_bhl, + fused_fftconv2d_bhl_chunked, + fused_fftconv2d_bhl_w_reshape, + fused_fftconv2d_bhl_w_reshape_chunked, + ) + + if use_chunked_fftconv: + self.fftconv_fn = fused_fftconv2d_bhl_w_reshape_chunked + self.fftconv_fn_bhl_input = fused_fftconv2d_bhl_chunked + else: + self.fftconv_fn = fused_fftconv2d_bhl_w_reshape + self.fftconv_fn_bhl_input = fused_fftconv2d_bhl elif self._is_tuple_mode: # Per-axis ``fft_padding`` (list of mode strings, e.g. # ["circular", "zero"]): route through the unified diff --git a/nvsubquadratic/ops/fftconv_custom.py b/nvsubquadratic/ops/fftconv_custom.py index f2cfe91f..b1cf49fa 100644 --- a/nvsubquadratic/ops/fftconv_custom.py +++ b/nvsubquadratic/ops/fftconv_custom.py @@ -38,6 +38,14 @@ - ``fftconv2d_blh`` / ``fftconv2d_blh_chunked``: aliases for the ``_w_reshape`` variants (BLH naming convention). +2D fused (non-causal, zero-padded, **native dtype**): + +- ``fused_fftconv2d_bhl`` / ``fused_fftconv2d_bhl_chunked`` and the matching + ``_w_reshape`` / ``_blh`` variants. These wrap + ``subquadratic_ops_torch.fused_fft_conv2d``, which runs the whole + rfft2 → multiply → irfft2 pipeline as a single cuFFTDx kernel in the + input dtype — no fp32 upcast — at the cost of a ``64x64`` spatial cap. + 1D causal: - ``causal_fftconv1d_bhl`` / ``causal_fftconv1d_bhl_chunked``: BHL layout @@ -48,9 +56,12 @@ - ``causal_fftconv1d_blh`` / ``causal_fftconv1d_blh_chunked``: aliases for the ``_w_reshape`` variants. -All functions accept any input dtype (bf16, fp16, fp32) and internally cast -to fp32 for the CUDA kernel, returning the output in the original dtype. -Shortcut semantics are identical to the torch.fft reference: +The ``fftconv2d*`` / ``causal_fftconv1d*`` functions accept any input dtype +(bf16, fp16, fp32) and internally cast to fp32 for the CUDA kernel, returning +the output in the original dtype. The ``fused_fftconv2d*`` functions instead +run the kernel *natively* in the input dtype, which is where their speedup +over the fp32 path comes from. Shortcut semantics are identical throughout, +and match the torch.fft reference: :math:`y \leftarrow y + \text{shortcut} \odot x`. The chunked variants process channels in groups of ``chunk_size`` to reduce @@ -80,9 +91,20 @@ "fftconv2d_bhl_w_reshape_chunked", "fftconv2d_blh", "fftconv2d_blh_chunked", + "fused_fftconv2d_bhl", + "fused_fftconv2d_bhl_chunked", + "fused_fftconv2d_bhl_w_reshape", + "fused_fftconv2d_bhl_w_reshape_chunked", + "fused_fftconv2d_blh", + "fused_fftconv2d_blh_chunked", + "fused_fftconv2d_max_spatial", + "fused_fftconv2d_supported", + "load_fused_fft_conv2d", + "resolve_fused_fft_size", ] import torch +import torch.nn.functional as F from einops import rearrange @@ -260,6 +282,310 @@ def fftconv2d_blh_chunked( return fftconv2d_bhl_w_reshape_chunked(x, kernel, shortcut, chunk_size) +# =========================================================================== +# 2D fused FFT conv — wraps subquadratic_ops_torch.fused_fft_conv2d +# +# Unlike ``fft_conv2d`` above, this kernel runs natively in fp32/fp16/bf16 and +# fuses the filter FFT, the spectral multiply, and the inverse FFT into a +# single cuFFTDx launch. Two contract differences drive the code below: +# +# 1. Square power-of-two FFT tile from a fixed table, with the input capped at +# ``fft_size // 2`` — so the spatial extent is capped at 64x64. +# 2. The 'same' crop is taken at offset ``fft_size // 2``, whereas the torch +# reference in :mod:`nvsubquadratic.ops.fftconv` crops at ``K // 2``. We +# reconcile the two by pre-padding the filter's top/left by the difference, +# which shifts the convolution back by exactly that amount. Without it the +# fused path would be offset by ``fft_size // 2 - K // 2`` pixels relative to +# every other backend. +# =========================================================================== + +_fused_fft_conv2d = None + +# FFT tile sizes the fused CUDA kernel is compiled for. +_FUSED_FFT_SIZES = (8, 16, 32, 64, 128) + +# The kernel requires ``max(X, Y) <= fft_size // 2``, so the largest tile +# (128) caps the spatial extent at 64x64. +_FUSED_MAX_SPATIAL = max(_FUSED_FFT_SIZES) // 2 + + +def _get_fused_fft_conv2d(): + """Return the ``fused_fft_conv2d`` callable, importing on first call.""" + global _fused_fft_conv2d + if _fused_fft_conv2d is None: + try: + from subquadratic_ops_torch.fused_fft_conv2d import fused_fft_conv2d + + _fused_fft_conv2d = fused_fft_conv2d + except ImportError as exc: + raise ImportError( + "subquadratic_ops_torch >= 0.2.2 is required for fft_backend='subq_ops_fused'. " + "Install the accelerated CUDA kernels with: pip install 'nvsubquadratic[cuda]'" + ) from exc + return _fused_fft_conv2d + + +def load_fused_fft_conv2d() -> None: + """Import the fused CUDA kernel now instead of on first call. + + Importing ``subquadratic_ops_torch.fused_fft_conv2d`` is what registers the + ``torch.ops.subquadratic_ops_torch.fused_*`` operators. The wrappers here + import it lazily, which is fine in eager mode — but a ``torch.compile`` + artifact restored from the on-disk FX cache references those operators + without any Python call happening first, and fails with a confusing + ``AttributeError`` on the op namespace. Callers that compile ahead of the + first eager call should invoke this to force registration up front. + + Raises: + ImportError: If ``subquadratic_ops_torch`` is not installed. + """ + _get_fused_fft_conv2d() + + +def _next_pow2(n: int) -> int: + """Smallest power of two >= ``n`` (with ``_next_pow2(0) == 1``).""" + return 1 << max(n - 1, 0).bit_length() + + +def fused_fftconv2d_max_spatial() -> int: + """Largest supported spatial extent per axis for the fused 2D kernel. + + Returns: + ``64`` — the fused kernel's largest FFT tile is 128 and it requires + ``max(X, Y) <= fft_size // 2``. + """ + return _FUSED_MAX_SPATIAL + + +def resolve_fused_fft_size( + x_dim: int, + y_dim: int, + k_x: int, + k_y: int, + fft_size: int | None = None, +) -> int: + """Pick the FFT tile size for a fused 2D conv, or validate an explicit one. + + The tile must satisfy two constraints simultaneously: + + * ``fft_size // 2 >= max(x_dim, y_dim)`` — the kernel's own input cap. + * ``fft_size // 2 >= ceil(k / 2)`` per axis — headroom for the top/left + pre-pad of ``fft_size // 2 - k // 2`` that realigns the crop with the + torch reference (see the section header). + + Args: + x_dim: Input height ``X``. + y_dim: Input width ``Y``. + k_x: Kernel height ``K_x``. + k_y: Kernel width ``K_y``. + fft_size: Explicit tile size to validate. When ``None``, the smallest + admissible power-of-two tile is chosen. + + Returns: + A tile size drawn from ``(8, 16, 32, 64, 128)``. + + Raises: + ValueError: If no admissible tile exists (shape too large for the fused + kernel), or if an explicit ``fft_size`` is not in the supported set + or is too small for the given shapes. + """ + # Half-extent each of the four dims needs the tile to cover. + needed_half = max(x_dim, y_dim, -(-k_x // 2), -(-k_y // 2)) + + if fft_size is None: + fft_size = 2 * _next_pow2(needed_half) + if fft_size > max(_FUSED_FFT_SIZES): + raise ValueError( + f"Shape is too large for the fused 2D FFT kernel: input {(x_dim, y_dim)} with " + f"kernel {(k_x, k_y)} needs fft_size={fft_size}, but the largest supported tile " + f"is {max(_FUSED_FFT_SIZES)} (spatial dims capped at {_FUSED_MAX_SPATIAL} per axis). " + "Use fft_backend='subq_ops' or 'torch_fft' for larger inputs." + ) + return max(fft_size, min(_FUSED_FFT_SIZES)) + + if fft_size not in _FUSED_FFT_SIZES: + raise ValueError(f"fft_size must be one of {_FUSED_FFT_SIZES}, got {fft_size}.") + if fft_size // 2 < needed_half: + raise ValueError( + f"fft_size={fft_size} is too small for input {(x_dim, y_dim)} with kernel " + f"{(k_x, k_y)}: requires fft_size >= {2 * _next_pow2(needed_half)}." + ) + return fft_size + + +def fused_fftconv2d_supported(x_dim: int, y_dim: int, k_x: int, k_y: int) -> bool: + """Whether the fused 2D kernel can handle this input/kernel shape combination. + + A cheap, exception-free probe over the same rules as + :func:`resolve_fused_fft_size` — useful for backend auto-selection and for + the ``torch.compile`` pattern-match guard. + + Args: + x_dim: Input height ``X``. + y_dim: Input width ``Y``. + k_x: Kernel height ``K_x``. + k_y: Kernel width ``K_y``. + + Returns: + ``True`` if an admissible FFT tile exists, ``False`` otherwise. + """ + try: + resolve_fused_fft_size(x_dim, y_dim, k_x, k_y) + except ValueError: + return False + return True + + +def _fused_conv2d_bhl(x: torch.Tensor, kernel: torch.Tensor, fft_size: int | None) -> torch.Tensor: + """Call the fused CUDA kernel on BHL tensors, in ``x``'s dtype. + + Pre-pads the filter's top/left by ``fft_size // 2 - K // 2`` per axis so the + fused kernel's fixed ``fft_size // 2`` crop lands on the same window the + torch reference produces with its ``K // 2`` crop. + """ + fused_fft_conv2d = _get_fused_fft_conv2d() + + _B, _H, x_dim, y_dim = x.shape + k_x, k_y = kernel.shape[-2], kernel.shape[-1] + fft_size = resolve_fused_fft_size(x_dim, y_dim, k_x, k_y, fft_size) + + pad_x = fft_size // 2 - k_x // 2 + pad_y = fft_size // 2 - k_y // 2 + k = F.pad(kernel, (pad_y, 0, pad_x, 0)) + + return fused_fft_conv2d(x.contiguous(), k.contiguous(), fft_size=fft_size) + + +def fused_fftconv2d_bhl( + x: torch.Tensor, + kernel: torch.Tensor, + shortcut: torch.Tensor | None = None, + fft_size: int | None = None, +) -> torch.Tensor: + """Fused 2D FFT convolution, BHL layout ``[B, H, X, Y]``, native dtype. + + Drop-in replacement for :func:`nvsubquadratic.ops.fftconv.fftconv2d_fp32_bhl` + that runs the CUDA kernel in ``x.dtype`` rather than upcasting to fp32. + Numerically equivalent to the reference up to dtype roundoff (~3e-7 + relative in fp32, ~2e-3 in bf16). + + Args: + x: Input tensor ``[B, H, X, Y]``, with ``X, Y <= 64``. + kernel: Kernel tensor ``[1|B, H, Kx, Ky]``. Cast to ``x.dtype`` if needed. + shortcut: Optional per-channel scale ``[H]``. + fft_size: Optional explicit FFT tile size from ``(8, 16, 32, 64, 128)``. + Defaults to the smallest admissible tile. + + Returns: + Output tensor ``[B, H, X, Y]`` in ``x.dtype``. + + Raises: + ValueError: If the shapes exceed what the fused kernel supports. + ImportError: If ``subquadratic_ops_torch`` is not installed. + """ + _B, H, _X, _Y = x.shape + + y = _fused_conv2d_bhl(x, kernel.to(x.dtype), fft_size) + + if shortcut is not None: + y = y + shortcut.to(x.dtype).view(1, H, 1, 1) * x + + return y + + +def fused_fftconv2d_bhl_w_reshape( + x: torch.Tensor, + kernel: torch.Tensor, + shortcut: torch.Tensor | None = None, + fft_size: int | None = None, +) -> torch.Tensor: + """Fused 2D FFT convolution for BLH inputs ``[B, X, Y, H]``. + + Reshapes to BHL, runs :func:`fused_fftconv2d_bhl`, reshapes back. + """ + x_bhl = rearrange(x, "b x y h -> b h x y") + kernel_bhl = rearrange(kernel, "b x y h -> b h x y") + y_bhl = fused_fftconv2d_bhl(x_bhl, kernel_bhl, shortcut, fft_size) + return rearrange(y_bhl, "b h x y -> b x y h") + + +def fused_fftconv2d_blh( + x: torch.Tensor, + kernel: torch.Tensor, + shortcut: torch.Tensor | None = None, + fft_size: int | None = None, +) -> torch.Tensor: + """Alias for :func:`fused_fftconv2d_bhl_w_reshape`.""" + return fused_fftconv2d_bhl_w_reshape(x, kernel, shortcut, fft_size) + + +def fused_fftconv2d_bhl_chunked( + x: torch.Tensor, + kernel: torch.Tensor, + shortcut: torch.Tensor | None = None, + chunk_size: int | None = None, + fft_size: int | None = None, +) -> torch.Tensor: + """Channel-chunked fused 2D FFT convolution, BHL layout. + + Processes channels in groups of ``chunk_size`` to cap the kernel's working + set. The fused kernel's spatial tile is bounded at 64x64, so its per-call + footprint is already small; chunking mainly helps at very large ``H``. + + Args: + x: Input tensor ``[B, H, X, Y]``. + kernel: Kernel tensor ``[1|B, H, Kx, Ky]``. + shortcut: Optional per-channel scale ``[H]``. + chunk_size: Channels per chunk (default 128). + fft_size: Optional explicit FFT tile size. + + Returns: + Output tensor ``[B, H, X, Y]`` in ``x.dtype``. + """ + if chunk_size is None: + chunk_size = _DEFAULT_CHUNK_SIZE + + _B, H, _X, _Y = x.shape + k = kernel.to(x.dtype) + + chunks = [] + for start in range(0, H, chunk_size): + end = min(start + chunk_size, H) + chunks.append(_fused_conv2d_bhl(x[:, start:end].contiguous(), k[:, start:end].contiguous(), fft_size)) + + y = torch.cat(chunks, dim=1) + + if shortcut is not None: + y = y + shortcut.to(x.dtype).view(1, H, 1, 1) * x + + return y + + +def fused_fftconv2d_bhl_w_reshape_chunked( + x: torch.Tensor, + kernel: torch.Tensor, + shortcut: torch.Tensor | None = None, + chunk_size: int | None = None, + fft_size: int | None = None, +) -> torch.Tensor: + """Channel-chunked fused 2D FFT convolution for BLH inputs ``[B, X, Y, H]``.""" + x_bhl = rearrange(x, "b x y h -> b h x y") + kernel_bhl = rearrange(kernel, "b x y h -> b h x y") + y_bhl = fused_fftconv2d_bhl_chunked(x_bhl, kernel_bhl, shortcut, chunk_size, fft_size) + return rearrange(y_bhl, "b h x y -> b x y h") + + +def fused_fftconv2d_blh_chunked( + x: torch.Tensor, + kernel: torch.Tensor, + shortcut: torch.Tensor | None = None, + chunk_size: int | None = None, + fft_size: int | None = None, +) -> torch.Tensor: + """Alias for :func:`fused_fftconv2d_bhl_w_reshape_chunked`.""" + return fused_fftconv2d_bhl_w_reshape_chunked(x, kernel, shortcut, chunk_size, fft_size) + + # =========================================================================== # 1D causal long FFT conv — wraps subquadratic_ops_torch.fft_causal_conv1d # =========================================================================== diff --git a/nvsubquadratic/ops/fftconv_lowering.py b/nvsubquadratic/ops/fftconv_lowering.py new file mode 100644 index 00000000..b0daa51e --- /dev/null +++ b/nvsubquadratic/ops/fftconv_lowering.py @@ -0,0 +1,533 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""``torch.compile`` lowering that rewrites 2D FFT convs onto the fused CUDA kernel. + +Motivation +---------- +Inductor cannot generate code for complex operators — it emits +``UserWarning: Torchinductor does not support code generation for complex +operators`` and falls back to eager cuFFT for the whole +``rfft2 → multiply → irfft2`` chain. So ``torch.compile`` on its own buys +essentially nothing for :func:`nvsubquadratic.ops.fftconv.fftconv2d_fp32_bhl`. +This pass detects that chain and swaps in +:func:`nvsubquadratic.ops.fftconv_custom.fused_fftconv2d_bhl`, which runs the +whole pipeline as a single cuFFTDx launch in the input dtype. + +It exists so that models already written against ``fft_backend="torch_fft"`` +pick up the kernel without a config change. Setting +``fft_backend="subq_ops_fused"`` explicitly is the more predictable route and +does not need this pass — see :mod:`nvsubquadratic.ops.fftconv_custom`. + +Enabling it +----------- +Pass it to ``torch.compile`` directly:: + + from nvsubquadratic.ops.fftconv_lowering import fused_fftconv2d_options + + compiled = torch.compile(model, options=fused_fftconv2d_options()) + +:func:`fused_fftconv2d_options` is scoped to that one compiled callable. The +:class:`fused_fftconv2d_lowering` context manager does the same thing by +patching global inductor config, for when a trainer or framework owns the +``torch.compile`` call and you cannot pass ``options``. + +Where the pass runs +------------------- +It is registered as an inductor **pre-grad** pass, which operates on the Dynamo +graph before AOTAutograd. That placement matters: the replacement node is a +plain forward call, and autograd for it is derived from the custom op's +registered backward. A post-grad pass would have to rewrite the forward and +the backward subgraphs consistently, which is far more fragile. + +Matched pattern +--------------- +The Dynamo graph for ``fftconv2d_fp32_bhl`` is fully static — the FFT sizes and +crop offsets are Python ints baked in as constants — so the pass can verify the +whole recipe rather than pattern-matching shapes structurally:: + + x ─[.to(float32)]─┐ + ├─ rfft2(s=S, dim=(2,3)) ─┐ + k ─[.to(float32)]─┘ ├─ mul ─ irfft2(s=S, dim=(2,3)) + ┘ + ─ [..., a:a+X, b:b+Y] ─ [.to(orig_dtype)] + +A rewrite only happens when every one of these holds: + +* ``S == (min(X + (Kx+1)//2, 2X), min(Y + (Ky+1)//2, 2Y))`` — the exact padding + recipe from ``fftconv.py``, so an unrelated FFT conv is never touched. +* ``a == Kx // 2`` and ``b == Ky // 2`` — the reference's 'same' crop. +* The shapes fit the fused kernel (spatial <= 64 per axis; see + :func:`~nvsubquadratic.ops.fftconv_custom.resolve_fused_fft_size`). +* The device is CUDA and the compute capability supports the required FFT tile. + +Anything else is left alone, so the fallback is the unmodified eager-cuFFT path. + +.. warning:: + ``fftconv.py`` has a ``COMPILE_COMPATIBLE`` flag that replaces the complex + multiply with a real-valued ``_complex_mul_real``. That variant produces a + different subgraph and is **not** matched; the pass reports it via + :func:`lowering_stats` rather than rewriting it. +""" + +from __future__ import annotations + + +__all__ = [ + "FusedFFTConv2dLowering", + "fused_fftconv2d_lowering", + "fused_fftconv2d_options", + "lowering_stats", + "reset_lowering_stats", +] + +import logging +import operator +from typing import Any + +import torch +from torch._inductor.custom_graph_pass import CustomGraphPass + +from nvsubquadratic.ops.fftconv_custom import ( + fused_fftconv2d_bhl, + load_fused_fft_conv2d, + resolve_fused_fft_size, +) + + +logger = logging.getLogger(__name__) + +# Bumped whenever the matching or replacement logic changes. Inductor mixes this +# into its cache key, so a stale compiled artifact is never reused across a +# change to this file. +_PASS_VERSION = "1" + +# The fused kernel's 128-tile needs more per-block shared memory than SM80/SM86 +# provide (the forward already fails on SM86; the backward fails on both). The +# kernel raises at runtime, which is fine for an explicit backend choice but not +# for a silent rewrite — so gate the largest tile on SM90+. +_FFT_SIZE_128_MIN_ARCH = (9, 0) + +_SUPPORTED_DTYPES = (torch.float32, torch.float16, torch.bfloat16) + +# Rewrite counters, keyed by outcome. Useful in tests and when a user asks "did +# it actually fire?", which is the failure mode a silent pass is prone to. +_STATS: dict[str, int] = {} + + +def lowering_stats() -> dict[str, int]: + """Return a snapshot of rewrite counters since the last reset. + + Keys are ``"rewritten"`` plus one ``"skipped:"`` entry per rejection + cause (e.g. ``"skipped:spatial-too-large"``). A pass that silently does + nothing is hard to debug; this is how you tell "did not match" from + "matched but rejected", and why. + + .. note:: + Counters only advance when the pass actually runs. Inductor caches + compiled artifacts on disk, and a cache hit skips every pre-grad pass — + so empty counters can mean "served from cache" rather than "did not + fire". Set ``torch._inductor.config.force_disable_caches = True`` when + using this to verify behaviour. + + Returns: + A copy of the counter dict. + """ + return dict(_STATS) + + +def reset_lowering_stats() -> None: + """Clear the rewrite counters returned by :func:`lowering_stats`.""" + _STATS.clear() + + +def _bump(key: str) -> None: + _STATS[key] = _STATS.get(key, 0) + 1 + + +def _skip(reason: str) -> None: + _bump(f"skipped:{reason}") + logger.debug("fused_fftconv2d lowering skipped a candidate: %s", reason) + + +def _meta_tensor(node: Any) -> torch.Tensor | None: + """Return the example tensor Dynamo attached to ``node``, if any.""" + if not isinstance(node, torch.fx.Node): + return None + value = node.meta.get("example_value") + if value is None: + value = node.meta.get("val") + return value if isinstance(value, torch.Tensor) else None + + +def _unwrap_cast(node: torch.fx.Node) -> torch.fx.Node: + """Strip a ``.to(float32)`` / ``.float()`` hop, returning the tensor behind it. + + The reference upcasts both operands before the FFT; the fused kernel takes + the original dtype directly, so the cast is exactly what we want to drop. + """ + if node.op == "call_method" and node.target in ("to", "float") and node.args: + source = node.args[0] + if isinstance(source, torch.fx.Node): + return source + if node.op == "call_function" and node.target is torch.Tensor.to and node.args: + source = node.args[0] + if isinstance(source, torch.fx.Node): + return source + return node + + +def _is_fft(node: Any, target) -> bool: + return isinstance(node, torch.fx.Node) and node.op == "call_function" and node.target is target + + +def _fft_kwarg(node: torch.fx.Node, name: str, position: int) -> Any: + """Read an FFT argument that may have been passed positionally or by keyword.""" + if name in node.kwargs: + return node.kwargs[name] + return node.args[position] if len(node.args) > position else None + + +def _as_int_pair(value: Any) -> tuple[int, int] | None: + if isinstance(value, (tuple, list)) and len(value) == 2: + try: + return (int(value[0]), int(value[1])) + except (TypeError, ValueError): + return None + return None + + +def _find_spectrum_operands(irfft: torch.fx.Node) -> tuple[torch.fx.Node, torch.fx.Node] | None: + """Recover the two ``rfft2`` nodes feeding an ``irfft2``. + + Handles both multiply spellings the reference can produce: the default + in-place ``fft_x.mul_(fft_kernel)`` (where ``irfft2``'s own input is the + *input* spectrum, mutated in place) and an out-of-place ``a * b``. + """ + spectrum = irfft.args[0] if irfft.args else irfft.kwargs.get("input") + if not isinstance(spectrum, torch.fx.Node): + return None + + # In-place: irfft2 reads the mutated input spectrum directly. + if _is_fft(spectrum, torch.fft.rfft2): + for user in spectrum.users: + if user.op == "call_method" and user.target == "mul_" and user.args and user.args[0] is spectrum: + other = user.args[1] if len(user.args) > 1 else None + if _is_fft(other, torch.fft.rfft2): + return spectrum, other + return None + + # Out-of-place: irfft2 reads the product node. + is_mul = (spectrum.op == "call_function" and spectrum.target in (operator.mul, torch.mul)) or ( + spectrum.op == "call_method" and spectrum.target == "mul" + ) + if is_mul and len(spectrum.args) == 2: + lhs, rhs = spectrum.args + if _is_fft(lhs, torch.fft.rfft2) and _is_fft(rhs, torch.fft.rfft2): + return lhs, rhs + return None + + +def _parse_crop(irfft: torch.fx.Node) -> tuple[torch.fx.Node, int, int, int, int] | None: + """Parse the ``[..., a:a+X, b:b+Y]`` crop that consumes ``irfft2``. + + Returns: + ``(getitem_node, a, X, b, Y)``, or ``None`` if the sole consumer is not + a trailing-two-axis slice of that exact form. + """ + users = list(irfft.users) + if len(users) != 1: + return None + getitem = users[0] + if getitem.op != "call_function" or getitem.target is not operator.getitem: + return None + + index = getitem.args[1] + if not isinstance(index, tuple) or len(index) != 3: + return None + ellipsis, slice_x, slice_y = index + if ellipsis is not Ellipsis or not isinstance(slice_x, slice) or not isinstance(slice_y, slice): + return None + if slice_x.step not in (None, 1) or slice_y.step not in (None, 1): + return None + if any(v is None for v in (slice_x.start, slice_x.stop, slice_y.start, slice_y.stop)): + return None + + start_x, start_y = int(slice_x.start), int(slice_y.start) + return getitem, start_x, int(slice_x.stop) - start_x, start_y, int(slice_y.stop) - start_y + + +def _arch_supports(device: torch.device, fft_size: int) -> bool: + """Whether ``device`` can run the fused kernel at this FFT tile size.""" + if device.type != "cuda" or not torch.cuda.is_available(): + return False + if fft_size < 128: + return True + return torch.cuda.get_device_capability(device.index) >= _FFT_SIZE_128_MIN_ARCH + + +def _try_rewrite(graph: torch.fx.Graph, irfft: torch.fx.Node, allow_reduced_precision: bool) -> bool: + """Attempt to replace one ``rfft2/mul/irfft2/crop`` chain. Returns whether it fired.""" + if _as_int_pair(_fft_kwarg(irfft, "dim", 2)) != (2, 3): + _skip("irfft-dim-not-2-3") + return False + fft_shape = _as_int_pair(_fft_kwarg(irfft, "s", 1)) + if fft_shape is None: + _skip("irfft-dynamic-shape") + return False + + operands = _find_spectrum_operands(irfft) + if operands is None: + _skip("no-rfft2-product") + return False + + crop = _parse_crop(irfft) + if crop is None: + _skip("no-same-crop") + return False + getitem, start_x, out_x, start_y, out_y = crop + + # Identify which spectrum is the input: its spatial extent equals the crop + # length. (For the in-place spelling the order is already known, but keying + # off shapes covers the out-of-place spelling too.) + x_rfft = kernel_rfft = None + for candidate, other in (operands, operands[::-1]): + meta = _meta_tensor(candidate.args[0] if candidate.args else None) + if meta is not None and meta.ndim == 4 and tuple(meta.shape[-2:]) == (out_x, out_y): + x_rfft, kernel_rfft = candidate, other + break + if x_rfft is None: + _skip("operand-roles-ambiguous") + return False + + for rfft in (x_rfft, kernel_rfft): + if _as_int_pair(_fft_kwarg(rfft, "s", 1)) != fft_shape or _as_int_pair(_fft_kwarg(rfft, "dim", 2)) != (2, 3): + _skip("rfft-args-mismatch") + return False + + x_node = _unwrap_cast(x_rfft.args[0]) + kernel_node = _unwrap_cast(kernel_rfft.args[0]) + x_meta, kernel_meta = _meta_tensor(x_node), _meta_tensor(kernel_node) + if x_meta is None or kernel_meta is None: + _skip("missing-shape-metadata") + return False + if x_meta.ndim != 4 or kernel_meta.ndim != 4: + _skip("not-4d") + return False + + dim_x, dim_y = int(x_meta.shape[-2]), int(x_meta.shape[-1]) + k_x, k_y = int(kernel_meta.shape[-2]), int(kernel_meta.shape[-1]) + + # The recipe check. This is what keeps the pass from touching an unrelated + # FFT convolution that merely has the same node types. + expected_shape = (min(dim_x + (k_x + 1) // 2, 2 * dim_x), min(dim_y + (k_y + 1) // 2, 2 * dim_y)) + if fft_shape != expected_shape or (start_x, start_y) != (k_x // 2, k_y // 2): + _skip("not-the-reference-recipe") + return False + if (out_x, out_y) != (dim_x, dim_y): + _skip("crop-is-not-same-size") + return False + if int(kernel_meta.shape[1]) != int(x_meta.shape[1]) or int(kernel_meta.shape[0]) not in (1, int(x_meta.shape[0])): + _skip("channel-or-batch-mismatch") + return False + + if x_meta.dtype not in _SUPPORTED_DTYPES: + _skip("unsupported-dtype") + return False + if not allow_reduced_precision and x_meta.dtype is not torch.float32: + _skip("reduced-precision-disabled") + return False + + try: + fft_size = resolve_fused_fft_size(dim_x, dim_y, k_x, k_y) + except ValueError: + _skip("spatial-too-large") + return False + if not _arch_supports(x_meta.device, fft_size): + _skip("arch-unsupported") + return False + + # The reference casts its fp32 result back to the input dtype. The fused + # kernel already returns the input dtype, so when that cast is present it is + # the node to replace — otherwise the rewrite would leave a dangling cast. + replace_target = getitem + consumers = list(getitem.users) + if len(consumers) == 1 and consumers[0].op == "call_method" and consumers[0].target == "to": + cast = consumers[0] + cast_meta = _meta_tensor(cast) + if cast_meta is not None and cast_meta.dtype == x_meta.dtype: + replace_target = cast + + with graph.inserting_before(replace_target): + fused = graph.call_function(fused_fftconv2d_bhl, args=(x_node, kernel_node)) + fused.meta.update(replace_target.meta) + replace_target.replace_all_uses_with(fused) + + _bump("rewritten") + logger.debug( + "fused_fftconv2d lowering rewrote a %sx%s conv (kernel %sx%s, fft_size=%s, dtype=%s)", + dim_x, + dim_y, + k_x, + k_y, + fft_size, + x_meta.dtype, + ) + return True + + +class FusedFFTConv2dLowering(CustomGraphPass): + """Inductor pre-grad pass replacing the 2D FFT-conv chain with the fused kernel. + + Install it with :func:`fused_fftconv2d_lowering`, which handles the inductor + config plumbing and restores any previously registered pass on exit. + + Args: + allow_reduced_precision: When ``True`` (default) fp16/bf16 graphs are + rewritten too, which changes the convolution from fp32-internal to + native-dtype — the main source of the speedup, and a real numerics + change (~2e-3 normwise in bf16). Set ``False`` to restrict the + rewrite to fp32 graphs, where it is numerically neutral. + """ + + def __init__(self, allow_reduced_precision: bool = True): + """Build the pass and force the fused kernel's operators to register. + + Args: + allow_reduced_precision: See the class docstring. + + Raises: + ImportError: If ``subquadratic_ops_torch`` is not installed. + """ + # Registration has to happen up front: a compiled artifact restored from + # inductor's on-disk cache skips this pass entirely and calls + # torch.ops.subquadratic_ops_torch.* directly, so relying on the + # wrappers' lazy import would fail with an op-namespace AttributeError + # on that first (cache-hit) call. Doing it here also turns a missing + # install into a clear ImportError at setup time. + load_fused_fft_conv2d() + self.allow_reduced_precision = allow_reduced_precision + + def __call__(self, graph: torch.fx.Graph) -> torch.fx.Graph: + """Rewrite every matching FFT-conv chain in ``graph``, in place.""" + candidates = [n for n in graph.nodes if _is_fft(n, torch.fft.irfft2)] + if not candidates: + return graph + + if any(_try_rewrite(graph, node, self.allow_reduced_precision) for node in candidates): + # Drops the now-unused rfft2/mul/irfft2/crop nodes. The in-place + # ``mul_`` is dead too: its only consumer was the irfft2 we replaced. + graph.eliminate_dead_code() + graph.lint() + return graph + + def uuid(self) -> str: + """Cache key contribution, so inductor never reuses artifacts across changes here.""" + return f"nvsubquadratic-fused-fftconv2d-lowering-v{_PASS_VERSION}-rp{int(self.allow_reduced_precision)}" + + +def fused_fftconv2d_options(allow_reduced_precision: bool = True) -> dict[str, Any]: + """Return a ``torch.compile(options=...)`` dict that enables the lowering. + + This is the preferred way to enable the pass. Inductor's ``options`` are + applied as a config patch scoped to that one compiled callable, so unlike + :class:`fused_fftconv2d_lowering` it neither mutates global inductor state + nor conflicts with a ``pre_grad_custom_pass`` registered elsewhere. + + Args: + allow_reduced_precision: See :class:`FusedFFTConv2dLowering`. + + Returns: + A dict to pass as ``torch.compile(..., options=...)``. + + Raises: + ImportError: If ``subquadratic_ops_torch`` is not installed. + + Example: + >>> import torch + >>> from nvsubquadratic.ops.fftconv_lowering import fused_fftconv2d_options + >>> compiled = torch.compile(model, options=fused_fftconv2d_options()) # doctest: +SKIP + >>> out = compiled(x) # doctest: +SKIP + """ + return {"pre_grad_custom_pass": FusedFFTConv2dLowering(allow_reduced_precision)} + + +class fused_fftconv2d_lowering: # Lowercase: only ever used as a context manager. + """Context manager that installs :class:`FusedFFTConv2dLowering` globally. + + Prefer :func:`fused_fftconv2d_options` when you control the + ``torch.compile`` call — it is scoped to a single callable instead of + patching global inductor config. Reach for this context manager when + something else compiles the model for you (a trainer, a framework entry + point) and you cannot pass ``options``. + + ``torch.compile`` caches compiled artifacts, so enter this **before** the + first compiled call on the model you want rewritten; a function already + compiled without the pass keeps its cached code. + + Args: + allow_reduced_precision: See :class:`FusedFFTConv2dLowering`. + + Example: + >>> import torch + >>> from nvsubquadratic.ops.fftconv_lowering import ( + ... fused_fftconv2d_lowering, + ... lowering_stats, + ... ) + >>> with fused_fftconv2d_lowering(): # doctest: +SKIP + ... compiled = torch.compile(model) + ... out = compiled(x) + >>> lowering_stats().get("rewritten", 0) # doctest: +SKIP + 1 + + .. note:: + Inductor exposes a single ``pre_grad_custom_pass`` slot. Any pass already + registered is saved and restored on exit, but the two do not compose + while this one is active. + """ + + def __init__(self, allow_reduced_precision: bool = True): + """Build the pass eagerly, so a missing install fails here, not at compile time. + + Args: + allow_reduced_precision: See :class:`FusedFFTConv2dLowering`. + """ + self.pass_ = FusedFFTConv2dLowering(allow_reduced_precision) + self._previous: Any = None + + def __enter__(self) -> FusedFFTConv2dLowering: + """Install the pass, saving whatever was registered before. + + Returns: + The installed :class:`FusedFFTConv2dLowering` instance. + """ + from torch._inductor import config as inductor_config + + self._previous = inductor_config.pre_grad_custom_pass + if self._previous is not None: + logger.warning( + "Replacing an already-registered inductor pre_grad_custom_pass (%r) for the " + "duration of fused_fftconv2d_lowering; it will be restored on exit.", + self._previous, + ) + inductor_config.pre_grad_custom_pass = self.pass_ + return self.pass_ + + def __exit__(self, *exc_info: object) -> None: + """Restore the previously registered pass, including on an exception.""" + from torch._inductor import config as inductor_config + + inductor_config.pre_grad_custom_pass = self._previous + self._previous = None diff --git a/pyproject.toml b/pyproject.toml index 7691a084..c2f94f57 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,8 +50,9 @@ classifiers = [ # project's CPU CI). The importable operators only need torch/einops/omegaconf/numpy # at runtime; the remaining deps power the training/data/experiment code. dependencies = [ - "torch>=2.10.0,<2.11.0", - "torchvision>=0.25.0,<0.26.0", + # CUDA 13.2 (cu132) wheels start at torch 2.12.0 / torchvision 0.27.0. + "torch>=2.12.0,<2.13.0", + "torchvision>=0.27.0,<0.28.0", "einops>=0.6.0", "numpy>=1.21.0", "pyarrow>=14.0.0,<20.0.0", # Constrained to be compatible with cudf/pylibcudf in NGC containers @@ -79,7 +80,9 @@ Issues = "https://github.com/NVIDIA-BioNeMo/nvSubquadratic/issues" # (e.g. a downstream project's CPU CI). The operators import it lazily and the # default fft_backend="torch_fft" path runs without it; install for the fast path: # pip install nvsubquadratic[cuda] -cuda = ["subquadratic-ops-torch-cu12>=0.2.0"] +# >=0.2.2 for fused_fft_conv2d, which backs fft_backend="subq_ops_fused" and the +# torch.compile lowering. fft_backend="subq_ops" itself still works on 0.2.0. +cuda = ["subquadratic-ops-torch-cu13>=0.2.2"] # QuACK fused RMSNorm kernel (Hopper/Blackwell only: H100, B200, B300). # On Ampere and older GPUs RMSNorm falls back to pure PyTorch automatically. @@ -87,7 +90,7 @@ quack = ["quack-kernels"] # NVIDIA DALI GPU data-loading pipeline, used by the ImageNet / The Well # examples. Not required by the importable library; ~400 MB download. -dali = ["nvidia-dali-cuda120==1.53.0"] +dali = ["nvidia-dali-cuda130==1.53.0"] # Megatron parallel state, needed only for context-parallel / distributed # training via nvsubquadratic.parallel.utils.init_parallel_state. The operators diff --git a/tests/conftest.py b/tests/conftest.py index 5279b710..3dcf688c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -24,17 +24,30 @@ # --------------------------------------------------------------------------- +# The kernels ship as one distribution per CUDA major version. Query every +# known name: pinning only one of them silently reports (0, 0, 0) on a machine +# that has the other installed, which turns the gates below into blanket xfails. +_SUBQ_OPS_DISTRIBUTIONS = ("subquadratic-ops-torch-cu13", "subquadratic-ops-torch-cu12") + + def _subq_ops_version() -> tuple[int, ...]: - """Return the installed subquadratic-ops-torch-cu12 version as an int tuple.""" - try: - from importlib.metadata import version + """Return the installed subquadratic-ops-torch version as an int tuple. + + Returns ``(0, 0, 0)`` when no known distribution is installed. + """ + from importlib.metadata import version - return tuple(int(x) for x in version("subquadratic-ops-torch-cu12").split(".")[:3]) - except Exception: - return (0, 0, 0) + for dist in _SUBQ_OPS_DISTRIBUTIONS: + try: + return tuple(int(x) for x in version(dist).split(".")[:3]) + except Exception: + continue + return (0, 0, 0) _SUBQ_OPS_MIN_VERSION = (0, 2, 0) +# fused_fft_conv2d (the native-dtype single-launch 2D kernel) landed in 0.2.2. +_SUBQ_OPS_FUSED_MIN_VERSION = (0, 2, 2) _subq_installed = _subq_ops_version() requires_subq_ops_v2 = pytest.mark.xfail( @@ -46,6 +59,17 @@ def _subq_ops_version() -> tuple[int, ...]: strict=False, ) +# Unlike requires_subq_ops_v2 this skips rather than xfails: the fused kernel is +# a strictly newer addition, so an older-but-working install is "not applicable" +# rather than "expected to fail". +requires_subq_ops_fused = pytest.mark.skipif( + _subq_installed < _SUBQ_OPS_FUSED_MIN_VERSION, + reason=( + f"subquadratic_ops_torch >= {'.'.join(str(x) for x in _SUBQ_OPS_FUSED_MIN_VERSION)} " + f"required for fused_fft_conv2d (installed: {'.'.join(str(x) for x in _subq_installed)})" + ), +) + @pytest.fixture def device(): diff --git a/tests/modules/test_ckconv_nd_subq_fused.py b/tests/modules/test_ckconv_nd_subq_fused.py new file mode 100644 index 00000000..087ca8ad --- /dev/null +++ b/tests/modules/test_ckconv_nd_subq_fused.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration tests for CKConvND with fft_backend='subq_ops_fused'. + +Companion to ``test_ckconv_nd_subq.py`` (which covers ``fft_backend="subq_ops"``). +The fused backend adds two things that backend does not have: native +fp16/bf16 execution, and a hard 64x64 spatial cap enforced at forward time +rather than at construction. + +Usage (requires GPU — run inside SLURM): + srun --gres=gpu:1 -c 16 --partition low \\ + conda run -n nv-subq python -m pytest tests/modules/test_ckconv_nd_subq_fused.py -v -o addopts="" +""" + +import pytest +import torch + +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.ckconv_nd import CKConvND +from nvsubquadratic.modules.kernels_nd import SIRENKernelND +from tests.conftest import requires_subq_ops_fused + + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + +pytestmark = [requires_subq_ops_fused, requires_cuda] + +HIDDEN_DIM = 32 +SPATIAL = 8 + +# Normwise relative-error budgets against the fp32 torch_fft backend. +L2_TOL = {torch.float32: 1e-6, torch.float16: 1e-3, torch.bfloat16: 8e-3} +L2_TOL_GRAD = {torch.float32: 1e-5, torch.float16: 3e-3, torch.bfloat16: 2e-2} + + +def _l2_rel(pred, ref): + pred64, ref64 = pred.double(), ref.double() + den = ref64.norm() + return ((pred64 - ref64).norm() / den).item() if den > 0 else (pred64 - ref64).norm().item() + + +def _assert_l2_close(pred, ref, dtype, tol_table=None, name=""): + tol = (tol_table or L2_TOL)[dtype] + rel = _l2_rel(pred, ref) + assert rel < tol, f"{name} L2 rel error {rel:.3e} exceeds tol {tol:.1e}" + + +def _make_ckconv(grid_type, fft_backend, use_chunked=False, spatial=SPATIAL, data_dim=2, **kw): + """Build a CKConvND with a small SIREN kernel, mirroring test_ckconv_nd_subq.py.""" + kernel_cfg = LazyConfig(SIRENKernelND)( + data_dim=data_dim, + out_dim=HIDDEN_DIM, + mlp_hidden_dim=16, + num_layers=2, + embedding_dim=16, + omega_0=10.0, + L_cache=spatial, + use_bias=True, + ) + kw.setdefault("fft_padding", "zero") + return CKConvND( + data_dim=data_dim, + hidden_dim=HIDDEN_DIM, + kernel_cfg=kernel_cfg, + mask_cfg=LazyConfig(torch.nn.Identity)(), + grid_type=grid_type, + fft_backend=fft_backend, + use_chunked_fftconv=use_chunked, + **kw, + ) + + +def _paired_models(grid_type, use_chunked=False, spatial=SPATIAL): + """Build torch_fft and subq_ops_fused models sharing identical weights.""" + reference = _make_ckconv(grid_type, "torch_fft", spatial=spatial).cuda() + fused = _make_ckconv(grid_type, "subq_ops_fused", use_chunked, spatial=spatial).cuda() + fused.load_state_dict(reference.state_dict()) + return reference, fused + + +# --------------------------------------------------------------------------- +# Forward / backward equivalence with torch_fft +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("grid_type", ["single", "double"]) +@pytest.mark.parametrize("use_chunked", [False, True]) +class TestForwardMatchesTorchFft: + def test_blh_layout(self, grid_type, use_chunked): + torch.manual_seed(42) + reference, fused = _paired_models(grid_type, use_chunked) + x = torch.randn(2, SPATIAL, SPATIAL, HIDDEN_DIM, device="cuda") + _assert_l2_close(fused(x), reference(x), torch.float32) + + def test_bhl_layout(self, grid_type, use_chunked): + torch.manual_seed(42) + reference, fused = _paired_models(grid_type, use_chunked) + x = torch.randn(2, HIDDEN_DIM, SPATIAL, SPATIAL, device="cuda") + _assert_l2_close(fused(x, is_bhl_input=True), reference(x, is_bhl_input=True), torch.float32) + + +def _cast_params(model, dtype): + """Cast parameters only, leaving buffers alone. + + ``model.to(dtype)`` would also cast ``SIRENKernelND.grid_cache``, which the + kernel asserts must stay fp32 (see the TODO at ``kernels_nd.py``) — the + positional embedding casts the grid to the weight dtype internally instead. + The buffer is non-persistent, so skipping it keeps ``state_dict`` parity. + """ + for param in model.parameters(): + param.data = param.data.to(dtype) + return model + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_forward_native_dtype(dtype): + """The fused backend runs natively in the input dtype. + + Both backends are cast to ``dtype`` so the SIREN generates the *same* + low-precision kernel for each; the only remaining difference is the + convolution itself, which torch_fft still evaluates in fp32 while the fused + kernel evaluates in ``dtype``. Comparing against an fp32 reference model + instead would fold in the SIREN's own low-precision error (~3e-2 in bf16), + which has nothing to do with this backend. + """ + torch.manual_seed(42) + reference, fused = _paired_models("double") + _cast_params(reference, dtype) + _cast_params(fused, dtype) + x = torch.randn(2, SPATIAL, SPATIAL, HIDDEN_DIM, device="cuda", dtype=dtype) + + out = fused(x) + assert out.dtype == dtype + _assert_l2_close(out, reference(x), dtype) + + +@pytest.mark.parametrize("grid_type", ["single", "double"]) +def test_backward_matches_torch_fft(grid_type): + """Gradients flow to the SIREN kernel parameters and match torch_fft.""" + torch.manual_seed(42) + reference, fused = _paired_models(grid_type) + x = torch.randn(2, SPATIAL, SPATIAL, HIDDEN_DIM, device="cuda") + grad_out = torch.randn(2, SPATIAL, SPATIAL, HIDDEN_DIM, device="cuda") + + reference(x).backward(grad_out) + fused(x).backward(grad_out) + + ref_grads = dict(reference.named_parameters()) + assert any(p.grad is not None for p in ref_grads.values()) + for name, param in fused.named_parameters(): + want = ref_grads[name].grad + if want is None: + assert param.grad is None, f"{name}: fused produced a grad where torch_fft did not" + continue + _assert_l2_close(param.grad, want, torch.float32, tol_table=L2_TOL_GRAD, name=name) + + +@pytest.mark.parametrize("spatial", [7, 16, 32, 64]) +def test_spatial_sizes_within_cap(spatial): + """Every supported spatial size, including the 64 boundary and a non-power-of-2.""" + torch.manual_seed(42) + reference, fused = _paired_models("double", spatial=spatial) + x = torch.randn(1, spatial, spatial, HIDDEN_DIM, device="cuda") + _assert_l2_close(fused(x), reference(x), torch.float32) + + +# --------------------------------------------------------------------------- +# Constraint validation +# --------------------------------------------------------------------------- + + +class TestRejectsInvalidConfigs: + """Unsupported configurations fail at construction, except the input-size cap.""" + + @pytest.mark.parametrize("data_dim", [1, 3]) + def test_rejects_non_2d(self, data_dim): + with pytest.raises(AssertionError, match="only supports data_dim=2"): + _make_ckconv("double", "subq_ops_fused", data_dim=data_dim) + + def test_rejects_causal(self): + # data_dim=1 is rejected first, so causal must be probed on the 2D path, + # where CKConvND's own "causal is 1D only" assert fires first. + with pytest.raises(AssertionError, match="Causal CKConvND only supports 1D"): + _make_ckconv("double", "subq_ops_fused", is_causal=True) + + def test_rejects_circular_padding(self): + with pytest.raises(AssertionError, match="only supports zero-padded"): + _make_ckconv("single", "subq_ops_fused", fft_padding="circular") + + def test_rejects_per_axis_padding(self): + with pytest.raises(ValueError, match="does not support a per-axis fft_padding"): + _make_ckconv(None, "subq_ops_fused", fft_padding=["circular", "zero"]) + + def test_rejects_unknown_backend(self): + with pytest.raises(AssertionError, match="Invalid fft_backend"): + _make_ckconv("double", "subq_ops_fusedd") + + def test_oversized_input_raises_at_forward(self): + """The 64x64 cap depends on the input, so it is enforced on the first call.""" + torch.manual_seed(42) + model = _make_ckconv("double", "subq_ops_fused", spatial=65).cuda() + x = torch.randn(1, 65, 65, HIDDEN_DIM, device="cuda") + with pytest.raises(ValueError, match="too large for the fused 2D FFT kernel"): + model(x) + + +def test_extra_repr_reports_backend(): + assert "fft_backend='subq_ops_fused'" in repr(_make_ckconv("double", "subq_ops_fused")) diff --git a/tests/ops/test_fftconv_lowering.py b/tests/ops/test_fftconv_lowering.py new file mode 100644 index 00000000..634fcbc0 --- /dev/null +++ b/tests/ops/test_fftconv_lowering.py @@ -0,0 +1,352 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the torch.compile lowering in ``nvsubquadratic.ops.fftconv_lowering``. + +A silent graph pass has two distinct failure modes, and this file covers both: + + - **False negative** — it quietly does not fire, and the user sees no speedup + and no error. Every positive test therefore asserts on + ``lowering_stats()["rewritten"]``, not just on numerics. + - **False positive** — it rewrites a subgraph that is not actually the + reference recipe, silently changing results. The guard tests assert the + specific skip reason *and* bit-exact agreement with eager, so a rewrite + that slipped through would show up as a numeric difference. + +Usage (requires GPU — run inside SLURM): + srun --gres=gpu:1 -c 16 --partition low \\ + conda run -n nv-subq python -m pytest tests/ops/test_fftconv_lowering.py -v -o addopts="" +""" + +import pytest +import torch + +from nvsubquadratic.ops.fftconv import fftconv2d_fp32_bhl +from nvsubquadratic.ops.fftconv_lowering import ( + FusedFFTConv2dLowering, + fused_fftconv2d_lowering, + fused_fftconv2d_options, + lowering_stats, + reset_lowering_stats, +) +from tests.conftest import requires_subq_ops_fused + + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + +pytestmark = [requires_subq_ops_fused, requires_cuda] + +L2_TOL = {torch.float32: 1e-6, torch.float16: 1e-3, torch.bfloat16: 8e-3} +L2_TOL_GRAD = {torch.float32: 1e-6, torch.float16: 2e-3, torch.bfloat16: 1.5e-2} + + +def _l2_rel(pred, ref): + pred64, ref64 = pred.double(), ref.double() + den = ref64.norm() + return ((pred64 - ref64).norm() / den).item() if den > 0 else (pred64 - ref64).norm().item() + + +def _assert_l2_close(pred, ref, dtype, tol_table=None, name=""): + tol = (tol_table or L2_TOL)[dtype] + rel = _l2_rel(pred, ref) + assert rel < tol, f"{name} L2 rel error {rel:.3e} exceeds tol {tol:.1e}" + + +@pytest.fixture(autouse=True) +def _clean_compile_state(): + """Give every test a fresh dynamo cache, no inductor caching, and zeroed counters. + + Disabling the caches is load-bearing, not hygiene. Inductor's FX graph cache + is on *disk* and survives across processes; on a cache hit it skips every + pre-grad pass, so the counters these tests assert on would stay empty and + the assertions would be measuring the cache rather than the pass. + """ + from torch._inductor import config as inductor_config + + previous = inductor_config.force_disable_caches + inductor_config.force_disable_caches = True + torch._dynamo.reset() + reset_lowering_stats() + try: + yield + finally: + torch._dynamo.reset() + reset_lowering_stats() + inductor_config.force_disable_caches = previous + + +def _compile_with_lowering(fn, *args, allow_reduced_precision=True): + """Compile via the ``options=`` route — the primary way to enable the pass.""" + options = fused_fftconv2d_options(allow_reduced_precision=allow_reduced_precision) + return torch.compile(fn, fullgraph=True, options=options)(*args) + + +def _bhl_inputs(spatial=16, kernel_size=31, dtype=torch.float32, channels=8, batch=2, kernel_batch=1): + torch.manual_seed(0) + x = torch.randn(batch, channels, spatial, spatial, device="cuda", dtype=dtype) + kernel = (torch.randn(kernel_batch, channels, kernel_size, kernel_size, device="cuda") * 0.05).to(dtype) + shortcut = torch.randn(channels, device="cuda", dtype=dtype) + return x, kernel, shortcut + + +# --------------------------------------------------------------------------- +# The pass fires and preserves semantics +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_rewrites_and_matches_eager(dtype): + x, kernel, shortcut = _bhl_inputs(dtype=dtype) + expected = fftconv2d_fp32_bhl(x, kernel, shortcut) + + out = _compile_with_lowering(fftconv2d_fp32_bhl, x, kernel, shortcut) + + assert lowering_stats().get("rewritten") == 1 + _assert_l2_close(out, expected, dtype) + + +def test_rewrites_without_shortcut(): + """The shortcut add is optional and sits outside the matched chain.""" + x, kernel, _ = _bhl_inputs() + expected = fftconv2d_fp32_bhl(x, kernel) + + out = _compile_with_lowering(fftconv2d_fp32_bhl, x, kernel) + + assert lowering_stats().get("rewritten") == 1 + _assert_l2_close(out, expected, torch.float32) + + +@pytest.mark.parametrize("spatial", [8, 16, 32, 64]) +def test_rewrites_across_supported_spatial_sizes(spatial): + x, kernel, shortcut = _bhl_inputs(spatial=spatial, kernel_size=2 * spatial - 1, channels=4, batch=1) + expected = fftconv2d_fp32_bhl(x, kernel, shortcut) + + out = _compile_with_lowering(fftconv2d_fp32_bhl, x, kernel, shortcut) + + assert lowering_stats().get("rewritten") == 1 + _assert_l2_close(out, expected, torch.float32) + + +def test_rewrites_film_per_sample_kernel(): + x, kernel, shortcut = _bhl_inputs(kernel_batch=2) + expected = fftconv2d_fp32_bhl(x, kernel, shortcut) + + out = _compile_with_lowering(fftconv2d_fp32_bhl, x, kernel, shortcut) + + assert lowering_stats().get("rewritten") == 1 + _assert_l2_close(out, expected, torch.float32) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_backward_through_lowered_graph(dtype): + """Autograd is derived from the custom op's registered backward. + + This is the payoff for running as a pre-grad pass rather than post-grad. + """ + x0, k0, s0 = _bhl_inputs() + grad_out = torch.randn_like(x0) + + def run(fn, cast_to): + x = x0.to(cast_to).detach().requires_grad_() + k = k0.to(cast_to).detach().requires_grad_() + s = s0.to(cast_to).detach().requires_grad_() + fn(x, k, s).backward(grad_out.to(cast_to)) + return x.grad, k.grad, s.grad + + expected = run(fftconv2d_fp32_bhl, torch.float32) + compiled = torch.compile(fftconv2d_fp32_bhl, fullgraph=True, options=fused_fftconv2d_options()) + got = run(compiled, dtype) + + assert lowering_stats().get("rewritten") == 1 + for actual, want, name in zip(got, expected, ("grad_x", "grad_kernel", "grad_shortcut")): + _assert_l2_close(actual, want, dtype, tol_table=L2_TOL_GRAD, name=name) + + +# --------------------------------------------------------------------------- +# Guards — the pass must decline, and declining must be a no-op +# --------------------------------------------------------------------------- + + +def _assert_declined(fn, args, reason): + """Assert the pass skipped for ``reason`` and left the eager path in place. + + The counters are the exact proof that no rewrite happened. The numeric check + is a second line of defence and is deliberately *not* bit-exact: inductor + perturbs results on its own through fusion and reduction ordering, so + compiled-without-our-pass differs from eager at the 1e-6 (fp32) level + regardless. What it does rule out is a rewrite having slipped through, which + would show up far above this budget. + """ + expected = fn(*args) + out = _compile_with_lowering(fn, *args) + + stats = lowering_stats() + assert stats.get("rewritten", 0) == 0, f"expected no rewrite, got {stats}" + assert stats.get(f"skipped:{reason}") == 1, f"expected skip reason {reason!r}, got {stats}" + _assert_l2_close(out, expected, out.dtype, name="declined-path") + + +@pytest.mark.parametrize("spatial", [65, 96, 128]) +def test_declines_oversized_spatial(spatial): + """Beyond the 64x64 cap the pass must leave the eager cuFFT path alone.""" + x, kernel, shortcut = _bhl_inputs(spatial=spatial, kernel_size=2 * spatial - 1, channels=4, batch=1) + _assert_declined(fftconv2d_fp32_bhl, (x, kernel, shortcut), "spatial-too-large") + + +def test_declines_on_cpu(): + torch.manual_seed(0) + x = torch.randn(1, 4, 16, 16) + kernel = torch.randn(1, 4, 31, 31) * 0.05 + _assert_declined(fftconv2d_fp32_bhl, (x, kernel), "arch-unsupported") + + +def test_declines_a_different_crop_convention(): + """A same-shaped FFT conv that is not the reference recipe must be untouched. + + This is the false-positive guard: identical node types, different crop. + """ + + def other_crop(x, kernel): + spec = torch.fft.rfft2(x, s=(32, 32), dim=(2, 3)) * torch.fft.rfft2(kernel, s=(32, 32), dim=(2, 3)) + return torch.fft.irfft2(spec, s=(32, 32), dim=(2, 3))[..., 0:16, 0:16] + + torch.manual_seed(0) + x = torch.randn(1, 4, 16, 16, device="cuda") + kernel = torch.randn(1, 4, 31, 31, device="cuda") * 0.05 + _assert_declined(other_crop, (x, kernel), "not-the-reference-recipe") + + +def test_declines_a_different_fft_size(): + """Right crop offset, wrong padding recipe — still not ours.""" + + def oversized_fft(x, kernel): + spec = torch.fft.rfft2(x, s=(64, 64), dim=(2, 3)) * torch.fft.rfft2(kernel, s=(64, 64), dim=(2, 3)) + return torch.fft.irfft2(spec, s=(64, 64), dim=(2, 3))[..., 15:31, 15:31] + + torch.manual_seed(0) + x = torch.randn(1, 4, 16, 16, device="cuda") + kernel = torch.randn(1, 4, 31, 31, device="cuda") * 0.05 + _assert_declined(oversized_fft, (x, kernel), "not-the-reference-recipe") + + +def test_declines_reduced_precision_when_disabled(): + """``allow_reduced_precision=False`` keeps bf16 graphs on the fp32-internal path.""" + x, kernel, shortcut = _bhl_inputs(dtype=torch.bfloat16) + expected = fftconv2d_fp32_bhl(x, kernel, shortcut) + + out = _compile_with_lowering(fftconv2d_fp32_bhl, x, kernel, shortcut, allow_reduced_precision=False) + + stats = lowering_stats() + assert stats.get("rewritten", 0) == 0 + assert stats.get("skipped:reduced-precision-disabled") == 1 + # The counters above are the real assertion. Numerics cannot distinguish the + # two paths here: both round the result to bf16 at the end, and that final + # cast alone accounts for ~3e-3 normwise — the same order as the native-bf16 + # kernel's own error. + _assert_l2_close(out, expected, torch.bfloat16) + + +def test_fp32_still_rewritten_when_reduced_precision_disabled(): + """The flag gates only reduced precision, where the rewrite changes numerics.""" + x, kernel, shortcut = _bhl_inputs() + expected = fftconv2d_fp32_bhl(x, kernel, shortcut) + + out = _compile_with_lowering(fftconv2d_fp32_bhl, x, kernel, shortcut, allow_reduced_precision=False) + + assert lowering_stats().get("rewritten") == 1 + _assert_l2_close(out, expected, torch.float32) + + +def test_no_fft_in_graph_is_untouched(): + """A graph with no irfft2 exits early without touching anything.""" + + def plain(x): + return x * 2 + 1 + + x = torch.randn(4, 4, device="cuda") + out = _compile_with_lowering(plain, x) + + assert lowering_stats() == {} + torch.testing.assert_close(out, plain(x), atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# Registration plumbing +# --------------------------------------------------------------------------- + + +class TestRegistration: + def test_options_route_does_not_touch_global_config(self): + """The ``options=`` patch is scoped to one callable, unlike the context manager.""" + from torch._inductor import config as inductor_config + + assert inductor_config.pre_grad_custom_pass is None + x, kernel, shortcut = _bhl_inputs() + + torch.compile(fftconv2d_fp32_bhl, fullgraph=True, options=fused_fftconv2d_options())(x, kernel, shortcut) + + assert lowering_stats().get("rewritten") == 1 + assert inductor_config.pre_grad_custom_pass is None + + def test_context_manager_route_also_rewrites(self): + """The global-config route stays supported for framework-owned compiles.""" + x, kernel, shortcut = _bhl_inputs() + expected = fftconv2d_fp32_bhl(x, kernel, shortcut) + + with fused_fftconv2d_lowering(): + out = torch.compile(fftconv2d_fp32_bhl, fullgraph=True)(x, kernel, shortcut) + + assert lowering_stats().get("rewritten") == 1 + _assert_l2_close(out, expected, torch.float32) + + def test_options_carries_the_precision_flag(self): + options = fused_fftconv2d_options(allow_reduced_precision=False) + assert options["pre_grad_custom_pass"].allow_reduced_precision is False + + def test_restores_previous_pass_on_exit(self): + from torch._inductor import config as inductor_config + + sentinel = FusedFFTConv2dLowering() + inductor_config.pre_grad_custom_pass = sentinel + try: + with fused_fftconv2d_lowering(): + assert inductor_config.pre_grad_custom_pass is not sentinel + assert inductor_config.pre_grad_custom_pass is sentinel + finally: + inductor_config.pre_grad_custom_pass = None + + def test_restores_on_exception(self): + from torch._inductor import config as inductor_config + + assert inductor_config.pre_grad_custom_pass is None + with pytest.raises(RuntimeError, match="boom"): + with fused_fftconv2d_lowering(): + raise RuntimeError("boom") + assert inductor_config.pre_grad_custom_pass is None + + def test_uuid_is_stable_and_flag_dependent(self): + """Inductor mixes uuid() into its cache key, so it must track behaviour.""" + assert FusedFFTConv2dLowering(True).uuid() == FusedFFTConv2dLowering(True).uuid() + assert FusedFFTConv2dLowering(True).uuid() != FusedFFTConv2dLowering(False).uuid() + + def test_not_installed_means_no_rewrite(self): + """Without the context manager, torch.compile leaves the chain as-is.""" + x, kernel, shortcut = _bhl_inputs() + expected = fftconv2d_fp32_bhl(x, kernel, shortcut) + + out = torch.compile(fftconv2d_fp32_bhl, fullgraph=True)(x, kernel, shortcut) + + assert lowering_stats() == {} + _assert_l2_close(out, expected, torch.float32) diff --git a/tests/ops/test_fused_fftconv2d.py b/tests/ops/test_fused_fftconv2d.py new file mode 100644 index 00000000..5a0fe5d5 --- /dev/null +++ b/tests/ops/test_fused_fftconv2d.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the fused 2D FFT-conv wrappers in ``nvsubquadratic.ops.fftconv_custom``. + +These wrap ``subquadratic_ops_torch.fused_fft_conv2d``, which differs from the +older ``fft_conv2d`` path in two ways that this file pins down: + + - It runs **natively** in fp32/fp16/bf16 rather than upcasting to fp32. + - It crops the 'same' window at ``fft_size // 2`` rather than at ``K // 2``. + The wrapper compensates with a top/left filter pad; the equivalence tests + below are what make that compensation load-bearing rather than incidental. + Dropping the pad shifts the output by ``fft_size // 2 - K // 2`` pixels, + which shows up here as a ~1.41 (= sqrt(2), i.e. fully decorrelated) error. + +Correctness is measured against ``fftconv2d_fp32_bhl`` / ``fftconv2d_fp32_blh`` +with a normwise (L2) relative error rather than elementwise ``assert_close``: +FFT roundoff scales with ``||output||``, not with each element, so near-zero +output elements make elementwise relative error meaningless. + +Usage (requires GPU — run inside SLURM): + srun --gres=gpu:1 -c 16 --partition low \\ + conda run -n nv-subq python -m pytest tests/ops/test_fused_fftconv2d.py -v -o addopts="" +""" + +import pytest +import torch + +from nvsubquadratic.ops.fftconv import fftconv2d_fp32_bhl, fftconv2d_fp32_blh +from nvsubquadratic.ops.fftconv_custom import ( + fused_fftconv2d_bhl, + fused_fftconv2d_bhl_chunked, + fused_fftconv2d_blh, + fused_fftconv2d_blh_chunked, + fused_fftconv2d_max_spatial, + fused_fftconv2d_supported, + resolve_fused_fft_size, +) +from tests.conftest import requires_subq_ops_fused + + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + +pytestmark = [requires_subq_ops_fused, requires_cuda] + +# Normwise relative-error budgets, calibrated against the observed error of the +# fused kernel vs the fp32 torch reference (fp32 ~3e-7, fp16 ~3e-4, bf16 ~3e-3) +# with ~3x headroom. Gradients accumulate over the batch, so they get more room. +L2_TOL = {torch.float32: 1e-6, torch.float16: 1e-3, torch.bfloat16: 8e-3} +L2_TOL_GRAD = {torch.float32: 1e-6, torch.float16: 2e-3, torch.bfloat16: 1.5e-2} + +HIDDEN_DIM = 16 +BATCH = 2 + + +def l2_rel(pred: torch.Tensor, ref: torch.Tensor) -> float: + """Normwise relative error, computed in fp64 so the metric adds no roundoff.""" + pred64, ref64 = pred.double(), ref.double() + den = ref64.norm() + return ((pred64 - ref64).norm() / den).item() if den > 0 else (pred64 - ref64).norm().item() + + +def assert_l2_close(pred, ref, dtype, tol_table=None, name=""): + tol = (tol_table or L2_TOL)[dtype] + rel = l2_rel(pred, ref) + assert rel < tol, f"{name} L2 rel error {rel:.3e} exceeds tol {tol:.1e}" + + +def _inputs(spatial, kernel_size, dtype, kernel_batch=1, seed=0): + """Build a BHL input/kernel/shortcut triple on CUDA.""" + torch.manual_seed(seed) + x = torch.randn(BATCH, HIDDEN_DIM, spatial, spatial, device="cuda", dtype=dtype) + # Scaled down so the convolution output stays in a sane range for fp16. + kernel = (torch.randn(kernel_batch, HIDDEN_DIM, kernel_size, kernel_size, device="cuda") * 0.05).to(dtype) + shortcut = torch.randn(HIDDEN_DIM, device="cuda", dtype=dtype) + return x, kernel, shortcut + + +# --------------------------------------------------------------------------- +# fft_size resolution +# --------------------------------------------------------------------------- + + +class TestResolveFftSize: + """The tile-size chooser and its guard rails.""" + + @pytest.mark.parametrize( + ("spatial", "expected"), + [(4, 8), (7, 16), (8, 16), (16, 32), (32, 64), (64, 128)], + ) + def test_default_for_double_grid_kernel(self, spatial, expected): + """CKConvND's double-grid kernel (K = 2N-1) resolves to a 2*next_pow2(N) tile.""" + assert resolve_fused_fft_size(spatial, spatial, 2 * spatial - 1, 2 * spatial - 1) == expected + + def test_default_is_smallest_admissible(self): + """A small kernel does not inflate the tile beyond what the input needs.""" + assert resolve_fused_fft_size(16, 16, 3, 3) == 32 + + def test_kernel_can_drive_the_tile_size(self): + """A kernel wider than 2x the input still gets enough headroom for the pre-pad.""" + # ceil(96 / 2) = 48 > 8, so the kernel, not the input, sets the tile. + assert resolve_fused_fft_size(8, 8, 96, 96) == 128 + + def test_non_square_uses_the_larger_axis(self): + assert resolve_fused_fft_size(8, 32, 15, 63) == 64 + + def test_minimum_tile_is_8(self): + """Tiny shapes clamp up to the smallest tile the kernel is compiled for.""" + assert resolve_fused_fft_size(1, 1, 1, 1) == 8 + + def test_rejects_oversized_spatial(self): + with pytest.raises(ValueError, match="too large for the fused 2D FFT kernel"): + resolve_fused_fft_size(65, 65, 129, 129) + + def test_rejects_unsupported_explicit_size(self): + with pytest.raises(ValueError, match="fft_size must be one of"): + resolve_fused_fft_size(8, 8, 15, 15, fft_size=48) + + def test_rejects_too_small_explicit_size(self): + with pytest.raises(ValueError, match="too small for input"): + resolve_fused_fft_size(32, 32, 63, 63, fft_size=16) + + def test_accepts_larger_explicit_size(self): + """Over-provisioning the tile is allowed (costs speed, not correctness).""" + assert resolve_fused_fft_size(8, 8, 15, 15, fft_size=128) == 128 + + def test_supported_probe_matches_resolve(self): + assert fused_fftconv2d_supported(64, 64, 127, 127) is True + assert fused_fftconv2d_supported(65, 65, 129, 129) is False + + def test_max_spatial(self): + assert fused_fftconv2d_max_spatial() == 64 + + +# --------------------------------------------------------------------------- +# Forward equivalence with the torch.fft reference +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +class TestForwardMatchesReference: + """The fused path reproduces fftconv2d_fp32_bhl up to dtype roundoff.""" + + @pytest.mark.parametrize("spatial", [7, 8, 16, 32, 64]) + def test_double_grid_kernel(self, spatial, dtype): + """K = 2N-1, the kernel size CKConvND generates on a double grid.""" + x, kernel, shortcut = _inputs(spatial, 2 * spatial - 1, dtype) + assert_l2_close(fused_fftconv2d_bhl(x, kernel, shortcut), fftconv2d_fp32_bhl(x, kernel, shortcut), dtype) + + @pytest.mark.parametrize("kernel_size", [1, 3, 8, 15, 16, 17, 32]) + def test_kernel_sizes(self, kernel_size, dtype): + """Both parities and both sides of K == N, since the pre-pad is K//2-dependent.""" + x, kernel, shortcut = _inputs(16, kernel_size, dtype) + assert_l2_close(fused_fftconv2d_bhl(x, kernel, shortcut), fftconv2d_fp32_bhl(x, kernel, shortcut), dtype) + + def test_film_per_sample_kernel(self, dtype): + """A per-sample (FiLM) kernel batch is supported on the fused path.""" + x, kernel, shortcut = _inputs(16, 31, dtype, kernel_batch=BATCH) + assert_l2_close(fused_fftconv2d_bhl(x, kernel, shortcut), fftconv2d_fp32_bhl(x, kernel, shortcut), dtype) + + def test_without_shortcut(self, dtype): + x, kernel, _ = _inputs(16, 31, dtype) + assert_l2_close(fused_fftconv2d_bhl(x, kernel), fftconv2d_fp32_bhl(x, kernel), dtype) + + def test_blh_layout(self, dtype): + """Channels-last entry point, against the channels-last reference.""" + x, kernel, shortcut = _inputs(16, 31, dtype) + x_blh = x.permute(0, 2, 3, 1).contiguous() + kernel_blh = kernel.permute(0, 2, 3, 1).contiguous() + assert_l2_close( + fused_fftconv2d_blh(x_blh, kernel_blh, shortcut), + fftconv2d_fp32_blh(x_blh, kernel_blh, shortcut), + dtype, + ) + + @pytest.mark.parametrize("chunk_size", [4, 8, HIDDEN_DIM, HIDDEN_DIM * 2]) + def test_chunked_matches_unchunked(self, chunk_size, dtype): + """Channel chunking is a pure memory optimisation, including a ragged last chunk.""" + x, kernel, shortcut = _inputs(16, 31, dtype) + chunked = fused_fftconv2d_bhl_chunked(x, kernel, shortcut, chunk_size=chunk_size) + torch.testing.assert_close(chunked, fused_fftconv2d_bhl(x, kernel, shortcut), atol=0, rtol=0) + + def test_chunked_blh(self, dtype): + x, kernel, shortcut = _inputs(16, 31, dtype) + x_blh = x.permute(0, 2, 3, 1).contiguous() + kernel_blh = kernel.permute(0, 2, 3, 1).contiguous() + assert_l2_close( + fused_fftconv2d_blh_chunked(x_blh, kernel_blh, shortcut, chunk_size=4), + fftconv2d_fp32_blh(x_blh, kernel_blh, shortcut), + dtype, + ) + + +# --------------------------------------------------------------------------- +# Crop alignment — the wrapper's reason for existing +# --------------------------------------------------------------------------- + + +class TestCropAlignment: + """The top/left pre-pad is what aligns the fused crop with every other backend.""" + + def test_raw_upstream_call_is_shifted(self): + """Without the pre-pad the upstream op is offset, not merely less accurate. + + This is the regression guard: if someone "simplifies" the wrapper by + calling fused_fft_conv2d directly, this is the error they would ship. + """ + pytest.importorskip("subquadratic_ops_torch") + from subquadratic_ops_torch.fused_fft_conv2d import fused_fft_conv2d + + spatial, kernel_size = 16, 31 + x, kernel, _ = _inputs(spatial, kernel_size, torch.float32) + ref = fftconv2d_fp32_bhl(x, kernel) + + unpadded = fused_fft_conv2d(x.contiguous(), kernel.contiguous(), fft_size=2 * spatial) + # ~sqrt(2): a shift by fft_size//2 - K//2 decorrelates the output entirely. + assert l2_rel(unpadded, ref) > 1.0 + + assert_l2_close(fused_fftconv2d_bhl(x, kernel), ref, torch.float32) + + +# --------------------------------------------------------------------------- +# Backward +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("kernel_batch", [1, BATCH]) +def test_backward_matches_reference(dtype, kernel_batch): + """Gradients w.r.t. x, kernel, and shortcut all match the fp32 reference.""" + spatial, kernel_size = 16, 31 + x0, k0, s0 = _inputs(spatial, kernel_size, torch.float32, kernel_batch=kernel_batch) + grad_out = torch.randn(BATCH, HIDDEN_DIM, spatial, spatial, device="cuda") + + def run(fn, cast_to): + x = x0.to(cast_to).detach().requires_grad_() + k = k0.to(cast_to).detach().requires_grad_() + s = s0.to(cast_to).detach().requires_grad_() + fn(x, k, s).backward(grad_out.to(cast_to)) + return x.grad, k.grad, s.grad + + ref = run(fftconv2d_fp32_bhl, torch.float32) + fused = run(fused_fftconv2d_bhl, dtype) + + for got, want, name in zip(fused, ref, ("grad_x", "grad_kernel", "grad_shortcut")): + assert_l2_close(got, want, dtype, tol_table=L2_TOL_GRAD, name=name) + + +def test_backward_flows_through_chunked(): + """The chunked path is differentiable and agrees with the unchunked one.""" + x0, k0, s0 = _inputs(16, 31, torch.float32) + + def run(fn, **kw): + x = x0.detach().requires_grad_() + k = k0.detach().requires_grad_() + s = s0.detach().requires_grad_() + fn(x, k, s, **kw).sum().backward() + return x.grad, k.grad, s.grad + + for got, want in zip(run(fused_fftconv2d_bhl_chunked, chunk_size=4), run(fused_fftconv2d_bhl)): + torch.testing.assert_close(got, want, atol=1e-5, rtol=1e-5) + + +# --------------------------------------------------------------------------- +# Dtype behaviour +# --------------------------------------------------------------------------- + + +class TestDtype: + """Native-dtype passthrough — the reason to prefer this over ``fftconv2d_bhl``.""" + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_output_dtype_matches_input(self, dtype): + x, kernel, shortcut = _inputs(16, 31, dtype) + assert fused_fftconv2d_bhl(x, kernel, shortcut).dtype == dtype + assert fused_fftconv2d_bhl_chunked(x, kernel, shortcut).dtype == dtype + + def test_kernel_dtype_is_coerced_to_input(self): + """A kernel in a different dtype is cast rather than raising from the CUDA op.""" + x, kernel, shortcut = _inputs(16, 31, torch.bfloat16) + out = fused_fftconv2d_bhl(x, kernel.float(), shortcut) + assert out.dtype == torch.bfloat16 + assert_l2_close(out, fused_fftconv2d_bhl(x, kernel, shortcut), torch.bfloat16) + + +# --------------------------------------------------------------------------- +# Error handling +# --------------------------------------------------------------------------- + + +def test_rejects_oversized_input(): + """Exceeding the 64x64 cap raises with a pointer to the other backends.""" + x = torch.randn(1, 4, 65, 65, device="cuda") + kernel = torch.randn(1, 4, 129, 129, device="cuda") + with pytest.raises(ValueError, match="too large for the fused 2D FFT kernel"): + fused_fftconv2d_bhl(x, kernel)