Add compact B200 FFT3D support to the shared kernel package

by ubospica - opened 24 days ago

base: refs/heads/main

←

from: refs/pr/1

Discussion Files changed

+26662

-10

This view is limited to 50 files because it contains too many changes. See the raw diff here.

Files changed (50) hide show

README.md +113 -10
build/torch-cuda/LICENSE.txt +21 -0
build/torch-cuda/__init__.py +60 -0
build/torch-cuda/_b200/__init__.py +151 -0
build/torch-cuda/_b200/c2r_cuda_z_inverse_cuda_kernel.py +63 -0
build/torch-cuda/_b200/c2r_cuda_z_inverse_fft_axis.cu +128 -0
build/torch-cuda/_b200/c2r_dense_half_inverse_wide_redirect.cu +85 -0
build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_cuda_kernel.py +63 -0
build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_fft_axis.cu +128 -0
build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_lto.py +110 -0
build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_native_dispatch.py +100 -0
build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_wide_redirect.cu +65 -0
build/torch-cuda/_b200/c2r_f64_dense_half_inverse_lto.py +110 -0
build/torch-cuda/_b200/c2r_f64_dense_half_inverse_wide_redirect.cu +67 -0
build/torch-cuda/_b200/c2r_f64_half_vector_inverse_wide_redirect.cu +65 -0
build/torch-cuda/_b200/c2r_f64_narrow_index_inverse_scalar_output_reconstruct.cu +45 -0
build/torch-cuda/_b200/c2r_f64_paired_x_inverse_cuda_kernel.py +64 -0
build/torch-cuda/_b200/c2r_f64_paired_x_inverse_fft_axis.cu +270 -0
build/torch-cuda/_b200/c2r_f64_paired_x_inverse_wide_redirect.cu +29 -0
build/torch-cuda/_b200/c2r_f64_paired_y_direct_inverse_lto.py +116 -0
build/torch-cuda/_b200/c2r_f64_paired_y_direct_inverse_wide_redirect.cu +114 -0
build/torch-cuda/_b200/c2r_f64_partial_y_inverse_wide_redirect.cu +114 -0
build/torch-cuda/_b200/c2r_f64_scalar_input_inverse_scalar_io.cu +26 -0
build/torch-cuda/_b200/c2r_paired_y_direct_inverse_wide_redirect.cu +109 -0
build/torch-cuda/_b200/c2r_paired_y_inverse_wide_redirect.cu +87 -0
build/torch-cuda/_b200/c2r_pitched_half_inverse_wide_redirect.cu +86 -0
build/torch-cuda/_b200/cuda_kernel.py +64 -0
build/torch-cuda/_b200/cuda_twiddle_forward_lto.py +113 -0
build/torch-cuda/_b200/cuda_twiddle_forward_native_dispatch.c +42 -0
build/torch-cuda/_b200/cuda_twiddle_forward_native_dispatch.py +101 -0
build/torch-cuda/_b200/cuda_z_inverse_cufft_support.py +105 -0
build/torch-cuda/_b200/cuda_z_inverse_lto.py +110 -0
build/torch-cuda/_b200/cuda_z_inverse_native_dispatch.c +42 -0
build/torch-cuda/_b200/cuda_z_inverse_native_dispatch.py +101 -0
build/torch-cuda/_b200/cuda_z_inverse_schedule.py +71 -0
build/torch-cuda/_b200/dense_half_inverse_abi.py +30 -0
build/torch-cuda/_b200/dense_half_inverse_kernels.py +10 -0
build/torch-cuda/_b200/dense_half_inverse_lto.py +110 -0
build/torch-cuda/_b200/dense_half_inverse_native_dispatch.c +43 -0
build/torch-cuda/_b200/dense_half_inverse_native_dispatch.py +104 -0
build/torch-cuda/_b200/dense_half_inverse_schedule.py +73 -0
build/torch-cuda/_b200/direct_dispatch.c +53 -0
build/torch-cuda/_b200/direct_dispatch.py +98 -0
build/torch-cuda/_b200/fft_axis.cu +169 -0
build/torch-cuda/_b200/forward_cuda_f32.py +151 -0
build/torch-cuda/_b200/forward_half.py +149 -0
build/torch-cuda/_b200/forward_native.py +182 -0
build/torch-cuda/_b200/forward_paired_y.py +153 -0
build/torch-cuda/_b200/forward_scalar.py +63 -0
build/torch-cuda/_b200/forward_small.py +195 -0

README.md CHANGED Viewed

@@ -10,10 +10,17 @@ tags:
   - fp64
   - blackwell
   - rtx-5090
 ---
 # KDA 3D FFTs on Packed Hermitian Half-Spectra
 CuTe-DSL forward (real-to-complex) and inverse (complex-to-real) unnormalized 3D
 FFTs of one real grid, in fp32 and fp64, for the NVIDIA RTX 5090. The spectrum
 is the packed half-spectrum that structured-grid spectral solvers use: x
@@ -28,6 +35,8 @@ about 3x.**
 ## Requirements
 - NVIDIA GPU with compute capability 12.0; tuned and tested on the RTX 5090
   only
 - CUDA 13 compatible driver
@@ -36,11 +45,14 @@ about 3x.**
 ## Usage
 ```python
 import torch
 from kernels import get_kernel
-fft3d = get_kernel("humanfia-lab/fft3d", version=1, trust_remote_code=True)
 u = torch.randn(96, 96, 144, device="cuda")  # [nz, ny, nx], fp32 or fp64
 spectrum = fft3d.r2c(u)  # [96, 96, 73, 2] == torch.view_as_real(torch.fft.rfftn(u))
@@ -60,24 +72,25 @@ spectrum = launch()  # ...then transform its current contents
   `launch()`, which transforms the buffer's current contents into one output
   tensor that it reuses on every call. Calls of one `launch` must be ordered
   (one stream). Keep `launch` alive as long as a CUDA graph captured from it.
-  fp64 r2c has no two-phase form; its `launch()` calls `r2c(x)`.
 - Calls on different CUDA streams may run concurrently, including calls of the
-  same shape.
 - The first call for a new shape compiles its kernels (CuTe DSL JIT), typically
-  in 1 to 30 s, up to about 4 minutes for a very long z axis over a tiny y-x
-  plane (such as 4096x8x16). Make that first call outside CUDA graph capture;
   later calls can be captured.
 - `c2r` expects a Hermitian-consistent half-spectrum, such as the output of
   `r2c` or Fourier-space data that keeps its symmetry. On other input the result
   depends on the order of the 1D passes, as with cuFFT; these kernels then
   match `torch.fft.irfftn`.
-- Tuned on the eight grids below and checked on dozens of others per kernel,
-  including odd and prime extents and up to 1024x1024x512. fp64 r2c raises
   `NotImplementedError` when a y or z axis needs a direct DFT longer than 6336
   points (prime factors above 13, or smooth lengths above 1536); fp64 c2r raises
   it for nz from 1280 to 4096, ny from 3072 to 4096 or nx / 2 from 3072 to 4096.
-## Accuracy
 fp32 kernels compute in fp32 from twiddle tables computed in float64 on the
 host: no reduced precision, no tensor cores, no approximate math. Against the
@@ -98,7 +111,7 @@ magnitude is unaffected.
 Repeated calls give bitwise identical results.
-## Performance
 Speedup over cuFFT's 3D plan:
@@ -141,7 +154,96 @@ timing and executed out of place; its C2R overwrites its input and is given a
 copy. `torch.fft` is `rfftn` / `irfftn` on the same library. Speeds differ by a
 few percent between RTX 5090 cards.
-## Source and licensing
 Written by Claude Opus 5.5 in four 12-hour runs of KDA (kernel design agents),
 one per KDA task (`fft3d_r2c`, `fft3d_c2r`, `fft3d_r2c_f64`, `fft3d_c2r_f64`),
@@ -151,6 +253,7 @@ and accepted by the KDA judge. Published from commit `e539498` of
 its tests and the fixes made after the judge's checks. MIT-licensed.
 ## Available functions
 - `c2r`
 - `prepare_c2r`
 - `prepare_r2c`

   - fp64
   - blackwell
   - rtx-5090
+  - b200
 ---
 # KDA 3D FFTs on Packed Hermitian Half-Spectra
+NVIDIA B200 and RTX 5090 share the `r2c`, `c2r`, `prepare_r2c`, and
+`prepare_c2r` interface in `build/torch-cuda`. The package selects the
+implementation from the CUDA device and tensor dtype.
+See [NVIDIA B200](#nvidia-b200) for B200 requirements and measurements; the
+accuracy and performance sections before it describe the RTX 5090.
 CuTe-DSL forward (real-to-complex) and inverse (complex-to-real) unnormalized 3D
 FFTs of one real grid, in fp32 and fp64, for the NVIDIA RTX 5090. The spectrum
 is the packed half-spectrum that structured-grid spectral solvers use: x
 ## Requirements
+For RTX 5090:
 - NVIDIA GPU with compute capability 12.0; tuned and tested on the RTX 5090
   only
 - CUDA 13 compatible driver
 ## Usage
+The pending unified package is available from the pull-request revision below.
+The existing `version=1` release remains the RTX 5090 package.
 ```python
 import torch
 from kernels import get_kernel
+fft3d = get_kernel("humanfia-lab/fft3d", revision="refs/pr/1", trust_remote_code=True)
 u = torch.randn(96, 96, 144, device="cuda")  # [nz, ny, nx], fp32 or fp64
 spectrum = fft3d.r2c(u)  # [96, 96, 73, 2] == torch.view_as_real(torch.fft.rfftn(u))
   `launch()`, which transforms the buffer's current contents into one output
   tensor that it reuses on every call. Calls of one `launch` must be ordered
   (one stream). Keep `launch` alive as long as a CUDA graph captured from it.
+  On RTX 5090, fp64 r2c has no two-phase form; its `launch()` calls `r2c(x)`.
 - Calls on different CUDA streams may run concurrently, including calls of the
+  same shape. On B200, prepare and warm each shape on its execution stream
+  before capturing a graph; a prepared launch stays bound to that stream.
 - The first call for a new shape compiles its kernels (CuTe DSL JIT), typically
+  on RTX 5090 in 1 to 30 s, up to about 4 minutes for a very long z axis over
+  a tiny y-x plane (such as 4096x8x16). Make that first call outside CUDA graph capture;
   later calls can be captured.
 - `c2r` expects a Hermitian-consistent half-spectrum, such as the output of
   `r2c` or Fourier-space data that keeps its symmetry. On other input the result
   depends on the order of the 1D passes, as with cuFFT; these kernels then
   match `torch.fft.irfftn`.
+- RTX 5090: tuned on the eight grids below and checked on dozens of others
+  per kernel, including odd and prime extents and up to 1024x1024x512. fp64 r2c raises
   `NotImplementedError` when a y or z axis needs a direct DFT longer than 6336
   points (prime factors above 13, or smooth lengths above 1536); fp64 c2r raises
   it for nz from 1280 to 4096, ny from 3072 to 4096 or nx / 2 from 3072 to 4096.
+## RTX 5090 accuracy
 fp32 kernels compute in fp32 from twiddle tables computed in float64 on the
 host: no reduced precision, no tensor cores, no approximate math. Against the
 Repeated calls give bitwise identical results.
+## RTX 5090 performance
 Speedup over cuFFT's 3D plan:
 copy. `torch.fft` is `rfftn` / `irfftn` on the same library. Speeds differ by a
 few percent between RTX 5090 cards.
+## NVIDIA B200
+The B200 implementation is in `build/torch-cuda/_b200` and uses the same
+public interface shown above. It uses Triton, CUDA C++, and cuFFT, and supports
+exactly the eight grids below, with no leading batch dimension. Other grids
+raise `ValueError`. Both directions are unnormalized, return fresh contiguous
+outputs, and preserve inputs. Both preparation functions return a callable
+bound to the input buffer, with one reused output.
+Requirements: B200 (compute capability 10.0), a CUDA 13 development environment
+with `nvcc` and a C++ compiler, and the PyTorch CUDA build with its matching
+Triton. The performance measurements used PyTorch 2.13.0+cu130 and Triton 3.7.1.
+Keep cuFFT, NVRTC, and nvJitLink from a compatible CUDA toolkit on the library
+search path; for CUDA 13.0, for example:
+```bash
+export CUDA_HOME=/usr/local/cuda-13.0
+export PATH="$CUDA_HOME/bin:$PATH"
+export LD_LIBRARY_PATH="$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}"
+```
+FP32 means 32-bit floating point and FP64 means 64-bit floating point; R2C is
+real-to-complex and C2R is complex-to-real. FP32 requests retain FP64 intermediate
+arithmetic. Against an FP64 reference, the elementwise limits are
+`0.01 + 0.001*abs(reference)` for FP32 and `1e-6 + 1e-9*abs(reference)` for FP64.
+**Across all 32 B200 workloads, the geometric mean is 1.1300x over direct cuFFT
+3D plans and 1.3644x over default `torch.fft` in the same run.** A separate
+complete run measured 1.1277x over direct 3D plans. Ratios below 1 mean the
+published kernel is slower. In particular, FP32 forward averages 0.90x: the
+B200 kernels are not faster on every workload.
+Speedup over a single three-dimensional cuFFT plan, including the inverse
+input-preservation copy:
+| Grid (nz x ny x nx) | R2C FP32 | C2R FP32 | R2C FP64 | C2R FP64 |
+|---|---:|---:|---:|---:|
+| 32x32x32 | 2.00x | 2.03x | 2.20x | 2.26x |
+| 64x64x64 | 0.88x | 1.07x | 1.11x | 1.26x |
+| 80x80x112 | 0.92x | 1.20x | 1.15x | 1.35x |
+| 96x96x96 | 1.04x | 1.25x | 1.40x | 1.56x |
+| 96x96x144 | 0.84x | 0.97x | 1.06x | 1.21x |
+| 128x128x128 | 0.74x | 0.88x | 1.13x | 1.31x |
+| 192x192x192 | 0.69x | 0.80x | 1.00x | 1.15x |
+| 256x256x256 | 0.58x | 0.76x | 1.00x | 1.34x |
+| Geometric mean | 0.90x | 1.07x | 1.21x | 1.40x |
+Latency in microseconds, **default torch.fft / direct cuFFT 3D / B200 kernel**:
+| Grid (nz x ny x nx) | R2C FP32 | C2R FP32 | R2C FP64 | C2R FP64 |
+|---|---:|---:|---:|---:|
+| 32x32x32 | 13.76 / 11.25 / 5.63 | 9.28 / 11.58 / 5.70 | 14.40 / 12.03 / 5.47 | 10.16 / 11.84 / 5.23 |
+| 64x64x64 | 12.22 / 9.33 / 10.61 | 11.58 / 11.87 / 11.06 | 14.46 / 11.71 / 10.56 | 12.37 / 14.64 / 11.58 |
+| 80x80x112 | 17.50 / 12.58 / 13.60 | 14.59 / 16.61 / 13.86 | 21.63 / 16.53 / 14.40 | 19.36 / 19.42 / 14.40 |
+| 96x96x96 | 21.41 / 15.49 / 14.85 | 17.26 / 18.75 / 15.01 | 29.18 / 21.79 / 15.52 | 22.66 / 25.50 / 16.32 |
+| 96x96x144 | 22.78 / 14.46 / 17.31 | 19.58 / 17.68 / 18.14 | 29.42 / 19.42 / 18.27 | 28.30 / 24.19 / 19.92 |
+| 128x128x128 | 27.68 / 17.63 / 23.71 | 31.01 / 21.49 / 24.38 | 38.30 / 28.80 / 25.47 | 35.62 / 34.69 / 26.43 |
+| 192x192x192 | 75.89 / 49.07 / 71.26 | 74.74 / 60.61 / 76.18 | 114.82 / 78.06 / 78.14 | 120.77 / 98.59 / 85.38 |
+| 256x256x256 | 167.39 / 85.92 / 148.10 | 139.47 / 106.83 / 140.30 | 259.99 / 160.02 / 160.51 | 195.57 / 206.66 / 154.58 |
+Measured October 14, 2026, on an exclusive NVIDIA B200, driver 610.57.04,
+PyTorch 2.13.0+cu130, Triton 3.7.1, and the same cuFFT 12.0 library shipped with
+PyTorch. Each plan is created before timing with `cufftMakePlan3d`. Timing uses
+NVIDIA's CUDA Profiling Tools Interface (CUPTI) on CUDA graph replay with a cold
+L2 cache, 3 warmups and 30 samples per trial, and the median of 3 trial medians.
+The order of implementations changes between trials. Allocation and planning
+are excluded from GPU execution time; inverse input copies are included.
+The previous B200 **1.410691x** figure compared against default `torch.fft`
+using a CUDA-event timing protocol. It was not a direct-3D-cuFFT score and is
+superseded as the native-baseline headline. It is not numerically interchangeable
+with these graph measurements. In
+[PyTorch 2.13's CUDA implementation](https://github.com/pytorch/pytorch/blob/v2.13.0/aten/src/ATen/native/cuda/SpectralOps.cpp#L305-L314),
+transform axes starting with `(0, 1)` take two separate cuFFT plans. Adding a
+batch dimension of size one and transforming axes `(1, 2, 3)` reaches the single
+3D-plan path through PyTorch. That path and direct cuFFT had matching GPU-kernel
+sequences in all eight profiled cases. A single 3D plan can still launch several
+GPU kernels, and the split path is faster for some inverse sizes.
+Inverse timings above use Hermitian-consistent spectra, as
+[required by cuFFT](https://docs.nvidia.com/cuda/cufft/index.html#type-of-transform).
+For arbitrary packed spectra accepted by the B200 implementation, directly
+substituting cuFFT is not equivalent. Including a boundary-plane projection
+with the input copy makes the direct baseline match that wider contract; the
+32-workload geometric mean is then **1.1125x**. All 144 candidate numerical and
+input-contract checks passed. Native FP32 cuFFT exceeds the historical strict
+per-element tolerance on several constant-input stress cases; the timed random
+inputs pass. The implementation retains its FP64 intermediate arithmetic.
+## RTX 5090 source and licensing
 Written by Claude Opus 5.5 in four 12-hour runs of KDA (kernel design agents),
 one per KDA task (`fft3d_r2c`, `fft3d_c2r`, `fft3d_r2c_f64`, `fft3d_c2r_f64`),
 its tests and the fixes made after the judge's checks. MIT-licensed.
 ## Available functions
 - `c2r`
 - `prepare_c2r`
 - `prepare_r2c`

build/torch-cuda/LICENSE.txt ADDED Viewed

	@@ -0,0 +1,21 @@

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

build/torch-cuda/__init__.py ADDED Viewed

	@@ -0,0 +1,60 @@

+"""3D FFTs on packed Hermitian half-spectra for NVIDIA B200 and RTX 5090.
+r2c(x): x is a contiguous float32 or float64 [nz, ny, nx] CUDA tensor (x fastest); returns a
+new [nz, ny, nx // 2 + 1, 2] tensor of the same dtype, torch.view_as_real(torch.fft.rfftn(x)):
+the unnormalized forward transform, packed with an x-row stride of exactly nx // 2 + 1.
+c2r(X): X is such a half-spectrum, a contiguous float32 or float64 [nz, ny, nf, 2]; returns a
+new real [nz, ny, 2 * (nf - 1)], torch.fft.irfftn(torch.view_as_complex(X), norm="forward"):
+the unnormalized inverse, so c2r(r2c(x)) is nz * ny * nx * x. X is not modified.
+prepare_r2c(x) / prepare_c2r(X) plan for exactly this input buffer and return launch(), which
+transforms the buffer's current contents and returns an output tensor it reuses on every
+call; keep launch alive as long as a CUDA graph captured from it. On RTX 5090, fp64 r2c has
+no two-phase form, so its launch() is r2c(x) and returns a new tensor each call.
+The implementation is selected by device and dtype, and imported on first use.
+"""
+import functools
+import importlib
+import torch
+__all__ = ["c2r", "prepare_c2r", "prepare_r2c", "r2c"]
+_PACKAGES = {
+    ("r2c", torch.float32): "_r2c_f32",
+    ("r2c", torch.float64): "_r2c_f64",
+    ("c2r", torch.float32): "_c2r_f32",
+    ("c2r", torch.float64): "_c2r_f64",
+}
+def _kernel(direction, tensor):
+    package = _PACKAGES.get((direction, tensor.dtype))
+    if package is None:
+        raise TypeError(f"{direction} takes float32 or float64, not {tensor.dtype}")
+    if tensor.device.type != "cuda":
+        raise ValueError(f"{direction} requires a CUDA tensor")
+    if torch.cuda.get_device_capability(tensor.device) == (10, 0):
+        from ._b200 import kernel
+        return kernel(direction, tensor.dtype)
+    return importlib.import_module(f".{package}", __name__)
+def r2c(x):
+    return _kernel("r2c", x).run(x)
+def c2r(X):
+    return _kernel("c2r", X).run(X)
+def prepare_r2c(x):
+    kernel = _kernel("r2c", x)
+    return kernel.prepare(x) if hasattr(kernel, "prepare") else functools.partial(kernel.run, x)
+def prepare_c2r(X):
+    return _kernel("c2r", X).prepare(X)

build/torch-cuda/_b200/__init__.py ADDED Viewed

	@@ -0,0 +1,151 @@

+"""Selected B200 FFT plans behind the shared public FFT3D interface."""
+import importlib
+import threading
+import torch
+from torch._C import _cuda_getCurrentRawStream
+from .prepared_call import Plan, load_module
+_SHAPES = frozenset((
+    (32, 32, 32), (64, 64, 64), (80, 80, 112), (96, 96, 96),
+    (96, 96, 144), (128, 128, 128), (192, 192, 192), (256, 256, 256),
+))
+def _plan(module, shape, dtype, device, stream, **options):
+    cls = importlib.import_module('.' + module, __name__).Plan
+    return cls(shape, dtype, device=device, stream=stream, **options)
+def _owner(direction, shape, dtype, device, stream):
+    double = dtype == torch.float64
+    options = dict(device=device, stream=stream)
+    if direction == 'r2c':
+        if shape == (32, 32, 32):
+            return _plan('forward_small', shape, dtype, **options, direction='r2c',
+                         rows_per_block=4, columns_per_block=4, num_warps=4,
+                         real_strategy='fused_xy', dispatch='native')
+        if shape == (64, 64, 64):
+            return _plan('forward_cuda_f32', shape, dtype, **options, columns=8, warps=8)
+        if shape in ((128, 128, 128), (96, 96, 144)) or (not double and shape == (192, 192, 192)):
+            return _plan('forward_paired_y', shape, dtype, **options,
+                         pair_order='yx', patient=shape == (96, 96, 144))
+        if double and shape in ((192, 192, 192), (256, 256, 256)):
+            aligned = _plan('forward_native', shape, dtype, **options, direction='r2c', variant='3d')
+            from .forward_scalar import Plan as ScalarPlan
+            return ScalarPlan(aligned)
+        return _plan('forward_half', shape, dtype, **options, direction='r2c',
+                     axis='z', patient=shape in ((80, 80, 112), (96, 96, 96)), xt=False)
+    if shape == (32, 32, 32):
+        return _plan('inverse_small_f64' if double else 'inverse_small_f32', shape, dtype,
+                     **options, direction='c2r', columns_per_block=4, num_warps=4)
+    if shape == (64, 64, 64):
+        return _plan('inverse_cuda_f64' if double else 'inverse_cuda_f32', shape, dtype,
+                     **options, columns=8, warps=8)
+    if shape == (80, 80, 112):
+        if not double:
+            return _plan('inverse_pitched', shape, dtype, **options, variant='xy_z_patient')
+        aligned = _plan('inverse_narrow', shape, dtype, **options, direction='c2r',
+                        strategy='xy_z', patient=True)
+        from .inverse_scalar import ScalarInputPlan
+        return ScalarInputPlan(aligned)
+    if shape == (96, 96, 96):
+        return _plan('inverse_partial_f32', shape, dtype, **options, variant='xy_z_patient')
+    if shape == (96, 96, 144):
+        return _plan('inverse_vector' if double else 'inverse_paired_y', shape, dtype,
+                     **options, variant='3d_patient' if double else 'xy_z_patient')
+    if shape == (128, 128, 128):
+        return _plan('inverse_dense_f64' if double else 'inverse_dense_f32', shape, dtype,
+                     **options, variant='xy_z')
+    if shape == (192, 192, 192):
+        return _plan('inverse_paired_y_f64' if double else 'inverse_paired_y_f32', shape, dtype,
+                     **options, variant='xy_z_patient' if double else 'xy_z')
+    return _plan('inverse_paired_x_f32', shape, dtype, **options,
+                 columns=8, warps=2, frequency=False, patient=False)
+def _shape(tensor, direction, dtype):
+    if tensor.dtype != dtype:
+        raise TypeError('Expected the requested floating-point dtype')
+    if tensor.device.type != 'cuda' or not tensor.is_contiguous():
+        raise ValueError('Expected a contiguous CUDA tensor')
+    if direction == 'r2c':
+        if tensor.ndim != 3:
+            raise ValueError('r2c expects [nz, ny, nx]')
+        shape = tuple(tensor.shape)
+    else:
+        if tensor.ndim != 4 or tensor.shape[-1] != 2:
+            raise ValueError('c2r expects [nz, ny, nx//2+1, 2]')
+        shape = (*tensor.shape[:2], 2 * (tensor.shape[2] - 1))
+    if shape not in _SHAPES:
+        raise ValueError(f'Unsupported B200 grid {shape}; expected one of {sorted(_SHAPES)}')
+    return shape
+class _Launch:
+    def __init__(self, tensor, plan):
+        self.tensor, self.plan = tensor, plan
+        self.output = torch.empty(plan.output_shape, dtype=plan.dtype, device=plan.device)
+    def __call__(self):
+        return self.plan.execute(self.tensor, self.output, stream=torch.cuda.current_stream(self.tensor.device))
+    def close(self):
+        """Release resources after all work, including captured graphs, completes."""
+        self.plan.close()
+    def __del__(self):
+        if not self.plan.closed:
+            try:
+                self.plan.stream.synchronize()
+                self.plan.close()
+            except (RuntimeError, AttributeError):
+                pass
+class _Kernel:
+    def __init__(self, direction, dtype):
+        self.direction, self.dtype = direction, dtype
+        self._plans = {}
+        self._lock = threading.RLock()
+        self._runner = None
+    def _prepare(self, tensor):
+        shape = _shape(tensor, self.direction, self.dtype)
+        stream = torch.cuda.current_stream(tensor.device)
+        owner = _owner(self.direction, shape, self.dtype, tensor.device, stream)
+        try:
+            return Plan(owner)
+        except BaseException:
+            owner.close()
+            raise
+    def _cached(self, tensor):
+        key = (tuple(tensor.shape), tensor.device.index, _cuda_getCurrentRawStream(tensor.device.index))
+        with self._lock:
+            plan = self._plans.get(key)
+            if plan is None:
+                plan = self._prepare(tensor)
+                self._plans[key] = plan
+        return plan.run
+    def run(self, tensor):
+        if self._runner is None:
+            with self._lock:
+                if self._runner is None:
+                    self._runner = load_module().bind_cache(self._cached,
+                        8 if self.dtype == torch.float64 else 4, 3 if self.direction == 'r2c' else 4)
+        return self._runner(tensor)
+    def prepare(self, tensor):
+        """Bind the current input buffer and reuse one output, matching the public API."""
+        return _Launch(tensor, self._prepare(tensor))
+_KERNELS = {(direction, dtype): _Kernel(direction, dtype)
+            for direction in ('r2c', 'c2r') for dtype in (torch.float32, torch.float64)}
+def kernel(direction, dtype):
+    return _KERNELS[direction, dtype]

build/torch-cuda/_b200/c2r_cuda_z_inverse_cuda_kernel.py ADDED Viewed

	@@ -0,0 +1,63 @@

+"""Compile and load fixed CUDA FFT kernels only during plan preparation."""
+import ctypes as ct
+import functools
+import hashlib
+from pathlib import Path
+from types import SimpleNamespace
+from .cuda_z_inverse_lto import library, version, _check
+from .cuda_z_inverse_native_dispatch import load_driver
+@functools.lru_cache(maxsize=128)
+def compile_kernel(shape, columns, warps, arch, device_ordinal):
+    # The ordinal keeps loaded CUDA modules within their owning device context.
+    nz, ny, nx = shape
+    base, odd = 1, nz
+    while odd % 2 == 0:
+        odd //= 2
+        base *= 2
+    if odd not in (1, 3, 5, 7) or base < warps*32//columns:
+        raise ValueError('Unsupported radix decomposition')
+    source = Path(__file__).with_name('c2r_cuda_z_inverse_fft_axis.cu').read_bytes()
+    options = ['--std=c++17', f'--gpu-architecture=sm_{arch}',
+               f'-DLENGTH={nz}', f'-DINNER={ny*(nx//2)}', f'-DBASE={base}',
+               f'-DLOG_BASE={base.bit_length()-1}', f'-DODD={odd}',
+               f'-DCOLUMNS={columns}', f'-DWARPS={warps}']
+    lib = library()
+    lib.nvrtcGetCUBINSize.argtypes = [ct.c_void_p, ct.POINTER(ct.c_size_t)]
+    lib.nvrtcGetCUBIN.argtypes = [ct.c_void_p, ct.c_void_p]
+    program = ct.c_void_p()
+    _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'fft_axis.cu', 0, None, None))
+    try:
+        opts = (ct.c_char_p * len(options))(*(s.encode() for s in options))
+        code = lib.nvrtcCompileProgram(program, len(options), opts)
+        size = ct.c_size_t()
+        _check(lib.nvrtcGetProgramLogSize(program, ct.byref(size)))
+        log = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetProgramLog(program, log))
+        if code:
+            raise RuntimeError(log.value.decode())
+        _check(lib.nvrtcGetCUBINSize(program, ct.byref(size)))
+        blob = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetCUBIN(program, blob))
+    finally:
+        _check(lib.nvrtcDestroyProgram(ct.byref(program)))
+    driver = load_driver()
+    driver.cuModuleLoadData.argtypes = [ct.POINTER(ct.c_void_p), ct.c_void_p]
+    driver.cuModuleGetFunction.argtypes = [ct.POINTER(ct.c_void_p), ct.c_void_p, ct.c_char_p]
+    driver.cuFuncGetAttribute.argtypes = [ct.POINTER(ct.c_int), ct.c_int, ct.c_void_p]
+    module, function = ct.c_void_p(), ct.c_void_p()
+    _check(driver.cuModuleLoadData(ct.byref(module), blob))
+    _check(driver.cuModuleGetFunction(ct.byref(function), module, b'fft_axis'))
+    registers, local = ct.c_int(), ct.c_int()
+    _check(driver.cuFuncGetAttribute(ct.byref(registers), 4, function))
+    _check(driver.cuFuncGetAttribute(ct.byref(local), 3, function))
+    # local is CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, not an instruction count.
+    metadata = dict(options=options, nvrtc=version(), local_memory_bytes=local.value, source_sha256=hashlib.sha256(source).hexdigest(),
+                    cubin_sha256=hashlib.sha256(blob.raw).hexdigest(), compiler_log=log.value.decode())
+    return SimpleNamespace(module=module, function=function.value, name='fft_axis',
+                           hash=metadata['cubin_sha256'], threads=32*warps, shared=0,
+                           n_regs=registers.value, n_spills=local.value,
+                           grid=((ny*(nx//2)+columns-1)//columns, 1, 1),
+                           compiler_metadata=metadata)

build/torch-cuda/_b200/c2r_cuda_z_inverse_fft_axis.cu ADDED Viewed

	@@ -0,0 +1,128 @@

+// SPDX-License-Identifier: Apache-2.0
+// Exact mathematical extents; each warp owns independent adjacent columns.
+__device__ __forceinline__ double2 add(double2 a, double2 b) {
+    return make_double2(a.x + b.x, a.y + b.y);
+}
+__device__ __forceinline__ double2 sub(double2 a, double2 b) {
+    return make_double2(a.x - b.x, a.y - b.y);
+}
+__device__ __forceinline__ double2 mul(double2 a, double2 b) {
+    return make_double2(a.x*b.x-a.y*b.y, a.x*b.y+a.y*b.x);
+}
+template<int H, int N, int G, int C, int Bound>
+__device__ __forceinline__ void stage(double2 (&a)[N/G], const double2 *twiddle,
+                                     int lane_z, double2 *exchange) {
+    if constexpr (H >= G && 4*H <= Bound) {
+        #pragma unroll
+        for (int r = 0; r < N/G; ++r) {
+            if ((r % (4*H/G)) < H/G) {
+                double2 a0 = a[r], a1 = a[r+H/G], a2 = a[r+2*H/G], a3 = a[r+3*H/G];
+                int k = (r*G+lane_z) % H;
+                double2 w1=twiddle[k*(N/(4*H))]; w1.y=-w1.y;
+                double2 w2=twiddle[2*k*(N/(4*H))]; w2.y=-w2.y;
+                double2 w3=twiddle[3*k*(N/(4*H))]; w3.y=-w3.y;
+                // Input is binary bit reversed, so the middle branches swap.
+                double2 b1=mul(a2,w1), b2=mul(a1,w2), b3=mul(a3,w3);
+                double2 even=add(a0,b2), even_diff=sub(a0,b2);
+                double2 odd=add(b1,b3), odd_diff=sub(b1,b3);
+                double2 turn=make_double2(-odd_diff.y,odd_diff.x);
+                a[r]=add(even,odd); a[r+H/G]=add(even_diff,turn);
+                a[r+2*H/G]=sub(even,odd); a[r+3*H/G]=sub(even_diff,turn);
+            }
+        }
+    } else if constexpr (H < G) {
+        if constexpr (H*C >= 32) {
+            #pragma unroll
+            for (int r = 0; r < N/G; ++r) exchange[r*blockDim.x+threadIdx.x]=a[r];
+            __syncthreads();
+        }
+        #pragma unroll
+        for (int r = 0; r < N/G; ++r) {
+            double2 v;
+            if constexpr (H*C < 32) {
+                v = make_double2(__shfl_xor_sync(0xffffffff, a[r].x, H*C),
+                                 __shfl_xor_sync(0xffffffff, a[r].y, H*C));
+            } else {
+                v = exchange[r*blockDim.x+(threadIdx.x ^ (H*C))];
+            }
+            double2 lo = (lane_z & H) ? v : a[r];
+            double2 hi = (lane_z & H) ? a[r] : v;
+            if constexpr (H > 1) {
+                double2 w = twiddle[(lane_z & (H-1)) * (N/(2*H))];
+                w.y = -w.y;
+                hi = mul(hi, w);
+            }
+            a[r] = (lane_z & H) ? sub(lo, hi) : add(lo, hi);
+        }
+        if constexpr (H*C >= 32) __syncthreads();
+    } else {
+        #pragma unroll
+        for (int r = 0; r < N/G; ++r) {
+            if ((r & (H/G)) == 0) {
+                double2 lo = a[r], hi = a[r + H/G];
+                double2 w = twiddle[((r*G + lane_z) & (H-1)) * (N/(2*H))];
+                w.y = -w.y;
+                hi = mul(hi, w);
+                a[r] = add(lo, hi);
+                a[r + H/G] = sub(lo, hi);
+            }
+        }
+    }
+    if constexpr (H >= G && 4*H <= Bound) {
+        if constexpr (4*H < Bound) stage<4*H, N, G, C, Bound>(a, twiddle, lane_z, exchange);
+    } else if constexpr (2*H < Bound) stage<2*H, N, G, C, Bound>(a, twiddle, lane_z, exchange);
+}
+extern "C" __global__ __launch_bounds__(WARPS*32)
+void fft_axis(const double2 *__restrict__ input, float *__restrict__ output,
+              const double2 *__restrict__ twiddle, void *, void *) {
+    constexpr int G = WARPS*32/COLUMNS, R = LENGTH/G;
+    __shared__ double2 exchange[LENGTH*COLUMNS];
+    int lane_z = threadIdx.x / COLUMNS;
+    int column = blockIdx.x * COLUMNS + threadIdx.x % COLUMNS;
+    double2 a[R];
+    #pragma unroll
+    for (int r = 0; r < R; ++r) {
+        int index = r*G + lane_z;
+        int source = (__brev((unsigned)(index % BASE)) >> (32-LOG_BASE))*ODD
+                   + index/BASE;
+        a[r] = column < INNER ? input[source*INNER + column] : make_double2(0., 0.);
+    }
+    stage<1, LENGTH, G, COLUMNS, BASE>(a, twiddle, lane_z, exchange);
+    if constexpr (ODD != 1) {
+        #pragma unroll
+        for (int r = 0; r < BASE/G; ++r) {
+            double2 b[ODD];
+            b[0] = a[r];
+            #pragma unroll
+            for (int t = 1; t < ODD; ++t) {
+                double2 w = twiddle[t*(r*G+lane_z)];
+                w.y = -w.y;
+                b[t] = mul(a[r+t*(BASE/G)], w);
+            }
+            #pragma unroll
+            for (int k = 0; k < ODD; ++k) {
+                double2 v = b[0];
+                #pragma unroll
+                for (int t = 1; t < ODD; ++t) {
+                    double2 w = twiddle[((k*t)%ODD)*BASE];
+                    w.y = -w.y;
+                    v = add(v, mul(b[t], w));
+                }
+                a[r+k*(BASE/G)] = v;
+            }
+        }
+    }
+    #pragma unroll
+    for (int r = 0; r < R; ++r) {
+        if (column < INNER) {
+            int offset = 2*((r*G+lane_z)*INNER+column);
+            if ((reinterpret_cast<unsigned long long>(output) & 7ULL) == 0) {
+                reinterpret_cast<float2 *>(output)[offset/2] = make_float2((float)a[r].x, (float)a[r].y);
+            } else {
+                output[offset] = (float)a[r].x;
+                output[offset+1] = (float)a[r].y;
+            }
+        }
+    }
+}

build/torch-cuda/_b200/c2r_dense_half_inverse_wide_redirect.cu ADDED Viewed

	@@ -0,0 +1,85 @@

+// SPDX-License-Identifier: Apache-2.0
+// Select index width from the complete physical packed extent at compilation.
+// The callback ABI remains 64-bit. Every valid caller index fits the selected
+// unsigned type; pointer arithmetic still computes full-width device addresses.
+#if (1ULL * FFT_NZ * FFT_NY * (FFT_N / 2 + 1)) <= 0x100000000ULL
+using fft_index_t = unsigned int;
+static_assert(sizeof(fft_index_t) == 4, "A 32-bit unsigned integer is required");
+#else
+using fft_index_t = unsigned long long;
+#endif
+#if FFT_DOUBLE
+using caller_scalar_t = double;
+#else
+using caller_scalar_t = float;
+#endif
+__device__ double2 fft_scalar_pair(const caller_scalar_t *input, fft_index_t index) {
+    double2 value;
+#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0
+    value = reinterpret_cast<const double2 *>(input)[index];
+#elif !FFT_DOUBLE && (FFT_INPUT_SHIFT % 2 == 0)
+    // The pair address is aligned for scalar residues zero and two.
+    const float2 packed = reinterpret_cast<const float2 *>(input)[index];
+    value.x = static_cast<double>(packed.x);
+    value.y = static_cast<double>(packed.y);
+#else
+    value.x = static_cast<double>(input[2ULL * index]);
+    value.y = static_cast<double>(input[2ULL * index + 1]);
+#endif
+    return value;
+}
+// The cuFFT logical input is dense. Callback arithmetic reads exact nf-stride
+// caller storage directly; no input repacking is performed.
+// Reconstruct twice the half-length complex spectrum from immutable packed
+// caller input. Endpoint projection preserves irfftn for arbitrary spectra.
+#if FFT_LOAD
+__device__ double2 fft_wide_load(void *data, unsigned long long offset,
+                                    void *caller_info, void *shared_memory) {
+    constexpr fft_index_t m = FFT_N / 2;
+    constexpr fft_index_t nf = m + 1;
+    const fft_index_t index = static_cast<fft_index_t>(offset);
+    const fft_index_t k = index % m;
+    const fft_index_t row = index / m;
+    const fft_index_t y = row % FFT_NY;
+    const fft_index_t z = row / FFT_NY;
+    const fft_index_t mirror_y = y == 0 ? 0 : FFT_NY - y;
+    const fft_index_t mirror_z = z == 0 ? 0 : FFT_NZ - z;
+    const fft_index_t mirror_row = mirror_z * FFT_NY + mirror_y;
+    const caller_scalar_t *input = static_cast<const caller_scalar_t *>(data) + FFT_INPUT_SHIFT;
+    double2 a = fft_scalar_pair(input, row * nf + k);
+    double2 b = fft_scalar_pair(input, mirror_row * nf + m - k);
+    if (k == 0) {
+        // Project each boundary plane against its Z/Y conjugate partner.
+        // b represents the conjugate of the projected Nyquist value.
+        const double2 a_mirror = fft_scalar_pair(input, mirror_row * nf);
+        const double2 b_mirror = fft_scalar_pair(input, row * nf + m);
+        a.x = 0.5 * (a.x + a_mirror.x);
+        a.y = 0.5 * (a.y - a_mirror.y);
+        b.x = 0.5 * (b.x + b_mirror.x);
+        b.y = 0.5 * (b.y - b_mirror.y);
+    }
+    const double2 w = static_cast<const double2 *>(caller_info)[k];
+    double2 result;
+    result.x = a.x + b.x - w.x * (a.y + b.y) + w.y * (a.x - b.x);
+    result.y = a.y - b.y + w.x * (a.x - b.x) + w.y * (a.y + b.y);
+    return result;
+}
+#else
+__device__ void fft_wide_store(void *physical_output, unsigned long long offset,
+                              double2 value, void *caller_info, void *shared_memory) {
+    caller_scalar_t *output = static_cast<caller_scalar_t *>(physical_output) + FFT_INPUT_SHIFT;
+#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0
+    // A nominal cuFFT output address is complex aligned.
+    reinterpret_cast<double2 *>(output)[offset] = value;
+#elif !FFT_DOUBLE && (FFT_INPUT_SHIFT % 2 == 0)
+    // Both components fit one aligned 64-bit store at these scalar residues.
+    float2 packed;
+    packed.x = static_cast<float>(value.x);
+    packed.y = static_cast<float>(value.y);
+    reinterpret_cast<float2 *>(output)[offset] = packed;
+#else
+    output[2 * offset] = static_cast<caller_scalar_t>(value.x);
+    output[2 * offset + 1] = static_cast<caller_scalar_t>(value.y);
+#endif
+}
+#endif

build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_cuda_kernel.py ADDED Viewed

	@@ -0,0 +1,63 @@

+"""Compile and load fixed CUDA FFT kernels only during plan preparation."""
+import ctypes as ct
+import functools
+import hashlib
+from pathlib import Path
+from types import SimpleNamespace
+from .c2r_f64_cuda_z_inverse_lto import library, version, _check
+from .c2r_f64_cuda_z_inverse_native_dispatch import load_driver
+@functools.lru_cache(maxsize=128)
+def compile_kernel(shape, columns, warps, arch, device_ordinal):
+    # The ordinal keeps loaded CUDA modules within their owning device context.
+    nz, ny, nx = shape
+    base, odd = 1, nz
+    while odd % 2 == 0:
+        odd //= 2
+        base *= 2
+    if odd not in (1, 3, 5, 7) or base < warps*32//columns:
+        raise ValueError('Unsupported radix decomposition')
+    source = Path(__file__).with_name('c2r_f64_cuda_z_inverse_fft_axis.cu').read_bytes()
+    options = ['--std=c++17', f'--gpu-architecture=sm_{arch}',
+               f'-DLENGTH={nz}', f'-DINNER={ny*(nx//2)}', f'-DBASE={base}',
+               f'-DLOG_BASE={base.bit_length()-1}', f'-DODD={odd}',
+               f'-DCOLUMNS={columns}', f'-DWARPS={warps}']
+    lib = library()
+    lib.nvrtcGetCUBINSize.argtypes = [ct.c_void_p, ct.POINTER(ct.c_size_t)]
+    lib.nvrtcGetCUBIN.argtypes = [ct.c_void_p, ct.c_void_p]
+    program = ct.c_void_p()
+    _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'fft_axis.cu', 0, None, None))
+    try:
+        opts = (ct.c_char_p * len(options))(*(s.encode() for s in options))
+        code = lib.nvrtcCompileProgram(program, len(options), opts)
+        size = ct.c_size_t()
+        _check(lib.nvrtcGetProgramLogSize(program, ct.byref(size)))
+        log = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetProgramLog(program, log))
+        if code:
+            raise RuntimeError(log.value.decode())
+        _check(lib.nvrtcGetCUBINSize(program, ct.byref(size)))
+        blob = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetCUBIN(program, blob))
+    finally:
+        _check(lib.nvrtcDestroyProgram(ct.byref(program)))
+    driver = load_driver()
+    driver.cuModuleLoadData.argtypes = [ct.POINTER(ct.c_void_p), ct.c_void_p]
+    driver.cuModuleGetFunction.argtypes = [ct.POINTER(ct.c_void_p), ct.c_void_p, ct.c_char_p]
+    driver.cuFuncGetAttribute.argtypes = [ct.POINTER(ct.c_int), ct.c_int, ct.c_void_p]
+    module, function = ct.c_void_p(), ct.c_void_p()
+    _check(driver.cuModuleLoadData(ct.byref(module), blob))
+    _check(driver.cuModuleGetFunction(ct.byref(function), module, b'fft_axis'))
+    registers, local = ct.c_int(), ct.c_int()
+    _check(driver.cuFuncGetAttribute(ct.byref(registers), 4, function))
+    _check(driver.cuFuncGetAttribute(ct.byref(local), 3, function))
+    # local is CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, not an instruction count.
+    metadata = dict(options=options, nvrtc=version(), local_memory_bytes=local.value, source_sha256=hashlib.sha256(source).hexdigest(),
+                    cubin_sha256=hashlib.sha256(blob.raw).hexdigest(), compiler_log=log.value.decode())
+    return SimpleNamespace(module=module, function=function.value, name='fft_axis',
+                           hash=metadata['cubin_sha256'], threads=32*warps, shared=0,
+                           n_regs=registers.value, n_spills=local.value,
+                           grid=((ny*(nx//2)+columns-1)//columns, 1, 1),
+                           compiler_metadata=metadata)

build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_fft_axis.cu ADDED Viewed

	@@ -0,0 +1,128 @@

+// SPDX-License-Identifier: Apache-2.0
+// Exact mathematical extents; each warp owns independent adjacent columns.
+__device__ __forceinline__ double2 add(double2 a, double2 b) {
+    return make_double2(a.x + b.x, a.y + b.y);
+}
+__device__ __forceinline__ double2 sub(double2 a, double2 b) {
+    return make_double2(a.x - b.x, a.y - b.y);
+}
+__device__ __forceinline__ double2 mul(double2 a, double2 b) {
+    return make_double2(a.x*b.x-a.y*b.y, a.x*b.y+a.y*b.x);
+}
+template<int H, int N, int G, int C, int Bound>
+__device__ __forceinline__ void stage(double2 (&a)[N/G], const double2 *twiddle,
+                                     int lane_z, double2 *exchange) {
+    if constexpr (H >= G && 4*H <= Bound) {
+        #pragma unroll
+        for (int r = 0; r < N/G; ++r) {
+            if ((r % (4*H/G)) < H/G) {
+                double2 a0 = a[r], a1 = a[r+H/G], a2 = a[r+2*H/G], a3 = a[r+3*H/G];
+                int k = (r*G+lane_z) % H;
+                double2 w1=twiddle[k*(N/(4*H))]; w1.y=-w1.y;
+                double2 w2=twiddle[2*k*(N/(4*H))]; w2.y=-w2.y;
+                double2 w3=twiddle[3*k*(N/(4*H))]; w3.y=-w3.y;
+                // Input is binary bit reversed, so the middle branches swap.
+                double2 b1=mul(a2,w1), b2=mul(a1,w2), b3=mul(a3,w3);
+                double2 even=add(a0,b2), even_diff=sub(a0,b2);
+                double2 odd=add(b1,b3), odd_diff=sub(b1,b3);
+                double2 turn=make_double2(-odd_diff.y,odd_diff.x);
+                a[r]=add(even,odd); a[r+H/G]=add(even_diff,turn);
+                a[r+2*H/G]=sub(even,odd); a[r+3*H/G]=sub(even_diff,turn);
+            }
+        }
+    } else if constexpr (H < G) {
+        if constexpr (H*C >= 32) {
+            #pragma unroll
+            for (int r = 0; r < N/G; ++r) exchange[r*blockDim.x+threadIdx.x]=a[r];
+            __syncthreads();
+        }
+        #pragma unroll
+        for (int r = 0; r < N/G; ++r) {
+            double2 v;
+            if constexpr (H*C < 32) {
+                v = make_double2(__shfl_xor_sync(0xffffffff, a[r].x, H*C),
+                                 __shfl_xor_sync(0xffffffff, a[r].y, H*C));
+            } else {
+                v = exchange[r*blockDim.x+(threadIdx.x ^ (H*C))];
+            }
+            double2 lo = (lane_z & H) ? v : a[r];
+            double2 hi = (lane_z & H) ? a[r] : v;
+            if constexpr (H > 1) {
+                double2 w = twiddle[(lane_z & (H-1)) * (N/(2*H))];
+                w.y = -w.y;
+                hi = mul(hi, w);
+            }
+            a[r] = (lane_z & H) ? sub(lo, hi) : add(lo, hi);
+        }
+        if constexpr (H*C >= 32) __syncthreads();
+    } else {
+        #pragma unroll
+        for (int r = 0; r < N/G; ++r) {
+            if ((r & (H/G)) == 0) {
+                double2 lo = a[r], hi = a[r + H/G];
+                double2 w = twiddle[((r*G + lane_z) & (H-1)) * (N/(2*H))];
+                w.y = -w.y;
+                hi = mul(hi, w);
+                a[r] = add(lo, hi);
+                a[r + H/G] = sub(lo, hi);
+            }
+        }
+    }
+    if constexpr (H >= G && 4*H <= Bound) {
+        if constexpr (4*H < Bound) stage<4*H, N, G, C, Bound>(a, twiddle, lane_z, exchange);
+    } else if constexpr (2*H < Bound) stage<2*H, N, G, C, Bound>(a, twiddle, lane_z, exchange);
+}
+extern "C" __global__ __launch_bounds__(WARPS*32)
+void fft_axis(const double2 *__restrict__ input, double *__restrict__ output,
+              const double2 *__restrict__ twiddle, void *, void *) {
+    constexpr int G = WARPS*32/COLUMNS, R = LENGTH/G;
+    __shared__ double2 exchange[LENGTH*COLUMNS];
+    int lane_z = threadIdx.x / COLUMNS;
+    int column = blockIdx.x * COLUMNS + threadIdx.x % COLUMNS;
+    double2 a[R];
+    #pragma unroll
+    for (int r = 0; r < R; ++r) {
+        int index = r*G + lane_z;
+        int source = (__brev((unsigned)(index % BASE)) >> (32-LOG_BASE))*ODD
+                   + index/BASE;
+        a[r] = column < INNER ? input[source*INNER + column] : make_double2(0., 0.);
+    }
+    stage<1, LENGTH, G, COLUMNS, BASE>(a, twiddle, lane_z, exchange);
+    if constexpr (ODD != 1) {
+        #pragma unroll
+        for (int r = 0; r < BASE/G; ++r) {
+            double2 b[ODD];
+            b[0] = a[r];
+            #pragma unroll
+            for (int t = 1; t < ODD; ++t) {
+                double2 w = twiddle[t*(r*G+lane_z)];
+                w.y = -w.y;
+                b[t] = mul(a[r+t*(BASE/G)], w);
+            }
+            #pragma unroll
+            for (int k = 0; k < ODD; ++k) {
+                double2 v = b[0];
+                #pragma unroll
+                for (int t = 1; t < ODD; ++t) {
+                    double2 w = twiddle[((k*t)%ODD)*BASE];
+                    w.y = -w.y;
+                    v = add(v, mul(b[t], w));
+                }
+                a[r+k*(BASE/G)] = v;
+            }
+        }
+    }
+    #pragma unroll
+    for (int r = 0; r < R; ++r) {
+        if (column < INNER) {
+            int offset = 2*((r*G+lane_z)*INNER+column);
+            if ((reinterpret_cast<unsigned long long>(output) & 15ULL) == 0) {
+                reinterpret_cast<double2 *>(output)[offset/2] = a[r];
+            } else {
+                output[offset] = a[r].x;
+                output[offset+1] = a[r].y;
+            }
+        }
+    }
+}

build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_lto.py ADDED Viewed

	@@ -0,0 +1,110 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Compile cuFFT LTO callbacks with the installed NVRTC, before execution.
+This module can compile offline without initializing CUDA or loading code on
+a GPU. Compiled blobs are cached by source, public shape, precision, NVRTC
+version and target. No input tensor values or pointers are cached.
+"""
+from __future__ import annotations
+import ctypes as ct
+import ctypes.util
+import functools
+import hashlib
+from pathlib import Path
+import threading
+_LOCK = threading.Lock()
+@functools.lru_cache(maxsize=1)
+def library():
+    candidates = [ctypes.util.find_library('nvrtc')]
+    for root in (Path('/usr/local/cuda/lib64'),
+                 Path('/usr/local/cuda/targets/x86_64-linux/lib')):
+        candidates.extend(str(p) for p in sorted(root.glob('libnvrtc.so*')))
+    candidates.extend(('libnvrtc.so.13', 'libnvrtc.so.12', 'libnvrtc.so'))
+    lib = None
+    for name in candidates:
+        if name:
+            try:
+                lib = ct.CDLL(name)
+                break
+            except OSError:
+                pass
+    if lib is None:
+        raise RuntimeError('The installed NVRTC library is unavailable')
+    signatures = {
+        'nvrtcVersion': [ct.POINTER(ct.c_int), ct.POINTER(ct.c_int)],
+        'nvrtcCreateProgram': [ct.POINTER(ct.c_void_p), ct.c_char_p, ct.c_char_p,
+                               ct.c_int, ct.POINTER(ct.c_char_p), ct.POINTER(ct.c_char_p)],
+        'nvrtcCompileProgram': [ct.c_void_p, ct.c_int, ct.POINTER(ct.c_char_p)],
+        'nvrtcGetProgramLogSize': [ct.c_void_p, ct.POINTER(ct.c_size_t)],
+        'nvrtcGetProgramLog': [ct.c_void_p, ct.c_void_p],
+        'nvrtcGetLTOIRSize': [ct.c_void_p, ct.POINTER(ct.c_size_t)],
+        'nvrtcGetLTOIR': [ct.c_void_p, ct.c_void_p],
+        'nvrtcDestroyProgram': [ct.POINTER(ct.c_void_p)],
+    }
+    for name, signature in signatures.items():
+        function = getattr(lib, name)
+        function.argtypes = signature
+        function.restype = ct.c_int
+    return lib
+def _check(result):
+    if result:
+        raise RuntimeError(f'NVRTC returned error {result}')
+def version():
+    major, minor = ct.c_int(), ct.c_int()
+    _check(library().nvrtcVersion(ct.byref(major), ct.byref(minor)))
+    return major.value, minor.value
+@functools.lru_cache(maxsize=128)
+def _compile(source, shape, double, arch, nvrtc_version, input_shift, kind):
+    nz, ny, nx = shape
+    lib = library()
+    options = ('--std=c++17', '--relocatable-device-code=true', '-dlto',
+               f'--gpu-architecture=compute_{arch}', f'-DFFT_N={nx}',
+               f'-DFFT_NZ={nz}', f'-DFFT_NY={ny}',
+               f'-DFFT_DOUBLE={int(double)}', f'-DFFT_INPUT_SHIFT={input_shift}',
+               f'-DFFT_LOAD={int(kind == "load")}')
+    program = ct.c_void_p()
+    _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'endpoint.cu', 0, None, None))
+    try:
+        option_array = (ct.c_char_p * len(options))(*(v.encode() for v in options))
+        result = lib.nvrtcCompileProgram(program, len(options), option_array)
+        log_size = ct.c_size_t()
+        _check(lib.nvrtcGetProgramLogSize(program, ct.byref(log_size)))
+        log = ct.create_string_buffer(log_size.value)
+        _check(lib.nvrtcGetProgramLog(program, log))
+        if result:
+            raise RuntimeError(f'NVRTC returned error {result}: {log.value.decode()}')
+        size = ct.c_size_t()
+        _check(lib.nvrtcGetLTOIRSize(program, ct.byref(size)))
+        data = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetLTOIR(program, data))
+        blob = data.raw
+        metadata = dict(shape=shape, double=bool(double), target=f'compute_{arch}',
+                        nvrtc=list(nvrtc_version), options=list(options),
+                        source_sha256=hashlib.sha256(source).hexdigest(),
+                        lto_ir_sha256=hashlib.sha256(blob).hexdigest(),
+                        lto_ir_bytes=len(blob), compiler_log=log.value.decode())
+        return blob, metadata
+    finally:
+        _check(lib.nvrtcDestroyProgram(ct.byref(program)))
+def compile_callback(shape, input_shift, kind, arch=100):
+    """Reconstruct a half-length spectrum in FP64 or store its inverse as FP64."""
+    if kind not in ('load', 'store'):
+        raise ValueError('Expected a load or store callback')
+    source = Path(__file__).with_name('c2r_f64_cuda_z_inverse_wide_redirect.cu').read_bytes()
+    with _LOCK:
+        blob, metadata = _compile(source, tuple(shape), True, int(arch), version(), input_shift, kind)
+    return (('fft_wide_' + kind).encode(), 1 if kind == 'load' else 5,
+            ct.create_string_buffer(blob, len(blob)), dict(metadata, kind=kind))

build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_native_dispatch.py ADDED Viewed

	@@ -0,0 +1,100 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Prepare a combined cuFFT and CUDA dispatcher before execution."""
+from __future__ import annotations
+import ctypes as ct
+import fcntl
+import hashlib
+import os
+from pathlib import Path
+import shlex
+import subprocess
+import tempfile
+import threading
+class Descriptor(ct.Structure):
+    _fields_ = [('complex_exec', ct.c_void_p), ('launch', ct.c_void_p),
+                ('handle', ct.c_int), ('intermediate', ct.c_uint64),
+                ('twiddle', ct.c_uint64), ('stream', ct.c_void_p),
+                ('function', ct.c_void_p), ('grid', ct.c_uint * 3),
+                ('threads', ct.c_uint), ('shared', ct.c_uint)]
+_LIBRARY = None
+_DRIVER = None
+_LOCK = threading.Lock()
+def load_library():
+    """Build with the installed host compiler and validate structure size."""
+    global _LIBRARY
+    with _LOCK:
+        if _LIBRARY is not None:
+            return _LIBRARY
+        source = Path(__file__).with_name('cuda_z_inverse_native_dispatch.c')
+        compiler = shlex.split(os.environ.get('CC', 'cc'))
+        flags = ['-std=c11', '-O3', '-fPIC', '-shared', '-Wall', '-Wextra', '-Werror']
+        key = hashlib.sha256(source.read_bytes() + repr((compiler, flags)).encode()).hexdigest()
+        cache = Path(tempfile.gettempdir()) / ('fft-half-triton-z-native-' + str(os.getuid())) / key
+        cache.mkdir(parents=True, exist_ok=True)
+        target = cache / 'dispatch.so'
+        with (cache / 'build.lock').open('w') as lock:
+            fcntl.flock(lock, fcntl.LOCK_EX)
+            if not target.is_file():
+                pending = cache / ('dispatch.' + str(os.getpid()) + '.so')
+                try:
+                    subprocess.run(compiler + flags + [str(source), '-o', str(pending)],
+                                   check=True, capture_output=True, text=True)
+                    pending.replace(target)
+                finally:
+                    pending.unlink(missing_ok=True)
+        lib = ct.CDLL(str(target))
+        lib.fft_execute.argtypes = [ct.POINTER(Descriptor), ct.c_uint64, ct.c_uint64]
+        lib.fft_execute.restype = ct.c_int
+        lib.fft_dispatch_size.argtypes = []
+        lib.fft_dispatch_size.restype = ct.c_size_t
+        if lib.fft_dispatch_size() != ct.sizeof(Descriptor):
+            raise RuntimeError('Native dispatcher structure size mismatch')
+        _LIBRARY = lib
+        return lib
+def load_driver():
+    """Resolve the driver function during preparation only."""
+    global _DRIVER
+    with _LOCK:
+        if _DRIVER is None:
+            _DRIVER = ct.CDLL('libcuda.so.1')
+            _DRIVER.cuLaunchKernel.restype = ct.c_int
+        return _DRIVER
+class PreparedDispatch:
+    """Immutable setup descriptor, with fresh caller pointers at each call.
+    The Plan retains the cuFFT handles and compiled CUDA kernel. Concurrent
+    execution requires independent plans because cuFFT workspace is reused.
+    """
+    def __init__(self, complex_exec, handle, compiled, grid, intermediate, twiddle, stream):
+        abi = dict(threads=compiled.threads, shared_bytes=compiled.shared)
+        if not compiled.function or compiled.module is None:
+            raise ValueError('The CUDA kernel must be loaded during setup')
+        if not 0 < handle < 2**31:
+            raise ValueError('A prepared cuFFT handle is required')
+        if len(grid) != 3 or any(not 0 < n < 2**32 for n in grid):
+            raise ValueError('Invalid launch grid')
+        if not intermediate or not twiddle:
+            raise ValueError('Prepared intermediate and twiddle storage is required')
+        self.library = load_library()
+        self.driver = load_driver()
+        self.descriptor = Descriptor(
+            complex_exec=ct.cast(complex_exec, ct.c_void_p).value,
+            launch=ct.cast(self.driver.cuLaunchKernel, ct.c_void_p).value,
+            handle=handle, intermediate=intermediate,
+            twiddle=twiddle, stream=stream, function=compiled.function,
+            grid=(ct.c_uint * 3)(*grid), threads=abi['threads'], shared=abi['shared_bytes'])
+        self.pointer = ct.pointer(self.descriptor)
+        self.function = self.library.fft_execute
+    def execute(self, input_ptr, output_ptr):
+        return self.function(self.pointer, input_ptr, output_ptr)

build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_wide_redirect.cu ADDED Viewed

	@@ -0,0 +1,65 @@

+// SPDX-License-Identifier: Apache-2.0
+// Select index width from the complete physical packed extent at compilation.
+// The callback ABI remains 64-bit. Every valid caller index fits the selected
+// unsigned type; pointer arithmetic still computes full-width device addresses.
+#if (1ULL * FFT_NZ * FFT_NY * (FFT_N / 2 + 1)) <= 0x100000000ULL
+using fft_index_t = unsigned int;
+static_assert(sizeof(fft_index_t) == 4, "A 32-bit unsigned integer is required");
+#else
+using fft_index_t = unsigned long long;
+#endif
+__device__ double2 fft_scalar_pair(const double *input, fft_index_t index) {
+    double2 value;
+#if FFT_INPUT_SHIFT == 0
+    value = reinterpret_cast<const double2 *>(input)[index];
+#else
+    value.x = input[2ULL * index];
+    value.y = input[2ULL * index + 1];
+#endif
+    return value;
+}
+// Reconstruct twice the half-length complex spectrum from immutable packed
+// caller input. Endpoint projection preserves irfftn for arbitrary spectra.
+#if FFT_LOAD
+__device__ double2 fft_wide_load(void *data, unsigned long long offset,
+                                    void *caller_info, void *shared_memory) {
+    constexpr fft_index_t m = FFT_N / 2;
+    constexpr fft_index_t nf = m + 1;
+    const fft_index_t index = static_cast<fft_index_t>(offset);
+    const fft_index_t k = index % m;
+    const fft_index_t row = index / m;
+    const fft_index_t y = row % FFT_NY;
+    const fft_index_t z = row / FFT_NY;
+    const fft_index_t mirror_y = y == 0 ? 0 : FFT_NY - y;
+    const fft_index_t mirror_z = z == 0 ? 0 : FFT_NZ - z;
+    const fft_index_t mirror_row = mirror_z * FFT_NY + mirror_y;
+    const double *input = static_cast<const double *>(data) + FFT_INPUT_SHIFT;
+    double2 a = fft_scalar_pair(input, row*nf+k);
+    double2 b = fft_scalar_pair(input, mirror_row * nf + m - k);
+    if (k == 0) {
+        // Project each boundary plane against its Z/Y conjugate partner.
+        // b represents the conjugate of the projected Nyquist value.
+        const double2 a_mirror = fft_scalar_pair(input, mirror_row * nf);
+        const double2 b_mirror = fft_scalar_pair(input, row * nf + m);
+        a.x = 0.5 * (a.x + a_mirror.x);
+        a.y = 0.5 * (a.y - a_mirror.y);
+        b.x = 0.5 * (b.x + b_mirror.x);
+        b.y = 0.5 * (b.y - b_mirror.y);
+    }
+    const unsigned long long table = static_cast<const unsigned long long *>(caller_info)[1];
+    const double2 w = reinterpret_cast<const double2 *>(table)[k];
+    double2 result;
+    result.x = a.x + b.x - w.x * (a.y + b.y) + w.y * (a.x - b.x);
+    result.y = a.y - b.y + w.x * (a.x - b.x) + w.y * (a.y + b.y);
+    return result;
+}
+#else
+__device__ void fft_wide_store(void *physical_output, unsigned long long offset,
+                              double2 value, void *caller_info, void *shared_memory) {
+    const unsigned long long address =
+        *static_cast<const unsigned long long *>(caller_info);
+    double *output = reinterpret_cast<double *>(address);
+    output[2 * offset] = static_cast<double>(value.x);
+    output[2 * offset + 1] = static_cast<double>(value.y);
+}
+#endif

build/torch-cuda/_b200/c2r_f64_dense_half_inverse_lto.py ADDED Viewed

	@@ -0,0 +1,110 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Compile cuFFT LTO callbacks with the installed NVRTC, before execution.
+This module can compile offline without initializing CUDA or loading code on
+a GPU. Compiled blobs are cached by source, public shape, precision, NVRTC
+version and target. No input tensor values or pointers are cached.
+"""
+from __future__ import annotations
+import ctypes as ct
+import ctypes.util
+import functools
+import hashlib
+from pathlib import Path
+import threading
+_LOCK = threading.Lock()
+@functools.lru_cache(maxsize=1)
+def library():
+    candidates = [ctypes.util.find_library('nvrtc')]
+    for root in (Path('/usr/local/cuda/lib64'),
+                 Path('/usr/local/cuda/targets/x86_64-linux/lib')):
+        candidates.extend(str(p) for p in sorted(root.glob('libnvrtc.so*')))
+    candidates.extend(('libnvrtc.so.13', 'libnvrtc.so.12', 'libnvrtc.so'))
+    lib = None
+    for name in candidates:
+        if name:
+            try:
+                lib = ct.CDLL(name)
+                break
+            except OSError:
+                pass
+    if lib is None:
+        raise RuntimeError('The installed NVRTC library is unavailable')
+    signatures = {
+        'nvrtcVersion': [ct.POINTER(ct.c_int), ct.POINTER(ct.c_int)],
+        'nvrtcCreateProgram': [ct.POINTER(ct.c_void_p), ct.c_char_p, ct.c_char_p,
+                               ct.c_int, ct.POINTER(ct.c_char_p), ct.POINTER(ct.c_char_p)],
+        'nvrtcCompileProgram': [ct.c_void_p, ct.c_int, ct.POINTER(ct.c_char_p)],
+        'nvrtcGetProgramLogSize': [ct.c_void_p, ct.POINTER(ct.c_size_t)],
+        'nvrtcGetProgramLog': [ct.c_void_p, ct.c_void_p],
+        'nvrtcGetLTOIRSize': [ct.c_void_p, ct.POINTER(ct.c_size_t)],
+        'nvrtcGetLTOIR': [ct.c_void_p, ct.c_void_p],
+        'nvrtcDestroyProgram': [ct.POINTER(ct.c_void_p)],
+    }
+    for name, signature in signatures.items():
+        function = getattr(lib, name)
+        function.argtypes = signature
+        function.restype = ct.c_int
+    return lib
+def _check(result):
+    if result:
+        raise RuntimeError(f'NVRTC returned error {result}')
+def version():
+    major, minor = ct.c_int(), ct.c_int()
+    _check(library().nvrtcVersion(ct.byref(major), ct.byref(minor)))
+    return major.value, minor.value
+@functools.lru_cache(maxsize=128)
+def _compile(source, shape, double, arch, nvrtc_version, input_shift, kind):
+    nz, ny, nx = shape
+    lib = library()
+    options = ('--std=c++17', '--relocatable-device-code=true', '-dlto',
+               f'--gpu-architecture=compute_{arch}', f'-DFFT_N={nx}',
+               f'-DFFT_NZ={nz}', f'-DFFT_NY={ny}',
+               f'-DFFT_DOUBLE={int(double)}', f'-DFFT_INPUT_SHIFT={input_shift}',
+               f'-DFFT_LOAD={int(kind == "load")}')
+    program = ct.c_void_p()
+    _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'endpoint.cu', 0, None, None))
+    try:
+        option_array = (ct.c_char_p * len(options))(*(v.encode() for v in options))
+        result = lib.nvrtcCompileProgram(program, len(options), option_array)
+        log_size = ct.c_size_t()
+        _check(lib.nvrtcGetProgramLogSize(program, ct.byref(log_size)))
+        log = ct.create_string_buffer(log_size.value)
+        _check(lib.nvrtcGetProgramLog(program, log))
+        if result:
+            raise RuntimeError(f'NVRTC returned error {result}: {log.value.decode()}')
+        size = ct.c_size_t()
+        _check(lib.nvrtcGetLTOIRSize(program, ct.byref(size)))
+        data = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetLTOIR(program, data))
+        blob = data.raw
+        metadata = dict(shape=shape, double=bool(double), target=f'compute_{arch}',
+                        nvrtc=list(nvrtc_version), options=list(options),
+                        source_sha256=hashlib.sha256(source).hexdigest(),
+                        lto_ir_sha256=hashlib.sha256(blob).hexdigest(),
+                        lto_ir_bytes=len(blob), compiler_log=log.value.decode())
+        return blob, metadata
+    finally:
+        _check(lib.nvrtcDestroyProgram(ct.byref(program)))
+def compile_callback(shape, input_shift, kind, arch=100):
+    """Reconstruct a half-length spectrum in FP64 or store its inverse as FP64."""
+    if kind not in ('load', 'store'):
+        raise ValueError('Expected a load or store callback')
+    source = Path(__file__).with_name('c2r_f64_dense_half_inverse_wide_redirect.cu').read_bytes()
+    with _LOCK:
+        blob, metadata = _compile(source, tuple(shape), True, int(arch), version(), input_shift, kind)
+    return (('fft_wide_' + kind).encode(), 1 if kind == 'load' else 5,
+            ct.create_string_buffer(blob, len(blob)), dict(metadata, kind=kind))

build/torch-cuda/_b200/c2r_f64_dense_half_inverse_wide_redirect.cu ADDED Viewed

	@@ -0,0 +1,67 @@

+// SPDX-License-Identifier: Apache-2.0
+// Select index width from the complete physical packed extent at compilation.
+// The callback ABI remains 64-bit. Every valid caller index fits the selected
+// unsigned type; pointer arithmetic still computes full-width device addresses.
+#if (1ULL * FFT_NZ * FFT_NY * (FFT_N / 2 + 1)) <= 0x100000000ULL
+using fft_index_t = unsigned int;
+static_assert(sizeof(fft_index_t) == 4, "A 32-bit unsigned integer is required");
+#else
+using fft_index_t = unsigned long long;
+#endif
+__device__ double2 fft_scalar_pair(const double *input, fft_index_t index) {
+    double2 value;
+#if FFT_INPUT_SHIFT == 0
+    value = reinterpret_cast<const double2 *>(input)[index];
+#else
+    value.x = input[2ULL * index];
+    value.y = input[2ULL * index + 1];
+#endif
+    return value;
+}
+// The cuFFT logical input is dense. Callback arithmetic reads exact nf-stride
+// caller storage directly; no input repacking is performed.
+// Reconstruct twice the half-length complex spectrum from immutable packed
+// caller input. Endpoint projection preserves irfftn for arbitrary spectra.
+#if FFT_LOAD
+__device__ double2 fft_wide_load(void *data, unsigned long long offset,
+                                    void *caller_info, void *shared_memory) {
+    constexpr fft_index_t m = FFT_N / 2;
+    constexpr fft_index_t nf = m + 1;
+    const fft_index_t index = static_cast<fft_index_t>(offset);
+    const fft_index_t k = index % m;
+    const fft_index_t row = index / m;
+    const fft_index_t y = row % FFT_NY;
+    const fft_index_t z = row / FFT_NY;
+    const fft_index_t mirror_y = y == 0 ? 0 : FFT_NY - y;
+    const fft_index_t mirror_z = z == 0 ? 0 : FFT_NZ - z;
+    const fft_index_t mirror_row = mirror_z * FFT_NY + mirror_y;
+    const double *input = static_cast<const double *>(data) + FFT_INPUT_SHIFT;
+    double2 a = fft_scalar_pair(input, row * nf + k);
+    double2 b = fft_scalar_pair(input, mirror_row * nf + m - k);
+    if (k == 0) {
+        // Project each boundary plane against its Z/Y conjugate partner.
+        // b represents the conjugate of the projected Nyquist value.
+        const double2 a_mirror = fft_scalar_pair(input, mirror_row * nf);
+        const double2 b_mirror = fft_scalar_pair(input, row * nf + m);
+        a.x = 0.5 * (a.x + a_mirror.x);
+        a.y = 0.5 * (a.y - a_mirror.y);
+        b.x = 0.5 * (b.x + b_mirror.x);
+        b.y = 0.5 * (b.y - b_mirror.y);
+    }
+    const unsigned long long table = static_cast<const unsigned long long *>(caller_info)[1];
+    const double2 w = reinterpret_cast<const double2 *>(table)[k];
+    double2 result;
+    result.x = a.x + b.x - w.x * (a.y + b.y) + w.y * (a.x - b.x);
+    result.y = a.y - b.y + w.x * (a.x - b.x) + w.y * (a.y + b.y);
+    return result;
+}
+#else
+__device__ void fft_wide_store(void *physical_output, unsigned long long offset,
+                              double2 value, void *caller_info, void *shared_memory) {
+    const unsigned long long address =
+        *static_cast<const unsigned long long *>(caller_info);
+    double *output = reinterpret_cast<double *>(address);
+    output[2 * offset] = static_cast<double>(value.x);
+    output[2 * offset + 1] = static_cast<double>(value.y);
+}
+#endif

build/torch-cuda/_b200/c2r_f64_half_vector_inverse_wide_redirect.cu ADDED Viewed

	@@ -0,0 +1,65 @@

+// SPDX-License-Identifier: Apache-2.0
+// Select index width from the complete physical packed extent at compilation.
+// The callback ABI remains 64-bit. Every valid caller index fits the selected
+// unsigned type; pointer arithmetic still computes full-width device addresses.
+#if (1ULL * FFT_NZ * FFT_NY * (FFT_N / 2 + 1)) <= 0x100000000ULL
+using fft_index_t = unsigned int;
+static_assert(sizeof(fft_index_t) == 4, "A 32-bit unsigned integer is required");
+#else
+using fft_index_t = unsigned long long;
+#endif
+__device__ double2 fft_scalar_pair(const double *input, fft_index_t index) {
+    double2 value;
+#if FFT_INPUT_SHIFT == 0
+    value = reinterpret_cast<const double2 *>(input)[index];
+#else
+    value.x = input[2ULL * index];
+    value.y = input[2ULL * index + 1];
+#endif
+    return value;
+}
+// Reconstruct twice the half-length complex spectrum from immutable packed
+// caller input. Endpoint projection preserves irfftn for arbitrary spectra.
+#if FFT_LOAD
+__device__ double2 fft_wide_load(void *data, unsigned long long offset,
+                                    void *caller_info, void *shared_memory) {
+    constexpr fft_index_t m = FFT_N / 2;
+    constexpr fft_index_t nf = m + 1;
+    const fft_index_t index = static_cast<fft_index_t>(offset);
+    const fft_index_t k = index % nf;
+    const fft_index_t row = index / nf;
+    const fft_index_t y = row % FFT_NY;
+    const fft_index_t z = row / FFT_NY;
+    const fft_index_t mirror_y = y == 0 ? 0 : FFT_NY - y;
+    const fft_index_t mirror_z = z == 0 ? 0 : FFT_NZ - z;
+    const fft_index_t mirror_row = mirror_z * FFT_NY + mirror_y;
+    const double *input = static_cast<const double *>(data) + FFT_INPUT_SHIFT;
+    double2 a = fft_scalar_pair(input, index);
+    double2 b = fft_scalar_pair(input, mirror_row * nf + m - k);
+    if (k == 0) {
+        // Project each boundary plane against its Z/Y conjugate partner.
+        // b represents the conjugate of the projected Nyquist value.
+        const double2 a_mirror = fft_scalar_pair(input, mirror_row * nf);
+        const double2 b_mirror = fft_scalar_pair(input, row * nf + m);
+        a.x = 0.5 * (a.x + a_mirror.x);
+        a.y = 0.5 * (a.y - a_mirror.y);
+        b.x = 0.5 * (b.x + b_mirror.x);
+        b.y = 0.5 * (b.y - b_mirror.y);
+    }
+    const unsigned long long table = static_cast<const unsigned long long *>(caller_info)[1];
+    const double2 w = reinterpret_cast<const double2 *>(table)[k];
+    double2 result;
+    result.x = a.x + b.x - w.x * (a.y + b.y) + w.y * (a.x - b.x);
+    result.y = a.y - b.y + w.x * (a.x - b.x) + w.y * (a.y + b.y);
+    return result;
+}
+#else
+__device__ void fft_wide_store(void *physical_output, unsigned long long offset,
+                              double2 value, void *caller_info, void *shared_memory) {
+    const unsigned long long address =
+        *static_cast<const unsigned long long *>(caller_info);
+    double *output = reinterpret_cast<double *>(address);
+    output[2 * offset] = static_cast<double>(value.x);
+    output[2 * offset + 1] = static_cast<double>(value.y);
+}
+#endif

build/torch-cuda/_b200/c2r_f64_narrow_index_inverse_scalar_output_reconstruct.cu ADDED Viewed

	@@ -0,0 +1,45 @@

+// SPDX-License-Identifier: Apache-2.0
+// Only completed Z/Y coefficients are read. All constants and callback plans
+// are prepared before execution; no caller address or data is cached.
+#if FFT_CALLBACK_KIND != 2
+__device__ double2 fft_half_load(void *data, unsigned long long offset,
+                               void *caller_info, void *shared_memory) {
+    constexpr unsigned long long m = FFT_N / 2;
+#if FFT_ROWMAJOR
+    const unsigned long long row = offset / (m + 1);
+    const unsigned long long k = offset % (m + 1);
+    const unsigned long long mirror = row * (m + 1) + m - k;
+#else
+    const unsigned long long k = offset / FFT_ROWS;
+    const unsigned long long row = offset % FFT_ROWS;
+    const unsigned long long mirror = (m - k) * FFT_ROWS + row;
+#endif
+    const double2 a = static_cast<const double2 *>(data)[offset];
+    const double2 b = static_cast<const double2 *>(data)[mirror];
+    const double2 w = static_cast<const double2 *>(caller_info)[k];
+    // Endpoint imaginary parts are ignored after the Z/Y inverse, matching
+    // irfftn for arbitrary input (including inconsistent endpoint planes).
+    const double ai = k == 0 ? 0.0 : a.y;
+    const double bi = k == 0 ? 0.0 : b.y;
+    double2 result;
+    result.x = a.x + b.x - w.x * (ai + bi) + w.y * (a.x - b.x);
+    result.y = ai - bi + w.x * (a.x - b.x) + w.y * (ai + bi);
+    return result;
+}
+#endif
+#if FFT_CALLBACK_KIND != 1
+__device__ void fft_half_store(void *data, unsigned long long offset,
+                                double2 value, void *caller_info,
+                                void *shared_memory) {
+    // The API receives an aligned base. A setup-specialized scalar offset
+    // writes the exact caller address; no preceding scalar is accessed.
+    // Re/Im of the M-point result are the even/odd N-point real outputs.
+    volatile double *output = static_cast<volatile double *>(data);
+    output[2 * offset + FFT_STORE_SHIFT] = value.x;
+    output[2 * offset + 1 + FFT_STORE_SHIFT] = value.y;
+}
+#endif

build/torch-cuda/_b200/c2r_f64_paired_x_inverse_cuda_kernel.py ADDED Viewed

	@@ -0,0 +1,64 @@

+"""Compile and load fixed CUDA FFT kernels only during plan preparation."""
+import ctypes as ct
+import functools
+import hashlib
+from pathlib import Path
+from types import SimpleNamespace
+from .paired_x_inverse_lto import library, version, _check
+from .paired_x_inverse_native_dispatch import load_driver
+@functools.lru_cache(maxsize=128)
+def compile_kernel(shape, columns, warps, arch, device_ordinal, frequency, double):
+    # The ordinal keeps loaded CUDA modules within their owning device context.
+    nz, ny, nx = shape
+    base, odd = 1, nx // 2
+    while odd % 2 == 0:
+        odd //= 2
+        base *= 2
+    if odd != 1 or base < 2*warps*32//columns:
+        raise ValueError('Unsupported radix decomposition')
+    source = Path(__file__).with_name('c2r_f64_paired_x_inverse_fft_axis.cu').read_bytes()
+    options = ['--std=c++17', f'--gpu-architecture=sm_{arch}',
+               f'-DLENGTH={nx//2}', f'-DINNER={nz*ny}', f'-DBASE={base}',
+               f'-DNZ={nz}', f'-DNY={ny}', f'-DFREQUENCY={int(frequency)}', f'-DDOUBLE={int(double)}',
+               f'-DLOG_BASE={base.bit_length()-1}', f'-DODD={odd}',
+               f'-DCOLUMNS={columns}', f'-DWARPS={warps}']
+    lib = library()
+    lib.nvrtcGetCUBINSize.argtypes = [ct.c_void_p, ct.POINTER(ct.c_size_t)]
+    lib.nvrtcGetCUBIN.argtypes = [ct.c_void_p, ct.c_void_p]
+    program = ct.c_void_p()
+    _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'fft_axis.cu', 0, None, None))
+    try:
+        opts = (ct.c_char_p * len(options))(*(s.encode() for s in options))
+        code = lib.nvrtcCompileProgram(program, len(options), opts)
+        size = ct.c_size_t()
+        _check(lib.nvrtcGetProgramLogSize(program, ct.byref(size)))
+        log = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetProgramLog(program, log))
+        if code:
+            raise RuntimeError(log.value.decode())
+        _check(lib.nvrtcGetCUBINSize(program, ct.byref(size)))
+        blob = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetCUBIN(program, blob))
+    finally:
+        _check(lib.nvrtcDestroyProgram(ct.byref(program)))
+    driver = load_driver()
+    driver.cuModuleLoadData.argtypes = [ct.POINTER(ct.c_void_p), ct.c_void_p]
+    driver.cuModuleGetFunction.argtypes = [ct.POINTER(ct.c_void_p), ct.c_void_p, ct.c_char_p]
+    driver.cuFuncGetAttribute.argtypes = [ct.POINTER(ct.c_int), ct.c_int, ct.c_void_p]
+    module, function = ct.c_void_p(), ct.c_void_p()
+    _check(driver.cuModuleLoadData(ct.byref(module), blob))
+    _check(driver.cuModuleGetFunction(ct.byref(function), module, b'fft_axis'))
+    registers, local = ct.c_int(), ct.c_int()
+    _check(driver.cuFuncGetAttribute(ct.byref(registers), 4, function))
+    _check(driver.cuFuncGetAttribute(ct.byref(local), 3, function))
+    # local is CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, not an instruction count.
+    metadata = dict(options=options, nvrtc=version(), local_memory_bytes=local.value, source_sha256=hashlib.sha256(source).hexdigest(),
+                    cubin_sha256=hashlib.sha256(blob.raw).hexdigest(), compiler_log=log.value.decode())
+    return SimpleNamespace(module=module, function=function.value, name='fft_axis',
+                           hash=metadata['cubin_sha256'], threads=32*warps, shared=0,
+                           n_regs=registers.value, n_spills=local.value,
+                           grid=((nz*ny//2+2+columns//2-1)//(columns//2), 1, 1),
+                           compiler_metadata=metadata)

build/torch-cuda/_b200/c2r_f64_paired_x_inverse_fft_axis.cu ADDED Viewed

	@@ -0,0 +1,270 @@

+#define FIRST_RADIX 4
+// SPDX-License-Identifier: Apache-2.0
+// Exact mathematical extents; each warp owns independent adjacent columns.
+__device__ __forceinline__ double2 add(double2 a, double2 b) {
+    return make_double2(a.x + b.x, a.y + b.y);
+}
+__device__ __forceinline__ double2 sub(double2 a, double2 b) {
+    return make_double2(a.x - b.x, a.y - b.y);
+}
+__device__ __forceinline__ double2 mul(double2 a, double2 b) {
+    return make_double2(a.x*b.x-a.y*b.y, a.x*b.y+a.y*b.x);
+}
+template<int H, int N, int G, int C, int Bound>
+__device__ __forceinline__ void stage(double2 (&a)[N/G], const double2 *twiddle,
+                                     int lane_z, double2 *exchange) {
+    if constexpr (H >= G && 4*H <= Bound) {
+        #pragma unroll
+        for (int r = 0; r < N/G; ++r) {
+            if ((r % (4*H/G)) < H/G) {
+                double2 a0 = a[r], a1 = a[r+H/G], a2 = a[r+2*H/G], a3 = a[r+3*H/G];
+                int k = (r*G+lane_z) % H;
+                double2 w1=twiddle[k*(N/(4*H))]; w1.y=-w1.y;
+                double2 w2=twiddle[2*k*(N/(4*H))]; w2.y=-w2.y;
+                double2 w3=twiddle[3*k*(N/(4*H))]; w3.y=-w3.y;
+                // Input is binary bit reversed, so the middle branches swap.
+                double2 b1=mul(a2,w1), b2=mul(a1,w2), b3=mul(a3,w3);
+                double2 even=add(a0,b2), even_diff=sub(a0,b2);
+                double2 odd=add(b1,b3), odd_diff=sub(b1,b3);
+                double2 turn=make_double2(-odd_diff.y,odd_diff.x);
+                a[r]=add(even,odd); a[r+H/G]=add(even_diff,turn);
+                a[r+2*H/G]=sub(even,odd); a[r+3*H/G]=sub(even_diff,turn);
+            }
+        }
+    } else if constexpr (H < G) {
+        if constexpr (H*C >= 32) {
+            #pragma unroll
+            for (int r = 0; r < N/G; ++r) exchange[r*blockDim.x+threadIdx.x]=a[r];
+            __syncthreads();
+        }
+        #pragma unroll
+        for (int r = 0; r < N/G; ++r) {
+            double2 v;
+            if constexpr (H*C < 32) {
+                v = make_double2(__shfl_xor_sync(0xffffffff, a[r].x, H*C),
+                                 __shfl_xor_sync(0xffffffff, a[r].y, H*C));
+            } else {
+                v = exchange[r*blockDim.x+(threadIdx.x ^ (H*C))];
+            }
+            double2 lo = (lane_z & H) ? v : a[r];
+            double2 hi = (lane_z & H) ? a[r] : v;
+            if constexpr (H > 1) {
+                double2 w = twiddle[(lane_z & (H-1)) * (N/(2*H))];
+                w.y = -w.y;
+                hi = mul(hi, w);
+            }
+            a[r] = (lane_z & H) ? sub(lo, hi) : add(lo, hi);
+        }
+        if constexpr (H*C >= 32) __syncthreads();
+    } else {
+        #pragma unroll
+        for (int r = 0; r < N/G; ++r) {
+            if ((r & (H/G)) == 0) {
+                double2 lo = a[r], hi = a[r + H/G];
+                double2 w = twiddle[((r*G + lane_z) & (H-1)) * (N/(2*H))];
+                w.y = -w.y;
+                hi = mul(hi, w);
+                a[r] = add(lo, hi);
+                a[r + H/G] = sub(lo, hi);
+            }
+        }
+    }
+    if constexpr (H >= G && 4*H <= Bound) {
+        if constexpr (4*H < Bound) stage<4*H, N, G, C, Bound>(a, twiddle, lane_z, exchange);
+    } else if constexpr (2*H < Bound) stage<2*H, N, G, C, Bound>(a, twiddle, lane_z, exchange);
+}
+#if DOUBLE
+using scalar_t = double;
+#else
+using scalar_t = float;
+#endif
+__device__ __forceinline__ double2 pair_load(const scalar_t *input, int offset) {
+#if DOUBLE
+    if ((reinterpret_cast<unsigned long long>(input) & 15ULL) == 0)
+        return reinterpret_cast<const double2 *>(input)[offset];
+#else
+    if ((reinterpret_cast<unsigned long long>(input) & 7ULL) == 0) {
+        float2 v = reinterpret_cast<const float2 *>(input)[offset];
+        return make_double2((double)v.x, (double)v.y);
+    }
+#endif
+    return make_double2((double)input[2*offset], (double)input[2*offset+1]);
+}
+__device__ __forceinline__ double2 reconstruct(const scalar_t *input, int row, int k,
+                                               const double2 *twiddle) {
+    constexpr int nf = LENGTH+1;
+    int y = row%NY, z = row/NY;
+    int mirror = (z == 0 ? 0 : NZ-z)*NY + (y == 0 ? 0 : NY-y);
+    double2 a = pair_load(input, row*nf+k);
+    double2 b = pair_load(input, mirror*nf+LENGTH-k);
+    if (k == 0) {
+        double2 am = pair_load(input, mirror*nf);
+        double2 bm = pair_load(input, row*nf+LENGTH);
+        a = make_double2(0.5*(a.x+am.x), 0.5*(a.y-am.y));
+        b = make_double2(0.5*(b.x+bm.x), 0.5*(b.y-bm.y));
+    }
+    double2 w = twiddle[k];
+    return make_double2(a.x+b.x-w.x*(a.y+b.y)+w.y*(a.x-b.x),
+                        a.y-b.y+w.x*(a.x-b.x)+w.y*(a.y+b.y));
+}
+// The table contains shape-only row representatives after the FFT roots.
+__device__ __forceinline__ int paired_row(const double2 *twiddle, int pair, int side) {
+    if (pair >= INNER/2 + 2) return INNER;
+    int row = pair <= NY/2 ? pair : pair + NY/2 - 1;
+    int z = row/NY, y = row%NY;
+    int mirror = (z == 0 ? 0 : NZ-z)*NY + (y == 0 ? 0 : NY-y);
+    return side ? mirror : row;
+}
+// This bijection is applied only within a reconstructed mathematical row.
+__device__ __forceinline__ int shared_index(int k) { return k ^ (k >> 4); }
+__device__ __forceinline__ void exact_four(double2 (&a)[4]) {
+    double2 even = add(a[0],a[2]), even_diff = sub(a[0],a[2]);
+    double2 odd = add(a[1],a[3]), odd_diff = sub(a[1],a[3]);
+    double2 turn = make_double2(-odd_diff.y,odd_diff.x);
+    a[0] = add(even,odd); a[1] = add(even_diff,turn);
+    a[2] = sub(even,odd); a[3] = sub(even_diff,turn);
+}
+template<int F>
+__device__ __forceinline__ void first_transform(double2 (&a)[F], const double2 *) {
+    if constexpr (F == 4) {
+        exact_four(a);
+    } else {
+        static_assert(F == 8);
+        double2 even[4] = {a[0],a[2],a[4],a[6]};
+        double2 odd[4] = {a[1],a[3],a[5],a[7]};
+        exact_four(even); exact_four(odd);
+        constexpr double c = 0x1.6a09e667f3bcdp-1;
+        #pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            double2 v;
+            if (j == 0) v = odd[j];
+            if (j == 1) v = make_double2(c*(odd[j].x-odd[j].y),c*(odd[j].x+odd[j].y));
+            if (j == 2) v = make_double2(-odd[j].y,odd[j].x);
+            if (j == 3) v = make_double2(-c*(odd[j].x+odd[j].y),c*(odd[j].x-odd[j].y));
+            a[j] = add(even[j],v); a[j+4] = sub(even[j],v);
+        }
+    }
+}
+extern "C" __global__ __launch_bounds__(WARPS*32)
+void fft_axis(const scalar_t *__restrict__ input, double2 *__restrict__ output,
+              const double2 *__restrict__ twiddle, unsigned long long *slot, scalar_t *caller_output) {
+    constexpr int G = WARPS*32/COLUMNS, R = LENGTH/G;
+    __shared__ double2 exchange[LENGTH*COLUMNS];
+    int lane_z = threadIdx.x % G;
+    int local_row = threadIdx.x / G;
+    int pair = blockIdx.x*(COLUMNS/2) + local_row/2;
+    int column = paired_row(twiddle, pair, local_row%2);
+    bool active = column < INNER && (local_row%2 == 0 ||
+                    column != paired_row(twiddle, pair, 0));
+    // Each group reconstructs F evenly spaced frequencies for both rows,
+    // completes their first log2(F) butterflies, and exchanges that result.
+    constexpr int F = FIRST_RADIX, D = LENGTH/F;
+    int pair_lane = threadIdx.x % (2*G);
+    int pair_local = threadIdx.x / (2*G);
+    int row_a = paired_row(twiddle, pair, 0);
+    int row_b = paired_row(twiddle, pair, 1);
+    #pragma unroll
+    for (int r = 0; r < (D+2*G-1)/(2*G); ++r) {
+        int k = r*(2*G) + pair_lane;
+        if (k < D) {
+            double2 aa[F], bb[F];
+            #pragma unroll
+            for (int j = 0; j < F; ++j) {
+                int frequency = k+j*D;
+                double2 ca = make_double2(0., 0.), cb = ca;
+                if (row_a < INNER) {
+                    double2 av = pair_load(input, row_a*(LENGTH+1)+frequency);
+                    double2 bv = pair_load(input, row_b*(LENGTH+1)+LENGTH-frequency);
+                    if (frequency == 0) {
+                        double2 am = pair_load(input, row_b*(LENGTH+1));
+                        double2 bm = pair_load(input, row_a*(LENGTH+1)+LENGTH);
+                        av = make_double2(0.5*(av.x+am.x), 0.5*(av.y-am.y));
+                        bv = make_double2(0.5*(bv.x+bm.x), 0.5*(bv.y-bm.y));
+                        double sx = av.x+bv.x, sy = av.y+bv.y;
+                        double dx = av.x-bv.x, dy = av.y-bv.y;
+                        ca = make_double2(sx-sy, dy+dx);
+                        cb = make_double2(sx+sy, dx-dy);
+                    } else {
+                        double2 w = twiddle[LENGTH+frequency];
+                        double sx = av.x+bv.x, sy = av.y+bv.y;
+                        double dx = av.x-bv.x, dy = av.y-bv.y;
+                        double tr = -w.x*sy+w.y*dx;
+                        double ti = w.x*dx+w.y*sy;
+                        ca = make_double2(sx+tr, dy+ti);
+                        cb = make_double2(sx-tr, ti-dy);
+                    }
+                }
+                aa[j] = ca;
+                if (k == 0) bb[(F-j)%F] = cb;
+                else bb[F-1-j] = cb;
+            }
+            first_transform<F>(aa, twiddle);
+            first_transform<F>(bb, twiddle);
+            int base_a = __brev((unsigned)k) >> (32-LOG_BASE);
+            int base_b = __brev((unsigned)(k == 0 ? 0 : D-k)) >> (32-LOG_BASE);
+            #pragma unroll
+            for (int j = 0; j < F; ++j) {
+                exchange[(2*pair_local)*LENGTH+shared_index(base_a+j)] = aa[j];
+                exchange[(2*pair_local+1)*LENGTH+shared_index(base_b+j)] = bb[j];
+            }
+        }
+    }
+    __syncthreads();
+    double2 a[R];
+    #pragma unroll
+    for (int r = 0; r < R; ++r) {
+        int index = r*G + lane_z;
+        a[r] = exchange[local_row*LENGTH+shared_index(index)];
+    }
+    __syncthreads();
+    stage<FIRST_RADIX, LENGTH, G, 1, BASE>(a, twiddle, lane_z, exchange);
+    if constexpr (ODD != 1) {
+        #pragma unroll
+        for (int r = 0; r < BASE/G; ++r) {
+            double2 b[ODD];
+            b[0] = a[r];
+            #pragma unroll
+            for (int t = 1; t < ODD; ++t) {
+                double2 w = twiddle[t*(r*G+lane_z)];
+                w.y = -w.y;
+                b[t] = mul(a[r+t*(BASE/G)], w);
+            }
+            #pragma unroll
+            for (int k = 0; k < ODD; ++k) {
+                double2 v = b[0];
+                #pragma unroll
+                for (int t = 1; t < ODD; ++t) {
+                    double2 w = twiddle[((k*t)%ODD)*BASE];
+                    w.y = -w.y;
+                    v = add(v, mul(b[t], w));
+                }
+                a[r+k*(BASE/G)] = v;
+            }
+        }
+    }
+#if FREQUENCY
+    #pragma unroll
+    for (int r = 0; r < R; ++r) {
+        exchange[(threadIdx.x/G)*LENGTH + r*G + lane_z] = a[r];
+    }
+    __syncthreads();
+    #pragma unroll
+    for (int r = 0; r < R; ++r) {
+        int index = r*(WARPS*32)+threadIdx.x;
+        int row_local = index%COLUMNS;
+        int k = index/COLUMNS;
+        int row = blockIdx.x*COLUMNS+row_local;
+        if (row < INNER) output[k*INNER+row] = exchange[row_local*LENGTH+k];
+    }
+#else
+    #pragma unroll
+    for (int r = 0; r < R; ++r) {
+        if (active) output[((column%NY)*NZ+column/NY)*LENGTH+r*G+lane_z] = a[r];
+    }
+#endif
+}

build/torch-cuda/_b200/c2r_f64_paired_x_inverse_wide_redirect.cu ADDED Viewed

	@@ -0,0 +1,29 @@

+// SPDX-License-Identifier: Apache-2.0
+// The nominal FP64 output remains full-sized for cuFFT internal stages.
+#if FFT_DOUBLE
+using scalar_t = double;
+#else
+using scalar_t = float;
+#endif
+#if FFT_MIDDLE
+__device__ void fft_wide_middle(void *, unsigned long long offset, double2 value,
+                                void *caller_info, void *) {
+    constexpr unsigned int m = FFT_N/2;
+    unsigned int index = static_cast<unsigned int>(offset);
+    unsigned int y = index%FFT_NY, k=(index/FFT_NY)%m, z=index/(FFT_NY*m);
+    static_cast<double2 *>(caller_info)[(z*FFT_NY+y)*m+k] = value;
+}
+#else
+__device__ void fft_wide_store(void *physical_output, unsigned long long offset,
+                               double2 value, void *, void *) {
+    scalar_t *output = static_cast<scalar_t *>(physical_output) + FFT_INPUT_SHIFT;
+#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0
+    reinterpret_cast<double2 *>(output)[offset] = value;
+#elif !FFT_DOUBLE && FFT_INPUT_SHIFT % 2 == 0
+    reinterpret_cast<float2 *>(output)[offset] = make_float2((float)value.x,(float)value.y);
+#else
+    output[2*offset] = static_cast<scalar_t>(value.x);
+    output[2*offset+1] = static_cast<scalar_t>(value.y);
+#endif
+}
+#endif

build/torch-cuda/_b200/c2r_f64_paired_y_direct_inverse_lto.py ADDED Viewed

	@@ -0,0 +1,116 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Compile cuFFT LTO callbacks with the installed NVRTC, before execution.
+This module can compile offline without initializing CUDA or loading code on
+a GPU. Compiled blobs are cached by source, public shape, precision, NVRTC
+version and target. No input tensor values or pointers are cached.
+"""
+from __future__ import annotations
+import ctypes as ct
+import ctypes.util
+import functools
+import hashlib
+from pathlib import Path
+import threading
+_LOCK = threading.Lock()
+@functools.lru_cache(maxsize=1)
+def library():
+    candidates = [ctypes.util.find_library('nvrtc')]
+    for root in (Path('/usr/local/cuda/lib64'),
+                 Path('/usr/local/cuda/targets/x86_64-linux/lib')):
+        candidates.extend(str(p) for p in sorted(root.glob('libnvrtc.so*')))
+    candidates.extend(('libnvrtc.so.13', 'libnvrtc.so.12', 'libnvrtc.so'))
+    lib = None
+    for name in candidates:
+        if name:
+            try:
+                lib = ct.CDLL(name)
+                break
+            except OSError:
+                pass
+    if lib is None:
+        raise RuntimeError('The installed NVRTC library is unavailable')
+    signatures = {
+        'nvrtcVersion': [ct.POINTER(ct.c_int), ct.POINTER(ct.c_int)],
+        'nvrtcCreateProgram': [ct.POINTER(ct.c_void_p), ct.c_char_p, ct.c_char_p,
+                               ct.c_int, ct.POINTER(ct.c_char_p), ct.POINTER(ct.c_char_p)],
+        'nvrtcCompileProgram': [ct.c_void_p, ct.c_int, ct.POINTER(ct.c_char_p)],
+        'nvrtcGetProgramLogSize': [ct.c_void_p, ct.POINTER(ct.c_size_t)],
+        'nvrtcGetProgramLog': [ct.c_void_p, ct.c_void_p],
+        'nvrtcGetLTOIRSize': [ct.c_void_p, ct.POINTER(ct.c_size_t)],
+        'nvrtcGetLTOIR': [ct.c_void_p, ct.c_void_p],
+        'nvrtcDestroyProgram': [ct.POINTER(ct.c_void_p)],
+    }
+    for name, signature in signatures.items():
+        function = getattr(lib, name)
+        function.argtypes = signature
+        function.restype = ct.c_int
+    return lib
+def _check(result):
+    if result:
+        raise RuntimeError(f'NVRTC returned error {result}')
+def version():
+    major, minor = ct.c_int(), ct.c_int()
+    _check(library().nvrtcVersion(ct.byref(major), ct.byref(minor)))
+    return major.value, minor.value
+@functools.lru_cache(maxsize=128)
+def _compile(source, shape, double, arch, nvrtc_version, input_shift, kind):
+    nz, ny, nx = shape
+    from .dense_half_inverse_schedule import twiddles
+    roots = twiddles(ny)
+    initializer = ','.join(value.hex() for value in roots)
+    declaration = ('__device__ __constant__ __align__(16) double fft_twiddles[] = {'
+                   + initializer + '};\n')
+    source = declaration.encode() + source
+    lib = library()
+    options = ('--std=c++17', '--relocatable-device-code=true', '-dlto',
+               f'--gpu-architecture=compute_{arch}', f'-DFFT_N={nx}',
+               f'-DFFT_NZ={nz}', f'-DFFT_NY={ny}',
+               f'-DFFT_DOUBLE={int(double)}', f'-DFFT_INPUT_SHIFT={input_shift}',
+               f'-DFFT_LOAD={int(kind == "load")}')
+    program = ct.c_void_p()
+    _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'endpoint.cu', 0, None, None))
+    try:
+        option_array = (ct.c_char_p * len(options))(*(v.encode() for v in options))
+        result = lib.nvrtcCompileProgram(program, len(options), option_array)
+        log_size = ct.c_size_t()
+        _check(lib.nvrtcGetProgramLogSize(program, ct.byref(log_size)))
+        log = ct.create_string_buffer(log_size.value)
+        _check(lib.nvrtcGetProgramLog(program, log))
+        if result:
+            raise RuntimeError(f'NVRTC returned error {result}: {log.value.decode()}')
+        size = ct.c_size_t()
+        _check(lib.nvrtcGetLTOIRSize(program, ct.byref(size)))
+        data = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetLTOIR(program, data))
+        blob = data.raw
+        metadata = dict(shape=shape, double=bool(double), target=f'compute_{arch}',
+                        nvrtc=list(nvrtc_version), options=list(options),
+                        source_sha256=hashlib.sha256(source).hexdigest(),
+                        lto_ir_sha256=hashlib.sha256(blob).hexdigest(),
+                        lto_ir_bytes=len(blob), compiler_log=log.value.decode())
+        return blob, metadata
+    finally:
+        _check(lib.nvrtcDestroyProgram(ct.byref(program)))
+def compile_callback(shape, input_shift, kind, arch=100, double=False):
+    """Reconstruct a half-length spectrum in FP64 or store its inverse as FP32."""
+    if kind not in ('load', 'store'):
+        raise ValueError('Expected a load or store callback')
+    source = Path(__file__).with_name('c2r_f64_paired_y_direct_inverse_wide_redirect.cu').read_bytes()
+    with _LOCK:
+        blob, metadata = _compile(source, tuple(shape), bool(double), int(arch), version(), input_shift, kind)
+    return (('fft_wide_' + kind).encode(), 1 if kind == 'load' else 5,
+            ct.create_string_buffer(blob, len(blob)), dict(metadata, kind=kind))

build/torch-cuda/_b200/c2r_f64_paired_y_direct_inverse_wide_redirect.cu ADDED Viewed

	@@ -0,0 +1,114 @@

+// SPDX-License-Identifier: Apache-2.0
+// Reconstruct a real-pair transform from arbitrary tightly packed X spectra.
+#define PACK_Y 1
+#if (1ULL * FFT_NZ * FFT_NY * FFT_N) <= 0x100000000ULL
+using index_t = unsigned int;
+#else
+using index_t = unsigned long long;
+#endif
+#if FFT_DOUBLE
+using scalar_t = double;
+#else
+using scalar_t = float;
+#endif
+static __device__ __forceinline__ index_t negative(index_t k, index_t n) {
+    return k == 0 ? 0 : n - k;
+}
+static __device__ __forceinline__ double2 pair_load(const scalar_t *input, index_t k) {
+#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0
+    return __ldcs(reinterpret_cast<const double2 *>(input)+k);
+#elif !FFT_DOUBLE && FFT_INPUT_SHIFT % 2 == 0
+    const float2 v = __ldcs(reinterpret_cast<const float2 *>(input)+k);
+    return make_double2(static_cast<double>(v.x), static_cast<double>(v.y));
+#else
+    return make_double2(static_cast<double>(__ldcs(input+2ULL*k)),
+                        static_cast<double>(__ldcs(input+2ULL*k+1)));
+#endif
+}
+static __device__ __forceinline__ double2 full_spectrum(
+        const scalar_t *input, index_t z, index_t y, index_t x) {
+    constexpr index_t nf = FFT_N/2 + 1;
+    const bool conjugate = x > FFT_N/2;
+    if (conjugate) {
+        z = negative(z, FFT_NZ);
+        y = negative(y, FFT_NY);
+        x = FFT_N-x;
+    }
+    double2 a = pair_load(input, (z*FFT_NY+y)*nf+x);
+    if (x == 0 || x == FFT_N/2) {
+        const double2 mirror = pair_load(input,
+            (negative(z,FFT_NZ)*FFT_NY+negative(y,FFT_NY))*nf+x);
+        a.x = 0.5*a.x+0.5*mirror.x;
+        a.y = 0.5*a.y-0.5*mirror.y;
+    }
+    if (conjugate) a.y = -a.y;
+    return a;
+}
+#if FFT_LOAD
+__device__ double2 fft_wide_load(void *data, unsigned long long offset,
+                                void *caller_info, void *) {
+    constexpr index_t hy = PACK_Y ? FFT_NY/2 : FFT_NY;
+    const index_t index = static_cast<index_t>(offset);
+    const index_t x = index % FFT_N;
+    const index_t y = (index / FFT_N) % hy;
+    const index_t z = index / (hy*FFT_N);
+    const scalar_t *input = static_cast<const scalar_t *>(data)+FFT_INPUT_SHIFT;
+    // Both packed frequencies use one common X reflection. Endpoint
+    // projection is needed only for the two original X boundary planes.
+    const bool conjugate = x > FFT_N/2;
+    const index_t px = conjugate ? FFT_N-x : x;
+    const index_t az = conjugate ? negative(z,FFT_NZ) : z;
+    const index_t ay = conjugate ? negative(y,FFT_NY) : y;
+#if PACK_Y
+    const index_t bz = az;
+    const index_t by = conjugate ? negative(y+FFT_NY/2,FFT_NY) : y+FFT_NY/2;
+    const index_t k = y;
+#else
+    const index_t bz = conjugate ? negative(z+FFT_NZ/2,FFT_NZ) : z+FFT_NZ/2;
+    const index_t by = ay;
+    const index_t k = z;
+#endif
+    constexpr index_t nf = FFT_N/2+1;
+    double2 a = pair_load(input,(az*FFT_NY+ay)*nf+px);
+    double2 b = pair_load(input,(bz*FFT_NY+by)*nf+px);
+    if (px == 0 || px == FFT_N/2) {
+        const double2 am = pair_load(input,
+            (negative(az,FFT_NZ)*FFT_NY+negative(ay,FFT_NY))*nf+px);
+        const double2 bm = pair_load(input,
+            (negative(bz,FFT_NZ)*FFT_NY+negative(by,FFT_NY))*nf+px);
+        a.x = 0.5*a.x+0.5*am.x;
+        a.y = 0.5*a.y-0.5*am.y;
+        b.x = 0.5*b.x+0.5*bm.x;
+        b.y = 0.5*b.y-0.5*bm.y;
+    }
+    if (conjugate) a.y = -a.y;
+    else b.y = -b.y;
+    const double2 w = reinterpret_cast<const double2 *>(fft_twiddles)[k];
+    // Preserve FP64 arithmetic while selecting the rounding sequence.
+    const double sr = a.x + b.x, dr = a.x - b.x;
+    const double si = a.y + b.y, di = a.y - b.y;
+    double2 result;
+    result.x = __fma_rn(w.y, dr, __dadd_rn(sr, -__dmul_rn(w.x, si)));
+    result.y = __fma_rn(w.y, si, __dadd_rn(di, __dmul_rn(w.x, dr)));
+    return result;
+}
+#else
+__device__ void fft_wide_store(void *data, unsigned long long offset,
+                               double2 value, void *, void *) {
+    constexpr index_t hy = PACK_Y ? FFT_NY/2 : FFT_NY;
+    const index_t index = static_cast<index_t>(offset);
+    const index_t x = index % FFT_N;
+    const index_t y = (index / FFT_N) % hy;
+    const index_t z = index / (hy*FFT_N);
+    scalar_t *output = static_cast<scalar_t *>(data)+FFT_INPUT_SHIFT;
+#if PACK_Y
+    const index_t even = (z*FFT_NY+2*y)*FFT_N+x;
+    constexpr index_t distance = FFT_N;
+#else
+    const index_t even = (2*z*FFT_NY+y)*FFT_N+x;
+    constexpr index_t distance = FFT_NY*FFT_N;
+#endif
+    __stcs(output+even, static_cast<scalar_t>(value.x));
+    __stcs(output+even+distance, static_cast<scalar_t>(value.y));
+}
+#endif

build/torch-cuda/_b200/c2r_f64_partial_y_inverse_wide_redirect.cu ADDED Viewed

	@@ -0,0 +1,114 @@

+// SPDX-License-Identifier: Apache-2.0
+// Select index width from the complete physical packed extent at compilation.
+// The callback ABI remains 64-bit. Every valid caller index fits the selected
+// unsigned type; pointer arithmetic still computes full-width device addresses.
+#if (1ULL * FFT_NZ * FFT_NY * (FFT_N / 2 + 1)) <= 0x100000000ULL
+using fft_index_t = unsigned int;
+static_assert(sizeof(fft_index_t) == 4, "A 32-bit unsigned integer is required");
+#else
+using fft_index_t = unsigned long long;
+#endif
+#if FFT_DOUBLE
+using caller_scalar_t = double;
+#else
+using caller_scalar_t = float;
+#endif
+static __device__ __forceinline__ double2 fft_scalar_pair(const caller_scalar_t *input, fft_index_t index) {
+    double2 value;
+#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0
+    value = reinterpret_cast<const double2 *>(input)[index];
+#elif !FFT_DOUBLE && (FFT_INPUT_SHIFT % 2 == 0)
+    // The pair address is aligned for scalar residues zero and two.
+    const float2 packed = reinterpret_cast<const float2 *>(input)[index];
+    value.x = static_cast<double>(packed.x);
+    value.y = static_cast<double>(packed.y);
+#else
+    value.x = static_cast<double>(input[2ULL * index]);
+    value.y = static_cast<double>(input[2ULL * index + 1]);
+#endif
+    return value;
+}
+// The cuFFT logical input is dense. Callback arithmetic reads exact nf-stride
+// caller storage directly; no input repacking is performed.
+// Reconstruct twice the half-length complex spectrum from immutable packed
+// caller input. Endpoint projection preserves irfftn for arbitrary spectra.
+#if FFT_LOAD
+__device__ double2 fft_wide_load(void *data, unsigned long long offset,
+                                    void *caller_info, void *shared_memory) {
+    constexpr fft_index_t m = FFT_N / 2;
+    constexpr fft_index_t nf = m + 1;
+    const fft_index_t index = static_cast<fft_index_t>(offset);
+    const fft_index_t k = index % m;
+    constexpr fft_index_t short_y = FFT_NY / FFT_RADIX;
+    const fft_index_t logical_row = index / m;
+    const fft_index_t j = logical_row % short_y;
+    const fft_index_t r = (logical_row / short_y) % FFT_RADIX;
+    const fft_index_t z = logical_row / FFT_NY;
+    const fft_index_t y = FFT_RADIX * j + r;
+    const fft_index_t row = z * FFT_NY + y;
+    const fft_index_t mirror_y = y == 0 ? 0 : FFT_NY - y;
+    const fft_index_t mirror_z = z == 0 ? 0 : FFT_NZ - z;
+    const fft_index_t mirror_row = mirror_z * FFT_NY + mirror_y;
+    const caller_scalar_t *input = static_cast<const caller_scalar_t *>(data) + FFT_INPUT_SHIFT;
+    double2 a = fft_scalar_pair(input, row * nf + k);
+    double2 b = fft_scalar_pair(input, mirror_row * nf + m - k);
+    if (k == 0) {
+        // Project each boundary plane against its Z/Y conjugate partner.
+        // b represents the conjugate of the projected Nyquist value.
+        const double2 a_mirror = fft_scalar_pair(input, mirror_row * nf);
+        const double2 b_mirror = fft_scalar_pair(input, row * nf + m);
+        a.x = 0.5 * (a.x + a_mirror.x);
+        a.y = 0.5 * (a.y - a_mirror.y);
+        b.x = 0.5 * (b.x + b_mirror.x);
+        b.y = 0.5 * (b.y - b_mirror.y);
+    }
+    const double2 w = static_cast<const double2 *>(caller_info)[k];
+    double2 result;
+    result.x = a.x + b.x - w.x * (a.y + b.y) + w.y * (a.x - b.x);
+    result.y = a.y - b.y + w.x * (a.x - b.x) + w.y * (a.y + b.y);
+    return result;
+}
+#elif FFT_FINISH
+// Complete the factored inverse Y transform while reading mathematical XY
+// results for the Z FFT. The intermediate has no padding or repack pass.
+__device__ double2 fft_wide_finish(void *data, unsigned long long offset,
+                                 void *caller_info, void *shared_memory) {
+    constexpr fft_index_t m = FFT_N / 2;
+    constexpr fft_index_t short_y = FFT_NY / FFT_RADIX;
+    const fft_index_t index = static_cast<fft_index_t>(offset);
+    const fft_index_t x = index % m;
+    const fft_index_t row = index / m;
+    const fft_index_t y = row % FFT_NY;
+    const fft_index_t z = row / FFT_NY;
+    const fft_index_t q = y % short_y;
+    const double2 *input = static_cast<const double2 *>(data);
+    const double2 *table = static_cast<const double2 *>(caller_info) + FFT_N;
+    double2 result = input[(z * FFT_NY + q) * m + x];
+    #pragma unroll
+    for (fft_index_t r = 1; r < FFT_RADIX; ++r) {
+        const double2 v = input[(z * FFT_NY + r * short_y + q) * m + x];
+        const double2 w = table[(r * y) % FFT_NY];
+        result.x += v.x * w.x + v.y * w.y;
+        result.y += v.y * w.x - v.x * w.y;
+    }
+    return result;
+}
+#else
+__device__ void fft_wide_store(void *physical_output, unsigned long long offset,
+                              double2 value, void *caller_info, void *shared_memory) {
+    caller_scalar_t *output = static_cast<caller_scalar_t *>(physical_output) + FFT_INPUT_SHIFT;
+#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0
+    // A nominal cuFFT output address is complex aligned.
+    reinterpret_cast<double2 *>(output)[offset] = value;
+#elif !FFT_DOUBLE && (FFT_INPUT_SHIFT % 2 == 0)
+    // Both components fit one aligned 64-bit store at these scalar residues.
+    float2 packed;
+    packed.x = static_cast<float>(value.x);
+    packed.y = static_cast<float>(value.y);
+    reinterpret_cast<float2 *>(output)[offset] = packed;
+#else
+    output[2 * offset] = static_cast<caller_scalar_t>(value.x);
+    output[2 * offset + 1] = static_cast<caller_scalar_t>(value.y);
+#endif
+}
+#endif

build/torch-cuda/_b200/c2r_f64_scalar_input_inverse_scalar_io.cu ADDED Viewed

	@@ -0,0 +1,26 @@

+// SPDX-License-Identifier: Apache-2.0
+// API addresses are rounded down to 16 bytes. Compile-time scalar residues
+// recover the exact caller address; no scalar preceding the caller is used.
+#if FFT_DOUBLE
+using scalar = double;
+#else
+using scalar = float;
+#endif
+#if IS_LOAD
+__device__ double2 fft_scalar_load(void *data, unsigned long long offset,
+                                  void *caller_info, void *shared_memory) {
+    const volatile scalar *input = static_cast<const volatile scalar *>(data);
+    double2 value;
+    value.x = static_cast<double>(input[2 * offset + SCALAR_SHIFT]);
+    value.y = static_cast<double>(input[2 * offset + SCALAR_SHIFT + 1]);
+    return value;
+}
+#else
+__device__ void fft_scalar_store(void *data, unsigned long long offset,
+                                 double value, void *caller_info,
+                                 void *shared_memory) {
+    static_cast<volatile scalar *>(data)[offset + SCALAR_SHIFT] =
+        static_cast<scalar>(value);
+}
+#endif

build/torch-cuda/_b200/c2r_paired_y_direct_inverse_wide_redirect.cu ADDED Viewed

	@@ -0,0 +1,109 @@

+// SPDX-License-Identifier: Apache-2.0
+// Reconstruct a real-pair transform from arbitrary tightly packed X spectra.
+#define PACK_Y 1
+#if (1ULL * FFT_NZ * FFT_NY * FFT_N) <= 0x100000000ULL
+using index_t = unsigned int;
+#else
+using index_t = unsigned long long;
+#endif
+#if FFT_DOUBLE
+using scalar_t = double;
+#else
+using scalar_t = float;
+#endif
+static __device__ __forceinline__ index_t negative(index_t k, index_t n) {
+    return k == 0 ? 0 : n - k;
+}
+static __device__ __forceinline__ double2 pair_load(const scalar_t *input, index_t k) {
+#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0
+    return reinterpret_cast<const double2 *>(input)[k];
+#elif !FFT_DOUBLE && FFT_INPUT_SHIFT % 2 == 0
+    const float2 v = reinterpret_cast<const float2 *>(input)[k];
+    return make_double2(static_cast<double>(v.x), static_cast<double>(v.y));
+#else
+    return make_double2(static_cast<double>(input[2ULL*k]),
+                        static_cast<double>(input[2ULL*k+1]));
+#endif
+}
+static __device__ __forceinline__ double2 full_spectrum(
+        const scalar_t *input, index_t z, index_t y, index_t x) {
+    constexpr index_t nf = FFT_N/2 + 1;
+    const bool conjugate = x > FFT_N/2;
+    if (conjugate) {
+        z = negative(z, FFT_NZ);
+        y = negative(y, FFT_NY);
+        x = FFT_N-x;
+    }
+    double2 a = pair_load(input, (z*FFT_NY+y)*nf+x);
+    if (x == 0 || x == FFT_N/2) {
+        const double2 mirror = pair_load(input,
+            (negative(z,FFT_NZ)*FFT_NY+negative(y,FFT_NY))*nf+x);
+        a.x = 0.5*a.x+0.5*mirror.x;
+        a.y = 0.5*a.y-0.5*mirror.y;
+    }
+    if (conjugate) a.y = -a.y;
+    return a;
+}
+#if FFT_LOAD
+__device__ double2 fft_wide_load(void *data, unsigned long long offset,
+                                void *caller_info, void *) {
+    constexpr index_t hy = PACK_Y ? FFT_NY/2 : FFT_NY;
+    const index_t index = static_cast<index_t>(offset);
+    const index_t x = index % FFT_N;
+    const index_t y = (index / FFT_N) % hy;
+    const index_t z = index / (hy*FFT_N);
+    const scalar_t *input = static_cast<const scalar_t *>(data)+FFT_INPUT_SHIFT;
+    // Both packed frequencies use one common X reflection. Endpoint
+    // projection is needed only for the two original X boundary planes.
+    const bool conjugate = x > FFT_N/2;
+    const index_t px = conjugate ? FFT_N-x : x;
+    const index_t az = conjugate ? negative(z,FFT_NZ) : z;
+    const index_t ay = conjugate ? negative(y,FFT_NY) : y;
+#if PACK_Y
+    const index_t bz = az;
+    const index_t by = conjugate ? negative(y+FFT_NY/2,FFT_NY) : y+FFT_NY/2;
+    const index_t k = y;
+#else
+    const index_t bz = conjugate ? negative(z+FFT_NZ/2,FFT_NZ) : z+FFT_NZ/2;
+    const index_t by = ay;
+    const index_t k = z;
+#endif
+    constexpr index_t nf = FFT_N/2+1;
+    double2 a = pair_load(input,(az*FFT_NY+ay)*nf+px);
+    double2 b = pair_load(input,(bz*FFT_NY+by)*nf+px);
+    if (px == 0 || px == FFT_N/2) {
+        const double2 am = pair_load(input,
+            (negative(az,FFT_NZ)*FFT_NY+negative(ay,FFT_NY))*nf+px);
+        const double2 bm = pair_load(input,
+            (negative(bz,FFT_NZ)*FFT_NY+negative(by,FFT_NY))*nf+px);
+        a.x = 0.5*a.x+0.5*am.x;
+        a.y = 0.5*a.y-0.5*am.y;
+        b.x = 0.5*b.x+0.5*bm.x;
+        b.y = 0.5*b.y-0.5*bm.y;
+    }
+    if (conjugate) a.y = -a.y;
+    else b.y = -b.y;
+    const double2 w = static_cast<const double2 *>(caller_info)[k];
+    return make_double2(a.x+b.x-w.x*(a.y+b.y)+w.y*(a.x-b.x),
+                        a.y-b.y+w.x*(a.x-b.x)+w.y*(a.y+b.y));
+}
+#else
+__device__ void fft_wide_store(void *data, unsigned long long offset,
+                               double2 value, void *, void *) {
+    constexpr index_t hy = PACK_Y ? FFT_NY/2 : FFT_NY;
+    const index_t index = static_cast<index_t>(offset);
+    const index_t x = index % FFT_N;
+    const index_t y = (index / FFT_N) % hy;
+    const index_t z = index / (hy*FFT_N);
+    scalar_t *output = static_cast<scalar_t *>(data)+FFT_INPUT_SHIFT;
+#if PACK_Y
+    const index_t even = (z*FFT_NY+2*y)*FFT_N+x;
+    constexpr index_t distance = FFT_N;
+#else
+    const index_t even = (2*z*FFT_NY+y)*FFT_N+x;
+    constexpr index_t distance = FFT_NY*FFT_N;
+#endif
+    output[even] = static_cast<scalar_t>(value.x);
+    output[even+distance] = static_cast<scalar_t>(value.y);
+}
+#endif

build/torch-cuda/_b200/c2r_paired_y_inverse_wide_redirect.cu ADDED Viewed

	@@ -0,0 +1,87 @@

+// SPDX-License-Identifier: Apache-2.0
+// Reconstruct a real-pair transform from arbitrary tightly packed X spectra.
+#define PACK_Y 1
+#if (1ULL * FFT_NZ * FFT_NY * FFT_N) <= 0x100000000ULL
+using index_t = unsigned int;
+#else
+using index_t = unsigned long long;
+#endif
+#if FFT_DOUBLE
+using scalar_t = double;
+#else
+using scalar_t = float;
+#endif
+static __device__ __forceinline__ index_t negative(index_t k, index_t n) {
+    return k == 0 ? 0 : n - k;
+}
+static __device__ __forceinline__ double2 pair_load(const scalar_t *input, index_t k) {
+#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0
+    return reinterpret_cast<const double2 *>(input)[k];
+#elif !FFT_DOUBLE && FFT_INPUT_SHIFT % 2 == 0
+    const float2 v = reinterpret_cast<const float2 *>(input)[k];
+    return make_double2(static_cast<double>(v.x), static_cast<double>(v.y));
+#else
+    return make_double2(static_cast<double>(input[2ULL*k]),
+                        static_cast<double>(input[2ULL*k+1]));
+#endif
+}
+static __device__ __forceinline__ double2 full_spectrum(
+        const scalar_t *input, index_t z, index_t y, index_t x) {
+    constexpr index_t nf = FFT_N/2 + 1;
+    const bool conjugate = x > FFT_N/2;
+    if (conjugate) {
+        z = negative(z, FFT_NZ);
+        y = negative(y, FFT_NY);
+        x = FFT_N-x;
+    }
+    double2 a = pair_load(input, (z*FFT_NY+y)*nf+x);
+    if (x == 0 || x == FFT_N/2) {
+        const double2 mirror = pair_load(input,
+            (negative(z,FFT_NZ)*FFT_NY+negative(y,FFT_NY))*nf+x);
+        a.x = 0.5*a.x+0.5*mirror.x;
+        a.y = 0.5*a.y-0.5*mirror.y;
+    }
+    if (conjugate) a.y = -a.y;
+    return a;
+}
+#if FFT_LOAD
+__device__ double2 fft_wide_load(void *data, unsigned long long offset,
+                                void *caller_info, void *) {
+    constexpr index_t hy = PACK_Y ? FFT_NY/2 : FFT_NY;
+    const index_t index = static_cast<index_t>(offset);
+    const index_t x = index % FFT_N;
+    const index_t y = (index / FFT_N) % hy;
+    const index_t z = index / (hy*FFT_N);
+    const scalar_t *input = static_cast<const scalar_t *>(data)+FFT_INPUT_SHIFT;
+    const double2 a = full_spectrum(input,z,y,x);
+#if PACK_Y
+    const double2 b = full_spectrum(input,negative(z,FFT_NZ),FFT_NY/2-y,negative(x,FFT_N));
+    const index_t k = y;
+#else
+    const double2 b = full_spectrum(input,FFT_NZ/2-z,negative(y,FFT_NY),negative(x,FFT_N));
+    const index_t k = z;
+#endif
+    const double2 w = static_cast<const double2 *>(caller_info)[k];
+    return make_double2(a.x+b.x-w.x*(a.y+b.y)+w.y*(a.x-b.x),
+                        a.y-b.y+w.x*(a.x-b.x)+w.y*(a.y+b.y));
+}
+#else
+__device__ void fft_wide_store(void *data, unsigned long long offset,
+                               double2 value, void *, void *) {
+    constexpr index_t hy = PACK_Y ? FFT_NY/2 : FFT_NY;
+    const index_t index = static_cast<index_t>(offset);
+    const index_t x = index % FFT_N;
+    const index_t y = (index / FFT_N) % hy;
+    const index_t z = index / (hy*FFT_N);
+    scalar_t *output = static_cast<scalar_t *>(data)+FFT_INPUT_SHIFT;
+#if PACK_Y
+    const index_t even = (z*FFT_NY+2*y)*FFT_N+x;
+    constexpr index_t distance = FFT_N;
+#else
+    const index_t even = (2*z*FFT_NY+y)*FFT_N+x;
+    constexpr index_t distance = FFT_NY*FFT_N;
+#endif
+    output[even] = static_cast<scalar_t>(value.x);
+    output[even+distance] = static_cast<scalar_t>(value.y);
+}
+#endif

build/torch-cuda/_b200/c2r_pitched_half_inverse_wide_redirect.cu ADDED Viewed

	@@ -0,0 +1,86 @@

+// SPDX-License-Identifier: Apache-2.0
+// Select index width from the complete physical packed extent at compilation.
+// The callback ABI remains 64-bit. Every valid caller index fits the selected
+// unsigned type; pointer arithmetic still computes full-width device addresses.
+#if (1ULL * FFT_NZ * FFT_NY * (FFT_N / 2 + 128)) <= 0x100000000ULL
+using fft_index_t = unsigned int;
+static_assert(sizeof(fft_index_t) == 4, "A 32-bit unsigned integer is required");
+#else
+using fft_index_t = unsigned long long;
+#endif
+#if FFT_DOUBLE
+using caller_scalar_t = double;
+#else
+using caller_scalar_t = float;
+#endif
+__device__ double2 fft_scalar_pair(const caller_scalar_t *input, fft_index_t index) {
+    double2 value;
+#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0
+    value = reinterpret_cast<const double2 *>(input)[index];
+#elif !FFT_DOUBLE && (FFT_INPUT_SHIFT % 2 == 0)
+    // The pair address is aligned for scalar residues zero and two.
+    const float2 packed = reinterpret_cast<const float2 *>(input)[index];
+    value.x = static_cast<double>(packed.x);
+    value.y = static_cast<double>(packed.y);
+#else
+    value.x = static_cast<double>(input[2ULL * index]);
+    value.y = static_cast<double>(input[2ULL * index + 1]);
+#endif
+    return value;
+}
+// The cuFFT logical input uses a planned virtual row stride. Callback arithmetic reads exact nf-stride
+// caller storage directly; no input repacking is performed.
+// Reconstruct twice the half-length complex spectrum from immutable packed
+// caller input. Endpoint projection preserves irfftn for arbitrary spectra.
+#if FFT_LOAD
+__device__ double2 fft_wide_load(void *data, unsigned long long offset,
+                                    void *caller_info, void *shared_memory) {
+    constexpr fft_index_t m = FFT_N / 2;
+    constexpr fft_index_t nf = m + 1;
+    constexpr fft_index_t pitch = m + 128;
+    const fft_index_t index = static_cast<fft_index_t>(offset);
+    const fft_index_t k = index % pitch;
+    const fft_index_t row = index / pitch;
+    const fft_index_t y = row % FFT_NY;
+    const fft_index_t z = row / FFT_NY;
+    const fft_index_t mirror_y = y == 0 ? 0 : FFT_NY - y;
+    const fft_index_t mirror_z = z == 0 ? 0 : FFT_NZ - z;
+    const fft_index_t mirror_row = mirror_z * FFT_NY + mirror_y;
+    const caller_scalar_t *input = static_cast<const caller_scalar_t *>(data) + FFT_INPUT_SHIFT;
+    double2 a = fft_scalar_pair(input, row * nf + k);
+    double2 b = fft_scalar_pair(input, mirror_row * nf + m - k);
+    if (k == 0) {
+        // Project each boundary plane against its Z/Y conjugate partner.
+        // b represents the conjugate of the projected Nyquist value.
+        const double2 a_mirror = fft_scalar_pair(input, mirror_row * nf);
+        const double2 b_mirror = fft_scalar_pair(input, row * nf + m);
+        a.x = 0.5 * (a.x + a_mirror.x);
+        a.y = 0.5 * (a.y - a_mirror.y);
+        b.x = 0.5 * (b.x + b_mirror.x);
+        b.y = 0.5 * (b.y - b_mirror.y);
+    }
+    const double2 w = static_cast<const double2 *>(caller_info)[k];
+    double2 result;
+    result.x = a.x + b.x - w.x * (a.y + b.y) + w.y * (a.x - b.x);
+    result.y = a.y - b.y + w.x * (a.x - b.x) + w.y * (a.y + b.y);
+    return result;
+}
+#else
+__device__ void fft_wide_store(void *physical_output, unsigned long long offset,
+                              double2 value, void *caller_info, void *shared_memory) {
+    caller_scalar_t *output = static_cast<caller_scalar_t *>(physical_output) + FFT_INPUT_SHIFT;
+#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0
+    // A nominal cuFFT output address is complex aligned.
+    reinterpret_cast<double2 *>(output)[offset] = value;
+#elif !FFT_DOUBLE && (FFT_INPUT_SHIFT % 2 == 0)
+    // Both components fit one aligned 64-bit store at these scalar residues.
+    float2 packed;
+    packed.x = static_cast<float>(value.x);
+    packed.y = static_cast<float>(value.y);
+    reinterpret_cast<float2 *>(output)[offset] = packed;
+#else
+    output[2 * offset] = static_cast<caller_scalar_t>(value.x);
+    output[2 * offset + 1] = static_cast<caller_scalar_t>(value.y);
+#endif
+}
+#endif

build/torch-cuda/_b200/cuda_kernel.py ADDED Viewed

	@@ -0,0 +1,64 @@

+"""Compile and load fixed CUDA FFT kernels only during plan preparation."""
+import ctypes as ct
+import functools
+import hashlib
+from pathlib import Path
+from types import SimpleNamespace
+from .cuda_twiddle_forward_lto import library, version, _check
+from .cuda_twiddle_forward_native_dispatch import load_driver
+@functools.lru_cache(maxsize=128)
+def compile_kernel(shape, columns, warps, arch, device_ordinal, double=False):
+    # The ordinal keeps loaded CUDA modules within their owning device context.
+    nz, ny, nx = shape
+    base, odd = 1, nz
+    while odd % 2 == 0:
+        odd //= 2
+        base *= 2
+    if odd not in (1, 3, 5, 7) or base < warps*32//columns:
+        raise ValueError('Unsupported radix decomposition')
+    source = Path(__file__).with_name('fft_axis.cu').read_bytes()
+    options = ['--std=c++17', f'--gpu-architecture=sm_{arch}',
+               f'-DLENGTH={nz}', f'-DINNER={ny*(nx//2+1)}', f'-DBASE={base}',
+               f'-DLOG_BASE={base.bit_length()-1}', f'-DODD={odd}',
+               f'-DCOLUMNS={columns}', f'-DWARPS={warps}', f'-DNX={nx}',
+               f'-DNY={ny}', f'-DFFT_DOUBLE={int(double)}']
+    lib = library()
+    lib.nvrtcGetCUBINSize.argtypes = [ct.c_void_p, ct.POINTER(ct.c_size_t)]
+    lib.nvrtcGetCUBIN.argtypes = [ct.c_void_p, ct.c_void_p]
+    program = ct.c_void_p()
+    _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'fft_axis.cu', 0, None, None))
+    try:
+        opts = (ct.c_char_p * len(options))(*(s.encode() for s in options))
+        code = lib.nvrtcCompileProgram(program, len(options), opts)
+        size = ct.c_size_t()
+        _check(lib.nvrtcGetProgramLogSize(program, ct.byref(size)))
+        log = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetProgramLog(program, log))
+        if code:
+            raise RuntimeError(log.value.decode())
+        _check(lib.nvrtcGetCUBINSize(program, ct.byref(size)))
+        blob = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetCUBIN(program, blob))
+    finally:
+        _check(lib.nvrtcDestroyProgram(ct.byref(program)))
+    driver = load_driver()
+    driver.cuModuleLoadData.argtypes = [ct.POINTER(ct.c_void_p), ct.c_void_p]
+    driver.cuModuleGetFunction.argtypes = [ct.POINTER(ct.c_void_p), ct.c_void_p, ct.c_char_p]
+    driver.cuFuncGetAttribute.argtypes = [ct.POINTER(ct.c_int), ct.c_int, ct.c_void_p]
+    module, function = ct.c_void_p(), ct.c_void_p()
+    _check(driver.cuModuleLoadData(ct.byref(module), blob))
+    _check(driver.cuModuleGetFunction(ct.byref(function), module, b'fft_axis'))
+    registers, local = ct.c_int(), ct.c_int()
+    _check(driver.cuFuncGetAttribute(ct.byref(registers), 4, function))
+    _check(driver.cuFuncGetAttribute(ct.byref(local), 3, function))
+    # local is CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, not an instruction count.
+    metadata = dict(options=options, nvrtc=version(), local_memory_bytes=local.value, source_sha256=hashlib.sha256(source).hexdigest(),
+                    cubin_sha256=hashlib.sha256(blob.raw).hexdigest(), compiler_log=log.value.decode())
+    return SimpleNamespace(module=module, function=function.value, name='fft_axis',
+                           hash=metadata['cubin_sha256'], threads=32*warps, shared=0,
+                           n_regs=registers.value, n_spills=local.value,
+                           grid=((ny*(nx//2+1)+columns-1)//columns, 1, 1),
+                           compiler_metadata=metadata)

build/torch-cuda/_b200/cuda_twiddle_forward_lto.py ADDED Viewed

	@@ -0,0 +1,113 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Compile cuFFT LTO callbacks with the installed NVRTC, before execution.
+This module can compile offline without initializing CUDA or loading code on
+a GPU. Compiled blobs are cached by source, public shape, precision, NVRTC
+version and target. No input tensor values or pointers are cached.
+"""
+from __future__ import annotations
+import ctypes as ct
+import ctypes.util
+import functools
+import hashlib
+from pathlib import Path
+import threading
+_LOCK = threading.Lock()
+@functools.lru_cache(maxsize=1)
+def library():
+    candidates = [ctypes.util.find_library('nvrtc')]
+    for root in (Path('/usr/local/cuda/lib64'),
+                 Path('/usr/local/cuda/targets/x86_64-linux/lib')):
+        candidates.extend(str(p) for p in sorted(root.glob('libnvrtc.so*')))
+    candidates.extend(('libnvrtc.so.13', 'libnvrtc.so.12', 'libnvrtc.so'))
+    lib = None
+    for name in candidates:
+        if name:
+            try:
+                lib = ct.CDLL(name)
+                break
+            except OSError:
+                pass
+    if lib is None:
+        raise RuntimeError('The installed NVRTC library is unavailable')
+    signatures = {
+        'nvrtcVersion': [ct.POINTER(ct.c_int), ct.POINTER(ct.c_int)],
+        'nvrtcCreateProgram': [ct.POINTER(ct.c_void_p), ct.c_char_p, ct.c_char_p,
+                               ct.c_int, ct.POINTER(ct.c_char_p), ct.POINTER(ct.c_char_p)],
+        'nvrtcCompileProgram': [ct.c_void_p, ct.c_int, ct.POINTER(ct.c_char_p)],
+        'nvrtcGetProgramLogSize': [ct.c_void_p, ct.POINTER(ct.c_size_t)],
+        'nvrtcGetProgramLog': [ct.c_void_p, ct.c_void_p],
+        'nvrtcGetLTOIRSize': [ct.c_void_p, ct.POINTER(ct.c_size_t)],
+        'nvrtcGetLTOIR': [ct.c_void_p, ct.c_void_p],
+        'nvrtcDestroyProgram': [ct.POINTER(ct.c_void_p)],
+    }
+    for name, signature in signatures.items():
+        function = getattr(lib, name)
+        function.argtypes = signature
+        function.restype = ct.c_int
+    return lib
+def _check(result):
+    if result:
+        raise RuntimeError(f'NVRTC returned error {result}')
+def version():
+    major, minor = ct.c_int(), ct.c_int()
+    _check(library().nvrtcVersion(ct.byref(major), ct.byref(minor)))
+    return major.value, minor.value
+@functools.lru_cache(maxsize=128)
+def _compile(source, double, arch, nvrtc_version):
+    lib = library()
+    options = ('--std=c++17', '--relocatable-device-code=true', '-dlto',
+               f'--gpu-architecture=compute_{arch}', f'-DFFT_DOUBLE={int(double)}')
+    program = ct.c_void_p()
+    _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'scalar_store.cu', 0, None, None))
+    try:
+        option_array = (ct.c_char_p * len(options))(*(v.encode() for v in options))
+        result = lib.nvrtcCompileProgram(program, len(options), option_array)
+        log_size = ct.c_size_t()
+        _check(lib.nvrtcGetProgramLogSize(program, ct.byref(log_size)))
+        log = ct.create_string_buffer(log_size.value)
+        _check(lib.nvrtcGetProgramLog(program, log))
+        if result:
+            raise RuntimeError(f'NVRTC returned error {result}: {log.value.decode()}')
+        size = ct.c_size_t()
+        _check(lib.nvrtcGetLTOIRSize(program, ct.byref(size)))
+        data = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetLTOIR(program, data))
+        blob = data.raw
+        metadata = dict(double=bool(double), target=f'compute_{arch}',
+                        nvrtc=list(nvrtc_version), options=list(options),
+                        source_sha256=hashlib.sha256(source).hexdigest(),
+                        lto_ir_sha256=hashlib.sha256(blob).hexdigest(),
+                        lto_ir_bytes=len(blob), compiler_log=log.value.decode())
+        return blob, metadata
+    finally:
+        _check(lib.nvrtcDestroyProgram(ct.byref(program)))
+def compile_callback(double, role, shift, arch=100, shape=(256,256,256), axis='z'):
+    """Compile every scalar residue at preparation, never during execution."""
+    if role not in ('pair', 'reconstruct', 'store') or not 0 <= shift < (2 if double else 4):
+        raise ValueError('Invalid callback role or scalar residue')
+    prefix = f'#define SCALAR_SHIFT {shift}\n#define IS_LOAD {int(role == "pair")}\n'
+    prefix += f'#define RECONSTRUCT {int(role == "reconstruct")}\n'
+    nz, ny, nx = shape
+    prefix += f'#define NZ {nz}\n#define NY {ny}\n#define NX {nx}\n'
+    prefix += f'#define NF {nx//2+1}\n#define FIRST_Y {int(axis == "y")}\n'
+    source = prefix.encode() + Path(__file__).with_name('scalar_io.cu').read_bytes()
+    with _LOCK:
+        blob, metadata = _compile(source, bool(double), int(arch), version())
+    metadata = dict(metadata, role=role, scalar_shift=shift,
+                    arithmetic='float64', caller_dtype='float64' if double else 'float32')
+    return [(b'fft_pair_load' if role == 'pair' else (b'fft_reconstruct_load' if role == 'reconstruct' else b'fft_scalar_store'),
+             5 if role == 'store' else 1, ct.create_string_buffer(blob, len(blob)), metadata)]

build/torch-cuda/_b200/cuda_twiddle_forward_native_dispatch.c ADDED Viewed

	@@ -0,0 +1,42 @@

+/* SPDX-License-Identifier: Apache-2.0 */
+/* Enqueue the forward Y/X cuFFT and Z-axis CUDA FFT in one host call.
+ * All handles, functions, grids, workspace and twiddles are prepared first.
+ * No data is copied, padded or reordered by this dispatcher.
+ */
+#include <stddef.h>
+#include <stdint.h>
+typedef int (*complex_fn)(int, void *, void *, int);
+typedef int (*launch_fn)(void *, unsigned int, unsigned int, unsigned int,
+                         unsigned int, unsigned int, unsigned int,
+                         unsigned int, void *, void **, void **);
+struct fft_dispatch {
+    complex_fn complex_exec;
+    launch_fn launch;
+    int handle;
+    uint64_t intermediate;
+    uint64_t twiddle;
+    void *stream;
+    void *function;
+    unsigned int grid[3];
+    unsigned int threads;
+    unsigned int shared;
+};
+int fft_execute(const struct fft_dispatch *plan, uint64_t input, uint64_t output) {
+    int result = plan->complex_exec(plan->handle, (void *)(uintptr_t)input,
+                                    (void *)(uintptr_t)plan->intermediate, -1);
+    if (result != 0) {
+        return result;
+    }
+    uint64_t values[5] = {plan->intermediate, output, plan->twiddle, 0, 0};
+    void *parameters[5] = {&values[0], &values[1], &values[2], &values[3], &values[4]};
+    result = plan->launch(plan->function, plan->grid[0], plan->grid[1], plan->grid[2],
+                           plan->threads, 1, 1, plan->shared, plan->stream,
+                           parameters, NULL);
+    /* Positive codes originate in cuFFT; negative codes in the CUDA driver. */
+    return -result;
+}
+size_t fft_dispatch_size(void) { return sizeof(struct fft_dispatch); }

build/torch-cuda/_b200/cuda_twiddle_forward_native_dispatch.py ADDED Viewed

	@@ -0,0 +1,101 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Prepare a combined cuFFT and CUDA dispatcher before execution."""
+from __future__ import annotations
+import ctypes as ct
+import fcntl
+import hashlib
+import os
+from pathlib import Path
+import shlex
+import subprocess
+import tempfile
+import threading
+class Descriptor(ct.Structure):
+    _fields_ = [('complex_exec', ct.c_void_p), ('launch', ct.c_void_p),
+                ('handle', ct.c_int), ('intermediate', ct.c_uint64),
+                ('twiddle', ct.c_uint64), ('stream', ct.c_void_p),
+                ('function', ct.c_void_p), ('grid', ct.c_uint * 3),
+                ('threads', ct.c_uint), ('shared', ct.c_uint)]
+_LIBRARY = None
+_DRIVER = None
+_LOCK = threading.Lock()
+def load_library():
+    """Build with the installed host compiler and validate structure size."""
+    global _LIBRARY
+    with _LOCK:
+        if _LIBRARY is not None:
+            return _LIBRARY
+        source = Path(__file__).with_name('cuda_twiddle_forward_native_dispatch.c')
+        compiler = shlex.split(os.environ.get('CC', 'cc'))
+        flags = ['-std=c11', '-O3', '-fPIC', '-shared', '-Wall', '-Wextra', '-Werror']
+        key = hashlib.sha256(source.read_bytes() + repr((compiler, flags)).encode()).hexdigest()
+        cache = Path(tempfile.gettempdir()) / ('fft-half-triton-z-native-' + str(os.getuid())) / key
+        cache.mkdir(parents=True, exist_ok=True)
+        target = cache / 'dispatch.so'
+        with (cache / 'build.lock').open('w') as lock:
+            fcntl.flock(lock, fcntl.LOCK_EX)
+            if not target.is_file():
+                pending = cache / ('dispatch.' + str(os.getpid()) + '.so')
+                try:
+                    subprocess.run(compiler + flags + [str(source), '-o', str(pending)],
+                                   check=True, capture_output=True, text=True)
+                    pending.replace(target)
+                finally:
+                    pending.unlink(missing_ok=True)
+        lib = ct.CDLL(str(target))
+        lib.fft_execute.argtypes = [ct.POINTER(Descriptor), ct.c_uint64, ct.c_uint64]
+        lib.fft_execute.restype = ct.c_int
+        lib.fft_dispatch_size.argtypes = []
+        lib.fft_dispatch_size.restype = ct.c_size_t
+        if lib.fft_dispatch_size() != ct.sizeof(Descriptor):
+            raise RuntimeError('Native dispatcher structure size mismatch')
+        _LIBRARY = lib
+        return lib
+def load_driver():
+    """Resolve the driver function during preparation only."""
+    global _DRIVER
+    with _LOCK:
+        if _DRIVER is None:
+            _DRIVER = ct.CDLL('libcuda.so.1')
+            _DRIVER.cuLaunchKernel.restype = ct.c_int
+        return _DRIVER
+class PreparedDispatch:
+    """Immutable setup descriptor, with fresh caller pointers at each call.
+    The Plan retains the cuFFT handles and compiled CUDA kernel. Concurrent
+    execution requires independent plans because cuFFT workspace is reused.
+    """
+    def __init__(self, complex_exec, handle, compiled, grid, intermediate, twiddle, stream):
+        abi = dict(threads=compiled.threads, shared_bytes=compiled.shared)
+        if not compiled.function or compiled.module is None:
+            raise ValueError('The CUDA kernel must be loaded during setup')
+        if not 0 < handle < 2**31:
+            raise ValueError('A prepared cuFFT handle is required')
+        if len(grid) != 3 or any(not 0 < n < 2**32 for n in grid):
+            raise ValueError('Invalid launch grid')
+        if not intermediate or not twiddle:
+            raise ValueError('Prepared intermediate and twiddle storage is required')
+        self.library = load_library()
+        self.driver = load_driver()
+        self.descriptor = Descriptor(
+            complex_exec=ct.cast(complex_exec, ct.c_void_p).value,
+            launch=ct.cast(self.driver.cuLaunchKernel, ct.c_void_p).value,
+            handle=handle, intermediate=intermediate,
+            twiddle=twiddle, stream=stream, function=compiled.function,
+            grid=(ct.c_uint * 3)(*grid), threads=abi['threads'], shared=abi['shared_bytes'])
+        self.pointer = ct.pointer(self.descriptor)
+        self.function = self.library.fft_execute
+    def execute(self, input_ptr, output_ptr):
+        return self.function(self.pointer, input_ptr, output_ptr)

build/torch-cuda/_b200/cuda_z_inverse_cufft_support.py ADDED Viewed

	@@ -0,0 +1,105 @@

+# SPDX-License-Identifier: Apache-2.0
+"""cuFFT ABI and preallocated handles with optional setup-time LTO callbacks."""
+from __future__ import annotations
+import ctypes as ct
+import ctypes.util
+from pathlib import Path
+import torch
+def _load_cufft():
+    candidates = [ctypes.util.find_library("cufft")]
+    for root in (Path(torch.__file__).resolve().parent.parent / "nvidia/cufft/lib",
+                 Path("/usr/local/cuda/lib64")):
+        candidates.extend(str(p) for p in sorted(root.glob("libcufft.so*")))
+    candidates.extend(("libcufft.so.12", "libcufft.so.11", "libcufft.so"))
+    for name in candidates:
+        if name:
+            try:
+                return ct.CDLL(name)
+            except OSError:
+                pass
+    raise RuntimeError("The installed cuFFT shared library could not be loaded")
+_lib = _load_cufft()
+_int_p = ct.POINTER(ct.c_int)
+_size_p = ct.POINTER(ct.c_size_t)
+_SIGNATURES = {
+    "cufftCreate": [_int_p],
+    "cufftDestroy": [ct.c_int],
+    "cufftSetAutoAllocation": [ct.c_int, ct.c_int],
+    "cufftSetWorkArea": [ct.c_int, ct.c_void_p],
+    "cufftSetStream": [ct.c_int, ct.c_void_p],
+    "cufftMakePlanMany": [ct.c_int, ct.c_int, _int_p, _int_p, ct.c_int,
+                          ct.c_int, _int_p, ct.c_int, ct.c_int, ct.c_int,
+                          ct.c_int, _size_p],
+    "cufftGetVersion": [_int_p],
+    "cufftXtSetJITCallback": [ct.c_int, ct.c_char_p, ct.c_void_p, ct.c_size_t,
+                             ct.c_int, ct.POINTER(ct.c_void_p)],
+    "cufftSetPlanPropertyInt64": [ct.c_int, ct.c_int, ct.c_longlong],
+}
+for _name in ("R2C", "D2Z", "C2R", "Z2D"):
+    _SIGNATURES["cufftExec" + _name] = [ct.c_int, ct.c_void_p, ct.c_void_p]
+for _name in ("C2C", "Z2Z"):
+    _SIGNATURES["cufftExec" + _name] = [ct.c_int, ct.c_void_p, ct.c_void_p, ct.c_int]
+for _name, _signature in _SIGNATURES.items():
+    _fn = getattr(_lib, _name)
+    _fn.argtypes = _signature
+    _fn.restype = ct.c_int
+def _check(result):
+    if result:
+        raise RuntimeError(f"cuFFT returned error {result}")
+def version():
+    value = ct.c_int()
+    _check(_lib.cufftGetVersion(ct.byref(value)))
+    return value.value
+class _Handle:
+    def __init__(self, sizes, kind, inembed, istride, idist, onembed, ostride,
+                 odist, batch, device, stream, callback=None, patient=False,
+                 caller_info=None):
+        handle = ct.c_int()
+        _check(_lib.cufftCreate(ct.byref(handle)))
+        self.value = handle.value
+        self.workspace = None
+        self.callback = callback
+        try:
+            _check(_lib.cufftSetAutoAllocation(self.value, 0))
+            if patient:
+                # NVFFT_PLAN_PROPERTY_INT64_PATIENT_JIT in cufft.h is enum 1.
+                _check(_lib.cufftSetPlanPropertyInt64(self.value, 1, 1))
+            if callback is not None:
+                for symbol, callback_type, blob, _metadata in callback:
+                    info = ct.c_void_p(caller_info) if caller_info is not None else None
+                    _check(_lib.cufftXtSetJITCallback(self.value, symbol, blob, len(blob),
+                                                    callback_type,
+                                                    ct.byref(info) if info is not None else None))
+            size = ct.c_size_t()
+            array = ct.c_int * len(sizes)
+            _check(_lib.cufftMakePlanMany(
+                self.value, len(sizes), array(*sizes), array(*inembed),
+                istride, idist, array(*onembed), ostride, odist, kind,
+                batch, ct.byref(size)))
+            self.workspace_bytes = size.value
+            self.workspace = torch.empty(size.value, dtype=torch.uint8, device=device)
+            _check(_lib.cufftSetWorkArea(self.value, self.workspace.data_ptr()))
+            _check(_lib.cufftSetStream(self.value, stream.cuda_stream))
+        except BaseException:
+            _lib.cufftDestroy(self.value)
+            self.value = None
+            raise
+    def close(self):
+        if self.value is not None:
+            _check(_lib.cufftDestroy(self.value))
+            self.value = None
+            self.workspace = None
+            self.callback = None

build/torch-cuda/_b200/cuda_z_inverse_lto.py ADDED Viewed

	@@ -0,0 +1,110 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Compile cuFFT LTO callbacks with the installed NVRTC, before execution.
+This module can compile offline without initializing CUDA or loading code on
+a GPU. Compiled blobs are cached by source, public shape, precision, NVRTC
+version and target. No input tensor values or pointers are cached.
+"""
+from __future__ import annotations
+import ctypes as ct
+import ctypes.util
+import functools
+import hashlib
+from pathlib import Path
+import threading
+_LOCK = threading.Lock()
+@functools.lru_cache(maxsize=1)
+def library():
+    candidates = [ctypes.util.find_library('nvrtc')]
+    for root in (Path('/usr/local/cuda/lib64'),
+                 Path('/usr/local/cuda/targets/x86_64-linux/lib')):
+        candidates.extend(str(p) for p in sorted(root.glob('libnvrtc.so*')))
+    candidates.extend(('libnvrtc.so.13', 'libnvrtc.so.12', 'libnvrtc.so'))
+    lib = None
+    for name in candidates:
+        if name:
+            try:
+                lib = ct.CDLL(name)
+                break
+            except OSError:
+                pass
+    if lib is None:
+        raise RuntimeError('The installed NVRTC library is unavailable')
+    signatures = {
+        'nvrtcVersion': [ct.POINTER(ct.c_int), ct.POINTER(ct.c_int)],
+        'nvrtcCreateProgram': [ct.POINTER(ct.c_void_p), ct.c_char_p, ct.c_char_p,
+                               ct.c_int, ct.POINTER(ct.c_char_p), ct.POINTER(ct.c_char_p)],
+        'nvrtcCompileProgram': [ct.c_void_p, ct.c_int, ct.POINTER(ct.c_char_p)],
+        'nvrtcGetProgramLogSize': [ct.c_void_p, ct.POINTER(ct.c_size_t)],
+        'nvrtcGetProgramLog': [ct.c_void_p, ct.c_void_p],
+        'nvrtcGetLTOIRSize': [ct.c_void_p, ct.POINTER(ct.c_size_t)],
+        'nvrtcGetLTOIR': [ct.c_void_p, ct.c_void_p],
+        'nvrtcDestroyProgram': [ct.POINTER(ct.c_void_p)],
+    }
+    for name, signature in signatures.items():
+        function = getattr(lib, name)
+        function.argtypes = signature
+        function.restype = ct.c_int
+    return lib
+def _check(result):
+    if result:
+        raise RuntimeError(f'NVRTC returned error {result}')
+def version():
+    major, minor = ct.c_int(), ct.c_int()
+    _check(library().nvrtcVersion(ct.byref(major), ct.byref(minor)))
+    return major.value, minor.value
+@functools.lru_cache(maxsize=128)
+def _compile(source, shape, double, arch, nvrtc_version, input_shift, kind):
+    nz, ny, nx = shape
+    lib = library()
+    options = ('--std=c++17', '--relocatable-device-code=true', '-dlto',
+               f'--gpu-architecture=compute_{arch}', f'-DFFT_N={nx}',
+               f'-DFFT_NZ={nz}', f'-DFFT_NY={ny}',
+               f'-DFFT_DOUBLE={int(double)}', f'-DFFT_INPUT_SHIFT={input_shift}',
+               f'-DFFT_LOAD={int(kind == "load")}')
+    program = ct.c_void_p()
+    _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'endpoint.cu', 0, None, None))
+    try:
+        option_array = (ct.c_char_p * len(options))(*(v.encode() for v in options))
+        result = lib.nvrtcCompileProgram(program, len(options), option_array)
+        log_size = ct.c_size_t()
+        _check(lib.nvrtcGetProgramLogSize(program, ct.byref(log_size)))
+        log = ct.create_string_buffer(log_size.value)
+        _check(lib.nvrtcGetProgramLog(program, log))
+        if result:
+            raise RuntimeError(f'NVRTC returned error {result}: {log.value.decode()}')
+        size = ct.c_size_t()
+        _check(lib.nvrtcGetLTOIRSize(program, ct.byref(size)))
+        data = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetLTOIR(program, data))
+        blob = data.raw
+        metadata = dict(shape=shape, double=bool(double), target=f'compute_{arch}',
+                        nvrtc=list(nvrtc_version), options=list(options),
+                        source_sha256=hashlib.sha256(source).hexdigest(),
+                        lto_ir_sha256=hashlib.sha256(blob).hexdigest(),
+                        lto_ir_bytes=len(blob), compiler_log=log.value.decode())
+        return blob, metadata
+    finally:
+        _check(lib.nvrtcDestroyProgram(ct.byref(program)))
+def compile_callback(shape, input_shift, kind, arch=100):
+    """Reconstruct a half-length spectrum in FP64 or store its inverse as FP32."""
+    if kind not in ('load', 'store'):
+        raise ValueError('Expected a load or store callback')
+    source = Path(__file__).with_name('wide_redirect.cu').read_bytes()
+    with _LOCK:
+        blob, metadata = _compile(source, tuple(shape), True, int(arch), version(), input_shift, kind)
+    return (('fft_wide_' + kind).encode(), 1 if kind == 'load' else 5,
+            ct.create_string_buffer(blob, len(blob)), dict(metadata, kind=kind))

build/torch-cuda/_b200/cuda_z_inverse_native_dispatch.c ADDED Viewed

	@@ -0,0 +1,42 @@

+/* SPDX-License-Identifier: Apache-2.0 */
+/* Enqueue the inverse Y/X cuFFT and Z-axis CUDA FFT in one host call.
+ * All handles, functions, grids, workspace and twiddles are prepared first.
+ * No data is copied, padded or reordered by this dispatcher.
+ */
+#include <stddef.h>
+#include <stdint.h>
+typedef int (*complex_fn)(int, void *, void *, int);
+typedef int (*launch_fn)(void *, unsigned int, unsigned int, unsigned int,
+                         unsigned int, unsigned int, unsigned int,
+                         unsigned int, void *, void **, void **);
+struct fft_dispatch {
+    complex_fn complex_exec;
+    launch_fn launch;
+    int handle;
+    uint64_t intermediate;
+    uint64_t twiddle;
+    void *stream;
+    void *function;
+    unsigned int grid[3];
+    unsigned int threads;
+    unsigned int shared;
+};
+int fft_execute(const struct fft_dispatch *plan, uint64_t input, uint64_t output) {
+    int result = plan->complex_exec(plan->handle, (void *)(uintptr_t)input,
+                                    (void *)(uintptr_t)plan->intermediate, 1);
+    if (result != 0) {
+        return result;
+    }
+    uint64_t values[5] = {plan->intermediate, output, plan->twiddle, 0, 0};
+    void *parameters[5] = {&values[0], &values[1], &values[2], &values[3], &values[4]};
+    result = plan->launch(plan->function, plan->grid[0], plan->grid[1], plan->grid[2],
+                           plan->threads, 1, 1, plan->shared, plan->stream,
+                           parameters, NULL);
+    /* Positive codes originate in cuFFT; negative codes in the CUDA driver. */
+    return -result;
+}
+size_t fft_dispatch_size(void) { return sizeof(struct fft_dispatch); }

build/torch-cuda/_b200/cuda_z_inverse_native_dispatch.py ADDED Viewed

	@@ -0,0 +1,101 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Prepare a combined cuFFT and CUDA dispatcher before execution."""
+from __future__ import annotations
+import ctypes as ct
+import fcntl
+import hashlib
+import os
+from pathlib import Path
+import shlex
+import subprocess
+import tempfile
+import threading
+class Descriptor(ct.Structure):
+    _fields_ = [('complex_exec', ct.c_void_p), ('launch', ct.c_void_p),
+                ('handle', ct.c_int), ('intermediate', ct.c_uint64),
+                ('twiddle', ct.c_uint64), ('stream', ct.c_void_p),
+                ('function', ct.c_void_p), ('grid', ct.c_uint * 3),
+                ('threads', ct.c_uint), ('shared', ct.c_uint)]
+_LIBRARY = None
+_DRIVER = None
+_LOCK = threading.Lock()
+def load_library():
+    """Build with the installed host compiler and validate structure size."""
+    global _LIBRARY
+    with _LOCK:
+        if _LIBRARY is not None:
+            return _LIBRARY
+        source = Path(__file__).with_name('cuda_z_inverse_native_dispatch.c')
+        compiler = shlex.split(os.environ.get('CC', 'cc'))
+        flags = ['-std=c11', '-O3', '-fPIC', '-shared', '-Wall', '-Wextra', '-Werror']
+        key = hashlib.sha256(source.read_bytes() + repr((compiler, flags)).encode()).hexdigest()
+        cache = Path(tempfile.gettempdir()) / ('fft-half-triton-z-native-' + str(os.getuid())) / key
+        cache.mkdir(parents=True, exist_ok=True)
+        target = cache / 'dispatch.so'
+        with (cache / 'build.lock').open('w') as lock:
+            fcntl.flock(lock, fcntl.LOCK_EX)
+            if not target.is_file():
+                pending = cache / ('dispatch.' + str(os.getpid()) + '.so')
+                try:
+                    subprocess.run(compiler + flags + [str(source), '-o', str(pending)],
+                                   check=True, capture_output=True, text=True)
+                    pending.replace(target)
+                finally:
+                    pending.unlink(missing_ok=True)
+        lib = ct.CDLL(str(target))
+        lib.fft_execute.argtypes = [ct.POINTER(Descriptor), ct.c_uint64, ct.c_uint64]
+        lib.fft_execute.restype = ct.c_int
+        lib.fft_dispatch_size.argtypes = []
+        lib.fft_dispatch_size.restype = ct.c_size_t
+        if lib.fft_dispatch_size() != ct.sizeof(Descriptor):
+            raise RuntimeError('Native dispatcher structure size mismatch')
+        _LIBRARY = lib
+        return lib
+def load_driver():
+    """Resolve the driver function during preparation only."""
+    global _DRIVER
+    with _LOCK:
+        if _DRIVER is None:
+            _DRIVER = ct.CDLL('libcuda.so.1')
+            _DRIVER.cuLaunchKernel.restype = ct.c_int
+        return _DRIVER
+class PreparedDispatch:
+    """Immutable setup descriptor, with fresh caller pointers at each call.
+    The Plan retains the cuFFT handles and compiled CUDA kernel. Concurrent
+    execution requires independent plans because cuFFT workspace is reused.
+    """
+    def __init__(self, complex_exec, handle, compiled, grid, intermediate, twiddle, stream):
+        abi = dict(threads=compiled.threads, shared_bytes=compiled.shared)
+        if not compiled.function or compiled.module is None:
+            raise ValueError('The CUDA kernel must be loaded during setup')
+        if not 0 < handle < 2**31:
+            raise ValueError('A prepared cuFFT handle is required')
+        if len(grid) != 3 or any(not 0 < n < 2**32 for n in grid):
+            raise ValueError('Invalid launch grid')
+        if not intermediate or not twiddle:
+            raise ValueError('Prepared intermediate and twiddle storage is required')
+        self.library = load_library()
+        self.driver = load_driver()
+        self.descriptor = Descriptor(
+            complex_exec=ct.cast(complex_exec, ct.c_void_p).value,
+            launch=ct.cast(self.driver.cuLaunchKernel, ct.c_void_p).value,
+            handle=handle, intermediate=intermediate,
+            twiddle=twiddle, stream=stream, function=compiled.function,
+            grid=(ct.c_uint * 3)(*grid), threads=abi['threads'], shared=abi['shared_bytes'])
+        self.pointer = ct.pointer(self.descriptor)
+        self.function = self.library.fft_execute
+    def execute(self, input_ptr, output_ptr):
+        return self.function(self.pointer, input_ptr, output_ptr)

build/torch-cuda/_b200/cuda_z_inverse_schedule.py ADDED Viewed

	@@ -0,0 +1,71 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Host-only mixed-radix schedules and accurately rounded twiddle tables."""
+from __future__ import annotations
+import math
+def factors(length):
+    """Return a deterministic DIT factorization; no measurements occur here."""
+    remaining = int(length)
+    if remaining < 1:
+        raise ValueError("FFT lengths must be positive")
+    result = []
+    for radix in (4, 2):
+        while remaining % radix == 0:
+            result.append(radix)
+            remaining //= radix
+    # Collapse the final pair of radix-3 stages for lengths 9 * 2**k.
+    # This gives the radix-9 kernel power-of-two branch lengths; other odd
+    # factorizations retain the established paired or generic network.
+    if remaining == 9:
+        result.append(9)
+        remaining = 1
+    for radix in (3, 5, 7):
+        while remaining % radix == 0:
+            result.append(radix)
+            remaining //= radix
+    if remaining != 1:
+        raise ValueError("This prototype supports lengths with factors 2, 3, 5, 7")
+    return tuple(result)
+def stage_sizes(radices):
+    sizes = []
+    length = 1
+    for radix in radices:
+        length *= radix
+        sizes.append(length)
+    return tuple(sizes)
+def permutation(length):
+    """Map DIT input positions to natural input indices, for CPU validation."""
+    radices = factors(length)
+    result = []
+    for index in range(length):
+        divisor, reverse = 1, 0
+        for radix in radices:
+            reverse += ((index // divisor) % radix) * (length // (divisor * radix))
+            divisor *= radix
+        result.append(reverse)
+    return result
+def twiddles(length):
+    """Interleaved exp(-2*pi*i*k/N), computed in host double precision."""
+    result = []
+    for index in range(length):
+        if index == 0:
+            value = (1.0, 0.0)
+        elif 2 * index == length:
+            value = (-1.0, 0.0)
+        elif 4 * index == length:
+            value = (0.0, -1.0)
+        elif 4 * index == 3 * length:
+            value = (0.0, 1.0)
+        else:
+            angle = -math.tau * index / length
+            value = (math.cos(angle), math.sin(angle))
+        result.extend(value)
+    return result

build/torch-cuda/_b200/dense_half_inverse_abi.py ADDED Viewed

	@@ -0,0 +1,30 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Validate the pointer-publication kernel before creating a native launcher."""
+import re
+def validate_kernel_abi(compiled):
+    metadata = compiled.metadata
+    if metadata.num_ctas != 1 or getattr(metadata, 'warp_size', 32) != 32:
+        raise ValueError('Only one-CTA kernels with warp size 32 are supported')
+    for name in ('global_scratch_size', 'profile_scratch_size',
+                 'launch_cooperative_grid', 'launch_pdl', 'tmem_size'):
+        if getattr(metadata, name, 0):
+            raise ValueError('Unsupported kernel requirement: ' + name)
+    if getattr(metadata, 'tensordesc_meta', None):
+        raise ValueError('Tensor descriptors are unsupported')
+    if metadata.num_warps != 1:
+        raise ValueError('Pointer publication uses exactly one warp')
+    if list(compiled.src.signature.values()) != ['*i64', 'u64']:
+        raise ValueError('Expected an int64 slot pointer and a uint64 output address')
+    match = re.search(r'\.entry\s+' + re.escape(metadata.name) + r'\s*\((.*?)\)',
+                      compiled.asm['ptx'], flags=re.DOTALL)
+    if match is None:
+        raise ValueError('Kernel entry signature was not found in PTX')
+    parameters = [part.strip() for part in match.group(1).split(',')]
+    if len(parameters) != 4 or any(not re.fullmatch(
+            r'\.param\s+\.u64(?:\s+\.ptr\s+\.global\s+\.align\s+\d+)?\s+\w+',
+            part) for part in parameters):
+        raise ValueError('Expected two 64-bit arguments and two null scratch pointers')
+    return dict(parameter_count=4, null_scratch_arguments=2,
+                threads=32, shared_bytes=metadata.shared)

build/torch-cuda/_b200/dense_half_inverse_kernels.py ADDED Viewed

	@@ -0,0 +1,10 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Stream-ordered pointer publication, with no tensor-data copy."""
+import triton
+import triton.language as tl
+@triton.jit(do_not_specialize=['output_address'],
+            do_not_specialize_on_alignment=['slot', 'output_address'])
+def publish_output(slot, output_address):
+    tl.store(slot, output_address.to(tl.uint64))

build/torch-cuda/_b200/dense_half_inverse_lto.py ADDED Viewed

	@@ -0,0 +1,110 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Compile cuFFT LTO callbacks with the installed NVRTC, before execution.
+This module can compile offline without initializing CUDA or loading code on
+a GPU. Compiled blobs are cached by source, public shape, precision, NVRTC
+version and target. No input tensor values or pointers are cached.
+"""
+from __future__ import annotations
+import ctypes as ct
+import ctypes.util
+import functools
+import hashlib
+from pathlib import Path
+import threading
+_LOCK = threading.Lock()
+@functools.lru_cache(maxsize=1)
+def library():
+    candidates = [ctypes.util.find_library('nvrtc')]
+    for root in (Path('/usr/local/cuda/lib64'),
+                 Path('/usr/local/cuda/targets/x86_64-linux/lib')):
+        candidates.extend(str(p) for p in sorted(root.glob('libnvrtc.so*')))
+    candidates.extend(('libnvrtc.so.13', 'libnvrtc.so.12', 'libnvrtc.so'))
+    lib = None
+    for name in candidates:
+        if name:
+            try:
+                lib = ct.CDLL(name)
+                break
+            except OSError:
+                pass
+    if lib is None:
+        raise RuntimeError('The installed NVRTC library is unavailable')
+    signatures = {
+        'nvrtcVersion': [ct.POINTER(ct.c_int), ct.POINTER(ct.c_int)],
+        'nvrtcCreateProgram': [ct.POINTER(ct.c_void_p), ct.c_char_p, ct.c_char_p,
+                               ct.c_int, ct.POINTER(ct.c_char_p), ct.POINTER(ct.c_char_p)],
+        'nvrtcCompileProgram': [ct.c_void_p, ct.c_int, ct.POINTER(ct.c_char_p)],
+        'nvrtcGetProgramLogSize': [ct.c_void_p, ct.POINTER(ct.c_size_t)],
+        'nvrtcGetProgramLog': [ct.c_void_p, ct.c_void_p],
+        'nvrtcGetLTOIRSize': [ct.c_void_p, ct.POINTER(ct.c_size_t)],
+        'nvrtcGetLTOIR': [ct.c_void_p, ct.c_void_p],
+        'nvrtcDestroyProgram': [ct.POINTER(ct.c_void_p)],
+    }
+    for name, signature in signatures.items():
+        function = getattr(lib, name)
+        function.argtypes = signature
+        function.restype = ct.c_int
+    return lib
+def _check(result):
+    if result:
+        raise RuntimeError(f'NVRTC returned error {result}')
+def version():
+    major, minor = ct.c_int(), ct.c_int()
+    _check(library().nvrtcVersion(ct.byref(major), ct.byref(minor)))
+    return major.value, minor.value
+@functools.lru_cache(maxsize=128)
+def _compile(source, shape, double, arch, nvrtc_version, input_shift, kind):
+    nz, ny, nx = shape
+    lib = library()
+    options = ('--std=c++17', '--relocatable-device-code=true', '-dlto',
+               f'--gpu-architecture=compute_{arch}', f'-DFFT_N={nx}',
+               f'-DFFT_NZ={nz}', f'-DFFT_NY={ny}',
+               f'-DFFT_DOUBLE={int(double)}', f'-DFFT_INPUT_SHIFT={input_shift}',
+               f'-DFFT_LOAD={int(kind == "load")}')
+    program = ct.c_void_p()
+    _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'endpoint.cu', 0, None, None))
+    try:
+        option_array = (ct.c_char_p * len(options))(*(v.encode() for v in options))
+        result = lib.nvrtcCompileProgram(program, len(options), option_array)
+        log_size = ct.c_size_t()
+        _check(lib.nvrtcGetProgramLogSize(program, ct.byref(log_size)))
+        log = ct.create_string_buffer(log_size.value)
+        _check(lib.nvrtcGetProgramLog(program, log))
+        if result:
+            raise RuntimeError(f'NVRTC returned error {result}: {log.value.decode()}')
+        size = ct.c_size_t()
+        _check(lib.nvrtcGetLTOIRSize(program, ct.byref(size)))
+        data = ct.create_string_buffer(size.value)
+        _check(lib.nvrtcGetLTOIR(program, data))
+        blob = data.raw
+        metadata = dict(shape=shape, double=bool(double), target=f'compute_{arch}',
+                        nvrtc=list(nvrtc_version), options=list(options),
+                        source_sha256=hashlib.sha256(source).hexdigest(),
+                        lto_ir_sha256=hashlib.sha256(blob).hexdigest(),
+                        lto_ir_bytes=len(blob), compiler_log=log.value.decode())
+        return blob, metadata
+    finally:
+        _check(lib.nvrtcDestroyProgram(ct.byref(program)))
+def compile_callback(shape, input_shift, kind, arch=100, double=False):
+    """Reconstruct a half-length spectrum in FP64 or store its inverse as FP32."""
+    if kind not in ('load', 'store'):
+        raise ValueError('Expected a load or store callback')
+    source = Path(__file__).with_name('c2r_dense_half_inverse_wide_redirect.cu').read_bytes()
+    with _LOCK:
+        blob, metadata = _compile(source, tuple(shape), bool(double), int(arch), version(), input_shift, kind)
+    return (('fft_wide_' + kind).encode(), 1 if kind == 'load' else 5,
+            ct.create_string_buffer(blob, len(blob)), dict(metadata, kind=kind))

build/torch-cuda/_b200/dense_half_inverse_native_dispatch.c ADDED Viewed

	@@ -0,0 +1,43 @@

+/* SPDX-License-Identifier: Apache-2.0 */
+/* Publish the caller output address and enqueue the prepared FP64 FFT chain.
+ * Only a scalar pointer is published. Caller tensor data is never repacked.
+ * Each stream owns an independent descriptor, slot, workspace and handles.
+ */
+#include <stddef.h>
+#include <stdint.h>
+typedef int (*complex_fn)(int, void *, void *, int);
+typedef int (*launch_fn)(void *, unsigned int, unsigned int, unsigned int,
+                        unsigned int, unsigned int, unsigned int,
+                        unsigned int, void *, void **, void **);
+struct fft_dispatch {
+    complex_fn first_exec;
+    complex_fn final_exec;
+    launch_fn launch;
+    int first_handle;
+    int final_handle;
+    uint64_t slot;
+    uint64_t scratch;
+    void *stream;
+    void *function;
+    unsigned int threads;
+    unsigned int shared;
+};
+int fft_execute(const struct fft_dispatch *plan, uint64_t input, uint64_t output) {
+    uint64_t values[4] = {plan->slot, output, 0, 0};
+    void *parameters[4] = {&values[0], &values[1], &values[2], &values[3]};
+    int result = plan->launch(plan->function, 1, 1, 1,
+                             plan->threads, 1, 1, plan->shared, plan->stream,
+                             parameters, NULL);
+    /* Positive codes originate in cuFFT; negative codes in the CUDA driver. */
+    if (result != 0) return -result;
+    result = plan->first_exec(plan->first_handle, (void *)(uintptr_t)input,
+                            (void *)(uintptr_t)plan->scratch, 1);
+    if (result != 0 || plan->final_handle == 0) return result;
+    return plan->final_exec(plan->final_handle, (void *)(uintptr_t)plan->scratch,
+                              (void *)(uintptr_t)plan->scratch, 1);
+}
+size_t fft_dispatch_size(void) { return sizeof(struct fft_dispatch); }

build/torch-cuda/_b200/dense_half_inverse_native_dispatch.py ADDED Viewed

	@@ -0,0 +1,104 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Prepare a combined cuFFT and Triton dispatcher before execution."""
+from __future__ import annotations
+import ctypes as ct
+import fcntl
+import hashlib
+import os
+from pathlib import Path
+import shlex
+import subprocess
+import tempfile
+import threading
+from .dense_half_inverse_abi import validate_kernel_abi
+class Descriptor(ct.Structure):
+    _fields_ = [('first_exec', ct.c_void_p), ('final_exec', ct.c_void_p),
+                ('launch', ct.c_void_p), ('first_handle', ct.c_int),
+                ('final_handle', ct.c_int), ('slot', ct.c_uint64),
+                ('scratch', ct.c_uint64), ('stream', ct.c_void_p),
+                ('function', ct.c_void_p), ('threads', ct.c_uint),
+                ('shared', ct.c_uint)]
+_LIBRARY = None
+_DRIVER = None
+_LOCK = threading.Lock()
+def load_library():
+    """Build with the installed host compiler and validate structure size."""
+    global _LIBRARY
+    with _LOCK:
+        if _LIBRARY is not None:
+            return _LIBRARY
+        source = Path(__file__).with_name('dense_half_inverse_native_dispatch.c')
+        compiler = shlex.split(os.environ.get('CC', 'cc'))
+        flags = ['-std=c11', '-O3', '-fPIC', '-shared', '-Wall', '-Wextra', '-Werror']
+        key = hashlib.sha256(source.read_bytes() + repr((compiler, flags)).encode()).hexdigest()
+        cache = Path(tempfile.gettempdir()) / ('fft-half-inverse-native-' + str(os.getuid())) / key
+        cache.mkdir(parents=True, exist_ok=True)
+        target = cache / 'dispatch.so'
+        with (cache / 'build.lock').open('w') as lock:
+            fcntl.flock(lock, fcntl.LOCK_EX)
+            if not target.is_file():
+                pending = cache / ('dispatch.' + str(os.getpid()) + '.so')
+                try:
+                    subprocess.run(compiler + flags + [str(source), '-o', str(pending)],
+                                   check=True, capture_output=True, text=True)
+                    pending.replace(target)
+                finally:
+                    pending.unlink(missing_ok=True)
+        lib = ct.CDLL(str(target))
+        lib.fft_execute.argtypes = [ct.POINTER(Descriptor), ct.c_uint64, ct.c_uint64]
+        lib.fft_execute.restype = ct.c_int
+        lib.fft_dispatch_size.argtypes = []
+        lib.fft_dispatch_size.restype = ct.c_size_t
+        if lib.fft_dispatch_size() != ct.sizeof(Descriptor):
+            raise RuntimeError('Native dispatcher structure size mismatch')
+        _LIBRARY = lib
+        return lib
+def load_driver():
+    """Resolve the driver function during preparation only."""
+    global _DRIVER
+    with _LOCK:
+        if _DRIVER is None:
+            _DRIVER = ct.CDLL('libcuda.so.1')
+            _DRIVER.cuLaunchKernel.restype = ct.c_int
+        return _DRIVER
+class PreparedDispatch:
+    """Immutable setup descriptor, with fresh caller pointers at each call.
+    The Plan retains the cuFFT handles and compiled Triton kernel. Concurrent
+    execution requires independent plans because cuFFT workspace is reused.
+    """
+    def __init__(self, first_exec, first_handle, final_exec, final_handle,
+                 compiled, slot, scratch, stream):
+        abi = validate_kernel_abi(compiled)
+        if not compiled.function or compiled.module is None:
+            raise ValueError('The Triton kernel must be loaded during setup')
+        if not 0 < first_handle < 2**31 or not 0 <= final_handle < 2**31:
+            raise ValueError('Valid prepared cuFFT handles are required')
+        if not slot or not scratch:
+            raise ValueError('Prepared pointer slot and FFT scratch are required')
+        self.library = load_library()
+        self.driver = load_driver()
+        self.descriptor = Descriptor(
+            first_exec=ct.cast(first_exec, ct.c_void_p).value,
+            final_exec=ct.cast(final_exec, ct.c_void_p).value,
+            launch=ct.cast(self.driver.cuLaunchKernel, ct.c_void_p).value,
+            first_handle=first_handle, final_handle=final_handle,
+            slot=slot, scratch=scratch, stream=stream, function=compiled.function,
+            threads=abi['threads'], shared=abi['shared_bytes'])
+        self.pointer = ct.pointer(self.descriptor)
+        self.function = self.library.fft_execute
+    def execute(self, input_ptr, output_ptr):
+        return self.function(self.pointer, input_ptr, output_ptr)

build/torch-cuda/_b200/dense_half_inverse_schedule.py ADDED Viewed

	@@ -0,0 +1,73 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Host-only mixed-radix schedules and accurately rounded twiddle tables."""
+from __future__ import annotations
+import math
+def factors(length, strategy='radix8'):
+    """Return a deterministic DIT factorization; no measurements occur here."""
+    remaining = int(length)
+    if remaining < 1:
+        raise ValueError("FFT lengths must be positive")
+    result = []
+    if strategy not in ('radix4', 'radix8'):
+        raise ValueError('Unknown factorization strategy')
+    for radix in ((8, 4, 2) if strategy == 'radix8' else (4, 2)):
+        while remaining % radix == 0:
+            result.append(radix)
+            remaining //= radix
+    # Collapse the final pair of radix-3 stages for lengths 9 * 2**k.
+    # This gives the radix-9 kernel power-of-two branch lengths; other odd
+    # factorizations retain the established paired or generic network.
+    if remaining == 9:
+        result.append(9)
+        remaining = 1
+    for radix in (3, 5, 7):
+        while remaining % radix == 0:
+            result.append(radix)
+            remaining //= radix
+    if remaining != 1:
+        raise ValueError("This prototype supports lengths with factors 2, 3, 5, 7")
+    return tuple(result)
+def stage_sizes(radices):
+    sizes = []
+    length = 1
+    for radix in radices:
+        length *= radix
+        sizes.append(length)
+    return tuple(sizes)
+def permutation(length):
+    """Map DIT input positions to natural input indices, for CPU validation."""
+    radices = factors(length)
+    result = []
+    for index in range(length):
+        divisor, reverse = 1, 0
+        for radix in radices:
+            reverse += ((index // divisor) % radix) * (length // (divisor * radix))
+            divisor *= radix
+        result.append(reverse)
+    return result
+def twiddles(length):
+    """Interleaved exp(-2*pi*i*k/N), computed in host double precision."""
+    result = []
+    for index in range(length):
+        if index == 0:
+            value = (1.0, 0.0)
+        elif 2 * index == length:
+            value = (-1.0, 0.0)
+        elif 4 * index == length:
+            value = (0.0, -1.0)
+        elif 4 * index == 3 * length:
+            value = (0.0, 1.0)
+        else:
+            angle = -math.tau * index / length
+            value = (math.cos(angle), math.sin(angle))
+        result.extend(value)
+    return result

build/torch-cuda/_b200/direct_dispatch.c ADDED Viewed

	@@ -0,0 +1,53 @@

+/* SPDX-License-Identifier: Apache-2.0 */
+/* Execute a prepared cuFFT chain without per-stage Python crossings.
+ *
+ * The caller owns all handles and buffers. Function addresses are resolved
+ * during setup by the installed cuFFT loader. This translation unit has no
+ * CUDA dependency of its own and performs no allocation or synchronization.
+ */
+#include <stddef.h>
+enum { FFT_INPUT = 0, FFT_OUTPUT = 1, FFT_SCRATCH = 2 };
+enum { FFT_REAL = 0, FFT_COMPLEX = 1 };
+typedef int (*fft_real_fn)(int, void *, void *);
+typedef int (*fft_complex_fn)(int, void *, void *, int);
+struct fft_stage {
+    void *function;
+    int handle;
+    int kind;
+    int direction;
+    int input_slot;
+    int output_slot;
+};
+struct fft_chain {
+    int count;
+    void *scratch;
+    struct fft_stage stages[3];
+};
+int fft_execute(const struct fft_chain *chain, void *input, void *output) {
+    void *buffers[3] = {input, output, chain->scratch};
+    for (int i = 0; i < chain->count; ++i) {
+        const struct fft_stage *stage = &chain->stages[i];
+        int result;
+        if (stage->kind == FFT_REAL) {
+            result = ((fft_real_fn)stage->function)(
+                stage->handle, buffers[stage->input_slot], buffers[stage->output_slot]);
+        } else {
+            result = ((fft_complex_fn)stage->function)(
+                stage->handle, buffers[stage->input_slot], buffers[stage->output_slot],
+                stage->direction);
+        }
+        if (result != 0) {
+            return result;
+        }
+    }
+    return 0;
+}
+/* Used only by the CPU ABI validation, never by FFT execution. */
+size_t fft_chain_size(void) { return sizeof(struct fft_chain); }
+size_t fft_stage_size(void) { return sizeof(struct fft_stage); }

build/torch-cuda/_b200/direct_dispatch.py ADDED Viewed

	@@ -0,0 +1,98 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Compile and prepare a small host dispatcher during FFT plan creation."""
+from __future__ import annotations
+import ctypes as ct
+import fcntl
+import hashlib
+import os
+from pathlib import Path
+import shlex
+import subprocess
+import tempfile
+import threading
+class Stage(ct.Structure):
+    _fields_ = [("function", ct.c_void_p), ("handle", ct.c_int),
+                ("kind", ct.c_int), ("direction", ct.c_int),
+                ("input_slot", ct.c_int), ("output_slot", ct.c_int)]
+class Chain(ct.Structure):
+    _fields_ = [("count", ct.c_int), ("scratch", ct.c_void_p),
+                ("stages", Stage * 3)]
+_LIBRARY = None
+_LOCK = threading.Lock()
+def load_library():
+    """Build/load before execution; use only the installed host C compiler."""
+    global _LIBRARY
+    if _LIBRARY is not None:
+        return _LIBRARY
+    with _LOCK:
+        if _LIBRARY is not None:
+            return _LIBRARY
+        source = Path(__file__).with_name('direct_dispatch.c')
+        compiler = shlex.split(os.environ.get('CC', 'cc'))
+        flags = ['-std=c11', '-O3', '-fPIC', '-shared', '-Wall', '-Wextra', '-Werror']
+        key = hashlib.sha256(source.read_bytes() + repr((compiler, flags)).encode()).hexdigest()
+        cache = Path(tempfile.gettempdir()) / ('fft-native-' + str(os.getuid())) / key
+        cache.mkdir(parents=True, exist_ok=True)
+        target = cache / 'dispatch.so'
+        # A process lock prevents concurrent preparations from loading a partial DSO.
+        with (cache / 'build.lock').open('w') as lock:
+            fcntl.flock(lock, fcntl.LOCK_EX)
+            if not target.is_file():
+                pending = cache / ('dispatch.' + str(os.getpid()) + '.so')
+                try:
+                    subprocess.run(compiler + flags + [str(source), '-o', str(pending)],
+                                   check=True, capture_output=True, text=True)
+                    pending.replace(target)
+                finally:
+                    pending.unlink(missing_ok=True)
+        lib = ct.CDLL(str(target))
+        lib.fft_execute.argtypes = [ct.POINTER(Chain), ct.c_void_p, ct.c_void_p]
+        lib.fft_execute.restype = ct.c_int
+        for name in ('fft_chain_size', 'fft_stage_size'):
+            getattr(lib, name).argtypes = []
+            getattr(lib, name).restype = ct.c_size_t
+        if lib.fft_chain_size() != ct.sizeof(Chain) or lib.fft_stage_size() != ct.sizeof(Stage):
+            raise RuntimeError('Native FFT dispatcher ABI mismatch')
+        _LIBRARY = lib
+        return lib
+class PreparedChain:
+    """Immutable stage descriptors plus a pre-resolved execution function.
+    Each stage is (function, handle, kind, direction, input_slot, output_slot).
+    kind=0 denotes a real transform, kind=1 a complex transform. Slots are
+    0=input, 1=output, 2=mathematical intermediate. Tensor values are never
+    cached; caller input/output pointers are supplied anew on each call.
+    """
+    def __init__(self, stages, scratch_ptr=0):
+        if not 1 <= len(stages) <= 3:
+            raise ValueError('A chain must contain one to three FFT stages')
+        self.library = load_library()
+        self.functions = tuple(stage[0] for stage in stages)
+        self.chain = Chain(count=len(stages), scratch=scratch_ptr)
+        for index, (function, handle, kind, direction, source, target) in enumerate(stages):
+            if kind not in (0, 1) or source not in (0, 1, 2) or target not in (0, 1, 2):
+                raise ValueError('Invalid FFT stage metadata')
+            if (source == 2 or target == 2) and not scratch_ptr:
+                raise ValueError('Scratch pointer required for this FFT chain')
+            if kind == 1 and direction not in (-1, 1):
+                raise ValueError('Complex FFT direction must be -1 or +1')
+            address = ct.cast(function, ct.c_void_p).value
+            if not address:
+                raise ValueError('A resolved cuFFT function address is required')
+            self.chain.stages[index] = Stage(address, handle, kind, direction, source, target)
+        self.pointer = ct.pointer(self.chain)
+        self.function = self.library.fft_execute
+    def execute(self, input_ptr, output_ptr):
+        return self.function(self.pointer, input_ptr, output_ptr)

build/torch-cuda/_b200/fft_axis.cu ADDED Viewed

	@@ -0,0 +1,169 @@

+// SPDX-License-Identifier: Apache-2.0
+// Exact mathematical extents; each warp owns independent adjacent columns.
+__device__ __forceinline__ double2 add(double2 a, double2 b) {
+    return make_double2(a.x + b.x, a.y + b.y);
+}
+__device__ __forceinline__ double2 sub(double2 a, double2 b) {
+    return make_double2(a.x - b.x, a.y - b.y);
+}
+__device__ __forceinline__ double2 mul(double2 a, double2 b) {
+    return make_double2(a.x*b.x-a.y*b.y, a.x*b.y+a.y*b.x);
+}
+template<int H, int N>
+__device__ __forceinline__ double2 rotate_forward(double2 value, int k,
+                                                const double2 *twiddle) {
+    if constexpr (H == 1) {
+        return value;
+    } else if constexpr (H == 2) {
+        return k == 0 ? value : make_double2(value.y, -value.x);
+    }  else {
+        return mul(value, twiddle[k*(N/(2*H))]);
+    }
+}
+template<int H, int N, int G, int C, int Bound>
+__device__ __forceinline__ void stage(double2 (&a)[N/G], const double2 *twiddle,
+                                     int lane_z, double2 *exchange) {
+    if constexpr (H >= G && 4*H <= Bound) {
+        #pragma unroll
+        for (int r = 0; r < N/G; ++r) {
+            if ((r % (4*H/G)) < H/G) {
+                double2 a0 = a[r], a1 = a[r+H/G], a2 = a[r+2*H/G], a3 = a[r+3*H/G];
+                int k = (r*G+lane_z) % H;
+                double2 w1=twiddle[k*(N/(4*H))];
+                double2 w2=twiddle[2*k*(N/(4*H))];
+                double2 w3=twiddle[3*k*(N/(4*H))];
+                // Input is binary bit reversed, so the middle branches swap.
+                double2 b1=mul(a2,w1), b2=mul(a1,w2), b3=mul(a3,w3);
+                double2 even=add(a0,b2), even_diff=sub(a0,b2);
+                double2 odd=add(b1,b3), odd_diff=sub(b1,b3);
+                double2 turn=make_double2(odd_diff.y,-odd_diff.x);
+                a[r]=add(even,odd); a[r+H/G]=add(even_diff,turn);
+                a[r+2*H/G]=sub(even,odd); a[r+3*H/G]=sub(even_diff,turn);
+            }
+        }
+    } else if constexpr (H < G) {
+        // Consecutive shared stages alternate disjoint mathematical buffers.
+        constexpr int exchange_base=(H & 0xAAAAAAAAU) ? N*C : 0;
+        if constexpr (H*C >= 32) {
+            #pragma unroll
+            for (int r = 0; r < N/G; ++r) exchange[exchange_base+r*blockDim.x+threadIdx.x]=a[r];
+            __syncthreads();
+        }
+        #pragma unroll
+        for (int r = 0; r < N/G; ++r) {
+            double2 v;
+            if constexpr (H*C < 32) {
+                v = make_double2(__shfl_xor_sync(0xffffffff, a[r].x, H*C),
+                                 __shfl_xor_sync(0xffffffff, a[r].y, H*C));
+            } else {
+                v = exchange[exchange_base+r*blockDim.x+(threadIdx.x ^ (H*C))];
+            }
+            double2 lo = (lane_z & H) ? v : a[r];
+            double2 hi = (lane_z & H) ? a[r] : v;
+            if constexpr (H > 1) {
+                hi = rotate_forward<H, N>(hi, (lane_z & (H-1)), twiddle);
+            }
+            a[r] = (lane_z & H) ? sub(lo, hi) : add(lo, hi);
+        }
+        // The next producer barrier protects the opposite bank before reuse.
+    } else {
+        #pragma unroll
+        for (int r = 0; r < N/G; ++r) {
+            if ((r & (H/G)) == 0) {
+                double2 lo = a[r], hi = a[r + H/G];
+                hi = rotate_forward<H, N>(hi, ((r*G + lane_z) & (H-1)), twiddle);
+                a[r] = add(lo, hi);
+                a[r + H/G] = sub(lo, hi);
+            }
+        }
+    }
+    if constexpr (H >= G && 4*H <= Bound) {
+        if constexpr (4*H < Bound) stage<4*H, N, G, C, Bound>(a, twiddle, lane_z, exchange);
+    } else if constexpr (2*H < Bound) stage<2*H, N, G, C, Bound>(a, twiddle, lane_z, exchange);
+}
+#if FFT_DOUBLE
+using output_scalar = double;
+#else
+using output_scalar = float;
+#endif
+// Reconstruct packed XY spectra from the transformed real-pair operand.
+static __device__ __forceinline__ double2 reconstruct(
+        const double2 *input, const double2 *twiddle, int z, int column) {
+    constexpr int m = NX/2, nf = m+1;
+    const int k = column % nf, y = column / nf;
+    const int mirror_y = y == 0 ? 0 : NY-y;
+    const double2 a = input[(z*NY+y)*m+k%m];
+    const double2 b = input[(z*NY+mirror_y)*m+(m-k)%m];
+    const double2 w = twiddle[LENGTH+k];
+    const double er = 0.5*(a.x+b.x), ei = 0.5*(a.y-b.y);
+    const double odd_r = 0.5*(a.y+b.y), odd_i = -0.5*(a.x-b.x);
+    return make_double2(er+w.x*odd_r-w.y*odd_i,
+                        ei+w.x*odd_i+w.y*odd_r);
+}
+extern "C" __global__ __launch_bounds__(WARPS*32)
+void fft_axis(const double2 *__restrict__ input, output_scalar *__restrict__ output,
+              const double2 *__restrict__ twiddle, void *, void *) {
+    constexpr int G = WARPS*32/COLUMNS, R = LENGTH/G;
+    __shared__ double2 exchange[2*LENGTH*COLUMNS];
+    int lane_z = threadIdx.x / COLUMNS;
+    int column = blockIdx.x * COLUMNS + threadIdx.x % COLUMNS;
+    double2 a[R];
+    #pragma unroll
+    for (int r = 0; r < R; ++r) {
+        int index = r*G + lane_z;
+        int source = (__brev((unsigned)(index % BASE)) >> (32-LOG_BASE))*ODD
+                   + index/BASE;
+        a[r] = column < INNER ? reconstruct(input, twiddle, source, column) : make_double2(0., 0.);
+    }
+    stage<1, LENGTH, G, COLUMNS, BASE>(a, twiddle, lane_z, exchange);
+    if constexpr (ODD != 1) {
+        #pragma unroll
+        for (int r = 0; r < BASE/G; ++r) {
+            double2 b[ODD];
+            b[0] = a[r];
+            #pragma unroll
+            for (int t = 1; t < ODD; ++t) {
+                double2 w = twiddle[t*(r*G+lane_z)];
+                b[t] = mul(a[r+t*(BASE/G)], w);
+            }
+            #pragma unroll
+            for (int k = 0; k < ODD; ++k) {
+                double2 v = b[0];
+                #pragma unroll
+                for (int t = 1; t < ODD; ++t) {
+                    double2 w = twiddle[((k*t)%ODD)*BASE];
+                    v = add(v, mul(b[t], w));
+                }
+                a[r+k*(BASE/G)] = v;
+            }
+        }
+    }
+    #pragma unroll
+    for (int r = 0; r < R; ++r) {
+        if (column < INNER) {
+            int offset = 2*((r*G+lane_z)*INNER+column);
+#if FFT_DOUBLE
+            if ((reinterpret_cast<unsigned long long>(output) & 15ULL) == 0) {
+                reinterpret_cast<double2 *>(output)[offset/2] = a[r];
+            } else {
+                output[offset] = a[r].x;
+                output[offset+1] = a[r].y;
+            }
+#else
+            if ((reinterpret_cast<unsigned long long>(output) & 7ULL) == 0) {
+                reinterpret_cast<float2 *>(output)[offset/2] = make_float2((float)a[r].x, (float)a[r].y);
+            } else {
+                output[offset] = (float)a[r].x;
+                output[offset+1] = (float)a[r].y;
+            }
+#endif
+        }
+    }
+}

build/torch-cuda/_b200/forward_cuda_f32.py ADDED Viewed

	@@ -0,0 +1,151 @@

+# SPDX-License-Identifier: Apache-2.0
+"""FP64 half-length XY forward followed by a CUDA register Z forward."""
+from __future__ import annotations
+import math
+import threading
+import torch
+from .cuda_z_inverse_cufft_support import _Handle, _check, _lib
+from .cuda_kernel import compile_kernel
+from .cuda_twiddle_forward_lto import compile_callback
+from .cuda_twiddle_forward_native_dispatch import PreparedDispatch
+from .cuda_z_inverse_schedule import twiddles
+class Plan:
+    def __init__(self, shape, dtype, direction='r2c', *, device=None, stream=None,
+                 variant='xy_z', dispatch='native', columns=8, warps=4):
+        self.shape = tuple(int(n) for n in shape)
+        if len(self.shape) != 3 or min(self.shape[:2]) < 8 or self.shape[2] < 16 or self.shape[2] % 2:
+            raise ValueError('Expected (nz >= 8, ny >= 8, nx >= 16)')
+        if dtype not in (torch.float32, torch.float64) or direction != 'r2c':
+            raise ValueError('Expected a real forward transform')
+        if variant not in ('xy_z', 'xy_z_patient'):
+            raise ValueError('Unknown variant')
+        if dispatch != 'native':
+            raise ValueError('Unknown dispatch mode')
+        self.dtype, self.direction, self.variant = dtype, direction, variant
+        self.scalar_bytes = 8 if dtype == torch.float64 else 4
+        self.residues = 16 // self.scalar_bytes
+        self.dispatch = dispatch
+        if columns not in (2, 4, 8, 16, 32) or warps not in (2, 4, 8):
+            raise ValueError("Unsupported column tile or warp count")
+        self.columns, self.warps = columns, warps
+        self.native = []
+        self.device = torch.device(device if device is not None else torch.cuda.current_device())
+        if self.device.type != 'cuda':
+            raise ValueError('A CUDA device is required')
+        if self.device.index is None:
+            self.device = torch.device('cuda', torch.cuda.current_device())
+        self.stream = stream if stream is not None else torch.cuda.current_stream(self.device)
+        if self.stream.device != self.device:
+            raise ValueError('The supplied stream must belong to the plan device')
+        self.stream_ptr = self.stream.cuda_stream
+        self.closed = False
+        self.handles = []
+        self.compiled_metadata = []
+        nz, ny, nx = self.shape
+        nf = nx // 2 + 1
+        self.input_shape, self.output_shape = self.shape, (nz, ny, nf, 2)
+        m = nx // 2
+        self.fft_function = _lib.cufftExecZ2Z
+        try:
+            with torch.cuda.device(self.device), torch.cuda.stream(self.stream):
+                capability = torch.cuda.get_device_capability(self.device)
+                arch = capability[0] * 10 + capability[1]
+                # This intermediate contains the transformed half-length XY data.
+                # The last Z stage writes directly to scalar caller storage.
+                self.scratch = torch.empty((nz, ny, m, 2), dtype=torch.float64,
+                                           device=self.device)
+                self.scratch_ptr = self.scratch.data_ptr()
+                patient = variant.endswith('_patient')
+                for residue in range(self.residues):
+                    load_callback = compile_callback(dtype == torch.float64, 'pair', residue, arch, self.shape)[0]
+                    self.compiled_metadata.append(dict(load_callback[3], input_residue=residue))
+                    handle = _Handle(
+                        (ny, m), 0x69, (ny, m), 1, ny * m,
+                        (ny, m), 1, ny * m, nz, self.device, self.stream,
+                        callback=[load_callback], patient=patient)
+                    self.handles.append(handle)
+                self.handle_values = tuple(h.value for h in self.handles)
+                self.z_table = torch.tensor(twiddles(nz) + twiddles(nx), dtype=torch.float64,
+                                            device=self.device)
+                compiled = compile_kernel(self.shape, columns, warps, arch, self.device.index, dtype == torch.float64)
+                self.grid = compiled.grid
+                self.compiled_kernel = compiled
+                self.runner = None
+                if dispatch == 'native':
+                    for handle in self.handle_values:
+                        self.native.append(PreparedDispatch(
+                            self.fft_function, handle, compiled, self.grid,
+                            self.scratch_ptr, self.z_table.data_ptr(), self.stream_ptr))
+                    self.native_function = self.native[0].function
+                    self.native_pointers = tuple(n.pointer for n in self.native)
+                self.compiled_metadata.append(dict(name=compiled.name, hash=compiled.hash,
+                                                   shared=compiled.shared, compiler=compiled.compiler_metadata,
+                                                   registers=compiled.n_regs,
+                                                   spills=compiled.n_spills,
+                                                   grid=self.grid, warps=warps))
+                dummy = torch.zeros(math.prod(self.input_shape) + 4, dtype=dtype,
+                                    device=self.device)
+                output = torch.empty(self.output_shape, dtype=dtype, device=self.device)
+                for residue in range(self.residues):
+                    self.execute_ptr(dummy.data_ptr() + residue * self.scalar_bytes, output.data_ptr(), self.stream_ptr)
+                self.stream.synchronize()
+        except BaseException:
+            for handle in self.handles:
+                handle.close()
+            raise
+    @property
+    def workspace_bytes(self):
+        return (sum(h.workspace_bytes for h in self.handles)
+                + self.scratch.numel() * 8
+                + self.z_table.numel() * 8)
+    def execute_ptr(self, input_ptr, output_ptr, stream_ptr):
+        """Enqueue the prepared cuFFT XY and CUDA Z transforms."""
+        if self.closed:
+            raise RuntimeError('The plan is closed')
+        if int(stream_ptr) != self.stream_ptr:
+            raise ValueError('Prepare an independent plan for a different stream')
+        a, b = int(input_ptr), int(output_ptr)
+        if not a or not b or a == b or a % self.scalar_bytes or b % self.scalar_bytes:
+            raise ValueError('Distinct non-null scalar-aligned pointers are required')
+        residue = (a & 15) // self.scalar_bytes
+        result = self.native_function(self.native_pointers[residue],
+                                      a - self.scalar_bytes * residue, b)
+        if result > 0:
+            _check(result)
+        if result < 0:
+            raise RuntimeError(f'cuLaunchKernel failed with CUDA driver error {-result}')
+    def execute(self, input, output, stream=None):
+        for tensor, shape in ((input, self.input_shape), (output, self.output_shape)):
+            if (tuple(tensor.shape) != shape or tensor.dtype != self.dtype
+                    or tensor.device != self.device or not tensor.is_contiguous()):
+                raise ValueError('Buffer metadata does not match the plan')
+        a, b = input.data_ptr(), output.data_ptr()
+        na, nb = input.numel() * input.element_size(), output.numel() * output.element_size()
+        if a < b + nb and b < a + na:
+            raise ValueError('Input and output must not overlap')
+        selected = self.stream if stream is None else stream
+        self.execute_ptr(a, b, selected.cuda_stream)
+        return output
+    def close(self):
+        """The caller must wait for every outstanding use before closing."""
+        if not self.closed:
+            for handle in self.handles:
+                handle.close()
+            self.scratch = self.z_table = self.runner = self.compiled_kernel = None
+            self.native = self.native_pointers = None
+            self.closed = True
+_CACHE = {}
+_LOCK = threading.Lock()

build/torch-cuda/_b200/forward_half.py ADDED Viewed

	@@ -0,0 +1,149 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Half-length FP64 XY FFT with reconstruction fused into the final Z load.
+Both cuFFT stages consume the caller's packed layout or a mathematical FFT
+intermediate. The final single-kernel plan writes scalar caller storage.
+"""
+from __future__ import annotations
+import threading
+import torch
+from .scalar_input_inverse_cufft_support import _Handle, _lib, _check
+from .half_forward_cufft_z_lto import compile_callback
+from .direct_dispatch import PreparedChain
+from .half_forward_cufft_z_plan_safety import require_single_kernel
+from .cuda_z_inverse_schedule import twiddles
+class Plan:
+    def __init__(self, shape, dtype, direction='r2c', *, device=None, stream=None,
+                 axis='z', patient=False, xt=False):
+        self.shape = tuple(int(v) for v in shape)
+        if (len(self.shape) != 3 or min(self.shape[:2]) < 8
+                or self.shape[2] < 16 or self.shape[2] % 2):
+            raise ValueError('Expected (nz >= 8, ny >= 8, even nx >= 16)')
+        if dtype not in (torch.float32, torch.float64) or direction != 'r2c':
+            raise ValueError('Expected a float32 or float64 forward transform')
+        if axis != 'z':
+            raise ValueError('This decomposition requires Z as the last axis')
+        self.dtype, self.direction = dtype, direction
+        self.itemsize = 8 if dtype == torch.float64 else 4
+        self.residues = 16 // self.itemsize
+        self.device = torch.device(device if device is not None else torch.cuda.current_device())
+        if self.device.type != 'cuda':
+            raise ValueError('A CUDA device is required')
+        if self.device.index is None:
+            self.device = torch.device('cuda', torch.cuda.current_device())
+        self.stream = stream if stream is not None else torch.cuda.current_stream(self.device)
+        if self.stream.device != self.device:
+            raise ValueError('Stream and plan device must match')
+        self.stream_ptr = self.stream.cuda_stream
+        self.variant = f'half_xy_reconstruct_z_p{int(patient)}_xt{int(xt)}'
+        self.dispatch = 'prepared_native'
+        self.closed, self.handles, self.chains = False, [], []
+        nz, ny, nx = self.shape
+        rows, nf = nz * ny, nx // 2 + 1
+        first, other = (nz, ny) if axis == 'z' else (ny, nz)
+        self.input_shape, self.output_shape = self.shape, (nz, ny, nf, 2)
+        self.compiled_metadata = []
+        try:
+            with torch.cuda.device(self.device), torch.cuda.stream(self.stream):
+                self.intermediate = torch.empty((nz, ny, nx // 2, 2), dtype=torch.float64, device=self.device)
+                self.table = torch.tensor(twiddles(nx), dtype=torch.float64, device=self.device)
+                cap = torch.cuda.get_device_capability(self.device)
+                arch = 10 * cap[0] + cap[1]
+                load_handles, store_handles = [], []
+                for shift in range(self.residues):
+                    cb = compile_callback(dtype == torch.float64, 'pair', shift, arch, self.shape, axis)
+                    self.compiled_metadata.extend(item[3] for item in cb)
+                    handle = _Handle((ny, nx // 2), 0x69, (ny, nx // 2), 1, ny * (nx // 2),
+                                     (ny, nx // 2), 1, ny * (nx // 2), nz, self.device, self.stream,
+                                     callback=cb, patient=patient)
+                    self.handles.append(handle)
+                    load_handles.append(handle)
+                    cb = (compile_callback(dtype == torch.float64, 'reconstruct', 0, arch, self.shape, axis)
+                          + compile_callback(dtype == torch.float64, 'store', shift, arch, self.shape, axis))
+                    self.compiled_metadata.extend(item[3] for item in cb)
+                    handle = _Handle((first,), 0x69, (first,), other * nf, 1,
+                                     (first,), other * nf, 1, other * nf, self.device, self.stream,
+                                     callback=cb, patient=patient, xt=xt, caller_info=self.table.data_ptr())
+                    self.handles.append(handle)
+                    store_handles.append(handle)
+                for input_handle in load_handles:
+                    for output_handle in store_handles:
+                        self.chains.append(PreparedChain([
+                            (_lib.cufftExecZ2Z, input_handle.value, 1, -1, 0, 2),
+                            (_lib.cufftExecZ2Z, output_handle.value, 1, -1, 2, 1)
+                        ], self.intermediate.data_ptr()))
+                self.native_function = self.chains[0].function
+                self.native_pointers = tuple(c.pointer for c in self.chains)
+                # Every prepared callback is loaded and warmed, including all
+                # address combinations. Dummy buffers are setup-only storage.
+                dummy_in = torch.zeros(rows * nx + self.residues,
+                                       dtype=dtype, device=self.device)
+                # Oversized nominal double-complex storage is setup-only. Every final
+                # handle must prove it has a single kernel before caller-sized use.
+                dummy_out = torch.empty(nz * ny * nf * 16 + 16,
+                                        dtype=torch.uint8, device=self.device)
+                self.final_plan_checks = []
+                for handle in store_handles:
+                    result = require_single_kernel(handle, self.intermediate, dummy_out,
+                                                   self.device, self.stream)
+                    self.final_plan_checks.append(result)
+                for i in range(self.residues):
+                    for o in range(self.residues):
+                        self.execute_ptr(dummy_in.data_ptr() + i * self.itemsize,
+                                         dummy_out.data_ptr() + o * self.itemsize,
+                                         self.stream_ptr)
+                self.stream.synchronize()
+        except BaseException:
+            for handle in reversed(self.handles):
+                handle.close()
+            self.closed = True
+            raise
+    @property
+    def workspace_bytes(self):
+        return sum(h.workspace_bytes for h in self.handles) + (self.intermediate.numel() + self.table.numel()) * 8
+    def execute_ptr(self, input_ptr, output_ptr, stream_ptr):
+        """Enqueue using current scalar-aligned addresses on the bound stream."""
+        if self.closed:
+            raise RuntimeError('The plan is closed')
+        if int(stream_ptr) != self.stream_ptr:
+            raise ValueError('Create a separate plan for the supplied stream')
+        a, b = int(input_ptr), int(output_ptr)
+        if not a or not b or a == b:
+            raise ValueError('Distinct non-null caller pointers are required')
+        if a % self.itemsize or b % self.itemsize:
+            raise ValueError('Caller buffers must be scalar aligned')
+        i, o = (a & 15) // self.itemsize, (b & 15) // self.itemsize
+        _check(self.native_function(self.native_pointers[i * self.residues + o],
+                                    a - (a & 15), b - (b & 15)))
+    def execute(self, input, output, stream=None):
+        for tensor, shape in ((input, self.input_shape), (output, self.output_shape)):
+            if (tuple(tensor.shape) != shape or tensor.dtype != self.dtype
+                    or tensor.device != self.device or not tensor.is_contiguous()):
+                raise ValueError('Buffer metadata does not match the plan')
+        a, b = input.data_ptr(), output.data_ptr()
+        na, nb = input.numel() * self.itemsize, output.numel() * self.itemsize
+        if a < b + nb and b < a + na:
+            raise ValueError('Input and output must not overlap')
+        self.execute_ptr(a, b, (self.stream if stream is None else stream).cuda_stream)
+        return output
+    def close(self):
+        """The caller must complete outstanding work before closing the plan."""
+        if not self.closed:
+            for handle in reversed(self.handles):
+                handle.close()
+            self.intermediate = self.table = self.chains = self.native_pointers = None
+            self.closed = True
+_CACHE = {}
+_LOCK = threading.Lock()

build/torch-cuda/_b200/forward_native.py ADDED Viewed

	@@ -0,0 +1,182 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Preplanned Z-then-XY cuFFT inverse with Hermitian boundary projection.
+The load callback projects the original Z/Y endpoint planes as the complex
+Z transform reads them. The separate-kernel alternative projects Y pairs
+after the Z transform. Both preserve arbitrary packed-input semantics.
+"""
+from __future__ import annotations
+import threading
+import torch
+import triton
+from .hermitian_projection_cufft_support import _Handle, _check, _lib
+from .hermitian_projection_kernels import project_y_endpoints
+from .hermitian_projection_lto import compile_callback
+from .direct_dispatch import PreparedChain
+class Plan:
+    """All preparation occurs in the constructor, on the supplied stream.
+    Execute calls only enqueue work. Callers own nonoverlapping input/output
+    buffers and must keep them and the plan alive until stream completion.
+    Prepare independent plans for concurrent streams or host threads.
+    """
+    def __init__(self, shape, dtype, direction, *, device=None, stream=None,
+                 variant=None):
+        self.shape = tuple(int(n) for n in shape)
+        if len(self.shape) != 3 or min(self.shape[:2]) < 8 or self.shape[2] < 16:
+            raise ValueError('Expected (nz >= 8, ny >= 8, nx >= 16)')
+        if dtype not in (torch.float32, torch.float64):
+            raise TypeError('Only float32 and float64 are supported')
+        if direction not in ('r2c', 'c2r'):
+            raise ValueError('direction must be r2c or c2r')
+        if direction == 'c2r' and self.shape[2] % 2:
+            raise ValueError('The inverse request infers an even nx from nf')
+        self.variant = variant or ('3d' if direction == 'r2c' else 'z_xy_load')
+        permitted = (('3d', 'xy_z', 'x_zy') if direction == 'r2c' else
+                     ('z_xy_load', 'z_xy_kernel'))
+        if self.variant not in permitted:
+            raise ValueError(f'Unknown {direction} variant: {self.variant}')
+        self.dtype, self.direction = dtype, direction
+        self.device = torch.device(device if device is not None else torch.cuda.current_device())
+        if self.device.type != 'cuda':
+            raise ValueError('A CUDA device is required')
+        if self.device.index is None:
+            self.device = torch.device('cuda', torch.cuda.current_device())
+        self.stream = stream if stream is not None else torch.cuda.current_stream(self.device)
+        if self.stream.device != self.device:
+            raise ValueError('The supplied stream is on a different device')
+        self.stream_ptr = self.stream.cuda_stream
+        self.closed = False
+        self.handles = []
+        self.intermediate = None
+        self._chain = None
+        self.runner = None
+        self.compiled_metadata = []
+        self.dispatch = 'native_cufft' if self.variant != 'z_xy_kernel' else 'cufft_projection_kernel'
+        nz, ny, nx = self.shape
+        nf = nx // 2 + 1
+        packed = (nz, ny, nf, 2)
+        inverse = direction == 'c2r'
+        self.input_shape, self.output_shape = (packed, self.shape) if inverse else (self.shape, packed)
+        double = dtype == torch.float64
+        rkind, ckind, ikind = (0x6A, 0x69, 0x6C) if double else (0x2A, 0x29, 0x2C)
+        self._real_exec = getattr(_lib, 'cufftExec' + (
+            ('Z2D' if double else 'C2R') if inverse else ('D2Z' if double else 'R2C')))
+        self._complex_exec = _lib.cufftExecZ2Z if double else _lib.cufftExecC2C
+        try:
+            with torch.cuda.device(self.device), torch.cuda.stream(self.stream):
+                def plan(sizes, kind, ie, istr, idist, oe, ostr, odist, batch,
+                         callback=None, patient=False):
+                    handle = _Handle(sizes, kind, ie, istr, idist, oe, ostr, odist,
+                                     batch, self.device, self.stream, callback, patient)
+                    self.handles.append(handle)
+                    return handle.value
+                if not inverse:
+                    if self.variant == '3d':
+                        self.h0 = plan(self.shape, rkind, self.shape, 1, nz * ny * nx,
+                                       (nz, ny, nf), 1, nz * ny * nf, 1)
+                    elif self.variant == 'xy_z':
+                        self.h0 = plan((ny, nx), rkind, (ny, nx), 1, ny * nx,
+                                       (ny, nf), 1, ny * nf, nz)
+                        self.h1 = plan((nz,), ckind, (nz,), ny * nf, 1,
+                                       (nz,), ny * nf, 1, ny * nf)
+                    else:
+                        self.h0 = plan((nx,), rkind, (nx,), 1, nx, (nf,), 1, nf, nz * ny)
+                        self.h1 = plan((nz, ny), ckind, (nz, ny), nf, 1,
+                                       (nz, ny), nf, 1, nf)
+                    stages = [(self._real_exec, self.h0, 0, 0, 0, 1)]
+                    if self.variant in ('xy_z', 'x_zy'):
+                        stages.append((self._complex_exec, self.h1, 1, -1, 1, 1))
+                    self._chain = PreparedChain(stages)
+                else:
+                    callback = None
+                    if self.variant == 'z_xy_load':
+                        capability = torch.cuda.get_device_capability(self.device)
+                        callback = compile_callback(self.shape, double,
+                                                    10 * capability[0] + capability[1])
+                        self.compiled_metadata.append(callback[3])
+                    self.intermediate = torch.empty(packed, dtype=dtype, device=self.device)
+                    self.tmp_ptr = self.intermediate.data_ptr()
+                    self.h0 = plan((nz,), ckind, (nz,), ny * nf, 1,
+                                   (nz,), ny * nf, 1, ny * nf, callback=callback)
+                    self.h1 = plan((ny, nx), ikind, (ny, nf), 1, ny * nf,
+                                   (ny, nx), 1, ny * nx, nz)
+                    if self.variant == 'z_xy_kernel':
+                        self.constants = (nz, ny, nf, 256)
+                        grid = (triton.cdiv(2 * nz * (ny // 2 + 1), 256), 1, 1)
+                        compiled = project_y_endpoints[grid](self.intermediate, *self.constants,
+                                                             num_warps=4)
+                        self.runner = compiled[grid]
+                        self.compiled_metadata.append(dict(name=compiled.name, hash=compiled.hash,
+                                                           shared=compiled.metadata.shared,
+                                                           registers=compiled.n_regs,
+                                                           spills=compiled.n_spills, grid=grid))
+                    else:
+                        self._chain = PreparedChain([
+                            (self._complex_exec, self.h0, 1, 1, 0, 2),
+                            (self._real_exec, self.h1, 0, 0, 2, 1)], self.tmp_ptr)
+                if self._chain is not None:
+                    self._native_function, self._native_pointer = self._chain.function, self._chain.pointer
+                dummy_in = torch.zeros(self.input_shape, dtype=dtype, device=self.device)
+                dummy_out = torch.empty(self.output_shape, dtype=dtype, device=self.device)
+                self.execute_ptr(dummy_in.data_ptr(), dummy_out.data_ptr(), self.stream_ptr)
+                self.stream.synchronize()
+        except BaseException:
+            for handle in self.handles:
+                handle.close()
+            raise
+    @property
+    def workspace_bytes(self):
+        return sum(h.workspace_bytes for h in self.handles) + (
+            self.intermediate.numel() * self.intermediate.element_size()
+            if self.intermediate is not None else 0)
+    def execute_ptr(self, input_ptr, output_ptr, stream_ptr):
+        """Enqueue prepared stages with exact packed pointers on the bound stream."""
+        if self.closed:
+            raise RuntimeError('The plan is closed')
+        if int(stream_ptr) != self.stream_ptr:
+            raise ValueError('Create a separate plan for the supplied stream')
+        if not input_ptr or not output_ptr or int(input_ptr) == int(output_ptr):
+            raise ValueError('Distinct non-null input/output pointers are required')
+        if self._chain is not None:
+            _check(self._native_function(self._native_pointer, input_ptr, output_ptr))
+        else:
+            _check(self._complex_exec(self.h0, input_ptr, self.tmp_ptr, 1))
+            self.runner(self.tmp_ptr, *self.constants, stream=self.stream_ptr)
+            _check(self._real_exec(self.h1, self.tmp_ptr, output_ptr))
+    def execute(self, input, output, stream=None):
+        """Check caller buffers, enqueue the transform, and return output."""
+        for tensor, shape in ((input, self.input_shape), (output, self.output_shape)):
+            if (tuple(tensor.shape) != shape or tensor.dtype != self.dtype
+                    or tensor.device != self.device or not tensor.is_contiguous()):
+                raise ValueError('Buffer metadata does not match the plan')
+        a, b = input.data_ptr(), output.data_ptr()
+        na, nb = input.numel() * input.element_size(), output.numel() * output.element_size()
+        if a < b + nb and b < a + na:
+            raise ValueError('Input and output buffers must not overlap')
+        selected = self.stream if stream is None else stream
+        self.execute_ptr(a, b, selected.cuda_stream)
+        return output
+    def close(self):
+        """Release setup resources after the caller has waited for all uses."""
+        if not self.closed:
+            for handle in self.handles:
+                handle.close()
+            self.intermediate = self.runner = self._chain = None
+            self.closed = True
+_CACHE = {}
+_LOCK = threading.Lock()

build/torch-cuda/_b200/forward_paired_y.py ADDED Viewed

	@@ -0,0 +1,153 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Y-paired half-length FP64 XY FFT with reconstruction fused into the final Z load.
+Both cuFFT stages consume the caller's packed layout or a mathematical FFT
+intermediate. The final single-kernel plan writes scalar caller storage.
+"""
+from __future__ import annotations
+import threading
+import torch
+from .scalar_input_inverse_cufft_support import _Handle, _lib, _check
+from .paired_y_forward_lto import compile_callback
+from .direct_dispatch import PreparedChain
+from .half_forward_cufft_z_plan_safety import require_single_kernel
+from .cuda_z_inverse_schedule import twiddles
+class Plan:
+    def __init__(self, shape, dtype, direction='r2c', *, device=None, stream=None,
+                 axis='z', patient=False, xt=False, pair_order='yx'):
+        self.shape = tuple(int(v) for v in shape)
+        if (len(self.shape) != 3 or min(self.shape[:2]) < 8
+                or self.shape[2] < 16 or self.shape[2] % 2 or self.shape[1] % 2):
+            raise ValueError('Expected (nz >= 8, ny >= 8, even nx >= 16)')
+        if dtype not in (torch.float32, torch.float64) or direction != 'r2c':
+            raise ValueError('Expected a float32 or float64 forward transform')
+        if axis != 'z':
+            raise ValueError('This decomposition requires Z as the last axis')
+        if pair_order not in ('yx', 'xy'):
+            raise ValueError('Expected yx or xy pair order')
+        self.pair_order = pair_order
+        self.dtype, self.direction = dtype, direction
+        self.itemsize = 8 if dtype == torch.float64 else 4
+        self.residues = 16 // self.itemsize
+        self.device = torch.device(device if device is not None else torch.cuda.current_device())
+        if self.device.type != 'cuda':
+            raise ValueError('A CUDA device is required')
+        if self.device.index is None:
+            self.device = torch.device('cuda', torch.cuda.current_device())
+        self.stream = stream if stream is not None else torch.cuda.current_stream(self.device)
+        if self.stream.device != self.device:
+            raise ValueError('Stream and plan device must match')
+        self.stream_ptr = self.stream.cuda_stream
+        self.variant = f'paired_y_{pair_order}_reconstruct_z_p{int(patient)}_xt{int(xt)}'
+        self.dispatch = 'prepared_native'
+        self.closed, self.handles, self.chains = False, [], []
+        nz, ny, nx = self.shape
+        rows, nf = nz * ny, nx // 2 + 1
+        first, other = (nz, ny) if axis == 'z' else (ny, nz)
+        self.input_shape, self.output_shape = self.shape, (nz, ny, nf, 2)
+        self.compiled_metadata = []
+        try:
+            with torch.cuda.device(self.device), torch.cuda.stream(self.stream):
+                first_shape = (ny // 2, nx) if pair_order == 'yx' else (nx, ny // 2)
+                self.intermediate = torch.empty((nz, *first_shape, 2), dtype=torch.float64, device=self.device)
+                self.table = torch.tensor(twiddles(ny), dtype=torch.float64, device=self.device)
+                cap = torch.cuda.get_device_capability(self.device)
+                arch = 10 * cap[0] + cap[1]
+                load_handles, store_handles = [], []
+                for shift in range(self.residues):
+                    cb = compile_callback(dtype == torch.float64, 'pair', shift, arch, self.shape, axis, pair_order)
+                    self.compiled_metadata.extend(item[3] for item in cb)
+                    handle = _Handle(first_shape, 0x69, first_shape, 1, ny * (nx // 2),
+                                     first_shape, 1, ny * (nx // 2), nz, self.device, self.stream,
+                                     callback=cb, patient=patient)
+                    self.handles.append(handle)
+                    load_handles.append(handle)
+                    cb = (compile_callback(dtype == torch.float64, 'reconstruct', 0, arch, self.shape, axis, pair_order)
+                          + compile_callback(dtype == torch.float64, 'store', shift, arch, self.shape, axis, pair_order))
+                    self.compiled_metadata.extend(item[3] for item in cb)
+                    handle = _Handle((first,), 0x69, (first,), other * nf, 1,
+                                     (first,), other * nf, 1, other * nf, self.device, self.stream,
+                                     callback=cb, patient=patient, xt=xt, caller_info=self.table.data_ptr())
+                    self.handles.append(handle)
+                    store_handles.append(handle)
+                for input_handle in load_handles:
+                    for output_handle in store_handles:
+                        self.chains.append(PreparedChain([
+                            (_lib.cufftExecZ2Z, input_handle.value, 1, -1, 0, 2),
+                            (_lib.cufftExecZ2Z, output_handle.value, 1, -1, 2, 1)
+                        ], self.intermediate.data_ptr()))
+                self.native_function = self.chains[0].function
+                self.native_pointers = tuple(c.pointer for c in self.chains)
+                # Every prepared callback is loaded and warmed, including all
+                # address combinations. Dummy buffers are setup-only storage.
+                dummy_in = torch.zeros(rows * nx + self.residues,
+                                       dtype=dtype, device=self.device)
+                # Oversized nominal double-complex storage is setup-only. Every final
+                # handle must prove it has a single kernel before caller-sized use.
+                dummy_out = torch.empty(nz * ny * nf * 16 + 16,
+                                        dtype=torch.uint8, device=self.device)
+                self.final_plan_checks = []
+                for handle in store_handles:
+                    result = require_single_kernel(handle, self.intermediate, dummy_out,
+                                                   self.device, self.stream)
+                    self.final_plan_checks.append(result)
+                for i in range(self.residues):
+                    for o in range(self.residues):
+                        self.execute_ptr(dummy_in.data_ptr() + i * self.itemsize,
+                                         dummy_out.data_ptr() + o * self.itemsize,
+                                         self.stream_ptr)
+                self.stream.synchronize()
+        except BaseException:
+            for handle in reversed(self.handles):
+                handle.close()
+            self.closed = True
+            raise
+    @property
+    def workspace_bytes(self):
+        return sum(h.workspace_bytes for h in self.handles) + (self.intermediate.numel() + self.table.numel()) * 8
+    def execute_ptr(self, input_ptr, output_ptr, stream_ptr):
+        """Enqueue using current scalar-aligned addresses on the bound stream."""
+        if self.closed:
+            raise RuntimeError('The plan is closed')
+        if int(stream_ptr) != self.stream_ptr:
+            raise ValueError('Create a separate plan for the supplied stream')
+        a, b = int(input_ptr), int(output_ptr)
+        if not a or not b or a == b:
+            raise ValueError('Distinct non-null caller pointers are required')
+        if a % self.itemsize or b % self.itemsize:
+            raise ValueError('Caller buffers must be scalar aligned')
+        i, o = (a & 15) // self.itemsize, (b & 15) // self.itemsize
+        _check(self.native_function(self.native_pointers[i * self.residues + o],
+                                    a - (a & 15), b - (b & 15)))
+    def execute(self, input, output, stream=None):
+        for tensor, shape in ((input, self.input_shape), (output, self.output_shape)):
+            if (tuple(tensor.shape) != shape or tensor.dtype != self.dtype
+                    or tensor.device != self.device or not tensor.is_contiguous()):
+                raise ValueError('Buffer metadata does not match the plan')
+        a, b = input.data_ptr(), output.data_ptr()
+        na, nb = input.numel() * self.itemsize, output.numel() * self.itemsize
+        if a < b + nb and b < a + na:
+            raise ValueError('Input and output must not overlap')
+        self.execute_ptr(a, b, (self.stream if stream is None else stream).cuda_stream)
+        return output
+    def close(self):
+        """The caller must complete outstanding work before closing the plan."""
+        if not self.closed:
+            for handle in reversed(self.handles):
+                handle.close()
+            self.intermediate = self.table = self.chains = self.native_pointers = None
+            self.closed = True
+_CACHE = {}
+_LOCK = threading.Lock()

build/torch-cuda/_b200/forward_scalar.py ADDED Viewed

	@@ -0,0 +1,63 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Prepare scalar-safe FP64 forward access alongside a measured aligned plan.
+The scalar route reads even/odd caller samples through prepared cuFFT load
+callbacks. Its exact-size half-length complex FFT intermediate contains
+transformed data. Reconstruction writes the exact caller packed output,
+using alignment-aware scalar/vector stores without a copy or repack.
+"""
+from .paired_vector_store_forward_backend import Plan as ScalarPlan
+class Plan:
+    def __init__(self, aligned):
+        self._aligned, self._scalar = aligned, None
+        self.closed = False
+        try:
+            for name in ('shape', 'dtype', 'direction', 'device', 'stream',
+                         'stream_ptr', 'input_shape', 'output_shape'):
+                setattr(self, name, getattr(aligned, name))
+            self.variant = 'scalar_forward_' + aligned.variant
+            self._scalar = ScalarPlan(self.shape, self.dtype, self.direction,
+                device=self.device, stream=self.stream, block=256, patient=False)
+        except BaseException:
+            aligned.close()
+            self.closed = True
+            raise
+    @property
+    def workspace_bytes(self):
+        return self._aligned.workspace_bytes + self._scalar.workspace_bytes
+    def execute_ptr(self, input_ptr, output_ptr, stream_ptr):
+        """Select an already prepared route by current caller addresses only."""
+        if self.closed:
+            raise RuntimeError('The plan is closed')
+        if int(stream_ptr) != self.stream_ptr:
+            raise ValueError('Create a separate plan for the supplied stream')
+        if not input_ptr or not output_ptr or int(input_ptr) == int(output_ptr):
+            raise ValueError('Distinct non-null caller pointers are required')
+        if int(input_ptr) % 8 or int(output_ptr) % 8:
+            raise ValueError('Caller FP64 buffers must be scalar aligned')
+        selected = self._scalar if (int(input_ptr) | int(output_ptr)) & 8 else self._aligned
+        selected.execute_ptr(input_ptr, output_ptr, stream_ptr)
+    def execute(self, input, output, stream=None):
+        for tensor, shape in ((input, self.input_shape), (output, self.output_shape)):
+            if (tuple(tensor.shape) != shape or tensor.dtype != self.dtype or
+                    tensor.device != self.device or not tensor.is_contiguous()):
+                raise ValueError('Buffer metadata does not match the plan')
+        a, b = input.data_ptr(), output.data_ptr()
+        if a < b + output.numel() * 8 and b < a + input.numel() * 8:
+            raise ValueError('Input and output must not overlap')
+        self.execute_ptr(a, b, (self.stream if stream is None else stream).cuda_stream)
+        return output
+    def close(self):
+        """Release prepared resources after the caller waits for queued work."""
+        if not self.closed:
+            try:
+                self._scalar.close()
+            finally:
+                self._aligned.close()
+                self.closed = True

build/torch-cuda/_b200/forward_small.py ADDED Viewed

	@@ -0,0 +1,195 @@

+# SPDX-License-Identifier: Apache-2.0
+"""Precompiled two-pass XY/Z FFT with FP64 arithmetic and caller-precision boundaries.
+This is an experimental implementation. CPU interpreter validation is not
+GPU validation; do not promote it without the full frozen and deployment
+checks. Execution can submit all stages through one prepared native chain.
+"""
+from __future__ import annotations
+import threading
+import torch
+import triton
+from .wide_native_kernels import complex_axis, packed_real_axis, real_axis, real_xy
+from .wide_native_schedule import factors, stage_sizes, twiddles
+from .grouped_dif_inverse_native_dispatch import PreparedChain, kernel_descriptor
+class Plan:
+    def __init__(self, shape, dtype, direction, *, device=None, stream=None,
+                 rows_per_block=4, columns_per_block=8, num_warps=4,
+                 real_strategy='fused_xy', dispatch='native'):
+        if dispatch not in ('native', 'python'):
+            raise ValueError('dispatch must be native or python')
+        self.native = None
+        self.descriptors = []
+        self.compiled_kernels = []
+        self.shape = tuple(int(n) for n in shape)
+        if len(self.shape) != 3 or min(self.shape[:2]) < 8 or self.shape[2] < 16:
+            raise ValueError('Expected (nz >= 8, ny >= 8, nx >= 16)')
+        if dtype not in (torch.float32, torch.float64):
+            raise TypeError('Only float32 and float64 are supported')
+        if direction not in ('r2c', 'c2r'):
+            raise ValueError('direction must be r2c or c2r')
+        if real_strategy != 'fused_xy':
+            raise ValueError('This wide variant requires the fused_xy strategy')
+        if direction == 'c2r' and self.shape[2] % 2:
+            raise ValueError('The inverse request infers even nx from nf')
+        for length in self.shape:
+            factors(length)
+        for block in (rows_per_block, columns_per_block, num_warps):
+            if block < 1 or block & (block - 1):
+                raise ValueError('Block sizes and warp count must be powers of two')
+        self.dtype, self.direction = dtype, direction
+        self.device = torch.device(device if device is not None else torch.cuda.current_device())
+        if self.device.type != 'cuda':
+            raise ValueError('A CUDA device is required')
+        if self.device.index is None:
+            self.device = torch.device('cuda', torch.cuda.current_device())
+        self.stream = stream if stream is not None else torch.cuda.current_stream(self.device)
+        if self.stream.device != self.device:
+            raise ValueError('The supplied stream is on a different device')
+        self.stream_ptr = self.stream.cuda_stream
+        # Odd forward lengths use the general real-axis network. All requested
+        # workloads have even nx and directly use the selected real strategy.
+        self.real_strategy = (real_strategy if self.shape[2] % 2 == 0 or real_strategy == 'fused_xy' else 'full')
+        if self.real_strategy == 'fused_xy' and (triton.next_power_of_2(self.shape[1]) * triton.next_power_of_2(self.shape[2]) > 16384):
+            raise ValueError('The fused experimental plane is limited to 16384 register slots')
+        self.variant = f'{self.real_strategy}_r{rows_per_block}_c{columns_per_block}_w{num_warps}'
+        self.dispatch = dispatch
+        self.closed = False
+        nz, ny, nx = self.shape
+        nf = nx // 2 + 1
+        packed = (nz, ny, nf, 2)
+        inverse = direction == 'c2r'
+        self.input_shape, self.output_shape = (packed, self.shape) if inverse else (self.shape, packed)
+        self.launches = []
+        self.compiled_metadata = []
+        self.intermediate = None
+        self.tables = {}
+        with torch.cuda.device(self.device), torch.cuda.stream(self.stream):
+            for length in set(self.shape):
+                self.tables[length] = torch.tensor(twiddles(length), dtype=torch.float64, device=self.device)
+            self.intermediate = torch.empty(packed, dtype=torch.float64, device=self.device)
+            dummy_in = torch.zeros(self.input_shape, dtype=dtype, device=self.device)
+            dummy_out = torch.empty(self.output_shape, dtype=dtype, device=self.device)
+            buffers = (dummy_in, dummy_out, self.intermediate)
+            def add(kernel, grid, input_slot, output_slot, length, constants):
+                tables = tuple(self.tables[n] for n in (length if isinstance(length, tuple) else (length,)))
+                stage_warps = num_warps if kernel == real_xy else 4
+                compiled = kernel[grid](buffers[input_slot], buffers[output_slot], *tables,
+                                        *constants, num_warps=stage_warps)
+                if compiled.metadata.global_scratch_size or compiled.metadata.profile_scratch_size:
+                    raise RuntimeError('Execution scratch allocation is not allowed')
+                self.compiled_kernels.append(compiled)
+                if dispatch == 'native':
+                    self.descriptors.append(kernel_descriptor(
+                        compiled, grid, input_slot, output_slot,
+                        tuple(t.data_ptr() for t in tables)))
+                # Construct and initialize the low-level launcher during setup.
+                runner = compiled[grid]
+                self.launches.append((runner, input_slot, output_slot, tuple(t.data_ptr() for t in tables), constants))
+                self.compiled_metadata.append(dict(name=compiled.name, hash=compiled.hash,
+                                                   shared=compiled.metadata.shared,
+                                                   registers=compiled.n_regs,
+                                                   spills=compiled.n_spills,
+                                                   grid=grid, num_warps=stage_warps))
+            def real(input_slot, output_slot):
+                length = nx // 2 if self.real_strategy == 'half' else nx
+                kernel = packed_real_axis if self.real_strategy == 'half' else real_axis
+                radices = factors(length)
+                constants = (nz * ny, nx, triton.next_power_of_2(length), rows_per_block,
+                             inverse, radices, stage_sizes(radices))
+                add(kernel, (triton.cdiv(nz * ny, rows_per_block), 1, 1),
+                    input_slot, output_slot, nx, constants)
+            def complex(input_slot, output_slot, length, inner, outer):
+                radices = factors(length)
+                constants = (length, inner, triton.next_power_of_2(length), columns_per_block,
+                             inverse, radices, stage_sizes(radices))
+                add(complex_axis, (triton.cdiv(inner, columns_per_block), outer, 1),
+                    input_slot, output_slot, length, constants)
+            def xy(input_slot, output_slot):
+                rx, ry = factors(nx), factors(ny)
+                constants = (ny, nx, triton.next_power_of_2(ny), triton.next_power_of_2(nx),
+                             inverse, rx, stage_sizes(rx), ry, stage_sizes(ry))
+                add(real_xy, (nz, 1, 1), input_slot, output_slot, (nx, ny), constants)
+            if self.real_strategy == 'fused_xy':
+                if inverse:
+                    complex(0, 2, nz, ny * nf, 1)
+                    xy(2, 1)
+                else:
+                    xy(0, 2)
+                    complex(2, 1, nz, ny * nf, 1)
+            elif inverse:
+                complex(0, 2, nz, ny * nf, 1)
+                complex(2, 2, ny, nf, nz)
+                real(2, 1)
+            else:
+                real(0, 1)
+                complex(1, 1, ny, nf, nz)
+                complex(1, 1, nz, ny * nf, 1)
+            self.tmp_ptr = self.intermediate.data_ptr()
+            if dispatch == 'native':
+                self.native = PreparedChain(self.descriptors, self.tmp_ptr, self.stream_ptr)
+            self.execute_ptr(dummy_in.data_ptr(), dummy_out.data_ptr(), self.stream_ptr)
+            self.stream.synchronize()
+    @property
+    def workspace_bytes(self):
+        return sum(t.numel() * t.element_size() for t in self.tables.values()) + (
+            self.intermediate.numel() * self.intermediate.element_size()
+            if self.intermediate is not None else 0)
+    def execute_ptr(self, input_ptr, output_ptr, stream_ptr):
+        """Enqueue the prepared stages. Caller manages device and buffer lifetime."""
+        if self.closed:
+            raise RuntimeError('The plan is closed')
+        if int(stream_ptr) != self.stream_ptr:
+            raise ValueError('Create a separate plan for the supplied stream')
+        if not input_ptr or not output_ptr or int(input_ptr) == int(output_ptr):
+            raise ValueError('Distinct non-null input/output pointers are required')
+        if self.native is not None:
+            result = self.native.execute(input_ptr, output_ptr)
+            if result:
+                raise RuntimeError(f'cuLaunchKernel failed with CUDA driver error {result}')
+            return
+        pointers = (input_ptr, output_ptr, self.tmp_ptr)
+        for runner, src, dst, table_ptrs, constants in self.launches:
+            runner(pointers[src], pointers[dst], *table_ptrs, *constants, stream=self.stream_ptr)
+    def execute(self, input, output, stream=None):
+        """Validate exact packed tensor buffers, enqueue, and return output."""
+        for tensor, shape in ((input, self.input_shape), (output, self.output_shape)):
+            if (tuple(tensor.shape) != shape or tensor.dtype != self.dtype
+                    or tensor.device != self.device or not tensor.is_contiguous()):
+                raise ValueError('Buffer metadata does not match the plan')
+        a, b = input.data_ptr(), output.data_ptr()
+        na, nb = input.numel() * input.element_size(), output.numel() * output.element_size()
+        if a < b + nb and b < a + na:
+            raise ValueError('Input and output buffers must not overlap')
+        selected = self.stream if stream is None else stream
+        self.execute_ptr(a, b, selected.cuda_stream)
+        return output
+    def close(self):
+        """Release resources only after the caller has waited for all uses."""
+        self.closed = True
+        self.native = None
+        self.descriptors.clear()
+        self.compiled_kernels.clear()
+        self.launches.clear()
+        self.tables.clear()
+        self.intermediate = None
+_CACHE = {}
+_LOCK = threading.Lock()