diff --git a/README.md b/README.md index 6996d2ba72fbcd712c0d961231570a5cf09017d1..0e371e83fe76e0e14af09f160c71dfa46b434085 100644 --- a/README.md +++ b/README.md @@ -10,10 +10,17 @@ tags: - 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 @@ -28,6 +35,8 @@ about 3x.** ## Requirements +For RTX 5090: + - 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 +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", version=1, trust_remote_code=True) +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)) @@ -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)`. + 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. + 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 - 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; + 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`. -- 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 +- 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. -## Accuracy +## 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 @@ -98,7 +111,7 @@ magnitude is unaffected. Repeated calls give bitwise identical results. -## Performance +## RTX 5090 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 +## 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`), @@ -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` diff --git a/build/torch-cuda/LICENSE.txt b/build/torch-cuda/LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..b4ae9ab93cdbd8ba37785bbe993d17d7c5619cc6 --- /dev/null +++ b/build/torch-cuda/LICENSE.txt @@ -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. diff --git a/build/torch-cuda/__init__.py b/build/torch-cuda/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..007ee97c124b7140c0b3d6580e1ca46eea890ac6 --- /dev/null +++ b/build/torch-cuda/__init__.py @@ -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) diff --git a/build/torch-cuda/_b200/__init__.py b/build/torch-cuda/_b200/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2717805d38bb07a4cf728e630242bee43d3827af --- /dev/null +++ b/build/torch-cuda/_b200/__init__.py @@ -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] diff --git a/build/torch-cuda/_b200/c2r_cuda_z_inverse_cuda_kernel.py b/build/torch-cuda/_b200/c2r_cuda_z_inverse_cuda_kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..c02d779bdd56f205ef1fdca5994a080a8549a529 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_cuda_z_inverse_cuda_kernel.py @@ -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) diff --git a/build/torch-cuda/_b200/c2r_cuda_z_inverse_fft_axis.cu b/build/torch-cuda/_b200/c2r_cuda_z_inverse_fft_axis.cu new file mode 100644 index 0000000000000000000000000000000000000000..936a6ae59a54b400448c344fb1a3e4495aa5116d --- /dev/null +++ b/build/torch-cuda/_b200/c2r_cuda_z_inverse_fft_axis.cu @@ -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 +__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(output) & 7ULL) == 0) { + reinterpret_cast(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; + } + } + } +} diff --git a/build/torch-cuda/_b200/c2r_dense_half_inverse_wide_redirect.cu b/build/torch-cuda/_b200/c2r_dense_half_inverse_wide_redirect.cu new file mode 100644 index 0000000000000000000000000000000000000000..7a1c9942f1b00df4d764cdb681e09371a0ee72fd --- /dev/null +++ b/build/torch-cuda/_b200/c2r_dense_half_inverse_wide_redirect.cu @@ -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(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(input)[index]; + value.x = static_cast(packed.x); + value.y = static_cast(packed.y); +#else + value.x = static_cast(input[2ULL * index]); + value.y = static_cast(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(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(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(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(physical_output) + FFT_INPUT_SHIFT; +#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0 + // A nominal cuFFT output address is complex aligned. + reinterpret_cast(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(value.x); + packed.y = static_cast(value.y); + reinterpret_cast(output)[offset] = packed; +#else + output[2 * offset] = static_cast(value.x); + output[2 * offset + 1] = static_cast(value.y); +#endif +} +#endif diff --git a/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_cuda_kernel.py b/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_cuda_kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..24961fe79a5216733d01e23722d7a63533853bf3 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_cuda_kernel.py @@ -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) diff --git a/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_fft_axis.cu b/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_fft_axis.cu new file mode 100644 index 0000000000000000000000000000000000000000..0f50ea8d68406e544b4118ba045d741c543ca2fd --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_fft_axis.cu @@ -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 +__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(output) & 15ULL) == 0) { + reinterpret_cast(output)[offset/2] = a[r]; + } else { + output[offset] = a[r].x; + output[offset+1] = a[r].y; + } + } + } +} diff --git a/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_lto.py b/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..3c689cc5b95be9aaf9d6b3f6b6e8b2718d4055b8 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_lto.py @@ -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)) diff --git a/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_native_dispatch.py b/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_native_dispatch.py new file mode 100644 index 0000000000000000000000000000000000000000..c68fc54a27240a30a3dbc436466f524378cec574 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_native_dispatch.py @@ -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) diff --git a/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_wide_redirect.cu b/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_wide_redirect.cu new file mode 100644 index 0000000000000000000000000000000000000000..a7417bbe788652b3c6e872298c37ac8e19736429 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_cuda_z_inverse_wide_redirect.cu @@ -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(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(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(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(caller_info)[1]; + const double2 w = reinterpret_cast(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(caller_info); + double *output = reinterpret_cast(address); + output[2 * offset] = static_cast(value.x); + output[2 * offset + 1] = static_cast(value.y); +} +#endif diff --git a/build/torch-cuda/_b200/c2r_f64_dense_half_inverse_lto.py b/build/torch-cuda/_b200/c2r_f64_dense_half_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..b97062bba17c77718014e2564f8d5e6542ebf15d --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_dense_half_inverse_lto.py @@ -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)) diff --git a/build/torch-cuda/_b200/c2r_f64_dense_half_inverse_wide_redirect.cu b/build/torch-cuda/_b200/c2r_f64_dense_half_inverse_wide_redirect.cu new file mode 100644 index 0000000000000000000000000000000000000000..8435a5156c4397662dcfb49ee7c903589af7d3a8 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_dense_half_inverse_wide_redirect.cu @@ -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(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(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(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(caller_info)[1]; + const double2 w = reinterpret_cast(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(caller_info); + double *output = reinterpret_cast(address); + output[2 * offset] = static_cast(value.x); + output[2 * offset + 1] = static_cast(value.y); +} +#endif diff --git a/build/torch-cuda/_b200/c2r_f64_half_vector_inverse_wide_redirect.cu b/build/torch-cuda/_b200/c2r_f64_half_vector_inverse_wide_redirect.cu new file mode 100644 index 0000000000000000000000000000000000000000..4a4b47d3bbf738a739da2eeceb5a682e6b4fb907 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_half_vector_inverse_wide_redirect.cu @@ -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(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(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(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(caller_info)[1]; + const double2 w = reinterpret_cast(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(caller_info); + double *output = reinterpret_cast(address); + output[2 * offset] = static_cast(value.x); + output[2 * offset + 1] = static_cast(value.y); +} +#endif diff --git a/build/torch-cuda/_b200/c2r_f64_narrow_index_inverse_scalar_output_reconstruct.cu b/build/torch-cuda/_b200/c2r_f64_narrow_index_inverse_scalar_output_reconstruct.cu new file mode 100644 index 0000000000000000000000000000000000000000..46ff8dc8f7019d2d8adeaf5fcf6352f6fc16da22 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_narrow_index_inverse_scalar_output_reconstruct.cu @@ -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(data)[offset]; + const double2 b = static_cast(data)[mirror]; + const double2 w = static_cast(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(data); + output[2 * offset + FFT_STORE_SHIFT] = value.x; + output[2 * offset + 1 + FFT_STORE_SHIFT] = value.y; +} + +#endif diff --git a/build/torch-cuda/_b200/c2r_f64_paired_x_inverse_cuda_kernel.py b/build/torch-cuda/_b200/c2r_f64_paired_x_inverse_cuda_kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..b62681ad8dde782b4a29a471a19df469f528651a --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_paired_x_inverse_cuda_kernel.py @@ -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) diff --git a/build/torch-cuda/_b200/c2r_f64_paired_x_inverse_fft_axis.cu b/build/torch-cuda/_b200/c2r_f64_paired_x_inverse_fft_axis.cu new file mode 100644 index 0000000000000000000000000000000000000000..e050c13cf4c42f1a86f44812e8783685f0eefd1a --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_paired_x_inverse_fft_axis.cu @@ -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 +__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(input) & 15ULL) == 0) + return reinterpret_cast(input)[offset]; +#else + if ((reinterpret_cast(input) & 7ULL) == 0) { + float2 v = reinterpret_cast(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 +__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(aa, twiddle); + first_transform(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(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 +} diff --git a/build/torch-cuda/_b200/c2r_f64_paired_x_inverse_wide_redirect.cu b/build/torch-cuda/_b200/c2r_f64_paired_x_inverse_wide_redirect.cu new file mode 100644 index 0000000000000000000000000000000000000000..da757bb594f3bd347c1a898cc29961a1c6ca3fa6 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_paired_x_inverse_wide_redirect.cu @@ -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(offset); + unsigned int y = index%FFT_NY, k=(index/FFT_NY)%m, z=index/(FFT_NY*m); + static_cast(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(physical_output) + FFT_INPUT_SHIFT; +#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0 + reinterpret_cast(output)[offset] = value; +#elif !FFT_DOUBLE && FFT_INPUT_SHIFT % 2 == 0 + reinterpret_cast(output)[offset] = make_float2((float)value.x,(float)value.y); +#else + output[2*offset] = static_cast(value.x); + output[2*offset+1] = static_cast(value.y); +#endif +} +#endif diff --git a/build/torch-cuda/_b200/c2r_f64_paired_y_direct_inverse_lto.py b/build/torch-cuda/_b200/c2r_f64_paired_y_direct_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..d65e590244a777f3169de995f644783d33a6a8b5 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_paired_y_direct_inverse_lto.py @@ -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)) diff --git a/build/torch-cuda/_b200/c2r_f64_paired_y_direct_inverse_wide_redirect.cu b/build/torch-cuda/_b200/c2r_f64_paired_y_direct_inverse_wide_redirect.cu new file mode 100644 index 0000000000000000000000000000000000000000..77aeb20ddbf48612aeb813f3f19acee64ee55679 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_paired_y_direct_inverse_wide_redirect.cu @@ -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(input)+k); +#elif !FFT_DOUBLE && FFT_INPUT_SHIFT % 2 == 0 + const float2 v = __ldcs(reinterpret_cast(input)+k); + return make_double2(static_cast(v.x), static_cast(v.y)); +#else + return make_double2(static_cast(__ldcs(input+2ULL*k)), + static_cast(__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(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(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(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(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(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(value.x)); + __stcs(output+even+distance, static_cast(value.y)); +} +#endif diff --git a/build/torch-cuda/_b200/c2r_f64_partial_y_inverse_wide_redirect.cu b/build/torch-cuda/_b200/c2r_f64_partial_y_inverse_wide_redirect.cu new file mode 100644 index 0000000000000000000000000000000000000000..3944bd94c35689abd7078b29f733a7d22234733b --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_partial_y_inverse_wide_redirect.cu @@ -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(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(input)[index]; + value.x = static_cast(packed.x); + value.y = static_cast(packed.y); +#else + value.x = static_cast(input[2ULL * index]); + value.y = static_cast(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(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(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(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(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(data); + const double2 *table = static_cast(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(physical_output) + FFT_INPUT_SHIFT; +#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0 + // A nominal cuFFT output address is complex aligned. + reinterpret_cast(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(value.x); + packed.y = static_cast(value.y); + reinterpret_cast(output)[offset] = packed; +#else + output[2 * offset] = static_cast(value.x); + output[2 * offset + 1] = static_cast(value.y); +#endif +} +#endif diff --git a/build/torch-cuda/_b200/c2r_f64_scalar_input_inverse_scalar_io.cu b/build/torch-cuda/_b200/c2r_f64_scalar_input_inverse_scalar_io.cu new file mode 100644 index 0000000000000000000000000000000000000000..bfcb9d3c8c03cfbb3de11a155d604ca968c02151 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_f64_scalar_input_inverse_scalar_io.cu @@ -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(data); + double2 value; + value.x = static_cast(input[2 * offset + SCALAR_SHIFT]); + value.y = static_cast(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(data)[offset + SCALAR_SHIFT] = + static_cast(value); +} +#endif diff --git a/build/torch-cuda/_b200/c2r_paired_y_direct_inverse_wide_redirect.cu b/build/torch-cuda/_b200/c2r_paired_y_direct_inverse_wide_redirect.cu new file mode 100644 index 0000000000000000000000000000000000000000..4aec976c7edf96107ca713b53da83311507c0878 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_paired_y_direct_inverse_wide_redirect.cu @@ -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(input)[k]; +#elif !FFT_DOUBLE && FFT_INPUT_SHIFT % 2 == 0 + const float2 v = reinterpret_cast(input)[k]; + return make_double2(static_cast(v.x), static_cast(v.y)); +#else + return make_double2(static_cast(input[2ULL*k]), + static_cast(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(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(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(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(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(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(value.x); + output[even+distance] = static_cast(value.y); +} +#endif diff --git a/build/torch-cuda/_b200/c2r_paired_y_inverse_wide_redirect.cu b/build/torch-cuda/_b200/c2r_paired_y_inverse_wide_redirect.cu new file mode 100644 index 0000000000000000000000000000000000000000..5eadf21859d79ab992d84729e63269c705f52fac --- /dev/null +++ b/build/torch-cuda/_b200/c2r_paired_y_inverse_wide_redirect.cu @@ -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(input)[k]; +#elif !FFT_DOUBLE && FFT_INPUT_SHIFT % 2 == 0 + const float2 v = reinterpret_cast(input)[k]; + return make_double2(static_cast(v.x), static_cast(v.y)); +#else + return make_double2(static_cast(input[2ULL*k]), + static_cast(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(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(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(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(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(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(value.x); + output[even+distance] = static_cast(value.y); +} +#endif diff --git a/build/torch-cuda/_b200/c2r_pitched_half_inverse_wide_redirect.cu b/build/torch-cuda/_b200/c2r_pitched_half_inverse_wide_redirect.cu new file mode 100644 index 0000000000000000000000000000000000000000..e3c952ab340184df581aef863bc5a3a5e73ae4c1 --- /dev/null +++ b/build/torch-cuda/_b200/c2r_pitched_half_inverse_wide_redirect.cu @@ -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(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(input)[index]; + value.x = static_cast(packed.x); + value.y = static_cast(packed.y); +#else + value.x = static_cast(input[2ULL * index]); + value.y = static_cast(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(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(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(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(physical_output) + FFT_INPUT_SHIFT; +#if FFT_DOUBLE && FFT_INPUT_SHIFT == 0 + // A nominal cuFFT output address is complex aligned. + reinterpret_cast(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(value.x); + packed.y = static_cast(value.y); + reinterpret_cast(output)[offset] = packed; +#else + output[2 * offset] = static_cast(value.x); + output[2 * offset + 1] = static_cast(value.y); +#endif +} +#endif diff --git a/build/torch-cuda/_b200/cuda_kernel.py b/build/torch-cuda/_b200/cuda_kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..2c00571898305c9d6c22561017a9ffe467676009 --- /dev/null +++ b/build/torch-cuda/_b200/cuda_kernel.py @@ -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) diff --git a/build/torch-cuda/_b200/cuda_twiddle_forward_lto.py b/build/torch-cuda/_b200/cuda_twiddle_forward_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..8d91c3c220ecba23b3809f04a6dffca9c19e08be --- /dev/null +++ b/build/torch-cuda/_b200/cuda_twiddle_forward_lto.py @@ -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)] diff --git a/build/torch-cuda/_b200/cuda_twiddle_forward_native_dispatch.c b/build/torch-cuda/_b200/cuda_twiddle_forward_native_dispatch.c new file mode 100644 index 0000000000000000000000000000000000000000..f7f706bf074029333c2fc3976a8e40d62d8b6410 --- /dev/null +++ b/build/torch-cuda/_b200/cuda_twiddle_forward_native_dispatch.c @@ -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 +#include + +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); } diff --git a/build/torch-cuda/_b200/cuda_twiddle_forward_native_dispatch.py b/build/torch-cuda/_b200/cuda_twiddle_forward_native_dispatch.py new file mode 100644 index 0000000000000000000000000000000000000000..829e34f70887e07e5ebc5d743f7f0e883c51056e --- /dev/null +++ b/build/torch-cuda/_b200/cuda_twiddle_forward_native_dispatch.py @@ -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) diff --git a/build/torch-cuda/_b200/cuda_z_inverse_cufft_support.py b/build/torch-cuda/_b200/cuda_z_inverse_cufft_support.py new file mode 100644 index 0000000000000000000000000000000000000000..15181d38d0f77974e08186ba1a27aa57cd3b9ab0 --- /dev/null +++ b/build/torch-cuda/_b200/cuda_z_inverse_cufft_support.py @@ -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 diff --git a/build/torch-cuda/_b200/cuda_z_inverse_lto.py b/build/torch-cuda/_b200/cuda_z_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..6def06b7cec7e7a3cc9b239e6644b098e9644a11 --- /dev/null +++ b/build/torch-cuda/_b200/cuda_z_inverse_lto.py @@ -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)) diff --git a/build/torch-cuda/_b200/cuda_z_inverse_native_dispatch.c b/build/torch-cuda/_b200/cuda_z_inverse_native_dispatch.c new file mode 100644 index 0000000000000000000000000000000000000000..1ceef7b00c65c09b8c244e20acbac40c391a77ca --- /dev/null +++ b/build/torch-cuda/_b200/cuda_z_inverse_native_dispatch.c @@ -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 +#include + +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); } diff --git a/build/torch-cuda/_b200/cuda_z_inverse_native_dispatch.py b/build/torch-cuda/_b200/cuda_z_inverse_native_dispatch.py new file mode 100644 index 0000000000000000000000000000000000000000..04bc40b4bc94c4fd5aca0536129ad60f67c04089 --- /dev/null +++ b/build/torch-cuda/_b200/cuda_z_inverse_native_dispatch.py @@ -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) diff --git a/build/torch-cuda/_b200/cuda_z_inverse_schedule.py b/build/torch-cuda/_b200/cuda_z_inverse_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..9fb2fbf67043d2f8851ed8917483661d1008dfb4 --- /dev/null +++ b/build/torch-cuda/_b200/cuda_z_inverse_schedule.py @@ -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 diff --git a/build/torch-cuda/_b200/dense_half_inverse_abi.py b/build/torch-cuda/_b200/dense_half_inverse_abi.py new file mode 100644 index 0000000000000000000000000000000000000000..6003def960b118f5486a3f2144564c5c43a406e1 --- /dev/null +++ b/build/torch-cuda/_b200/dense_half_inverse_abi.py @@ -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) diff --git a/build/torch-cuda/_b200/dense_half_inverse_kernels.py b/build/torch-cuda/_b200/dense_half_inverse_kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..507db2edb34e8223300bd56310c2d16bbccd3815 --- /dev/null +++ b/build/torch-cuda/_b200/dense_half_inverse_kernels.py @@ -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)) diff --git a/build/torch-cuda/_b200/dense_half_inverse_lto.py b/build/torch-cuda/_b200/dense_half_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..03001393f8b60f755629e5ff548d8a00ef5b87a2 --- /dev/null +++ b/build/torch-cuda/_b200/dense_half_inverse_lto.py @@ -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)) diff --git a/build/torch-cuda/_b200/dense_half_inverse_native_dispatch.c b/build/torch-cuda/_b200/dense_half_inverse_native_dispatch.c new file mode 100644 index 0000000000000000000000000000000000000000..493d0086ecb87b56b16a0d916992dc7b604e9b27 --- /dev/null +++ b/build/torch-cuda/_b200/dense_half_inverse_native_dispatch.c @@ -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 +#include + +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); } diff --git a/build/torch-cuda/_b200/dense_half_inverse_native_dispatch.py b/build/torch-cuda/_b200/dense_half_inverse_native_dispatch.py new file mode 100644 index 0000000000000000000000000000000000000000..08fc3475894965bcfce0182220e0995c63357af9 --- /dev/null +++ b/build/torch-cuda/_b200/dense_half_inverse_native_dispatch.py @@ -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) diff --git a/build/torch-cuda/_b200/dense_half_inverse_schedule.py b/build/torch-cuda/_b200/dense_half_inverse_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..40d54655123afcb885078acf59dede8fbca21808 --- /dev/null +++ b/build/torch-cuda/_b200/dense_half_inverse_schedule.py @@ -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 diff --git a/build/torch-cuda/_b200/direct_dispatch.c b/build/torch-cuda/_b200/direct_dispatch.c new file mode 100644 index 0000000000000000000000000000000000000000..5c0c9acc2011132a27b22bf7edc71b76c905c59b --- /dev/null +++ b/build/torch-cuda/_b200/direct_dispatch.c @@ -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 + +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); } diff --git a/build/torch-cuda/_b200/direct_dispatch.py b/build/torch-cuda/_b200/direct_dispatch.py new file mode 100644 index 0000000000000000000000000000000000000000..6fd0baf8f413fb5885cc15554fe7e9897c2344e9 --- /dev/null +++ b/build/torch-cuda/_b200/direct_dispatch.py @@ -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) diff --git a/build/torch-cuda/_b200/fft_axis.cu b/build/torch-cuda/_b200/fft_axis.cu new file mode 100644 index 0000000000000000000000000000000000000000..3aa0dc790c5c4438770e29fc0e49d5b29006175b --- /dev/null +++ b/build/torch-cuda/_b200/fft_axis.cu @@ -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 +__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 +__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(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(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(output) & 15ULL) == 0) { + reinterpret_cast(output)[offset/2] = a[r]; + } else { + output[offset] = a[r].x; + output[offset+1] = a[r].y; + } +#else + if ((reinterpret_cast(output) & 7ULL) == 0) { + reinterpret_cast(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 + } + } +} diff --git a/build/torch-cuda/_b200/forward_cuda_f32.py b/build/torch-cuda/_b200/forward_cuda_f32.py new file mode 100644 index 0000000000000000000000000000000000000000..84a631fb7c0b0f26531ab147de8ebd33e7416bf6 --- /dev/null +++ b/build/torch-cuda/_b200/forward_cuda_f32.py @@ -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() diff --git a/build/torch-cuda/_b200/forward_half.py b/build/torch-cuda/_b200/forward_half.py new file mode 100644 index 0000000000000000000000000000000000000000..0502b8c1a9e6e89ca0cfd5a0a49bc0035fa8f59d --- /dev/null +++ b/build/torch-cuda/_b200/forward_half.py @@ -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() diff --git a/build/torch-cuda/_b200/forward_native.py b/build/torch-cuda/_b200/forward_native.py new file mode 100644 index 0000000000000000000000000000000000000000..2e5eab829112f01b21e136c965721673ae6a15d5 --- /dev/null +++ b/build/torch-cuda/_b200/forward_native.py @@ -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() diff --git a/build/torch-cuda/_b200/forward_paired_y.py b/build/torch-cuda/_b200/forward_paired_y.py new file mode 100644 index 0000000000000000000000000000000000000000..149c1ae118b36c130c42f21fbd8be97bd7b1b566 --- /dev/null +++ b/build/torch-cuda/_b200/forward_paired_y.py @@ -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() diff --git a/build/torch-cuda/_b200/forward_scalar.py b/build/torch-cuda/_b200/forward_scalar.py new file mode 100644 index 0000000000000000000000000000000000000000..54df0ac3695819caac36bc02773cf4c48dec2e91 --- /dev/null +++ b/build/torch-cuda/_b200/forward_scalar.py @@ -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 diff --git a/build/torch-cuda/_b200/forward_small.py b/build/torch-cuda/_b200/forward_small.py new file mode 100644 index 0000000000000000000000000000000000000000..ba68b6baf6a05d4eade55fcd9231c9f1b2ad2394 --- /dev/null +++ b/build/torch-cuda/_b200/forward_small.py @@ -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() diff --git a/build/torch-cuda/_b200/grouped_dif_inverse_abi.py b/build/torch-cuda/_b200/grouped_dif_inverse_abi.py new file mode 100644 index 0000000000000000000000000000000000000000..9775251a2d92711f831db7dfe08311d31270a32f --- /dev/null +++ b/build/torch-cuda/_b200/grouped_dif_inverse_abi.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reject kernel ABIs that the small native launcher cannot represent.""" +from __future__ import annotations + +import re + + +def validate_kernel_abi(compiled): + """Inspect compiler output on the host; this function does not load CUDA.""" + 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 not in (1, 2, 4, 8, 16, 32): + raise ValueError('Unsupported warp count') + signature = compiled.src.signature + runtime = [value for value in signature.values() if value != 'constexpr'] + if (len(runtime) not in (3, 4) + or any(value not in ('*fp32', '*fp64') for value in runtime[:2]) + or '*fp64' not in runtime[:2] + or any(value != '*fp64' for value in runtime[2:])): + raise ValueError('Expected scalar float boundaries, an FP64 intermediate, and FP64 twiddles') + expected_count = len(runtime) + 2 + ptx = compiled.asm['ptx'] + match = re.search(r'\.entry\s+' + re.escape(metadata.name) + r'\s*\((.*?)\)', + ptx, flags=re.DOTALL) + if match is None: + raise ValueError('Kernel entry signature was not found in compiler PTX') + parameters = [part.strip() for part in match.group(1).split(',')] + # Current installed Triton appends global and profiling scratch pointers. + # Require this precise ABI instead of guessing across compiler versions. + if len(parameters) != expected_count 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 tensor pointers followed by two scalar 64-bit scratch pointers') + return dict(parameter_count=expected_count, tensor_arguments=len(runtime), null_scratch_arguments=2, + num_ctas=1, threads=32 * metadata.num_warps, + shared_bytes=metadata.shared) diff --git a/build/torch-cuda/_b200/grouped_dif_inverse_kernels.py b/build/torch-cuda/_b200/grouped_dif_inverse_kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..1ed9a9da329826bfb0738212b0a4f9e33055ce09 --- /dev/null +++ b/build/torch-cuda/_b200/grouped_dif_inverse_kernels.py @@ -0,0 +1,784 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Three-pass FFT with fused radix-9 and paired radix-2/3/4/5/7 arithmetic. + +Power-of-two padding exists only in register tensors required by Triton. +There is no padded device array, transpose pass, or input repacking pass. +Twiddle tables contain transform constants, never cached input/output data. +""" +import triton +import triton.language as tl + + +@triton.jit +def _input_index(index, N: tl.constexpr, RADICES: tl.constexpr, SIZES: tl.constexpr): + reverse = tl.full(index.shape, 0, tl.int32) + for stage in tl.static_range(len(RADICES)): + reverse += ((index // (SIZES[stage] // RADICES[stage])) % RADICES[stage]) * (N // SIZES[stage]) + return reverse + + +@triton.jit +def _fft_generic_stage(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADIX: tl.constexpr, SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Keep each stage's constants in its own JIT function scope. + + Annotated constexpr assignments inside a static loop are rejected on + the second unrolled iteration by Triton's CUDA frontend. Passing the + radix and size as constexpr arguments also keeps inner loops static. + """ + index = tl.arange(0, P) + PREVIOUS: tl.constexpr = SIZE // RADIX + base = index // SIZE * SIZE + index % PREVIOUS + # Padding lanes are never stored. Clamp their local gather indices. + source = tl.minimum(base, P - 1) + if STRIDED: + gather = tl.broadcast_to(source[:, None], (P, B)) + out_real = tl.gather(real, gather, 0) + out_imag = tl.gather(imag, gather, 0) + else: + gather = tl.broadcast_to(source[None, :], (B, P)) + out_real = tl.gather(real, gather, 1) + out_imag = tl.gather(imag, gather, 1) + for term in tl.static_range(1, RADIX): + source = tl.minimum(base + term * PREVIOUS, P - 1) + exponent = (term * (index % SIZE) * (N // SIZE)) % N + if N * TWIDDLE_STEP == 32: + wr, wi = _constant32_twiddle(exponent * TWIDDLE_STEP) + else: + wr = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP) + wi = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP + 1) + if INVERSE: + wi = -wi + if STRIDED: + gather = tl.broadcast_to(source[:, None], (P, B)) + vr = tl.gather(real, gather, 0) + vi = tl.gather(imag, gather, 0) + wr, wi = wr[:, None], wi[:, None] + else: + gather = tl.broadcast_to(source[None, :], (B, P)) + vr = tl.gather(real, gather, 1) + vi = tl.gather(imag, gather, 1) + wr, wi = wr[None, :], wi[None, :] + out_real += vr * wr - vi * wi + out_imag += vr * wi + vi * wr + return out_real, out_imag + + +@triton.jit +def _constant32_twiddle(exponent): + """Reconstruct accurately rounded 32nd roots using fixed constants.""" + k = exponent % 32 + c = tl.minimum(k % 16, 16 - k % 16) + q = tl.where(c <= 8, c, 16-c) + real = tl.full(q.shape, 0., tl.float64) + real = tl.where(q == 0, tl.full((), 1.0, tl.float64), real) + real = tl.where(q == 1, tl.full((), 0.9807852804032304, tl.float64), real) + real = tl.where(q == 2, tl.full((), 0.9238795325112867, tl.float64), real) + real = tl.where(q == 3, tl.full((), 0.8314696123025452, tl.float64), real) + real = tl.where(q == 4, tl.full((), 0.7071067811865476, tl.float64), real) + real = tl.where(q == 5, tl.full((), 0.5555702330196023, tl.float64), real) + real = tl.where(q == 6, tl.full((), 0.38268343236508984, tl.float64), real) + real = tl.where(q == 7, tl.full((), 0.19509032201612833, tl.float64), real) + real = tl.where(q == 8, tl.full((), 0.0, tl.float64), real) + real = tl.where((k > 8) & (k < 24), -real, real) + t = (k + 24) % 32 + q = tl.minimum(t % 16, 16 - t % 16) + imag = tl.full(q.shape, 0., tl.float64) + imag = tl.where(q == 0, tl.full((), 1.0, tl.float64), imag) + imag = tl.where(q == 1, tl.full((), 0.9807852804032304, tl.float64), imag) + imag = tl.where(q == 2, tl.full((), 0.9238795325112867, tl.float64), imag) + imag = tl.where(q == 3, tl.full((), 0.8314696123025452, tl.float64), imag) + imag = tl.where(q == 4, tl.full((), 0.7071067811865476, tl.float64), imag) + imag = tl.where(q == 5, tl.full((), 0.5555702330196023, tl.float64), imag) + imag = tl.where(q == 6, tl.full((), 0.38268343236508984, tl.float64), imag) + imag = tl.where(q == 7, tl.full((), 0.19509032201612833, tl.float64), imag) + imag = tl.where(q == 8, tl.full((), 0.0, tl.float64), imag) + imag = tl.where((t > 8) & (t < 24), imag, -imag) + return real, imag + + +@triton.jit +def _stage_twiddle(real, imag, Twiddle, N: tl.constexpr, + SIZE: tl.constexpr, RADIX: tl.constexpr, + TERM: tl.constexpr, STRIDED: tl.constexpr, + INVERSE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Apply one branch's twiddle; broadcast only across independent FFTs.""" + if SIZE == RADIX: + # All first-stage twiddles are one. No table access or multiplication. + return real, imag + else: + index = tl.arange(0, SIZE // RADIX) + exponent = TERM * index * (N // SIZE) + if N * TWIDDLE_STEP == 32: + wr, wi = _constant32_twiddle(exponent * TWIDDLE_STEP) + else: + wr = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP) + wi = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP + 1) + if INVERSE: + wi = -wi + if STRIDED: + wr, wi = wr[None, :, None], wi[None, :, None] + else: + wr, wi = wr[None, None, :], wi[None, None, :] + return real * wr - imag * wi, real * wi + imag * wr + + +@triton.jit +def _fft_radix2(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + H: tl.constexpr = SIZE // 2 + G: tl.constexpr = P // SIZE + if STRIDED: + r = tl.permute(tl.reshape(real, (G, 2, H, B)), (0, 2, 3, 1)) + i = tl.permute(tl.reshape(imag, (G, 2, H, B)), (0, 2, 3, 1)) + else: + r = tl.permute(tl.reshape(real, (B, G, 2, H)), (0, 1, 3, 2)) + i = tl.permute(tl.reshape(imag, (B, G, 2, H)), (0, 1, 3, 2)) + ar, br = tl.split(r) + ai, bi = tl.split(i) + br, bi = _stage_twiddle(br, bi, Twiddle, N, SIZE, 2, 1, + STRIDED, INVERSE, TWIDDLE_STEP) + r = tl.join(ar + br, ar - br) + i = tl.join(ai + bi, ai - bi) + if STRIDED: + return (tl.reshape(tl.permute(r, (0, 3, 1, 2)), (P, B)), + tl.reshape(tl.permute(i, (0, 3, 1, 2)), (P, B))) + else: + return (tl.reshape(tl.permute(r, (0, 1, 3, 2)), (B, P)), + tl.reshape(tl.permute(i, (0, 1, 3, 2)), (B, P))) + + +@triton.jit +def _fft_radix4(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + H: tl.constexpr = SIZE // 4 + G: tl.constexpr = P // SIZE + if STRIDED: + r = tl.permute(tl.reshape(real, (G, 4, H, B)), (0, 2, 3, 1)) + i = tl.permute(tl.reshape(imag, (G, 4, H, B)), (0, 2, 3, 1)) + r = tl.reshape(r, (G, H, B, 2, 2)) + i = tl.reshape(i, (G, H, B, 2, 2)) + else: + r = tl.permute(tl.reshape(real, (B, G, 4, H)), (0, 1, 3, 2)) + i = tl.permute(tl.reshape(imag, (B, G, 4, H)), (0, 1, 3, 2)) + r = tl.reshape(r, (B, G, H, 2, 2)) + i = tl.reshape(i, (B, G, H, 2, 2)) + even_r, odd_r = tl.split(r) + even_i, odd_i = tl.split(i) + ar, cr = tl.split(even_r) + ai, ci = tl.split(even_i) + br, dr = tl.split(odd_r) + bi, di = tl.split(odd_i) + br, bi = _stage_twiddle(br, bi, Twiddle, N, SIZE, 4, 1, + STRIDED, INVERSE, TWIDDLE_STEP) + cr, ci = _stage_twiddle(cr, ci, Twiddle, N, SIZE, 4, 2, + STRIDED, INVERSE, TWIDDLE_STEP) + dr, di = _stage_twiddle(dr, di, Twiddle, N, SIZE, 4, 3, + STRIDED, INVERSE, TWIDDLE_STEP) + ac_r, ac_i = ar + cr, ai + ci + bd_r, bd_i = br + dr, bi + di + diff_ac_r, diff_ac_i = ar - cr, ai - ci + diff_bd_r, diff_bd_i = br - dr, bi - di + # Multiplication by +/- i is a swap and sign, exactly in either precision. + if INVERSE: + turn_r, turn_i = -diff_bd_i, diff_bd_r + else: + turn_r, turn_i = diff_bd_i, -diff_bd_r + r0, i0 = ac_r + bd_r, ac_i + bd_i + r2, i2 = ac_r - bd_r, ac_i - bd_i + r1, i1 = diff_ac_r + turn_r, diff_ac_i + turn_i + r3, i3 = diff_ac_r - turn_r, diff_ac_i - turn_i + r = tl.join(tl.join(r0, r2), tl.join(r1, r3)) + i = tl.join(tl.join(i0, i2), tl.join(i1, i3)) + if STRIDED: + r = tl.reshape(r, (G, H, B, 4)) + i = tl.reshape(i, (G, H, B, 4)) + return (tl.reshape(tl.permute(r, (0, 3, 1, 2)), (P, B)), + tl.reshape(tl.permute(i, (0, 3, 1, 2)), (P, B))) + else: + r = tl.reshape(r, (B, G, H, 4)) + i = tl.reshape(i, (B, G, H, 4)) + return (tl.reshape(tl.permute(r, (0, 1, 3, 2)), (B, P)), + tl.reshape(tl.permute(i, (0, 1, 3, 2)), (B, P))) + + +@triton.jit +def _odd_branch(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADIX: tl.constexpr, SIZE: tl.constexpr, TERM: tl.constexpr, + HP: tl.constexpr, GP: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Read one branch into a smaller register tensor and apply its twiddle.""" + H: tl.constexpr = SIZE // RADIX + RP: tl.constexpr = triton.next_power_of_2(RADIX) + index = tl.arange(0, GP * HP) + local = index % HP + source = tl.minimum(index // HP * SIZE + local + TERM * H, P - 1) + if GP == 1 and HP == H and RP * HP == P: + # A final odd stage with power-of-two branch length is already a + # padded radix tile. Split its minor radix bits instead of gathering. + if STRIDED: + r = tl.permute(tl.reshape(real, (RP, HP, B)), (1, 2, 0)) + i = tl.permute(tl.reshape(imag, (RP, HP, B)), (1, 2, 0)) + if RP == 4: + r, i = tl.reshape(r, (HP, B, 2, 2)), tl.reshape(i, (HP, B, 2, 2)) + else: + r = tl.reshape(r, (HP, B, 2, 2, 2)) + i = tl.reshape(i, (HP, B, 2, 2, 2)) + else: + r = tl.permute(tl.reshape(real, (B, RP, HP)), (0, 2, 1)) + i = tl.permute(tl.reshape(imag, (B, RP, HP)), (0, 2, 1)) + if RP == 4: + r, i = tl.reshape(r, (B, HP, 2, 2)), tl.reshape(i, (B, HP, 2, 2)) + else: + r = tl.reshape(r, (B, HP, 2, 2, 2)) + i = tl.reshape(i, (B, HP, 2, 2, 2)) + lo_r, hi_r = tl.split(r) + lo_i, hi_i = tl.split(i) + if TERM % 2 == 0: + r, i = lo_r, lo_i + else: + r, i = hi_r, hi_i + lo_r, hi_r = tl.split(r) + lo_i, hi_i = tl.split(i) + if (TERM // 2) % 2 == 0: + r, i = lo_r, lo_i + else: + r, i = hi_r, hi_i + if RP == 8: + lo_r, hi_r = tl.split(r) + lo_i, hi_i = tl.split(i) + if TERM // 4 == 0: + r, i = lo_r, lo_i + else: + r, i = hi_r, hi_i + elif STRIDED: + gather = tl.broadcast_to(source[:, None], (GP * HP, B)) + r = tl.gather(real, gather, 0) + i = tl.gather(imag, gather, 0) + else: + gather = tl.broadcast_to(source[None, :], (B, GP * HP)) + r = tl.gather(real, gather, 1) + i = tl.gather(imag, gather, 1) + if TERM != 0 and SIZE != RADIX: + exponent = TERM * local * (N // SIZE) + valid = (local < H) & (index // HP < N // SIZE) + wr = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP, valid, 0) + wi = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP + 1, valid, 0) + if INVERSE: + wi = -wi + if STRIDED: + wr, wi = wr[:, None], wi[:, None] + else: + wr, wi = wr[None, :], wi[None, :] + return r * wr - i * wi, r * wi + i * wr + else: + return r, i + + +@triton.jit +def _join4(a, b, c, d): + return tl.reshape(tl.join(tl.join(a, c), tl.join(b, d)), a.shape + (4,)) + + +@triton.jit +def _join8(a, b, c, d, e, f, g, h): + return tl.reshape(tl.join(_join4(a, c, e, g), _join4(b, d, f, h)), + a.shape + (8,)) + + +@triton.jit +def _odd_output(outputs, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, RADIX: tl.constexpr, + SIZE: tl.constexpr, HP: tl.constexpr, GP: tl.constexpr): + """Join the small branches and gather natural ordering only in registers.""" + RP: tl.constexpr = triton.next_power_of_2(RADIX) + H: tl.constexpr = SIZE // RADIX + zero = tl.full(outputs[0].shape, 0, outputs[0].dtype) + # outputs = [DC, +1, -1, +2, -2, ...]. Natural branch order follows. + if RADIX == 3: + data = _join4(outputs[0], outputs[1], outputs[2], zero) + elif RADIX == 5: + data = _join8(outputs[0], outputs[1], outputs[3], outputs[4], outputs[2], + zero, zero, zero) + else: + data = _join8(outputs[0], outputs[1], outputs[3], outputs[5], + outputs[6], outputs[4], outputs[2], zero) + if STRIDED: + data = tl.reshape(data, (GP, HP, B, RP)) + data = tl.reshape(tl.permute(data, (0, 3, 1, 2)), (GP * RP * HP, B)) + else: + data = tl.reshape(data, (B, GP, HP, RP)) + data = tl.reshape(tl.permute(data, (0, 1, 3, 2)), (B, GP * RP * HP)) + if GP == 1 and HP == H and RP * HP == P: + return data + index = tl.arange(0, P) + source = (index // SIZE * RP * HP + (index % SIZE) // H * HP + index % H) + source = tl.minimum(source, GP * RP * HP - 1) + if STRIDED: + return tl.gather(data, tl.broadcast_to(source[:, None], (P, B)), 0) + else: + return tl.gather(data, tl.broadcast_to(source[None, :], (B, P)), 1) + + +@triton.jit +def _fft_odd(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADIX: tl.constexpr, SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Paired DFT on compact branch tensors, retaining the input precision. + + v[j] + v[r-j] contributes cosine terms to a conjugate output pair; + v[j] - v[r-j] contributes sine terms with opposite signs. This reduces + duplicated arithmetic across output branches. Padding is register-only. + """ + HP: tl.constexpr = triton.next_power_of_2(SIZE // RADIX) + GP: tl.constexpr = triton.next_power_of_2(N // SIZE) + ar, ai = _odd_branch(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADIX, SIZE, 0, HP, GP, TWIDDLE_STEP) + sums_r, sums_i, diffs_r, diffs_i = (), (), (), () + dc_r, dc_i = ar, ai + for j in tl.static_range(1, (RADIX + 1) // 2): + br, bi = _odd_branch(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADIX, SIZE, j, HP, GP, TWIDDLE_STEP) + cr, ci = _odd_branch(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADIX, SIZE, RADIX - j, HP, GP, TWIDDLE_STEP) + sr, si, dr, di = br + cr, bi + ci, br - cr, bi - ci + sums_r, sums_i = sums_r + (sr,), sums_i + (si,) + diffs_r, diffs_i = diffs_r + (dr,), diffs_i + (di,) + dc_r, dc_i = dc_r + sr, dc_i + si + outputs_r, outputs_i = (dc_r,), (dc_i,) + for k in tl.static_range(1, (RADIX + 1) // 2): + even_r, even_i = ar, ai + odd_r = tl.full(ar.shape, 0, ar.dtype) + odd_i = tl.full(ai.shape, 0, ai.dtype) + for j in tl.static_range(1, (RADIX + 1) // 2): + exponent = (j * k % RADIX) * (N // RADIX) + c = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP) + s = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP + 1) + if INVERSE: + s = -s + even_r = even_r + c * sums_r[j - 1] + even_i = even_i + c * sums_i[j - 1] + odd_r = odd_r + s * diffs_r[j - 1] + odd_i = odd_i + s * diffs_i[j - 1] + outputs_r = outputs_r + (even_r - odd_i, even_r + odd_i) + outputs_i = outputs_i + (even_i + odd_r, even_i - odd_r) + return (_odd_output(outputs_r, N, P, B, STRIDED, RADIX, SIZE, HP, GP), + _odd_output(outputs_i, N, P, B, STRIDED, RADIX, SIZE, HP, GP)) + + +@triton.jit +def _radix9_branch(data, H: tl.constexpr, B: tl.constexpr, + STRIDED: tl.constexpr, TERM: tl.constexpr): + """Select one of sixteen padded register branches using four splits.""" + if STRIDED: + value = tl.permute(tl.reshape(data, (16, H, B)), (1, 2, 0)) + value = tl.reshape(value, (H, B, 2, 2, 2, 2)) + else: + value = tl.permute(tl.reshape(data, (B, 16, H)), (0, 2, 1)) + value = tl.reshape(value, (B, H, 2, 2, 2, 2)) + for bit in tl.static_range(4): + lo, hi = tl.split(value) + if (TERM >> bit) & 1: + value = hi + else: + value = lo + return value + + +@triton.jit +def _dft3(ar, ai, br, bi, cr, ci, sine): + """Three-point DFT with an exactly representable -1/2 cosine.""" + sr, si = br + cr, bi + ci + dr, di = br - cr, bi - ci + er, ei = ar - 0.5 * sr, ai - 0.5 * si + return (ar + sr, ai + si, + er - sine * di, ei + sine * dr, + er + sine * di, ei - sine * dr) + + +@triton.jit +def _join16(outputs): + zero = tl.full(outputs[0].shape, 0, outputs[0].dtype) + evens = _join8(outputs[0], outputs[2], outputs[4], outputs[6], outputs[8], + zero, zero, zero) + odds = _join8(outputs[1], outputs[3], outputs[5], outputs[7], + zero, zero, zero, zero) + return tl.reshape(tl.join(evens, odds), outputs[0].shape + (16,)) + + +@triton.jit +def _fft_radix9(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Fuse two final radix-3 stages without an intermediate register gather. + + After the stage twiddles, compute three DFTs over indices (b, b+3, + b+6), multiply branch b and frequency k1 by W9**(b*k1), then compute + three DFTs over b. Natural output frequency is k1 + 3*k2. + Padding is confined to register tensors; all global rows stay packed. + """ + H: tl.constexpr = SIZE // 9 + tl.static_assert(SIZE == N and H == triton.next_power_of_2(H) and P == 16 * H) + local = tl.arange(0, H) + values_r, values_i = (), () + for j in tl.static_range(9): + r = _radix9_branch(real, H, B, STRIDED, j) + i = _radix9_branch(imag, H, B, STRIDED, j) + if j != 0 and H != 1: + wr = tl.load(Twiddle + 2 * j * local * TWIDDLE_STEP) + wi = tl.load(Twiddle + 2 * j * local * TWIDDLE_STEP + 1) + if INVERSE: + wi = -wi + if STRIDED: + wr, wi = wr[:, None], wi[:, None] + else: + wr, wi = wr[None, :], wi[None, :] + r, i = r * wr - i * wi, r * wi + i * wr + values_r, values_i = values_r + (r,), values_i + (i,) + sine = tl.load(Twiddle + 2 * (N // 3) * TWIDDLE_STEP + 1) + if INVERSE: + sine = -sine + columns = () + for b in tl.static_range(3): + columns = columns + (_dft3(values_r[b], values_i[b], + values_r[b + 3], values_i[b + 3], + values_r[b + 6], values_i[b + 6], sine),) + rows = () + for k1 in tl.static_range(3): + branches_r, branches_i = (), () + for b in tl.static_range(3): + r, i = columns[b][2 * k1], columns[b][2 * k1 + 1] + if b != 0 and k1 != 0: + wr = tl.load(Twiddle + 2 * (b * k1 * H) * TWIDDLE_STEP) + wi = tl.load(Twiddle + 2 * (b * k1 * H) * TWIDDLE_STEP + 1) + if INVERSE: + wi = -wi + r, i = r * wr - i * wi, r * wi + i * wr + branches_r, branches_i = branches_r + (r,), branches_i + (i,) + rows = rows + (_dft3(branches_r[0], branches_i[0], + branches_r[1], branches_i[1], + branches_r[2], branches_i[2], sine),) + outputs_r, outputs_i = (), () + for k2 in tl.static_range(3): + for k1 in tl.static_range(3): + outputs_r = outputs_r + (rows[k1][2 * k2],) + outputs_i = outputs_i + (rows[k1][2 * k2 + 1],) + r, i = _join16(outputs_r), _join16(outputs_i) + if STRIDED: + return (tl.reshape(tl.permute(r, (2, 0, 1)), (P, B)), + tl.reshape(tl.permute(i, (2, 0, 1)), (P, B))) + else: + return (tl.reshape(tl.permute(r, (0, 2, 1)), (B, P)), + tl.reshape(tl.permute(i, (0, 2, 1)), (B, P))) + + +@triton.jit +def _fft_stage(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADIX: tl.constexpr, SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + # Factors 4 and 2 precede odd factors, so these reshapes divide P exactly. + if RADIX == 4 and P % SIZE == 0: + return _fft_radix4(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + SIZE, TWIDDLE_STEP) + elif RADIX == 2 and P % SIZE == 0: + return _fft_radix2(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + SIZE, TWIDDLE_STEP) + elif RADIX == 9: + return _fft_radix9(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + SIZE, TWIDDLE_STEP) + elif RADIX == 3 or RADIX == 5 or RADIX == 7: + return _fft_odd(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADIX, SIZE, TWIDDLE_STEP) + else: + return _fft_generic_stage(real, imag, Twiddle, N, P, B, STRIDED, + INVERSE, RADIX, SIZE, TWIDDLE_STEP) + + +@triton.jit +def _fft(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr, TWIDDLE_STEP: tl.constexpr = 1): + for stage in tl.static_range(len(RADICES)): + real, imag = _fft_stage(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADICES[stage], SIZES[stage], TWIDDLE_STEP) + return real, imag + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def real_axis(Input, Output, Twiddle, ROWS: tl.constexpr, N: tl.constexpr, + P: tl.constexpr, B: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr): + rows = tl.program_id(0) * B + tl.arange(0, B) + index = tl.arange(0, P) + source = _input_index(index, N, RADICES, SIZES) + NF: tl.constexpr = N // 2 + 1 + mask = (rows[:, None] < ROWS) & (index[None, :] < N) + if INVERSE: + half_index = tl.minimum(source, N - source) + offsets = 2 * (rows[:, None] * NF + half_index[None, :]) + real = tl.load(Input + offsets, mask, 0) + imag = tl.load(Input + offsets + 1, mask, 0) + imag = tl.where(source[None, :] >= NF, -imag, imag) + # Real end points after inverse Z/Y implement irfftn's convention + # even when the original packed input is not Hermitian-consistent. + imag = tl.where((source[None, :] == 0) | (2 * source[None, :] == N), 0, imag) + else: + real = tl.load(Input + rows[:, None] * N + source[None, :], mask, 0) + imag = tl.full((B, P), 0, real.dtype) + real, imag = _fft(real, imag, Twiddle, N, P, B, False, INVERSE, RADICES, SIZES) + if INVERSE: + tl.store(Output + rows[:, None] * N + index[None, :], real, mask) + else: + offsets = 2 * (rows[:, None] * NF + index[None, :]) + output_mask = (rows[:, None] < ROWS) & (index[None, :] < NF) + tl.store(Output + offsets, real, output_mask) + tl.store(Output + offsets + 1, imag, output_mask) + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def packed_real_axis(Input, Output, Twiddle, ROWS: tl.constexpr, N: tl.constexpr, + P: tl.constexpr, B: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr): + """Even-length real FFT using an N/2 complex FFT and fused reconstruction. + + RADICES, SIZES and P describe M=N/2; Twiddle contains N complex constants. + Even/odd samples are paired only in registers; global layout stays exact. + The inverse constructs twice the packed complex spectrum so that the + length-M unnormalized inverse returns the required length-N scaling. + """ + tl.static_assert(N % 2 == 0) + M: tl.constexpr = N // 2 + NF: tl.constexpr = M + 1 + rows = tl.program_id(0) * B + tl.arange(0, B) + index = tl.arange(0, P) + source = _input_index(index, M, RADICES, SIZES) + mask = (rows[:, None] < ROWS) & (index[None, :] < M) + if INVERSE: + offsets = 2 * (rows[:, None] * NF + source[None, :]) + mirrors = 2 * (rows[:, None] * NF + M - source[None, :]) + ar = tl.load(Input + offsets, mask, 0) + ai = tl.load(Input + offsets + 1, mask, 0) + br = tl.load(Input + mirrors, mask, 0) + bi = tl.load(Input + mirrors + 1, mask, 0) + # irfftn ignores real-axis endpoint imaginary parts AFTER Z/Y inverse. + # Arbitrary complex packed input is valid; do not zero planes earlier. + ai = tl.where(source[None, :] == 0, 0, ai) + bi = tl.where(source[None, :] == 0, 0, bi) + wr = tl.load(Twiddle + 2 * source, index < M, 0)[None, :] + wi = tl.load(Twiddle + 2 * source + 1, index < M, 0)[None, :] + # 2 Z[k] = X[k] + conj(X[M-k]) + # + i conj(W_N^k) (X[k] - conj(X[M-k])). + real = ar + br - wr * (ai + bi) + wi * (ar - br) + imag = ai - bi + wr * (ar - br) + wi * (ai + bi) + else: + offsets = rows[:, None] * N + 2 * source[None, :] + real = tl.load(Input + offsets, mask, 0) + imag = tl.load(Input + offsets + 1, mask, 0) + real, imag = _fft(real, imag, Twiddle, M, P, B, False, INVERSE, + RADICES, SIZES, 2) + if INVERSE: + offsets = rows[:, None] * N + 2 * index[None, :] + tl.store(Output + offsets, real, mask) + tl.store(Output + offsets + 1, imag, mask) + else: + mirror_index = tl.where(index < M, (M - index) % M, 0) + mirror = tl.broadcast_to(mirror_index[None, :], (B, P)) + br = tl.gather(real, mirror, 1) + bi = tl.gather(imag, mirror, 1) + wr = tl.load(Twiddle + 2 * index, index < M, 0)[None, :] + wi = tl.load(Twiddle + 2 * index + 1, index < M, 0)[None, :] + # X[k] = (Z[k] + conj(Z[-k]))/2 + # - i W_N^k (Z[k] - conj(Z[-k]))/2. + out_real = 0.5 * (real + br + wr * (imag + bi) + wi * (real - br)) + out_imag = 0.5 * (imag - bi + wi * (imag + bi) - wr * (real - br)) + offsets = 2 * (rows[:, None] * NF + index[None, :]) + tl.store(Output + offsets, out_real, mask) + tl.store(Output + offsets + 1, out_imag, mask) + # The k=0 lane owns the extra, tightly packed Nyquist coefficient. + endpoints = (rows[:, None] < ROWS) & (index[None, :] == 0) + nyquist = 2 * (rows[:, None] * NF + M) + tl.store(Output + nyquist + 0 * index[None, :], real - imag, endpoints) + tl.store(Output + nyquist + 0 * index[None, :] + 1, 0, endpoints) + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def complex_axis(Input, Output, Twiddle, N: tl.constexpr, INNER: tl.constexpr, + P: tl.constexpr, B: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr): + # Program dimension 1 indexes independent outer batches (Z for Y FFTs). + columns = tl.program_id(0) * B + tl.arange(0, B) + index = tl.arange(0, P) + source = _input_index(index, N, RADICES, SIZES) + base = tl.program_id(1) * N * INNER + offsets = 2 * (base + source[:, None] * INNER + columns[None, :]) + mask = (index[:, None] < N) & (columns[None, :] < INNER) + real = tl.load(Input + offsets, mask, 0).to(tl.float64) + imag = tl.load(Input + offsets + 1, mask, 0).to(tl.float64) + real, imag = _fft(real, imag, Twiddle, N, P, B, True, INVERSE, RADICES, SIZES) + offsets = 2 * (base + index[:, None] * INNER + columns[None, :]) + tl.store(Output + offsets, real, mask) + tl.store(Output + offsets + 1, imag, mask) + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def real_xy(Input, Output, TwiddleX, TwiddleY, + NY: tl.constexpr, NX: tl.constexpr, PY: tl.constexpr, PX: tl.constexpr, + INVERSE: tl.constexpr, RX: tl.constexpr, SX: tl.constexpr, + RY: tl.constexpr, SY: tl.constexpr, Q: tl.constexpr): + """Fuse the real X transform with one short Y residue transform. + + Both FFT networks operate on registers/shared memory in the same CTA. + Only caller layout is accessed globally: real NY*NX or packed NY*NF*2. + An inverse call consumes a plane AFTER the Z inverse. Arbitrary boundary + values are projected only after the Y inverse, matching irfftn. + """ + z = tl.program_id(0) + residue = tl.program_id(1) + M: tl.constexpr = NY // Q + y = tl.arange(0, PY) + x = tl.arange(0, PX) + sy = _input_index(y, M, RY, SY) + sx = x + dy, dx = y, _reverse_bits(x, NX) + NF: tl.constexpr = NX // 2 + 1 + if INVERSE: + valid = (y[:, None] < M) & (x[None, :] < NF) + values_r, values_i = (), () + for part in tl.static_range(Q): + offsets = 2 * ((z * NY + sy[:, None] + part * M) * NF + x[None, :]) + ar = tl.load(Input + offsets, valid, 0).to(tl.float64) + ai = tl.load(Input + offsets + 1, valid, 0).to(tl.float64) + values_r, values_i = values_r + (ar,), values_i + (ai,) + er, ei = values_r[0] + values_r[2], values_i[0] + values_i[2] + dr, di = values_r[0] - values_r[2], values_i[0] - values_i[2] + or_, oi = values_r[1] + values_r[3], values_i[1] + values_i[3] + pr, pi = values_r[1] - values_r[3], values_i[1] - values_i[3] + real = tl.where(residue == 0, er + or_, tl.where(residue == 1, dr - pi, + tl.where(residue == 2, er - or_, dr + pi))) + imag = tl.where(residue == 0, ei + oi, tl.where(residue == 1, di + pr, + tl.where(residue == 2, ei - oi, di - pr))) + exponent = sy * residue % NY + if NY == 32: + wr, wi = _constant32_twiddle(exponent) + else: + wr = tl.load(TwiddleY + 2 * exponent) + wi = tl.load(TwiddleY + 2 * exponent + 1) + wi = -wi + real, imag = real * wr[:, None] - imag * wi[:, None], real * wi[:, None] + imag * wr[:, None] + real, imag = _fft(real, imag, TwiddleY, M, PY, PX, True, True, RY, SY, Q) + # Reconstruct the N/2 complex inverse only after the complete Y inverse. + # Its real and imaginary parts are the even and odd real output samples. + H: tl.constexpr = NX // 2 + h = tl.arange(0, H) + hs = h + aindex = tl.broadcast_to(hs[None, :], (PY, H)) + bindex = tl.broadcast_to((H - hs)[None, :], (PY, H)) + ar = tl.gather(real, aindex, 1) + ai = tl.gather(imag, aindex, 1) + br = tl.gather(real, bindex, 1) + bi = tl.gather(imag, bindex, 1) + ai = tl.where(hs[None, :] == 0, 0., ai) + bi = tl.where(hs[None, :] == 0, 0., bi) + if NX == 32: + wr, wi = _constant32_twiddle(hs) + else: + wr = tl.load(TwiddleX + 2 * hs) + wi = tl.load(TwiddleX + 2 * hs + 1) + real = ar + br - wr[None, :] * (ai + bi) + wi[None, :] * (ar - br) + imag = ai - bi + wr[None, :] * (ar - br) + wi[None, :] * (ai + bi) + real, imag = _dif(real, imag, TwiddleX, H, H, PY, False, True, 2) + offsets = (z * NY + Q * dy[:, None] + residue) * NX + 2 * _reverse_bits(h, H)[None, :] + tl.store(Output + offsets, real, y[:, None] < M) + tl.store(Output + offsets + 1, imag, y[:, None] < M) + else: + offsets = (z * NY + Q * sy[:, None] + residue) * NX + sx[None, :] + valid = (y[:, None] < M) & (x[None, :] < NX) + real = tl.load(Input + offsets, valid, 0).to(tl.float64) + imag = tl.full((PY, PX), 0, real.dtype) + real, imag = _fft(real, imag, TwiddleX, NX, PX, PY, False, False, RX, SX) + real, imag = _fft(real, imag, TwiddleY, M, PY, PX, True, False, RY, SY, Q) + offsets = 2 * ((z * NY + residue * M + y[:, None]) * NF + x[None, :]) + valid = (y[:, None] < M) & (x[None, :] < NF) + tl.store(Output + offsets, real, valid) + tl.store(Output + offsets + 1, imag, valid) + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def combined_z(Input, Output, Twiddle, TwiddleY, + N: tl.constexpr, INNER: tl.constexpr, + P: tl.constexpr, B: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr, + NY: tl.constexpr, NX: tl.constexpr, Q: tl.constexpr): + """Finish forward Y in the Z loads, without an extra global stage.""" + NF: tl.constexpr = NX // 2 + 1 + M: tl.constexpr = NY // Q + columns = tl.program_id(0) * B + tl.arange(0, B) + ky = columns // NF + kx = columns % NF + index = tl.arange(0, P) + source = _input_index(index, N, RADICES, SIZES) + mask = (index[:, None] < N) & (columns[None, :] < INNER) + real = tl.full((P, B), 0., tl.float64) + imag = tl.full((P, B), 0., tl.float64) + for residue in tl.static_range(Q): + offsets = 2 * ((source[:, None] * NY + residue * M + ky[None, :] % M) * NF + kx[None, :]) + ar = tl.load(Input + offsets, mask, 0).to(tl.float64) + ai = tl.load(Input + offsets + 1, mask, 0).to(tl.float64) + exponent = ky * residue % NY + if NY == 32: + wr, wi = _constant32_twiddle(exponent) + else: + wr = tl.load(TwiddleY + 2 * exponent) + wi = tl.load(TwiddleY + 2 * exponent + 1) + real += ar * wr[None, :] - ai * wi[None, :] + imag += ar * wi[None, :] + ai * wr[None, :] + real, imag = _fft(real, imag, Twiddle, N, P, B, True, False, RADICES, SIZES) + offsets = 2 * (index[:, None] * INNER + columns[None, :]) + tl.store(Output + offsets, real, mask) + tl.store(Output + offsets + 1, imag, mask) + +@triton.jit +def _reverse_bits(index, N: tl.constexpr): + result = tl.full(index.shape, 0, tl.int32) + for bit in tl.static_range(2 if N == 4 else 3 if N == 8 else 4 if N == 16 else 5): + result = result * 2 + ((index >> bit) & 1) + return result + + +@triton.jit +def _dif_stage(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + H: tl.constexpr = SIZE // 2 + G: tl.constexpr = P // SIZE + if STRIDED: + r = tl.permute(tl.reshape(real, (G, 2, H, B)), (0, 2, 3, 1)) + i = tl.permute(tl.reshape(imag, (G, 2, H, B)), (0, 2, 3, 1)) + else: + r = tl.permute(tl.reshape(real, (B, G, 2, H)), (0, 1, 3, 2)) + i = tl.permute(tl.reshape(imag, (B, G, 2, H)), (0, 1, 3, 2)) + ar, br = tl.split(r) + ai, bi = tl.split(i) + sr, si = ar + br, ai + bi + dr, di = _stage_twiddle(ar - br, ai - bi, Twiddle, N, SIZE, 2, 1, + STRIDED, INVERSE, TWIDDLE_STEP) + r, i = tl.join(sr, dr), tl.join(si, di) + if STRIDED: + return (tl.reshape(tl.permute(r, (0, 3, 1, 2)), (P, B)), + tl.reshape(tl.permute(i, (0, 3, 1, 2)), (P, B))) + else: + return (tl.reshape(tl.permute(r, (0, 1, 3, 2)), (B, P)), + tl.reshape(tl.permute(i, (0, 1, 3, 2)), (B, P))) + + +@triton.jit +def _dif(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, TWIDDLE_STEP: tl.constexpr = 1): + tl.static_assert(N == P and (N == 4 or N == 8 or N == 16 or N == 32)) + for stage in tl.static_range(2 if N == 4 else 3 if N == 8 else 4 if N == 16 else 5): + real, imag = _dif_stage(real, imag, Twiddle, N, P, B, + STRIDED, INVERSE, N >> stage, TWIDDLE_STEP) + return real, imag diff --git a/build/torch-cuda/_b200/grouped_dif_inverse_native_dispatch.c b/build/torch-cuda/_b200/grouped_dif_inverse_native_dispatch.c new file mode 100644 index 0000000000000000000000000000000000000000..cf38e9be18023bf1258f79b5c869f5f78aafa268 --- /dev/null +++ b/build/torch-cuda/_b200/grouped_dif_inverse_native_dispatch.c @@ -0,0 +1,54 @@ +/* SPDX-License-Identifier: Apache-2.0 */ +/* Launch a prepared Triton FFT chain. No CUDA headers or link dependency. + * cuLaunchKernel is resolved before execution. The three or four tensor arguments + * and two unused compiler scratch arguments are validated against the PTX + * signature during preparation. All per-call parameters live on the stack. + */ +#include +#include + +typedef int (*launch_fn)(void *, unsigned int, unsigned int, unsigned int, + unsigned int, unsigned int, unsigned int, + unsigned int, void *, void **, void **); + +struct fft_stage { + void *function; + uint64_t twiddle[2]; + int table_count; + unsigned int grid[3]; + unsigned int threads; + unsigned int shared; + int input_slot; + int output_slot; +}; + +struct fft_chain { + launch_fn launch; + uint64_t scratch; + void *stream; + int count; + struct fft_stage stages[3]; +}; + +int fft_execute(const struct fft_chain *chain, uint64_t input, uint64_t output) { + uint64_t buffers[3] = {input, output, chain->scratch}; + for (int i = 0; i < chain->count; ++i) { + const struct fft_stage *stage = &chain->stages[i]; + uint64_t values[6] = {buffers[stage->input_slot], buffers[stage->output_slot], + stage->twiddle[0], 0, 0, 0}; + if (stage->table_count == 2) values[3] = stage->twiddle[1]; + void *parameters[6] = {&values[0], &values[1], &values[2], + &values[3], &values[4], &values[5]}; + int result = chain->launch(stage->function, + stage->grid[0], stage->grid[1], stage->grid[2], + stage->threads, 1, 1, stage->shared, chain->stream, + parameters, NULL); + if (result != 0) { + return result; + } + } + return 0; +} + +size_t fft_chain_size(void) { return sizeof(struct fft_chain); } +size_t fft_stage_size(void) { return sizeof(struct fft_stage); } diff --git a/build/torch-cuda/_b200/grouped_dif_inverse_native_dispatch.py b/build/torch-cuda/_b200/grouped_dif_inverse_native_dispatch.py new file mode 100644 index 0000000000000000000000000000000000000000..c2722cc753cfd2c65a704f035a60950c45550a27 --- /dev/null +++ b/build/torch-cuda/_b200/grouped_dif_inverse_native_dispatch.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Prepare a native cuLaunchKernel dispatcher before the execution phase.""" +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 .grouped_dif_inverse_abi import validate_kernel_abi + + +class Stage(ct.Structure): + _fields_ = [('function', ct.c_void_p), ('twiddle', ct.c_uint64 * 2), + ('table_count', ct.c_int), + ('grid', ct.c_uint * 3), ('threads', ct.c_uint), ('shared', ct.c_uint), + ('input_slot', ct.c_int), ('output_slot', ct.c_int)] + + +class Chain(ct.Structure): + _fields_ = [('launch', ct.c_void_p), ('scratch', ct.c_uint64), + ('stream', ct.c_void_p), ('count', ct.c_int), ('stages', Stage * 3)] + + +_LIBRARY = None +_DRIVER = None +_LOCK = threading.Lock() + + +def load_library(): + """Build and load using the installed host compiler; no GPU work here.""" + global _LIBRARY + with _LOCK: + if _LIBRARY is not None: + return _LIBRARY + source = Path(__file__).with_name('grouped_dif_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-native-fused-xy-' + 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(Chain), ct.c_uint64, ct.c_uint64] + lib.fft_execute.restype = ct.c_int + for name, cls in (('fft_chain_size', Chain), ('fft_stage_size', Stage)): + function = getattr(lib, name) + function.argtypes = [] + function.restype = ct.c_size_t + if function() != ct.sizeof(cls): + raise RuntimeError('Native dispatcher ABI mismatch: ' + name) + _LIBRARY = lib + return lib + + +def load_driver(): + """Resolve the CUDA launch function during plan preparation only.""" + global _DRIVER + with _LOCK: + if _DRIVER is None: + driver = ct.CDLL('libcuda.so.1') + # Access once now; no name resolution occurs in execute(). + driver.cuLaunchKernel.restype = ct.c_int + _DRIVER = driver + return _DRIVER + + +def kernel_descriptor(compiled, grid, source, target, table_ptrs): + abi = validate_kernel_abi(compiled) + if len(table_ptrs) != abi['tensor_arguments'] - 2: + raise ValueError('Twiddle pointer count differs from the compiled ABI') + if not compiled.function or compiled.module is None: + raise ValueError('Kernel must be loaded and initialized during setup') + return dict(function=compiled.function, grid=tuple(grid), threads=abi['threads'], + shared=abi['shared_bytes'], input_slot=source, output_slot=target, + twiddle=tuple(table_ptrs), table_count=len(table_ptrs)) + + +class PreparedChain: + """Immutable descriptors; new input and output pointers on every call. + + launch_function is injectable only for the CPU ABI test. Production + callers resolve the installed CUDA driver. CompiledKernel owners must + stay alive until all asynchronous launches have completed. + """ + def __init__(self, stages, scratch_ptr, stream_ptr, *, launch_function=None): + if not 1 <= len(stages) <= 3: + raise ValueError('One to three stages are required') + self.library = load_library() + self.driver = load_driver() if launch_function is None else None + self.launch = self.driver.cuLaunchKernel if launch_function is None else launch_function + address = ct.cast(self.launch, ct.c_void_p).value + if not address: + raise ValueError('A resolved launch function is required') + self.chain = Chain(launch=address, scratch=scratch_ptr, stream=stream_ptr, + count=len(stages)) + for index, stage in enumerate(stages): + source, target = stage['input_slot'], stage['output_slot'] + if source not in (0, 1, 2) or target not in (1, 2): + raise ValueError('Invalid stage routing, or attempted input overwrite') + if (source == 2 or target == 2) and not scratch_ptr: + raise ValueError('Intermediate storage is required') + grid = tuple(stage['grid']) + if len(grid) != 3 or any(not 0 < n < 2**32 for n in grid): + raise ValueError('Invalid launch grid') + tables = tuple(stage['twiddle']) + if (stage['table_count'] not in (1, 2) or len(tables) != stage['table_count'] + or any(not 0 < p < 2**64 for p in tables)): + raise ValueError('One or two valid twiddle pointers are required') + if (not stage['function'] + or stage['threads'] not in (32, 64, 128, 256, 512, 1024) + or not 0 <= stage['shared'] < 2**32): + raise ValueError('Invalid stage metadata') + values = dict(stage, grid=(ct.c_uint * 3)(*grid), + twiddle=(ct.c_uint64 * 2)(*(tables + (0,) * (2 - len(tables))))) + self.chain.stages[index] = Stage(**values) + 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) diff --git a/build/torch-cuda/_b200/grouped_dif_inverse_schedule.py b/build/torch-cuda/_b200/grouped_dif_inverse_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..5ecb400ba29972d52ad1de636a118131e38822e5 --- /dev/null +++ b/build/torch-cuda/_b200/grouped_dif_inverse_schedule.py @@ -0,0 +1,74 @@ +# 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") + # The 32-point radix-2 network avoids costly radix-4 layout exchanges. + if remaining in (8, 16, 32): + return (2,) * (remaining.bit_length() - 1) + 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 diff --git a/build/torch-cuda/_b200/half_forward_cufft_z_lto.py b/build/torch-cuda/_b200/half_forward_cufft_z_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..4ea7c33cf959fea015a3655d4511abf0b772e071 --- /dev/null +++ b/build/torch-cuda/_b200/half_forward_cufft_z_lto.py @@ -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('r2c_half_forward_cufft_z_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)] diff --git a/build/torch-cuda/_b200/half_forward_cufft_z_plan_safety.py b/build/torch-cuda/_b200/half_forward_cufft_z_plan_safety.py new file mode 100644 index 0000000000000000000000000000000000000000..0a332465c8fe6d562be03df15f27e408bf9841c8 --- /dev/null +++ b/build/torch-cuda/_b200/half_forward_cufft_z_plan_safety.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reject multi-kernel final plans before accepting caller-sized output. + +cuFFT may use output as intermediate storage in multi-kernel plans. Capture +and inspect each final prepared handle with full nominal double-complex +output storage. Only a graph containing exactly one kernel is supported. +All discovery, stream rebinding and warmup happen at preparation time. +""" +import ctypes as ct +import ctypes.util +import functools +from pathlib import Path + +import torch + +from .scalar_input_inverse_cufft_support import _lib, _check + + +@functools.lru_cache(maxsize=1) +def runtime(): + candidates = [ctypes.util.find_library('cudart')] + for root in (Path(torch.__file__).resolve().parent.parent / 'nvidia/cuda_runtime/lib', + Path('/usr/local/cuda/lib64'), Path('/usr/local/cuda/targets/x86_64-linux/lib')): + candidates.extend(str(p) for p in sorted(root.glob('libcudart.so*'))) + candidates += ['libcudart.so.13', 'libcudart.so.12', 'libcudart.so'] + for path in candidates: + if not path: + continue + try: + lib = ct.CDLL(path) + break + except OSError: + continue + else: + raise RuntimeError('Installed CUDA runtime unavailable') + signatures = { + 'cudaRuntimeGetVersion': [ct.POINTER(ct.c_int)], + 'cudaStreamBeginCapture': [ct.c_void_p, ct.c_int], + 'cudaStreamEndCapture': [ct.c_void_p, ct.POINTER(ct.c_void_p)], + 'cudaGraphGetNodes': [ct.c_void_p, ct.POINTER(ct.c_void_p), ct.POINTER(ct.c_size_t)], + 'cudaGraphNodeGetType': [ct.c_void_p, ct.POINTER(ct.c_int)], + 'cudaGraphDestroy': [ct.c_void_p], + } + for name, signature in signatures.items(): + function = getattr(lib, name) + function.argtypes, function.restype = signature, ct.c_int + return lib + + +def check(status): + if status: + raise RuntimeError(f'CUDA runtime returned error {status}') + + +def require_single_kernel(handle, intermediate, nominal_output, device, target_stream): + """Inspect a handle at setup, retaining neither the graph nor probe stream.""" + lib = runtime() + version = ct.c_int() + check(lib.cudaRuntimeGetVersion(ct.byref(version))) + probe_stream = torch.cuda.Stream(device=device) + graph = ct.c_void_p() + capturing = False + _check(_lib.cufftSetStream(handle.value, probe_stream.cuda_stream)) + try: + _check(_lib.cufftExecZ2Z(handle.value, intermediate.data_ptr(), nominal_output.data_ptr(), -1)) + probe_stream.synchronize() + check(lib.cudaStreamBeginCapture(probe_stream.cuda_stream, 1)) + capturing = True + _check(_lib.cufftExecZ2Z(handle.value, intermediate.data_ptr(), nominal_output.data_ptr(), -1)) + check(lib.cudaStreamEndCapture(probe_stream.cuda_stream, ct.byref(graph))) + capturing = False + count = ct.c_size_t() + check(lib.cudaGraphGetNodes(graph, None, ct.byref(count))) + nodes = (ct.c_void_p * count.value)() + check(lib.cudaGraphGetNodes(graph, nodes, ct.byref(count))) + types = [] + for node in nodes: + kind = ct.c_int() + check(lib.cudaGraphNodeGetType(node, ct.byref(kind))) + types.append(kind.value) + result = dict(runtime_version=version.value, graph_nodes=count.value, + node_types=types, workspace_bytes=handle.workspace_bytes, + passed=count.value == 1 and types == [0]) + if not result['passed']: + raise ValueError('Final cuFFT plan is not one kernel: ' + repr(result)) + return result + finally: + if capturing: + lib.cudaStreamEndCapture(probe_stream.cuda_stream, ct.byref(graph)) + if graph.value: + check(lib.cudaGraphDestroy(graph)) + _check(_lib.cufftSetStream(handle.value, target_stream.cuda_stream)) diff --git a/build/torch-cuda/_b200/half_vector_inverse_lto.py b/build/torch-cuda/_b200/half_vector_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..7c711c164b0cf1ede63713e03bdb0d3c40a0c306 --- /dev/null +++ b/build/torch-cuda/_b200/half_vector_inverse_lto.py @@ -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_half_vector_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)) diff --git a/build/torch-cuda/_b200/hermitian_projection_cufft_support.py b/build/torch-cuda/_b200/hermitian_projection_cufft_support.py new file mode 100644 index 0000000000000000000000000000000000000000..cb2f282fa44fd0caac0ce5550f9bda1bae250826 --- /dev/null +++ b/build/torch-cuda/_b200/hermitian_projection_cufft_support.py @@ -0,0 +1,102 @@ +# 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): + 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: + symbol, callback_type, blob, _metadata = callback + _check(_lib.cufftXtSetJITCallback(self.value, symbol, blob, len(blob), + callback_type, 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 diff --git a/build/torch-cuda/_b200/hermitian_projection_kernels.py b/build/torch-cuda/_b200/hermitian_projection_kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..ecd262f5e7350a89798a23864fc80c369a45cb9f --- /dev/null +++ b/build/torch-cuda/_b200/hermitian_projection_kernels.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Project the two Y endpoint lines after the Z inverse, without repacking.""" +import triton +import triton.language as tl + + +@triton.jit +def project_y_endpoints(data, NZ: tl.constexpr, NY: tl.constexpr, + NF: tl.constexpr, BLOCK: tl.constexpr): + # Each lane owns both members of a conjugate pair. No other lane reads + # or writes those addresses, so the projection is safe in place. + index = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + HALF_Y: tl.constexpr = NY // 2 + 1 + z = index // (2 * HALF_Y) + y = (index // 2) % HALF_Y + kx = (index % 2) * (NF - 1) + mirror_y = tl.where(y == 0, 0, NY - y) + a = 2 * ((z * NY + y) * NF + kx) + b = 2 * ((z * NY + mirror_y) * NF + kx) + valid = z < NZ + ar = tl.load(data + a, valid, 0) + ai = tl.load(data + a + 1, valid, 0) + br = tl.load(data + b, valid, 0) + bi = tl.load(data + b + 1, valid, 0) + real = 0.5 * ar + 0.5 * br + imag = 0.5 * ai - 0.5 * bi + tl.store(data + a, real, valid) + tl.store(data + a + 1, imag, valid) + tl.store(data + b, real, valid & (y != mirror_y)) + tl.store(data + b + 1, -imag, valid & (y != mirror_y)) diff --git a/build/torch-cuda/_b200/hermitian_projection_lto.py b/build/torch-cuda/_b200/hermitian_projection_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..78c39a77c3d5ab859bf63c2cd7aed21e9f4acb09 --- /dev/null +++ b/build/torch-cuda/_b200/hermitian_projection_lto.py @@ -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): + nz, ny, nx = shape + nf = nx // 2 + 1 + lib = library() + options = ('--std=c++17', '--relocatable-device-code=true', '-dlto', + f'--gpu-architecture=compute_{arch}', f'-DFFT_NF={nf}', + f'-DFFT_NZ={nz}', f'-DFFT_NY={ny}', f'-DFFT_DOUBLE={int(double)}') + program = ct.c_void_p() + _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'projection.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=list(shape), nf=nf, 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, double, arch=100): + """Return the callback symbol, enum, owned host LTO buffer and metadata.""" + shape = tuple(int(n) for n in shape) + if len(shape) != 3 or min(shape) < 1 or shape[2] % 2: + raise ValueError('Expected a positive real-space shape with even nx') + source = Path(__file__).with_name('projection.cu').read_bytes() + with _LOCK: + blob, metadata = _compile(source, shape, bool(double), int(arch), version()) + return (b'fft_hermitian_load', 1 if double else 0, + ct.create_string_buffer(blob, len(blob)), dict(metadata, kind='load')) diff --git a/build/torch-cuda/_b200/inverse_cuda_f32.py b/build/torch-cuda/_b200/inverse_cuda_f32.py new file mode 100644 index 0000000000000000000000000000000000000000..7bd7aa4fae7f51c4e4ae214f0c287d1627028067 --- /dev/null +++ b/build/torch-cuda/_b200/inverse_cuda_f32.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Full-FP64 half-length XY inverse followed by a CUDA register Z inverse.""" +from __future__ import annotations + +import math +import threading + +import torch +import triton + +from .cuda_z_inverse_cufft_support import _Handle, _check, _lib +from .c2r_cuda_z_inverse_cuda_kernel import compile_kernel +from .cuda_z_inverse_lto import compile_callback +from .cuda_z_inverse_native_dispatch import PreparedDispatch +from .cuda_z_inverse_schedule import twiddles, factors, stage_sizes + + +class Plan: + def __init__(self, shape, dtype, direction='c2r', *, 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 != torch.float32 or direction != 'c2r': + raise ValueError('Expected an FP32 inverse 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.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 = (nz, ny, nf, 2), self.shape + 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] + self.table = torch.tensor(twiddles(nx), dtype=torch.float64, device=self.device) + self.slot = torch.tensor([0, self.table.data_ptr()], dtype=torch.int64, device=self.device) + # This intermediate contains the transformed half-length XY data. + # The last Z stage writes directly to caller FP32 storage. + self.scratch = torch.empty((nz, ny, m, 2), dtype=torch.float64, + device=self.device) + self.slot_ptr, self.scratch_ptr = self.slot.data_ptr(), self.scratch.data_ptr() + patient = variant.endswith('_patient') + for residue in range(4): + load_callback = compile_callback(self.shape, residue, 'load', arch) + 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, caller_info=self.slot_ptr) + self.handles.append(handle) + self.handle_values = tuple(h.value for h in self.handles) + self.z_table = torch.tensor(twiddles(nz), dtype=torch.float64, + device=self.device) + radices = factors(nz) + self.constants = (nz, ny * m, triton.next_power_of_2(nz), columns, + True, radices, stage_sizes(radices)) + compiled = compile_kernel(self.shape, columns, warps, arch, self.device.index) + 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(4): + self.execute_ptr(dummy.data_ptr() + residue * 4, 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 + 16 + self.table.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 % 4 or b % 4: + raise ValueError('Distinct non-null scalar-aligned pointers are required') + residue = (a & 15) // 4 + if self.dispatch == 'native': + result = self.native_function(self.native_pointers[residue], a - 4 * residue, b) + if result > 0: + _check(result) + if result < 0: + raise RuntimeError(f'cuLaunchKernel failed with CUDA driver error {-result}') + return + _check(self.fft_function(self.handle_values[residue], a - 4 * residue, self.scratch_ptr, 1)) + self.runner(self.scratch_ptr, b, self.z_table.data_ptr(), *self.constants, + stream=self.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() + 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.slot = self.table = self.z_table = self.runner = self.compiled_kernel = None + self.native = self.native_pointers = None + self.closed = True + + + + +_CACHE = {} +_LOCK = threading.Lock() diff --git a/build/torch-cuda/_b200/inverse_cuda_f64.py b/build/torch-cuda/_b200/inverse_cuda_f64.py new file mode 100644 index 0000000000000000000000000000000000000000..1e4864341511d1baeb4a5b988b25d1f9ef3e616b --- /dev/null +++ b/build/torch-cuda/_b200/inverse_cuda_f64.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Full-FP64 half-length XY inverse followed by a CUDA register Z inverse.""" +from __future__ import annotations + +import math +import threading + +import torch +import triton + +from .cuda_z_inverse_cufft_support import _Handle, _check, _lib +from .c2r_f64_cuda_z_inverse_cuda_kernel import compile_kernel +from .c2r_f64_cuda_z_inverse_lto import compile_callback +from .c2r_f64_cuda_z_inverse_native_dispatch import PreparedDispatch +from .cuda_z_inverse_schedule import twiddles, factors, stage_sizes + + +class Plan: + def __init__(self, shape, dtype, direction='c2r', *, 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 != torch.float64 or direction != 'c2r': + raise ValueError('Expected an FP64 inverse 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.dispatch = dispatch + if columns not in (2, 4, 8, 16, 32) or warps not in (1, 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 = (nz, ny, nf, 2), self.shape + 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] + self.table = torch.tensor(twiddles(nx), dtype=torch.float64, device=self.device) + self.slot = torch.tensor([0, self.table.data_ptr()], dtype=torch.int64, device=self.device) + # This intermediate contains the transformed half-length XY data. + # The last Z stage writes directly to caller FP64 storage. + self.scratch = torch.empty((nz, ny, m, 2), dtype=torch.float64, + device=self.device) + self.slot_ptr, self.scratch_ptr = self.slot.data_ptr(), self.scratch.data_ptr() + patient = variant.endswith('_patient') + for residue in range(2): + load_callback = compile_callback(self.shape, residue, 'load', arch) + 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, caller_info=self.slot_ptr) + self.handles.append(handle) + self.handle_values = tuple(h.value for h in self.handles) + self.z_table = torch.tensor(twiddles(nz), dtype=torch.float64, + device=self.device) + radices = factors(nz) + self.constants = (nz, ny * m, triton.next_power_of_2(nz), columns, + True, radices, stage_sizes(radices)) + compiled = compile_kernel(self.shape, columns, warps, arch, self.device.index) + 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(2): + self.execute_ptr(dummy.data_ptr() + residue * 8, 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 + 16 + self.table.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 % 8 or b % 8: + raise ValueError('Distinct non-null scalar-aligned pointers are required') + residue = (a & 15) // 8 + if self.dispatch == 'native': + result = self.native_function(self.native_pointers[residue], a - 8 * residue, b) + if result > 0: + _check(result) + if result < 0: + raise RuntimeError(f'cuLaunchKernel failed with CUDA driver error {-result}') + return + _check(self.fft_function(self.handle_values[residue], a - 8 * residue, self.scratch_ptr, 1)) + self.runner(self.scratch_ptr, b, self.z_table.data_ptr(), *self.constants, + stream=self.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() + 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.slot = self.table = self.z_table = self.runner = self.compiled_kernel = None + self.native = self.native_pointers = None + self.closed = True + + + + +_CACHE = {} +_LOCK = threading.Lock() diff --git a/build/torch-cuda/_b200/inverse_dense_f32.py b/build/torch-cuda/_b200/inverse_dense_f32.py new file mode 100644 index 0000000000000000000000000000000000000000..849058071c71bd87a7d0aeb2ac8b49854d07d9e8 --- /dev/null +++ b/build/torch-cuda/_b200/inverse_dense_f32.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FP64 half-length inverse with a checked direct scalar Z output stage.""" +from __future__ import annotations + +import math +import threading + +import torch + +from .cuda_z_inverse_cufft_support import _Handle, _lib, _check +from .dense_half_inverse_lto import compile_callback +from .direct_dispatch import PreparedChain +from .paired_x_inverse_plan_safety import require_single_kernel +from .dense_half_inverse_schedule import twiddles + + +class Plan: + """Prepare every scalar alignment before enqueue-only caller execution.""" + def __init__(self, shape, dtype, direction='c2r', *, device=None, stream=None, + variant='xy_z', dispatch='native'): + 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, even nx >= 16)') + if dtype not in (torch.float32, torch.float64) or direction != 'c2r': + raise ValueError('Expected a float32 or float64 inverse transform') + if variant not in ('xy_z', 'xy_z_patient') or dispatch != 'native': + raise ValueError('Expected a native XY/Z decomposition') + self.dtype, self.direction, self.variant, self.dispatch = dtype, direction, variant, dispatch + 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.closed, self.handles, self.chains = False, [], [] + self.compiled_metadata, self.final_plan_checks = [], [] + nz, ny, nx = self.shape + m, nf = nx // 2, nx // 2 + 1 + self.input_shape, self.output_shape = (nz, ny, nf, 2), self.shape + patient = variant.endswith('_patient') + try: + with torch.cuda.device(self.device), torch.cuda.stream(self.stream): + self.intermediate = torch.empty((nz, ny, m, 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] + loads, stores = [], [] + for residue in range(self.residues): + load = compile_callback(self.shape, residue, 'load', arch, + double=dtype == torch.float64) + first = _Handle((ny, m), 0x69, (ny, m), 1, ny * m, + (ny, m), 1, ny * m, nz, self.device, self.stream, + callback=[load], patient=patient, + caller_info=self.table.data_ptr()) + self.handles.append(first) + loads.append(first) + store = compile_callback(self.shape, residue, 'store', arch, + double=dtype == torch.float64) + final = _Handle((nz,), 0x69, (nz,), ny * m, 1, + (nz,), ny * m, 1, ny * m, self.device, self.stream, + callback=[store], patient=patient) + self.handles.append(final) + stores.append(final) + self.compiled_metadata.extend((load[3], store[3])) + # Inspect all final plans using full nominal FP64 storage. + # Multi-kernel cuFFT may use output as an intermediate and is + # unsuitable for scalar-offset or narrower caller storage. + nominal = torch.empty(math.prod(self.shape) * 8 + 16, + dtype=torch.uint8, device=self.device) + for final in stores: + self.final_plan_checks.append(require_single_kernel( + final, self.intermediate, nominal, self.device, self.stream)) + for first in loads: + for final in stores: + self.chains.append(PreparedChain([ + (_lib.cufftExecZ2Z, first.value, 1, 1, 0, 2), + (_lib.cufftExecZ2Z, final.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) + dummy = torch.zeros(math.prod(self.input_shape) + self.residues, + dtype=dtype, device=self.device) + for i in range(self.residues): + for o in range(self.residues): + self.execute_ptr(dummy.data_ptr() + self.itemsize * i, + nominal.data_ptr() + self.itemsize * o, 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() * 8 + self.table.numel() * 8) + + def execute_ptr(self, input_ptr, output_ptr, stream_ptr): + """Enqueue with new caller addresses using only prepared state.""" + 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.itemsize or b % self.itemsize: + raise ValueError('Distinct non-null scalar-aligned pointers are required') + 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 wait for outstanding uses before closing.""" + 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() diff --git a/build/torch-cuda/_b200/inverse_dense_f64.py b/build/torch-cuda/_b200/inverse_dense_f64.py new file mode 100644 index 0000000000000000000000000000000000000000..4c1e187813538478f2abe92bc27309f68347ca6c --- /dev/null +++ b/build/torch-cuda/_b200/inverse_dense_f64.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Half-length FP64 inverse with direct aligned and redirected scalar output.""" +from __future__ import annotations + +import math +import threading + +import torch +import triton +from triton.backends.compiler import GPUTarget +from triton.compiler import ASTSource + +from .cuda_z_inverse_cufft_support import _Handle, _check, _lib +from .dense_half_inverse_kernels import publish_output +from .c2r_f64_dense_half_inverse_lto import compile_callback +from .dense_half_inverse_native_dispatch import PreparedDispatch +from .direct_dispatch import PreparedChain +from .dense_half_inverse_schedule import twiddles + + +class Plan: + def __init__(self, shape, dtype, direction='c2r', *, device=None, stream=None, + variant='3d', dispatch='native'): + 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 != torch.float64 or direction != 'c2r': + raise ValueError('Expected an FP64 inverse transform') + if variant not in ('3d', '3d_patient', 'xy_z', 'xy_z_patient'): + raise ValueError('Unknown variant') + if dispatch not in ('native', 'python'): + raise ValueError('Unknown dispatch mode') + self.dtype, self.direction, self.variant = dtype, direction, variant + self.dispatch = dispatch + 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.direct_handles = [] + self.direct_native = [] + self.z_handle = None + self.compiled_metadata = [] + nz, ny, nx = self.shape + nf = nx // 2 + 1 + self.input_shape, self.output_shape = (nz, ny, nf, 2), self.shape + m = nx // 2 + self.fft_function = self.z_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] + self.table = torch.tensor(twiddles(nx), dtype=torch.float64, device=self.device) + self.slot = torch.tensor([0, self.table.data_ptr()], dtype=torch.int64, device=self.device) + # cuFFT may alternate output and workspace during a transform. + # Allocate the entire FP64 extent even though the final callback + # writes to caller FP64 storage. No raw input is copied here. + self.scratch = torch.empty((nz, ny, m, 2), dtype=torch.float64, + device=self.device) + self.slot_ptr, self.scratch_ptr = self.slot.data_ptr(), self.scratch.data_ptr() + store_callback = compile_callback(self.shape, 0, 'store', arch) + self.compiled_metadata.append(store_callback[3]) + patient = variant.endswith('_patient') + for residue in range(2): + load_callback = compile_callback(self.shape, residue, 'load', arch) + self.compiled_metadata.append(dict(load_callback[3], input_residue=residue)) + if variant.startswith('3d'): + handle = _Handle( + (nz, ny, m), 0x69, (nz, ny, m), 1, nz * ny * m, + (nz, ny, m), 1, nz * ny * m, 1, self.device, self.stream, + callback=[load_callback, store_callback], patient=patient, + caller_info=self.slot_ptr) + else: + 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, caller_info=self.slot_ptr) + self.handles.append(handle) + self.handle_values = tuple(h.value for h in self.handles) + if variant.startswith('xy_z'): + self.z_handle = _Handle( + (nz,), 0x69, (nz,), ny * m, 1, (nz,), ny * m, 1, + ny * m, self.device, self.stream, callback=[store_callback], + patient=patient, caller_info=self.slot_ptr) + self.z_handle_value = self.z_handle.value + # Complex-aligned caller output can hold every cuFFT stage. + # Scalar-offset output uses the fully prepared redirected route. + for residue in range(2): + load_callback = compile_callback(self.shape, residue, 'load', arch) + if variant.startswith('3d'): + direct = _Handle( + (nz, ny, m), 0x69, (nz, ny, m), 1, nz * ny * m, + (nz, ny, m), 1, nz * ny * m, 1, self.device, self.stream, + callback=[load_callback], patient=patient, + caller_info=self.slot_ptr) + else: + direct = _Handle( + (ny, m), 0x69, (ny, m), 1, ny * m, + (ny, m), 1, ny * m, nz, self.device, self.stream, + callback=[load_callback], patient=patient, + caller_info=self.slot_ptr) + self.direct_handles.append(direct) + final_direct = None + if variant.startswith('xy_z'): + final_direct = _Handle( + (nz,), 0x69, (nz,), ny * m, 1, (nz,), ny * m, 1, + ny * m, self.device, self.stream, patient=patient) + self.direct_handles.append(final_direct) + for first in self.direct_handles[:2]: + stages = [(self.fft_function, first.value, 1, 1, 0, 1)] + if final_direct is not None: + stages.append((self.z_function, final_direct.value, 1, 1, 1, 1)) + self.direct_native.append(PreparedChain(stages, 0)) + self.direct_function = self.direct_native[0].function + self.direct_pointers = tuple(p.pointer for p in self.direct_native) + source = ASTSource(publish_output, {'slot': '*i64', 'output_address': 'u64'}) + compiled = triton.compile(source, target=GPUTarget('cuda', arch, 32), + options={'num_warps': 1}) + for name in ('global_scratch_size', 'profile_scratch_size'): + if getattr(compiled.metadata, name, 0): + raise ValueError('Execution scratch is unsupported: ' + name) + self.compiled_kernel = compiled + self.runner = compiled[(1, 1, 1)] + if dispatch == 'native': + for handle in self.handle_values: + self.native.append(PreparedDispatch( + self.fft_function, handle, self.z_function, + self.z_handle_value if self.z_handle is not None else 0, + compiled, self.slot_ptr, self.scratch_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.metadata.shared)) + dummy = torch.zeros(math.prod(self.input_shape) + 4, dtype=dtype, + device=self.device) + output = torch.empty(math.prod(self.output_shape) + 1, dtype=dtype, device=self.device) + for residue in range(2): + for output_residue in range(2): + self.execute_ptr(dummy.data_ptr() + residue * 8, + output.data_ptr() + output_residue * 8, self.stream_ptr) + self.stream.synchronize() + except BaseException: + for handle in self.handles + self.direct_handles: + handle.close() + if self.z_handle is not None: + self.z_handle.close() + raise + + @property + def workspace_bytes(self): + return (sum(h.workspace_bytes for h in self.handles + self.direct_handles) + + (self.z_handle.workspace_bytes if self.z_handle is not None else 0) + + self.scratch.numel() * 8 + 16 + self.table.numel() * 8) + + def execute_ptr(self, input_ptr, output_ptr, stream_ptr): + """Enqueue prepared FFTs for the public scalar alignment metadata.""" + 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 % 8 or b % 8: + raise ValueError('Distinct non-null scalar-aligned pointers are required') + residue = (a & 15) // 8 + if b % 16 == 0: + _check(self.direct_function(self.direct_pointers[residue], a - 8 * residue, b)) + return + if self.dispatch == 'native': + result = self.native_function(self.native_pointers[residue], a - 8 * residue, b) + if result > 0: + _check(result) + if result < 0: + raise RuntimeError(f'cuLaunchKernel failed with CUDA driver error {-result}') + return + self.runner(self.slot_ptr, b, stream=self.stream_ptr) + _check(self.fft_function(self.handle_values[residue], a - 8 * residue, self.scratch_ptr, 1)) + if self.z_handle is not None: + _check(self.z_function(self.z_handle_value, self.scratch_ptr, self.scratch_ptr, 1)) + + 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 + self.direct_handles: + handle.close() + if self.z_handle is not None: + self.z_handle.close() + self.scratch = self.slot = self.table = self.runner = self.compiled_kernel = None + self.native = self.native_pointers = self.direct_native = self.direct_pointers = None + self.closed = True + + + + +_CACHE = {} +_LOCK = threading.Lock() diff --git a/build/torch-cuda/_b200/inverse_narrow.py b/build/torch-cuda/_b200/inverse_narrow.py new file mode 100644 index 0000000000000000000000000000000000000000..a4dda4c1c808438a4c620e06ddfef00236b82347 --- /dev/null +++ b/build/torch-cuda/_b200/inverse_narrow.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Prepared complex 3D inverse with half-length reconstruction on input load.""" +from __future__ import annotations + +import threading +import torch + +from .narrow_index_inverse_cufft_support import _Handle, _lib, _check +from .narrow_index_inverse_lto import compile_callbacks +from .direct_dispatch import PreparedChain +from .dense_half_inverse_schedule import twiddles, factors +from . import scalar_output_backend as fallback_backend + + +class Plan: + """One caller stream; prepare every branch before enqueue-only execution. + + cuFFT complex outputs require complex alignment. For a scalar-offset real + output, use the prepared one-dimensional half-length inverse with its + scalar store callback. Do not shift the + pointer supplied to a multidimensional complex plan: its internal stages + may write the output before the final callback. No caller buffer is padded. + """ + def __init__(self, shape, dtype, direction='c2r', *, device=None, stream=None, + strategy='3d', patient=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 != torch.float64 or direction != 'c2r': + raise ValueError('This experiment supports FP64 inverse transforms') + if strategy not in ('3d', 'xy_z', 'x_zy'): + raise ValueError('Unknown decomposition') + for n in self.shape: + factors(n) + self.dtype, self.direction = dtype, direction + self.variant = f'full_half_{strategy}_p{int(patient)}' + 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.closed, self.handles, self.offset_plan = False, [], None + nz,ny,nx = self.shape + m,nf,rows = nx//2,nx//2+1,nz*ny + self.input_shape, self.output_shape = (nz,ny,nf,2), self.shape + self._execute = _lib.cufftExecZ2Z + try: + with torch.cuda.device(self.device), torch.cuda.stream(self.stream): + self.table = torch.tensor(twiddles(nx),dtype=dtype,device=self.device) + cap = torch.cuda.get_device_capability(self.device) + callbacks = compile_callbacks(self.shape,10*cap[0]+cap[1]) + self.compiled_metadata = [v[3] for v in callbacks] + def handle(sizes,ie,istride,idist,oe,ostride,odist,batch,load=False): + result = _Handle(sizes,0x69,ie,istride,idist,oe,ostride,odist,batch, + self.device,self.stream,callback=callbacks if load else None, + caller_info=self.table.data_ptr() if load else None, + patient=bool(patient)) + self.handles.append(result) + return result + if strategy == '3d': + handle((nz,ny,m),(nz,ny,nf),1,rows*nf,(nz,ny,m),1,rows*m,1,True) + elif strategy == 'xy_z': + handle((ny,m),(ny,nf),1,ny*nf,(ny,m),1,ny*m,nz,True) + handle((nz,),(nz,),ny*m,1,(nz,),ny*m,1,ny*m) + else: + handle((m,),(nf,),1,nf,(m,),1,m,rows,True) + handle((nz,ny),(nz,ny),m,1,(nz,ny),m,1,m) + stages = [(self._execute,self.handles[0].value,1,1,0,1)] + stages.extend((self._execute,h.value,1,1,1,1) for h in self.handles[1:]) + self.native = PreparedChain(stages,0) + self.native_function, self.native_pointer = self.native.function, self.native.pointer + self.offset_plan = fallback_backend.Plan(self.shape,dtype,'c2r',device=self.device, + stream=self.stream,layout='frequency',patient=False) + dummy_input = torch.zeros(self.input_shape,dtype=dtype,device=self.device) + dummy_output = torch.empty(rows*nx+1,dtype=dtype,device=self.device) + for shift in (0,1): + self.execute_ptr(dummy_input.data_ptr(),dummy_output.data_ptr()+8*shift,self.stream_ptr) + self.stream.synchronize() + except BaseException: + for h in reversed(self.handles): + h.close() + if self.offset_plan is not None: + self.offset_plan.close() + raise + + @property + def workspace_bytes(self): + return (sum(h.workspace_bytes for h in self.handles) + self.table.numel()*8 + + self.offset_plan.workspace_bytes) + + def execute_ptr(self,input_ptr,output_ptr,stream_ptr): + """Enqueue prepared operations using new caller addresses on each call.""" + 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(output_ptr) % 8: + raise ValueError('The real output must be double-aligned') + if int(output_ptr) & 8: + self.offset_plan.execute_ptr(input_ptr,output_ptr,stream_ptr) + else: + _check(self.native_function(self.native_pointer,input_ptr,output_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() + 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') + self.execute_ptr(a,b,(self.stream if stream is None else stream).cuda_stream) + return output + + def close(self): + if not self.closed: + for h in reversed(self.handles): + h.close() + self.offset_plan.close() + self.table = self.native = self.offset_plan = None + self.closed = True + + + + +_CACHE = {} +_LOCK = threading.Lock() diff --git a/build/torch-cuda/_b200/inverse_paired_x_f32.py b/build/torch-cuda/_b200/inverse_paired_x_f32.py new file mode 100644 index 0000000000000000000000000000000000000000..fd6164929bbc4a7e0fe6308da475dfa38b9a0a1f --- /dev/null +++ b/build/torch-cuda/_b200/inverse_paired_x_f32.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reconstruct and invert half-X in CUDA, then finish separate Y/Z using FP64 cuFFT.""" +from __future__ import annotations +import math +import torch +from .cuda_z_inverse_cufft_support import _Handle, _lib, _check +from .c2r_f64_paired_x_inverse_cuda_kernel import compile_kernel +from .paired_x_inverse_lto import compile_callback +from .paired_x_inverse_native_dispatch import PreparedDispatch +from .cuda_z_inverse_schedule import twiddles +from .paired_x_inverse_plan_safety import require_single_kernel + + +class Plan: + def __init__(self, shape, dtype, direction='c2r', *, device=None, stream=None, + columns=4, warps=4, frequency=False, patient=False): + self.shape = tuple(int(n) for n in shape) + nz, ny, nx = self.shape + if min(nz, ny) < 8 or nx < 16 or nx % 2: + raise ValueError('Expected a valid even-X grid') + if dtype not in (torch.float32, torch.float64) or direction != 'c2r': + raise ValueError('Expected a float32 or float64 inverse') + self.dtype = dtype + self.direction = direction + self.itemsize = 8 if dtype == torch.float64 else 4 + self.device = torch.device(device if device is not None else torch.cuda.current_device()) + if self.device.type != 'cuda': + raise ValueError('Expected a CUDA device') + 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.closed = False + self.handles = [] + self.compiled_metadata = [] + self.variant = f'cuda_half_x_c{columns}_w{warps}_f{int(frequency)}_p{int(patient)}' + self.dispatch = 'native' + m, rows = nx//2, nz*ny + if nz % 2 or ny % 2 or columns % 2 or frequency: + raise ValueError('This paired-row experiment requires even axes and natural output') + row_pairs = [] + for row in range(rows): + z, y = divmod(row, ny) + mirror = ((-z) % nz)*ny + (-y) % ny + if row <= mirror: + row_pairs.append((row, mirror)) + assert len(row_pairs) == rows//2 + 2 + self.input_shape, self.output_shape = (nz, ny, m+1, 2), self.shape + try: + with torch.cuda.device(self.device), torch.cuda.stream(self.stream): + self.intermediate = torch.empty((nz, ny, m, 2), dtype=torch.float64, device=self.device) + self.middle = torch.empty_like(self.intermediate) + self.nominal = torch.empty_like(self.intermediate) + self.table = torch.tensor(twiddles(m)+twiddles(nx)+[v for pair in row_pairs for v in pair], dtype=torch.float64, device=self.device) + self.slot = torch.zeros((1,), dtype=torch.int64, device=self.device) + cap = torch.cuda.get_device_capability(self.device) + arch = 10*cap[0]+cap[1] + compiled = compile_kernel(self.shape, columns, warps, arch, + self.device.index, frequency, dtype == torch.float64) + self.compiled_kernel = compiled + middle_store = compile_callback(self.shape, 0, 'middle', arch, + double=dtype == torch.float64, frequency=frequency) + middle = _Handle((ny,), 0x69, (ny,), 1 if frequency else nz*m, + ny if frequency else 1, (ny,), 1, ny, + nz*m, self.device, self.stream, callback=[middle_store], + patient=patient, caller_info=self.middle.data_ptr()) + self.handles.append(middle) + finals, store_metadata = [], [] + self.final_plan_checks = [] + probe = torch.empty(math.prod(self.shape)*8+16, dtype=torch.uint8, device=self.device) + for residue in range(16//self.itemsize): + store = compile_callback(self.shape, residue, 'store', arch, + double=dtype == torch.float64, frequency=False) + final = _Handle((nz,), 0x69, (nz,), ny*m, 1, (nz,), ny*m, 1, + ny*m, self.device, self.stream, callback=[store],patient=patient) + self.handles.append(final) + self.final_plan_checks.append(require_single_kernel(final, self.middle, probe, + self.device, self.stream)) + finals.append(final.value) + store_metadata.append(store[3]) + self.native = PreparedDispatch(_lib.cufftExecZ2Z, finals[0], + compiled, compiled.grid, self.intermediate.data_ptr(), + self.table.data_ptr(), self.stream_ptr, self.slot.data_ptr(), + middle.value, self.middle.data_ptr(), self.nominal.data_ptr(), + finals, self.itemsize) + self.native_function, self.native_pointer = self.native.function, self.native.pointer + self.compiled_metadata = store_metadata + [middle_store[3],dict(name=compiled.name, + compiler=compiled.compiler_metadata,registers=compiled.n_regs, + local_memory_bytes=compiled.n_spills,grid=compiled.grid)] + dummy = torch.zeros(math.prod(self.input_shape)+4, dtype=dtype, device=self.device) + output = torch.empty(math.prod(self.output_shape)+4, dtype=dtype, device=self.device) + for i in range(16//self.itemsize): + for o in range(16//self.itemsize): + self.execute_ptr(dummy.data_ptr()+i*self.itemsize, + output.data_ptr()+o*self.itemsize, self.stream_ptr) + self.stream.synchronize() + except BaseException: + for h in reversed(self.handles): + h.close() + self.closed = True + raise + + @property + def workspace_bytes(self): + return sum(h.workspace_bytes for h in self.handles)+8*(self.intermediate.numel()+self.middle.numel()+self.nominal.numel()+self.table.numel()+1) + + def execute_ptr(self, input_ptr, output_ptr, stream_ptr): + 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.itemsize or b%self.itemsize: + raise ValueError('Distinct scalar-aligned pointers are required') + result = self.native_function(self.native_pointer, a, b) + if result > 0: + _check(result) + if result < 0: + raise RuntimeError(f'CUDA driver returned {-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()*self.itemsize, output.numel()*self.itemsize + if a < b+nb and b < a+na: + raise ValueError('Buffers must not overlap') + self.execute_ptr(a, b, (self.stream if stream is None else stream).cuda_stream) + return output + + def close(self): + if not self.closed: + for h in reversed(self.handles): + h.close() + self.nominal = self.middle = self.intermediate = self.table = self.slot = self.native = self.native_pointer = None + self.compiled_kernel = None + self.closed = True diff --git a/build/torch-cuda/_b200/inverse_paired_y.py b/build/torch-cuda/_b200/inverse_paired_y.py new file mode 100644 index 0000000000000000000000000000000000000000..32beabf837e9d7a709f0268c75e21b9b24479c51 --- /dev/null +++ b/build/torch-cuda/_b200/inverse_paired_y.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FP64 inverse with real samples paired along the Y axis.""" +from __future__ import annotations + +import math +import threading + +import torch + +from .cuda_z_inverse_cufft_support import _Handle, _lib, _check +from .paired_y_inverse_lto import compile_callback +from .direct_dispatch import PreparedChain +from .paired_x_inverse_plan_safety import require_single_kernel +from .dense_half_inverse_schedule import twiddles + + +class Plan: + """Prepare every scalar alignment before enqueue-only caller execution.""" + def __init__(self, shape, dtype, direction='c2r', *, device=None, stream=None, + variant='xy_z', dispatch='native'): + 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, even nx >= 16)') + if dtype not in (torch.float32, torch.float64) or direction != 'c2r': + raise ValueError('Expected a float32 or float64 inverse transform') + if variant not in ('xy_z', 'xy_z_patient') or dispatch != 'native': + raise ValueError('Expected a native XY/Z decomposition') + self.dtype, self.direction, self.variant, self.dispatch = dtype, direction, variant, dispatch + 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.closed, self.handles, self.chains = False, [], [] + self.compiled_metadata, self.final_plan_checks = [], [] + nz, ny, nx = self.shape + m, nf = nx // 2, nx // 2 + 1 + hz, hy = nz, ny // 2 + if ny % 2: + raise ValueError("The paired axis must be even") + self.input_shape, self.output_shape = (nz, ny, nf, 2), self.shape + patient = variant.endswith('_patient') + try: + with torch.cuda.device(self.device), torch.cuda.stream(self.stream): + self.intermediate = torch.empty((hz, hy, nx, 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] + loads, stores = [], [] + for residue in range(self.residues): + load = compile_callback(self.shape, residue, 'load', arch, + double=dtype == torch.float64) + first = _Handle((hy, nx), 0x69, (hy, nx), 1, hy * nx, + (hy, nx), 1, hy * nx, hz, self.device, self.stream, + callback=[load], patient=patient, + caller_info=self.table.data_ptr()) + self.handles.append(first) + loads.append(first) + store = compile_callback(self.shape, residue, 'store', arch, + double=dtype == torch.float64) + final = _Handle((hz,), 0x69, (hz,), hy * nx, 1, + (hz,), hy * nx, 1, hy * nx, self.device, self.stream, + callback=[store], patient=patient) + self.handles.append(final) + stores.append(final) + self.compiled_metadata.extend((load[3], store[3])) + # Inspect all final plans using full nominal FP64 storage. + # Multi-kernel cuFFT may use output as an intermediate and is + # unsuitable for scalar-offset or narrower caller storage. + nominal = torch.empty(math.prod(self.shape) * 8 + 16, + dtype=torch.uint8, device=self.device) + for final in stores: + self.final_plan_checks.append(require_single_kernel( + final, self.intermediate, nominal, self.device, self.stream)) + for first in loads: + for final in stores: + self.chains.append(PreparedChain([ + (_lib.cufftExecZ2Z, first.value, 1, 1, 0, 2), + (_lib.cufftExecZ2Z, final.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) + dummy = torch.zeros(math.prod(self.input_shape) + self.residues, + dtype=dtype, device=self.device) + for i in range(self.residues): + for o in range(self.residues): + self.execute_ptr(dummy.data_ptr() + self.itemsize * i, + nominal.data_ptr() + self.itemsize * o, 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() * 8 + self.table.numel() * 8) + + def execute_ptr(self, input_ptr, output_ptr, stream_ptr): + """Enqueue with new caller addresses using only prepared state.""" + 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.itemsize or b % self.itemsize: + raise ValueError('Distinct non-null scalar-aligned pointers are required') + 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 wait for outstanding uses before closing.""" + 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() diff --git a/build/torch-cuda/_b200/inverse_paired_y_f32.py b/build/torch-cuda/_b200/inverse_paired_y_f32.py new file mode 100644 index 0000000000000000000000000000000000000000..c935ec58af9c23d717f0fe0838225406ba650587 --- /dev/null +++ b/build/torch-cuda/_b200/inverse_paired_y_f32.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FP64 inverse with real samples paired along the Y axis.""" +from __future__ import annotations + +import math +import threading + +import torch + +from .cuda_z_inverse_cufft_support import _Handle, _lib, _check +from .paired_y_direct_inverse_lto import compile_callback +from .direct_dispatch import PreparedChain +from .paired_x_inverse_plan_safety import require_single_kernel +from .dense_half_inverse_schedule import twiddles + + +class Plan: + """Prepare every scalar alignment before enqueue-only caller execution.""" + def __init__(self, shape, dtype, direction='c2r', *, device=None, stream=None, + variant='xy_z', dispatch='native'): + 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, even nx >= 16)') + if dtype not in (torch.float32, torch.float64) or direction != 'c2r': + raise ValueError('Expected a float32 or float64 inverse transform') + if variant not in ('xy_z', 'xy_z_patient') or dispatch != 'native': + raise ValueError('Expected a native XY/Z decomposition') + self.dtype, self.direction, self.variant, self.dispatch = dtype, direction, variant, dispatch + 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.closed, self.handles, self.chains = False, [], [] + self.compiled_metadata, self.final_plan_checks = [], [] + nz, ny, nx = self.shape + m, nf = nx // 2, nx // 2 + 1 + hz, hy = nz, ny // 2 + if ny % 2: + raise ValueError("The paired axis must be even") + self.input_shape, self.output_shape = (nz, ny, nf, 2), self.shape + patient = variant.endswith('_patient') + try: + with torch.cuda.device(self.device), torch.cuda.stream(self.stream): + self.intermediate = torch.empty((hz, hy, nx, 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] + loads, stores = [], [] + for residue in range(self.residues): + load = compile_callback(self.shape, residue, 'load', arch, + double=dtype == torch.float64) + first = _Handle((hy, nx), 0x69, (hy, nx), 1, hy * nx, + (hy, nx), 1, hy * nx, hz, self.device, self.stream, + callback=[load], patient=patient, + caller_info=self.table.data_ptr()) + self.handles.append(first) + loads.append(first) + store = compile_callback(self.shape, residue, 'store', arch, + double=dtype == torch.float64) + final = _Handle((hz,), 0x69, (hz,), hy * nx, 1, + (hz,), hy * nx, 1, hy * nx, self.device, self.stream, + callback=[store], patient=patient) + self.handles.append(final) + stores.append(final) + self.compiled_metadata.extend((load[3], store[3])) + # Inspect all final plans using full nominal FP64 storage. + # Multi-kernel cuFFT may use output as an intermediate and is + # unsuitable for scalar-offset or narrower caller storage. + nominal = torch.empty(math.prod(self.shape) * 8 + 16, + dtype=torch.uint8, device=self.device) + for final in stores: + self.final_plan_checks.append(require_single_kernel( + final, self.intermediate, nominal, self.device, self.stream)) + for first in loads: + for final in stores: + self.chains.append(PreparedChain([ + (_lib.cufftExecZ2Z, first.value, 1, 1, 0, 2), + (_lib.cufftExecZ2Z, final.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) + dummy = torch.zeros(math.prod(self.input_shape) + self.residues, + dtype=dtype, device=self.device) + for i in range(self.residues): + for o in range(self.residues): + self.execute_ptr(dummy.data_ptr() + self.itemsize * i, + nominal.data_ptr() + self.itemsize * o, 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() * 8 + self.table.numel() * 8) + + def execute_ptr(self, input_ptr, output_ptr, stream_ptr): + """Enqueue with new caller addresses using only prepared state.""" + 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.itemsize or b % self.itemsize: + raise ValueError('Distinct non-null scalar-aligned pointers are required') + 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 wait for outstanding uses before closing.""" + 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() diff --git a/build/torch-cuda/_b200/inverse_paired_y_f64.py b/build/torch-cuda/_b200/inverse_paired_y_f64.py new file mode 100644 index 0000000000000000000000000000000000000000..cf564a3df5c75128a4adbeebb7237c54a058fc89 --- /dev/null +++ b/build/torch-cuda/_b200/inverse_paired_y_f64.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FP64 inverse with real samples paired along the Y axis.""" +from __future__ import annotations + +import math +import threading + +import torch + +from .cuda_z_inverse_cufft_support import _Handle, _lib, _check +from .c2r_f64_paired_y_direct_inverse_lto import compile_callback +from .direct_dispatch import PreparedChain +from .paired_x_inverse_plan_safety import require_single_kernel +from .dense_half_inverse_schedule import twiddles + + +class Plan: + """Prepare every scalar alignment before enqueue-only caller execution.""" + def __init__(self, shape, dtype, direction='c2r', *, device=None, stream=None, + variant='xy_z', dispatch='native'): + 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, even nx >= 16)') + if dtype not in (torch.float32, torch.float64) or direction != 'c2r': + raise ValueError('Expected a float32 or float64 inverse transform') + if variant not in ('xy_z', 'xy_z_patient') or dispatch != 'native': + raise ValueError('Expected a native XY/Z decomposition') + self.dtype, self.direction, self.variant, self.dispatch = dtype, direction, variant, dispatch + 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.closed, self.handles, self.chains = False, [], [] + self.compiled_metadata, self.final_plan_checks = [], [] + nz, ny, nx = self.shape + m, nf = nx // 2, nx // 2 + 1 + hz, hy = nz, ny // 2 + if ny % 2: + raise ValueError("The paired axis must be even") + self.input_shape, self.output_shape = (nz, ny, nf, 2), self.shape + patient = variant.endswith('_patient') + try: + with torch.cuda.device(self.device), torch.cuda.stream(self.stream): + self.intermediate = torch.empty((hz, hy, nx, 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] + loads, stores = [], [] + for residue in range(self.residues): + load = compile_callback(self.shape, residue, 'load', arch, + double=dtype == torch.float64) + first = _Handle((hy, nx), 0x69, (hy, nx), 1, hy * nx, + (hy, nx), 1, hy * nx, hz, self.device, self.stream, + callback=[load], patient=patient, + caller_info=self.table.data_ptr()) + self.handles.append(first) + loads.append(first) + store = compile_callback(self.shape, residue, 'store', arch, + double=dtype == torch.float64) + final = _Handle((hz,), 0x69, (hz,), hy * nx, 1, + (hz,), hy * nx, 1, hy * nx, self.device, self.stream, + callback=[store], patient=patient) + self.handles.append(final) + stores.append(final) + self.compiled_metadata.extend((load[3], store[3])) + # Inspect all final plans using full nominal FP64 storage. + # Multi-kernel cuFFT may use output as an intermediate and is + # unsuitable for scalar-offset or narrower caller storage. + nominal = torch.empty(math.prod(self.shape) * 8 + 16, + dtype=torch.uint8, device=self.device) + for final in stores: + self.final_plan_checks.append(require_single_kernel( + final, self.intermediate, nominal, self.device, self.stream)) + for first in loads: + for final in stores: + self.chains.append(PreparedChain([ + (_lib.cufftExecZ2Z, first.value, 1, 1, 0, 2), + (_lib.cufftExecZ2Z, final.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) + dummy = torch.zeros(math.prod(self.input_shape) + self.residues, + dtype=dtype, device=self.device) + for i in range(self.residues): + for o in range(self.residues): + self.execute_ptr(dummy.data_ptr() + self.itemsize * i, + nominal.data_ptr() + self.itemsize * o, 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() * 8 + self.table.numel() * 8) + + def execute_ptr(self, input_ptr, output_ptr, stream_ptr): + """Enqueue with new caller addresses using only prepared state.""" + 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.itemsize or b % self.itemsize: + raise ValueError('Distinct non-null scalar-aligned pointers are required') + 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 wait for outstanding uses before closing.""" + 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() diff --git a/build/torch-cuda/_b200/inverse_partial_f32.py b/build/torch-cuda/_b200/inverse_partial_f32.py new file mode 100644 index 0000000000000000000000000000000000000000..cee3be004cef3ca946fb19f5993fed1d45940780 --- /dev/null +++ b/build/torch-cuda/_b200/inverse_partial_f32.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FP64 half-length inverse with a checked direct scalar Z output stage.""" +from __future__ import annotations + +import math +import threading + +import torch + +from .cuda_z_inverse_cufft_support import _Handle, _lib, _check +from .partial_y_inverse_lto import compile_callback +from .direct_dispatch import PreparedChain +from .paired_x_inverse_plan_safety import require_single_kernel +from .dense_half_inverse_schedule import twiddles + + +class Plan: + """Prepare every scalar alignment before enqueue-only caller execution.""" + def __init__(self, shape, dtype, direction='c2r', *, device=None, stream=None, + variant='xy_z', dispatch='native'): + 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, even nx >= 16)') + if dtype not in (torch.float32, torch.float64) or direction != 'c2r': + raise ValueError('Expected a float32 or float64 inverse transform') + if variant not in ('xy_z', 'xy_z_patient') or dispatch != 'native': + raise ValueError('Expected a native XY/Z decomposition') + self.dtype, self.direction, self.variant, self.dispatch = dtype, direction, variant, dispatch + 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.closed, self.handles, self.chains = False, [], [] + self.compiled_metadata, self.final_plan_checks = [], [] + nz, ny, nx = self.shape + m, nf = nx // 2, nx // 2 + 1 + radix = 2 + if ny % radix: + raise ValueError("Y length must be divisible by radix") + short_y = ny // radix + self.input_shape, self.output_shape = (nz, ny, nf, 2), self.shape + patient = variant.endswith('_patient') + try: + with torch.cuda.device(self.device), torch.cuda.stream(self.stream): + self.intermediate = torch.empty((nz, ny, m, 2), dtype=torch.float64, + device=self.device) + self.table = torch.tensor(twiddles(nx) + twiddles(ny), dtype=torch.float64, device=self.device) + cap = torch.cuda.get_device_capability(self.device) + arch = 10 * cap[0] + cap[1] + loads, stores = [], [] + for residue in range(self.residues): + load = compile_callback(self.shape, residue, 'load', arch, + double=dtype == torch.float64) + first = _Handle((short_y, m), 0x69, (short_y, m), 1, short_y * m, + (short_y, m), 1, short_y * m, nz * radix, self.device, self.stream, + callback=[load], patient=patient, + caller_info=self.table.data_ptr()) + self.handles.append(first) + loads.append(first) + store = compile_callback(self.shape, residue, 'store', arch, + double=dtype == torch.float64) + finish = compile_callback(self.shape, 0, 'finish', arch, double=dtype == torch.float64) + final = _Handle((nz,), 0x69, (nz,), ny * m, 1, + (nz,), ny * m, 1, ny * m, self.device, self.stream, + callback=[finish, store], patient=patient, caller_info=self.table.data_ptr()) + self.handles.append(final) + stores.append(final) + self.compiled_metadata.extend((load[3], finish[3], store[3])) + # Inspect all final plans using full nominal FP64 storage. + # Multi-kernel cuFFT may use output as an intermediate and is + # unsuitable for scalar-offset or narrower caller storage. + nominal = torch.empty(math.prod(self.shape) * 8 + 16, + dtype=torch.uint8, device=self.device) + for final in stores: + self.final_plan_checks.append(require_single_kernel( + final, self.intermediate, nominal, self.device, self.stream)) + for first in loads: + for final in stores: + self.chains.append(PreparedChain([ + (_lib.cufftExecZ2Z, first.value, 1, 1, 0, 2), + (_lib.cufftExecZ2Z, final.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) + dummy = torch.zeros(math.prod(self.input_shape) + self.residues, + dtype=dtype, device=self.device) + for i in range(self.residues): + for o in range(self.residues): + self.execute_ptr(dummy.data_ptr() + self.itemsize * i, + nominal.data_ptr() + self.itemsize * o, 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() * 8 + self.table.numel() * 8) + + def execute_ptr(self, input_ptr, output_ptr, stream_ptr): + """Enqueue with new caller addresses using only prepared state.""" + 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.itemsize or b % self.itemsize: + raise ValueError('Distinct non-null scalar-aligned pointers are required') + 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 wait for outstanding uses before closing.""" + 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() diff --git a/build/torch-cuda/_b200/inverse_pitched.py b/build/torch-cuda/_b200/inverse_pitched.py new file mode 100644 index 0000000000000000000000000000000000000000..25a5a87929335cd0d8587bf65654bb3558b226cc --- /dev/null +++ b/build/torch-cuda/_b200/inverse_pitched.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FP64 half-length inverse with a checked direct scalar Z output stage.""" +from __future__ import annotations + +import math +import threading + +import torch + +from .cuda_z_inverse_cufft_support import _Handle, _lib, _check +from .pitched_half_inverse_lto import compile_callback +from .direct_dispatch import PreparedChain +from .paired_x_inverse_plan_safety import require_single_kernel +from .dense_half_inverse_schedule import twiddles + + +class Plan: + """Prepare every scalar alignment before enqueue-only caller execution.""" + def __init__(self, shape, dtype, direction='c2r', *, device=None, stream=None, + variant='xy_z', dispatch='native'): + 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, even nx >= 16)') + if dtype not in (torch.float32, torch.float64) or direction != 'c2r': + raise ValueError('Expected a float32 or float64 inverse transform') + if variant not in ('xy_z', 'xy_z_patient') or dispatch != 'native': + raise ValueError('Expected a native XY/Z decomposition') + self.dtype, self.direction, self.variant, self.dispatch = dtype, direction, variant, dispatch + 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.closed, self.handles, self.chains = False, [], [] + self.compiled_metadata, self.final_plan_checks = [], [] + nz, ny, nx = self.shape + m, nf = nx // 2, nx // 2 + 1 + pitch = m + 128 + self.input_shape, self.output_shape = (nz, ny, nf, 2), self.shape + patient = variant.endswith('_patient') + try: + with torch.cuda.device(self.device), torch.cuda.stream(self.stream): + self.intermediate = torch.empty((nz, ny, m, 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] + loads, stores = [], [] + for residue in range(self.residues): + load = compile_callback(self.shape, residue, 'load', arch, + double=dtype == torch.float64) + first = _Handle((ny, m), 0x69, (ny, pitch), 1, ny * pitch, + (ny, m), 1, ny * m, nz, self.device, self.stream, + callback=[load], patient=patient, + caller_info=self.table.data_ptr()) + self.handles.append(first) + loads.append(first) + store = compile_callback(self.shape, residue, 'store', arch, + double=dtype == torch.float64) + final = _Handle((nz,), 0x69, (nz,), ny * m, 1, + (nz,), ny * m, 1, ny * m, self.device, self.stream, + callback=[store], patient=patient) + self.handles.append(final) + stores.append(final) + self.compiled_metadata.extend((load[3], store[3])) + # Inspect all final plans using full nominal FP64 storage. + # Multi-kernel cuFFT may use output as an intermediate and is + # unsuitable for scalar-offset or narrower caller storage. + nominal = torch.empty(math.prod(self.shape) * 8 + 16, + dtype=torch.uint8, device=self.device) + for final in stores: + self.final_plan_checks.append(require_single_kernel( + final, self.intermediate, nominal, self.device, self.stream)) + for first in loads: + for final in stores: + self.chains.append(PreparedChain([ + (_lib.cufftExecZ2Z, first.value, 1, 1, 0, 2), + (_lib.cufftExecZ2Z, final.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) + dummy = torch.zeros(math.prod(self.input_shape) + self.residues, + dtype=dtype, device=self.device) + for i in range(self.residues): + for o in range(self.residues): + self.execute_ptr(dummy.data_ptr() + self.itemsize * i, + nominal.data_ptr() + self.itemsize * o, 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() * 8 + self.table.numel() * 8) + + def execute_ptr(self, input_ptr, output_ptr, stream_ptr): + """Enqueue with new caller addresses using only prepared state.""" + 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.itemsize or b % self.itemsize: + raise ValueError('Distinct non-null scalar-aligned pointers are required') + 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 wait for outstanding uses before closing.""" + 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() diff --git a/build/torch-cuda/_b200/inverse_scalar.py b/build/torch-cuda/_b200/inverse_scalar.py new file mode 100644 index 0000000000000000000000000000000000000000..d338017e8d2899451557cc1a3ebca73e68dbf08e --- /dev/null +++ b/build/torch-cuda/_b200/inverse_scalar.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Add prepared scalar-input support while retaining measured aligned paths.""" +from . import scalar_input_inverse_backend as scalar_backend + + +class ScalarInputPlan: + def __init__(self, aligned): + self._aligned, self._scalar, self.closed = aligned, None, False + try: + for name in ('shape','dtype','direction','device','stream','stream_ptr','input_shape','output_shape'): + setattr(self,name,getattr(aligned,name)) + self.itemsize = 8 if str(self.dtype)=='torch.float64' else 4 + self.variant = getattr(aligned,'variant','aligned')+'_scalar_input' + self.dispatch = getattr(aligned,'dispatch','prepared') + self._scalar = scalar_backend.Plan(self.shape,self.dtype,self.direction, + device=self.device,stream=self.stream) + 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): + """Use the prepared branch for the caller's natural scalar alignment.""" + 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') + selected=self._scalar if a%(2*self.itemsize) else self._aligned + selected.execute_ptr(a,b,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() + na,nb=input.numel()*self.itemsize,output.numel()*self.itemsize + if a= 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)) + if inverse: + add(complex_axis, (triton.cdiv(inner, columns_per_block), outer, 1), + input_slot, output_slot, length, constants) + else: + add(combined_z, (triton.cdiv(inner, columns_per_block), outer, 1), + input_slot, output_slot, (length, ny), constants + (ny, nx, 2)) + + def xy(input_slot, output_slot): + rx, ry = factors(nx), factors(ny // 2) + constants = (ny, nx, triton.next_power_of_2(ny // 2), triton.next_power_of_2(nx), + inverse, rx, stage_sizes(rx), ry, stage_sizes(ry), 2) + add(real_xy, (nz, 2, 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() diff --git a/build/torch-cuda/_b200/inverse_small_f64.py b/build/torch-cuda/_b200/inverse_small_f64.py new file mode 100644 index 0000000000000000000000000000000000000000..ef047ca0e98d27787576e30b81ee176f403f8328 --- /dev/null +++ b/build/torch-cuda/_b200/inverse_small_f64.py @@ -0,0 +1,203 @@ +# 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 .grouped_dif_inverse_kernels import combined_z, complex_axis, packed_real_axis, real_axis, real_xy +from .grouped_dif_inverse_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 != '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') + if direction != 'c2r' or self.shape != (32, 32, 32): + raise ValueError('This prepared implementation handles the 32-cubed inverse') + 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') + if self.shape != (32, 32, 32): + raise ValueError('This experimental composition requires 32 cubed') + 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)) + if inverse: + add(complex_axis, (triton.cdiv(inner, columns_per_block), outer, 1), + input_slot, output_slot, length, constants) + else: + add(combined_z, (triton.cdiv(inner, columns_per_block), outer, 1), + input_slot, output_slot, (length, ny), constants + (ny, nx, 4)) + + def xy(input_slot, output_slot): + rx, ry = factors(nx // 2), factors(ny // 4) + constants = (ny, nx, triton.next_power_of_2(ny // 4), triton.next_power_of_2(nx), + inverse, rx, stage_sizes(rx), ry, stage_sizes(ry), 4) + add(real_xy, (nz, 4, 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() diff --git a/build/torch-cuda/_b200/inverse_vector.py b/build/torch-cuda/_b200/inverse_vector.py new file mode 100644 index 0000000000000000000000000000000000000000..a8a94d972e2426838b7dcfd8f8b4930b287832a8 --- /dev/null +++ b/build/torch-cuda/_b200/inverse_vector.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Half-length FP64 inverse with direct aligned and redirected scalar output.""" +from __future__ import annotations + +import math +import threading + +import torch +import triton +from triton.backends.compiler import GPUTarget +from triton.compiler import ASTSource + +from .cuda_z_inverse_cufft_support import _Handle, _check, _lib +from .dense_half_inverse_kernels import publish_output +from .half_vector_inverse_lto import compile_callback +from .dense_half_inverse_native_dispatch import PreparedDispatch +from .direct_dispatch import PreparedChain +from .dense_half_inverse_schedule import twiddles + + +class Plan: + def __init__(self, shape, dtype, direction='c2r', *, device=None, stream=None, + variant='3d', dispatch='native'): + 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 != torch.float64 or direction != 'c2r': + raise ValueError('Expected an FP64 inverse transform') + if variant not in ('3d', '3d_patient', 'xy_z', 'xy_z_patient'): + raise ValueError('Unknown variant') + if dispatch not in ('native', 'python'): + raise ValueError('Unknown dispatch mode') + self.dtype, self.direction, self.variant = dtype, direction, variant + self.dispatch = dispatch + 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.direct_handles = [] + self.direct_native = [] + self.z_handle = None + self.compiled_metadata = [] + nz, ny, nx = self.shape + nf = nx // 2 + 1 + self.input_shape, self.output_shape = (nz, ny, nf, 2), self.shape + m = nx // 2 + self.fft_function = self.z_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] + self.table = torch.tensor(twiddles(nx), dtype=torch.float64, device=self.device) + self.slot = torch.tensor([0, self.table.data_ptr()], dtype=torch.int64, device=self.device) + # cuFFT may alternate output and workspace during a transform. + # Allocate the entire FP64 extent even though the final callback + # writes to caller FP64 storage. No raw input is copied here. + self.scratch = torch.empty((nz, ny, m, 2), dtype=torch.float64, + device=self.device) + self.slot_ptr, self.scratch_ptr = self.slot.data_ptr(), self.scratch.data_ptr() + store_callback = compile_callback(self.shape, 0, 'store', arch) + self.compiled_metadata.append(store_callback[3]) + patient = variant.endswith('_patient') + for residue in range(2): + load_callback = compile_callback(self.shape, residue, 'load', arch) + self.compiled_metadata.append(dict(load_callback[3], input_residue=residue)) + if variant.startswith('3d'): + handle = _Handle( + (nz, ny, m), 0x69, (nz, ny, nf), 1, nz * ny * nf, + (nz, ny, m), 1, nz * ny * m, 1, self.device, self.stream, + callback=[load_callback, store_callback], patient=patient, + caller_info=self.slot_ptr) + else: + handle = _Handle( + (ny, m), 0x69, (ny, nf), 1, ny * nf, + (ny, m), 1, ny * m, nz, self.device, self.stream, + callback=[load_callback], patient=patient, caller_info=self.slot_ptr) + self.handles.append(handle) + self.handle_values = tuple(h.value for h in self.handles) + if variant.startswith('xy_z'): + self.z_handle = _Handle( + (nz,), 0x69, (nz,), ny * m, 1, (nz,), ny * m, 1, + ny * m, self.device, self.stream, callback=[store_callback], + patient=patient, caller_info=self.slot_ptr) + self.z_handle_value = self.z_handle.value + # Complex-aligned caller output can hold every cuFFT stage. + # Scalar-offset output uses the fully prepared redirected route. + for residue in range(2): + load_callback = compile_callback(self.shape, residue, 'load', arch) + if variant.startswith('3d'): + direct = _Handle( + (nz, ny, m), 0x69, (nz, ny, nf), 1, nz * ny * nf, + (nz, ny, m), 1, nz * ny * m, 1, self.device, self.stream, + callback=[load_callback], patient=patient, + caller_info=self.slot_ptr) + else: + direct = _Handle( + (ny, m), 0x69, (ny, nf), 1, ny * nf, + (ny, m), 1, ny * m, nz, self.device, self.stream, + callback=[load_callback], patient=patient, + caller_info=self.slot_ptr) + self.direct_handles.append(direct) + final_direct = None + if variant.startswith('xy_z'): + final_direct = _Handle( + (nz,), 0x69, (nz,), ny * m, 1, (nz,), ny * m, 1, + ny * m, self.device, self.stream, patient=patient) + self.direct_handles.append(final_direct) + for first in self.direct_handles[:2]: + stages = [(self.fft_function, first.value, 1, 1, 0, 1)] + if final_direct is not None: + stages.append((self.z_function, final_direct.value, 1, 1, 1, 1)) + self.direct_native.append(PreparedChain(stages, 0)) + self.direct_function = self.direct_native[0].function + self.direct_pointers = tuple(p.pointer for p in self.direct_native) + source = ASTSource(publish_output, {'slot': '*i64', 'output_address': 'u64'}) + compiled = triton.compile(source, target=GPUTarget('cuda', arch, 32), + options={'num_warps': 1}) + for name in ('global_scratch_size', 'profile_scratch_size'): + if getattr(compiled.metadata, name, 0): + raise ValueError('Execution scratch is unsupported: ' + name) + self.compiled_kernel = compiled + self.runner = compiled[(1, 1, 1)] + if dispatch == 'native': + for handle in self.handle_values: + self.native.append(PreparedDispatch( + self.fft_function, handle, self.z_function, + self.z_handle_value if self.z_handle is not None else 0, + compiled, self.slot_ptr, self.scratch_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.metadata.shared)) + dummy = torch.zeros(math.prod(self.input_shape) + 4, dtype=dtype, + device=self.device) + output = torch.empty(math.prod(self.output_shape) + 1, dtype=dtype, device=self.device) + for residue in range(2): + for output_residue in range(2): + self.execute_ptr(dummy.data_ptr() + residue * 8, + output.data_ptr() + output_residue * 8, self.stream_ptr) + self.stream.synchronize() + except BaseException: + for handle in self.handles + self.direct_handles: + handle.close() + if self.z_handle is not None: + self.z_handle.close() + raise + + @property + def workspace_bytes(self): + return (sum(h.workspace_bytes for h in self.handles + self.direct_handles) + + (self.z_handle.workspace_bytes if self.z_handle is not None else 0) + + self.scratch.numel() * 8 + 16 + self.table.numel() * 8) + + def execute_ptr(self, input_ptr, output_ptr, stream_ptr): + """Enqueue prepared FFTs for the public scalar alignment metadata.""" + 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 % 8 or b % 8: + raise ValueError('Distinct non-null scalar-aligned pointers are required') + residue = (a & 15) // 8 + if b % 16 == 0: + _check(self.direct_function(self.direct_pointers[residue], a - 8 * residue, b)) + return + if self.dispatch == 'native': + result = self.native_function(self.native_pointers[residue], a - 8 * residue, b) + if result > 0: + _check(result) + if result < 0: + raise RuntimeError(f'cuLaunchKernel failed with CUDA driver error {-result}') + return + self.runner(self.slot_ptr, b, stream=self.stream_ptr) + _check(self.fft_function(self.handle_values[residue], a - 8 * residue, self.scratch_ptr, 1)) + if self.z_handle is not None: + _check(self.z_function(self.z_handle_value, self.scratch_ptr, self.scratch_ptr, 1)) + + 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 + self.direct_handles: + handle.close() + if self.z_handle is not None: + self.z_handle.close() + self.scratch = self.slot = self.table = self.runner = self.compiled_kernel = None + self.native = self.native_pointers = self.direct_native = self.direct_pointers = None + self.closed = True + + + + +_CACHE = {} +_LOCK = threading.Lock() diff --git a/build/torch-cuda/_b200/narrow_index_inverse_cufft_support.py b/build/torch-cuda/_b200/narrow_index_inverse_cufft_support.py new file mode 100644 index 0000000000000000000000000000000000000000..4cb966041cab341a95565e1f04e4a198bb9a894c --- /dev/null +++ b/build/torch-cuda/_b200/narrow_index_inverse_cufft_support.py @@ -0,0 +1,119 @@ +# 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], + "cufftXtMakePlanMany": [ct.c_int, ct.c_int, ct.POINTER(ct.c_longlong), + ct.POINTER(ct.c_longlong), ct.c_longlong, ct.c_longlong, ct.c_int, + ct.POINTER(ct.c_longlong), ct.c_longlong, ct.c_longlong, ct.c_int, + ct.c_longlong, _size_p, ct.c_int], + "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, xt=False): + 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)) + for item in (() if callback is None else callback): + symbol, callback_type, blob, _metadata = item + 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) + if xt: + wide = ct.c_longlong * len(sizes) + # This experiment uses complex double for both stages. + if kind != 0x69: + raise ValueError('Only Z2Z is supported by this Xt wrapper') + _check(_lib.cufftXtMakePlanMany( + self.value, len(sizes), wide(*sizes), wide(*inembed), + istride, idist, 5, wide(*onembed), ostride, odist, + 5, batch, ct.byref(size), 5)) + else: + _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 diff --git a/build/torch-cuda/_b200/narrow_index_inverse_lto.py b/build/torch-cuda/_b200/narrow_index_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..69f556fb3de96a553769291fbdae324c9b1bc499 --- /dev/null +++ b/build/torch-cuda/_b200/narrow_index_inverse_lto.py @@ -0,0 +1,107 @@ +# 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, nz, ny, n, arch, nvrtc_version): + lib = library() + options = ('--std=c++17', '--relocatable-device-code=true', '-dlto', + f'--gpu-architecture=compute_{arch}', f'-DFFT_N={n}', + f'-DFFT_NZ={nz}', f'-DFFT_NY={ny}') + program = ct.c_void_p() + _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'reconstruct.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=(nz,ny,n), 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_callbacks(shape, arch=100): + """Compile the reconstruction load callback for the public real shape.""" + nz,ny,n = map(int,shape) + if min(nz,ny) < 8 or n < 16 or n % 2: + raise ValueError('Invalid transform shape') + source = Path(__file__).with_name('reconstruct.cu').read_bytes() + with _LOCK: + blob,metadata = _compile(source,nz,ny,n,int(arch),version()) + return [(b'fft_full_half_load',1,ct.create_string_buffer(blob,len(blob)),metadata)] diff --git a/build/torch-cuda/_b200/pair_load.cu b/build/torch-cuda/_b200/pair_load.cu new file mode 100644 index 0000000000000000000000000000000000000000..f5383e4c62234ecb63570b19b871b3a16ece3cf6 --- /dev/null +++ b/build/torch-cuda/_b200/pair_load.cu @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: Apache-2.0 +// Read adjacent even/odd caller samples as one complex FFT operand. +// Each precompiled shift handles one scalar residue of a 16-byte address. +__device__ double2 fft_pair_load(void *data, unsigned long long offset, + void *caller_info, void *shared_memory) { +#if FFT_DOUBLE + const double *input = static_cast(data); +#else + const float *input = static_cast(data); +#endif + double2 value; + value.x = static_cast(input[2 * offset + FFT_INPUT_SHIFT]); + value.y = static_cast(input[2 * offset + FFT_INPUT_SHIFT + 1]); + return value; +} diff --git a/build/torch-cuda/_b200/paired_vector_store_forward_abi.py b/build/torch-cuda/_b200/paired_vector_store_forward_abi.py new file mode 100644 index 0000000000000000000000000000000000000000..bda49188db1bcaadfdec3aef8f4afe71f6dd47da --- /dev/null +++ b/build/torch-cuda/_b200/paired_vector_store_forward_abi.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reject kernel ABIs that the small native launcher cannot represent.""" +from __future__ import annotations + +import re + + +def validate_kernel_abi(compiled): + """Inspect compiler output on the host; this function does not load CUDA.""" + 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 not in (1, 2, 4, 8, 16, 32): + raise ValueError('Unsupported warp count') + signature = compiled.src.signature + runtime = [value for value in signature.values() if value != 'constexpr'] + if runtime not in (['*fp64'] * 3, ['*fp64', '*fp32', '*fp64']): + raise ValueError('Expected exactly FP64 intermediate/table and FP32 or FP64 output') + ptx = compiled.asm['ptx'] + match = re.search(r'\.entry\s+' + re.escape(metadata.name) + r'\s*\((.*?)\)', + ptx, flags=re.DOTALL) + if match is None: + raise ValueError('Kernel entry signature was not found in compiler PTX') + parameters = [part.strip() for part in match.group(1).split(',')] + # Current installed Triton appends global and profiling scratch pointers. + # Require this precise ABI instead of guessing across compiler versions. + if len(parameters) != 5 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 five scalar 64-bit pointer parameters') + return dict(parameter_count=5, tensor_arguments=3, null_scratch_arguments=2, + num_ctas=1, threads=32 * metadata.num_warps, + shared_bytes=metadata.shared) diff --git a/build/torch-cuda/_b200/paired_vector_store_forward_backend.py b/build/torch-cuda/_b200/paired_vector_store_forward_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..bb6c9e9a0e4be35d3397e87e9423f532c5f593cd --- /dev/null +++ b/build/torch-cuda/_b200/paired_vector_store_forward_backend.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Prepared half-length complex forward FFT with direct caller packed output. + +All arithmetic is FP64; caller precision is preserved at the boundaries. +The intermediate holds transformed data, never a repacked raw input. +""" +import threading +import torch +import triton + +from .narrow_index_inverse_cufft_support import _Handle, _lib, _check +from .paired_vector_store_forward_kernels import recombine +from .paired_vector_store_forward_lto import compile_callback +from .paired_vector_store_forward_native_dispatch import PreparedDispatch +from .dense_half_inverse_schedule import factors, twiddles + + +class Plan: + def __init__(self, shape, dtype, direction='r2c', *, device=None, stream=None, + block=256, patient=False): + 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, even nx >= 16)') + if dtype not in (torch.float32, torch.float64) or direction != 'r2c': + raise ValueError('Only FP32/FP64 forward transforms are supported') + if block not in (128, 256, 512): + raise ValueError('Unsupported reconstruction tile') + for n in self.shape: + factors(n) + 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('Stream and plan device must match') + self.stream_ptr = self.stream.cuda_stream + self.itemsize = 8 if dtype == torch.float64 else 4 + self.closed, self.handles, self.native = False, [], [] + self.variant = f'paired_vector_store_forward_b{block}_p{int(patient)}' + nz, ny, nx = self.shape + m, rows = nx // 2, nz * ny + self.input_shape, self.output_shape = self.shape, (nz, ny, m + 1, 2) + self.grid = (triton.cdiv(rows * (m // 2 + 1), block), 1, 1) + try: + with torch.cuda.device(self.device), torch.cuda.stream(self.stream): + self.intermediate = torch.empty((nz, ny, m, 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) + dummy_in = torch.zeros(rows * nx + 3, dtype=dtype, device=self.device) + dummy_out = torch.empty(rows * (m + 1) * 2 + 1, dtype=dtype, device=self.device) + self.compiled_metadata = [] + for shift in range(16 // self.itemsize): + callback = compile_callback(dtype == torch.float64, shift, 10 * cap[0] + cap[1]) + self.compiled_metadata.append(callback[3]) + handle = _Handle((nz, ny, m), 0x69, (nz, ny, m), 1, rows * m, + (nz, ny, m), 1, rows * m, 1, self.device, self.stream, + callback=[callback], patient=bool(patient)) + self.handles.append(handle) + _check(_lib.cufftExecZ2Z(self.handles[0].value, dummy_in.data_ptr(), + self.intermediate.data_ptr(), -1)) + self.compiled = recombine[self.grid](self.intermediate, dummy_out, self.table, + nz, ny, nx, block, num_warps=4, enable_fp_fusion=False) + for handle in self.handles: + self.native.append(PreparedDispatch(_lib.cufftExecZ2Z, handle.value, self.compiled, + self.grid, self.intermediate.data_ptr(), self.table.data_ptr(), self.stream_ptr)) + self.native_function = self.native[0].function + self.native_pointers = tuple(n.pointer for n in self.native) + for shift in range(16 // self.itemsize): + for out_shift in (0, 1): + self.execute_ptr(dummy_in.data_ptr() + shift * self.itemsize, + dummy_out.data_ptr() + out_shift * self.itemsize, self.stream_ptr) + self.stream.synchronize() + except BaseException: + for handle in reversed(self.handles): + handle.close() + raise + + @property + def workspace_bytes(self): + return sum(h.workspace_bytes for h in self.handles) + 8 * (self.intermediate.numel() + self.table.numel()) + + def execute_ptr(self, input_ptr, output_ptr, stream_ptr): + 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) % self.itemsize or int(output_ptr) % self.itemsize: + raise ValueError('Caller buffers must be scalar aligned') + residue = int(input_ptr) % 16 + result = self.native_function(self.native_pointers[residue // self.itemsize], + int(input_ptr) - residue, output_ptr) + 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() + if a < b + output.numel() * self.itemsize and b < a + input.numel() * self.itemsize: + 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): + if not self.closed: + for handle in reversed(self.handles): + handle.close() + self.intermediate = self.table = self.compiled = self.native = self.native_pointers = None + self.closed = True + + + + +_CACHE = {} +_LOCK = threading.Lock() diff --git a/build/torch-cuda/_b200/paired_vector_store_forward_kernels.py b/build/torch-cuda/_b200/paired_vector_store_forward_kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..97c0f601d39f8eae581e3d9459b66b9c14db0f8a --- /dev/null +++ b/build/torch-cuda/_b200/paired_vector_store_forward_kernels.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Use explicit complex loads and alignment-aware stores for paired reconstruction. + +All vector loads address private FP64 transformed data or twiddle tables. +Naturally aligned caller outputs use complex vector stores; scalar-offset +outputs use predicated scalar stores in the same prepared kernel. +""" +import triton +import triton.language as tl + + +@triton.jit +def store_pair(pointer, first, second, valid): + """Store two adjacent scalars, checking vector alignment at execution. + + The caller can change to a scalar-offset output between launches. Every + active pair uses either one naturally aligned vector store or two scalar + stores. Inactive pairs access no memory. The assembly is side-effecting. + """ + if pointer.dtype.element_ty == tl.float64: + tl.inline_asm_elementwise( + "{ .reg .pred p, a, v, s; .reg .b64 rem; " + "setp.ne.u32 p, $4, 0; and.b64 rem, $1, 15; setp.eq.u64 a, rem, 0; " + "and.pred v, p, a; not.pred a, a; and.pred s, p, a; " + "@v st.global.v2.f64 [$1], {$2, $3}; " + "@s st.global.f64 [$1], $2; @s st.global.f64 [$1+8], $3; " + "mov.u32 $0, 0; }", + constraints="=r,l,d,d,r", args=[pointer.to(tl.uint64), first, second, valid.to(tl.int32)], + dtype=tl.int32, is_pure=False, pack=1) + else: + tl.inline_asm_elementwise( + "{ .reg .pred p, a, v, s; .reg .b64 rem; " + "setp.ne.u32 p, $4, 0; and.b64 rem, $1, 7; setp.eq.u64 a, rem, 0; " + "and.pred v, p, a; not.pred a, a; and.pred s, p, a; " + "@v st.global.v2.f32 [$1], {$2, $3}; " + "@s st.global.f32 [$1], $2; @s st.global.f32 [$1+4], $3; " + "mov.u32 $0, 0; }", + constraints="=r,l,f,f,r", args=[pointer.to(tl.uint64), first, second, valid.to(tl.int32)], + dtype=tl.int32, is_pure=False, pack=1) + + +@triton.jit +def load_complex(pointer, valid): + """Predicated 16-byte read of an aligned private complex FP64 value.""" + return tl.inline_asm_elementwise( + "{ .reg .pred p; setp.ne.u32 p, $3, 0; mov.b64 $0, 0; mov.b64 $1, 0; @p ld.global.v2.f64 {$0, $1}, [$2]; }", + constraints="=d,=d,l,r", args=[pointer, valid.to(tl.int32)], + dtype=(tl.float64, tl.float64), is_pure=False, pack=1) + + +@triton.jit(do_not_specialize_on_alignment=["Output"]) +def recombine(Input, Output, Twiddle, NZ: tl.constexpr, NY: tl.constexpr, + NX: tl.constexpr, BLOCK: tl.constexpr): + m: tl.constexpr = NX // 2 + nf: tl.constexpr = m + 1 + pairs: tl.constexpr = m // 2 + 1 + index = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = index < NZ * NY * pairs + kx = index % pairs + row = index // pairs + ky = row % NY + kz = row // NY + partner_row = ((NZ - kz) % NZ) * NY + (NY - ky) % NY + direct = row * m + kx + mirror = partner_row * m + (m - kx) % m + ar, ai = load_complex(Input + 2 * direct, mask) + br, raw_bi = load_complex(Input + 2 * mirror, mask) + bi = -raw_bi + wr, wi = load_complex(Twiddle + 2 * kx, tl.full((BLOCK,), True, tl.int1)) + even_r, even_i = 0.5 * (ar + br), 0.5 * (ai + bi) + odd_r, odd_i = 0.5 * (ai - bi), -0.5 * (ar - br) + rotated_r = wr * odd_r - wi * odd_i + rotated_i = wr * odd_i + wi * odd_r + output_index = row * nf + kx + store_pair(Output + 2 * output_index, + (even_r + rotated_r).to(Output.dtype.element_ty), + (even_i + rotated_i).to(Output.dtype.element_ty), mask) + # X(-z,-y,m-k) = conj(E-W*O) for 0 < k < m/2. + # k=0 instead emits both endpoints in the same row: E+O and E-O. + # The self-paired middle column is evaluated once per row without a second store. + second_row = tl.where(kx == 0, row, partner_row) + second_index = second_row * nf + m - kx + second_mask = mask & (2 * kx < m) + second_i = tl.where(kx == 0, even_i - rotated_i, rotated_i - even_i) + store_pair(Output + 2 * second_index, + (even_r - rotated_r).to(Output.dtype.element_ty), + second_i.to(Output.dtype.element_ty), second_mask) diff --git a/build/torch-cuda/_b200/paired_vector_store_forward_lto.py b/build/torch-cuda/_b200/paired_vector_store_forward_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..5cb836b8d5dcd315d50589abac45b02514b627f7 --- /dev/null +++ b/build/torch-cuda/_b200/paired_vector_store_forward_lto.py @@ -0,0 +1,105 @@ +# 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, nf, double, arch, nvrtc_version, input_shift): + lib = library() + options = ('--std=c++17', '--relocatable-device-code=true', '-dlto', + f'--gpu-architecture=compute_{arch}', f'-DFFT_NF={nf}', + f'-DFFT_DOUBLE={int(double)}', f'-DFFT_INPUT_SHIFT={input_shift}') + 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(nf=nf, 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(input_double, input_shift, arch=100): + """Load paired caller samples directly as FP64 complex FFT operands.""" + source = Path(__file__).with_name('pair_load.cu').read_bytes() + with _LOCK: + blob, metadata = _compile(source, 0, bool(input_double), int(arch), version(), input_shift) + return (b'fft_pair_load', 1, ct.create_string_buffer(blob, len(blob)), + dict(metadata, kind='caller real pairs to FP64 complex load')) diff --git a/build/torch-cuda/_b200/paired_vector_store_forward_native_dispatch.c b/build/torch-cuda/_b200/paired_vector_store_forward_native_dispatch.c new file mode 100644 index 0000000000000000000000000000000000000000..15ccddd3e1726017e3c36891031bb65b877aec94 --- /dev/null +++ b/build/torch-cuda/_b200/paired_vector_store_forward_native_dispatch.c @@ -0,0 +1,42 @@ +/* SPDX-License-Identifier: Apache-2.0 */ +/* Enqueue the forward half-length cuFFT and packed-spectrum reconstruction 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 +#include + +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); } diff --git a/build/torch-cuda/_b200/paired_vector_store_forward_native_dispatch.py b/build/torch-cuda/_b200/paired_vector_store_forward_native_dispatch.py new file mode 100644 index 0000000000000000000000000000000000000000..f23d78b202b92096c4db1dfb8de88fd564f9ae92 --- /dev/null +++ b/build/torch-cuda/_b200/paired_vector_store_forward_native_dispatch.py @@ -0,0 +1,103 @@ +# 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 .paired_vector_store_forward_abi import validate_kernel_abi + + +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('paired_vector_store_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-frequency-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, complex_exec, handle, compiled, grid, intermediate, twiddle, 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 < 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) diff --git a/build/torch-cuda/_b200/paired_x_inverse_lto.py b/build/torch-cuda/_b200/paired_x_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..2e25b8906c3c5f5b915151e1ccfed9ffabfde0fc --- /dev/null +++ b/build/torch-cuda/_b200/paired_x_inverse_lto.py @@ -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, frequency): + 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")}', f'-DFFT_FREQUENCY={int(frequency)}', f'-DFFT_MIDDLE={int(kind == "middle")}') + 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, frequency=False): + """Reconstruct a half-length spectrum in FP64 or store its inverse as FP32.""" + if kind not in ('load', 'store', 'middle'): + raise ValueError('Expected a load or store callback') + source = Path(__file__).with_name('c2r_f64_paired_x_inverse_wide_redirect.cu').read_bytes() + with _LOCK: + blob, metadata = _compile(source, tuple(shape), bool(double), int(arch), version(), input_shift, kind, frequency) + return (('fft_wide_' + kind).encode(), 1 if kind == 'load' else 5, + ct.create_string_buffer(blob, len(blob)), dict(metadata, kind=kind)) diff --git a/build/torch-cuda/_b200/paired_x_inverse_native_dispatch.c b/build/torch-cuda/_b200/paired_x_inverse_native_dispatch.c new file mode 100644 index 0000000000000000000000000000000000000000..f2452ff666d89031ddb05564d628506a1096e009 --- /dev/null +++ b/build/torch-cuda/_b200/paired_x_inverse_native_dispatch.c @@ -0,0 +1,50 @@ +/* SPDX-License-Identifier: Apache-2.0 */ +/* Enqueue the CUDA half-X FFT and Y/Z cuFFT 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 +#include + +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; + int middle_handle; + uint64_t middle; + uint64_t nominal; + uint64_t intermediate; + uint64_t twiddle; + uint64_t slot; + void *stream; + void *function; + unsigned int grid[3]; + unsigned int threads; + unsigned int shared; + unsigned int scalar; + int finals[4]; +}; + +int fft_execute(const struct fft_dispatch *plan, uint64_t input, uint64_t output) { + int result; + uint64_t values[5] = {input, plan->intermediate, plan->twiddle, plan->slot, output}; + 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. */ + if (result != 0) return -result; + result = plan->complex_exec(plan->middle_handle, (void *)(uintptr_t)plan->intermediate, + (void *)(uintptr_t)plan->nominal, 1); + if (result != 0) return result; + unsigned int residue = (unsigned int)((output & 15) / plan->scalar); + return plan->complex_exec(plan->finals[residue], (void *)(uintptr_t)plan->middle, + (void *)(uintptr_t)(output & ~((uint64_t)15)), 1); +} + +size_t fft_dispatch_size(void) { return sizeof(struct fft_dispatch); } diff --git a/build/torch-cuda/_b200/paired_x_inverse_native_dispatch.py b/build/torch-cuda/_b200/paired_x_inverse_native_dispatch.py new file mode 100644 index 0000000000000000000000000000000000000000..32d9335b763fe0ef70aae43d33e6f5acd9bcbcc9 --- /dev/null +++ b/build/torch-cuda/_b200/paired_x_inverse_native_dispatch.py @@ -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), ('middle_handle', ct.c_int), ('middle', ct.c_uint64), ('nominal', ct.c_uint64), ('intermediate', ct.c_uint64), + ('twiddle', ct.c_uint64), ('slot', 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), ('scalar', ct.c_uint), ('finals', ct.c_int * 4)] + + +_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('paired_x_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, slot, middle_handle, middle, nominal, finals, scalar): + 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, middle_handle=middle_handle, middle=middle, nominal=nominal, intermediate=intermediate, + twiddle=twiddle, slot=slot, stream=stream, function=compiled.function, + grid=(ct.c_uint * 3)(*grid), threads=abi['threads'], shared=abi['shared_bytes'], scalar=scalar, finals=(ct.c_int*4)(*(list(finals)+[0]*(4-len(finals))))) + 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) diff --git a/build/torch-cuda/_b200/paired_x_inverse_plan_safety.py b/build/torch-cuda/_b200/paired_x_inverse_plan_safety.py new file mode 100644 index 0000000000000000000000000000000000000000..40145fea029e703cc0668ae4c7c625e4dcb6cf8e --- /dev/null +++ b/build/torch-cuda/_b200/paired_x_inverse_plan_safety.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reject multi-kernel final plans before accepting caller-sized output. + +cuFFT may use output as intermediate storage in multi-kernel plans. Capture +and inspect each final prepared handle with full nominal double-complex +output storage. Only a graph containing exactly one kernel is supported. +All discovery, stream rebinding and warmup happen at preparation time. +""" +import ctypes as ct +import ctypes.util +import functools +from pathlib import Path + +import torch + +from .cuda_z_inverse_cufft_support import _lib, _check + + +@functools.lru_cache(maxsize=1) +def runtime(): + candidates = [ctypes.util.find_library('cudart')] + for root in (Path(torch.__file__).resolve().parent.parent / 'nvidia/cuda_runtime/lib', + Path('/usr/local/cuda/lib64'), Path('/usr/local/cuda/targets/x86_64-linux/lib')): + candidates.extend(str(p) for p in sorted(root.glob('libcudart.so*'))) + candidates += ['libcudart.so.13', 'libcudart.so.12', 'libcudart.so'] + for path in candidates: + if not path: + continue + try: + lib = ct.CDLL(path) + break + except OSError: + continue + else: + raise RuntimeError('Installed CUDA runtime unavailable') + signatures = { + 'cudaRuntimeGetVersion': [ct.POINTER(ct.c_int)], + 'cudaStreamBeginCapture': [ct.c_void_p, ct.c_int], + 'cudaStreamEndCapture': [ct.c_void_p, ct.POINTER(ct.c_void_p)], + 'cudaGraphGetNodes': [ct.c_void_p, ct.POINTER(ct.c_void_p), ct.POINTER(ct.c_size_t)], + 'cudaGraphNodeGetType': [ct.c_void_p, ct.POINTER(ct.c_int)], + 'cudaGraphDestroy': [ct.c_void_p], + } + for name, signature in signatures.items(): + function = getattr(lib, name) + function.argtypes, function.restype = signature, ct.c_int + return lib + + +def check(status): + if status: + raise RuntimeError(f'CUDA runtime returned error {status}') + + +def require_single_kernel(handle, intermediate, nominal_output, device, target_stream): + """Inspect a handle at setup, retaining neither the graph nor probe stream.""" + lib = runtime() + version = ct.c_int() + check(lib.cudaRuntimeGetVersion(ct.byref(version))) + probe_stream = torch.cuda.Stream(device=device) + graph = ct.c_void_p() + capturing = False + _check(_lib.cufftSetStream(handle.value, probe_stream.cuda_stream)) + try: + _check(_lib.cufftExecZ2Z(handle.value, intermediate.data_ptr(), nominal_output.data_ptr(), 1)) + probe_stream.synchronize() + check(lib.cudaStreamBeginCapture(probe_stream.cuda_stream, 1)) + capturing = True + _check(_lib.cufftExecZ2Z(handle.value, intermediate.data_ptr(), nominal_output.data_ptr(), 1)) + check(lib.cudaStreamEndCapture(probe_stream.cuda_stream, ct.byref(graph))) + capturing = False + count = ct.c_size_t() + check(lib.cudaGraphGetNodes(graph, None, ct.byref(count))) + nodes = (ct.c_void_p * count.value)() + check(lib.cudaGraphGetNodes(graph, nodes, ct.byref(count))) + types = [] + for node in nodes: + kind = ct.c_int() + check(lib.cudaGraphNodeGetType(node, ct.byref(kind))) + types.append(kind.value) + result = dict(runtime_version=version.value, graph_nodes=count.value, + node_types=types, workspace_bytes=handle.workspace_bytes, + passed=count.value == 1 and types == [0]) + if not result['passed']: + raise ValueError('Final cuFFT plan is not one kernel: ' + repr(result)) + return result + finally: + if capturing: + lib.cudaStreamEndCapture(probe_stream.cuda_stream, ct.byref(graph)) + if graph.value: + check(lib.cudaGraphDestroy(graph)) + _check(_lib.cufftSetStream(handle.value, target_stream.cuda_stream)) diff --git a/build/torch-cuda/_b200/paired_y_direct_inverse_lto.py b/build/torch-cuda/_b200/paired_y_direct_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..39ddb0584571de22f9de7e6d632dde6ff8a49633 --- /dev/null +++ b/build/torch-cuda/_b200/paired_y_direct_inverse_lto.py @@ -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_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)) diff --git a/build/torch-cuda/_b200/paired_y_forward_lto.py b/build/torch-cuda/_b200/paired_y_forward_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..04d53a5cf8bc6c697a5ec8cd66da268512f2ec23 --- /dev/null +++ b/build/torch-cuda/_b200/paired_y_forward_lto.py @@ -0,0 +1,114 @@ +# 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', pair_order='yx'): + """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' + prefix += f'#define PAIR_YX {int(pair_order == "yx")}\n' + source = prefix.encode() + Path(__file__).with_name('r2c_paired_y_forward_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)] diff --git a/build/torch-cuda/_b200/paired_y_inverse_lto.py b/build/torch-cuda/_b200/paired_y_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..1efab82f7f317a74162d24d59cf4916efbe9da8e --- /dev/null +++ b/build/torch-cuda/_b200/paired_y_inverse_lto.py @@ -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_paired_y_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)) diff --git a/build/torch-cuda/_b200/partial_y_inverse_lto.py b/build/torch-cuda/_b200/partial_y_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..a6870f14b9529e915092f88356359769e26a820e --- /dev/null +++ b/build/torch-cuda/_b200/partial_y_inverse_lto.py @@ -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")}', f'-DFFT_FINISH={int(kind == "finish")}', '-DFFT_RADIX=2') + 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', 'finish'): + raise ValueError('Expected a load or store callback') + source = Path(__file__).with_name('c2r_f64_partial_y_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(), 5 if kind == 'store' else 1, + ct.create_string_buffer(blob, len(blob)), dict(metadata, kind=kind)) diff --git a/build/torch-cuda/_b200/pitched_half_inverse_lto.py b/build/torch-cuda/_b200/pitched_half_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..97ab3355a7da31ee0f8a71e9a789e7976320d565 --- /dev/null +++ b/build/torch-cuda/_b200/pitched_half_inverse_lto.py @@ -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_pitched_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)) diff --git a/build/torch-cuda/_b200/prepared_call.cpp b/build/torch-cuda/_b200/prepared_call.cpp new file mode 100644 index 0000000000000000000000000000000000000000..a33728a4875149421da9ac34e08f5c6a9d77ecea --- /dev/null +++ b/build/torch-cuda/_b200/prepared_call.cpp @@ -0,0 +1,305 @@ +/* SPDX-License-Identifier: Apache-2.0 */ +/* Bind prepared FFT descriptors to a Python fast-call entry point. */ +#define FFT_STRING_IMPL(x) #x +#define FFT_STRING(x) FFT_STRING_IMPL(x) +#define FFT_JOIN_IMPL(x, y) x##y +#define FFT_JOIN(x, y) FFT_JOIN_IMPL(x, y) +#ifndef FFT_CALL_MODULE +#define FFT_CALL_MODULE _fft_tensor_call +#endif +#define PY_SSIZE_T_CLEAN +#include +#include +#include +#include +#include +#include +#include +#include + +typedef int (*execute_fn)(const void *, uint64_t, uint64_t); +typedef struct { execute_fn fn; const void *descriptor; unsigned int input_floor, output_floor; } Entry; +typedef struct { + PyObject *owner; + Entry entries[16]; + uint64_t stream, input_bytes, output_bytes; + unsigned int scalar, residues; + int closed, release_gil; + int input_rank, output_rank, device_index, tensor_bound; + int64_t input_shape[4], output_shape[4]; +} Prepared; + +static void destroy(PyObject *capsule) { + Prepared *p = (Prepared *)PyCapsule_GetPointer(capsule, "fft.prepared"); + if (p) { Py_DECREF(p->owner); PyMem_Free(p); } +} + +static PyObject *execute(PyObject *capsule, PyObject *const *args, Py_ssize_t nargs) { + Prepared *p = (Prepared *)PyCapsule_GetPointer(capsule, "fft.prepared"); + if (!p) return NULL; + if (nargs != 3) { PyErr_SetString(PyExc_TypeError, "Expected input, output and stream pointers"); return NULL; } + if (p->closed) { PyErr_SetString(PyExc_RuntimeError, "The plan is closed"); return NULL; } + uint64_t a = PyLong_AsUnsignedLongLong(args[0]); + uint64_t b = PyLong_AsUnsignedLongLong(args[1]); + uint64_t stream = PyLong_AsUnsignedLongLong(args[2]); + if (PyErr_Occurred()) return NULL; + if (stream != p->stream) { PyErr_SetString(PyExc_ValueError, "Prepare an independent plan for a different stream"); return NULL; } + if (!a || !b || (a % p->scalar) || (b % p->scalar)) { + PyErr_SetString(PyExc_ValueError, "Non-null scalar-aligned pointers are required"); return NULL; + } + if ((a <= b && b - a < p->input_bytes) || (b < a && a - b < p->output_bytes)) { + PyErr_SetString(PyExc_ValueError, "Input and output must not overlap"); return NULL; + } + unsigned int i = (a & 15) / p->scalar, o = (b & 15) / p->scalar; + const Entry *e = &p->entries[i * p->residues + o]; + uint64_t source = a - (e->input_floor ? (a & 15) : 0); + uint64_t target = b - (e->output_floor ? (b & 15) : 0); + int status; + if (p->release_gil) { + Py_BEGIN_ALLOW_THREADS + status = e->fn(e->descriptor, source, target); + Py_END_ALLOW_THREADS + } else { + status = e->fn(e->descriptor, source, target); + } + if (status) { PyErr_Format(PyExc_RuntimeError, "Prepared native FFT execution failed with status %d", status); return NULL; } + Py_RETURN_NONE; +} +static PyMethodDef execute_method = {"execute_ptr", (PyCFunction)(void (*)(void))execute, METH_FASTCALL, "Enqueue the prepared FFT chain with fresh caller pointers."}; + +static PyObject *bind(PyObject *module, PyObject *args) { + (void)module; + PyObject *owner, *table; + unsigned int scalar; + unsigned long long stream, input_bytes, output_bytes; + int release_gil; + if (!PyArg_ParseTuple(args, "OOIKKKp", &owner, &table, &scalar, &stream, &input_bytes, &output_bytes, &release_gil)) return NULL; + if ((scalar != 4 && scalar != 8) || !PyTuple_Check(table) || PyTuple_Size(table) != (16/scalar)*(16/scalar)) { + PyErr_SetString(PyExc_ValueError, "Invalid prepared alignment table"); return NULL; + } + Prepared *p = (Prepared *)PyMem_Calloc(1, sizeof(*p)); + if (!p) return PyErr_NoMemory(); + p->owner = Py_NewRef(owner); p->scalar = scalar; p->residues = 16/scalar; + p->stream = stream; p->input_bytes = input_bytes; p->output_bytes = output_bytes; p->release_gil = release_gil; + for (unsigned int i=0; iresidues*p->residues; ++i) { + unsigned long long fn, descriptor; + unsigned int input_floor, output_floor; + if (!PyArg_ParseTuple(PyTuple_GetItem(table, i), "KKII", &fn, &descriptor, &input_floor, &output_floor)) { + Py_DECREF(p->owner); PyMem_Free(p); return NULL; + } + if (!fn || !descriptor || input_floor > 1 || output_floor > 1) { + Py_DECREF(p->owner); PyMem_Free(p); PyErr_SetString(PyExc_ValueError, "Invalid prepared native descriptor"); return NULL; + } + p->entries[i]=(Entry){(execute_fn)(uintptr_t)fn, (void *)(uintptr_t)descriptor, input_floor, output_floor}; + } + PyObject *capsule = PyCapsule_New(p, "fft.prepared", destroy); + if (!capsule) { Py_DECREF(p->owner); PyMem_Free(p); return NULL; } + PyObject *result = PyCFunction_New(&execute_method, capsule); + Py_DECREF(capsule); + return result; +} +static PyObject *close_binding(PyObject *module, PyObject *function) { + (void)module; + if (!PyCFunction_Check(function) || PyCFunction_GetFunction(function) != execute_method.ml_meth) { + PyErr_SetString(PyExc_TypeError, "Expected a prepared FFT callable"); return NULL; + } + Prepared *p = (Prepared *)PyCapsule_GetPointer(PyCFunction_GetSelf(function), "fft.prepared"); + if (!p) return NULL; + p->closed=1; + Py_RETURN_NONE; +} + +/* Tensor allocation belongs exclusively to the value-returning wrapper. */ +static PyObject *run_prepared(Prepared *p, const at::Tensor &input) { + if (p->closed) { PyErr_SetString(PyExc_RuntimeError, "The plan is closed"); return NULL; } + if (!input.is_cuda() || input.get_device() != p->device_index || !input.is_contiguous() || + input.dim() != p->input_rank || input.scalar_type() != (p->scalar == 8 ? at::kDouble : at::kFloat)) { + PyErr_SetString(PyExc_ValueError, "Input metadata does not match the prepared plan"); return NULL; + } + for (int d = 0; d < p->input_rank; ++d) { + if (input.size(d) != p->input_shape[d]) { + PyErr_SetString(PyExc_ValueError, "Input shape does not match the prepared plan"); return NULL; + } + } + // Allocate fresh contiguous CUDA storage through PyTorch's caching allocator. + // This remains exclusive to the allocating wrapper, outside execute_ptr. + at::Tensor output = at::detail::empty_cuda( + at::IntArrayRef(p->output_shape, p->output_rank), input.scalar_type(), + input.device(), at::MemoryFormat::Contiguous); + uint64_t a = (uint64_t)(uintptr_t)input.const_data_ptr(); + uint64_t b = (uint64_t)(uintptr_t)output.mutable_data_ptr(); + unsigned int i = (a & 15) / p->scalar, o = (b & 15) / p->scalar; + const Entry *e = &p->entries[i * p->residues + o]; + uint64_t source = a - (e->input_floor ? (a & 15) : 0); + uint64_t target = b - (e->output_floor ? (b & 15) : 0); + int status; + if (p->release_gil) { + Py_BEGIN_ALLOW_THREADS + status = e->fn(e->descriptor, source, target); + Py_END_ALLOW_THREADS + } else { + status = e->fn(e->descriptor, source, target); + } + if (status) { PyErr_Format(PyExc_RuntimeError, "Prepared native FFT execution failed with status %d", status); return NULL; } + return THPVariable_Wrap(output); + +} + +static PyObject *run_tensor(PyObject *capsule, PyObject *const *args, Py_ssize_t nargs) { + Prepared *p = (Prepared *)PyCapsule_GetPointer(capsule, "fft.prepared"); + if (!p) return NULL; + if (nargs != 2) { PyErr_SetString(PyExc_TypeError, "Expected input tensor and stream pointer"); return NULL; } + if (p->closed) { PyErr_SetString(PyExc_RuntimeError, "The plan is closed"); return NULL; } + if (!p->tensor_bound) { PyErr_SetString(PyExc_RuntimeError, "Tensor boundary was not prepared"); return NULL; } + uint64_t stream = PyLong_AsUnsignedLongLong(args[1]); + if (PyErr_Occurred()) return NULL; + if (stream != p->stream) { PyErr_SetString(PyExc_ValueError, "Prepare an independent plan for a different stream"); return NULL; } + if (!THPVariable_Check(args[0])) { PyErr_SetString(PyExc_TypeError, "Expected a torch tensor"); return NULL; } + try { + const at::Tensor &input = THPVariable_Unpack(args[0]); + return run_prepared(p, input); + } catch (const std::exception &error) { + PyErr_SetString(PyExc_RuntimeError, error.what()); return NULL; + } +} +static PyMethodDef run_method = {"run", (PyCFunction)(void (*)(void))run_tensor, METH_FASTCALL, + "Allocate a fresh output and enqueue the prepared FFT chain."}; + +static PyObject *bind_tensor(PyObject *module, PyObject *args) { + (void)module; + PyObject *function, *input_shape, *output_shape; + int device_index; + if (!PyArg_ParseTuple(args, "OOOi", &function, &input_shape, &output_shape, &device_index)) return NULL; + if (!PyCFunction_Check(function) || PyCFunction_GetFunction(function) != execute_method.ml_meth) { + PyErr_SetString(PyExc_TypeError, "Expected a prepared FFT pointer callable"); return NULL; + } + PyObject *capsule = PyCFunction_GetSelf(function); + Prepared *p = (Prepared *)PyCapsule_GetPointer(capsule, "fft.prepared"); + if (!p) return NULL; + if (p->closed || p->tensor_bound) { PyErr_SetString(PyExc_RuntimeError, "The tensor boundary must be bound once during preparation"); return NULL; } + if (device_index < 0 || !PyTuple_Check(input_shape) || !PyTuple_Check(output_shape)) { + PyErr_SetString(PyExc_ValueError, "Expected shape tuples and a CUDA device index"); return NULL; + } + int nr = (int)PyTuple_Size(input_shape), nw = (int)PyTuple_Size(output_shape); + if (nr < 1 || nr > 4 || nw < 1 || nw > 4) { PyErr_SetString(PyExc_ValueError, "Unsupported tensor rank"); return NULL; } + uint64_t input_bytes = p->scalar, output_bytes = p->scalar; + for (int d = 0; d < nr; ++d) { + int64_t value = PyLong_AsLongLong(PyTuple_GetItem(input_shape, d)); + if (PyErr_Occurred()) return NULL; + if (value <= 0 || (uint64_t)value > UINT64_MAX / input_bytes) { PyErr_SetString(PyExc_ValueError, "Invalid input extent"); return NULL; } + p->input_shape[d] = value; input_bytes *= value; + } + for (int d = 0; d < nw; ++d) { + int64_t value = PyLong_AsLongLong(PyTuple_GetItem(output_shape, d)); + if (PyErr_Occurred()) return NULL; + if (value <= 0 || (uint64_t)value > UINT64_MAX / output_bytes) { PyErr_SetString(PyExc_ValueError, "Invalid output extent"); return NULL; } + p->output_shape[d] = value; output_bytes *= value; + } + if (input_bytes != p->input_bytes || output_bytes != p->output_bytes) { + PyErr_SetString(PyExc_ValueError, "Tensor extents must match the prepared FFT"); return NULL; + } + p->input_rank = nr; p->output_rank = nw; p->device_index = device_index; p->tensor_bound = 1; + return PyCFunction_New(&run_method, capsule); +} + +/* Cache prepared resources using public metadata only. Input addresses and values + * are never part of the key. The GIL serializes lookup and preparation; native + * execution follows the independently bound stream of each cached plan. */ +struct CachedPlan { + int device, rank; + uint64_t stream; + int64_t shape[4]; + PyObject *callable; + Prepared *prepared; +}; +struct PlanCache { + PyObject *factory; + unsigned int scalar, rank; + std::vector entries; + size_t recent; + PlanCache(PyObject *f, unsigned int s, unsigned int r) + : factory(Py_NewRef(f)), scalar(s), rank(r), recent(0) {} + ~PlanCache() { + for (auto &entry : entries) Py_DECREF(entry.callable); + Py_DECREF(factory); + } +}; +static void destroy_cache(PyObject *capsule) { + auto *cache = (PlanCache *)PyCapsule_GetPointer(capsule, "fft.plan_cache"); + if (cache) delete cache; +} +static bool matches(const CachedPlan &entry, const at::Tensor &input, uint64_t stream) { + return entry.device == input.get_device() && entry.stream == stream && + entry.rank == input.dim() && + std::memcmp(entry.shape, input.sizes().data(), entry.rank * sizeof(int64_t)) == 0; +} +static PyObject *run_cached(PyObject *capsule, PyObject *argument) { + auto *cache = (PlanCache *)PyCapsule_GetPointer(capsule, "fft.plan_cache"); + if (!cache) return NULL; + if (!THPVariable_Check(argument)) { + PyErr_SetString(PyExc_TypeError, "Expected a torch tensor"); return NULL; + } + try { + const at::Tensor &input = THPVariable_Unpack(argument); + if (!input.is_cuda() || !input.is_contiguous() || input.dim() != cache->rank || + input.scalar_type() != (cache->scalar == 8 ? at::kDouble : at::kFloat) || + (cache->rank == 4 && input.size(3) != 2)) { + PyErr_SetString(PyExc_ValueError, "Expected a contiguous CUDA tensor in the request shape and precision"); return NULL; + } + const uint64_t stream = (uint64_t)(uintptr_t)c10::cuda::getCurrentCUDAStream(input.get_device()).stream(); + if (!cache->entries.empty()) { + auto &recent = cache->entries[cache->recent]; + if (matches(recent, input, stream)) return run_prepared(recent.prepared, input); + for (size_t i = 0; i < cache->entries.size(); ++i) { + if (matches(cache->entries[i], input, stream)) { + cache->recent = i; + return run_prepared(cache->entries[i].prepared, input); + } + } + } + PyObject *callable = PyObject_CallOneArg(cache->factory, argument); + if (!callable) return NULL; + if (!PyCFunction_Check(callable) || PyCFunction_GetFunction(callable) != run_method.ml_meth) { + Py_DECREF(callable); PyErr_SetString(PyExc_TypeError, "Plan factory must return a native prepared tensor callable"); return NULL; + } + auto *p = (Prepared *)PyCapsule_GetPointer(PyCFunction_GetSelf(callable), "fft.prepared"); + if (!p) { Py_DECREF(callable); return NULL; } + if (p->closed || !p->tensor_bound || p->scalar != cache->scalar || p->stream != stream || + p->device_index != input.get_device() || p->input_rank != input.dim() || + std::memcmp(p->input_shape, input.sizes().data(), p->input_rank * sizeof(int64_t)) != 0) { + Py_DECREF(callable); PyErr_SetString(PyExc_ValueError, "Plan factory returned incompatible metadata or stream"); return NULL; + } + CachedPlan entry = {}; + entry.device = input.get_device(); entry.rank = input.dim(); entry.stream = stream; + std::memcpy(entry.shape, input.sizes().data(), entry.rank * sizeof(int64_t)); + entry.callable = callable; entry.prepared = p; + try { cache->entries.push_back(entry); } + catch (...) { Py_DECREF(callable); throw; } + cache->recent = cache->entries.size() - 1; + return run_prepared(p, input); + } catch (const std::exception &error) { + PyErr_SetString(PyExc_RuntimeError, error.what()); return NULL; + } +} +static PyMethodDef cached_method = {"run_cached", run_cached, METH_O, + "Return a fresh FFT output using plans cached by shape, dtype, device and stream."}; +static PyObject *bind_cache(PyObject *module, PyObject *args) { + (void)module; + PyObject *factory; + unsigned int scalar, rank; + if (!PyArg_ParseTuple(args, "OII", &factory, &scalar, &rank)) return NULL; + if (!PyCallable_Check(factory) || (scalar != 4 && scalar != 8) || (rank != 3 && rank != 4)) { + PyErr_SetString(PyExc_ValueError, "Expected a plan factory, scalar size and FFT input rank"); return NULL; + } + PlanCache *cache; + try { cache = new PlanCache(factory, scalar, rank); } + catch (const std::exception &error) { PyErr_SetString(PyExc_RuntimeError, error.what()); return NULL; } + PyObject *capsule = PyCapsule_New(cache, "fft.plan_cache", destroy_cache); + if (!capsule) { delete cache; return NULL; } + PyObject *result = PyCFunction_New(&cached_method, capsule); + Py_DECREF(capsule); + return result; +} +static PyMethodDef methods[] = {{"bind_cache", bind_cache, METH_VARARGS, "Create a metadata-only prepared plan cache."}, {"bind_tensor", bind_tensor, METH_VARARGS, "Prepare the allocating tensor boundary."}, {"bind", bind, METH_VARARGS, "Bind prepared state."}, {"close", close_binding, METH_O, "Invalidate a prepared callable before releasing its resources."}, {NULL,NULL,0,NULL}}; +static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, FFT_STRING(FFT_CALL_MODULE), NULL, -1, methods, NULL, NULL, NULL, NULL}; +PyMODINIT_FUNC FFT_JOIN(PyInit_, FFT_CALL_MODULE)(void) { return PyModule_Create(&module); } diff --git a/build/torch-cuda/_b200/prepared_call.py b/build/torch-cuda/_b200/prepared_call.py new file mode 100644 index 0000000000000000000000000000000000000000..1acb47dab9b3548896c3460cc6da5d6170f3bc3a --- /dev/null +++ b/build/torch-cuda/_b200/prepared_call.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Prepare an allocation-free native caller for immutable FFT dispatch tables.""" +from __future__ import annotations +import ctypes as ct +import fcntl +import hashlib +import importlib.util +import math +import os +from pathlib import Path +import shlex +import subprocess +import sys +import sysconfig +import tempfile +import threading +import torch + +_MODULE = None +_LOCK = threading.Lock() + +def load_module(): + global _MODULE + if _MODULE is not None: + return _MODULE + with _LOCK: + if _MODULE is not None: + return _MODULE + source = Path(__file__).with_name('prepared_call.cpp') + compiler = shlex.split(os.environ.get('CXX', 'c++')) + torch_root = Path(torch.__file__).resolve().parent + flags = ['-std=c++17', '-O3', '-DNDEBUG', '-fPIC', '-shared', '-Wall', '-Wextra', + '-I' + sysconfig.get_path('include'), '-isystem', str(torch_root / 'include'), + '-D_GLIBCXX_USE_CXX11_ABI=' + str(int(torch._C._GLIBCXX_USE_CXX11_ABI))] + from torch.utils.cpp_extension import include_paths + for include in include_paths(device_type='cuda'): + flags.extend(['-isystem', include]) + links = ['-L' + str(torch_root / 'lib'), '-Wl,-rpath,' + str(torch_root / 'lib'), + '-ltorch_python', '-ltorch_cpu', '-ltorch_cuda', '-lc10', '-lc10_cuda'] + key = hashlib.sha256(source.read_bytes() + repr((compiler, flags, links, sys.version, torch.__version__)).encode()).hexdigest() + cache = Path(tempfile.gettempdir()) / ('fft-tensor-call-' + str(os.getuid())) / key + cache.mkdir(parents=True, exist_ok=True) + module_name = '_fft_tensor_call_' + key[:24] + target = cache / (module_name + sysconfig.get_config_var('EXT_SUFFIX')) + with (cache / 'build.lock').open('w') as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + if not target.is_file(): + pending = cache / ('pending.' + str(os.getpid()) + '.so') + try: + subprocess.run(compiler + flags + ['-DFFT_CALL_MODULE=' + module_name, str(source), '-o', str(pending)] + links, + check=True, capture_output=True, text=True) + pending.replace(target) + finally: + pending.unlink(missing_ok=True) + spec = importlib.util.spec_from_file_location(module_name, target) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + _MODULE = module + return module + + +def _entry(plan, i, o, scalar): + """Resolve only known immutable descriptor formats during preparation.""" + if hasattr(plan, '_aligned'): + use_scalar = (i or o) if plan.direction == 'r2c' else i + return _entry(plan._scalar if use_scalar else plan._aligned, i, o, scalar) + if getattr(plan, 'offset_plan', None) is not None and o: + # These inverse plans use a separate output-only callback table. + selected = plan.offset_plan + return selected.native_function, selected.native_pointers[o], 0, 1 + if getattr(plan, '_offset_plan', None) is not None and o: + selected = plan._offset_plan + return selected.native_function, selected.native_pointers[o], 0, 1 + if hasattr(plan, 'direct_pointers'): + if o == 0: + return plan.direct_function, plan.direct_pointers[i], 1, 0 + return plan.native_function, plan.native_pointers[i], 1, 0 + if hasattr(plan, 'native_pointers'): + pointers = plan.native_pointers + residues = 16 // scalar + if len(pointers) == residues * residues: + return plan.native_function, pointers[i * residues + o], 1, 1 + if len(pointers) == residues: + return plan.native_function, pointers[i], 1, 0 + raise ValueError('Unsupported prepared pointer table') + if hasattr(plan, '_native_pointer'): + return plan._native_function, plan._native_pointer, 0, 0 + if getattr(plan, 'native', None) is not None: + return plan.native.function, plan.native.pointer, 0, 0 + raise ValueError('No supported immutable native dispatcher') + + +class Plan: + """Own a prepared FFT plan and its native, stream-bound caller. + + The owner remains alive while an extracted execute_ptr callable exists. + Closing invalidates the callable before the original resources are freed. + As with the underlying plans, queued work must finish before close. + """ + def __init__(self, owner, *, release_gil=False): + self._owner = owner + self.closed = False + for name in ('shape', 'dtype', 'direction', 'device', 'stream', 'stream_ptr', + 'input_shape', 'output_shape'): + setattr(self, name, getattr(owner, name)) + self.itemsize = 8 if str(self.dtype) == 'torch.float64' else 4 + table = [] + for i in range(16 // self.itemsize): + for o in range(16 // self.itemsize): + fn, ptr, inf, outf = _entry(owner, i, o, self.itemsize) + table.append((ct.cast(fn, ct.c_void_p).value, ct.cast(ptr, ct.c_void_p).value, inf, outf)) + self._module = load_module() + self.execute_ptr = self._module.bind(owner, tuple(table), self.itemsize, self.stream_ptr, + math.prod(self.input_shape) * self.itemsize, math.prod(self.output_shape) * self.itemsize, + release_gil) + self.run = self._module.bind_tensor(self.execute_ptr, tuple(self.input_shape), + tuple(self.output_shape), self.device.index) + + @property + def workspace_bytes(self): + return self._owner.workspace_bytes + + 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') + selected = self.stream if stream is None else stream + self.execute_ptr(input.data_ptr(), output.data_ptr(), selected.cuda_stream) + return output + + def close(self): + if not self.closed: + self._module.close(self.execute_ptr) + self._owner.close() + self.closed = True diff --git a/build/torch-cuda/_b200/projection.cu b/build/torch-cuda/_b200/projection.cu new file mode 100644 index 0000000000000000000000000000000000000000..65a6ed68b02e7c615ed8b4208de522b92560471d --- /dev/null +++ b/build/torch-cuda/_b200/projection.cu @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: Apache-2.0 +// Load-only Hermitian projection on the original packed endpoint planes. +// The first transform is out-of-place C2C, so all mirror loads read the +// immutable caller input. No callback depends on block execution order. +#if FFT_DOUBLE +typedef double2 fft_complex; +typedef double fft_real; +#else +typedef float2 fft_complex; +typedef float fft_real; +#endif + +__device__ fft_complex fft_hermitian_load(void *data, unsigned long long offset, + void *caller_info, void *shared_memory) { + const fft_complex *input = static_cast(data); + fft_complex value = input[offset]; + const unsigned long long kx = offset % FFT_NF; + if (kx == 0 || kx == FFT_NF - 1) { + const unsigned long long ky = (offset / FFT_NF) % FFT_NY; + const unsigned long long kz = offset / (FFT_NY * FFT_NF); + const unsigned long long my = ky == 0 ? 0 : FFT_NY - ky; + const unsigned long long mz = kz == 0 ? 0 : FFT_NZ - kz; + const fft_complex mirror = input[(mz * FFT_NY + my) * FFT_NF + kx]; + // Scale each operand first to avoid overflow in the sum/difference. + const fft_real half = static_cast(0.5); + value.x = half * value.x + half * mirror.x; + value.y = half * value.y - half * mirror.y; + } + return value; +} diff --git a/build/torch-cuda/_b200/r2c_half_forward_cufft_z_scalar_io.cu b/build/torch-cuda/_b200/r2c_half_forward_cufft_z_scalar_io.cu new file mode 100644 index 0000000000000000000000000000000000000000..ecfda56fe5233ab5dceb0a742ed7ca7bb1de3802 --- /dev/null +++ b/build/torch-cuda/_b200/r2c_half_forward_cufft_z_scalar_io.cu @@ -0,0 +1,86 @@ +// Use streaming cache hints only for the selected complete grid. +#define FFT_STREAM_BOUNDARY (!FFT_DOUBLE && NZ == 256 && NY == 256 && NX == 256) +// SPDX-License-Identifier: Apache-2.0 +// Half-length complex operands are loaded directly from caller real pairs. +// The final stage reconstructs XY spectra from mathematical FFT results. +#if FFT_DOUBLE +using scalar = double; +#else +using scalar = float; +#endif + +#if IS_LOAD +__device__ double2 fft_pair_load(void *data, unsigned long long offset, + void *caller_info, void *shared_memory) { + const scalar *input = static_cast(data) + SCALAR_SHIFT; +#if FFT_DOUBLE && SCALAR_SHIFT == 0 + return reinterpret_cast(input)[offset]; +#elif !FFT_DOUBLE && SCALAR_SHIFT % 2 == 0 + const float2 pair = reinterpret_cast(input)[offset]; + return make_double2(static_cast(pair.x), static_cast(pair.y)); +#else + return make_double2(static_cast(input[2 * offset]), + static_cast(input[2 * offset + 1])); +#endif +} +#elif RECONSTRUCT +// Select an index width from the full logical packed extent at prepare time. +// Device addresses remain full width, independently of the logical index. +#if (1ULL * NZ * NY * NF) <= 0x100000000ULL +using index_t = unsigned int; +#else +using index_t = unsigned long long; +#endif +__device__ double2 fft_reconstruct_load(void *data, unsigned long long offset, + void *caller_info, void *shared_memory) { + constexpr index_t m = NX / 2; + const index_t index = static_cast(offset); + const index_t k = index % NF; + const index_t row = index / NF; + const index_t y = row % NY; + const index_t z = row / NY; + const index_t mirror_y = y == 0 ? 0 : NY - y; + const double2 *input = static_cast(data); + const double2 a = input[row * m + k % m]; + const double2 b = input[(z * NY + mirror_y) * m + (m - k) % m]; + const double2 w = static_cast(caller_info)[k]; + const double er = 0.5 * (a.x + b.x); + const double ei = 0.5 * (a.y - b.y); + const double odd_r = 0.5 * (a.y + b.y); + const double 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); +} +#else +__device__ void fft_scalar_store(void *data, unsigned long long offset, + double2 value, void *caller_info, + void *shared_memory) { + scalar *output = static_cast(data) + SCALAR_SHIFT; +#if FFT_DOUBLE && SCALAR_SHIFT == 0 +#if FFT_STREAM_BOUNDARY + __stcs(reinterpret_cast(output) + offset, value); +#else + reinterpret_cast(output)[offset] = value; +#endif +#elif !FFT_DOUBLE && SCALAR_SHIFT % 2 == 0 +#if FFT_STREAM_BOUNDARY + __stcs(reinterpret_cast(output) + offset, + make_float2(static_cast(value.x), static_cast(value.y))); +#else + reinterpret_cast(output)[offset] = + make_float2(static_cast(value.x), static_cast(value.y)); +#endif +#else +#if FFT_STREAM_BOUNDARY + __stcs(output + 2 * offset, static_cast(value.x)); +#else + output[2 * offset] = static_cast(value.x); +#endif +#if FFT_STREAM_BOUNDARY + __stcs(output + 2 * offset + 1, static_cast(value.y)); +#else + output[2 * offset + 1] = static_cast(value.y); +#endif +#endif +} +#endif diff --git a/build/torch-cuda/_b200/r2c_paired_y_forward_scalar_io.cu b/build/torch-cuda/_b200/r2c_paired_y_forward_scalar_io.cu new file mode 100644 index 0000000000000000000000000000000000000000..b0771541148af721f8207ac3383f8071066b38bd --- /dev/null +++ b/build/torch-cuda/_b200/r2c_paired_y_forward_scalar_io.cu @@ -0,0 +1,97 @@ +// Use streaming cache hints only for the selected complete grid. +#define FFT_STREAM_BOUNDARY (!FFT_DOUBLE && NZ == 256 && NY == 256 && NX == 256) +// SPDX-License-Identifier: Apache-2.0 +// Y-paired complex operands are loaded directly from caller real rows. +// The final stage reconstructs XY spectra from mathematical FFT results. +#if FFT_DOUBLE +using scalar = double; +#else +using scalar = float; +#endif + +#if IS_LOAD +__device__ double2 fft_pair_load(void *data, unsigned long long offset, + void *caller_info, void *shared_memory) { + constexpr unsigned int hy = NY / 2; + const unsigned int index = static_cast(offset); + const unsigned int z = index / (hy * NX); +#if PAIR_YX + const unsigned int x = index % NX; + const unsigned int y = index / NX % hy; +#else + const unsigned int y = index % hy; + const unsigned int x = index / hy % NX; +#endif + const scalar *input = static_cast(data) + SCALAR_SHIFT; + const unsigned int address = (z * NY + 2 * y) * NX + x; + return make_double2(static_cast(input[address]), + static_cast(input[address + NX])); +} +#elif RECONSTRUCT +// Select an index width from the full logical packed extent at prepare time. +// Device addresses remain full width, independently of the logical index. +#if (1ULL * NZ * NY * NF) <= 0x100000000ULL +using index_t = unsigned int; +#else +using index_t = unsigned long long; +#endif +__device__ double2 fft_reconstruct_load(void *data, unsigned long long offset, + void *caller_info, void *shared_memory) { + constexpr index_t hy = NY / 2; + const index_t index = static_cast(offset); + const index_t x = index % NF; + const index_t row = index / NF; + const index_t y = row % NY; + const index_t z = row / NY; + const index_t ky = y % hy; + const index_t my = ky == 0 ? 0 : hy - ky; + const index_t mx = x == 0 ? 0 : NX - x; + const double2 *input = static_cast(data); +#if PAIR_YX + const double2 a = input[(z * hy + ky) * NX + x]; + const double2 b = input[(z * hy + my) * NX + mx]; +#else + const double2 a = input[(z * NX + x) * hy + ky]; + const double2 b = input[(z * NX + mx) * hy + my]; +#endif + const double2 w = static_cast(caller_info)[y]; + const double er = 0.5 * (a.x + b.x); + const double ei = 0.5 * (a.y - b.y); + const double odd_r = 0.5 * (a.y + b.y); + const double 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); +} +#else +__device__ void fft_scalar_store(void *data, unsigned long long offset, + double2 value, void *caller_info, + void *shared_memory) { + scalar *output = static_cast(data) + SCALAR_SHIFT; +#if FFT_DOUBLE && SCALAR_SHIFT == 0 +#if FFT_STREAM_BOUNDARY + __stcs(reinterpret_cast(output) + offset, value); +#else + reinterpret_cast(output)[offset] = value; +#endif +#elif !FFT_DOUBLE && SCALAR_SHIFT % 2 == 0 +#if FFT_STREAM_BOUNDARY + __stcs(reinterpret_cast(output) + offset, + make_float2(static_cast(value.x), static_cast(value.y))); +#else + reinterpret_cast(output)[offset] = + make_float2(static_cast(value.x), static_cast(value.y)); +#endif +#else +#if FFT_STREAM_BOUNDARY + __stcs(output + 2 * offset, static_cast(value.x)); +#else + output[2 * offset] = static_cast(value.x); +#endif +#if FFT_STREAM_BOUNDARY + __stcs(output + 2 * offset + 1, static_cast(value.y)); +#else + output[2 * offset + 1] = static_cast(value.y); +#endif +#endif +} +#endif diff --git a/build/torch-cuda/_b200/reconstruct.cu b/build/torch-cuda/_b200/reconstruct.cu new file mode 100644 index 0000000000000000000000000000000000000000..9716807d998f6e1a068b9cbb1ba7af2076bbbd87 --- /dev/null +++ b/build/torch-cuda/_b200/reconstruct.cu @@ -0,0 +1,43 @@ +// 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 +// Reconstruct twice the half-length complex spectrum from immutable packed +// caller input. Endpoint projection preserves irfftn for arbitrary spectra. +__device__ double2 fft_full_half_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(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 double2 *input = static_cast(data); + double2 a = input[index]; + double2 b = 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 = input[mirror_row * nf]; + const double2 b_mirror = 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(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; +} diff --git a/build/torch-cuda/_b200/scalar_input_inverse_backend.py b/build/torch-cuda/_b200/scalar_input_inverse_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..bf06fc26e997586f16d329e74e3c3022b6892da8 --- /dev/null +++ b/build/torch-cuda/_b200/scalar_input_inverse_backend.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Prepared scalar-aligned inverse FFT with direct caller-buffer callbacks. + +Both caller precisions use FP64 arithmetic, with conversion at the initial +loads and final stores. The only intermediate is transformed Z/Y data in +exact [nf,nz,ny,2] layout. This experimental adapter is not a release. +""" +from __future__ import annotations + +import threading +import torch + +from .scalar_input_inverse_cufft_support import _Handle, _lib, _check +from .scalar_input_inverse_lto import compile_callback +from .direct_dispatch import PreparedChain + + +class Plan: + def __init__(self, shape, dtype, direction='c2r', *, device=None, stream=None): + 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 != 'c2r': + raise ValueError('Expected a float32 or float64 inverse transform') + 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, self.dispatch = 'scalar_io_zy_x_fp64', 'prepared_native' + self.closed, self.handles, self.chains = False, [], [] + nz, ny, nx = self.shape + rows, nf = nz * ny, nx // 2 + 1 + self.input_shape, self.output_shape = (nz, ny, nf, 2), self.shape + self.compiled_metadata = [] + try: + with torch.cuda.device(self.device), torch.cuda.stream(self.stream): + self.intermediate = torch.empty((nf, nz, ny, 2), 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, 'load', shift, arch) + self.compiled_metadata.extend(item[3] for item in cb) + handle = _Handle((nz, ny), 0x69, (nz, ny), nf, 1, + (nz, ny), 1, rows, nf, self.device, self.stream, + callback=cb, xt=True) + self.handles.append(handle) + load_handles.append(handle) + cb = compile_callback(dtype == torch.float64, 'store', shift, arch) + self.compiled_metadata.extend(item[3] for item in cb) + handle = _Handle((nx,), 0x6C, (nf,), rows, 1, + (nx,), 1, nx, rows, self.device, self.stream, + callback=cb) + 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.cufftExecZ2D, output_handle.value, 0, 0, 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(nz * ny * nf * 2 + self.residues, + dtype=dtype, device=self.device) + dummy_out = torch.empty(rows * nx + self.residues, + dtype=dtype, device=self.device) + 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() * 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.chains = self.native_pointers = None + self.closed = True + + + + +_CACHE = {} +_LOCK = threading.Lock() diff --git a/build/torch-cuda/_b200/scalar_input_inverse_cufft_support.py b/build/torch-cuda/_b200/scalar_input_inverse_cufft_support.py new file mode 100644 index 0000000000000000000000000000000000000000..5d180d4ed416d7abda67ce6fa2b3efe88d4bc92c --- /dev/null +++ b/build/torch-cuda/_b200/scalar_input_inverse_cufft_support.py @@ -0,0 +1,120 @@ +# 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], + "cufftXtMakePlanMany": [ct.c_int, ct.c_int, ct.POINTER(ct.c_longlong), + ct.POINTER(ct.c_longlong), ct.c_longlong, ct.c_longlong, ct.c_int, + ct.POINTER(ct.c_longlong), ct.c_longlong, ct.c_longlong, ct.c_int, + ct.c_longlong, _size_p, ct.c_int], + "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, xt=False): + 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)) + for item in (() if callback is None else callback): + symbol, callback_type, blob, _metadata = item + 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) + if xt: + wide = ct.c_longlong * len(sizes) + # Complex transform types preserve caller precision. + if kind not in (0x29, 0x69): + raise ValueError('Only C2C and Z2Z use the Xt wrapper') + precision = 5 if kind == 0x69 else 4 + _check(_lib.cufftXtMakePlanMany( + self.value, len(sizes), wide(*sizes), wide(*inembed), + istride, idist, precision, wide(*onembed), ostride, odist, + precision, batch, ct.byref(size), precision)) + else: + _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 diff --git a/build/torch-cuda/_b200/scalar_input_inverse_lto.py b/build/torch-cuda/_b200/scalar_input_inverse_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..abc71bbd2b63694917f4db6e8235af90bfc524d0 --- /dev/null +++ b/build/torch-cuda/_b200/scalar_input_inverse_lto.py @@ -0,0 +1,109 @@ +# 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): + """Compile every scalar residue at preparation, never during execution.""" + if role not in ('load', '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 == "load")}\n' + source = prefix.encode() + Path(__file__).with_name('c2r_f64_scalar_input_inverse_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_scalar_load' if role == 'load' else b'fft_scalar_store', + 1 if role == 'load' else 7, ct.create_string_buffer(blob, len(blob)), metadata)] diff --git a/build/torch-cuda/_b200/scalar_io.cu b/build/torch-cuda/_b200/scalar_io.cu new file mode 100644 index 0000000000000000000000000000000000000000..630ca2307e2e24a811c909c950b9b2922e16c959 --- /dev/null +++ b/build/torch-cuda/_b200/scalar_io.cu @@ -0,0 +1,67 @@ +// SPDX-License-Identifier: Apache-2.0 +// Half-length complex operands are loaded directly from caller real pairs. +// The final stage reconstructs XY spectra from mathematical FFT results. +#if FFT_DOUBLE +using scalar = double; +#else +using scalar = float; +#endif + +#if IS_LOAD +__device__ double2 fft_pair_load(void *data, unsigned long long offset, + void *caller_info, void *shared_memory) { + const scalar *input = static_cast(data) + SCALAR_SHIFT; +#if FFT_DOUBLE && SCALAR_SHIFT == 0 + return reinterpret_cast(input)[offset]; +#elif !FFT_DOUBLE && SCALAR_SHIFT % 2 == 0 + const float2 pair = reinterpret_cast(input)[offset]; + return make_double2(static_cast(pair.x), static_cast(pair.y)); +#else + return make_double2(static_cast(input[2 * offset]), + static_cast(input[2 * offset + 1])); +#endif +} +#elif RECONSTRUCT +// Select an index width from the full logical packed extent at prepare time. +// Device addresses remain full width, independently of the logical index. +#if (1ULL * NZ * NY * NF) <= 0x100000000ULL +using index_t = unsigned int; +#else +using index_t = unsigned long long; +#endif +__device__ double2 fft_reconstruct_load(void *data, unsigned long long offset, + void *caller_info, void *shared_memory) { + constexpr index_t m = NX / 2; + const index_t index = static_cast(offset); + const index_t k = index % NF; + const index_t row = index / NF; + const index_t y = row % NY; + const index_t z = row / NY; + const index_t mirror_y = y == 0 ? 0 : NY - y; + const double2 *input = static_cast(data); + const double2 a = input[row * m + k % m]; + const double2 b = input[(z * NY + mirror_y) * m + (m - k) % m]; + const double2 w = static_cast(caller_info)[k]; + const double er = 0.5 * (a.x + b.x); + const double ei = 0.5 * (a.y - b.y); + const double odd_r = 0.5 * (a.y + b.y); + const double 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); +} +#else +__device__ void fft_scalar_store(void *data, unsigned long long offset, + double2 value, void *caller_info, + void *shared_memory) { + scalar *output = static_cast(data) + SCALAR_SHIFT; +#if FFT_DOUBLE && SCALAR_SHIFT == 0 + reinterpret_cast(output)[offset] = value; +#elif !FFT_DOUBLE && SCALAR_SHIFT % 2 == 0 + reinterpret_cast(output)[offset] = + make_float2(static_cast(value.x), static_cast(value.y)); +#else + output[2 * offset] = static_cast(value.x); + output[2 * offset + 1] = static_cast(value.y); +#endif +} +#endif diff --git a/build/torch-cuda/_b200/scalar_output_backend.py b/build/torch-cuda/_b200/scalar_output_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..8e314442c7e845b1a77b8316a58718b21e944b15 --- /dev/null +++ b/build/torch-cuda/_b200/scalar_output_backend.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Two prepared FP64 cuFFT stages with callback-fused half-length inversion.""" +from __future__ import annotations + +import threading +import torch + +from .narrow_index_inverse_cufft_support import _Handle, _lib, _check +from .scalar_output_lto import compile_callbacks +from .direct_dispatch import PreparedChain +from .dense_half_inverse_schedule import twiddles, factors + + +class Plan: + """Setup-only planning/JIT, then enqueue-only caller-buffer execution. + + One plan binds one stream. Serialize submissions; use independent plans + for concurrent streams. The caller keeps plan and buffers alive until all + work completes, and waits externally before calling close. + """ + def __init__(self, shape, dtype, direction='c2r', *, device=None, stream=None, + layout='frequency', patient=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 != torch.float64 or direction != 'c2r': + raise ValueError('This experiment supports FP64 inverse transforms') + if layout not in ('frequency', 'rowmajor'): + raise ValueError('Unknown intermediate layout') + for n in self.shape: + factors(n) + self.dtype, self.direction = dtype, direction + self.variant = f'half_cufft_aligned_{layout}_p{int(patient)}' + 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.closed, self.handles = False, [] + nz, ny, nx = self.shape + rows, m, nf = nz * ny, nx // 2, nx // 2 + 1 + self.input_shape, self.output_shape = (nz, ny, nf, 2), self.shape + self._execute = _lib.cufftExecZ2Z + self.layout = layout + rowmajor = layout == 'rowmajor' + try: + with torch.cuda.device(self.device), torch.cuda.stream(self.stream): + self.intermediate = torch.empty((nz, ny, nf, 2) if rowmajor else (nf, nz, ny, 2), + dtype=dtype, device=self.device) + self.table = torch.tensor(twiddles(nx), dtype=dtype, device=self.device) + self.tmp_ptr = self.intermediate.data_ptr() + capability = torch.cuda.get_device_capability(self.device) + self.compiled_metadata = [] + # Read exact caller packing and write transformed values directly + # into their mathematical intermediate order, without repacking. + h0 = _Handle((nz, ny), 0x69, (nz, ny), nf, 1, (nz, ny), + nf if rowmajor else 1, 1 if rowmajor else rows, + nf, self.device, self.stream, xt=True) + self.handles.append(h0) + self.native = [] + for shift in (0, 1): + callbacks = compile_callbacks(nx, rows, rowmajor, 10 * capability[0] + capability[1], + store_shift=None if shift == 0 else 1) + self.compiled_metadata.extend(callback[3] for callback in callbacks) + h1 = _Handle((m,), 0x69, (nf,), 1 if rowmajor else rows, + nf if rowmajor else 1, (m,), 1, m, rows, + self.device, self.stream, callback=callbacks, + caller_info=self.table.data_ptr(), patient=bool(patient)) + self.handles.append(h1) + self.native.append(PreparedChain([ + (self._execute, h0.value, 1, 1, 0, 2), + (self._execute, h1.value, 1, 1, 2, 1)], self.tmp_ptr)) + self.native_function = self.native[0].function + self.native_pointers = tuple(chain.pointer for chain in self.native) + dummy_input = torch.zeros(self.input_shape, dtype=dtype, device=self.device) + dummy_output = torch.empty(nz * ny * nx + 1, dtype=dtype, device=self.device) + for shift in (0, 1): + self.execute_ptr(dummy_input.data_ptr(), dummy_output.data_ptr() + 8 * shift, + self.stream_ptr) + self.stream.synchronize() + except BaseException: + for handle in reversed(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() + + self.table.numel() * self.table.element_size()) + + def execute_ptr(self, input_ptr, output_ptr, stream_ptr): + """Enqueue on the prepared stream; raw callers guarantee buffer metadata.""" + 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(output_ptr) % 8: + raise ValueError('The real output must be double-aligned') + residue = int(output_ptr) & 8 + _check(self.native_function(self.native_pointers[residue >> 3], input_ptr, int(output_ptr) - residue)) + + 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') + self.execute_ptr(a, b, (self.stream if stream is None else stream).cuda_stream) + return output + + def close(self): + if not self.closed: + for handle in reversed(self.handles): + handle.close() + self.intermediate = self.table = self.native = self.native_pointers = None + self.closed = True + + + + +_CACHE = {} +_LOCK = threading.Lock() diff --git a/build/torch-cuda/_b200/scalar_output_lto.py b/build/torch-cuda/_b200/scalar_output_lto.py new file mode 100644 index 0000000000000000000000000000000000000000..f8f8cef714ac202c9f641b00a82e9dbfe5ab7a9a --- /dev/null +++ b/build/torch-cuda/_b200/scalar_output_lto.py @@ -0,0 +1,114 @@ +# 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, n, rows, rowmajor, arch, nvrtc_version, kind, shift): + lib = library() + options = ('--std=c++17', '--relocatable-device-code=true', '-dlto', + f'--gpu-architecture=compute_{arch}', f'-DFFT_N={n}', + f'-DFFT_ROWS={rows}', f'-DFFT_ROWMAJOR={int(rowmajor)}', + f'-DFFT_CALLBACK_KIND={kind}', f'-DFFT_STORE_SHIFT={shift}') + program = ct.c_void_p() + _check(lib.nvrtcCreateProgram(ct.byref(program), source, b'reconstruct.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(n=n, rows=rows, rowmajor=bool(rowmajor), callback_kind=kind, store_shift=shift, 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_callbacks(n, rows, rowmajor, arch=100, *, store_shift=None): + """Compile distinct callback modules; None means a load-only plan.""" + if n < 16 or n % 2 or rows < 1 or store_shift not in (None, 0, 1): + raise ValueError('Invalid shape or scalar store shift') + source = Path(__file__).with_name('c2r_f64_narrow_index_inverse_scalar_output_reconstruct.cu').read_bytes() + result = [] + choices = [('load', 1, 1, 0)] + if store_shift is not None: + choices.append(('store', 5, 2, store_shift)) + with _LOCK: + for kind, enum, macro, shift in choices: + blob, metadata = _compile(source, int(n), int(rows), bool(rowmajor), int(arch), version(), macro, shift) + result.append((('fft_half_' + kind).encode(), enum, + ct.create_string_buffer(blob, len(blob)), dict(metadata, kind=kind))) + return result diff --git a/build/torch-cuda/_b200/wide_native_kernels.py b/build/torch-cuda/_b200/wide_native_kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..a98d9eccdefa66d2fe317f860d8603c2cd82a575 --- /dev/null +++ b/build/torch-cuda/_b200/wide_native_kernels.py @@ -0,0 +1,672 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Three-pass FFT with fused radix-9 and paired radix-2/3/4/5/7 arithmetic. + +Power-of-two padding exists only in register tensors required by Triton. +There is no padded device array, transpose pass, or input repacking pass. +Twiddle tables contain transform constants, never cached input/output data. +""" +import triton +import triton.language as tl + + +@triton.jit +def _input_index(index, N: tl.constexpr, RADICES: tl.constexpr, SIZES: tl.constexpr): + reverse = tl.full(index.shape, 0, tl.int32) + for stage in tl.static_range(len(RADICES)): + reverse += ((index // (SIZES[stage] // RADICES[stage])) % RADICES[stage]) * (N // SIZES[stage]) + return reverse + + +@triton.jit +def _fft_generic_stage(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADIX: tl.constexpr, SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Keep each stage's constants in its own JIT function scope. + + Annotated constexpr assignments inside a static loop are rejected on + the second unrolled iteration by Triton's CUDA frontend. Passing the + radix and size as constexpr arguments also keeps inner loops static. + """ + index = tl.arange(0, P) + PREVIOUS: tl.constexpr = SIZE // RADIX + base = index // SIZE * SIZE + index % PREVIOUS + # Padding lanes are never stored. Clamp their local gather indices. + source = tl.minimum(base, P - 1) + if STRIDED: + gather = tl.broadcast_to(source[:, None], (P, B)) + out_real = tl.gather(real, gather, 0) + out_imag = tl.gather(imag, gather, 0) + else: + gather = tl.broadcast_to(source[None, :], (B, P)) + out_real = tl.gather(real, gather, 1) + out_imag = tl.gather(imag, gather, 1) + for term in tl.static_range(1, RADIX): + source = tl.minimum(base + term * PREVIOUS, P - 1) + exponent = (term * (index % SIZE) * (N // SIZE)) % N + if N * TWIDDLE_STEP == 32: + wr, wi = _constant32_twiddle(exponent * TWIDDLE_STEP) + else: + wr = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP) + wi = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP + 1) + if INVERSE: + wi = -wi + if STRIDED: + gather = tl.broadcast_to(source[:, None], (P, B)) + vr = tl.gather(real, gather, 0) + vi = tl.gather(imag, gather, 0) + wr, wi = wr[:, None], wi[:, None] + else: + gather = tl.broadcast_to(source[None, :], (B, P)) + vr = tl.gather(real, gather, 1) + vi = tl.gather(imag, gather, 1) + wr, wi = wr[None, :], wi[None, :] + out_real += vr * wr - vi * wi + out_imag += vr * wi + vi * wr + return out_real, out_imag + + +@triton.jit +def _constant32_twiddle(exponent): + """Reconstruct accurately rounded 32nd roots using fixed constants.""" + k = exponent % 32 + c = tl.minimum(k % 16, 16 - k % 16) + q = tl.where(c <= 8, c, 16-c) + real = tl.full(q.shape, 0., tl.float64) + real = tl.where(q == 0, tl.full((), 1.0, tl.float64), real) + real = tl.where(q == 1, tl.full((), 0.9807852804032304, tl.float64), real) + real = tl.where(q == 2, tl.full((), 0.9238795325112867, tl.float64), real) + real = tl.where(q == 3, tl.full((), 0.8314696123025452, tl.float64), real) + real = tl.where(q == 4, tl.full((), 0.7071067811865476, tl.float64), real) + real = tl.where(q == 5, tl.full((), 0.5555702330196023, tl.float64), real) + real = tl.where(q == 6, tl.full((), 0.38268343236508984, tl.float64), real) + real = tl.where(q == 7, tl.full((), 0.19509032201612833, tl.float64), real) + real = tl.where(q == 8, tl.full((), 0.0, tl.float64), real) + real = tl.where((k > 8) & (k < 24), -real, real) + t = (k + 24) % 32 + q = tl.minimum(t % 16, 16 - t % 16) + imag = tl.full(q.shape, 0., tl.float64) + imag = tl.where(q == 0, tl.full((), 1.0, tl.float64), imag) + imag = tl.where(q == 1, tl.full((), 0.9807852804032304, tl.float64), imag) + imag = tl.where(q == 2, tl.full((), 0.9238795325112867, tl.float64), imag) + imag = tl.where(q == 3, tl.full((), 0.8314696123025452, tl.float64), imag) + imag = tl.where(q == 4, tl.full((), 0.7071067811865476, tl.float64), imag) + imag = tl.where(q == 5, tl.full((), 0.5555702330196023, tl.float64), imag) + imag = tl.where(q == 6, tl.full((), 0.38268343236508984, tl.float64), imag) + imag = tl.where(q == 7, tl.full((), 0.19509032201612833, tl.float64), imag) + imag = tl.where(q == 8, tl.full((), 0.0, tl.float64), imag) + imag = tl.where((t > 8) & (t < 24), imag, -imag) + return real, imag + + +@triton.jit +def _stage_twiddle(real, imag, Twiddle, N: tl.constexpr, + SIZE: tl.constexpr, RADIX: tl.constexpr, + TERM: tl.constexpr, STRIDED: tl.constexpr, + INVERSE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Apply one branch's twiddle; broadcast only across independent FFTs.""" + if SIZE == RADIX: + # All first-stage twiddles are one. No table access or multiplication. + return real, imag + else: + index = tl.arange(0, SIZE // RADIX) + exponent = TERM * index * (N // SIZE) + if N * TWIDDLE_STEP == 32: + wr, wi = _constant32_twiddle(exponent * TWIDDLE_STEP) + else: + wr = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP) + wi = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP + 1) + if INVERSE: + wi = -wi + if STRIDED: + wr, wi = wr[None, :, None], wi[None, :, None] + else: + wr, wi = wr[None, None, :], wi[None, None, :] + return real * wr - imag * wi, real * wi + imag * wr + + +@triton.jit +def _fft_radix2(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + H: tl.constexpr = SIZE // 2 + G: tl.constexpr = P // SIZE + if STRIDED: + r = tl.permute(tl.reshape(real, (G, 2, H, B)), (0, 2, 3, 1)) + i = tl.permute(tl.reshape(imag, (G, 2, H, B)), (0, 2, 3, 1)) + else: + r = tl.permute(tl.reshape(real, (B, G, 2, H)), (0, 1, 3, 2)) + i = tl.permute(tl.reshape(imag, (B, G, 2, H)), (0, 1, 3, 2)) + ar, br = tl.split(r) + ai, bi = tl.split(i) + br, bi = _stage_twiddle(br, bi, Twiddle, N, SIZE, 2, 1, + STRIDED, INVERSE, TWIDDLE_STEP) + r = tl.join(ar + br, ar - br) + i = tl.join(ai + bi, ai - bi) + if STRIDED: + return (tl.reshape(tl.permute(r, (0, 3, 1, 2)), (P, B)), + tl.reshape(tl.permute(i, (0, 3, 1, 2)), (P, B))) + else: + return (tl.reshape(tl.permute(r, (0, 1, 3, 2)), (B, P)), + tl.reshape(tl.permute(i, (0, 1, 3, 2)), (B, P))) + + +@triton.jit +def _fft_radix4(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + H: tl.constexpr = SIZE // 4 + G: tl.constexpr = P // SIZE + if STRIDED: + r = tl.permute(tl.reshape(real, (G, 4, H, B)), (0, 2, 3, 1)) + i = tl.permute(tl.reshape(imag, (G, 4, H, B)), (0, 2, 3, 1)) + r = tl.reshape(r, (G, H, B, 2, 2)) + i = tl.reshape(i, (G, H, B, 2, 2)) + else: + r = tl.permute(tl.reshape(real, (B, G, 4, H)), (0, 1, 3, 2)) + i = tl.permute(tl.reshape(imag, (B, G, 4, H)), (0, 1, 3, 2)) + r = tl.reshape(r, (B, G, H, 2, 2)) + i = tl.reshape(i, (B, G, H, 2, 2)) + even_r, odd_r = tl.split(r) + even_i, odd_i = tl.split(i) + ar, cr = tl.split(even_r) + ai, ci = tl.split(even_i) + br, dr = tl.split(odd_r) + bi, di = tl.split(odd_i) + br, bi = _stage_twiddle(br, bi, Twiddle, N, SIZE, 4, 1, + STRIDED, INVERSE, TWIDDLE_STEP) + cr, ci = _stage_twiddle(cr, ci, Twiddle, N, SIZE, 4, 2, + STRIDED, INVERSE, TWIDDLE_STEP) + dr, di = _stage_twiddle(dr, di, Twiddle, N, SIZE, 4, 3, + STRIDED, INVERSE, TWIDDLE_STEP) + ac_r, ac_i = ar + cr, ai + ci + bd_r, bd_i = br + dr, bi + di + diff_ac_r, diff_ac_i = ar - cr, ai - ci + diff_bd_r, diff_bd_i = br - dr, bi - di + # Multiplication by +/- i is a swap and sign, exactly in either precision. + if INVERSE: + turn_r, turn_i = -diff_bd_i, diff_bd_r + else: + turn_r, turn_i = diff_bd_i, -diff_bd_r + r0, i0 = ac_r + bd_r, ac_i + bd_i + r2, i2 = ac_r - bd_r, ac_i - bd_i + r1, i1 = diff_ac_r + turn_r, diff_ac_i + turn_i + r3, i3 = diff_ac_r - turn_r, diff_ac_i - turn_i + r = tl.join(tl.join(r0, r2), tl.join(r1, r3)) + i = tl.join(tl.join(i0, i2), tl.join(i1, i3)) + if STRIDED: + r = tl.reshape(r, (G, H, B, 4)) + i = tl.reshape(i, (G, H, B, 4)) + return (tl.reshape(tl.permute(r, (0, 3, 1, 2)), (P, B)), + tl.reshape(tl.permute(i, (0, 3, 1, 2)), (P, B))) + else: + r = tl.reshape(r, (B, G, H, 4)) + i = tl.reshape(i, (B, G, H, 4)) + return (tl.reshape(tl.permute(r, (0, 1, 3, 2)), (B, P)), + tl.reshape(tl.permute(i, (0, 1, 3, 2)), (B, P))) + + +@triton.jit +def _odd_branch(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADIX: tl.constexpr, SIZE: tl.constexpr, TERM: tl.constexpr, + HP: tl.constexpr, GP: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Read one branch into a smaller register tensor and apply its twiddle.""" + H: tl.constexpr = SIZE // RADIX + RP: tl.constexpr = triton.next_power_of_2(RADIX) + index = tl.arange(0, GP * HP) + local = index % HP + source = tl.minimum(index // HP * SIZE + local + TERM * H, P - 1) + if GP == 1 and HP == H and RP * HP == P: + # A final odd stage with power-of-two branch length is already a + # padded radix tile. Split its minor radix bits instead of gathering. + if STRIDED: + r = tl.permute(tl.reshape(real, (RP, HP, B)), (1, 2, 0)) + i = tl.permute(tl.reshape(imag, (RP, HP, B)), (1, 2, 0)) + if RP == 4: + r, i = tl.reshape(r, (HP, B, 2, 2)), tl.reshape(i, (HP, B, 2, 2)) + else: + r = tl.reshape(r, (HP, B, 2, 2, 2)) + i = tl.reshape(i, (HP, B, 2, 2, 2)) + else: + r = tl.permute(tl.reshape(real, (B, RP, HP)), (0, 2, 1)) + i = tl.permute(tl.reshape(imag, (B, RP, HP)), (0, 2, 1)) + if RP == 4: + r, i = tl.reshape(r, (B, HP, 2, 2)), tl.reshape(i, (B, HP, 2, 2)) + else: + r = tl.reshape(r, (B, HP, 2, 2, 2)) + i = tl.reshape(i, (B, HP, 2, 2, 2)) + lo_r, hi_r = tl.split(r) + lo_i, hi_i = tl.split(i) + if TERM % 2 == 0: + r, i = lo_r, lo_i + else: + r, i = hi_r, hi_i + lo_r, hi_r = tl.split(r) + lo_i, hi_i = tl.split(i) + if (TERM // 2) % 2 == 0: + r, i = lo_r, lo_i + else: + r, i = hi_r, hi_i + if RP == 8: + lo_r, hi_r = tl.split(r) + lo_i, hi_i = tl.split(i) + if TERM // 4 == 0: + r, i = lo_r, lo_i + else: + r, i = hi_r, hi_i + elif STRIDED: + gather = tl.broadcast_to(source[:, None], (GP * HP, B)) + r = tl.gather(real, gather, 0) + i = tl.gather(imag, gather, 0) + else: + gather = tl.broadcast_to(source[None, :], (B, GP * HP)) + r = tl.gather(real, gather, 1) + i = tl.gather(imag, gather, 1) + if TERM != 0 and SIZE != RADIX: + exponent = TERM * local * (N // SIZE) + valid = (local < H) & (index // HP < N // SIZE) + wr = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP, valid, 0) + wi = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP + 1, valid, 0) + if INVERSE: + wi = -wi + if STRIDED: + wr, wi = wr[:, None], wi[:, None] + else: + wr, wi = wr[None, :], wi[None, :] + return r * wr - i * wi, r * wi + i * wr + else: + return r, i + + +@triton.jit +def _join4(a, b, c, d): + return tl.reshape(tl.join(tl.join(a, c), tl.join(b, d)), a.shape + (4,)) + + +@triton.jit +def _join8(a, b, c, d, e, f, g, h): + return tl.reshape(tl.join(_join4(a, c, e, g), _join4(b, d, f, h)), + a.shape + (8,)) + + +@triton.jit +def _odd_output(outputs, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, RADIX: tl.constexpr, + SIZE: tl.constexpr, HP: tl.constexpr, GP: tl.constexpr): + """Join the small branches and gather natural ordering only in registers.""" + RP: tl.constexpr = triton.next_power_of_2(RADIX) + H: tl.constexpr = SIZE // RADIX + zero = tl.full(outputs[0].shape, 0, outputs[0].dtype) + # outputs = [DC, +1, -1, +2, -2, ...]. Natural branch order follows. + if RADIX == 3: + data = _join4(outputs[0], outputs[1], outputs[2], zero) + elif RADIX == 5: + data = _join8(outputs[0], outputs[1], outputs[3], outputs[4], outputs[2], + zero, zero, zero) + else: + data = _join8(outputs[0], outputs[1], outputs[3], outputs[5], + outputs[6], outputs[4], outputs[2], zero) + if STRIDED: + data = tl.reshape(data, (GP, HP, B, RP)) + data = tl.reshape(tl.permute(data, (0, 3, 1, 2)), (GP * RP * HP, B)) + else: + data = tl.reshape(data, (B, GP, HP, RP)) + data = tl.reshape(tl.permute(data, (0, 1, 3, 2)), (B, GP * RP * HP)) + if GP == 1 and HP == H and RP * HP == P: + return data + index = tl.arange(0, P) + source = (index // SIZE * RP * HP + (index % SIZE) // H * HP + index % H) + source = tl.minimum(source, GP * RP * HP - 1) + if STRIDED: + return tl.gather(data, tl.broadcast_to(source[:, None], (P, B)), 0) + else: + return tl.gather(data, tl.broadcast_to(source[None, :], (B, P)), 1) + + +@triton.jit +def _fft_odd(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADIX: tl.constexpr, SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Paired DFT on compact branch tensors, retaining the input precision. + + v[j] + v[r-j] contributes cosine terms to a conjugate output pair; + v[j] - v[r-j] contributes sine terms with opposite signs. This reduces + duplicated arithmetic across output branches. Padding is register-only. + """ + HP: tl.constexpr = triton.next_power_of_2(SIZE // RADIX) + GP: tl.constexpr = triton.next_power_of_2(N // SIZE) + ar, ai = _odd_branch(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADIX, SIZE, 0, HP, GP, TWIDDLE_STEP) + sums_r, sums_i, diffs_r, diffs_i = (), (), (), () + dc_r, dc_i = ar, ai + for j in tl.static_range(1, (RADIX + 1) // 2): + br, bi = _odd_branch(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADIX, SIZE, j, HP, GP, TWIDDLE_STEP) + cr, ci = _odd_branch(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADIX, SIZE, RADIX - j, HP, GP, TWIDDLE_STEP) + sr, si, dr, di = br + cr, bi + ci, br - cr, bi - ci + sums_r, sums_i = sums_r + (sr,), sums_i + (si,) + diffs_r, diffs_i = diffs_r + (dr,), diffs_i + (di,) + dc_r, dc_i = dc_r + sr, dc_i + si + outputs_r, outputs_i = (dc_r,), (dc_i,) + for k in tl.static_range(1, (RADIX + 1) // 2): + even_r, even_i = ar, ai + odd_r = tl.full(ar.shape, 0, ar.dtype) + odd_i = tl.full(ai.shape, 0, ai.dtype) + for j in tl.static_range(1, (RADIX + 1) // 2): + exponent = (j * k % RADIX) * (N // RADIX) + c = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP) + s = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP + 1) + if INVERSE: + s = -s + even_r = even_r + c * sums_r[j - 1] + even_i = even_i + c * sums_i[j - 1] + odd_r = odd_r + s * diffs_r[j - 1] + odd_i = odd_i + s * diffs_i[j - 1] + outputs_r = outputs_r + (even_r - odd_i, even_r + odd_i) + outputs_i = outputs_i + (even_i + odd_r, even_i - odd_r) + return (_odd_output(outputs_r, N, P, B, STRIDED, RADIX, SIZE, HP, GP), + _odd_output(outputs_i, N, P, B, STRIDED, RADIX, SIZE, HP, GP)) + + +@triton.jit +def _radix9_branch(data, H: tl.constexpr, B: tl.constexpr, + STRIDED: tl.constexpr, TERM: tl.constexpr): + """Select one of sixteen padded register branches using four splits.""" + if STRIDED: + value = tl.permute(tl.reshape(data, (16, H, B)), (1, 2, 0)) + value = tl.reshape(value, (H, B, 2, 2, 2, 2)) + else: + value = tl.permute(tl.reshape(data, (B, 16, H)), (0, 2, 1)) + value = tl.reshape(value, (B, H, 2, 2, 2, 2)) + for bit in tl.static_range(4): + lo, hi = tl.split(value) + if (TERM >> bit) & 1: + value = hi + else: + value = lo + return value + + +@triton.jit +def _dft3(ar, ai, br, bi, cr, ci, sine): + """Three-point DFT with an exactly representable -1/2 cosine.""" + sr, si = br + cr, bi + ci + dr, di = br - cr, bi - ci + er, ei = ar - 0.5 * sr, ai - 0.5 * si + return (ar + sr, ai + si, + er - sine * di, ei + sine * dr, + er + sine * di, ei - sine * dr) + + +@triton.jit +def _join16(outputs): + zero = tl.full(outputs[0].shape, 0, outputs[0].dtype) + evens = _join8(outputs[0], outputs[2], outputs[4], outputs[6], outputs[8], + zero, zero, zero) + odds = _join8(outputs[1], outputs[3], outputs[5], outputs[7], + zero, zero, zero, zero) + return tl.reshape(tl.join(evens, odds), outputs[0].shape + (16,)) + + +@triton.jit +def _fft_radix9(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Fuse two final radix-3 stages without an intermediate register gather. + + After the stage twiddles, compute three DFTs over indices (b, b+3, + b+6), multiply branch b and frequency k1 by W9**(b*k1), then compute + three DFTs over b. Natural output frequency is k1 + 3*k2. + Padding is confined to register tensors; all global rows stay packed. + """ + H: tl.constexpr = SIZE // 9 + tl.static_assert(SIZE == N and H == triton.next_power_of_2(H) and P == 16 * H) + local = tl.arange(0, H) + values_r, values_i = (), () + for j in tl.static_range(9): + r = _radix9_branch(real, H, B, STRIDED, j) + i = _radix9_branch(imag, H, B, STRIDED, j) + if j != 0 and H != 1: + wr = tl.load(Twiddle + 2 * j * local * TWIDDLE_STEP) + wi = tl.load(Twiddle + 2 * j * local * TWIDDLE_STEP + 1) + if INVERSE: + wi = -wi + if STRIDED: + wr, wi = wr[:, None], wi[:, None] + else: + wr, wi = wr[None, :], wi[None, :] + r, i = r * wr - i * wi, r * wi + i * wr + values_r, values_i = values_r + (r,), values_i + (i,) + sine = tl.load(Twiddle + 2 * (N // 3) * TWIDDLE_STEP + 1) + if INVERSE: + sine = -sine + columns = () + for b in tl.static_range(3): + columns = columns + (_dft3(values_r[b], values_i[b], + values_r[b + 3], values_i[b + 3], + values_r[b + 6], values_i[b + 6], sine),) + rows = () + for k1 in tl.static_range(3): + branches_r, branches_i = (), () + for b in tl.static_range(3): + r, i = columns[b][2 * k1], columns[b][2 * k1 + 1] + if b != 0 and k1 != 0: + wr = tl.load(Twiddle + 2 * (b * k1 * H) * TWIDDLE_STEP) + wi = tl.load(Twiddle + 2 * (b * k1 * H) * TWIDDLE_STEP + 1) + if INVERSE: + wi = -wi + r, i = r * wr - i * wi, r * wi + i * wr + branches_r, branches_i = branches_r + (r,), branches_i + (i,) + rows = rows + (_dft3(branches_r[0], branches_i[0], + branches_r[1], branches_i[1], + branches_r[2], branches_i[2], sine),) + outputs_r, outputs_i = (), () + for k2 in tl.static_range(3): + for k1 in tl.static_range(3): + outputs_r = outputs_r + (rows[k1][2 * k2],) + outputs_i = outputs_i + (rows[k1][2 * k2 + 1],) + r, i = _join16(outputs_r), _join16(outputs_i) + if STRIDED: + return (tl.reshape(tl.permute(r, (2, 0, 1)), (P, B)), + tl.reshape(tl.permute(i, (2, 0, 1)), (P, B))) + else: + return (tl.reshape(tl.permute(r, (0, 2, 1)), (B, P)), + tl.reshape(tl.permute(i, (0, 2, 1)), (B, P))) + + +@triton.jit +def _fft_stage(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADIX: tl.constexpr, SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + # Factors 4 and 2 precede odd factors, so these reshapes divide P exactly. + if RADIX == 4 and P % SIZE == 0: + return _fft_radix4(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + SIZE, TWIDDLE_STEP) + elif RADIX == 2 and P % SIZE == 0: + return _fft_radix2(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + SIZE, TWIDDLE_STEP) + elif RADIX == 9: + return _fft_radix9(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + SIZE, TWIDDLE_STEP) + elif RADIX == 3 or RADIX == 5 or RADIX == 7: + return _fft_odd(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADIX, SIZE, TWIDDLE_STEP) + else: + return _fft_generic_stage(real, imag, Twiddle, N, P, B, STRIDED, + INVERSE, RADIX, SIZE, TWIDDLE_STEP) + + +@triton.jit +def _fft(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr, TWIDDLE_STEP: tl.constexpr = 1): + for stage in tl.static_range(len(RADICES)): + real, imag = _fft_stage(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADICES[stage], SIZES[stage], TWIDDLE_STEP) + return real, imag + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def real_axis(Input, Output, Twiddle, ROWS: tl.constexpr, N: tl.constexpr, + P: tl.constexpr, B: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr): + rows = tl.program_id(0) * B + tl.arange(0, B) + index = tl.arange(0, P) + source = _input_index(index, N, RADICES, SIZES) + NF: tl.constexpr = N // 2 + 1 + mask = (rows[:, None] < ROWS) & (index[None, :] < N) + if INVERSE: + half_index = tl.minimum(source, N - source) + offsets = 2 * (rows[:, None] * NF + half_index[None, :]) + real = tl.load(Input + offsets, mask, 0) + imag = tl.load(Input + offsets + 1, mask, 0) + imag = tl.where(source[None, :] >= NF, -imag, imag) + # Real end points after inverse Z/Y implement irfftn's convention + # even when the original packed input is not Hermitian-consistent. + imag = tl.where((source[None, :] == 0) | (2 * source[None, :] == N), 0, imag) + else: + real = tl.load(Input + rows[:, None] * N + source[None, :], mask, 0) + imag = tl.full((B, P), 0, real.dtype) + real, imag = _fft(real, imag, Twiddle, N, P, B, False, INVERSE, RADICES, SIZES) + if INVERSE: + tl.store(Output + rows[:, None] * N + index[None, :], real, mask) + else: + offsets = 2 * (rows[:, None] * NF + index[None, :]) + output_mask = (rows[:, None] < ROWS) & (index[None, :] < NF) + tl.store(Output + offsets, real, output_mask) + tl.store(Output + offsets + 1, imag, output_mask) + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def packed_real_axis(Input, Output, Twiddle, ROWS: tl.constexpr, N: tl.constexpr, + P: tl.constexpr, B: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr): + """Even-length real FFT using an N/2 complex FFT and fused reconstruction. + + RADICES, SIZES and P describe M=N/2; Twiddle contains N complex constants. + Even/odd samples are paired only in registers; global layout stays exact. + The inverse constructs twice the packed complex spectrum so that the + length-M unnormalized inverse returns the required length-N scaling. + """ + tl.static_assert(N % 2 == 0) + M: tl.constexpr = N // 2 + NF: tl.constexpr = M + 1 + rows = tl.program_id(0) * B + tl.arange(0, B) + index = tl.arange(0, P) + source = _input_index(index, M, RADICES, SIZES) + mask = (rows[:, None] < ROWS) & (index[None, :] < M) + if INVERSE: + offsets = 2 * (rows[:, None] * NF + source[None, :]) + mirrors = 2 * (rows[:, None] * NF + M - source[None, :]) + ar = tl.load(Input + offsets, mask, 0) + ai = tl.load(Input + offsets + 1, mask, 0) + br = tl.load(Input + mirrors, mask, 0) + bi = tl.load(Input + mirrors + 1, mask, 0) + # irfftn ignores real-axis endpoint imaginary parts AFTER Z/Y inverse. + # Arbitrary complex packed input is valid; do not zero planes earlier. + ai = tl.where(source[None, :] == 0, 0, ai) + bi = tl.where(source[None, :] == 0, 0, bi) + wr = tl.load(Twiddle + 2 * source, index < M, 0)[None, :] + wi = tl.load(Twiddle + 2 * source + 1, index < M, 0)[None, :] + # 2 Z[k] = X[k] + conj(X[M-k]) + # + i conj(W_N^k) (X[k] - conj(X[M-k])). + real = ar + br - wr * (ai + bi) + wi * (ar - br) + imag = ai - bi + wr * (ar - br) + wi * (ai + bi) + else: + offsets = rows[:, None] * N + 2 * source[None, :] + real = tl.load(Input + offsets, mask, 0) + imag = tl.load(Input + offsets + 1, mask, 0) + real, imag = _fft(real, imag, Twiddle, M, P, B, False, INVERSE, + RADICES, SIZES, 2) + if INVERSE: + offsets = rows[:, None] * N + 2 * index[None, :] + tl.store(Output + offsets, real, mask) + tl.store(Output + offsets + 1, imag, mask) + else: + mirror_index = tl.where(index < M, (M - index) % M, 0) + mirror = tl.broadcast_to(mirror_index[None, :], (B, P)) + br = tl.gather(real, mirror, 1) + bi = tl.gather(imag, mirror, 1) + wr = tl.load(Twiddle + 2 * index, index < M, 0)[None, :] + wi = tl.load(Twiddle + 2 * index + 1, index < M, 0)[None, :] + # X[k] = (Z[k] + conj(Z[-k]))/2 + # - i W_N^k (Z[k] - conj(Z[-k]))/2. + out_real = 0.5 * (real + br + wr * (imag + bi) + wi * (real - br)) + out_imag = 0.5 * (imag - bi + wi * (imag + bi) - wr * (real - br)) + offsets = 2 * (rows[:, None] * NF + index[None, :]) + tl.store(Output + offsets, out_real, mask) + tl.store(Output + offsets + 1, out_imag, mask) + # The k=0 lane owns the extra, tightly packed Nyquist coefficient. + endpoints = (rows[:, None] < ROWS) & (index[None, :] == 0) + nyquist = 2 * (rows[:, None] * NF + M) + tl.store(Output + nyquist + 0 * index[None, :], real - imag, endpoints) + tl.store(Output + nyquist + 0 * index[None, :] + 1, 0, endpoints) + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def complex_axis(Input, Output, Twiddle, N: tl.constexpr, INNER: tl.constexpr, + P: tl.constexpr, B: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr): + # Program dimension 1 indexes independent outer batches (Z for Y FFTs). + columns = tl.program_id(0) * B + tl.arange(0, B) + index = tl.arange(0, P) + source = _input_index(index, N, RADICES, SIZES) + base = tl.program_id(1) * N * INNER + offsets = 2 * (base + source[:, None] * INNER + columns[None, :]) + mask = (index[:, None] < N) & (columns[None, :] < INNER) + real = tl.load(Input + offsets, mask, 0).to(tl.float64) + imag = tl.load(Input + offsets + 1, mask, 0).to(tl.float64) + real, imag = _fft(real, imag, Twiddle, N, P, B, True, INVERSE, RADICES, SIZES) + offsets = 2 * (base + index[:, None] * INNER + columns[None, :]) + tl.store(Output + offsets, real, mask) + tl.store(Output + offsets + 1, imag, mask) + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def real_xy(Input, Output, TwiddleX, TwiddleY, + NY: tl.constexpr, NX: tl.constexpr, PY: tl.constexpr, PX: tl.constexpr, + INVERSE: tl.constexpr, RX: tl.constexpr, SX: tl.constexpr, + RY: tl.constexpr, SY: tl.constexpr): + """Fuse the real X and complex Y FFTs for one complete Z plane. + + Both FFT networks operate on registers/shared memory in the same CTA. + Only caller layout is accessed globally: real NY*NX or packed NY*NF*2. + An inverse call consumes a plane AFTER the Z inverse. Arbitrary boundary + values are projected only after the Y inverse, matching irfftn. + """ + z = tl.program_id(0) + y = tl.arange(0, PY) + x = tl.arange(0, PX) + sy = _input_index(y, NY, RY, SY) + sx = _input_index(x, NX, RX, SX) + NF: tl.constexpr = NX // 2 + 1 + if INVERSE: + offsets = 2 * ((z * NY + sy[:, None]) * NF + x[None, :]) + valid = (y[:, None] < NY) & (x[None, :] < NF) + real = tl.load(Input + offsets, valid, 0).to(tl.float64) + imag = tl.load(Input + offsets + 1, valid, 0).to(tl.float64) + real, imag = _fft(real, imag, TwiddleY, NY, PY, PX, True, True, RY, SY) + # Complete and permute each X spectrum only after inverse Y. No + # original endpoint plane is discarded or assumed Hermitian. + half = tl.minimum(sx, NX - sx) + half = tl.where(x < NX, half, 0) + gather = tl.broadcast_to(half[None, :], (PY, PX)) + real = tl.gather(real, gather, 1) + imag = tl.gather(imag, gather, 1) + imag = tl.where(sx[None, :] >= NF, -imag, imag) + imag = tl.where((sx[None, :] == 0) | (2 * sx[None, :] == NX), 0, imag) + real, imag = _fft(real, imag, TwiddleX, NX, PX, PY, False, True, RX, SX) + offsets = (z * NY + y[:, None]) * NX + x[None, :] + tl.store(Output + offsets, real, (y[:, None] < NY) & (x[None, :] < NX)) + else: + offsets = (z * NY + sy[:, None]) * NX + sx[None, :] + valid = (y[:, None] < NY) & (x[None, :] < NX) + real = tl.load(Input + offsets, valid, 0).to(tl.float64) + imag = tl.full((PY, PX), 0, real.dtype) + real, imag = _fft(real, imag, TwiddleX, NX, PX, PY, False, False, RX, SX) + real, imag = _fft(real, imag, TwiddleY, NY, PY, PX, True, False, RY, SY) + offsets = 2 * ((z * NY + y[:, None]) * NF + x[None, :]) + valid = (y[:, None] < NY) & (x[None, :] < NF) + tl.store(Output + offsets, real, valid) + tl.store(Output + offsets + 1, imag, valid) diff --git a/build/torch-cuda/_b200/wide_native_schedule.py b/build/torch-cuda/_b200/wide_native_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..b32bdabca052fd8d8550200d5c1b1d33a76d9fbf --- /dev/null +++ b/build/torch-cuda/_b200/wide_native_schedule.py @@ -0,0 +1,74 @@ +# 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") + # The 32-point radix-2 network avoids costly radix-4 layout exchanges. + if remaining == 32: + return (2, 2, 2, 2, 2) + 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 diff --git a/build/torch-cuda/_b200/wide_redirect.cu b/build/torch-cuda/_b200/wide_redirect.cu new file mode 100644 index 0000000000000000000000000000000000000000..7fa7677cdd2b843008b30590a4c3cd719f9c327f --- /dev/null +++ b/build/torch-cuda/_b200/wide_redirect.cu @@ -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 float *input, fft_index_t index) { + double2 value; +#if FFT_INPUT_SHIFT % 2 == 0 + const float2 pair = reinterpret_cast(input)[index]; + value.x = static_cast(pair.x); + value.y = static_cast(pair.y); +#else + value.x = static_cast(input[2ULL * index]); + value.y = static_cast(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(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 float *input = static_cast(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(caller_info)[1]; + const double2 w = reinterpret_cast(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(caller_info); + float *output = reinterpret_cast(address); + output[2 * offset] = static_cast(value.x); + output[2 * offset + 1] = static_cast(value.y); +} +#endif diff --git a/build/torch-cuda/_b200/y_residue_inverse_kernels.py b/build/torch-cuda/_b200/y_residue_inverse_kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..639a52a14d453dbfbc8a75446439336fd6d93185 --- /dev/null +++ b/build/torch-cuda/_b200/y_residue_inverse_kernels.py @@ -0,0 +1,722 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Three-pass FFT with fused radix-9 and paired radix-2/3/4/5/7 arithmetic. + +Power-of-two padding exists only in register tensors required by Triton. +There is no padded device array, transpose pass, or input repacking pass. +Twiddle tables contain transform constants, never cached input/output data. +""" +import triton +import triton.language as tl + + +@triton.jit +def _input_index(index, N: tl.constexpr, RADICES: tl.constexpr, SIZES: tl.constexpr): + reverse = tl.full(index.shape, 0, tl.int32) + for stage in tl.static_range(len(RADICES)): + reverse += ((index // (SIZES[stage] // RADICES[stage])) % RADICES[stage]) * (N // SIZES[stage]) + return reverse + + +@triton.jit +def _fft_generic_stage(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADIX: tl.constexpr, SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Keep each stage's constants in its own JIT function scope. + + Annotated constexpr assignments inside a static loop are rejected on + the second unrolled iteration by Triton's CUDA frontend. Passing the + radix and size as constexpr arguments also keeps inner loops static. + """ + index = tl.arange(0, P) + PREVIOUS: tl.constexpr = SIZE // RADIX + base = index // SIZE * SIZE + index % PREVIOUS + # Padding lanes are never stored. Clamp their local gather indices. + source = tl.minimum(base, P - 1) + if STRIDED: + gather = tl.broadcast_to(source[:, None], (P, B)) + out_real = tl.gather(real, gather, 0) + out_imag = tl.gather(imag, gather, 0) + else: + gather = tl.broadcast_to(source[None, :], (B, P)) + out_real = tl.gather(real, gather, 1) + out_imag = tl.gather(imag, gather, 1) + for term in tl.static_range(1, RADIX): + source = tl.minimum(base + term * PREVIOUS, P - 1) + exponent = (term * (index % SIZE) * (N // SIZE)) % N + if N * TWIDDLE_STEP == 32: + wr, wi = _constant32_twiddle(exponent * TWIDDLE_STEP) + else: + wr = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP) + wi = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP + 1) + if INVERSE: + wi = -wi + if STRIDED: + gather = tl.broadcast_to(source[:, None], (P, B)) + vr = tl.gather(real, gather, 0) + vi = tl.gather(imag, gather, 0) + wr, wi = wr[:, None], wi[:, None] + else: + gather = tl.broadcast_to(source[None, :], (B, P)) + vr = tl.gather(real, gather, 1) + vi = tl.gather(imag, gather, 1) + wr, wi = wr[None, :], wi[None, :] + out_real += vr * wr - vi * wi + out_imag += vr * wi + vi * wr + return out_real, out_imag + + +@triton.jit +def _constant32_twiddle(exponent): + """Reconstruct accurately rounded 32nd roots using fixed constants.""" + k = exponent % 32 + c = tl.minimum(k % 16, 16 - k % 16) + q = tl.where(c <= 8, c, 16-c) + real = tl.full(q.shape, 0., tl.float64) + real = tl.where(q == 0, tl.full((), 1.0, tl.float64), real) + real = tl.where(q == 1, tl.full((), 0.9807852804032304, tl.float64), real) + real = tl.where(q == 2, tl.full((), 0.9238795325112867, tl.float64), real) + real = tl.where(q == 3, tl.full((), 0.8314696123025452, tl.float64), real) + real = tl.where(q == 4, tl.full((), 0.7071067811865476, tl.float64), real) + real = tl.where(q == 5, tl.full((), 0.5555702330196023, tl.float64), real) + real = tl.where(q == 6, tl.full((), 0.38268343236508984, tl.float64), real) + real = tl.where(q == 7, tl.full((), 0.19509032201612833, tl.float64), real) + real = tl.where(q == 8, tl.full((), 0.0, tl.float64), real) + real = tl.where((k > 8) & (k < 24), -real, real) + t = (k + 24) % 32 + q = tl.minimum(t % 16, 16 - t % 16) + imag = tl.full(q.shape, 0., tl.float64) + imag = tl.where(q == 0, tl.full((), 1.0, tl.float64), imag) + imag = tl.where(q == 1, tl.full((), 0.9807852804032304, tl.float64), imag) + imag = tl.where(q == 2, tl.full((), 0.9238795325112867, tl.float64), imag) + imag = tl.where(q == 3, tl.full((), 0.8314696123025452, tl.float64), imag) + imag = tl.where(q == 4, tl.full((), 0.7071067811865476, tl.float64), imag) + imag = tl.where(q == 5, tl.full((), 0.5555702330196023, tl.float64), imag) + imag = tl.where(q == 6, tl.full((), 0.38268343236508984, tl.float64), imag) + imag = tl.where(q == 7, tl.full((), 0.19509032201612833, tl.float64), imag) + imag = tl.where(q == 8, tl.full((), 0.0, tl.float64), imag) + imag = tl.where((t > 8) & (t < 24), imag, -imag) + return real, imag + + +@triton.jit +def _stage_twiddle(real, imag, Twiddle, N: tl.constexpr, + SIZE: tl.constexpr, RADIX: tl.constexpr, + TERM: tl.constexpr, STRIDED: tl.constexpr, + INVERSE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Apply one branch's twiddle; broadcast only across independent FFTs.""" + if SIZE == RADIX: + # All first-stage twiddles are one. No table access or multiplication. + return real, imag + else: + index = tl.arange(0, SIZE // RADIX) + exponent = TERM * index * (N // SIZE) + if N * TWIDDLE_STEP == 32: + wr, wi = _constant32_twiddle(exponent * TWIDDLE_STEP) + else: + wr = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP) + wi = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP + 1) + if INVERSE: + wi = -wi + if STRIDED: + wr, wi = wr[None, :, None], wi[None, :, None] + else: + wr, wi = wr[None, None, :], wi[None, None, :] + return real * wr - imag * wi, real * wi + imag * wr + + +@triton.jit +def _fft_radix2(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + H: tl.constexpr = SIZE // 2 + G: tl.constexpr = P // SIZE + if STRIDED: + r = tl.permute(tl.reshape(real, (G, 2, H, B)), (0, 2, 3, 1)) + i = tl.permute(tl.reshape(imag, (G, 2, H, B)), (0, 2, 3, 1)) + else: + r = tl.permute(tl.reshape(real, (B, G, 2, H)), (0, 1, 3, 2)) + i = tl.permute(tl.reshape(imag, (B, G, 2, H)), (0, 1, 3, 2)) + ar, br = tl.split(r) + ai, bi = tl.split(i) + br, bi = _stage_twiddle(br, bi, Twiddle, N, SIZE, 2, 1, + STRIDED, INVERSE, TWIDDLE_STEP) + r = tl.join(ar + br, ar - br) + i = tl.join(ai + bi, ai - bi) + if STRIDED: + return (tl.reshape(tl.permute(r, (0, 3, 1, 2)), (P, B)), + tl.reshape(tl.permute(i, (0, 3, 1, 2)), (P, B))) + else: + return (tl.reshape(tl.permute(r, (0, 1, 3, 2)), (B, P)), + tl.reshape(tl.permute(i, (0, 1, 3, 2)), (B, P))) + + +@triton.jit +def _fft_radix4(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + H: tl.constexpr = SIZE // 4 + G: tl.constexpr = P // SIZE + if STRIDED: + r = tl.permute(tl.reshape(real, (G, 4, H, B)), (0, 2, 3, 1)) + i = tl.permute(tl.reshape(imag, (G, 4, H, B)), (0, 2, 3, 1)) + r = tl.reshape(r, (G, H, B, 2, 2)) + i = tl.reshape(i, (G, H, B, 2, 2)) + else: + r = tl.permute(tl.reshape(real, (B, G, 4, H)), (0, 1, 3, 2)) + i = tl.permute(tl.reshape(imag, (B, G, 4, H)), (0, 1, 3, 2)) + r = tl.reshape(r, (B, G, H, 2, 2)) + i = tl.reshape(i, (B, G, H, 2, 2)) + even_r, odd_r = tl.split(r) + even_i, odd_i = tl.split(i) + ar, cr = tl.split(even_r) + ai, ci = tl.split(even_i) + br, dr = tl.split(odd_r) + bi, di = tl.split(odd_i) + br, bi = _stage_twiddle(br, bi, Twiddle, N, SIZE, 4, 1, + STRIDED, INVERSE, TWIDDLE_STEP) + cr, ci = _stage_twiddle(cr, ci, Twiddle, N, SIZE, 4, 2, + STRIDED, INVERSE, TWIDDLE_STEP) + dr, di = _stage_twiddle(dr, di, Twiddle, N, SIZE, 4, 3, + STRIDED, INVERSE, TWIDDLE_STEP) + ac_r, ac_i = ar + cr, ai + ci + bd_r, bd_i = br + dr, bi + di + diff_ac_r, diff_ac_i = ar - cr, ai - ci + diff_bd_r, diff_bd_i = br - dr, bi - di + # Multiplication by +/- i is a swap and sign, exactly in either precision. + if INVERSE: + turn_r, turn_i = -diff_bd_i, diff_bd_r + else: + turn_r, turn_i = diff_bd_i, -diff_bd_r + r0, i0 = ac_r + bd_r, ac_i + bd_i + r2, i2 = ac_r - bd_r, ac_i - bd_i + r1, i1 = diff_ac_r + turn_r, diff_ac_i + turn_i + r3, i3 = diff_ac_r - turn_r, diff_ac_i - turn_i + r = tl.join(tl.join(r0, r2), tl.join(r1, r3)) + i = tl.join(tl.join(i0, i2), tl.join(i1, i3)) + if STRIDED: + r = tl.reshape(r, (G, H, B, 4)) + i = tl.reshape(i, (G, H, B, 4)) + return (tl.reshape(tl.permute(r, (0, 3, 1, 2)), (P, B)), + tl.reshape(tl.permute(i, (0, 3, 1, 2)), (P, B))) + else: + r = tl.reshape(r, (B, G, H, 4)) + i = tl.reshape(i, (B, G, H, 4)) + return (tl.reshape(tl.permute(r, (0, 1, 3, 2)), (B, P)), + tl.reshape(tl.permute(i, (0, 1, 3, 2)), (B, P))) + + +@triton.jit +def _odd_branch(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADIX: tl.constexpr, SIZE: tl.constexpr, TERM: tl.constexpr, + HP: tl.constexpr, GP: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Read one branch into a smaller register tensor and apply its twiddle.""" + H: tl.constexpr = SIZE // RADIX + RP: tl.constexpr = triton.next_power_of_2(RADIX) + index = tl.arange(0, GP * HP) + local = index % HP + source = tl.minimum(index // HP * SIZE + local + TERM * H, P - 1) + if GP == 1 and HP == H and RP * HP == P: + # A final odd stage with power-of-two branch length is already a + # padded radix tile. Split its minor radix bits instead of gathering. + if STRIDED: + r = tl.permute(tl.reshape(real, (RP, HP, B)), (1, 2, 0)) + i = tl.permute(tl.reshape(imag, (RP, HP, B)), (1, 2, 0)) + if RP == 4: + r, i = tl.reshape(r, (HP, B, 2, 2)), tl.reshape(i, (HP, B, 2, 2)) + else: + r = tl.reshape(r, (HP, B, 2, 2, 2)) + i = tl.reshape(i, (HP, B, 2, 2, 2)) + else: + r = tl.permute(tl.reshape(real, (B, RP, HP)), (0, 2, 1)) + i = tl.permute(tl.reshape(imag, (B, RP, HP)), (0, 2, 1)) + if RP == 4: + r, i = tl.reshape(r, (B, HP, 2, 2)), tl.reshape(i, (B, HP, 2, 2)) + else: + r = tl.reshape(r, (B, HP, 2, 2, 2)) + i = tl.reshape(i, (B, HP, 2, 2, 2)) + lo_r, hi_r = tl.split(r) + lo_i, hi_i = tl.split(i) + if TERM % 2 == 0: + r, i = lo_r, lo_i + else: + r, i = hi_r, hi_i + lo_r, hi_r = tl.split(r) + lo_i, hi_i = tl.split(i) + if (TERM // 2) % 2 == 0: + r, i = lo_r, lo_i + else: + r, i = hi_r, hi_i + if RP == 8: + lo_r, hi_r = tl.split(r) + lo_i, hi_i = tl.split(i) + if TERM // 4 == 0: + r, i = lo_r, lo_i + else: + r, i = hi_r, hi_i + elif STRIDED: + gather = tl.broadcast_to(source[:, None], (GP * HP, B)) + r = tl.gather(real, gather, 0) + i = tl.gather(imag, gather, 0) + else: + gather = tl.broadcast_to(source[None, :], (B, GP * HP)) + r = tl.gather(real, gather, 1) + i = tl.gather(imag, gather, 1) + if TERM != 0 and SIZE != RADIX: + exponent = TERM * local * (N // SIZE) + valid = (local < H) & (index // HP < N // SIZE) + wr = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP, valid, 0) + wi = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP + 1, valid, 0) + if INVERSE: + wi = -wi + if STRIDED: + wr, wi = wr[:, None], wi[:, None] + else: + wr, wi = wr[None, :], wi[None, :] + return r * wr - i * wi, r * wi + i * wr + else: + return r, i + + +@triton.jit +def _join4(a, b, c, d): + return tl.reshape(tl.join(tl.join(a, c), tl.join(b, d)), a.shape + (4,)) + + +@triton.jit +def _join8(a, b, c, d, e, f, g, h): + return tl.reshape(tl.join(_join4(a, c, e, g), _join4(b, d, f, h)), + a.shape + (8,)) + + +@triton.jit +def _odd_output(outputs, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, RADIX: tl.constexpr, + SIZE: tl.constexpr, HP: tl.constexpr, GP: tl.constexpr): + """Join the small branches and gather natural ordering only in registers.""" + RP: tl.constexpr = triton.next_power_of_2(RADIX) + H: tl.constexpr = SIZE // RADIX + zero = tl.full(outputs[0].shape, 0, outputs[0].dtype) + # outputs = [DC, +1, -1, +2, -2, ...]. Natural branch order follows. + if RADIX == 3: + data = _join4(outputs[0], outputs[1], outputs[2], zero) + elif RADIX == 5: + data = _join8(outputs[0], outputs[1], outputs[3], outputs[4], outputs[2], + zero, zero, zero) + else: + data = _join8(outputs[0], outputs[1], outputs[3], outputs[5], + outputs[6], outputs[4], outputs[2], zero) + if STRIDED: + data = tl.reshape(data, (GP, HP, B, RP)) + data = tl.reshape(tl.permute(data, (0, 3, 1, 2)), (GP * RP * HP, B)) + else: + data = tl.reshape(data, (B, GP, HP, RP)) + data = tl.reshape(tl.permute(data, (0, 1, 3, 2)), (B, GP * RP * HP)) + if GP == 1 and HP == H and RP * HP == P: + return data + index = tl.arange(0, P) + source = (index // SIZE * RP * HP + (index % SIZE) // H * HP + index % H) + source = tl.minimum(source, GP * RP * HP - 1) + if STRIDED: + return tl.gather(data, tl.broadcast_to(source[:, None], (P, B)), 0) + else: + return tl.gather(data, tl.broadcast_to(source[None, :], (B, P)), 1) + + +@triton.jit +def _fft_odd(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADIX: tl.constexpr, SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Paired DFT on compact branch tensors, retaining the input precision. + + v[j] + v[r-j] contributes cosine terms to a conjugate output pair; + v[j] - v[r-j] contributes sine terms with opposite signs. This reduces + duplicated arithmetic across output branches. Padding is register-only. + """ + HP: tl.constexpr = triton.next_power_of_2(SIZE // RADIX) + GP: tl.constexpr = triton.next_power_of_2(N // SIZE) + ar, ai = _odd_branch(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADIX, SIZE, 0, HP, GP, TWIDDLE_STEP) + sums_r, sums_i, diffs_r, diffs_i = (), (), (), () + dc_r, dc_i = ar, ai + for j in tl.static_range(1, (RADIX + 1) // 2): + br, bi = _odd_branch(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADIX, SIZE, j, HP, GP, TWIDDLE_STEP) + cr, ci = _odd_branch(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADIX, SIZE, RADIX - j, HP, GP, TWIDDLE_STEP) + sr, si, dr, di = br + cr, bi + ci, br - cr, bi - ci + sums_r, sums_i = sums_r + (sr,), sums_i + (si,) + diffs_r, diffs_i = diffs_r + (dr,), diffs_i + (di,) + dc_r, dc_i = dc_r + sr, dc_i + si + outputs_r, outputs_i = (dc_r,), (dc_i,) + for k in tl.static_range(1, (RADIX + 1) // 2): + even_r, even_i = ar, ai + odd_r = tl.full(ar.shape, 0, ar.dtype) + odd_i = tl.full(ai.shape, 0, ai.dtype) + for j in tl.static_range(1, (RADIX + 1) // 2): + exponent = (j * k % RADIX) * (N // RADIX) + c = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP) + s = tl.load(Twiddle + 2 * exponent * TWIDDLE_STEP + 1) + if INVERSE: + s = -s + even_r = even_r + c * sums_r[j - 1] + even_i = even_i + c * sums_i[j - 1] + odd_r = odd_r + s * diffs_r[j - 1] + odd_i = odd_i + s * diffs_i[j - 1] + outputs_r = outputs_r + (even_r - odd_i, even_r + odd_i) + outputs_i = outputs_i + (even_i + odd_r, even_i - odd_r) + return (_odd_output(outputs_r, N, P, B, STRIDED, RADIX, SIZE, HP, GP), + _odd_output(outputs_i, N, P, B, STRIDED, RADIX, SIZE, HP, GP)) + + +@triton.jit +def _radix9_branch(data, H: tl.constexpr, B: tl.constexpr, + STRIDED: tl.constexpr, TERM: tl.constexpr): + """Select one of sixteen padded register branches using four splits.""" + if STRIDED: + value = tl.permute(tl.reshape(data, (16, H, B)), (1, 2, 0)) + value = tl.reshape(value, (H, B, 2, 2, 2, 2)) + else: + value = tl.permute(tl.reshape(data, (B, 16, H)), (0, 2, 1)) + value = tl.reshape(value, (B, H, 2, 2, 2, 2)) + for bit in tl.static_range(4): + lo, hi = tl.split(value) + if (TERM >> bit) & 1: + value = hi + else: + value = lo + return value + + +@triton.jit +def _dft3(ar, ai, br, bi, cr, ci, sine): + """Three-point DFT with an exactly representable -1/2 cosine.""" + sr, si = br + cr, bi + ci + dr, di = br - cr, bi - ci + er, ei = ar - 0.5 * sr, ai - 0.5 * si + return (ar + sr, ai + si, + er - sine * di, ei + sine * dr, + er + sine * di, ei - sine * dr) + + +@triton.jit +def _join16(outputs): + zero = tl.full(outputs[0].shape, 0, outputs[0].dtype) + evens = _join8(outputs[0], outputs[2], outputs[4], outputs[6], outputs[8], + zero, zero, zero) + odds = _join8(outputs[1], outputs[3], outputs[5], outputs[7], + zero, zero, zero, zero) + return tl.reshape(tl.join(evens, odds), outputs[0].shape + (16,)) + + +@triton.jit +def _fft_radix9(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + """Fuse two final radix-3 stages without an intermediate register gather. + + After the stage twiddles, compute three DFTs over indices (b, b+3, + b+6), multiply branch b and frequency k1 by W9**(b*k1), then compute + three DFTs over b. Natural output frequency is k1 + 3*k2. + Padding is confined to register tensors; all global rows stay packed. + """ + H: tl.constexpr = SIZE // 9 + tl.static_assert(SIZE == N and H == triton.next_power_of_2(H) and P == 16 * H) + local = tl.arange(0, H) + values_r, values_i = (), () + for j in tl.static_range(9): + r = _radix9_branch(real, H, B, STRIDED, j) + i = _radix9_branch(imag, H, B, STRIDED, j) + if j != 0 and H != 1: + wr = tl.load(Twiddle + 2 * j * local * TWIDDLE_STEP) + wi = tl.load(Twiddle + 2 * j * local * TWIDDLE_STEP + 1) + if INVERSE: + wi = -wi + if STRIDED: + wr, wi = wr[:, None], wi[:, None] + else: + wr, wi = wr[None, :], wi[None, :] + r, i = r * wr - i * wi, r * wi + i * wr + values_r, values_i = values_r + (r,), values_i + (i,) + sine = tl.load(Twiddle + 2 * (N // 3) * TWIDDLE_STEP + 1) + if INVERSE: + sine = -sine + columns = () + for b in tl.static_range(3): + columns = columns + (_dft3(values_r[b], values_i[b], + values_r[b + 3], values_i[b + 3], + values_r[b + 6], values_i[b + 6], sine),) + rows = () + for k1 in tl.static_range(3): + branches_r, branches_i = (), () + for b in tl.static_range(3): + r, i = columns[b][2 * k1], columns[b][2 * k1 + 1] + if b != 0 and k1 != 0: + wr = tl.load(Twiddle + 2 * (b * k1 * H) * TWIDDLE_STEP) + wi = tl.load(Twiddle + 2 * (b * k1 * H) * TWIDDLE_STEP + 1) + if INVERSE: + wi = -wi + r, i = r * wr - i * wi, r * wi + i * wr + branches_r, branches_i = branches_r + (r,), branches_i + (i,) + rows = rows + (_dft3(branches_r[0], branches_i[0], + branches_r[1], branches_i[1], + branches_r[2], branches_i[2], sine),) + outputs_r, outputs_i = (), () + for k2 in tl.static_range(3): + for k1 in tl.static_range(3): + outputs_r = outputs_r + (rows[k1][2 * k2],) + outputs_i = outputs_i + (rows[k1][2 * k2 + 1],) + r, i = _join16(outputs_r), _join16(outputs_i) + if STRIDED: + return (tl.reshape(tl.permute(r, (2, 0, 1)), (P, B)), + tl.reshape(tl.permute(i, (2, 0, 1)), (P, B))) + else: + return (tl.reshape(tl.permute(r, (0, 2, 1)), (B, P)), + tl.reshape(tl.permute(i, (0, 2, 1)), (B, P))) + + +@triton.jit +def _fft_stage(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADIX: tl.constexpr, SIZE: tl.constexpr, TWIDDLE_STEP: tl.constexpr): + # Factors 4 and 2 precede odd factors, so these reshapes divide P exactly. + if RADIX == 4 and P % SIZE == 0: + return _fft_radix4(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + SIZE, TWIDDLE_STEP) + elif RADIX == 2 and P % SIZE == 0: + return _fft_radix2(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + SIZE, TWIDDLE_STEP) + elif RADIX == 9: + return _fft_radix9(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + SIZE, TWIDDLE_STEP) + elif RADIX == 3 or RADIX == 5 or RADIX == 7: + return _fft_odd(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADIX, SIZE, TWIDDLE_STEP) + else: + return _fft_generic_stage(real, imag, Twiddle, N, P, B, STRIDED, + INVERSE, RADIX, SIZE, TWIDDLE_STEP) + + +@triton.jit +def _fft(real, imag, Twiddle, N: tl.constexpr, P: tl.constexpr, + B: tl.constexpr, STRIDED: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr, TWIDDLE_STEP: tl.constexpr = 1): + for stage in tl.static_range(len(RADICES)): + real, imag = _fft_stage(real, imag, Twiddle, N, P, B, STRIDED, INVERSE, + RADICES[stage], SIZES[stage], TWIDDLE_STEP) + return real, imag + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def real_axis(Input, Output, Twiddle, ROWS: tl.constexpr, N: tl.constexpr, + P: tl.constexpr, B: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr): + rows = tl.program_id(0) * B + tl.arange(0, B) + index = tl.arange(0, P) + source = _input_index(index, N, RADICES, SIZES) + NF: tl.constexpr = N // 2 + 1 + mask = (rows[:, None] < ROWS) & (index[None, :] < N) + if INVERSE: + half_index = tl.minimum(source, N - source) + offsets = 2 * (rows[:, None] * NF + half_index[None, :]) + real = tl.load(Input + offsets, mask, 0) + imag = tl.load(Input + offsets + 1, mask, 0) + imag = tl.where(source[None, :] >= NF, -imag, imag) + # Real end points after inverse Z/Y implement irfftn's convention + # even when the original packed input is not Hermitian-consistent. + imag = tl.where((source[None, :] == 0) | (2 * source[None, :] == N), 0, imag) + else: + real = tl.load(Input + rows[:, None] * N + source[None, :], mask, 0) + imag = tl.full((B, P), 0, real.dtype) + real, imag = _fft(real, imag, Twiddle, N, P, B, False, INVERSE, RADICES, SIZES) + if INVERSE: + tl.store(Output + rows[:, None] * N + index[None, :], real, mask) + else: + offsets = 2 * (rows[:, None] * NF + index[None, :]) + output_mask = (rows[:, None] < ROWS) & (index[None, :] < NF) + tl.store(Output + offsets, real, output_mask) + tl.store(Output + offsets + 1, imag, output_mask) + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def packed_real_axis(Input, Output, Twiddle, ROWS: tl.constexpr, N: tl.constexpr, + P: tl.constexpr, B: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr): + """Even-length real FFT using an N/2 complex FFT and fused reconstruction. + + RADICES, SIZES and P describe M=N/2; Twiddle contains N complex constants. + Even/odd samples are paired only in registers; global layout stays exact. + The inverse constructs twice the packed complex spectrum so that the + length-M unnormalized inverse returns the required length-N scaling. + """ + tl.static_assert(N % 2 == 0) + M: tl.constexpr = N // 2 + NF: tl.constexpr = M + 1 + rows = tl.program_id(0) * B + tl.arange(0, B) + index = tl.arange(0, P) + source = _input_index(index, M, RADICES, SIZES) + mask = (rows[:, None] < ROWS) & (index[None, :] < M) + if INVERSE: + offsets = 2 * (rows[:, None] * NF + source[None, :]) + mirrors = 2 * (rows[:, None] * NF + M - source[None, :]) + ar = tl.load(Input + offsets, mask, 0) + ai = tl.load(Input + offsets + 1, mask, 0) + br = tl.load(Input + mirrors, mask, 0) + bi = tl.load(Input + mirrors + 1, mask, 0) + # irfftn ignores real-axis endpoint imaginary parts AFTER Z/Y inverse. + # Arbitrary complex packed input is valid; do not zero planes earlier. + ai = tl.where(source[None, :] == 0, 0, ai) + bi = tl.where(source[None, :] == 0, 0, bi) + wr = tl.load(Twiddle + 2 * source, index < M, 0)[None, :] + wi = tl.load(Twiddle + 2 * source + 1, index < M, 0)[None, :] + # 2 Z[k] = X[k] + conj(X[M-k]) + # + i conj(W_N^k) (X[k] - conj(X[M-k])). + real = ar + br - wr * (ai + bi) + wi * (ar - br) + imag = ai - bi + wr * (ar - br) + wi * (ai + bi) + else: + offsets = rows[:, None] * N + 2 * source[None, :] + real = tl.load(Input + offsets, mask, 0) + imag = tl.load(Input + offsets + 1, mask, 0) + real, imag = _fft(real, imag, Twiddle, M, P, B, False, INVERSE, + RADICES, SIZES, 2) + if INVERSE: + offsets = rows[:, None] * N + 2 * index[None, :] + tl.store(Output + offsets, real, mask) + tl.store(Output + offsets + 1, imag, mask) + else: + mirror_index = tl.where(index < M, (M - index) % M, 0) + mirror = tl.broadcast_to(mirror_index[None, :], (B, P)) + br = tl.gather(real, mirror, 1) + bi = tl.gather(imag, mirror, 1) + wr = tl.load(Twiddle + 2 * index, index < M, 0)[None, :] + wi = tl.load(Twiddle + 2 * index + 1, index < M, 0)[None, :] + # X[k] = (Z[k] + conj(Z[-k]))/2 + # - i W_N^k (Z[k] - conj(Z[-k]))/2. + out_real = 0.5 * (real + br + wr * (imag + bi) + wi * (real - br)) + out_imag = 0.5 * (imag - bi + wi * (imag + bi) - wr * (real - br)) + offsets = 2 * (rows[:, None] * NF + index[None, :]) + tl.store(Output + offsets, out_real, mask) + tl.store(Output + offsets + 1, out_imag, mask) + # The k=0 lane owns the extra, tightly packed Nyquist coefficient. + endpoints = (rows[:, None] < ROWS) & (index[None, :] == 0) + nyquist = 2 * (rows[:, None] * NF + M) + tl.store(Output + nyquist + 0 * index[None, :], real - imag, endpoints) + tl.store(Output + nyquist + 0 * index[None, :] + 1, 0, endpoints) + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def complex_axis(Input, Output, Twiddle, N: tl.constexpr, INNER: tl.constexpr, + P: tl.constexpr, B: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr): + # Program dimension 1 indexes independent outer batches (Z for Y FFTs). + columns = tl.program_id(0) * B + tl.arange(0, B) + index = tl.arange(0, P) + source = _input_index(index, N, RADICES, SIZES) + base = tl.program_id(1) * N * INNER + offsets = 2 * (base + source[:, None] * INNER + columns[None, :]) + mask = (index[:, None] < N) & (columns[None, :] < INNER) + real = tl.load(Input + offsets, mask, 0).to(tl.float64) + imag = tl.load(Input + offsets + 1, mask, 0).to(tl.float64) + real, imag = _fft(real, imag, Twiddle, N, P, B, True, INVERSE, RADICES, SIZES) + offsets = 2 * (base + index[:, None] * INNER + columns[None, :]) + tl.store(Output + offsets, real, mask) + tl.store(Output + offsets + 1, imag, mask) + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def real_xy(Input, Output, TwiddleX, TwiddleY, + NY: tl.constexpr, NX: tl.constexpr, PY: tl.constexpr, PX: tl.constexpr, + INVERSE: tl.constexpr, RX: tl.constexpr, SX: tl.constexpr, + RY: tl.constexpr, SY: tl.constexpr, Q: tl.constexpr): + """Fuse the real X transform with one short Y residue transform. + + Both FFT networks operate on registers/shared memory in the same CTA. + Only caller layout is accessed globally: real NY*NX or packed NY*NF*2. + An inverse call consumes a plane AFTER the Z inverse. Arbitrary boundary + values are projected only after the Y inverse, matching irfftn. + """ + z = tl.program_id(0) + residue = tl.program_id(1) + M: tl.constexpr = NY // Q + y = tl.arange(0, PY) + x = tl.arange(0, PX) + sy = _input_index(y, M, RY, SY) + sx = _input_index(x, NX, RX, SX) + NF: tl.constexpr = NX // 2 + 1 + if INVERSE: + valid = (y[:, None] < M) & (x[None, :] < NF) + values_r, values_i = (), () + for part in tl.static_range(Q): + offsets = 2 * ((z * NY + sy[:, None] + part * M) * NF + x[None, :]) + ar = tl.load(Input + offsets, valid, 0).to(tl.float64) + ai = tl.load(Input + offsets + 1, valid, 0).to(tl.float64) + values_r, values_i = values_r + (ar,), values_i + (ai,) + real = tl.where(residue == 0, values_r[0] + values_r[1], values_r[0] - values_r[1]) + imag = tl.where(residue == 0, values_i[0] + values_i[1], values_i[0] - values_i[1]) + exponent = sy * residue % NY + if NY == 32: + wr, wi = _constant32_twiddle(exponent) + else: + wr = tl.load(TwiddleY + 2 * exponent) + wi = tl.load(TwiddleY + 2 * exponent + 1) + wi = -wi + real, imag = real * wr[:, None] - imag * wi[:, None], real * wi[:, None] + imag * wr[:, None] + real, imag = _fft(real, imag, TwiddleY, M, PY, PX, True, True, RY, SY, Q) + # Complete and permute each X spectrum only after inverse Y. No + # original endpoint plane is discarded or assumed Hermitian. + half = tl.minimum(sx, NX - sx) + half = tl.where(x < NX, half, 0) + gather = tl.broadcast_to(half[None, :], (PY, PX)) + real = tl.gather(real, gather, 1) + imag = tl.gather(imag, gather, 1) + imag = tl.where(sx[None, :] >= NF, -imag, imag) + imag = tl.where((sx[None, :] == 0) | (2 * sx[None, :] == NX), 0, imag) + real, imag = _fft(real, imag, TwiddleX, NX, PX, PY, False, True, RX, SX) + offsets = (z * NY + Q * y[:, None] + residue) * NX + x[None, :] + tl.store(Output + offsets, real, (y[:, None] < M) & (x[None, :] < NX)) + else: + offsets = (z * NY + Q * sy[:, None] + residue) * NX + sx[None, :] + valid = (y[:, None] < M) & (x[None, :] < NX) + real = tl.load(Input + offsets, valid, 0).to(tl.float64) + imag = tl.full((PY, PX), 0, real.dtype) + real, imag = _fft(real, imag, TwiddleX, NX, PX, PY, False, False, RX, SX) + real, imag = _fft(real, imag, TwiddleY, M, PY, PX, True, False, RY, SY, Q) + offsets = 2 * ((z * NY + residue * M + y[:, None]) * NF + x[None, :]) + valid = (y[:, None] < M) & (x[None, :] < NF) + tl.store(Output + offsets, real, valid) + tl.store(Output + offsets + 1, imag, valid) + + +@triton.jit(do_not_specialize_on_alignment=["Input", "Output"]) +def combined_z(Input, Output, Twiddle, TwiddleY, + N: tl.constexpr, INNER: tl.constexpr, + P: tl.constexpr, B: tl.constexpr, INVERSE: tl.constexpr, + RADICES: tl.constexpr, SIZES: tl.constexpr, + NY: tl.constexpr, NX: tl.constexpr, Q: tl.constexpr): + """Finish forward Y in the Z loads, without an extra global stage.""" + NF: tl.constexpr = NX // 2 + 1 + M: tl.constexpr = NY // Q + columns = tl.program_id(0) * B + tl.arange(0, B) + ky = columns // NF + kx = columns % NF + index = tl.arange(0, P) + source = _input_index(index, N, RADICES, SIZES) + mask = (index[:, None] < N) & (columns[None, :] < INNER) + real = tl.full((P, B), 0., tl.float64) + imag = tl.full((P, B), 0., tl.float64) + for residue in tl.static_range(Q): + offsets = 2 * ((source[:, None] * NY + residue * M + ky[None, :] % M) * NF + kx[None, :]) + ar = tl.load(Input + offsets, mask, 0).to(tl.float64) + ai = tl.load(Input + offsets + 1, mask, 0).to(tl.float64) + exponent = ky * residue % NY + if NY == 32: + wr, wi = _constant32_twiddle(exponent) + else: + wr = tl.load(TwiddleY + 2 * exponent) + wi = tl.load(TwiddleY + 2 * exponent + 1) + real += ar * wr[None, :] - ai * wi[None, :] + imag += ar * wi[None, :] + ai * wr[None, :] + real, imag = _fft(real, imag, Twiddle, N, P, B, True, False, RADICES, SIZES) + offsets = 2 * (index[:, None] * INNER + columns[None, :]) + tl.store(Output + offsets, real, mask) + tl.store(Output + offsets + 1, imag, mask) diff --git a/build/torch-cuda/_c2r_f32/__init__.py b/build/torch-cuda/_c2r_f32/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f37bcfdee76d2fc0fefc8e007a6000c171b723a5 --- /dev/null +++ b/build/torch-cuda/_c2r_f32/__init__.py @@ -0,0 +1,5 @@ +"""fp32 3D C2R (inverse real) FFT on a packed Hermitian half spectrum; see ``kernel``.""" + +from .kernel import prepare, run + +__all__ = ["prepare", "run"] diff --git a/build/torch-cuda/_c2r_f32/fftgen.py b/build/torch-cuda/_c2r_f32/fftgen.py new file mode 100644 index 0000000000000000000000000000000000000000..377c0147d07903fa0d9e67defd6922e3273cb20b --- /dev/null +++ b/build/torch-cuda/_c2r_f32/fftgen.py @@ -0,0 +1,243 @@ +"""Trace-time DFT generators and memory helpers shared by all kernels. + +A complex value is a ``(re, im)`` tuple of ``cutlass.Float32`` proxies. All +loops here run in the Python interpreter while a kernel is traced, so every +DFT below becomes straight-line code with compile-time twiddle constants. + +Sign convention: ``sign=+1`` is the inverse (unnormalized) DFT used by C2R, +``X[k] = sum_j x[j] * exp(+2*pi*i*j*k/n)``. + +The memory helpers load / store complex values (8 or 16 bytes) at a Float32 +pointer plus element offset, optionally with an L2 cache-eviction policy. +""" + +from __future__ import annotations + +import math +from fractions import Fraction + +from cutlass import Float32, Int64, cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import vector + +SQRT1_2 = 0.70710678118654752440 + + +def F(v): + return Float32(v) + + +def cadd(a, b): + return (a[0] + b[0], a[1] + b[1]) + + +def csub(a, b): + return (a[0] - b[0], a[1] - b[1]) + + +def cneg(a): + return (-a[0], -a[1]) + + +def mul_i(a, s): + """Multiply by ``s * i`` with ``s = +1/-1``.""" + if s > 0: + return (-a[1], a[0]) + return (a[1], -a[0]) + + +def cmul(a, w): + """General complex multiply ``a * w`` (both runtime values).""" + return (a[0] * w[0] - a[1] * w[1], a[0] * w[1] + a[1] * w[0]) + + +def twiddle_value(num, den, sign): + """exp(sign * 2*pi*i * num/den) in float64.""" + ang = sign * 2.0 * math.pi * num / den + return (math.cos(ang), math.sin(ang)) + + +def cmul_tw(a, num, den, sign): + """Multiply ``a`` by the constant twiddle exp(sign*2*pi*i*num/den).""" + fr = Fraction(num, den) % 1 + if fr == 0: + return a + s = 1 if sign > 0 else -1 + if fr == Fraction(1, 2): + return cneg(a) + if fr == Fraction(1, 4): + return mul_i(a, s) + if fr == Fraction(3, 4): + return mul_i(a, -s) + eighths = fr * 8 + if eighths.denominator == 1: + # (+-1 +- i)/sqrt(2) + k = int(eighths) + c = F(SQRT1_2) + # angle = s * k * pi/4 ; cos, sin signs: + cs = {1: (1, 1), 3: (-1, 1), 5: (-1, -1), 7: (1, -1)}[k] + sc, ss = cs[0], cs[1] * s + re, im = a + # (re + i im)(sc + i ss)/sqrt2 = (re*sc - im*ss) + i(re*ss + im*sc) + if sc > 0 and ss > 0: + return ((re - im) * c, (re + im) * c) + if sc > 0 and ss < 0: + return ((re + im) * c, (im - re) * c) + if sc < 0 and ss > 0: + return ((-re - im) * c, (re - im) * c) + return ((im - re) * c, (-re - im) * c) + wr, wi = twiddle_value(fr.numerator, fr.denominator, sign) + return (a[0] * F(wr) - a[1] * F(wi), a[0] * F(wi) + a[1] * F(wr)) + + +def _smallest_factor(n): + if n % 4 == 0 and n != 4: + return 4 + for p in (2, 3, 5, 7, 11, 13): + if n % p == 0: + return p + p = 17 + while p * p <= n: + if n % p == 0: + return p + p += 2 + return n + + +def dft_prime(xs, sign): + """Direct DFT for an odd prime length using the symmetric pairing.""" + p = len(xs) + h = (p - 1) // 2 + s = [cadd(xs[j], xs[p - j]) for j in range(1, h + 1)] + d = [csub(xs[j], xs[p - j]) for j in range(1, h + 1)] + x0 = xs[0] + out = [None] * p + tot = x0 + for j in range(h): + tot = cadd(tot, s[j]) + out[0] = tot + for k in range(1, h + 1): + ar, ai = x0 + br = None + bi = None + for j in range(1, h + 1): + c = math.cos(2.0 * math.pi * j * k / p) + sn = math.sin(2.0 * math.pi * j * k / p) + sj = s[j - 1] + dj = d[j - 1] + ar = ar + sj[0] * F(c) + ai = ai + sj[1] * F(c) + tr = dj[0] * F(sn) + ti = dj[1] * F(sn) + br = tr if br is None else br + tr + bi = ti if bi is None else bi + ti + # X[k] = A + sign*i*B ; X[p-k] = A - sign*i*B + ib = mul_i((br, bi), sign) + out[k] = (ar + ib[0], ai + ib[1]) + out[p - k] = (ar - ib[0], ai - ib[1]) + return out + + +def dft(xs, sign=1): + """Unrolled DFT of a list of complex values (any length).""" + n = len(xs) + if n == 1: + return list(xs) + if n == 2: + return [cadd(xs[0], xs[1]), csub(xs[0], xs[1])] + if n == 4: + t0 = cadd(xs[0], xs[2]) + t1 = csub(xs[0], xs[2]) + t2 = cadd(xs[1], xs[3]) + t3 = mul_i(csub(xs[1], xs[3]), sign) + return [cadd(t0, t2), cadd(t1, t3), csub(t0, t2), csub(t1, t3)] + p = _smallest_factor(n) + if p == n: + return dft_prime(xs, sign) + q = n // p + # inner p-point DFTs over k1 on subsequences xs[k2 + q*k1] + inner = [] + for k2 in range(q): + sub = dft([xs[k2 + q * k1] for k1 in range(p)], sign) + inner.append([cmul_tw(sub[n1], n1 * k2, n, sign) for n1 in range(p)]) + out = [None] * n + for n1 in range(p): + res = dft([inner[k2][n1] for k2 in range(q)], sign) + for n2 in range(q): + out[n1 + p * n2] = res[n2] + return out + + +def factor2(n, max_radix=16): + """Split n = n1 * n2 with n1 >= n2, both <= max_radix if possible. + + Prefers balanced splits; returns None if impossible. + """ + best = None + for n2 in range(1, n + 1): + if n % n2: + continue + n1 = n // n2 + if n1 < n2: + break + if n1 <= max_radix: + cand = (n1, n2) + if best is None or abs(n1 - n2) < abs(best[0] - best[1]): + best = cand + return best + + +# --------------------------------------------------------------------------- +# Memory helpers (pointer + element offset in Float32 units) +# --------------------------------------------------------------------------- + + +def _v2f32(): + return ir.VectorType.get([2], Float32.mlir_type) + + +def ld_c(ptr, off, **kw): + """Load one complex (8 bytes) at float offset ``off``.""" + v = cute.arch.load(ptr + off, _v2f32(), **kw) + return (Float32(vector.extract(v, [], [0])), Float32(vector.extract(v, [], [1]))) + + +def st_c(ptr, off, c, **kw): + """Store one complex (8 bytes) at float offset ``off``.""" + v = vector.from_elements(_v2f32(), [Float32(c[0]).ir_value(), Float32(c[1]).ir_value()]) + cute.arch.store(ptr + off, v, **kw) + + +def make_policy(kind): + """64-bit L2 cache policy: kind in {'evict_first', 'evict_last', 'evict_unchanged'}.""" + return cute.arch.inline_ptx( + "createpolicy.fractional.L2::" + kind + ".b64 {$w0}, 1.0;", + write_only_types=[Int64], + ) + + +def ld_c_hint(ptr, off, policy): + """Global 8-byte complex load with an L2 cache policy.""" + return cute.arch.inline_ptx( + "ld.global.L2::cache_hint.v2.f32 {{$w0}, {$w1}}, [{$r0}], {$r1};", + write_only_types=[Float32, Float32], + read_only_args=[(ptr + off).toint(), policy], + ) + + +def ld_c2_hint(ptr, off, policy): + """Global 16-byte load of two complex values (16B aligned) with an L2 cache policy.""" + v = cute.arch.inline_ptx( + "ld.global.L2::cache_hint.v4.f32 {{$w0}, {$w1}, {$w2}, {$w3}}, [{$r0}], {$r1};", + write_only_types=[Float32, Float32, Float32, Float32], + read_only_args=[(ptr + off).toint(), policy], + ) + return (v[0], v[1]), (v[2], v[3]) + + +def st_c_hint(ptr, off, c, policy): + """Global 8-byte complex store with an L2 cache policy.""" + cute.arch.inline_ptx( + "st.global.L2::cache_hint.v2.f32 [{$r0}], {{$r1}, {$r2}}, {$r3};", + read_only_args=[(ptr + off).toint(), Float32(c[0]), Float32(c[1]), policy], + ) diff --git a/build/torch-cuda/_c2r_f32/generic.py b/build/torch-cuda/_c2r_f32/generic.py new file mode 100644 index 0000000000000000000000000000000000000000..83d2e9a5cf3019fb3792c2d40e4d5dbb10b419d2 --- /dev/null +++ b/build/torch-cuda/_c2r_f32/generic.py @@ -0,0 +1,288 @@ +"""Shape-agnostic fallback for shapes neither the two-pass nor the in-place path can plan. + +Every axis is transformed by a sequence of out-of-place mixed-radix stages +(``_stage_kernel``) followed by a digit-reversal permutation +(``_perm_kernel``). A radix-p stage evaluates its p-point DFTs as direct +sums with twiddles from a float64-derived table, so any length works, +including large primes (at O(N*p) cost per stage). The x axis is done as a +full-length complex transform of the Hermitian completion of each row. + +All sizes are runtime arguments, so these kernels compile once per process. +Every launch takes its own two full-size ping-pong work buffers from torch's +caching allocator on the current stream, so launches of the same shape that +are in flight at the same time (other streams) never share a workspace. +""" + +from __future__ import annotations + +import cutlass +import torch +from cutlass import Float32, Int32, cute + +from .fftgen import ld_c, st_c +from .tables import twiddle_line + +_BLOCK = 256 + + +def _factor(n): + fs = [] + while n % 4 == 0: + fs.append(4) + n //= 4 + while n % 2 == 0: + fs.append(2) + n //= 2 + p = 3 + while p * p <= n: + while n % p == 0: + fs.append(p) + n //= p + p += 2 + if n > 1: + fs.append(n) + return fs + + +@cute.kernel +def _stage_kernel( + src: cute.Tensor, + dst: cute.Tensor, + W: cute.Tensor, + N: Int32, + S: Int32, + total: Int32, + p: Int32, + Nprev: Int32, + Np: Int32, + L: Int32, + NdivP: Int32, +): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + t = bid * _BLOCK + tid + if t < total: + sp = src.iterator + wp = W.iterator + i = t % S + rest = t // S + nout = rest % N + o = rest // N + P = nout // Nprev + r = nout - P * Nprev + n1 = r // Np + k2 = r - n1 * Np + base = (o * N + P * Nprev + k2) * S + i + step = (n1 * NdivP) % N + acc_r = Float32(0.0) + acc_i = Float32(0.0) + idx = Int32(0) + for k1 in range(p): + v = ld_c(sp, (base + k1 * Np * S) * 2) + w = ld_c(wp, idx * 2) + acc_r = acc_r + v[0] * w[0] - v[1] * w[1] + acc_i = acc_i + v[0] * w[1] + v[1] * w[0] + idx = idx + step + idx = cutlass.select_(idx >= N, idx - N, idx) + tw = ld_c(wp, (((n1 * k2) % Nprev) * L) * 2) + out_r = acc_r * tw[0] - acc_i * tw[1] + out_i = acc_r * tw[1] + acc_i * tw[0] + st_c(dst.iterator, ((o * N + nout) * S + i) * 2, (out_r, out_i)) + + +@cute.kernel +def _perm_kernel(src: cute.Tensor, dst: cute.Tensor, rad: cute.Tensor, nst: Int32, N: Int32, S: Int32, total: Int32): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + t = bid * _BLOCK + tid + if t < total: + i = t % S + rest = t // S + f = rest % N + o = rest // N + pos = Int32(0) + ff = f + lc = N + for s in range(nst): + pr = rad[s] + d = ff % pr + ff = ff // pr + lc = lc // pr + pos = pos + d * lc + v = ld_c(src.iterator, ((o * N + pos) * S + i) * 2) + st_c(dst.iterator, t * 2, v) + + +@cute.kernel +def _perm_real_kernel(src: cute.Tensor, dst: cute.Tensor, rad: cute.Tensor, nst: Int32, N: Int32, total: Int32): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + t = bid * _BLOCK + tid + if t < total: + f = t % N + o = t // N + pos = Int32(0) + ff = f + lc = N + for s in range(nst): + pr = rad[s] + d = ff % pr + ff = ff // pr + lc = lc // pr + pos = pos + d * lc + v = ld_c(src.iterator, (o * N + pos) * 2) + dst[t] = v[0] + + +@cute.kernel +def _herm_kernel(src: cute.Tensor, dst: cute.Tensor, nf: Int32, nx: Int32, total: Int32): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + t = bid * _BLOCK + tid + if t < total: + row = t // nx + k = t - row * nx + M = nx // 2 + kk = cutlass.select_(k <= M, k, nx - k) + v = ld_c(src.iterator, (row * nf + kk) * 2) + edge = (k == 0) | (k == M) + im = cutlass.select_(k <= M, v[1], -v[1]) + im = cutlass.select_(edge, Float32(0.0), im) + st_c(dst.iterator, t * 2, (v[0], im)) + + +def _grid(total): + return (total + _BLOCK - 1) // _BLOCK + + +@cute.jit +def _launch_stage( + src: cute.Tensor, + dst: cute.Tensor, + W: cute.Tensor, + N: Int32, + S: Int32, + total: Int32, + p: Int32, + Nprev: Int32, + Np: Int32, + L: Int32, + NdivP: Int32, + nblk: Int32, + stream, +): + _stage_kernel(src, dst, W, N, S, total, p, Nprev, Np, L, NdivP).launch( + grid=(nblk, 1, 1), block=(_BLOCK, 1, 1), stream=stream + ) + + +@cute.jit +def _launch_perm( + src: cute.Tensor, + dst: cute.Tensor, + rad: cute.Tensor, + nst: Int32, + N: Int32, + S: Int32, + total: Int32, + nblk: Int32, + stream, +): + _perm_kernel(src, dst, rad, nst, N, S, total).launch(grid=(nblk, 1, 1), block=(_BLOCK, 1, 1), stream=stream) + + +@cute.jit +def _launch_perm_real( + src: cute.Tensor, dst: cute.Tensor, rad: cute.Tensor, nst: Int32, N: Int32, total: Int32, nblk: Int32, stream +): + _perm_real_kernel(src, dst, rad, nst, N, total).launch(grid=(nblk, 1, 1), block=(_BLOCK, 1, 1), stream=stream) + + +@cute.jit +def _launch_herm(src: cute.Tensor, dst: cute.Tensor, nf: Int32, nx: Int32, total: Int32, nblk: Int32, stream): + _herm_kernel(src, dst, nf, nx, total).launch(grid=(nblk, 1, 1), block=(_BLOCK, 1, 1), stream=stream) + + +_COMPILED = {} + + +def _compiled(): + if _COMPILED: + return _COMPILED + n = cute.sym_int() + m = cute.sym_int() + k = cute.sym_int() + + def fk(s): + return cute.runtime.make_fake_compact_tensor(Float32, (s,), assumed_align=8) + + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + rad = cute.runtime.make_fake_compact_tensor(Int32, (k,), assumed_align=4) + i = Int32(1) + opts = "--enable-tvm-ffi" + _COMPILED["stage"] = cute.compile( + _launch_stage, fk(n), fk(m), fk(k), i, i, i, i, i, i, i, i, i, stream, options=opts + ) + _COMPILED["perm"] = cute.compile(_launch_perm, fk(n), fk(m), rad, i, i, i, i, i, stream, options=opts) + _COMPILED["perm_real"] = cute.compile(_launch_perm_real, fk(n), fk(m), rad, i, i, i, i, stream, options=opts) + _COMPILED["herm"] = cute.compile(_launch_herm, fk(n), fk(m), i, i, i, i, stream, options=opts) + return _COMPILED + + +class GenericPlan: + """z, then y, then Hermitian completion + full-length x, each as radix stages + a permutation.""" + + def __init__(self, nz, ny, nx, device): + self.k = _compiled() + self.nz, self.ny, self.nx = nz, ny, nx + self.nf = nx // 2 + 1 + self.nbig = nz * ny * max(nx, self.nf) # complex elements per work buffer + self.axes = {} + for name, N in (("z", nz), ("y", ny), ("x", nx)): + rad = _factor(N) if N > 1 else [] + self.axes[name] = ( + N, + rad, + torch.tensor(rad if rad else [1], dtype=torch.int32, device=device), + twiddle_line(N, N).to(device), + ) + + def _axis(self, name, src, B, S, out_bufs): + """Transform along an axis; returns (buffer holding result, free buffer).""" + N, rad, rad_t, W = self.axes[name] + total = B * N * S + nblk = _grid(total) + cur = src + a, b = out_bufs + dst = a + L = 1 + for p in rad: + Nprev = N // L + Np = Nprev // p + self.k["stage"](cur, dst, W, N, S, total, p, Nprev, Np, L, N // p, nblk) + L *= p + cur, dst = dst, (b if dst is a else a) + return cur, dst, rad_t, len(rad), total, nblk + + def __call__(self, x_flat, out_flat): + nz, ny, nx, nf = self.nz, self.ny, self.nx, self.nf + # per-launch workspace (stream-ordered reuse by the caching allocator) + A, Bf = (torch.empty(self.nbig * 2, dtype=torch.float32, device=x_flat.device) for _ in range(2)) + # z axis: element (kz, c) at kz*S + c with S = ny*nf + cur, free, rad_t, nst, total, nblk = self._axis("z", x_flat, 1, ny * nf, (A, Bf)) + self.k["perm"](cur, free, rad_t, nst, nz, ny * nf, total, nblk) + cur = free + free = A if cur is Bf else Bf + # y axis: batch nz, stride nf + cur, free, rad_t, nst, total, nblk = self._axis("y", cur, nz, nf, (free, cur)) + self.k["perm"](cur, free, rad_t, nst, ny, nf, total, nblk) + cur = free + free = A if cur is Bf else Bf + # Hermitian completion of every row to nx complex values + rows = nz * ny + tot = rows * nx + self.k["herm"](cur, free, nf, nx, tot, _grid(tot)) + cur = free + free = A if cur is Bf else Bf + cur, free, rad_t, nst, total, nblk = self._axis("x", cur, rows, 1, (free, cur)) + self.k["perm_real"](cur, out_flat, rad_t, nst, nx, total, nblk) diff --git a/build/torch-cuda/_c2r_f32/inplace.py b/build/torch-cuda/_c2r_f32/inplace.py new file mode 100644 index 0000000000000000000000000000000000000000..c747595cea564a29af829e177efbaab8a11be6c2 --- /dev/null +++ b/build/torch-cuda/_c2r_f32/inplace.py @@ -0,0 +1,669 @@ +"""Kernel builders for the in-place Z -> Y -> X fallback pipeline. + +ZPassIP : inverse FFT along z over the flattened (ky, kx) columns; writes + the kx < nx/2 part into the output buffer (complex [nz][ny][M]) + and the Nyquist column into a side buffer nyq [ny][nz]. +YXFusedIP : one CTA per z-plane, y-FFT then C2R along x, in place. +YPassIP / +XPassIP : the same two transforms as separate in-place passes, for planes + that do not fit in shared memory. + +The passes are chained with programmatic dependent launch (each pass triggers +its dependents early and waits on its predecessor before touching ``u``). +All sizes are compile-time constants; the DFT butterflies are unrolled at +trace time by ``fftgen``. +""" + +from __future__ import annotations + +import cutlass +import torch +from cutlass import Float32, cute + +from .fftgen import cmul, dft, factor2, ld_c, ld_c_hint, make_policy, st_c, st_c_hint +from .tables import twiddle_line, twiddle_table + + +def _cdiv(a, b): + return (a + b - 1) // b + + +def _odd_stride(n): + return n + 1 if n % 2 == 0 else n + 2 + + +# =========================================================================== +# In-place pipeline: the z-FFT output (kx < M) lives in the output buffer u, +# viewed as complex [nz][NY][M]; the Nyquist column kx = M lives in a side +# buffer nyq laid out [NY][nz] (complex). +# =========================================================================== + + +def _load_table_to_smem(tid, T, twp, S, n_floats): + """Cooperative copy of a float table from global to shared memory. + + Plain trace-time helper: out-of-range threads clamp to the last entry and + rewrite the same value, so no dynamic branch is needed. + """ + n = n_floats // 2 + for rnd in range(_cdiv(n, T)): + e = tid + rnd * T + if (rnd + 1) * T > n: + e = cutlass.min(e, n - 1) + st_c(S, e * 2, ld_c(twp, e * 2)) + + +class ZPassIP: + """z-FFT over flattened (ky, kx) columns, writing u-aliased output + nyq.""" + + def __init__(self, nz, NY, NX, TC, nthreads): + self.N, self.NY, self.NX = nz, NY, NX + self.M = NX // 2 + self.NF = self.M + 1 + self.C = NY * self.NF + self.TC = TC + self.N1, self.N2 = factor2(nz) + self.T = nthreads + assert self.C % TC == 0 + self.grid = self.C // TC + self.table = twiddle_table(nz, self.N1, self.N2) + + def kernel(self): + N, C, TC, N1, N2, T = self.N, self.C, self.TC, self.N1, self.N2, self.T + NY, M, NF = self.NY, self.M, self.NF + PL = NY * M # complex per plane of u-aliased A + nslots = T // TC + r1 = _cdiv(N2, nslots) + r2 = _cdiv(N1, nslots) + S_ELEMS = N * TC * 2 + TWF = self.table.numel() + + @cute.kernel + def zpass_ip_kernel(x: cute.Tensor, u: cute.Tensor, nyq: cute.Tensor, tw: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + cute.arch.griddepcontrol_launch_dependents() + col = tid % TC + slot = tid // TC + gcol = bid * TC + col + xp = x.iterator + up = u.iterator + nqp = nyq.iterator + alloc = cutlass.utils.SmemAllocator() + sp = alloc.allocate_array(Float32, S_ELEMS, byte_alignment=16) + stw = alloc.allocate_array(Float32, TWF, byte_alignment=16) + lpol = make_policy("evict_first") + loaded = [None] * r1 + for it in cutlass.range_constexpr(r1): + k2 = slot + it * nslots + if cutlass.const_expr((it + 1) * nslots > N2): + k2 = cutlass.min(k2, N2 - 1) + vals = [None] * N1 + for k1 in cutlass.range_constexpr(N1): + vals[k1] = ld_c_hint(xp, ((k2 + N2 * k1) * C + gcol) * 2, lpol) + loaded[it] = vals + _load_table_to_smem(tid, T, tw.iterator, stw, TWF) + ky = gcol // NF + kx = gcol - ky * NF + uoff = ky * M + kx + noff = ky * N + cute.arch.sync_threads() + for it in cutlass.range_constexpr(r1): + k2 = slot + it * nslots + if cutlass.const_expr((it + 1) * nslots > N2): + k2 = cutlass.min(k2, N2 - 1) + res = dft(loaded[it], 1) + for n1 in cutlass.range_constexpr(1, N1): + res[n1] = cmul(res[n1], ld_c(stw, (k2 * N1 + n1) * 2)) + if cutlass.const_expr((it + 1) * nslots > N2): + if slot + it * nslots < N2: + for n1 in cutlass.range_constexpr(N1): + st_c(sp, ((k2 + N2 * n1) * TC + col) * 2, res[n1]) + else: + for n1 in cutlass.range_constexpr(N1): + st_c(sp, ((k2 + N2 * n1) * TC + col) * 2, res[n1]) + cute.arch.sync_threads() + for it in cutlass.range_constexpr(r2): + n1 = slot + it * nslots + if cutlass.const_expr((it + 1) * nslots > N1): + n1 = cutlass.min(n1, N1 - 1) + vals = [None] * N2 + for k2 in cutlass.range_constexpr(N2): + vals[k2] = ld_c(sp, ((k2 + N2 * n1) * TC + col) * 2) + res = dft(vals, 1) + if cutlass.const_expr((it + 1) * nslots > N1): + if slot + it * nslots < N1: + if kx < M: + for n2 in cutlass.range_constexpr(N2): + st_c(up, ((n1 + N1 * n2) * PL + uoff) * 2, res[n2]) + else: + for n2 in cutlass.range_constexpr(N2): + st_c(nqp, (noff + n1 + N1 * n2) * 2, res[n2]) + else: + if kx < M: + for n2 in cutlass.range_constexpr(N2): + st_c(up, ((n1 + N1 * n2) * PL + uoff) * 2, res[n2]) + else: + for n2 in cutlass.range_constexpr(N2): + st_c(nqp, (noff + n1 + N1 * n2) * 2, res[n2]) + + return zpass_ip_kernel + + +class YPassIP: + """In-place y-FFT on column blocks of u (and on the nyq matrix). + + CTAs [0, nz*NB) handle u planes (block of MC columns); CTAs beyond that + handle the Nyquist matrix nyq [NY][nz] in blocks of MC "columns" (= z). + """ + + def __init__(self, nz, NY, NX, MC): + self.nz, self.NY, self.NX = nz, NY, NX + M = NX // 2 + self.M = M + self.MC = MC + assert M % MC == 0 + self.NB = M // MC + self.NQ = _cdiv(nz, MC) + self.Y1, self.Y2 = factor2(NY) + self.T = MC * self.Y2 + self.grid = nz * self.NB + self.NQ + self.table = twiddle_table(NY, self.Y1, self.Y2) + + def kernel(self): + nz, NY, M, MC, NB = self.nz, self.NY, self.M, self.MC, self.NB + Y1, Y2, T = self.Y1, self.Y2, self.T + PL = NY * M + S_ELEMS = NY * MC * 2 + TWF = self.table.numel() + r2 = _cdiv(Y1, Y2) + NUB = nz * NB + + @cute.kernel + def ypass_ip_kernel(u: cute.Tensor, nyq: cute.Tensor, tw: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + j = tid % MC + k2 = tid // MC + alloc = cutlass.utils.SmemAllocator() + sp = alloc.allocate_array(Float32, S_ELEMS, byte_alignment=16) + stw = alloc.allocate_array(Float32, TWF, byte_alignment=16) + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + if bid < NUB: + z = bid // NB + q = bid - z * NB + gp = u.iterator + base = z * PL + q * MC + j + vals = [None] * Y1 + for k1 in cutlass.range_constexpr(Y1): + vals[k1] = ld_c(gp, (base + (k2 + Y2 * k1) * M) * 2) + _load_table_to_smem(tid, T, tw.iterator, stw, TWF) + cute.arch.sync_threads() + res = dft(vals, 1) + for n1 in cutlass.range_constexpr(1, Y1): + res[n1] = cmul(res[n1], ld_c(stw, (k2 * Y1 + n1) * 2)) + for n1 in cutlass.range_constexpr(Y1): + st_c(sp, ((k2 + Y2 * n1) * MC + j) * 2, res[n1]) + cute.arch.sync_threads() + for it in cutlass.range_constexpr(r2): + n1 = k2 + it * Y2 + if cutlass.const_expr((it + 1) * Y2 > Y1): + n1 = cutlass.min(n1, Y1 - 1) + vals2 = [None] * Y2 + for kk in cutlass.range_constexpr(Y2): + vals2[kk] = ld_c(sp, ((kk + Y2 * n1) * MC + j) * 2) + res2 = dft(vals2, 1) + if cutlass.const_expr((it + 1) * Y2 > Y1): + if k2 + it * Y2 < Y1: + for n2 in cutlass.range_constexpr(Y2): + st_c(gp, (base + (n1 + Y1 * n2) * M) * 2, res2[n2]) + else: + for n2 in cutlass.range_constexpr(Y2): + st_c(gp, (base + (n1 + Y1 * n2) * M) * 2, res2[n2]) + else: + qb = bid - NUB + zc = qb * MC + j + zcc = cutlass.min(zc, nz - 1) + gp = nyq.iterator + base = zcc + vals = [None] * Y1 + for k1 in cutlass.range_constexpr(Y1): + vals[k1] = ld_c(gp, (base + (k2 + Y2 * k1) * nz) * 2) + _load_table_to_smem(tid, T, tw.iterator, stw, TWF) + cute.arch.sync_threads() + res = dft(vals, 1) + for n1 in cutlass.range_constexpr(1, Y1): + res[n1] = cmul(res[n1], ld_c(stw, (k2 * Y1 + n1) * 2)) + for n1 in cutlass.range_constexpr(Y1): + st_c(sp, ((k2 + Y2 * n1) * MC + j) * 2, res[n1]) + cute.arch.sync_threads() + for it in cutlass.range_constexpr(r2): + n1 = k2 + it * Y2 + if cutlass.const_expr((it + 1) * Y2 > Y1): + n1 = cutlass.min(n1, Y1 - 1) + vals2 = [None] * Y2 + for kk in cutlass.range_constexpr(Y2): + vals2[kk] = ld_c(sp, ((kk + Y2 * n1) * MC + j) * 2) + res2 = dft(vals2, 1) + ok = zc < nz + if cutlass.const_expr((it + 1) * Y2 > Y1): + ok = ok & (k2 + it * Y2 < Y1) + if ok: + for n2 in cutlass.range_constexpr(Y2): + st_c(gp, (base + (n1 + Y1 * n2) * nz) * 2, res2[n2]) + + return ypass_ip_kernel + + +class XPassIP: + """In-place C2R along x on rows of u (complex [rows][M]) + nyq [NY][nz].""" + + def __init__(self, nz, NY, NX, RB, nthreads): + self.nz, self.NY, self.NX = nz, NY, NX + M = NX // 2 + self.M = M + self.RB = RB + self.rows = nz * NY + assert self.rows % RB == 0 + self.Ms = _odd_stride(M) + self.X1, self.X2 = factor2(M) + self.T = nthreads + self.grid = self.rows // RB + tX = twiddle_table(M, self.X1, self.X2) + tP = twiddle_line(NX, M // 2 + 1) + self.oX = 0 + self.oP = tX.numel() + self.table = torch.cat([tX, tP]) + + def kernel(self): + nz, NY, M, Ms, RB, T = self.nz, self.NY, self.M, self.Ms, self.RB, self.T + X1, X2 = self.X1, self.X2 + oX, oP = self.oX, self.oP + TWF = self.table.numel() + NR = RB + n_load = RB * M + n_pre = NR * (M // 2 + 1) + n_xs1 = NR * X2 + n_xs2 = NR * X1 + n_out = NR * M + SMF = RB * Ms * 2 + + @cute.kernel + def xpass_ip_kernel(u: cute.Tensor, nyq: cute.Tensor, tw: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + cute.arch.griddepcontrol_wait() + row0 = bid * RB + up = u.iterator + row0 * M * 2 + alloc = cutlass.utils.SmemAllocator() + P = alloc.allocate_array(Float32, SMF, byte_alignment=16) + stw = alloc.allocate_array(Float32, TWF, byte_alignment=16) + NQ = alloc.allocate_array(Float32, RB * 2, byte_alignment=16) + spol = make_policy("evict_first") + # ---- load rows ---- + for rnd in cutlass.range_constexpr(_cdiv(n_load, T)): + e = tid + rnd * T + ec = e + if cutlass.const_expr((rnd + 1) * T > n_load): + ec = cutlass.min(e, n_load - 1) + r = ec // M + m = ec - r * M + v = ld_c(up, ec * 2) + if cutlass.const_expr((rnd + 1) * T > n_load): + if e < n_load: + st_c(P, (r * Ms + m) * 2, v) + else: + st_c(P, (r * Ms + m) * 2, v) + for rnd in cutlass.range_constexpr(_cdiv(RB, T)): + r = tid + rnd * T + rc = cutlass.min(r, RB - 1) + g = row0 + rc + zz = g // NY + yy = g - zz * NY + v = ld_c(nyq.iterator, (yy * nz + zz) * 2) + if r < RB: + st_c(NQ, rc * 2, v) + _load_table_to_smem(tid, T, tw.iterator, stw, TWF) + cute.arch.sync_threads() + + # ---- C2R pre-processing: tasks (r, k), k in [0, M/2] ---- + for rnd in cutlass.range_constexpr(_cdiv(n_pre, T)): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > n_pre): + tc = cutlass.min(t, n_pre - 1) + r = tc % NR + k = tc // NR + bk = ld_c(P, (r * Ms + k) * 2) + bmk = ld_c(P, (r * Ms + (M - k)) * 2) + nq = ld_c(NQ, r * 2) + is0 = k == 0 + bkr = bk[0] + bki = cutlass.select_(is0, Float32(0.0), bk[1]) + bmr = cutlass.select_(is0, nq[0], bmk[0]) + bmi = cutlass.select_(is0, Float32(0.0), bmk[1]) + sr = bkr + bmr + si = bki - bmi + dr = bkr - bmr + di = bki + bmi + w = ld_c(stw, oP + k * 2) + wr = dr * w[0] - di * w[1] + wi = dr * w[1] + di * w[0] + if cutlass.const_expr((rnd + 1) * T > n_pre): + if t < n_pre: + st_c(P, (r * Ms + k) * 2, (sr - wi, si + wr)) + st_c(P, (r * Ms + (M - k)) * 2, (sr + wi, wr - si)) + else: + st_c(P, (r * Ms + k) * 2, (sr - wi, si + wr)) + st_c(P, (r * Ms + (M - k)) * 2, (sr + wi, wr - si)) + cute.arch.sync_threads() + + # ---- x stage 1: tasks (r, k2) ---- + for rnd in cutlass.range_constexpr(_cdiv(n_xs1, T)): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > n_xs1): + tc = cutlass.min(t, n_xs1 - 1) + r = tc % NR + k2 = tc // NR + vals = [None] * X1 + for k1 in cutlass.range_constexpr(X1): + vals[k1] = ld_c(P, (r * Ms + k2 + X2 * k1) * 2) + res = dft(vals, 1) + for n1 in cutlass.range_constexpr(1, X1): + res[n1] = cmul(res[n1], ld_c(stw, oX + (k2 * X1 + n1) * 2)) + if cutlass.const_expr((rnd + 1) * T > n_xs1): + if t < n_xs1: + for n1 in cutlass.range_constexpr(X1): + st_c(P, (r * Ms + k2 + X2 * n1) * 2, res[n1]) + else: + for n1 in cutlass.range_constexpr(X1): + st_c(P, (r * Ms + k2 + X2 * n1) * 2, res[n1]) + cute.arch.sync_threads() + + # ---- x stage 2 (held across a barrier) ---- + nr2 = _cdiv(n_xs2, T) + held2 = [None] * nr2 + for rnd in cutlass.range_constexpr(nr2): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > n_xs2): + tc = cutlass.min(t, n_xs2 - 1) + r = tc % NR + n1 = tc // NR + vals = [None] * X2 + for k2 in cutlass.range_constexpr(X2): + vals[k2] = ld_c(P, (r * Ms + k2 + X2 * n1) * 2) + held2[rnd] = dft(vals, 1) + cute.arch.sync_threads() + for rnd in cutlass.range_constexpr(nr2): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > n_xs2): + tc = cutlass.min(t, n_xs2 - 1) + r = tc % NR + n1 = tc // NR + res = held2[rnd] + if cutlass.const_expr((rnd + 1) * T > n_xs2): + if t < n_xs2: + for n2 in cutlass.range_constexpr(X2): + st_c(P, (r * Ms + n1 + X1 * n2) * 2, res[n2]) + else: + for n2 in cutlass.range_constexpr(X2): + st_c(P, (r * Ms + n1 + X1 * n2) * 2, res[n2]) + cute.arch.sync_threads() + + # ---- copy out ---- + for rnd in cutlass.range_constexpr(_cdiv(n_out, T)): + e = tid + rnd * T + ec = e + if cutlass.const_expr((rnd + 1) * T > n_out): + ec = cutlass.min(e, n_out - 1) + r = ec // M + m = ec - r * M + v = ld_c(P, (r * Ms + m) * 2) + if cutlass.const_expr((rnd + 1) * T > n_out): + if e < n_out: + st_c_hint(up, ec * 2, v, spol) + else: + st_c_hint(up, ec * 2, v, spol) + + return xpass_ip_kernel + + +# --------------------------------------------------------------------------- +# Fused in-place YX pass: warp teams + named barriers, bank-conflict-free +# --------------------------------------------------------------------------- + + +def _named_barrier(bid, nthreads): + cute.arch.barrier(barrier_id=bid, number_of_threads=nthreads) + + +class YXFusedIP: + """One CTA per z-plane; bank-conflict-free warp-team schedule. + + y phase: team t owns columns [16t, 16t+16); in every shared-memory access + a half-warp touches 16 consecutive columns of one row. + x phase: team owns 16 consecutive shared-memory rows (odd row stride); + a half-warp's lanes are 16 different rows at the same slot. + """ + + def __init__(self, nz, NY, NX, twy, twx): + self.nz, self.NY, self.NX = nz, NY, NX + M = NX // 2 + self.M = M + self.Ms = _odd_stride(M) + self.Y1, self.Y2 = factor2(NY) + self.X1, self.X2 = factor2(M) + self.NTY = _cdiv(M, 16) + self.TWY = twy + self.W = self.NTY * twy + assert self.W <= 32 + self.NTX = _cdiv(NY, 16) # row groups + if twx is None: + twx = max(1, min(4, self.W // self.NTX)) + twx = min(twx, self.W) # a team cannot exceed the CTA + self.TWX = twx + self.NXT = max(1, self.W // twx) # x teams + assert self.NTY <= 15 and self.NXT <= 15 + self.T = 32 * self.W + self.grid = nz + tY = twiddle_table(NY, self.Y1, self.Y2) + tX = twiddle_table(M, self.X1, self.X2) + tP = twiddle_line(NX, M // 2 + 1) + self.oY = 0 + self.oX = tY.numel() + self.oP = tY.numel() + tX.numel() + self.table = torch.cat([tY, tX, tP]) + + def kernel(self): + nz, NY, M, Ms = self.nz, self.NY, self.M, self.Ms + Y1, Y2, X1, X2 = self.Y1, self.Y2, self.X1, self.X2 + TWY, T = self.TWY, self.T + NTX, TWX, NXT = self.NTX, self.TWX, self.NXT + oY, oX, oP = self.oY, self.oX, self.oP + TWF = self.table.numel() + SMF = NY * Ms * 2 + PL = NY * M + HWY = 2 * TWY # half-warps per y team + HWX = 2 * TWX # half-warps per x team + r1y = _cdiv(Y2, HWY) + r2y = _cdiv(Y1, HWY) + r1x = _cdiv(X2, HWX) + r2x = _cdiv(X1, HWX) + n_ld = NY * M + r_ld = _cdiv(n_ld, T) + n_mrg = NY // 2 + 1 + r_mrg = _cdiv(n_mrg, 32 * TWY) + TXT = 32 * TWX + n_out = 16 * M + r_out = _cdiv(n_out, TXT) + rgx = _cdiv(NTX, NXT) + + @cute.kernel + def yx_fused2_kernel(u: cute.Tensor, nyq: cute.Tensor, tw: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + z, _, _ = cute.arch.block_idx() + lane = tid % 32 + w = tid // 32 + l16 = lane % 16 + up = u.iterator + z * PL * 2 + alloc = cutlass.utils.SmemAllocator() + P = alloc.allocate_array(Float32, SMF, byte_alignment=16) + stw = alloc.allocate_array(Float32, TWF, byte_alignment=16) + spol = make_policy("evict_first") + _load_table_to_smem(tid, T, tw.iterator, stw, TWF) + cute.arch.griddepcontrol_wait() + # ---------------- load plane (8-byte, conflict-free) ---------------- + for rnd in cutlass.range_constexpr(r_ld): + e = tid + rnd * T + ec = e + if cutlass.const_expr((rnd + 1) * T > n_ld): + ec = cutlass.min(e, n_ld - 1) + r = ec // M + m = ec - r * M + v0 = ld_c(up, ec * 2) + if cutlass.const_expr((rnd + 1) * T > n_ld): + if e < n_ld: + st_c(P, (r * Ms + m) * 2, v0) + else: + st_c(P, (r * Ms + m) * 2, v0) + for rnd in cutlass.range_constexpr(_cdiv(NY, T)): + r = tid + rnd * T + rc = cutlass.min(r, NY - 1) + v = ld_c(nyq.iterator, (rc * nz + z) * 2) + if r < NY: + st_c(P, (rc * Ms + M) * 2, v) + cute.arch.sync_threads() + + # ---------------- y phase ---------------- + ty = w // TWY # y team + wt = w - ty * TWY + hw = wt * 2 + lane // 16 # half-warp index in team + ttid = wt * 32 + lane # thread index in team + jj = ty * 16 + l16 + colok = jj < M + j = cutlass.min(jj, M - 1) + if ty == 0: + for i in cutlass.range_constexpr(r_mrg): + rr = ttid + i * 32 * TWY + rc = cutlass.min(rr, n_mrg - 1) + rm = (NY - rc) % NY + a0 = ld_c(P, (rc * Ms) * 2) + am = ld_c(P, (rm * Ms) * 2) + b0 = ld_c(P, (rc * Ms + M) * 2) + bm = ld_c(P, (rm * Ms + M) * 2) + har = (a0[0] + am[0]) * Float32(0.5) + hai = (a0[1] - am[1]) * Float32(0.5) + hbr = (b0[0] + bm[0]) * Float32(0.5) + hbi = (b0[1] - bm[1]) * Float32(0.5) + if rr < n_mrg: + st_c(P, (rc * Ms) * 2, (har - hbi, hai + hbr)) + st_c(P, (rm * Ms) * 2, (har + hbi, hbr - hai)) + _named_barrier(1 + ty, 32 * TWY) + for i in cutlass.range_constexpr(r1y): + k2u = hw + i * HWY + k2 = cutlass.min(k2u, Y2 - 1) + vals = [None] * Y1 + for k1 in cutlass.range_constexpr(Y1): + vals[k1] = ld_c(P, ((k2 + Y2 * k1) * Ms + j) * 2) + res = dft(vals, 1) + for n1 in cutlass.range_constexpr(1, Y1): + res[n1] = cmul(res[n1], ld_c(stw, oY + (k2 * Y1 + n1) * 2)) + if colok & (k2u < Y2): + for n1 in cutlass.range_constexpr(Y1): + st_c(P, ((k2 + Y2 * n1) * Ms + j) * 2, res[n1]) + _named_barrier(1 + ty, 32 * TWY) + for i in cutlass.range_constexpr(r2y): + n1u = hw + i * HWY + n1 = cutlass.min(n1u, Y1 - 1) + vals2 = [None] * Y2 + for kk in cutlass.range_constexpr(Y2): + vals2[kk] = ld_c(P, ((kk + Y2 * n1) * Ms + j) * 2) + res2 = dft(vals2, 1) + if colok & (n1u < Y1): + for n2 in cutlass.range_constexpr(Y2): + st_c(P, ((n2 + Y2 * n1) * Ms + j) * 2, res2[n2]) + cute.arch.sync_threads() + + # ---------------- x phase ---------------- + tx = w // TWX + wx = w - tx * TWX + hwx = wx * 2 + lane // 16 + xtid = wx * 32 + lane + if tx < NXT: + for gi in cutlass.range_constexpr(rgx): + gg = tx + gi * NXT + if gg < NTX: + xrraw = gg * 16 + l16 + xrok = xrraw < NY + xr = cutlass.min(xrraw, NY - 1) + # stage 1 with fused C2R pre-processing (read only) + xheld = [None] * r1x + for i in cutlass.range_constexpr(r1x): + xk21u = hwx + i * HWX + xk21 = cutlass.min(xk21u, X2 - 1) + xvals = [None] * X1 + for xk1 in cutlass.range_constexpr(X1): + xkk = xk21 + X2 * xk1 + xbk = ld_c(P, (xr * Ms + xkk) * 2) + xbmk = ld_c(P, (xr * Ms + (M - xkk)) * 2) + if cutlass.const_expr(xk1 == 0): + xis0 = xk21 == 0 + xbki = cutlass.select_(xis0, Float32(0.0), xbk[1]) + xbmr = cutlass.select_(xis0, xbk[1], xbmk[0]) + xbmi = cutlass.select_(xis0, Float32(0.0), xbmk[1]) + else: + xbki = xbk[1] + xbmr = xbmk[0] + xbmi = xbmk[1] + xsr = xbk[0] + xbmr + xsi = xbki - xbmi + xdr = xbk[0] - xbmr + xdi = xbki + xbmi + xkh = cutlass.min(xkk, M - xkk) + xwv = ld_c(stw, oP + xkh * 2) + xwvr = cutlass.select_(xkk > M - xkk, -xwv[0], xwv[0]) + xwr = xdr * xwvr - xdi * xwv[1] + xwi = xdr * xwv[1] + xdi * xwvr + xvals[xk1] = (xsr - xwi, xsi + xwr) + xres = dft(xvals, 1) + for xn1 in cutlass.range_constexpr(1, X1): + xres[xn1] = cmul(xres[xn1], ld_c(stw, oX + (xk21 * X1 + xn1) * 2)) + xheld[i] = xres + _named_barrier(1 + tx, TXT) + for i in cutlass.range_constexpr(r1x): + xk21u = hwx + i * HWX + xk21 = cutlass.min(xk21u, X2 - 1) + if xrok & (xk21u < X2): + for xn1 in cutlass.range_constexpr(X1): + st_c(P, (xr * Ms + xn1 + X1 * xk21) * 2, xheld[i][xn1]) + _named_barrier(1 + tx, TXT) + # stage 2: in place, natural-order output + for i in cutlass.range_constexpr(r2x): + xn1u = hwx + i * HWX + xn1 = cutlass.min(xn1u, X1 - 1) + xvals2 = [None] * X2 + for xkk2 in cutlass.range_constexpr(X2): + xvals2[xkk2] = ld_c(P, (xr * Ms + xn1 + X1 * xkk2) * 2) + xres2 = dft(xvals2, 1) + if xrok & (xn1u < X1): + for xn2 in cutlass.range_constexpr(X2): + st_c(P, (xr * Ms + xn1 + X1 * xn2) * 2, xres2[xn2]) + _named_barrier(1 + tx, TXT) + # copy out the group's 16 rows (8-byte, coalesced) + for i in cutlass.range_constexpr(r_out): + xe = xtid + i * TXT + xec = cutlass.min(xe, n_out - 1) + xrl = xec // M + xm = xec - xrl * M + xrs = gg * 16 + xrl + xrsc = cutlass.min(xrs, NY - 1) + xy = xrsc // Y2 + Y1 * (xrsc % Y2) + xv = ld_c(P, (xrsc * Ms + xm) * 2) + if (xe < n_out) & (xrs < NY): + st_c_hint(up, (xy * M + xm) * 2, xv, spol) + + return yx_fused2_kernel diff --git a/build/torch-cuda/_c2r_f32/kernel.py b/build/torch-cuda/_c2r_f32/kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..71b025a3fe0cafb31c8620b8d0cb8276e56d7c2c --- /dev/null +++ b/build/torch-cuda/_c2r_f32/kernel.py @@ -0,0 +1,284 @@ +"""fp32 3D complex-to-real inverse FFT (cuFFT C2R convention) in CuTe DSL for the RTX 5090. + +Contract + input : float32 CUDA tensor [nz, ny, nf, 2], the packed Hermitian half spectrum + (interleaved re/im, nf = nx/2 + 1); never modified. Contiguous 16-byte + aligned inputs are read in place, others through a contiguous copy. + output: a fresh float32 [nz, ny, nx], nx = 2 * (nf - 1), equal to + torch.fft.irfftn(torch.view_as_complex(input), s=(nz, ny, nx), norm="forward") + (unnormalized inverse transform). + run(input) computes it; prepare(input) does the planning and output allocation + once and returns launch(), which takes no arguments, recomputes the output + from the input buffer's current contents and returns it. + +Routing (one plan per (shape, device), compiled on first use and cached): + * Two-pass (twopass.py; every benchmark grid except 32^3): the C2R + pre-processing is applied on the input side (row (kz, ky) paired with the + conjugate of its mirror row (-kz, -ky)), which turns the problem into a + plain 3D complex FFT of size nz x ny x nx/2. Pass 1 (DRAM bound) loads a + mirror pair of row groups, does the pre-processing, the whole x-FFT and the + first 4-step stage across rows; pass 2 (L2 bound) finishes y and z in place + on column strips. + * In-place fallback (inplace.py; tiny grids and shapes without a two-pass + config), kernels chained with programmatic dependent launch: + 1. Z pass - inverse FFT along z over all (ky, kx) columns; kx < nx/2 is + written straight into the output buffer (viewed as complex + [nz, ny, nx/2]) and the Nyquist column kx = nx/2 into a small side + buffer, so no full-size workspace is touched. + 2. If a z-plane fits in shared memory: fused in-place YX pass, one CTA per + plane (y-FFT, then C2R along x as a half-length complex FFT). + Otherwise: in-place Y pass on column blocks, then in-place X pass on + row blocks. + * Any shape neither path can plan goes through generic.py (radix stages with + runtime sizes, two full-size work buffers allocated per launch). +""" + +from __future__ import annotations + +import threading + +import torch +from cutlass import Float32, cute + +from .fftgen import factor2 +from .inplace import XPassIP, YPassIP, YXFusedIP, ZPassIP +from .twopass import TwoPassPlan, two_pass_config + +_SMEM_LIMIT = 99 * 1024 +_LOCK = threading.Lock() +_PLANS: dict = {} + + +def _fake(n): + return cute.runtime.make_fake_compact_tensor(Float32, (n,), assumed_align=16) + + +def _pick(n, choices): + for c in choices: + if n % c == 0: + return c + return None + + +def _fast_config(nz, ny, nx): + """Kernel parameters for the in-place Z -> Y -> X pipeline, or None.""" + if nx % 2: + return None + M = nx // 2 + nf = M + 1 + C = ny * nf + fz = factor2(nz) + fy = factor2(ny) + fx = factor2(M) + if fz is None or fy is None or fx is None or M < 2: + return None + TC = _pick(C, (32, 16, 8)) + if TC is None: + return None + if nz * ny * nf <= 64 * 1024: + # tiny grids: more, narrower column tiles spread the z pass over SMs + TC = _pick(C, (8, 16, 32)) + zT = TC * fz[1] + while zT > 512 and zT % 2 == 0 and zT // 2 >= TC: + zT //= 2 + if zT > 1024 or nz * TC * 8 + fz[0] * fz[1] * 8 > _SMEM_LIMIT: + return None + zcfg = {"nz": nz, "NY": ny, "NX": nx, "TC": TC, "nthreads": zT} + Ms = M + 1 if M % 2 == 0 else M + 2 + smem_yx = ny * Ms * 8 + (fy[0] * fy[1] + fx[0] * fx[1] + M // 2 + 1) * 8 + 64 + nty = (M + 15) // 16 + if smem_yx <= _SMEM_LIMIT and fy[1] <= 32 and fx[1] <= 32 and nty <= 15: + # warp-team sizes (measured on RTX 5090): wide planes prefer 2-warp + # y teams and 4-warp x teams + if M >= 56: + twy, twx = 2, 4 + elif M >= 48: + twy, twx = 4, 4 + else: + twy, twx = 4, None + if nty * twy > 32: + twy = max(1, 32 // nty) + w = nty * twy + if twx is not None: + twx = min(twx, w) + if w // twx > 15: + twx = None + if twx is None and w // max(1, min(4, w // ((ny + 15) // 16))) > 15: + return None + return {"z": zcfg, "yx": {"nz": nz, "NY": ny, "NX": nx, "twy": twy, "twx": twx}} + MC = _pick(M, (16, 8, 32, 4)) + RB = _pick(nz * ny, (16, 8, 32, 4)) + if MC is None or RB is None or MC * fy[1] > 1024: + return None + smem_y = ny * MC * 8 + fy[0] * fy[1] * 8 + smem_x = RB * Ms * 8 + (fx[0] * fx[1] + M // 2 + 1) * 8 + RB * 8 + if max(smem_y, smem_x) > _SMEM_LIMIT: + return None + return { + "z": zcfg, + "y": {"nz": nz, "NY": ny, "NX": nx, "MC": MC}, + "x": {"nz": nz, "NY": ny, "NX": nx, "RB": RB, "nthreads": 128}, + } + + +class _FusedPlan: + """Z pass + fused in-place YX pass (planes that fit in shared memory).""" + + def __init__(self, nz, ny, nx, device, cfg): + nf = nx // 2 + 1 + zp = ZPassIP(**cfg["z"]) + yp = YXFusedIP(**cfg["yx"]) + zk, yk = zp.kernel(), yp.kernel() + + @cute.jit + def launch(x: cute.Tensor, u: cute.Tensor, nq: cute.Tensor, tz: cute.Tensor, ty: cute.Tensor, stream): + zk(x, u, nq, tz).launch(grid=(zp.grid, 1, 1), block=(zp.T, 1, 1), stream=stream) + yk(u, nq, ty).launch(grid=(yp.grid, 1, 1), block=(yp.T, 1, 1), stream=stream, use_pdl=True) + + self.fn = cute.compile( + launch, + _fake(nz * ny * nf * 2), + _fake(nz * ny * nx), + _fake(ny * nz * 2), + _fake(zp.table.numel()), + _fake(yp.table.numel()), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + self.tz = zp.table.to(device) + self.ty = yp.table.to(device) + self.nyq_numel = ny * nz * 2 + + def __call__(self, x_flat, out_flat, nyq): + self.fn(x_flat, out_flat, nyq, self.tz, self.ty) + + +class _FastPlan: + """Z pass + in-place Y pass + in-place X pass (planes larger than shared memory).""" + + def __init__(self, nz, ny, nx, device, cfg): + nf = nx // 2 + 1 + zp = ZPassIP(**cfg["z"]) + yp = YPassIP(**cfg["y"]) + xp = XPassIP(**cfg["x"]) + zk, yk, xk = zp.kernel(), yp.kernel(), xp.kernel() + + @cute.jit + def launch( + x: cute.Tensor, + u: cute.Tensor, + nq: cute.Tensor, + tz: cute.Tensor, + ty: cute.Tensor, + tx: cute.Tensor, + stream, + ): + zk(x, u, nq, tz).launch(grid=(zp.grid, 1, 1), block=(zp.T, 1, 1), stream=stream) + yk(u, nq, ty).launch(grid=(yp.grid, 1, 1), block=(yp.T, 1, 1), stream=stream, use_pdl=True) + xk(u, nq, tx).launch(grid=(xp.grid, 1, 1), block=(xp.T, 1, 1), stream=stream, use_pdl=True) + + self.fn = cute.compile( + launch, + _fake(nz * ny * nf * 2), + _fake(nz * ny * nx), + _fake(ny * nz * 2), + _fake(zp.table.numel()), + _fake(yp.table.numel()), + _fake(xp.table.numel()), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + self.tz = zp.table.to(device) + self.ty = yp.table.to(device) + self.tx = xp.table.to(device) + self.nyq_numel = ny * nz * 2 + + def __call__(self, x_flat, out_flat, nyq): + self.fn(x_flat, out_flat, nyq, self.tz, self.ty, self.tx) + + +def _get_plan(nz, ny, nx, device): + key = (nz, ny, nx, device.index) + plan = _PLANS.get(key) + if plan is not None: + return plan + with _LOCK: + plan = _PLANS.get(key) + if plan is None: + cfg = _fast_config(nz, ny, nx) + cfg2 = two_pass_config(nz, ny, nx, _SMEM_LIMIT) + if cfg2 is not None: + plan = TwoPassPlan(nz, ny, nx, device, **cfg2) + elif cfg is not None and "yx" in cfg: + plan = _FusedPlan(nz, ny, nx, device, cfg) + elif cfg is not None: + plan = _FastPlan(nz, ny, nx, device, cfg) + else: + from .generic import GenericPlan # only shapes no fast path can plan need it + + plan = GenericPlan(nz, ny, nx, device) + _PLANS[key] = plan + return plan + + +def _check(input): + if input.dim() != 4 or input.shape[-1] != 2: + raise ValueError("input must be [nz, ny, nf, 2]") + if input.dtype != torch.float32: + raise TypeError("input must be float32") + if not input.is_cuda: + raise ValueError("input must be a CUDA tensor") + + +def _prepare(input): + _check(input) + nz, ny, nf, _ = input.shape + nx = 2 * (nf - 1) + plan = _get_plan(nz, ny, nx, input.device) + out = torch.empty((nz, ny, nx), dtype=torch.float32, device=input.device) + out_flat = out.view(-1) + # per-launch side buffer (Nyquist column) so independent launches never share state + extra = () + if hasattr(plan, "nyq_numel"): + extra = (torch.empty(plan.nyq_numel, dtype=torch.float32, device=input.device),) + if input.is_contiguous() and input.data_ptr() % 16 == 0: + x_flat = input.view(-1) + + def launch(): + plan(x_flat, out_flat, *extra) + return out + elif input.is_contiguous(): + # kernels assume a 16-byte aligned input; stage misaligned views + stage = torch.empty(input.numel(), dtype=torch.float32, device=input.device) + + def launch(): + stage.copy_(input.view(-1)) + plan(stage, out_flat, *extra) + return out + else: + + def launch(): + plan(input.contiguous().view(-1), out_flat, *extra) + return out + + return launch + + +def prepare(input): + """Plan for ``input`` and return ``launch()`` (compiles and launches on the input's device).""" + dev = input.device + if dev.type == "cuda" and dev.index is not None and dev.index != torch.cuda.current_device(): + with torch.cuda.device(dev): + launch = _prepare(input) + + def guarded_launch(): + with torch.cuda.device(dev): + return launch() + + return guarded_launch + return _prepare(input) + + +def run(input): + """Return the unnormalized 3D C2R inverse FFT of ``input`` as a new [nz, ny, nx] tensor.""" + return prepare(input)() diff --git a/build/torch-cuda/_c2r_f32/tables.py b/build/torch-cuda/_c2r_f32/tables.py new file mode 100644 index 0000000000000000000000000000000000000000..e640ce716d8409508f4cc61f145688265c7a7380 --- /dev/null +++ b/build/torch-cuda/_c2r_f32/tables.py @@ -0,0 +1,33 @@ +"""Host-side twiddle tables, built once per plan. + +Every table is a flat float32 tensor of interleaved (re, im) pairs on the CPU; +the plans copy it to the device. Angles are formed in float64 as 2*pi*k/N +(k reduced mod N where noted) and evaluated with the C library's cos/sin +(``math``), then rounded to float32. ``math`` is used rather than torch.cos / +torch.sin because the latter (vectorized, 1-ulp) do not round every value the +same way. +""" + +from __future__ import annotations + +import math + +import torch + + +def _cis(angles): + vals = [] + for a in angles: + vals.append(math.cos(a)) + vals.append(math.sin(a)) + return torch.tensor(vals, dtype=torch.float64).to(torch.float32) + + +def twiddle_line(N, count): + """t[k] = exp(2*pi*i*k/N), k in [0, count).""" + return _cis(2.0 * math.pi * k / N for k in range(count)) + + +def twiddle_table(N, A, B): + """4-step twiddles t[b*A + a] = exp(2*pi*i*((a*b) mod N)/N), a in [0, A), b in [0, B).""" + return _cis(2.0 * math.pi * ((a * b) % N) / N for b in range(B) for a in range(A)) diff --git a/build/torch-cuda/_c2r_f32/twopass.py b/build/torch-cuda/_c2r_f32/twopass.py new file mode 100644 index 0000000000000000000000000000000000000000..9208227faa71a4eb4fe0e5a3857d77f9cdaf9951 --- /dev/null +++ b/build/torch-cuda/_c2r_f32/twopass.py @@ -0,0 +1,1115 @@ +"""Two-pass formulation of the 3D C2R transform (the main path). + +The C2R pre-processing that turns a length-nx real inverse DFT into a +length-M (M = nx/2) complex one commutes with the z/y transforms once it is +applied to the *input* with the (kz, ky) index mirrored: + + W[kz,ky,k] = a + b + i*w^k*(a - b), a = X[kz,ky,k], + b = conj(X[-kz,-ky,M-k]), w = e^{2 pi i/nx} + (k = 0 uses the Hermitian parts of the k = 0 and k = M bins) + +and u[z,y,2m] + i*u[z,y,2m+1] is the plain 3D complex inverse DFT of W. All +three axes are then ordinary complex FFTs that can be reordered and split. +With nz = Az*Bz, ny = Ay*By (4-step split, kz = Bz*k1z + k2z, z = n1z + Az*n2z): + +RowPassA (pass 1, DRAM bound): one CTA per mirror pair of groups + (k2z, k2y) / (-k2z, -k2y), i.e. 2*Az*Ay input rows, bulk-loaded with + 16-byte loads. Applies the pre-processing, the full M-point x-FFT, + the Az x Ay first-stage DFT across rows and the 4-step twiddles, and + writes the rows to u[n1z + Az*k2z, n1y + Ay*k2y, :]. With more than + two waves of CTAs, each CTA also bulk-prefetches into L2 the rows of the + CTA two waves later, so the last waves do not queue behind the DRAM stream. +StripPass (pass 2, L2 bound, in place): one CTA per (n1z, n1y, column strip): + the Bz x By second-stage DFT on rows u[n1z + Az*k2z, n1y + Ay*k2y]. +StripPassR the same with each Bz-point z column held in registers, for tiles + larger than shared memory (Bz <= 64). + +``two_pass_config`` picks the group split (Az, Ay), strip width and thread +counts for a shape (or rejects it); ``TwoPassPlan`` compiles the two kernels. +""" + +from __future__ import annotations + +import math + +import cutlass +import torch +from cutlass import Float32, Int32, cute + +from .fftgen import cmul, dft, factor2, ld_c, ld_c2_hint, make_policy, st_c +from .tables import twiddle_line, twiddle_table + + +def _cdiv(a, b): + return (a + b - 1) // b + + +def _fake(n, dt=Float32): + return cute.runtime.make_fake_compact_tensor(dt, (n,), assumed_align=16) + + +def _wavefronts(addr_lists): + """Shared-memory wavefronts for 8-byte accesses (half-warp model). + + addr_lists: iterable of per-instruction lists of 32 word addresses (8-byte + units) or None for inactive lanes. + """ + tot = 0 + for addrs in addr_lists: + for h in (0, 16): + banks = {} + for a in addrs[h : h + 16]: + if a is None: + continue + banks.setdefault(a % 16, set()).add(a) + tot += max((len(v) for v in banks.values()), default=0) + return tot + + +def _pair_regions(Bz, By): + Hz, Hy = (Bz - 1) // 2, (By - 1) // 2 + nsz = 2 if Bz % 2 == 0 else 1 # self-mirrored k2z values: 0 (, Bz/2) + nsy = 2 if By % 2 == 0 else 1 + return Hz, Hy, nsz, nsy + + +def pair_table(Bz, By): + """CTA -> (g0, g1, self) with g = k2z*By + k2y: mirror group pairs, then the + self-mirrored groups paired up. Same order as ``pair_of`` (computed on device). + + Flat int32 tensor; its length fixes the pass-1 grid. It is also passed to + the pass-1 kernel, whose device code recomputes the entries with ``pair_of`` + instead of loading them.""" + Hz, Hy, nsz, nsy = _pair_regions(Bz, By) + out = [] + for z in range(1, Hz + 1): # k2z in (0, Bz/2): mirror k2z' = Bz - k2z + for y in range(By): + out.append((z * By + y, (Bz - z) * By + (By - y) % By, 0)) + for zi in range(nsz): # k2z self-mirrored: pair k2y with By - k2y + z = zi * (Bz // 2) + for y in range(1, Hy + 1): + out.append((z * By + y, z * By + By - y, 0)) + selfs = [(zi * (Bz // 2)) * By + yi * (By // 2) for zi in range(nsz) for yi in range(nsy)] + for i in range(0, len(selfs), 2): + out.append((selfs[i], selfs[min(i + 1, len(selfs) - 1)], 1)) + return torch.tensor(out, dtype=torch.int32).reshape(-1) + + +def pair_of(c, Bz, By): + """Device: CTA index c -> (k2z0, k2y0, k2z1, k2y1, self flag), the ``pair_table`` entry, + computed arithmetically (no dependent global load before the row loads).""" + Hz, Hy, nsz, nsy = _pair_regions(Bz, By) + n1 = Hz * By + n2 = nsz * Hy + nself = nsz * nsy + # region 1 + a_z = 1 + c // By + a_y = c - (c // By) * By + # region 2 + j2 = cutlass.max(c - n1, 0) + q2 = j2 // max(Hy, 1) + b_z = q2 * (Bz // 2) + b_y = 1 + j2 - q2 * max(Hy, 1) + # region 3 + j3 = cutlass.max(c - n1 - n2, 0) + i0 = 2 * j3 + i1 = cutlass.min(2 * j3 + 1, nself - 1) + s0z = (i0 // nsy) * (Bz // 2) + s0y = (i0 - (i0 // nsy) * nsy) * (By // 2) + s1z = (i1 // nsy) * (Bz // 2) + s1y = (i1 - (i1 // nsy) * nsy) * (By // 2) + r1 = c < n1 + r2 = c < n1 + n2 + k2z0 = cutlass.select_(r1, a_z, cutlass.select_(r2, b_z, s0z)) + k2y0 = cutlass.select_(r1, a_y, cutlass.select_(r2, b_y, s0y)) + k2z1 = cutlass.select_(r1, Bz - a_z, cutlass.select_(r2, b_z, s1z)) + k2y1 = cutlass.select_(r1, (By - a_y) % By, cutlass.select_(r2, By - b_y, s1y)) + sf = cutlass.select_(r2, 0, 1) + return k2z0, k2y0, k2z1, k2y1, sf + + +class RowPassA: + """Pass 1, pair variant: one CTA owns a group and its mirror group (2R rows). + + Rows are bulk-loaded once with 16-byte loads, the pre-processing reads + both rows of every mirror pair from shared memory. + """ + + def __init__(self, nz, ny, nx, Az, Ay, nthreads): + M = nx // 2 + self.nz, self.ny, self.nx, self.M = nz, ny, nx, M + assert nz % Az == 0 and ny % Ay == 0 + self.Az, self.Ay = Az, Ay + self.Bz, self.By = nz // Az, ny // Ay + self.R = Az * Ay + fx = factor2(M) + assert fx is not None + self.X1, self.X2 = fx + self.T = nthreads + self.ptab = pair_table(self.Bz, self.By) + self.grid = self.ptab.numel() // 3 + t1 = twiddle_table(M, self.X1, self.X2) + pw = twiddle_line(nx, M) + tz = twiddle_line(nz, nz) + ty = twiddle_line(ny, ny) + self.o_t1 = 0 + self.o_pw = self.o_t1 + t1.numel() + self.o_tz = self.o_pw + pw.numel() + self.o_ty = self.o_tz + tz.numel() + self.table = torch.cat([t1, pw, tz, ty]) + self.NCH = (M + 1) // 2 + 1 + # phase 3 (the Az x Ay first-stage DFT) fused into one register pass for small groups + self.fuse3 = self.R <= 16 + # paired phase 1 (one task serves a row and its mirror row) pays off once + # the unpaired version needs more than one task per thread + self.pair1 = 2 * self.R * self.X2 > nthreads + # split loading into two mirror-closed halves (x stages of half A overlap + # half B's loads) when each half still has a phase-1 task per thread + self.halves = self.R * self.X2 >= nthreads and self.R % 2 == 0 + self.Ms = self._choose_ms() + self.n_tws = (t1.numel() + pw.numel()) // 2 # complex twiddles staged in smem + self.smem_bytes = (2 * self.R * self.Ms + self.n_tws + 2 * (Az + Ay)) * 8 + 64 + + def _choose_ms(self): + """Shared-memory row stride (complex) with the fewest bank wavefronts in phases 1-2.""" + M, R, T, X1, X2 = self.M, self.R, self.T, self.X1, self.X2 + NR = 2 * R + best = None + for Ms in range(M + 1, M + 18): + insts = [] + + def pos(r, n1, k2, Ms=Ms): + v = n1 + k2 + if v >= X1: + v -= X1 + return r * Ms + X1 * k2 + v + + n1t = NR * X2 + for w0 in range(0, min(T, n1t), 32): + for k1 in range(X1): + insts.append( + [ + None if w0 + ln >= n1t else ((w0 + ln) // X2) * Ms + (w0 + ln) % X2 + X2 * k1 + for ln in range(32) + ] + ) + for n1 in range(X1): + insts.append( + [None if w0 + ln >= n1t else pos((w0 + ln) // X2, n1, (w0 + ln) % X2) for ln in range(32)] + ) + n2t = NR * X1 + for w0 in range(0, min(T, n2t), 32): + for k2 in range(X2): + insts.append( + [None if w0 + ln >= n2t else pos((w0 + ln) // X1, (w0 + ln) % X1, k2) for ln in range(32)] + ) + wf = _wavefronts(insts) + if best is None or wf < best[0]: + best = (wf, Ms) + return best[1] + + def kernel(self, pf): + """Build the pass-1 kernel. + + pf > 0: once its own rows have landed in shared memory, CTA b bulk-prefetches + into L2 the input rows of CTA b + pf (the CTA that will take its place in a + later wave), so DRAM stays busy through the compute phases of the current wave. + """ + nz, ny, M, R = self.nz, self.ny, self.M, self.R + Az, Ay, Bz, By = self.Az, self.Ay, self.Bz, self.By + X1, X2, T, Ms = self.X1, self.X2, self.T, self.Ms + NF = M + 1 + NCH = self.NCH + NR = 2 * R + TOT = nz * ny * NF + o_t1, o_pw, o_tz, o_ty = self.o_t1, self.o_pw, self.o_tz, self.o_ty + SMF = NR * Ms * 2 + n0t = NR * NCH + n1t = NR * X2 + n2t = NR * X1 + n3a = 2 * Az * M + n3b = 2 * Ay * M + r1 = _cdiv(n1t, T) + n1p = R * X2 # paired phase-1 tasks (slot-0 rows) + r1p = _cdiv(n1p, T) + PAIR1 = self.pair1 + HALVES = self.halves + RH = R // 2 + NTW = self.n_tws + PF = pf + GRID = self.grid + PFB = _cdiv(NF + 1, 2) * 16 # bytes of a row's 16-byte-aligned cover + fuse3 = self.fuse3 + n3f = 2 * M + + def rot(v): + return cutlass.select_(v >= X1, v - X1, v) + + @cute.kernel + def rowpassa_kernel(x: cute.Tensor, u: cute.Tensor, tw: cute.Tensor, pt: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + cute.arch.griddepcontrol_launch_dependents() + xp = x.iterator + up = u.iterator + twp = tw.iterator + alloc = cutlass.utils.SmemAllocator() + P = alloc.allocate_array(Float32, SMF, byte_alignment=16) + TWZ = alloc.allocate_array(Float32, 2 * Az * 2, byte_alignment=16) + TWY = alloc.allocate_array(Float32, 2 * Ay * 2, byte_alignment=16) + STW = alloc.allocate_array(Float32, NTW * 2, byte_alignment=16) + lpol = make_policy("evict_first") + k2z0, k2y0, k2z1, k2y1, sf = pair_of(bid, Bz, By) + + def prefetch_next(): + # bulk L2 prefetch of the 2R input rows of the pair of CTA bid + PF + nb_ = bid + PF + if (nb_ < GRID) & (tid < NR): + p0z, p0y, p1z, p1y, _sf = pair_of(nb_, Bz, By) + s_ = tid // R + r_ = tid - s_ * R + h2z = cutlass.select_(s_ == 0, p0z, p1z) + h2y = cutlass.select_(s_ == 0, p0y, p1y) + q1z = r_ // Ay + q1y = r_ - q1z * Ay + rb_ = ((Bz * q1z + h2z) * ny + (By * q1y + h2y)) * NF + st_ = rb_ - (rb_ & 1) + nbytes = (cutlass.min(PFB // 8, TOT - st_) * 8) & -16 + cute.arch.inline_ptx( + "cp.async.bulk.prefetch.L2.global [{$r0}], {$r1};", + read_only_args=[(xp + st_ * 2).toint(), nbytes], + ) + + def tables(): + # twiddle tables -> smem (overlaps the row loads) + for rnd in cutlass.range_constexpr(_cdiv(NTW, T)): + q = tid + rnd * T + qc = q + if cutlass.const_expr((rnd + 1) * T > NTW): + qc = cutlass.min(q, NTW - 1) + st_c(STW, qc * 2, ld_c(twp, qc * 2)) + if tid < 2 * Ay: + ss = tid // Ay + jy = tid - ss * Ay + kk = cutlass.select_(ss == 0, k2y0, k2y1) + st_c(TWY, tid * 2, ld_c(twp, o_ty + ((kk * jy) % ny) * 2)) + if cutlass.const_expr(32 + 2 * Az <= T): + if tid >= 32: # noqa: SIM102 + if tid < 32 + 2 * Az: + q = tid - 32 + ss = q // Az + jz = q - ss * Az + kk = cutlass.select_(ss == 0, k2z0, k2z1) + st_c(TWZ, q * 2, ld_c(twp, o_tz + ((kk * jz) % nz) * 2)) + else: + # the 2*Az entries do not fit in threads [32, T) (Az > 48 with 128 + # threads): rotate by 32 so that warp 0 takes the tail entries + for rnd in cutlass.range_constexpr(_cdiv(2 * Az, T)): + q = (tid + (T - 32)) % T + rnd * T + if q < 2 * Az: + ss = q // Az + jz = q - ss * Az + kk = cutlass.select_(ss == 0, k2z0, k2z1) + st_c(TWZ, q * 2, ld_c(twp, o_tz + ((kk * jz) % nz) * 2)) + + def phases_012(): + # ---- phase 0: bulk-load 2R rows (16-byte loads, all issued first) ---- + r0 = _cdiv(n0t, T) + ldv = [None] * r0 + for rnd in cutlass.range_constexpr(r0): + c = tid + rnd * T + cc = c + if cutlass.const_expr((rnd + 1) * T > n0t): + cc = cutlass.min(c, n0t - 1) + row = cc // NCH + j = cc - row * NCH + s = row // R + r = row - s * R + k1z = r // Ay + k1y = r - k1z * Ay + k2z = cutlass.select_(s == 0, k2z0, k2z1) + k2y = cutlass.select_(s == 0, k2y0, k2y1) + rb = ((Bz * k1z + k2z) * ny + (By * k1y + k2y)) * NF + dl = rb % 2 + e = cutlass.min(rb - dl + 2 * j, TOT - 2) + ldv[rnd] = (ld_c2_hint(xp, e * 2, lpol), row, 2 * j - dl) + tables() + for rnd in cutlass.range_constexpr(r0): + (v0, v1), row, k0 = ldv[rnd] + c = tid + rnd * T + ok = c < n0t + if ok & (k0 >= 0): + st_c(P, (row * Ms + k0) * 2, v0) + if ok & (k0 + 1 <= M): + st_c(P, (row * Ms + k0 + 1) * 2, v1) + cute.arch.sync_threads() + if cutlass.const_expr(PF > 0): + prefetch_next() + + # ---- phase 1: pre-processing + x stage 1 (held across a barrier) ---- + if (sf == 0) & PAIR1: + # Mirror groups in different slots: one task (slot-0 row r, residue k2) + # reads the pair (r, k2 + X2*k1) / (pi(r), M - k2 - X2*k1) once and + # produces both rows' W values: W_r[k] = s + p, + # W_pi(r)[M-k] = conj(s - p) (k = 0: Hermitian parts, same formulas). + heldp = [None] * r1p + for rnd in cutlass.range_constexpr(r1p): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > n1p): + tc = cutlass.min(t, n1p - 1) + k2 = tc % X2 + r = tc // X2 + k1z = r // Ay + k1y = r - k1z * Ay + mk1z = cutlass.select_(k2z0 == 0, cutlass.select_(k1z == 0, 0, Az - k1z), Az - 1 - k1z) + mk1y = cutlass.select_(k2y0 == 0, cutlass.select_(k1y == 0, 0, Ay - k1y), Ay - 1 - k1y) + ob = r * Ms + mb = (R + mk1z * Ay + mk1y) * Ms + is0 = k2 == 0 + wo = [None] * X1 + wm = [None] * X1 + for k1 in cutlass.range_constexpr(X1): + a = ld_c(P, (ob + k2 + X2 * k1) * 2) + cm = ld_c(P, (mb + M - k2 - X2 * k1) * 2) + ar, ai = a + br = cm[0] + bi = -cm[1] + if cutlass.const_expr(k1 == 0): + e1 = ld_c(P, cutlass.select_(is0, mb, ob + k2) * 2) + e2 = ld_c(P, cutlass.select_(is0, ob + M, ob + k2) * 2) + ar = cutlass.select_(is0, (ar + e1[0]) * Float32(0.5), ar) + ai = cutlass.select_(is0, (ai - e1[1]) * Float32(0.5), ai) + br = cutlass.select_(is0, (e2[0] + br) * Float32(0.5), br) + bi = cutlass.select_(is0, (e2[1] + bi) * Float32(0.5), bi) + sr = ar + br + si = ai + bi + dr = ar - br + di = ai - bi + w = ld_c(STW, o_pw + (k2 + X2 * k1) * 2) + pr = dr * w[0] - di * w[1] + pi = dr * w[1] + di * w[0] + # p = i*w*d = (-pi, pr) + wo[k1] = (sr - pi, si + pr) + wm[k1] = (sr + pi, pr - si) # conj(s - p) + k2m = cutlass.select_(is0, 0, X2 - k2) + vm = [None] * X1 + for j in cutlass.range_constexpr(X1): + ja = (X1 - j) % X1 # k2 == 0: k1 = (X1 - j) mod X1 + jb = X1 - 1 - j # k2 > 0 : k1 = X1 - 1 - j + vm[j] = ( + cutlass.select_(is0, wm[ja][0], wm[jb][0]), + cutlass.select_(is0, wm[ja][1], wm[jb][1]), + ) + ro = dft(wo, 1) + rm = dft(vm, 1) + for n1 in cutlass.range_constexpr(1, X1): + ro[n1] = cmul(ro[n1], ld_c(STW, o_t1 + (k2 * X1 + n1) * 2)) + rm[n1] = cmul(rm[n1], ld_c(STW, o_t1 + (k2m * X1 + n1) * 2)) + heldp[rnd] = (ro, rm, mb, k2m) + cute.arch.sync_threads() + for rnd in cutlass.range_constexpr(r1p): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > n1p): + tc = cutlass.min(t, n1p - 1) + k2 = tc % X2 + r = tc // X2 + ro, rm, mb, k2m = heldp[rnd] + base_o = r * Ms + X1 * k2 + base_m = mb + X1 * k2m + if cutlass.const_expr((rnd + 1) * T > n1p): + if t < n1p: + for n1 in cutlass.range_constexpr(X1): + st_c(P, (base_o + rot(k2 + n1)) * 2, ro[n1]) + st_c(P, (base_m + rot(k2m + n1)) * 2, rm[n1]) + else: + for n1 in cutlass.range_constexpr(X1): + st_c(P, (base_o + rot(k2 + n1)) * 2, ro[n1]) + st_c(P, (base_m + rot(k2m + n1)) * 2, rm[n1]) + cute.arch.sync_threads() + else: + # ---- phase 1: pre-processing + x stage 1 (held across a barrier) ---- + held = [None] * r1 + for rnd in cutlass.range_constexpr(r1): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > n1t): + tc = cutlass.min(t, n1t - 1) + k2 = tc % X2 + row = tc // X2 + s = row // R + r = row - s * R + k1z = r // Ay + k1y = r - k1z * Ay + k2z = cutlass.select_(s == 0, k2z0, k2z1) + k2y = cutlass.select_(s == 0, k2y0, k2y1) + mk1z = cutlass.select_(k2z == 0, cutlass.select_(k1z == 0, 0, Az - k1z), Az - 1 - k1z) + mk1y = cutlass.select_(k2y == 0, cutlass.select_(k1y == 0, 0, Ay - k1y), Ay - 1 - k1y) + ms = cutlass.select_(sf == 0, 1 - s, s) + mrow = ms * R + mk1z * Ay + mk1y + ob = row * Ms + mb = mrow * Ms + is0 = k2 == 0 + vals = [None] * X1 + for k1 in cutlass.range_constexpr(X1): + a = ld_c(P, (ob + k2 + X2 * k1) * 2) + cm = ld_c(P, (mb + M - k2 - X2 * k1) * 2) + ar, ai = a + br = cm[0] + bi = -cm[1] + if cutlass.const_expr(k1 == 0): + e1 = ld_c(P, cutlass.select_(is0, mb, ob + k2) * 2) + e2 = ld_c(P, cutlass.select_(is0, ob + M, ob + k2) * 2) + ar = cutlass.select_(is0, (ar + e1[0]) * Float32(0.5), ar) + ai = cutlass.select_(is0, (ai - e1[1]) * Float32(0.5), ai) + br = cutlass.select_(is0, (e2[0] + br) * Float32(0.5), br) + bi = cutlass.select_(is0, (e2[1] + bi) * Float32(0.5), bi) + sr = ar + br + si = ai + bi + dr = ar - br + di = ai - bi + w = ld_c(STW, o_pw + (k2 + X2 * k1) * 2) + pr = dr * w[0] - di * w[1] + pi = dr * w[1] + di * w[0] + vals[k1] = (sr - pi, si + pr) + res = dft(vals, 1) + for n1 in cutlass.range_constexpr(1, X1): + res[n1] = cmul(res[n1], ld_c(STW, o_t1 + (k2 * X1 + n1) * 2)) + held[rnd] = res + cute.arch.sync_threads() + for rnd in cutlass.range_constexpr(r1): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > n1t): + tc = cutlass.min(t, n1t - 1) + k2 = tc % X2 + row = tc // X2 + base = row * Ms + X1 * k2 + if cutlass.const_expr((rnd + 1) * T > n1t): + if t < n1t: + for n1 in cutlass.range_constexpr(X1): + st_c(P, (base + rot(k2 + n1)) * 2, held[rnd][n1]) + else: + for n1 in cutlass.range_constexpr(X1): + st_c(P, (base + rot(k2 + n1)) * 2, held[rnd][n1]) + cute.arch.sync_threads() + + # ---- phase 2: x stage 2 (in place) ---- + for rnd in cutlass.range_constexpr(_cdiv(n2t, T)): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > n2t): + tc = cutlass.min(t, n2t - 1) + n1 = tc % X1 + row = tc // X1 + vals = [None] * X2 + for k2 in cutlass.range_constexpr(X2): + vals[k2] = ld_c(P, (row * Ms + X1 * k2 + rot(n1 + k2)) * 2) + res = dft(vals, 1) + if cutlass.const_expr((rnd + 1) * T > n2t): + if t < n2t: + for n2 in cutlass.range_constexpr(X2): + st_c(P, (row * Ms + X1 * n2 + rot(n1 + n2)) * 2, res[n2]) + else: + for n2 in cutlass.range_constexpr(X2): + st_c(P, (row * Ms + X1 * n2 + rot(n1 + n2)) * 2, res[n2]) + cute.arch.sync_threads() + + def phases_012_halves(): + # Mirror-group CTAs (sf == 0): rows are loaded in two halves, each a + # set of complete mirror pairs (slot-0 rows [h0, h0+Rh) and their + # mirrors in slot 1); x stages of half A overlap the loads of half B. + def half_row(h, q): + """Row q of half h -> (own smem row, mirror smem row, kz, ky). + + sf == 0: half h = slot-0 rows [h*RH, (h+1)*RH) + their slot-1 mirrors; + sf == 1 (self-mirrored slots): half h = slot h. + """ + is_s0 = q < RH + r = h * RH + cutlass.select_(is_s0, q, q - RH) + k1z = r // Ay + k1y = r - k1z * Ay + mk1z = cutlass.select_(k2z0 == 0, cutlass.select_(k1z == 0, 0, Az - k1z), Az - 1 - k1z) + mk1y = cutlass.select_(k2y0 == 0, cutlass.select_(k1y == 0, 0, Ay - k1y), Ay - 1 - k1y) + pr_ = R + mk1z * Ay + mk1y + own0 = cutlass.select_(is_s0, r, pr_) + mir0 = cutlass.select_(is_s0, pr_, r) + kz0 = cutlass.select_(is_s0, Bz * k1z + k2z0, Bz * mk1z + k2z1) + ky0 = cutlass.select_(is_s0, By * k1y + k2y0, By * mk1y + k2y1) + k2zh = k2z0 if h == 0 else k2z1 + k2yh = k2y0 if h == 0 else k2y1 + q1z = q // Ay + q1y = q - q1z * Ay + m1z = cutlass.select_(k2zh == 0, cutlass.select_(q1z == 0, 0, Az - q1z), Az - 1 - q1z) + m1y = cutlass.select_(k2yh == 0, cutlass.select_(q1y == 0, 0, Ay - q1y), Ay - 1 - q1y) + s0 = sf == 0 + own = cutlass.select_(s0, own0, h * R + q) + mir = cutlass.select_(s0, mir0, h * R + m1z * Ay + m1y) + kz = cutlass.select_(s0, kz0, Bz * q1z + k2zh) + ky = cutlass.select_(s0, ky0, By * q1y + k2yh) + return own, mir, kz, ky + + def issue(h): + nt = R * NCH + out = [None] * _cdiv(nt, T) + for rnd in cutlass.range_constexpr(_cdiv(nt, T)): + c = tid + rnd * T + cc = c + if cutlass.const_expr((rnd + 1) * T > nt): + cc = cutlass.min(c, nt - 1) + q = cc // NCH + j = cc - q * NCH + own, _mir, kz, ky = half_row(h, q) + rb = (kz * ny + ky) * NF + dl = rb % 2 + e = cutlass.min(rb - dl + 2 * j, TOT - 2) + out[rnd] = (ld_c2_hint(xp, e * 2, lpol), own, 2 * j - dl, c < nt) + return out + + def stage(lds): + for (v0, v1), row, k0, ok in lds: + if ok & (k0 >= 0): + st_c(P, (row * Ms + k0) * 2, v0) + if ok & (k0 + 1 <= M): + st_c(P, (row * Ms + k0 + 1) * 2, v1) + + def ph1_compute(h): + nt = R * X2 + held = [None] * _cdiv(nt, T) + for rnd in cutlass.range_constexpr(_cdiv(nt, T)): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > nt): + tc = cutlass.min(t, nt - 1) + k2 = tc % X2 + q = tc // X2 + own, mir, _kz, _ky = half_row(h, q) + ob = own * Ms + mb = mir * Ms + is0 = k2 == 0 + vals = [None] * X1 + for k1 in cutlass.range_constexpr(X1): + a_ = ld_c(P, (ob + k2 + X2 * k1) * 2) + cm = ld_c(P, (mb + M - k2 - X2 * k1) * 2) + ar, ai = a_ + br = cm[0] + bi = -cm[1] + if cutlass.const_expr(k1 == 0): + e1 = ld_c(P, cutlass.select_(is0, mb, ob + k2) * 2) + e2 = ld_c(P, cutlass.select_(is0, ob + M, ob + k2) * 2) + ar = cutlass.select_(is0, (ar + e1[0]) * Float32(0.5), ar) + ai = cutlass.select_(is0, (ai - e1[1]) * Float32(0.5), ai) + br = cutlass.select_(is0, (e2[0] + br) * Float32(0.5), br) + bi = cutlass.select_(is0, (e2[1] + bi) * Float32(0.5), bi) + sr = ar + br + si = ai + bi + dr = ar - br + di = ai - bi + w = ld_c(STW, o_pw + (k2 + X2 * k1) * 2) + pr = dr * w[0] - di * w[1] + pi = dr * w[1] + di * w[0] + vals[k1] = (sr - pi, si + pr) + res = dft(vals, 1) + for n1 in cutlass.range_constexpr(1, X1): + res[n1] = cmul(res[n1], ld_c(STW, o_t1 + (k2 * X1 + n1) * 2)) + held[rnd] = (res, own * Ms + X1 * k2, k2, t < nt) + return held + + def ph1_write(held): + for res, base, k2, ok in held: + if ok: + for n1 in cutlass.range_constexpr(X1): + st_c(P, (base + rot(k2 + n1)) * 2, res[n1]) + + def ph2(h): + nt = R * X1 + for rnd in cutlass.range_constexpr(_cdiv(nt, T)): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > nt): + tc = cutlass.min(t, nt - 1) + n1 = tc % X1 + q = tc // X1 + own, _mir, _kz, _ky = half_row(h, q) + vals = [None] * X2 + for k2 in cutlass.range_constexpr(X2): + vals[k2] = ld_c(P, (own * Ms + X1 * k2 + rot(n1 + k2)) * 2) + res = dft(vals, 1) + if t < nt: + for n2 in cutlass.range_constexpr(X2): + st_c(P, (own * Ms + X1 * n2 + rot(n1 + n2)) * 2, res[n2]) + + ldA = issue(0) + ldB = issue(1) + tables() + stage(ldA) + cute.arch.sync_threads() + hA = ph1_compute(0) + cute.arch.sync_threads() + ph1_write(hA) + stage(ldB) + cute.arch.sync_threads() + if cutlass.const_expr(PF > 0): + prefetch_next() + ph2(0) + hB = ph1_compute(1) + cute.arch.sync_threads() + ph1_write(hB) + cute.arch.sync_threads() + ph2(1) + cute.arch.sync_threads() + + if cutlass.const_expr(HALVES): + phases_012_halves() + else: + phases_012() + + # ---- fused phase 3: Az x Ay DFT per column in registers -> global ---- + if cutlass.const_expr(fuse3): + for rnd in cutlass.range_constexpr(_cdiv(n3f, T)): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > n3f): + tc = cutlass.min(t, n3f - 1) + m = tc % M + s = tc // M + q = m // X1 + pm = X1 * q + rot(m - X1 * q + q) + k2z = cutlass.select_(s == 0, k2z0, k2z1) + k2y = cutlass.select_(s == 0, k2y0, k2y1) + grid2 = [[None] * Ay for _ in range(Az)] + for k1z in cutlass.range_constexpr(Az): + vy = [None] * Ay + for k1y in cutlass.range_constexpr(Ay): + vy[k1y] = ld_c(P, ((s * R + k1z * Ay + k1y) * Ms + pm) * 2) + ry = dft(vy, 1) + for n1y in cutlass.range_constexpr(Ay): + grid2[k1z][n1y] = ry[n1y] + ob = ((Az * k2z) * ny + Ay * k2y) * M + m + for n1y in cutlass.range_constexpr(Ay): + vz = [grid2[k1z][n1y] for k1z in range(Az)] + rz = dft(vz, 1) + twy = ld_c(TWY, (s * Ay + n1y) * 2) + for n1z in cutlass.range_constexpr(Az): + v = rz[n1z] + if cutlass.const_expr(n1z > 0): + v = cmul(v, ld_c(TWZ, (s * Az + n1z) * 2)) + if cutlass.const_expr(n1y > 0): + v = cmul(v, twy) + if cutlass.const_expr((rnd + 1) * T > n3f): + if t < n3f: + st_c(up, (ob + (n1z * ny + n1y) * M) * 2, v) + else: + st_c(up, (ob + (n1z * ny + n1y) * M) * 2, v) + + # ---- phase 3a: y first-stage DFT across rows (in place) ---- + if cutlass.const_expr((Ay > 1) and not fuse3): + for rnd in cutlass.range_constexpr(_cdiv(n3a, T)): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > n3a): + tc = cutlass.min(t, n3a - 1) + m = tc % M + sz = tc // M + s = sz // Az + q = m // X1 + pm = X1 * q + rot(m - X1 * q + q) + rb = sz * Ay + vals = [None] * Ay + for k1y in cutlass.range_constexpr(Ay): + vals[k1y] = ld_c(P, ((rb + k1y) * Ms + pm) * 2) + res = dft(vals, 1) + for n1y in cutlass.range_constexpr(1, Ay): + res[n1y] = cmul(res[n1y], ld_c(TWY, (s * Ay + n1y) * 2)) + if cutlass.const_expr((rnd + 1) * T > n3a): + if t < n3a: + for n1y in cutlass.range_constexpr(Ay): + st_c(P, ((rb + n1y) * Ms + pm) * 2, res[n1y]) + else: + for n1y in cutlass.range_constexpr(Ay): + st_c(P, ((rb + n1y) * Ms + pm) * 2, res[n1y]) + cute.arch.sync_threads() + + # ---- phase 3b: z first-stage DFT across rows -> global ---- + for rnd in cutlass.range_constexpr(0 if fuse3 else _cdiv(n3b, T)): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > n3b): + tc = cutlass.min(t, n3b - 1) + m = tc % M + sy = tc // M + s = sy // Ay + n1y = sy - s * Ay + q = m // X1 + pm = X1 * q + rot(m - X1 * q + q) + k2z = cutlass.select_(s == 0, k2z0, k2z1) + k2y = cutlass.select_(s == 0, k2y0, k2y1) + vals = [None] * Az + for k1z in cutlass.range_constexpr(Az): + vals[k1z] = ld_c(P, ((s * R + k1z * Ay + n1y) * Ms + pm) * 2) + res = dft(vals, 1) + for n1z in cutlass.range_constexpr(1, Az): + res[n1z] = cmul(res[n1z], ld_c(TWZ, (s * Az + n1z) * 2)) + ob = ((Az * k2z) * ny + (n1y + Ay * k2y)) * M + m + if cutlass.const_expr((rnd + 1) * T > n3b): + if t < n3b: + for n1z in cutlass.range_constexpr(Az): + st_c(up, (ob + n1z * ny * M) * 2, res[n1z]) + else: + for n1z in cutlass.range_constexpr(Az): + st_c(up, (ob + n1z * ny * M) * 2, res[n1z]) + + return rowpassa_kernel + + +class StripPass: + """Pass 2: in-place Bz x By second-stage DFT on column strips of XC columns.""" + + def __init__(self, nz, ny, nx, Az, Ay, XC): + M = nx // 2 + self.nz, self.ny, self.nx, self.M = nz, ny, nx, M + self.Az, self.Ay = Az, Ay + self.Bz, self.By = nz // Az, ny // Ay + assert M % XC == 0 + self.XC = XC + self.NCB = M // XC + self.grid = Az * Ay * self.NCB + nt = max(self.Bz, self.By) * XC + self.T = min(1024, _cdiv(nt, 32) * 32) + self.pad = self._choose_pad() + self.RS = self.By * XC + self.pad # complex per k2z slab in smem + + def _choose_pad(self): + """Slab padding (complex) with the fewest shared-memory bank wavefronts.""" + Bz, By, XC, T = self.Bz, self.By, self.XC, self.T + best = None + for pad in range(17): + RS = By * XC + pad + insts = [] + nY = Bz * XC + for w0 in range(0, min(T, nY), 32): + for n2y in range(By): + a = [] + for ln in range(32): + t = w0 + ln + a.append(None if t >= nY else (t // XC) * RS + n2y * XC + t % XC) + insts.append(a) + nZ = By * XC + for w0 in range(0, min(T, nZ), 32): + for k2z in range(Bz): + a = [] + for ln in range(32): + t = w0 + ln + a.append(None if t >= nZ else k2z * RS + (t // XC) * XC + t % XC) + insts.append(a) + wf = _wavefronts(insts) + if best is None or wf < best[0]: + best = (wf, pad) + return best[1] + + def kernel(self): + ny, M = self.ny, self.M + Az, Ay, Bz, By, XC, NCB, T = self.Az, self.Ay, self.Bz, self.By, self.XC, self.NCB, self.T + RS = self.RS + SMF = Bz * RS * 2 + nY = Bz * XC + nZ = By * XC + RZ = Az * ny * M # complex stride between k2z rows + RY = Ay * M # complex stride between k2y rows + + @cute.kernel + def strippass_kernel(u: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + up = u.iterator + alloc = cutlass.utils.SmemAllocator() + S = alloc.allocate_array(Float32, SMF, byte_alignment=16) + cb = bid % NCB + rest = bid // NCB + n1y = rest % Ay + n1z = rest // Ay + base = (n1z * ny + n1y) * M + cb * XC + cute.arch.griddepcontrol_wait() + # ---- y stage: tasks (k2z, c) ---- + for rnd in cutlass.range_constexpr(_cdiv(nY, T)): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > nY): + tc = cutlass.min(t, nY - 1) + c = tc % XC + k2z = tc // XC + gb = base + k2z * RZ + c + vals = [None] * By + for k2y in cutlass.range_constexpr(By): + vals[k2y] = ld_c(up, (gb + k2y * RY) * 2) + res = dft(vals, 1) + sb = k2z * RS + c + if cutlass.const_expr((rnd + 1) * T > nY): + if t < nY: + for n2y in cutlass.range_constexpr(By): + st_c(S, (sb + n2y * XC) * 2, res[n2y]) + else: + for n2y in cutlass.range_constexpr(By): + st_c(S, (sb + n2y * XC) * 2, res[n2y]) + cute.arch.sync_threads() + # ---- z stage: tasks (n2y, c) ---- + for rnd in cutlass.range_constexpr(_cdiv(nZ, T)): + t = tid + rnd * T + tc = t + if cutlass.const_expr((rnd + 1) * T > nZ): + tc = cutlass.min(t, nZ - 1) + c = tc % XC + n2y = tc // XC + vals = [None] * Bz + for k2z in cutlass.range_constexpr(Bz): + vals[k2z] = ld_c(S, (k2z * RS + n2y * XC + c) * 2) + res = dft(vals, 1) + gb = base + n2y * RY + c + if cutlass.const_expr((rnd + 1) * T > nZ): + if t < nZ: + for n2z in cutlass.range_constexpr(Bz): + st_c(up, (gb + n2z * RZ) * 2, res[n2z]) + else: + for n2z in cutlass.range_constexpr(Bz): + st_c(up, (gb + n2z * RZ) * 2, res[n2z]) + + return strippass_kernel + + +class StripPassR: + """Pass 2, register-heavy variant for tiles larger than shared memory. + + Each thread loads a whole Bz-point z column (k2y, c) and transforms it in + registers; the y-DFTs then run in NH chunks of n2z through a shared buffer + of (Bz/NH) x By x XC complex values. + """ + + def __init__(self, nz, ny, nx, Az, Ay, XC, NH): + M = nx // 2 + self.nz, self.ny, self.nx, self.M = nz, ny, nx, M + self.Az, self.Ay = Az, Ay + self.Bz, self.By = nz // Az, ny // Ay + assert M % XC == 0 and self.Bz % NH == 0 + self.XC, self.NH = XC, NH + self.BH = self.Bz // NH + self.NCB = M // XC + self.grid = Az * Ay * self.NCB + self.T = max(self.By * XC, self.BH * XC) + self.T = _cdiv(self.T, 32) * 32 + self.pad = self._choose_pad() + self.RS = self.By * XC + self.pad # complex per n2z slab + + def _choose_pad(self): + """Slab padding (complex) with the fewest shared-memory bank wavefronts.""" + By, XC, BH, T = self.By, self.XC, self.BH, self.T + best = None + for pad in range(17): + RS = By * XC + pad + insts = [] + n1 = By * XC # writes: task (k2y, c) -> S[nl][k2y][c] + for w0 in range(0, min(T, n1), 32): + for nl in range(BH): + insts.append([None if w0 + ln >= n1 else nl * RS + (w0 + ln) for ln in range(32)]) + n2 = BH * XC # reads: task (nl, c) -> S[nl][k2y][c] + for w0 in range(0, min(T, n2), 32): + for k2y in range(By): + insts.append( + [ + None if w0 + ln >= n2 else ((w0 + ln) // XC) * RS + k2y * XC + (w0 + ln) % XC + for ln in range(32) + ] + ) + wf = _wavefronts(insts) + if best is None or wf < best[0]: + best = (wf, pad) + return best[1] + + def kernel(self): + ny, M = self.ny, self.M + Az, Ay, Bz, By, XC, NCB = self.Az, self.Ay, self.Bz, self.By, self.XC, self.NCB + NH, BH, RS = self.NH, self.BH, self.RS + SMF = BH * RS * 2 + n1t = By * XC + n2t = BH * XC + RZ = Az * ny * M + RY = Ay * M + + @cute.kernel + def strippassr_kernel(u: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + up = u.iterator + alloc = cutlass.utils.SmemAllocator() + S = alloc.allocate_array(Float32, SMF, byte_alignment=16) + cb = bid % NCB + rest = bid // NCB + n1y = rest % Ay + n1z = rest // Ay + base = (n1z * ny + n1y) * M + cb * XC + cute.arch.griddepcontrol_wait() + # ---- z stage: task (k2y, c), whole Bz-point column in registers ---- + t1 = cutlass.min(tid, n1t - 1) + c1 = t1 % XC + k2y = t1 // XC + g1 = base + k2y * RY + c1 + vals = [None] * Bz + for k2z in cutlass.range_constexpr(Bz): + vals[k2z] = ld_c(up, (g1 + k2z * RZ) * 2) + zr = dft(vals, 1) + # ---- y stage in NH chunks of n2z through shared memory ---- + t2 = cutlass.min(tid, n2t - 1) + c2 = t2 % XC + nl = t2 // XC + for h in cutlass.range_constexpr(NH): + if cutlass.const_expr(h > 0): + cute.arch.sync_threads() + if tid < n1t: + for q in cutlass.range_constexpr(BH): + st_c(S, (q * RS + k2y * XC + c1) * 2, zr[h * BH + q]) + cute.arch.sync_threads() + vy = [None] * By + for kk in cutlass.range_constexpr(By): + vy[kk] = ld_c(S, (nl * RS + kk * XC + c2) * 2) + yr = dft(vy, 1) + gb = base + (h * BH + nl) * RZ + c2 + if tid < n2t: + for n2y in cutlass.range_constexpr(By): + st_c(up, (gb + n2y * RY) * 2, yr[n2y]) + + return strippassr_kernel + + +class TwoPassPlan: + """RowPassA -> StripPass / StripPassR, chained with programmatic dependent launch.""" + + def __init__(self, nz, ny, nx, device, Az, Ay, XC, row_threads, strip_nh, strip_minb): + nf = nx // 2 + 1 + rp = RowPassA(nz, ny, nx, Az, Ay, row_threads) + # pass-1 CTA b prefetches (bulk L2 prefetch) the rows of CTA b + 2*#SMs (one + # wave ahead at 2 CTAs/SM, two at 1 CTA/SM): the later waves then find their + # input in L2 instead of queueing behind the whole DRAM stream + # (192^3 -2.4%, 256^3 -2.5% measured; shapes with fewer CTAs: off) + nsm = torch.cuda.get_device_properties(device).multi_processor_count + per_sm = max(1, min((100 * 1024) // (rp.smem_bytes + 1024), 1536 // row_threads)) + row_pf = 2 * nsm + if rp.grid <= nsm * per_sm or row_pf >= rp.grid: + row_pf = 0 # single wave (all CTAs resident): nothing to prefetch + if strip_nh: + sp = StripPassR(nz, ny, nx, Az, Ay, XC, strip_nh) + else: + sp = StripPass(nz, ny, nx, Az, Ay, XC) + rk, sk = rp.kernel(row_pf), sp.kernel() + + @cute.jit + def launch(x: cute.Tensor, u: cute.Tensor, tw: cute.Tensor, gt: cute.Tensor, stream): + rk(x, u, tw, gt).launch(grid=(rp.grid, 1, 1), block=(rp.T, 1, 1), stream=stream) + if cutlass.const_expr(strip_minb > 0): + sk(u).launch( + grid=(sp.grid, 1, 1), block=(sp.T, 1, 1), stream=stream, use_pdl=True, min_blocks_per_mp=strip_minb + ) + else: + sk(u).launch(grid=(sp.grid, 1, 1), block=(sp.T, 1, 1), stream=stream, use_pdl=True) + + self.fn = cute.compile( + launch, + _fake(nz * ny * nf * 2), + _fake(nz * ny * nx), + _fake(rp.table.numel()), + _fake(rp.ptab.numel(), Int32), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + self.tw = rp.table.to(device) + self.gt = rp.ptab.to(device) + + def __call__(self, x_flat, out_flat): + self.fn(x_flat, out_flat, self.tw, self.gt) + + +def _divisors(n): + return [d for d in range(1, n + 1) if n % d == 0] + + +def two_pass_config(nz, ny, nx, smem_limit=99 * 1024): + """TwoPassPlan parameters for a shape, or None if the two-pass path does not apply. + + Pass 2 keeps whole By-point (<= 32) and Bz-point DFTs in registers: Bz <= 32 + with the shared-memory StripPass, Bz <= 64 with the register-heavy + StripPassR. Pass 1 must hold the 2*Az*Ay rows of a group pair in shared + memory; group pairs small enough for two pass-1 CTAs per SM are preferred + (pass 1 is latency bound with one CTA per SM). Group sizes and thread + counts follow RTX 5090 measurements. + """ + if nx % 2: + return None + M = nx // 2 + if M < 2 or factor2(M) is None: + return None + rows = nz * ny + if (rows * (M + 1)) % 2: + return None # 16-byte row loads need an even input length + if rows * (M + 1) < 40000: + return None # tiny grids: the fused Z + YX path is faster + X1, X2 = factor2(M) + n_tws = X1 * X2 + M + + def row_smem(Az, Ay): + return (2 * Az * Ay * (M + 17) + n_tws + 2 * (Az + Ay)) * 8 + 64 + + def search(bz_max, row_cap): + r_target = max(16, rows / (bz_max * 32)) + best = None + for Az in _divisors(nz): + if nz // Az > bz_max: + continue + for Ay in _divisors(ny): + if ny // Ay > 32 or Az * Ay > 64 or row_smem(Az, Ay) > row_cap: + continue + R = Az * Ay + if bz_max > 32: # register strip pass: small By measured best + score = (-Ay, abs(math.log(R / r_target))) + else: + score = (abs(math.log(R / r_target)) + 0.1 * abs(math.log(Az / Ay)),) + if best is None or score < best[0]: + best = (score, Az, Ay) + return best + + two_per_sm = smem_limit // 2 - 1024 + found = None + for bz_max, cap in ((32, two_per_sm), (64, two_per_sm), (32, smem_limit), (64, smem_limit)): + best = search(bz_max, cap) + if best is not None: + found = (best[1], best[2]) + break + if found is None: + return None + Az, Ay = found + Bz, By = nz // Az, ny // Ay + xcs = [c for c in _divisors(M) if c <= 16] + XC = min(xcs, key=lambda c: (abs(math.log2(c / 8.0)), -c)) + if Bz <= 32: + if Bz * (By * XC + 16) * 8 + 64 > smem_limit: + cands = [c for c in xcs if Bz * (By * c + 16) * 8 + 64 <= smem_limit] + if not cands: + return None + XC = max(cands) + nh = 0 + else: + XC = max([c for c in xcs if By * c <= 256] or [0]) + if XC == 0: + return None + fits = [h for h in _divisors(Bz) if (Bz // h) * (By * XC + 16) * 8 + 64 <= smem_limit] + if not fits: + return None + small = [h for h in fits if (Bz // h) * (By * XC + 16) * 8 <= 40 * 1024] + nh = min(small) if small else min(fits) + R = Az * Ay + # 128 threads for small groups, else 256 (192^3: 256 beats 192 by ~1-3%) + row_threads = 128 if R * X2 <= 128 else 256 + # register strip pass: ask for two CTAs per SM when that still leaves + # >= 168 registers per thread (192^3: 175 -> 168 regs, -2% measured) + minb = 0 + if nh and _cdiv(By * XC, 32) * 32 <= 192 and (Bz // nh) * (By * XC + 16) * 8 + 1024 <= smem_limit // 2: + minb = 2 + return {"Az": Az, "Ay": Ay, "XC": XC, "row_threads": row_threads, "strip_nh": nh, "strip_minb": minb} diff --git a/build/torch-cuda/_c2r_f64/__init__.py b/build/torch-cuda/_c2r_f64/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cc2056f38e267f2ed7387e1cc1846e2f771bf109 --- /dev/null +++ b/build/torch-cuda/_c2r_f64/__init__.py @@ -0,0 +1,5 @@ +"""fp64 inverse complex-to-real 3D FFT (RTX 5090, CuTe DSL); see kernel.py for the contract.""" + +from .kernel import prepare, run + +__all__ = ["prepare", "run"] diff --git a/build/torch-cuda/_c2r_f64/fftgen.py b/build/torch-cuda/_c2r_f64/fftgen.py new file mode 100644 index 0000000000000000000000000000000000000000..a473b05f7c16d1f82dac8125f5eecaafcbde09f1 --- /dev/null +++ b/build/torch-cuda/_c2r_f64/fftgen.py @@ -0,0 +1,424 @@ +"""Trace-time FFT codelet generator. + +Values are complex numbers whose real and imaginary components each carry a compile-time +coefficient: component = coef * stored_value. Multiplying by a real constant, by +-1 or +by +-i only rewrites coefficients (no FP64 op); a complex twiddle costs 2 FMAs +(Linzer-Feig style, the scale goes into the coefficients); an addition costs one ADD or +FMA per component. Coefficients are only materialized (one DMUL, or a free sign flip) +when a value leaves a codelet. + +dft() emits a DFT of any length (inverse sign) through a backend: kernel.DslBk emits +CuTe-DSL float64 operations inside the kernels, CountBackend counts them for the planner. +""" + +import math + + +class Real: + __slots__ = ("c", "v") + + def __init__(self, v, c=1.0): + self.v = v # stored backend value + self.c = c # compile-time coefficient (python float) + + +class Cx: + __slots__ = ("im", "re", "tw") + + def __init__(self, re, im, tw=None): + self.re = re + self.im = im + self.tw = tw # pending runtime twiddle (wr, wi) backend values, or None + + +def _is_unit(c): + return c == 1.0 or c == -1.0 + + +def radd(bk, a, b, sb=1.0): + """a + sb*b for Real a, b (sb a python float).""" + ca, cb = a.c, b.c * sb + if cb == 0.0: + return a + if ca == 0.0: + return Real(b.v, cb) + # Keep the coefficient of the operand with unit coefficient if possible, + # otherwise the larger one (keeps |ratio| <= 1). + if _is_unit(ca) and not _is_unit(cb): + keep_a = True + elif _is_unit(cb) and not _is_unit(ca): + keep_a = False + else: + keep_a = abs(ca) >= abs(cb) + if keep_a: + r = cb / ca + if r == 1.0: + return Real(bk.add(a.v, b.v), ca) + if r == -1.0: + return Real(bk.sub(a.v, b.v), ca) + return Real(bk.fma(r, b.v, a.v), ca) + r = ca / cb + if r == 1.0: + return Real(bk.add(b.v, a.v), cb) + if r == -1.0: + return Real(bk.sub(b.v, a.v), cb) + return Real(bk.fma(r, a.v, b.v), cb) + + +def rpair(bk, a, b): + """(a + b, a - b) for Real a, b (shares work when routed to one addsub).""" + ca, cb = a.c, b.c + if cb == 0.0: + return a, a + if ca == 0.0: + return Real(b.v, cb), Real(b.v, -cb) + if _is_unit(ca) and not _is_unit(cb): + keep_a = True + elif _is_unit(cb) and not _is_unit(ca): + keep_a = False + else: + keep_a = abs(ca) >= abs(cb) + if keep_a: + r = cb / ca + if r == 1.0 or r == -1.0: + s_, d_ = bk.addsub(a.v, b.v) + if r == 1.0: + return Real(s_, ca), Real(d_, ca) + return Real(d_, ca), Real(s_, ca) + return Real(bk.fma(r, b.v, a.v), ca), Real(bk.fma(-r, b.v, a.v), ca) + r = ca / cb + if r == 1.0 or r == -1.0: + s_, d_ = bk.addsub(a.v, b.v) + if r == 1.0: + # a + b = cb (va + vb), a - b = cb (va - vb) + return Real(s_, cb), Real(d_, cb) + # a + b = cb (vb - va) = -cb (va - vb), a - b = -cb (va + vb) + return Real(d_, -cb), Real(s_, -cb) + # a + b = cb (r va + vb), a - b = cb (r va - vb) + return Real(bk.fma(r, a.v, b.v), cb), Real(bk.fma(r, a.v, bk.negv(b.v)), cb) + + +def cpair(bk, a, b): + """(a + b, a - b) for complex a, b.""" + a, b = resolve(bk, a), resolve(bk, b) + sr, dr = rpair(bk, a.re, b.re) + si, di = rpair(bk, a.im, b.im) + return Cx(sr, si), Cx(dr, di) + + +def resolve(bk, a): + """Apply a pending runtime twiddle (4 ops).""" + if a.tw is None: + return a + return cmul_rt(bk, Cx(a.re, a.im), a.tw[0], a.tw[1]) + + +def cadd(bk, a, b): + a, b = resolve(bk, a), resolve(bk, b) + return Cx(radd(bk, a.re, b.re), radd(bk, a.im, b.im)) + + +def _twfma(bk, a, b): + """a + w*b with b's pending twiddle w; a, b materialized. 4 FMAs.""" + wr, wi = b.tw + br = materialize_real(bk, b.re) + bi = materialize_real(bk, b.im) + ar = materialize_real(bk, a.re) + ai = materialize_real(bk, a.im) + sr = bk.fma_rt(wr, br, bk.fma_rt_neg(wi, bi, ar)) + si = bk.fma_rt(wi, br, bk.fma_rt(wr, bi, ai)) + return (ar, ai), (sr, si) + + +def butterfly(bk, a, b): + """(a + b, a - b), fusing a pending runtime twiddle on b (or a) into FMAs.""" + if a.tw is None and b.tw is None: + return cpair(bk, a, b) + tm = bk.two_minus + if b.tw is None: + # (w a + b, w a - b) = (s, s - 2b) with s = b + w a + (br, bi), (sr, si) = _twfma(bk, b, a) + if tm is not None: + d = Cx(Real(bk.neg(tm(br, sr))), Real(bk.neg(tm(bi, si)))) + else: + d = Cx(Real(bk.fma(-2.0, br, sr)), Real(bk.fma(-2.0, bi, si))) + return Cx(Real(sr), Real(si)), d + a = resolve(bk, a) + (ar, ai), (sr, si) = _twfma(bk, a, b) + if tm is not None: + d = Cx(Real(tm(ar, sr)), Real(tm(ai, si))) + else: + d = Cx(Real(bk.fma(2.0, ar, bk.negv(sr))), Real(bk.fma(2.0, ai, bk.negv(si)))) + return Cx(Real(sr), Real(si)), d + + +def cscale(a, s): + assert a.tw is None + return Cx(Real(a.re.v, a.re.c * s), Real(a.im.v, a.im.c * s)) + + +def cmul_i(a): + """i * a.""" + assert a.tw is None + return Cx(Real(a.im.v, -a.im.c), Real(a.re.v, a.re.c)) + + +def cmul_const(bk, a, ang): + """a * exp(i*ang) with ang a python float (compile-time twiddle).""" + a = resolve(bk, a) + # Reduce to an exact quarter turn when possible. + q = ang / (math.pi / 2) + qr = round(q) + if abs(q - qr) < 1e-12: + k = qr % 4 + for _ in range(k): + a = cmul_i(a) + return a + c, s = math.cos(ang), math.sin(ang) + if abs(abs(c) - abs(s)) < 1e-12: + # Odd multiple of pi/4: one exact add/sub pair (cos and sin of pi/4 differ in the + # last bit as doubles, which would otherwise turn them into FP64 FMAs). + h = math.copysign(math.sqrt(0.5), c) + sm, df = rpair(bk, Real(a.re.v, a.re.c * h), Real(a.im.v, a.im.c * h)) + if s * c > 0: + return Cx(df, sm) # re = c(ar - ai), im = c(ar + ai) + return Cx(sm, Real(df.v, -df.c)) # re = c(ar + ai), im = -c(ar - ai) + # re = c*a.re - s*a.im ; im = s*a.re + c*a.im + re = radd(bk, Real(a.re.v, a.re.c * c), Real(a.im.v, -a.im.c * s)) + im = radd(bk, Real(a.re.v, a.re.c * s), Real(a.im.v, a.im.c * c)) + return Cx(re, im) + + +def cmul_rt(bk, a, wr, wi): + """a * (wr + i wi) with runtime twiddle (backend values). 4 ops.""" + ar = materialize_real(bk, a.re) + ai = materialize_real(bk, a.im) + re = bk.fma_rt(wr, ar, bk.mul_rt_neg(wi, ai)) + im = bk.fma_rt(wi, ar, bk.mul_rt(wr, ai)) + return Cx(Real(re), Real(im)) + + +def materialize_real(bk, r): + if r.c == 1.0: + return r.v + if r.c == -1.0: + return bk.neg(r.v) + return bk.mul(r.c, r.v) + + +def materialize(bk, a): + a = resolve(bk, a) + return Cx(Real(materialize_real(bk, a.re)), Real(materialize_real(bk, a.im))) + + +# --------------------------------------------------------------------------- +# DFT codelets, inverse sign: y_k = sum_j x_j exp(+2 pi i j k / R) +# --------------------------------------------------------------------------- + + +def factorize(n): + f = [] + d = 2 + while d * d <= n: + while n % d == 0: + f.append(d) + n //= d + d += 1 + if n > 1: + f.append(n) + return f + + +def dft(bk, xs): + R = len(xs) + if R == 1: + return list(xs) + if R == 2: + return list(butterfly(bk, xs[0], xs[1])) + if R == 4: + a, b = butterfly(bk, xs[0], xs[2]) + c, d = butterfly(bk, xs[1], xs[3]) + id_ = cmul_i(d) + y0, y2 = cpair(bk, a, c) + y1, y3 = cpair(bk, b, id_) + return [y0, y1, y2, y3] + if R & (R - 1) == 0: + return _split_radix(bk, xs) + if len(factorize(R)) == 1: + return _dft_prime(bk, xs) + return _mixed(bk, xs) + + +def _split_radix(bk, xs): + R = len(xs) + U = dft(bk, xs[0::2]) + Z = dft(bk, xs[1::4]) + Zp = dft(bk, xs[3::4]) + y = [None] * R + q = R // 4 + for k in range(q): + a1 = cmul_const(bk, Z[k], 2 * math.pi * k / R) + a3 = cmul_const(bk, Zp[k], 2 * math.pi * 3 * k / R) + s, d = cpair(bk, a1, a3) + idd = cmul_i(d) + y[k], y[k + 2 * q] = cpair(bk, U[k], s) + y[k + q], y[k + 3 * q] = cpair(bk, U[k + q], idd) + return y + + +def _scaled(bk, a, s): + """s * a materialized with unit coefficients (one multiply per component).""" + a = resolve(bk, a) + return Cx(Real(bk.mul(s * a.re.c, a.re.v)), Real(bk.mul(s * a.im.c, a.im.v))) + + +def _dft_prime(bk, xs): + p = len(xs) + if p == 3: + s, d = butterfly(bk, xs[1], xs[2]) + x0 = resolve(bk, xs[0]) + y0 = cadd(bk, x0, s) + m = cadd(bk, x0, cscale(s, -0.5)) + t = cmul_i(cscale(d, math.sqrt(3.0) / 2)) + y1, y2 = cpair(bk, m, t) + return [y0, y1, y2] + # p >= 5: the sine sums are materialized so that the final y_m, y_{p-m} combinations + # are plain add/sub pairs (one FP64 multiply traded for an offloadable pair). + h = (p - 1) // 2 + S, D = [], [] + for k in range(1, h + 1): + sk, dk = butterfly(bk, xs[k], xs[p - k]) + S.append(sk) + D.append(dk) + x0 = resolve(bk, xs[0]) + y = [None] * p + acc = x0 + for k in range(h): + acc = cadd(bk, acc, S[k]) + y[0] = acc + for m in range(1, h + 1): + a = x0 + for k in range(1, h + 1): + a = cadd(bk, a, cscale(S[k - 1], math.cos(2 * math.pi * k * m / p))) + b = None + for k in range(1, h + 1): + sk = math.sin(2 * math.pi * k * m / p) + if b is None: + b = _scaled(bk, D[k - 1], sk) + continue + b = cadd(bk, b, cscale(D[k - 1], sk)) + ib = cmul_i(b) + y[m], y[p - m] = cpair(bk, a, ib) + return y + + +def _best_split(R): + f = factorize(R) + # Prefer coprime splits (prime-factor algorithm, no twiddles). + p2 = 1 + while R % (p2 * 2) == 0: + p2 *= 2 + if 1 < p2 < R: + return p2, R // p2 + # otherwise split off the smallest prime + return f[0], R // f[0] + + +def _mixed(bk, xs): + R = len(xs) + R1, R2 = _best_split(R) + if math.gcd(R1, R2) == 1: + # Good-Thomas PFA: n = (R2*n1 + R1*n2) mod R ; k = k1*R2*inv(R2,R1) + k2*R1*inv(R1,R2) + sub = [] + for n2 in range(R2): + sub.append(dft(bk, [xs[(R2 * n1 + R1 * n2) % R] for n1 in range(R1)])) + cols = [] + for k1 in range(R1): + cols.append(dft(bk, [sub[n2][k1] for n2 in range(R2)])) + e1 = R2 * pow(R2, -1, R1) + e2 = R1 * pow(R1, -1, R2) + y = [None] * R + for k1 in range(R1): + for k2 in range(R2): + y[(k1 * e1 + k2 * e2) % R] = cols[k1][k2] + return y + # Cooley-Tukey DIT: x[R2*n1 + n2] ; X[k1 + R1*k2] + sub = [] + for n2 in range(R2): + sub.append(dft(bk, [xs[R2 * n1 + n2] for n1 in range(R1)])) + y = [None] * R + for k1 in range(R1): + col = [cmul_const(bk, sub[n2][k1], 2 * math.pi * n2 * k1 / R) for n2 in range(R2)] + out = dft(bk, col) + for k2 in range(R2): + y[k1 + R1 * k2] = out[k2] + return y + + +# --------------------------------------------------------------------------- +# Op counting (planner cost model and soft/native balance) +# --------------------------------------------------------------------------- + + +class CountBackend: + """Counts FP64 ops (sign flips are free) and add/sub pairs; values are dummies.""" + + two_minus = None + + def __init__(self): + self.ops = 0 + self.pairs = 0 + + def add(self, a, b): + self.ops += 1 + return 0.0 + + sub = add + + def fma(self, c, x, y): + self.ops += 1 + return 0.0 + + def mul(self, c, x): + self.ops += 1 + return 0.0 + + def neg(self, x): + return -x + + negv = neg + + def addsub(self, a, b): + self.ops += 2 + self.pairs += 1 + return 0.0, 0.0 + + fma_rt = fma + fma_rt_neg = fma + mul_rt = mul + mul_rt_neg = mul + + +def _count(R, twiddled): + bk = CountBackend() + xs = [Cx(Real(0.0), Real(0.0), (0.0, 0.0) if (twiddled and j > 0) else None) for j in range(R)] + for y in dft(bk, xs): + materialize(bk, y) + return bk + + +def codelet_cost(R, twiddled=False): + """FP64 ops of a length-R codelet including materialization; twiddles on inputs 1.. .""" + return _count(R, twiddled).ops + + +_MIX = {} + + +def codelet_mix(R, twiddled): + """(add/sub pairs, other FP64 ops) of a length-R codelet; twiddles on inputs 1.. .""" + key = (R, twiddled) + if key not in _MIX: + bk = _count(R, twiddled) + _MIX[key] = (bk.pairs, bk.ops - 2 * bk.pairs) + return _MIX[key] diff --git a/build/torch-cuda/_c2r_f64/kernel.py b/build/torch-cuda/_c2r_f64/kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..c4310006e78c3e9ea79be2f1bf7c4d4910138d55 --- /dev/null +++ b/build/torch-cuda/_c2r_f64/kernel.py @@ -0,0 +1,2903 @@ +"""fp64 3D complex-to-real inverse FFT for the RTX 5090 (FP64 at 1/64 rate), in CuTe DSL. + +run(input): input is a contiguous float64 [nz, ny, nf, 2] packed Hermitian half-spectrum +(interleaved re/im); returns a new contiguous float64 [nz, ny, nx], nx = 2*(nf-1), equal to +torch.fft.irfftn(torch.view_as_complex(input), s=(nz, ny, nx), norm="forward") (unnormalized, +cuFFT Z2D semantics). The input is not modified. prepare(input) binds a contiguous input +and a fresh output once and returns launch(), which recomputes from the input buffer's +current contents into that output and returns it. The first call for a (shape, device) +builds the plan, the twiddle tables and the JIT-compiled kernels, cached for later calls. + +How it works. With M = nx/2, the half-spectrum columns k and M-k (index-reversed in kz, ky) +are paired into one complex [nz, ny, M] array whose plain inverse 3D DFT is the real output +viewed as complex (u[2n] + i*u[2n+1]); the three 1D passes then run in any order: + * x-first (the tuned benchmark shapes except 32^3): pairing + x FFT on row pairs (p, -p), + then the y and z FFTs in place in the output, or for some ny == nz shapes two 2D + vector-radix passes; + * z-first (any other shape whose nz, ny and M each split into at most 3 radices from + 2..16): pairing fused into the z FFT, then y and x (y+x fused per z plane for 32^3); + * direct (all other shapes): pairing, then three O(n^2) DFT passes. +Each FFT pass keeps a tile of sequences in shared memory and runs mixed-radix DIT stages +(optionally Good-Thomas) built from fftgen codelets. Plans (radices, tile, threads, pass +order, launch details) are tuned per benchmark shape; other shapes use a heuristic planner. + +All arithmetic is IEEE float64. Part of the additions run through softfp, a software +add.rn.f64 on the integer pipes, which execute concurrently with the slow FP64 unit; it is +exact for finite normal or zero operands and results, so subnormal, Inf or NaN data (or +results that overflow) give wrong results. +""" + +# Nested ifs are deliberate: inside the kernels an outer compile-time if (cutlass.const_expr) +# and every runtime condition stay separate branches. +# ruff: noqa: SIM102 + +import math + +import cuda.bindings.driver as cuda +import cutlass +import torch +from cutlass import cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import Float64, Int32, T, dsl_user_op + +from . import fftgen as G +from . import softfp, twiddles + +_DYN = -2147483648 +NUM_SMS = 170 + + +# --------------------------------------------------------------------------- +# Per-shape tuning (benchmark shapes), keyed (nz, ny, M) with M = nx / 2 +# --------------------------------------------------------------------------- + +# (radices, S, NT[, soft, pfa[, twiddles from L1 instead of shared memory]]) per pass, keyed +# (pass, nz, ny, M): passes "z", "y", "x" of the z-first pipeline and "fx", "fy", "fz" of the +# x-first one ("fy"/"fz" fall back to "y"/"z"). Any other pass uses the heuristic planner. +_OVERRIDES = { + ("z", 32, 32, 16): ([4, 2, 4], 2, 32, (3, 4, 1, 2)), + ("y", 32, 32, 16): ([4, 2, 4], 2, 32), + ("x", 32, 32, 16): ([4, 2, 2], 2, 32), + ("z", 64, 64, 32): ([4, 4, 4], 16, 256), + ("y", 64, 64, 32): ([4, 4, 4], 16, 256), + ("x", 64, 64, 32): ([4, 2, 4], 4, 32), + ("z", 96, 96, 48): ([8, 4, 3], 16, 192, None, True), + ("y", 96, 96, 48): ([8, 4, 3], 16, 192, None, True), + ("x", 96, 96, 48): ([4, 6, 2], 8, 32), + ("z", 128, 128, 64): ([8, 8, 2], 8, 128), + ("y", 128, 128, 64): ([8, 2, 8], 8, 128), + ("x", 128, 128, 64): ([4, 4, 4], 2, 32), + ("z", 192, 192, 96): ([12, 4, 4], 8, 128), + ("y", 192, 192, 96): ([8, 8, 3], 8, 192, None, True), + ("x", 192, 192, 96): ([12, 2, 4], 8, 64), + ("z", 256, 256, 128): ([8, 8, 4], 8, 256), + ("y", 256, 256, 128): ([8, 4, 8], 8, 256), + ("x", 256, 256, 128): ([8, 16], 8, 64), + ("z", 80, 80, 56): ([4, 4, 5], 16, 128, None, True), + ("y", 80, 80, 56): ([4, 4, 5], 16, 128, None, True), + ("x", 80, 80, 56): ([2, 2, 14], 8, 32), + ("z", 96, 96, 72): ([4, 8, 3], 16, 192, None, True), + ("y", 96, 96, 72): ([8, 4, 3], 16, 192, None, True), + ("x", 96, 96, 72): ([6, 6, 2], 8, 32), + ("fx", 64, 64, 32): ([4, 2, 4], 4, 32, (3, 4, 1, 2)), + ("fy", 64, 64, 32): ([4, 4, 4], 16, 256), + ("fz", 64, 64, 32): ([4, 4, 4], 16, 256), + ("fx", 96, 96, 48): ([8, 6], 4, 32, None, False, True), + ("fy", 96, 96, 48): ([4, 8, 3], 16, 192, None, True), + ("fz", 96, 96, 48): ([8, 4, 3], 16, 192, None, True), + ("fx", 128, 128, 64): ([4, 4, 4], 2, 32, None, False, True), + ("fy", 128, 128, 64): ([8, 2, 8], 8, 128), + ("fz", 128, 128, 64): ([8, 8, 2], 8, 128), + ("fx", 192, 192, 96): ([4, 4, 6], 4, 64, None, False, True), + ("fy", 192, 192, 96): ([8, 8, 3], 8, 192, None, True, True), + ("fz", 192, 192, 96): ([8, 8, 3], 8, 192, None, True, True), + ("fx", 256, 256, 128): ([4, 8, 4], 4, 128, (1, 1, 0, 1), False, True), + ("fy", 256, 256, 128): ([8, 4, 8], 4, 128, None, False, True), + ("fz", 256, 256, 128): ([8, 4, 8], 4, 128, None, False, True), + ("fx", 80, 80, 56): ([8, 7], 4, 32, None, True), + ("fx", 96, 96, 72): ([6, 2, 6], 8, 96), + ("fy", 80, 80, 56): ([4, 4, 5], 16, 128, None, True), + ("fz", 80, 80, 56): ([4, 4, 5], 16, 128, None, True), +} + +# shapes where the x-first pass order is faster +_XFIRST_SHAPES = { + (64, 64, 32), + (96, 96, 48), + (128, 128, 64), + (192, 192, 96), + (256, 256, 128), + (80, 80, 56), + (96, 96, 72), +} + +# shapes where a fused y+x pass (one block per z plane) beats separate passes (z-first) +_FUSED_YX = {(32, 32, 16): {"ry": [4, 4, 2], "rx": [4, 4], "NT": 128, "soft": (3, 4, 1, 2)}} + +# x-first shapes whose y/z passes run as 2D vector-radix passes (see YZ2D), ny == nz +_YZ2D = { + (64, 64, 32): { + "A": 8, + "rc": [8], + "N1": 1, + "SB": 4, + "NTB": 32, + "SC": 4, + "NTC": 32, + "softB": (3, 4, 1, 2), + "softC": (3, 4, 1, 2), + }, + (96, 96, 48): {"A": 8, "rc": [4], "N1": 3, "SB": 4, "NTB": 32, "SC": 4, "NTC": 192, "UC": 1}, + (128, 128, 64): {"A": 8, "rc": [4, 4], "N1": 1, "SB": 4, "NTB": 32, "SC": 2, "NTC": 128}, + (80, 80, 56): {"A": 8, "rc": [2], "N1": 5, "SB": 4, "NTB": 32, "SC": 8, "NTC": 160, "UC": 1}, + (96, 96, 72): {"A": 8, "rc": [4], "N1": 3, "SB": 4, "NTB": 32, "SC": 4, "NTC": 192, "UC": 1}, +} + +# x-first shapes whose x pass runs software pipelined (kxp): CTAs per SM, raw tile buffers +_XPIPE_SHAPES = {(128, 128, 64): {"cps": 16, "NB": 2}} + +# x-first shapes whose later passes use programmatic dependent launch +_PDL_SHAPES = {(64, 64, 32), (128, 128, 64), (80, 80, 56)} + +# x-first shapes whose x pass streams the input (L1 no-allocate, L2 evict-first) +_XSTREAM_SHAPES = {(256, 256, 128), (96, 96, 48), (192, 192, 96)} + +# shapes whose first pass prefetches the twiddle tables of the later passes into L2 +_TWPF_SHAPES = {(64, 64, 32), (80, 80, 56), (32, 32, 16)} + +# x-first shapes whose pairing loop maps lanes over row slots first (see kxa, kxp) +_PAIRMAP_SHAPES = {(128, 128, 64): 1} + +# x-first shapes whose arrays exceed L2: the z pass stores with an L2 evict-first policy +# (its own finished output goes first, so z's misses do not evict y output not yet read) +_ZPOL_SHAPES = {(256, 256, 128): "evict_first"} + +# x-first shapes whose y pass visits the planes in reverse order of the x pass's completion +_YORD_SHAPES = {(256, 256, 128)} + +# share of the fused-twiddle 2x-y operations computed in software, per x-first shape +_SOFT2M_SHAPES = { + (192, 192, 96): 1.0, + (128, 128, 64): 1.0, + (96, 96, 72): 1.0, + (96, 96, 48): 1.0, + (256, 256, 128): 1.0, +} + +# integer-pipe cost of a software butterfly assumed by the soft/native balance, per shape +_INTCYC_SHAPES = {(192, 192, 96): 50.0, (256, 256, 128): 60.0} + +# software butterfly variant (softfp.sbfly3/4/5) per x-first shape (default 3) +_SOFTBF_SHAPES = {(96, 96, 48): 4, (192, 192, 96): 4, (256, 256, 128): 4, (80, 80, 56): 5} + + +# --------------------------------------------------------------------------- +# Low-level helpers +# --------------------------------------------------------------------------- + + +@dsl_user_op +def ld2(base, idx, *, loc=None, ip=None): + """Load complex (2 x f64, 16B aligned) at complex index idx.""" + v2 = ir.VectorType.get([2], T.f64()) + p = llvm.getelementptr( + base.type, base, [Int32(idx).ir_value(loc=loc, ip=ip)], [_DYN], v2, no_wrap_flags="None", loc=loc, ip=ip + ) + v = llvm.load(v2, p, alignment=16, loc=loc, ip=ip) + re = llvm.extractelement(v, Int32(0).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + im = llvm.extractelement(v, Int32(1).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + return Float64(re), Float64(im) + + +@dsl_user_op +def _ptr_add(base, idx, *, loc=None, ip=None): + """base + idx complex elements (generic llvm pointer arithmetic).""" + v2 = ir.VectorType.get([2], T.f64()) + return llvm.getelementptr( + base.type, base, [Int32(idx).ir_value(loc=loc, ip=ip)], [_DYN], v2, no_wrap_flags="None", loc=loc, ip=ip + ) + + +@dsl_user_op +def ld2_hint(base, idx, hint, *, loc=None, ip=None): + """ld2 from global memory with an L1 eviction-priority qualifier (e.g. 'L1::evict_last').""" + v2 = ir.VectorType.get([2], T.f64()) + p = llvm.getelementptr( + base.type, base, [Int32(idx).ir_value(loc=loc, ip=ip)], [_DYN], v2, no_wrap_flags="None", loc=loc, ip=ip + ) + pi = llvm.ptrtoint(T.i64(), p, loc=loc, ip=ip) + res = llvm.inline_asm( + llvm.StructType.get_literal([T.f64(), T.f64()]), + [pi], + "ld.global." + hint + ".v2.f64 {$0, $1}, [$2];", + "=d,=d,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float64(llvm.extractvalue(T.f64(), res, [0], loc=loc, ip=ip)), + Float64(llvm.extractvalue(T.f64(), res, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def ld2_hint_v(base, idx, hint, *, loc=None, ip=None): + """ld2_hint that the compiler may not hoist out of a loop (side-effecting asm).""" + v2 = ir.VectorType.get([2], T.f64()) + p = llvm.getelementptr( + base.type, base, [Int32(idx).ir_value(loc=loc, ip=ip)], [_DYN], v2, no_wrap_flags="None", loc=loc, ip=ip + ) + pi = llvm.ptrtoint(T.i64(), p, loc=loc, ip=ip) + res = llvm.inline_asm( + llvm.StructType.get_literal([T.f64(), T.f64()]), + [pi], + "ld.global." + hint + ".v2.f64 {$0, $1}, [$2];", + "=d,=d,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float64(llvm.extractvalue(T.f64(), res, [0], loc=loc, ip=ip)), + Float64(llvm.extractvalue(T.f64(), res, [1], loc=loc, ip=ip)), + ) + + +_TWV = [False] # set while tracing a kernel with a runtime tile loop (kxp) + + +@dsl_user_op +def tid_v(*, loc=None, ip=None): + """%tid.x through a side-effecting asm: address math derived from it stays inside a runtime + loop instead of being hoisted (and kept live in registers) by the compiler.""" + r = llvm.inline_asm( + T.i32(), + [], + "mov.u32 $0, %tid.x;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int32(r) + + +@dsl_user_op +def ld2_stream(base, idx, *, loc=None, ip=None): + """ld2 of read-once input: no L1 allocation, L2 evict-first policy.""" + v2 = ir.VectorType.get([2], T.f64()) + p = llvm.getelementptr( + base.type, base, [Int32(idx).ir_value(loc=loc, ip=ip)], [_DYN], v2, no_wrap_flags="None", loc=loc, ip=ip + ) + pi = llvm.ptrtoint(T.i64(), p, loc=loc, ip=ip) + res = llvm.inline_asm( + llvm.StructType.get_literal([T.f64(), T.f64()]), + [pi], + "{.reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol, 1.0; " + "ld.global.L1::no_allocate.L2::cache_hint.v2.f64 {$0, $1}, [$2], pol;}", + "=d,=d,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float64(llvm.extractvalue(T.f64(), res, [0], loc=loc, ip=ip)), + Float64(llvm.extractvalue(T.f64(), res, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def st2(base, idx, re, im, *, loc=None, ip=None): + """Store complex (2 x f64, 16B aligned) at complex index idx.""" + v2 = ir.VectorType.get([2], T.f64()) + p = llvm.getelementptr( + base.type, base, [Int32(idx).ir_value(loc=loc, ip=ip)], [_DYN], v2, no_wrap_flags="None", loc=loc, ip=ip + ) + v = llvm.mlir_undef(v2, loc=loc, ip=ip) + v = llvm.insertelement(v, Float64(re).ir_value(loc=loc, ip=ip), Int32(0).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + v = llvm.insertelement(v, Float64(im).ir_value(loc=loc, ip=ip), Int32(1).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + llvm.store(v, p, alignment=16, loc=loc, ip=ip) + + +@dsl_user_op +def mk_policy(kind, *, loc=None, ip=None): + """64-bit L2 cache policy (createpolicy.fractional, fraction 1.0) for st2p, e.g. kind 'evict_first'.""" + r = llvm.inline_asm( + T.i64(), + [], + "createpolicy.fractional.L2::" + kind + ".b64 $0, 1.0;", + "=l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Int64(r) + + +@dsl_user_op +def st2p(base, idx, re, im, pol, *, loc=None, ip=None): + """st2 to global memory with an L2 cache-policy hint (see mk_policy).""" + v2 = ir.VectorType.get([2], T.f64()) + p = llvm.getelementptr( + base.type, base, [Int32(idx).ir_value(loc=loc, ip=ip)], [_DYN], v2, no_wrap_flags="None", loc=loc, ip=ip + ) + pi = llvm.ptrtoint(T.i64(), p, loc=loc, ip=ip) + llvm.inline_asm( + None, + [ + pi, + Float64(re).ir_value(loc=loc, ip=ip), + Float64(im).ir_value(loc=loc, ip=ip), + cutlass.Int64(pol).ir_value(loc=loc, ip=ip), + ], + "st.global.L2::cache_hint.v2.f64 [$0], {$1, $2}, $3;", + "l,d,d,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def gst(op, idx, re, im, pol): + """Global complex store, with an L2 policy if pol is not None.""" + if pol is None: + st2(op, idx, re, im) + else: + st2p(op, idx, re, im, pol) + + +@dsl_user_op +def pf_l2(base, idx, *, loc=None, ip=None): + """prefetch.global.L2 of the line holding complex element idx.""" + v2 = ir.VectorType.get([2], T.f64()) + p = llvm.getelementptr( + base.type, base, [Int32(idx).ir_value(loc=loc, ip=ip)], [_DYN], v2, no_wrap_flags="None", loc=loc, ip=ip + ) + llvm.inline_asm( + None, + [p], + "prefetch.global.L2 [$0];", + "l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def fneg_int(x, *, loc=None, ip=None): + """Sign flip through an integer op (a plain fneg would cost a DADD).""" + r = llvm.inline_asm( + T.f64(), + [Float64(x).ir_value(loc=loc, ip=ip)], + "{.reg .b32 lo, hi; mov.b64 {lo, hi}, $1; xor.b32 hi, hi, 0x80000000; mov.b64 $0, {lo, hi};}", + "=d,d", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Float64(r) + + +@dsl_user_op +def st1(base, idx, v, *, loc=None, ip=None): + """Store one f64 at f64 index idx.""" + p = llvm.getelementptr( + base.type, base, [Int32(idx).ir_value(loc=loc, ip=ip)], [_DYN], T.f64(), no_wrap_flags="None", loc=loc, ip=ip + ) + llvm.store(Float64(v).ir_value(loc=loc, ip=ip), p, alignment=8, loc=loc, ip=ip) + + +@dsl_user_op +def fxor_sign(x, m, *, loc=None, ip=None): + """x with its sign bit xor'ed with m (0 or 0x80000000), on the integer pipe.""" + r = llvm.inline_asm( + T.f64(), + [Float64(x).ir_value(loc=loc, ip=ip), Int32(m).ir_value(loc=loc, ip=ip)], + "{.reg .b32 lo, hi; mov.b64 {lo, hi}, $1; xor.b32 hi, hi, $2; mov.b64 $0, {lo, hi};}", + "=d,d,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Float64(r) + + +# Bulk asynchronous copies (global -> shared) with mbarrier completion, for the +# software-pipelined x pass. Shared addresses are 32-bit shared-window offsets. + + +def _asm(ret, args, txt, cons, loc=None, ip=None): + return llvm.inline_asm( + ret, + args, + txt, + cons, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def saddr(ptr, byte_off, *, loc=None, ip=None): + """32-bit shared-window address of (shared llvm ptr + byte_off).""" + p = llvm.getelementptr( + ptr.type, ptr, [Int32(byte_off).ir_value(loc=loc, ip=ip)], [_DYN], T.i8(), no_wrap_flags="None", loc=loc, ip=ip + ) + return Int32(llvm.ptrtoint(T.i32(), p, loc=loc, ip=ip)) + + +@dsl_user_op +def gaddr(ptr, idx, *, loc=None, ip=None): + """64-bit address of complex element idx of a global llvm ptr.""" + v2 = ir.VectorType.get([2], T.f64()) + p = llvm.getelementptr( + ptr.type, ptr, [Int32(idx).ir_value(loc=loc, ip=ip)], [_DYN], v2, no_wrap_flags="None", loc=loc, ip=ip + ) + return cutlass.Int64(llvm.ptrtoint(T.i64(), p, loc=loc, ip=ip)) + + +@dsl_user_op +def mbar_init(a, cnt, *, loc=None, ip=None): + _asm(None, [Int32(a).ir_value(loc=loc, ip=ip)], f"mbarrier.init.shared::cta.b64 [$0], {cnt};", "r", loc=loc, ip=ip) + + +@dsl_user_op +def mbar_fence_init(*, loc=None, ip=None): + _asm(None, [], "fence.mbarrier_init.release.cluster;", "", loc=loc, ip=ip) + + +@dsl_user_op +def fence_async_shared(*, loc=None, ip=None): + """Order this thread's generic-proxy shared accesses before later async-proxy ones.""" + _asm(None, [], "fence.proxy.async.shared::cta;", "", loc=loc, ip=ip) + + +@dsl_user_op +def mbar_expect_tx(a, nbytes, *, loc=None, ip=None): + _asm( + None, + [Int32(a).ir_value(loc=loc, ip=ip), Int32(nbytes).ir_value(loc=loc, ip=ip)], + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [$0], $1;", + "r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def bulk_g2s(dst, src, nbytes, mbar, *, loc=None, ip=None): + """cp.async.bulk global -> shared, completing nbytes on the mbarrier; read-once data (L2 evict_first).""" + _asm( + None, + [ + Int32(dst).ir_value(loc=loc, ip=ip), + cutlass.Int64(src).ir_value(loc=loc, ip=ip), + Int32(nbytes).ir_value(loc=loc, ip=ip), + Int32(mbar).ir_value(loc=loc, ip=ip), + ], + "{.reg .b64 pol; createpolicy.fractional.L2::evict_first.b64 pol, 1.0; " + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint [$0], [$1], $2, [$3], pol;}", + "r,l,r,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def mbar_wait(a, parity, *, loc=None, ip=None): + """Spin until the mbarrier phase of the given parity has completed.""" + _asm( + None, + [Int32(a).ir_value(loc=loc, ip=ip), Int32(parity).ir_value(loc=loc, ip=ip)], + "{.reg .pred p; WAIT_${:uid}: mbarrier.try_wait.parity.shared::cta.b64 p, [$0], $1; @!p bra WAIT_${:uid};}", + "r,r", + loc=loc, + ip=ip, + ) + + +# --------------------------------------------------------------------------- +# Codelet backend: FP64 ops, part of the additions routed to softfp +# --------------------------------------------------------------------------- + +INT_PAIR_CYC = 80.0 # integer-pipe cycles of one software butterfly assumed by balance() +SOFT_DEFAULT = (1, 1, 1, 2) # all balanced pairs and every second single addition in software + + +class DslBk: + """Emits fp64 ops. Some plain additions run on the integer-pipe float64 + adder (softfp), concurrently with the 1/64-rate FP64 unit. Pairs a+-b are + routed per codelet so that integer and FP64 work balance (see balance()).""" + + def __init__(self): + self.set_soft((0, 1, 0, 1)) + + def set_soft(self, frac): + """frac = (pairs num, den, singles num, den[, share of fused 2x-y in software + [, integer cycles per software butterfly[, softfp butterfly variant]]]).""" + self.pn, self.pd, self.sn, self.sd = frac[:4] + self.ks = 0 + self.frac = self.pn / self.pd + self.acc = 0.0 + f2 = frac[4] if len(frac) > 4 else 0.0 + self.frac2 = f2 if self.pn else 0.0 + self.intcyc = frac[5] if len(frac) > 5 else INT_PAIR_CYC + self.bf = int(frac[6]) if len(frac) > 6 else 3 + self.acc2 = 0.0 + self.two_minus = self._two_minus if self.frac2 > 0 else None + + def balance(self, pairs, other): + """Route ~p of the codelet's pairs to software: INT_CYC*p = 64*(2(A-p)+F).""" + if self.pn == 0 or pairs == 0: + self.frac = 0.0 + else: + p = 64.0 * (2 * pairs + other) / (self.intcyc + 128.0) + self.frac = min(1.0, p / pairs) * (self.pn / self.pd) + self.acc = 0.0 + + def _soft_single(self): + self.ks += 1 + return (self.ks % self.sd) < self.sn + + def _soft_pair(self): + self.acc += self.frac + if self.acc >= 0.5: + self.acc -= 1.0 + return True + return False + + def add(self, a, b): + if self._soft_single(): + return softfp.sadd(a, b) + return a + b + + def sub(self, a, b): + if self._soft_single(): + return softfp.ssub(a, b) + return a - b + + def addsub(self, a, b): + if self._soft_pair(): + if self.bf == 4: + return softfp.sbfly4(a, b) + if self.bf == 5: + return softfp.sbfly5(a, b) + return softfp.sbfly3(a, b) + return a + b, a - b + + def fma(self, c, x, y): + return cute.fma(Float64(c), x, y) + + def _two_minus(self, x, y): + """2x - y: software (exact doubling) for a fraction of the fused butterflies.""" + self.acc2 += self.frac2 + if self.acc2 >= 0.5: + self.acc2 -= 1.0 + return softfp.stwominus(x, y) + return cute.fma(Float64(2.0), x, -y) + + def mul(self, c, x): + return x * Float64(c) + + def neg(self, x): + return fneg_int(x) + + def negv(self, x): + # only used as an FMA addend, where the negation folds into the DFMA + return -x + + def fma_rt(self, c, x, y): + return cute.fma(c, x, y) + + def fma_rt_neg(self, c, x, y): + return cute.fma(-c, x, y) + + def mul_rt(self, c, x): + return c * x + + def mul_rt_neg(self, c, x): + return -(c * x) + + +_BK = DslBk() + + +def _set_soft(frac): + _BK.set_soft(frac) + + +def _dft(xs): + tw = any(x.tw is not None for x in xs) + _BK.balance(*G.codelet_mix(len(xs), tw)) + return G.dft(_BK, xs) + + +def _pair_combine(ar, ai, br, bi, wr, wi): + """Z1 = A + i w B, Z2 = conj(A - i w B) with A = a + conj b, B = a - conj b: + four single additions (A, B), then six FMAs.""" + Ar = _BK.add(ar, br) + Ai = _BK.sub(ai, bi) + Br = _BK.sub(ar, br) + Bi = _BK.add(ai, bi) + z1r = cute.fma(-wi, Br, cute.fma(-wr, Bi, Ar)) + z1i = cute.fma(wr, Br, cute.fma(-wi, Bi, Ai)) + z2r = cute.fma(Float64(2.0), Ar, -z1r) + z2i = cute.fma(Float64(-2.0), Ai, z1i) + return z1r, z1i, z2r, z2i + + +def _out(y): + m = G.materialize(_BK, y) + return m.re.v, m.im.v + + +def _imin(a, b): + return cutlass.min(a, b) + + +def _imax(a, b): + return cutlass.max(a, b) + + +# --------------------------------------------------------------------------- +# Planning +# --------------------------------------------------------------------------- + +_RADICES = [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] +_COST = {} +SMEM_TILE_MAX = 32768 + + +def _ccost(R, tw): + key = (R, tw) + if key not in _COST: + _COST[key] = G.codelet_cost(R, tw) + return _COST[key] + + +def _radix_lists(n, max_stages=3): + out = [] + + def rec(rem, acc): + if rem == 1: + out.append(list(acc)) + return + if len(acc) >= max_stages: + return + for R in _RADICES: + if rem % R == 0: + acc.append(R) + rec(rem // R, acc) + acc.pop() + + rec(n, []) + return out + + +class Stages: + """In-place mixed-radix DIT schedule for length n. + + With pfa=True the last radix N1 is coprime to N2 = n/N1 and the plan is the + Good-Thomas prime-factor split: the first stages are a DIT of length N2 on + each of the N1 rows (natural j -> row (j mod N1)*N2^-1, column + (j mod N2)*N1^-1), and the last stage is a twiddle-free length-N1 DFT across + rows; output (row q, column k2) is X[CRT(q, k2)]. + """ + + def __init__(self, n, radices, pfa=False): + self.n = n + self.R = list(radices) + self.pfa = pfa + s = len(self.R) + self.B = [1] * (s + 1) # B[i] = R_1..R_i (B[0] = 1) + for i in range(s): + self.B[i + 1] = self.B[i] * self.R[i] + assert self.B[s] == n + self.D = [1] * s # D[i] = R_{i+1}..R_s + for i in range(s - 2, -1, -1): + self.D[i] = self.D[i + 1] * self.R[i + 1] + if pfa: + assert s >= 2 + self.N1 = self.R[-1] + self.N2 = n // self.N1 + assert math.gcd(self.N1, self.N2) == 1 + self.u1 = pow(self.N2 % self.N1, -1, self.N1) + self.u2 = pow(self.N1 % self.N2, -1, self.N2) + self.sub = Stages(self.N2, self.R[:-1]) + + def _pinv_dit(self, j): + p = None + for i, R in enumerate(self.R): + r = (j // self.D[i]) % R if self.D[i] > 1 else j % R + t = r * self.B[i] if self.B[i] > 1 else r + p = t if p is None else p + t + return p + + def pinv(self, j): + """Natural index j (DSL Int32) -> in-place storage position.""" + if not self.pfa: + return self._pinv_dit(j) + N1, N2 = self.N1, self.N2 + j1 = ((j % N1) * self.u1) % N1 + j2 = ((j % N2) * self.u2) % N2 + p2 = self.sub._pinv_dit(j2) if len(self.sub.R) > 1 else j2 + return j1 * N2 + p2 + + def blk1(self, j1): + """Stage-1 codelet with inputs x[j1 + (n/R1) r] -> its DIT block index.""" + if len(self.R) == 1: + return j1 * 0 + b = None + for i in range(1, len(self.R)): + R = self.R[i] + r = (j1 // self.D[i]) % R if self.D[i] > 1 else j1 % R + mul = self.B[i] // self.R[0] + t = r * mul if mul > 1 else r + b = t if b is None else b + t + return b + + # -- stage-1 direct global loads: codelet c in [0, n/R0), slot r ---------- + def stage1_rot(self, c): + """PFA: rotation t of the stage-1 codelet's input slots (0 for DIT).""" + if not self.pfa: + return None + q = self.N2 // self.R[0] + c2 = ((c % self.N2) * self.u2) % self.N2 + return c2 // q + + def stage1_src(self, c, rslot, rot): + """Natural index feeding input slot rslot of stage-1 codelet c.""" + R0 = self.R[0] + if rot is None: + return c + (self.n // R0) * rslot + r = (rot * (R0 - 1) + rslot) % R0 # (rslot - rot) mod R0 + return c + (self.n // R0) * r + + def stage1_base(self, c): + """Storage position of slot 0 of stage-1 codelet c.""" + R0 = self.R[0] + if not self.pfa: + return self.blk1(c) * R0 + N1, N2 = self.N1, self.N2 + j1 = ((c % N1) * self.u1) % N1 + c2 = ((c % N2) * self.u2) % N2 + c2p = c2 % (N2 // R0) if N2 // R0 > 1 else c2 * 0 + b2 = self.sub.blk1(c2p) + return j1 * N2 + b2 * R0 + + # -- last stage: codelet k1, output q -> natural index -------------------- + def out_idx(self, k1, q): + if not self.pfa: + return k1 + q * (self.n // self.R[-1]) + n = self.n + a = (self.N2 * self.u1) % n + b = (self.N1 * self.u2) % n + return (q * a + k1 * b) % n + + +class PassPlan: + """Tile of S sequences of length n, NT threads, in-place DIT stages.""" + + def __init__(self, n, radices, S, NT, pre=False, soft=None, pfa=False): + self.soft = SOFT_DEFAULT if soft is None else soft + self.n = n + self.st = Stages(n, radices, pfa) + self.R = self.st.R + self.S = S + self.LD = S + self.PS = S # smem stride between FFT positions + self.SS = 1 # smem stride between sequences + self.SQ = 0 # if > 0: sequence s at (s // SQ) * SS2 + s % SQ (padded rows) + self.SS2 = 0 + self.NT = NT + self.gtw = False # stage twiddles read from global memory (L1) instead of shared memory + self.G = [S * (n // R) for R in self.R] + self.rounds = [(g + NT - 1) // NT for g in self.G] + self.exact = [g % NT == 0 for g in self.G] + self.pre = pre + self.smem_elems = n * S + c = 0 + for i, R in enumerate(self.R): + tw = i > 0 and not (pfa and i == len(self.R) - 1) + c += self.rounds[i] * NT * _ccost(R, tw) + if pre: + units = (S // 2) * n + c += ((units + NT - 1) // NT) * NT * 10 + self.ops_per_tile = c + + def regs(self): + return max(r * R for r, R in zip(self.rounds, self.R)) + + +def est_regs(pp): + return 4 * pp.regs() + 48 + + +def plan_candidates(n, nseq, pre=False, exact_only=True): + cands = [] + lists = [(rs, False) for rs in _radix_lists(n)] + lists += [(rs, True) for rs, _ in lists if len(rs) >= 2 and math.gcd(rs[-1], n // rs[-1]) == 1] + for rs, pfa in lists: + for S in (16, 8, 4, 2, 1): + if pre and S < 2: + continue + if n * S * 16 > SMEM_TILE_MAX: + continue + for NT in (32, 64, 96, 128, 192, 256): + pp = PassPlan(n, rs, S, NT, pre, None, pfa) + if pp.regs() > 32: + continue + if exact_only and not all(pp.exact): + continue + if NT > max(pp.G) and NT > 32: + continue + cands.append(pp) + return cands + + +def plan_score(pp, nseq): + n, S, NT = pp.n, pp.S, pp.NT + ntiles = (nseq + S - 1) // S + work = pp.ops_per_tile * ntiles + n * S * ntiles * (len(pp.R) - 1) * 2 + bps = max(1, min(24, 65536 // (NT * est_regs(pp)), (100 * 1024) // (n * S * 16 + 1024))) + warps = bps * NT // 32 + eff = min(1.0, 0.55 + 0.04 * warps) + slots = NUM_SMS * bps + waves = ntiles / slots + tail = math.ceil(waves) / waves if waves > 0 else 1.0 + if waves < 1: + tail = 1.0 / max(waves, 1e-3) ** 0.5 + return work / eff * min(tail, 2.0) + + +def plan_pass(n, nseq, pre=False, key=None): + if key is not None and key not in _OVERRIDES and key[0] in ("fx", "fy", "fz"): + alt = (key[0][1],) + tuple(key[1:]) + if alt in _OVERRIDES: + key = alt + if key is not None and key in _OVERRIDES: + v = _OVERRIDES[key] + rs, S, NT = v[:3] + soft = tuple(v[3]) if len(v) > 3 and v[3] is not None else None + pfa = bool(v[4]) if len(v) > 4 else False + pp = PassPlan(n, rs, S, NT, pre, soft, pfa) + pp.gtw = bool(v[5]) if len(v) > 5 else False + return pp + cands = plan_candidates(n, nseq, pre, True) or plan_candidates(n, nseq, pre, False) + if not cands: + raise NotImplementedError(f"no plan for n={n}") + return min(cands, key=lambda pp: plan_score(pp, nseq)) + + +# --------------------------------------------------------------------------- +# Shared FFT engine (trace-time) +# --------------------------------------------------------------------------- + + +def _codelet_loc(plan, i, g): + """Codelet g of stage i (i >= 1) -> (seq s, k1, base position).""" + S = plan.S + st = plan.st + m = st.B[i] + Bsz = st.B[i + 1] + if S == 1: + s = g * 0 + h = g + else: + s = g % S + h = g // S + k1 = h % m + if Bsz == plan.n: + base = k1 + else: + base = (h // m) * Bsz + k1 + return s, k1, base + + +def _load_codelet(sm, tw, tw_off, plan, i, s, k1, base): + st = plan.st + R = st.R[i] + m = st.B[i] + step = plan.n // st.B[i + 1] + nott = st.pfa and i == len(st.R) - 1 + xs = [] + for r in range(R): + re, im = ld2(sm, (base + r * m) * plan.PS + _soff(plan, s)) + t = None + if r > 0 and not nott: + if plan.gtw: + if _TWV[0]: + t = ld2_hint_v(tw, tw_off + k1 * (r * step), "L1::evict_last") + else: + t = ld2_hint(tw, tw_off + k1 * (r * step), "L1::evict_last") + else: + t = ld2(tw, tw_off + k1 * (r * step)) + xs.append(G.Cx(G.Real(re), G.Real(im), t)) + return xs + + +@cute.jit +def _special_unit( + sm, + px: cutlass.Constexpr, + LDx: cutlass.Constexpr, + h: cutlass.Constexpr, + has_mid: cutlass.Constexpr, + M: cutlass.Constexpr, + NUs: cutlass.Constexpr, + un, + i, + q, + ro, + rp, + a0r, + a0i, + b0r, + b0i, + aMr, + aMi, + bMr, + bMi, + er, + ei, +): + """One special-column unit of the x-first pairing pass (see kxa).""" + w0, w1, w2, w3 = a0r, b0r, aMr, bMr + w4, w5, w6, w7 = a0i, b0i, aMi, bMi + if cutlass.const_expr(has_mid): + ec = er + if q % 2 == 1: + ec = ei + if q >= 4: + w0 = ec + w1 = Float64(0.0) + w2 = Float64(0.0) + w3 = Float64(0.0) + w4 = Float64(0.0) + w5 = Float64(0.0) + w6 = Float64(0.0) + w7 = Float64(0.0) + w = [w0, w1, w2, w3, w4, w5, w6, w7] + # sign patterns (bit j negates part j): q0 0x90 q1 0xAC q2 0x60 q3 0x5C, q5/q7 0x01 + pk = Int32(0x5C60AC90) + if q >= 4: + pk = Int32(0x01000100) + msk = (pk >> ((q % 4) * 8)) & 0xFF + w = [fxor_sign(w[j], ((msk >> j) & 1) << 31) for j in range(8)] + tot = ((w[0] + w[1]) + (w[2] + w[3])) + ((w[4] + w[5]) + (w[6] + w[7])) + scale = Float64(0.5) + if q >= 4: + scale = Float64(2.0) + val = tot * scale + # own slot i, partner slot h+i (same slot for a self-paired row) + slot = i + if q % 4 >= 2: + if ro != rp: + slot = h + i + pos = slot + if cutlass.const_expr(has_mid): + pm = Int32(px.st.pinv(M // 2) * LDx) + if q >= 4: + pos = pm + slot + if un < NUs: + st1(sm, 2 * pos + q % 2, val) + + +def _copy_tw(smtw, tw, off, n, NT, tid): + """Copy the pass twiddle table W_n^e (n complex) into shared memory.""" + for it in range((n + NT - 1) // NT): + e = tid + NT * it + if n % NT != 0: + e = _imin(e, n - 1) + re, im = ld2(tw, off + e) + st2(smtw, e, re, im) + + +def _soff(plan, s): + """Shared-memory offset of sequence s.""" + if plan.SQ: + return (s // plan.SQ) * plan.SS2 + s % plan.SQ + return s * plan.SS + + +def _split_rounds(rounds, exact, NT=0, Gn=0): + """(batched rounds, tail?): a partial last round with at least one idle warp runs + separately so that warps without codelets skip it (a clamped duplicate costs full + FP64 issue slots).""" + if exact or rounds < 1 or NT - (Gn - (rounds - 1) * NT) < 32: + return rounds, False + return rounds - 1, True + + +def _warp_live(tid, NT, a, Gn): + """Warp-uniform: does the warp of tid own any codelet in round a.""" + return (tid // 32) * 32 + NT * a < Gn + + +@cute.jit +def _mid_rounds( + sm, + tw, + tw_off: cutlass.Constexpr, + plan: cutlass.Constexpr, + i: cutlass.Constexpr, + tid, + a0: cutlass.Constexpr, + a1: cutlass.Constexpr, +): + R = plan.R[i] + m = plan.st.B[i] + NT = plan.NT + Gn = plan.G[i] + locs = [] + xss = [] + oks = [] + for a in cutlass.range_constexpr(a0, a1): + g = tid + NT * a + oks.append(g < Gn) + if cutlass.const_expr(not plan.exact[i]): + g = _imin(g, Gn - 1) + s, k1, base = _codelet_loc(plan, i, g) + locs.append((s, base)) + xss.append(_load_codelet(sm, tw, tw_off, plan, i, s, k1, base)) + yss = [_dft(xs) for xs in xss] + for j in cutlass.range_constexpr(a1 - a0): + a = a0 + j + s, base = locs[j] + outs = [_out(y) for y in yss[j]] + if cutlass.const_expr(plan.exact[i] or (a + 1) * NT <= Gn): + for q in cutlass.range_constexpr(R): + st2(sm, (base + q * m) * plan.PS + _soff(plan, s), outs[q][0], outs[q][1]) + else: + if oks[j]: + for q in cutlass.range_constexpr(R): + st2(sm, (base + q * m) * plan.PS + _soff(plan, s), outs[q][0], outs[q][1]) + + +@cute.jit +def _mid_stage(sm, tw, tw_off: cutlass.Constexpr, plan: cutlass.Constexpr, i: cutlass.Constexpr, tid): + """Stage i >= 1 that reads and writes shared memory (in place).""" + full, tail = _split_rounds(plan.rounds[i], plan.exact[i], plan.NT, plan.G[i]) + if cutlass.const_expr(full > 0): + _mid_rounds(sm, tw, tw_off, plan, i, tid, 0, full) + if cutlass.const_expr(tail): + if _warp_live(tid, plan.NT, full, plan.G[i]): + _mid_rounds(sm, tw, tw_off, plan, i, tid, full, full + 1) + cute.arch.sync_threads() + + +@cute.jit +def _first_rounds(sm, plan: cutlass.Constexpr, tid, a0: cutlass.Constexpr, a1: cutlass.Constexpr): + R = plan.R[0] + NT = plan.NT + Gn = plan.G[0] + S = plan.S + locs = [] + xss = [] + oks = [] + for a in cutlass.range_constexpr(a0, a1): + g = tid + NT * a + oks.append(g < Gn) + if cutlass.const_expr(not plan.exact[0]): + g = _imin(g, Gn - 1) + if cutlass.const_expr(S == 1): + s = g * 0 + b = g + else: + s = g % S + b = g // S + base = b * R + xs = [] + for r in cutlass.range_constexpr(R): + re, im = ld2(sm, (base + r) * plan.PS + _soff(plan, s)) + xs.append(G.Cx(G.Real(re), G.Real(im))) + locs.append((s, base)) + xss.append(xs) + yss = [_dft(xs) for xs in xss] + for j in cutlass.range_constexpr(a1 - a0): + a = a0 + j + s, base = locs[j] + outs = [_out(y) for y in yss[j]] + if cutlass.const_expr(plan.exact[0] or (a + 1) * NT <= Gn): + for q in cutlass.range_constexpr(R): + st2(sm, (base + q) * plan.PS + _soff(plan, s), outs[q][0], outs[q][1]) + else: + if oks[j]: + for q in cutlass.range_constexpr(R): + st2(sm, (base + q) * plan.PS + _soff(plan, s), outs[q][0], outs[q][1]) + + +@cute.jit +def _first_stage_from_smem(sm, plan: cutlass.Constexpr, tid): + """Stage 0 (no twiddles) read from shared memory, results back to smem.""" + full, tail = _split_rounds(plan.rounds[0], plan.exact[0], plan.NT, plan.G[0]) + if cutlass.const_expr(full > 0): + _first_rounds(sm, plan, tid, 0, full) + if cutlass.const_expr(tail): + if _warp_live(tid, plan.NT, full, plan.G[0]): + _first_rounds(sm, plan, tid, full, full + 1) + + +# Stages as a runtime loop over groups of U codelet rounds (bounded code size and +# registers for tiles with many rounds per stage); masked partial last round. +def _round_split(Gn, NT, U): + """(full rounds, loop iterations of U rounds, leftover full rounds, partial tail?)""" + nfull = Gn // NT + return nfull, nfull // U, nfull % U, Gn % NT != 0 + + +@cute.jit +def _mid_rounds_d( + sm, + tw, + tw_off: cutlass.Constexpr, + plan: cutlass.Constexpr, + i: cutlass.Constexpr, + tid, + abase, + cnt: cutlass.Constexpr, + masked: cutlass.Constexpr, +): + R = plan.R[i] + m = plan.st.B[i] + NT = plan.NT + Gn = plan.G[i] + locs = [] + xss = [] + oks = [] + for j in cutlass.range_constexpr(cnt): + g = tid + NT * (abase + j) + if cutlass.const_expr(masked): + oks.append(g < Gn) + g = _imin(g, Gn - 1) + s, k1, base = _codelet_loc(plan, i, g) + locs.append((s, base)) + xss.append(_load_codelet(sm, tw, tw_off, plan, i, s, k1, base)) + yss = [_dft(xs) for xs in xss] + for j in cutlass.range_constexpr(cnt): + s, base = locs[j] + outs = [_out(y) for y in yss[j]] + if cutlass.const_expr(masked): + if oks[j]: + for q in cutlass.range_constexpr(R): + st2(sm, (base + q * m) * plan.PS + _soff(plan, s), outs[q][0], outs[q][1]) + else: + for q in cutlass.range_constexpr(R): + st2(sm, (base + q * m) * plan.PS + _soff(plan, s), outs[q][0], outs[q][1]) + + +@cute.jit +def _mid_stage_d( + sm, tw, tw_off: cutlass.Constexpr, plan: cutlass.Constexpr, i: cutlass.Constexpr, tid, U: cutlass.Constexpr +): + nfull, nit, nrem, tail = _round_split(plan.G[i], plan.NT, U) + if cutlass.const_expr(nit > 1): + for it in cutlass.range(nit, unroll=1): + _mid_rounds_d(sm, tw, tw_off, plan, i, tid, it * U, U, False) + elif cutlass.const_expr(nit == 1): + _mid_rounds_d(sm, tw, tw_off, plan, i, tid, 0, U, False) + if cutlass.const_expr(nrem > 0): + _mid_rounds_d(sm, tw, tw_off, plan, i, tid, nit * U, nrem, False) + if cutlass.const_expr(tail): + if _warp_live(tid, plan.NT, nfull, plan.G[i]): + _mid_rounds_d(sm, tw, tw_off, plan, i, tid, nfull, 1, True) + cute.arch.sync_threads() + + +@cute.jit +def _first_rounds_d(sm, plan: cutlass.Constexpr, tid, abase, cnt: cutlass.Constexpr, masked: cutlass.Constexpr): + R = plan.R[0] + NT = plan.NT + Gn = plan.G[0] + S = plan.S + locs = [] + xss = [] + oks = [] + for j in cutlass.range_constexpr(cnt): + g = tid + NT * (abase + j) + if cutlass.const_expr(masked): + oks.append(g < Gn) + g = _imin(g, Gn - 1) + if cutlass.const_expr(S == 1): + s = g * 0 + b = g + else: + s = g % S + b = g // S + base = b * R + xs = [] + for r in cutlass.range_constexpr(R): + re, im = ld2(sm, (base + r) * plan.PS + _soff(plan, s)) + xs.append(G.Cx(G.Real(re), G.Real(im))) + locs.append((s, base)) + xss.append(xs) + yss = [_dft(xs) for xs in xss] + for j in cutlass.range_constexpr(cnt): + s, base = locs[j] + outs = [_out(y) for y in yss[j]] + if cutlass.const_expr(masked): + if oks[j]: + for q in cutlass.range_constexpr(R): + st2(sm, (base + q) * plan.PS + _soff(plan, s), outs[q][0], outs[q][1]) + else: + for q in cutlass.range_constexpr(R): + st2(sm, (base + q) * plan.PS + _soff(plan, s), outs[q][0], outs[q][1]) + + +@cute.jit +def _first_stage_d(sm, plan: cutlass.Constexpr, tid, U: cutlass.Constexpr): + nfull, nit, nrem, tail = _round_split(plan.G[0], plan.NT, U) + if cutlass.const_expr(nit > 1): + for it in cutlass.range(nit, unroll=1): + _first_rounds_d(sm, plan, tid, it * U, U, False) + elif cutlass.const_expr(nit == 1): + _first_rounds_d(sm, plan, tid, 0, U, False) + if cutlass.const_expr(nrem > 0): + _first_rounds_d(sm, plan, tid, nit * U, nrem, False) + if cutlass.const_expr(tail): + if _warp_live(tid, plan.NT, nfull, plan.G[0]): + _first_rounds_d(sm, plan, tid, nfull, 1, True) + + +# --------------------------------------------------------------------------- +# z-first pipeline: pairing + z FFT, then y and x (or fused y+x) in place +# --------------------------------------------------------------------------- + + +def make_kernels(nz, ny, M, fused=None): + nf = M + 1 + nrows = nz * ny + pz = plan_pass(nz, ny * M, pre=True, key=("z", nz, ny, M)) + py = plan_pass(ny, nz * M, key=("y", nz, ny, M)) + px = plan_pass(M, nrows, key=("x", nz, ny, M)) + TWZ, TWY, TWX, TWP = 0, nz, nz + ny, nz + ny + M + + Np = (M - 1) // 2 # column pairs (k, M-k), 1 <= k <= Np + has_mid = M % 2 == 0 # self-paired column M/2 + Sz = pz.S + ntz = (M + Sz - 1) // Sz # z tiles per ky + spec_tile = (2 * Np) // Sz + spec_slot = (2 * Np) % Sz # slot of column 0 (column M/2 follows) + + # ---------------- pass z: pre-combination + kz FFT ------------------- + NTz, LDz = pz.NT, pz.LD + Hz = Sz // 2 + upairs = Hz * nz + NITp = (upairs + NTz - 1) // NTz + NITs = (nz + NTz - 1) // NTz + sz_last = len(pz.R) - 1 + + gz = ny * ntz + ZTW_PF = min((ny + M + 7) // 8, gz) if (nz, ny, M) in _TWPF_SHAPES else 0 + + @cute.kernel + def kz(X: cute.Tensor, O: cute.Tensor, TW: cute.Tensor): + _set_soft(pz.soft) + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + ky = bid // ntz + tile = bid % ntz + kyn = (ny - ky) % ny + alloc = cutlass.utils.SmemAllocator() + sbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * pz.smem_elems, byte_alignment=16) + tbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * nz, byte_alignment=16) + sm = sbuf.llvm_ptr + xp = X.iterator.llvm_ptr + op = O.iterator.llvm_ptr + tw = TW.iterator.llvm_ptr + smtw = tbuf.llvm_ptr + if cutlass.const_expr(ZTW_PF > 0): + # warm L2 with the twiddle tables of the later passes (cold after the flush) + if tid == 0: + if bid < ZTW_PF: + pf_l2(tw, TWY + bid * 8) + _copy_tw(smtw, tw, TWZ, nz, NTz, tid) + # pairs: gather + va = [] + vb = [] + vw = [] + meta = [] + for it in cutlass.range_constexpr(NITp): + u = tid + NTz * it + if cutlass.const_expr(upairs % NTz != 0): + u = _imin(u, upairs - 1) + i = u % Hz + p = u // Hz + pi = tile * Hz + i + k = _imin(pi, Np - 1) + 1 + pn = (nz - p) % nz + va.append(ld2(xp, (p * ny + ky) * nf + k)) + vb.append(ld2(xp, (pn * ny + kyn) * nf + (M - k))) + vw.append(ld2(tw, TWP + k)) + meta.append((i, p, pn, pi)) + for it in cutlass.range_constexpr(NITp): + ar, ai = va[it] + br, bi = vb[it] + wr, wi = vw[it] + i, p, pn, pi = meta[it] + z1r, z1i, z2r, z2i = _pair_combine(ar, ai, br, bi, wr, wi) + if pi < Np: + st2(sm, pz.st.pinv(p) * LDz + 2 * i, z1r, z1i) + st2(sm, pz.st.pinv(pn) * LDz + 2 * i + 1, z2r, z2i) + # special columns 0 (and M/2) live in tile spec_tile + if tile == spec_tile: + v0 = [] + for it in cutlass.range_constexpr(NITs): + p = tid + NTz * it + if cutlass.const_expr(nz % NTz != 0): + p = _imin(p, nz - 1) + pn = (nz - p) % nz + a = ld2(xp, (p * ny + ky) * nf) + b = ld2(xp, (pn * ny + kyn) * nf + M) + c = ld2(xp, (p * ny + ky) * nf + M) + d = ld2(xp, (pn * ny + kyn) * nf) + e = (Float64(0.0), Float64(0.0)) + if cutlass.const_expr(has_mid): + e = ld2(xp, (pn * ny + kyn) * nf + M // 2) + v0.append((p, a, b, c, d, e)) + for it in cutlass.range_constexpr(NITs): + p, a, b, c, d, e = v0[it] + Pr = a[0] + d[0] + Pi = a[1] - d[1] + Qr = c[0] + b[0] + Qi = c[1] - b[1] + u1 = Pr + Qr + v1 = Pi - Qi + x1 = Pi + Qi + y1 = Pr - Qr + zr = (u1 - v1) * 0.5 + zi = (x1 + y1) * 0.5 + pp = pz.st.pinv(p) * LDz + st2(sm, pp + spec_slot, zr, zi) + if cutlass.const_expr(has_mid): + st2(sm, pp + spec_slot + 1, e[0] + e[0], (-e[1]) - e[1]) + cute.arch.sync_threads() + if cutlass.const_expr(sz_last > 0): + _first_stage_from_smem(sm, pz, tid) + cute.arch.sync_threads() + for i in cutlass.range_constexpr(1, sz_last): + _mid_stage(sm, smtw, 0, pz, i, tid) + # last stage: smem -> global + i = sz_last + R = pz.R[i] + m = pz.st.B[i] + locs = [] + xss = [] + for a in cutlass.range_constexpr(pz.rounds[i]): + g = tid + NTz * a + if cutlass.const_expr(not pz.exact[i]): + g = _imin(g, pz.G[i] - 1) + s, k1, base = _codelet_loc(pz, i, g) + locs.append((s, k1)) + if cutlass.const_expr(i == 0): + xss.append( + [G.Cx(G.Real(v[0]), G.Real(v[1])) for v in [ld2(sm, (base + r * m) * LDz + s) for r in range(R)]] + ) + else: + xss.append(_load_codelet(sm, smtw, 0, pz, i, s, k1, base)) + yss = [_dft(xs) for xs in xss] + for a in cutlass.range_constexpr(pz.rounds[i]): + s, k1 = locs[a] + item = tile * Sz + s + pi = item // 2 + odd = item % 2 + row = ky + col = pi + 1 + if odd == 1: + row = kyn + col = M - (pi + 1) + if item >= 2 * Np: + row = ky + col = (item - 2 * Np) * (M // 2) + outs = [_out(y) for y in yss[a]] + if item < M: + for q in cutlass.range_constexpr(R): + st2(op, (pz.st.out_idx(k1, q) * ny + row) * M + col, outs[q][0], outs[q][1]) + + # ---------------- generic column/row pass (y and x) -------------------- + def make_seq_pass(pp, tw_off, kind): + S, NT, LD = pp.S, pp.NT, pp.LD + n = pp.n + R0 = pp.R[0] + nlast = len(pp.R) - 1 + if kind == "y": + nty = (M + S - 1) // S + + def addr(bid, s, j): + z = bid // nty + k = _imin((bid % nty) * S + s, M - 1) + return (z * ny + j) * M + k + + else: + + def addr(bid, s, j): + r = _imin(bid * S + s, nrows - 1) + return r * M + j + + @cute.kernel + def kern(O: cute.Tensor, TW: cute.Tensor): + _set_soft(pp.soft) + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + alloc = cutlass.utils.SmemAllocator() + sbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * pp.smem_elems, byte_alignment=16) + tbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * n, byte_alignment=16) + sm = sbuf.llvm_ptr + op = O.iterator.llvm_ptr + tw = TW.iterator.llvm_ptr + smtw = tbuf.llvm_ptr + _copy_tw(smtw, tw, tw_off, n, NT, tid) + # stage 0: global -> regs -> smem + locs = [] + xss = [] + for a in cutlass.range_constexpr(pp.rounds[0]): + g = tid + NT * a + if cutlass.const_expr(not pp.exact[0]): + g = _imin(g, pp.G[0] - 1) + if cutlass.const_expr(S == 1): + s = g * 0 + j1 = g + else: + s = g % S + j1 = g // S + xs = [] + rot = pp.st.stage1_rot(j1) + for r in cutlass.range_constexpr(R0): + re, im = ld2(op, addr(bid, s, pp.st.stage1_src(j1, r, rot))) + xs.append(G.Cx(G.Real(re), G.Real(im))) + if cutlass.const_expr(nlast == 0): + b = j1 * R0 + else: + b = pp.st.stage1_base(j1) + locs.append((s, j1, b)) + xss.append(xs) + yss = [_dft(xs) for xs in xss] + if cutlass.const_expr(nlast == 0): + for a in cutlass.range_constexpr(pp.rounds[0]): + s, j1, b = locs[a] + outs = [_out(y) for y in yss[a]] + for q in cutlass.range_constexpr(R0): + st2(op, addr(bid, s, q), outs[q][0], outs[q][1]) + else: + for a in cutlass.range_constexpr(pp.rounds[0]): + s, j1, b = locs[a] + outs = [_out(y) for y in yss[a]] + if cutlass.const_expr(pp.exact[0] or (a + 1) * NT <= pp.G[0]): + for q in cutlass.range_constexpr(R0): + st2(sm, (b + q) * LD + s, outs[q][0], outs[q][1]) + else: + if tid + NT * a < pp.G[0]: + for q in cutlass.range_constexpr(R0): + st2(sm, (b + q) * LD + s, outs[q][0], outs[q][1]) + cute.arch.sync_threads() + for i in cutlass.range_constexpr(1, nlast): + _mid_stage(sm, smtw, 0, pp, i, tid) + # last stage: smem -> global + i = nlast + R = pp.R[i] + locs2 = [] + xss2 = [] + for a in cutlass.range_constexpr(pp.rounds[i]): + g = tid + NT * a + if cutlass.const_expr(not pp.exact[i]): + g = _imin(g, pp.G[i] - 1) + s, k1, base = _codelet_loc(pp, i, g) + locs2.append((s, k1)) + xss2.append(_load_codelet(sm, smtw, 0, pp, i, s, k1, base)) + yss2 = [_dft(xs) for xs in xss2] + for a in cutlass.range_constexpr(pp.rounds[i]): + s, k1 = locs2[a] + outs = [_out(y) for y in yss2[a]] + for q in cutlass.range_constexpr(R): + st2(op, addr(bid, s, pp.st.out_idx(k1, q)), outs[q][0], outs[q][1]) + + return kern + + ky_ = make_seq_pass(py, TWY, "y") + kx = make_seq_pass(px, TWX, "x") + gy = nz * ((M + py.S - 1) // py.S) + gx = (nrows + px.S - 1) // px.S + + # ---------------- optional fused y+x pass: one block per z plane -------- + if fused is not None: + NTp = fused["NT"] + RS = M + 1 + qy = PassPlan(ny, fused["ry"], M, NTp) + qy.PS, qy.SS = RS, 1 + qx = PassPlan(M, fused["rx"], ny, NTp) + qx.PS, qx.SS = 1, RS + assert not qy.st.pfa and not qx.st.pfa + ptot = ny * M + NITl = (ptot + NTp - 1) // NTp + + @cute.kernel + def kyx(O: cute.Tensor, TW: cute.Tensor): + _set_soft(fused.get("soft", SOFT_DEFAULT)) + tid, _, _ = cute.arch.thread_idx() + z = cute.arch.block_idx()[0] + alloc = cutlass.utils.SmemAllocator() + sbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * ny * RS, byte_alignment=16) + tyb = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * ny, byte_alignment=16) + txb = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * M, byte_alignment=16) + sm = sbuf.llvm_ptr + op = O.iterator.llvm_ptr + tw = TW.iterator.llvm_ptr + ty = tyb.llvm_ptr + tx = txb.llvm_ptr + _copy_tw(ty, tw, TWY, ny, NTp, tid) + _copy_tw(tx, tw, TWX, M, NTp, tid) + vals = [] + for it in cutlass.range_constexpr(NITl): + e = tid + NTp * it + if cutlass.const_expr(ptot % NTp != 0): + e = _imin(e, ptot - 1) + vals.append((e, ld2(op, z * ptot + e))) + for it in cutlass.range_constexpr(NITl): + e, (re, im) = vals[it] + y = e // M + k = e % M + st2(sm, qy.st.pinv(y) * RS + qx.st.pinv(k), re, im) + cute.arch.sync_threads() + for qq in (qy, qx): + tws = ty if qq is qy else tx + _first_stage_from_smem(sm, qq, tid) + cute.arch.sync_threads() + for i in cutlass.range_constexpr(1, len(qq.R)): + _mid_stage(sm, tws, 0, qq, i, tid) + vals = [] + for it in cutlass.range_constexpr(NITl): + e = tid + NTp * it + if cutlass.const_expr(ptot % NTp != 0): + e = _imin(e, ptot - 1) + y = e // M + k = e % M + vals.append((e, ld2(sm, y * RS + k))) + for it in cutlass.range_constexpr(NITl): + e, (re, im) = vals[it] + st2(op, z * ptot + e, re, im) + + @cute.jit + def launch(X: cute.Tensor, O: cute.Tensor, TW: cute.Tensor, stream: cuda.CUstream): + kz(X, O, TW).launch(grid=[gz, 1, 1], block=[NTz, 1, 1], stream=stream) + if cutlass.const_expr(fused is not None): + kyx(O, TW).launch(grid=[nz, 1, 1], block=[fused["NT"], 1, 1], stream=stream) + if cutlass.const_expr(fused is None): + ky_(O, TW).launch(grid=[gy, 1, 1], block=[py.NT, 1, 1], stream=stream) + kx(O, TW).launch(grid=[gx, 1, 1], block=[px.NT, 1, 1], stream=stream) + + return launch + + +# --------------------------------------------------------------------------- +# x-first pipeline. After the pre-combination the transform is a plain +# separable inverse 3D DFT, so the passes can run in any order. Doing the +# pre-combination + x FFT first reads the cold input as contiguous rows (row p +# is paired with row -p, both contiguous), then y, then z (warm data, no +# reflected partner reads). +# --------------------------------------------------------------------------- + + +def make_kernels_xfirst(nz, ny, M, yz2d=None): + nf = M + 1 + PDL = (nz, ny, M) in _PDL_SHAPES + nrows = nz * ny + px = plan_pass(M, nrows, pre=True, key=("fx", nz, ny, M)) + if yz2d is None: + py = plan_pass(ny, nz * M, key=("fy", nz, ny, M)) + pz = plan_pass(nz, ny * M, key=("fz", nz, ny, M)) + else: + py = pz = None + assert nz == ny + yzd = YZ2D(ny, yz2d["A"], yz2d["rc"], yz2d.get("N1", 1)) + + # per-shape soft/native balance extensions of the plans' soft tuples (see DslBk.set_soft) + f2 = _SOFT2M_SHAPES.get((nz, ny, M)) + ic = _INTCYC_SHAPES.get((nz, ny, M)) + bfv = _SOFTBF_SHAPES.get((nz, ny, M)) + if bfv is not None and ic is None: + ic = INT_PAIR_CYC + if f2 is not None or ic is not None: + ext = (0.0 if f2 is None else f2,) + (() if ic is None else (ic,)) + (() if bfv is None else (bfv,)) + for pp_ in (px, py, pz): + if pp_ is not None: + pp_.soft = tuple(pp_.soft[:4]) + ext + if yz2d is not None: + yz2d = dict(yz2d) + for key_ in ("softB", "softC"): + yz2d[key_] = tuple((yz2d.get(key_) or SOFT_DEFAULT)[:4]) + ext + + def srow(row): + """Storage row of logical row kz*ny+ky (pass-B input layout for the 2D passes).""" + if yz2d is None: + return row + return yzd.p_in(row // ny) * ny + yzd.p_in(row % ny) + + TWZ, TWY, TWX, TWP = 0, nz, nz + ny, nz + ny + M + Np = (M - 1) // 2 + has_mid = M % 2 == 0 + Sx, NTx, LDx = px.S, px.NT, px.LD + assert Sx % 2 == 0 + h = Sx // 2 + + P0 = 1 + ny // 2 # pairs in a self-paired plane + NMID = (nz - 1) // 2 # planes kz with kz < nz - kz (kz >= 1) + # row pairs (r, rp), r <= rp, rp = partner of r: plane 0, planes 1..NMID, then plane nz/2 if nz is even + npairs = P0 + NMID * ny + (P0 if nz % 2 == 0 else 0) + ntx = (npairs + h - 1) // h + + # ---------------- pass A: pairing + x FFT on row pairs ----------------- + upairs = Np * 2 * h + NITp = (upairs + NTx - 1) // NTx + pmap = _PAIRMAP_SHAPES.get((nz, ny, M), 0) + PM_GI = min(2 * h, 8) # lanes over row slots (pmap 1) + PM_NI = (2 * h) // PM_GI + PM_KW = 32 // PM_GI + if pmap == 1: + assert (2 * h) % PM_GI == 0 and NTx % 32 == 0 + NITp = (PM_NI * ((Np + PM_KW - 1) // PM_KW) * 32 + NTx - 1) // NTx + UPP = 8 if has_mid else 4 # special-column units per row pair + NUs = UPP * h + NITs = (NUs + NTx - 1) // NTx + sx_last = len(px.R) - 1 + + @cute.jit + def pair_rows(i0): + """Row pair i in closed form (no table load on the critical path): plane 0, + planes 1..NMID, then plane nz/2 if nz is even.""" + i = _imin(i0, npairs - 1) + j = _imax(i - P0, 0) + kz = 1 + j // ny + ky = j % ny + if cutlass.const_expr(nz % 2 == 0): + jm = j - NMID * ny + if jm >= 0: + kz = i * 0 + nz // 2 + ky = jm + if i < P0: + kz = i * 0 + ky = i + r = kz * ny + ky + rp = ((nz - kz) % nz) * ny + (ny - ky) % ny + return r, rp + + xgtw = px.gtw + if (nz, ny, M) in _TWPF_SHAPES: + ntw_ = nz + ny + 2 * M + if yz2d is not None: + ntw_ += yzd.nc + ny * nz + TW_PF = min((ntw_ + 7) // 8, ntx) # 128-byte lines, one per block + else: + TW_PF = 0 + XS = (nz, ny, M) in _XSTREAM_SHAPES + YORD = (nz, ny, M) in _YORD_SHAPES + POLZ = _ZPOL_SHAPES.get((nz, ny, M)) + + @cute.kernel + def kxa(X: cute.Tensor, O: cute.Tensor, TW: cute.Tensor): + _set_soft(px.soft) + xload = ld2_stream if XS else ld2 + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + alloc = cutlass.utils.SmemAllocator() + sbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * px.smem_elems, byte_alignment=16) + sm = sbuf.llvm_ptr + xp = X.iterator.llvm_ptr + op = O.iterator.llvm_ptr + tw = TW.iterator.llvm_ptr + if cutlass.const_expr(TW_PF > 0): + # warm L2 with the twiddle tables of the later passes (cold after the flush) + if tid == 0: + if bid < TW_PF: + pf_l2(tw, bid * 8) + if cutlass.const_expr(xgtw): + smtw = _ptr_add(tw, TWX) + else: + tbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * M, byte_alignment=16) + smtw = tbuf.llvm_ptr + _copy_tw(smtw, tw, TWX, M, NTx, tid) + va = [] + vb = [] + vw = [] + meta = [] + for it in cutlass.range_constexpr(NITp): + u = tid + NTx * it + if cutlass.const_expr(upairs % NTx != 0): + u = _imin(u, upairs - 1) + if cutlass.const_expr(pmap == 1): + # PM_GI consecutive lanes cover PM_GI row slots of one column (no shared-memory + # bank conflicts on the stores); a warp covers 32 / PM_GI consecutive columns + ul = tid + NTx * it + lane = ul % 32 + wv = ul // 32 + i = lane % PM_GI + PM_GI * (wv % PM_NI) + k = _imin(lane // PM_GI + PM_KW * (wv // PM_NI), Np - 1) + 1 + else: + # lanes run along a row: contiguous global loads + i = u // Np + k = u % Np + 1 + # slot i < h: (own, partner) = (ro, rp); slot h + i: (rp, ro) + ii = i % h + r1, r2 = pair_rows(bid * h + ii) + ro = r1 + rp = r2 + if i >= h: + ro = r2 + rp = r1 + va.append(xload(xp, ro * nf + k)) + vb.append(xload(xp, rp * nf + (M - k))) + vw.append(ld2(tw, TWP + k)) + meta.append((i, k, ii)) + # Special columns 0 and M/2, one real output per lane (no divergent FP64): + # with 2*H0 = X0 + conj(X0'), 2*HM = XM + conj(XM') (' = partner row), + # own col 0 = (H0 + HM) + i (H0 - HM), partner col 0 = its conj-swapped twin, + # own col M/2 = 2 conj(XM/2'), partner col M/2 = 2 conj(XM/2). + # Unit q of a row pair forms a signed sum of the 8 parts of X0, X0', XM, XM' + # (q < 4, then * 0.5) or a single +-XM/2 part (q >= 4, then * 2). + v0 = [] + for it in cutlass.range_constexpr(NITs): + un = tid + NTx * it + if cutlass.const_expr(NUs % NTx != 0): + un = _imin(un, NUs - 1) + i = un // UPP + q = un % UPP + ro, rp = pair_rows(bid * h + i) + a0 = xload(xp, ro * nf) + b0 = xload(xp, rp * nf) + aM = xload(xp, ro * nf + M) + bM = xload(xp, rp * nf + M) + e = (Float64(0.0), Float64(0.0)) + if cutlass.const_expr(has_mid): + rm = rp + if q >= 6: + rm = ro + e = xload(xp, rm * nf + M // 2) + v0.append((un, i, q, ro, rp, a0, b0, aM, bM, e)) + for it in cutlass.range_constexpr(NITp): + ar, ai = va[it] + br, bi = vb[it] + wr, wi = vw[it] + i, k, ii = meta[it] + z1r, z1i, z2r, z2i = _pair_combine(ar, ai, br, bi, wr, wi) + # Z1 -> own row (slot i), Z2 -> partner row (the other slot of the pair) + sp = i + h + if i >= h: + sp = i - h + st2(sm, px.st.pinv(k) * LDx + i, z1r, z1i) + st2(sm, px.st.pinv(M - k) * LDx + sp, z2r, z2i) + for it in cutlass.range_constexpr(NITs): + un, i, q, ro, rp, a0, b0, aM, bM, e = v0[it] + args = (un, i, q, ro, rp, a0[0], a0[1], b0[0], b0[1], aM[0], aM[1], bM[0], bM[1], e[0], e[1]) + if cutlass.const_expr(NUs >= NTx * (it + 1)): + _special_unit(sm, px, LDx, h, has_mid, M, NUs, *args) + else: + # skip warps without units (warp-uniform branch) + if tid - tid % 32 + NTx * it < NUs: + _special_unit(sm, px, LDx, h, has_mid, M, NUs, *args) + cute.arch.sync_threads() + if cutlass.const_expr(sx_last > 0): + _first_stage_from_smem(sm, px, tid) + cute.arch.sync_threads() + for i in cutlass.range_constexpr(1, sx_last): + _mid_stage(sm, smtw, 0, px, i, tid) + i = sx_last + R = px.R[i] + locs = [] + xss = [] + for a in cutlass.range_constexpr(px.rounds[i]): + g = tid + NTx * a + if cutlass.const_expr(not px.exact[i]): + g = _imin(g, px.G[i] - 1) + s, k1, base = _codelet_loc(px, i, g) + locs.append((s, k1)) + xss.append(_load_codelet(sm, smtw, 0, px, i, s, k1, base)) + yss = [_dft(xs) for xs in xss] + if cutlass.const_expr(PDL): + cute.arch.griddepcontrol_launch_dependents() + for a in cutlass.range_constexpr(px.rounds[i]): + s, k1 = locs[a] + outs = [_out(y) for y in yss[a]] + ii = s % h + ro, rp = pair_rows(bid * h + ii) + row = ro + if s >= h: + row = rp + if s < h or ro != rp: + row = srow(row) + for q in cutlass.range_constexpr(R): + st2(op, row * M + px.st.out_idx(k1, q), outs[q][0], outs[q][1]) + + # ---------------- pass A, software pipelined: persistent CTAs, bulk async copies -------- + # Each CTA loops over tiles bid, bid + GX, ...; the raw input rows of the next NB - 1 tiles + # are in flight (cp.async.bulk -> shared, mbarrier completion) while the current tile is + # paired and transformed in place in its buffer. + XP = _XPIPE_SHAPES.get((nz, ny, M)) + if XP is not None: + NB = XP["NB"] + GX = min(NUM_SMS * XP["cps"], ntx) + RAWE = 2 * h * nf # complex elements of one raw tile (2h rows) + BUFE = (max(RAWE, px.smem_elems) + 7) // 8 * 8 + TXB = RAWE * 16 + assert 2 * h <= 32 + + @cute.kernel + def kxp(X: cute.Tensor, O: cute.Tensor, TW: cute.Tensor): + _set_soft(px.soft) + _TWV[0] = True + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + alloc = cutlass.utils.SmemAllocator() + sbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * NB * BUFE, byte_alignment=128) + mbuf = alloc.allocate_array(element_type=cutlass.Int64, num_elems=NB, byte_alignment=8) + sm0 = sbuf.llvm_ptr + mbp = mbuf.llvm_ptr + xp = X.iterator.llvm_ptr + op = O.iterator.llvm_ptr + tw = TW.iterator.llvm_ptr + if cutlass.const_expr(TW_PF > 0): + if tid == 0: + if bid < TW_PF: + pf_l2(tw, bid * 8) + if cutlass.const_expr(xgtw): + smtw = _ptr_add(tw, TWX) + else: + tbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * M, byte_alignment=16) + smtw = tbuf.llvm_ptr + _copy_tw(smtw, tw, TWX, M, NTx, tid) + if tid == 0: + for b in cutlass.range_constexpr(NB): + mbar_init(saddr(mbp, 8 * b), 1) + mbar_fence_init() + cute.arch.sync_threads() + nmine = (ntx - bid + GX - 1) // GX + lane = tid % 32 + # prologue: the first NB tiles of this CTA + if tid < 32: + for j in cutlass.range_constexpr(NB): + t = bid + j * GX + if t < ntx: + if lane == 0: + mbar_expect_tx(saddr(mbp, 8 * j), TXB) + cute.arch.sync_warp() + if lane < 2 * h: + r1, r2 = pair_rows(t * h + lane % h) + row = r1 + if lane >= h: + row = r2 + bulk_g2s( + saddr(sm0, (j * BUFE + lane * nf) * 16), + gaddr(xp, row * nf), + nf * 16, + saddr(mbp, 8 * j), + ) + for it in cutlass.range(nmine, unroll=1): + tid = tid_v() + b = it % NB + t = bid + it * GX + sm = _ptr_add(sm0, b * BUFE) + mbar_wait(saddr(mbp, 8 * b), (it // NB) % 2) + # raw rows: slot j < h = first row of pair j, slot h + j = its partner row + va = [] + vb = [] + vw = [] + meta = [] + for itp in cutlass.range_constexpr(NITp): + u = tid + NTx * itp + if cutlass.const_expr(upairs % NTx != 0): + u = _imin(u, upairs - 1) + if cutlass.const_expr(pmap == 1): + ul = tid + NTx * itp + ln = ul % 32 + wv = ul // 32 + i = ln % PM_GI + PM_GI * (wv % PM_NI) + k = _imin(ln // PM_GI + PM_KW * (wv // PM_NI), Np - 1) + 1 + else: + i = u // Np + k = u % Np + 1 + sp = i + h + if i >= h: + sp = i - h + va.append(ld2(sm, i * nf + k)) + vb.append(ld2(sm, sp * nf + (M - k))) + vw.append(ld2_hint_v(tw, TWP + k, "L1::evict_last")) + meta.append((i, k, sp)) + v0 = [] + for its in cutlass.range_constexpr(NITs): + un = tid + NTx * its + if cutlass.const_expr(NUs % NTx != 0): + un = _imin(un, NUs - 1) + i = un // UPP + q = un % UPP + ro, rp = pair_rows(t * h + i) + a0 = ld2(sm, i * nf) + b0 = ld2(sm, (i + h) * nf) + aM = ld2(sm, i * nf + M) + bM = ld2(sm, (i + h) * nf + M) + e = (Float64(0.0), Float64(0.0)) + if cutlass.const_expr(has_mid): + sm_ = i + h + if q >= 6: + sm_ = i + e = ld2(sm, sm_ * nf + M // 2) + v0.append((un, i, q, ro, rp, a0, b0, aM, bM, e)) + cute.arch.sync_threads() + for itp in cutlass.range_constexpr(NITp): + ar, ai = va[itp] + br, bi = vb[itp] + wr, wi = vw[itp] + i, k, sp = meta[itp] + z1r, z1i, z2r, z2i = _pair_combine(ar, ai, br, bi, wr, wi) + st2(sm, px.st.pinv(k) * LDx + i, z1r, z1i) + st2(sm, px.st.pinv(M - k) * LDx + sp, z2r, z2i) + for its in cutlass.range_constexpr(NITs): + un, i, q, ro, rp, a0, b0, aM, bM, e = v0[its] + args = (un, i, q, ro, rp, a0[0], a0[1], b0[0], b0[1], aM[0], aM[1], bM[0], bM[1], e[0], e[1]) + if cutlass.const_expr(NUs >= NTx * (its + 1)): + _special_unit(sm, px, LDx, h, has_mid, M, NUs, *args) + else: + if tid - tid % 32 + NTx * its < NUs: + _special_unit(sm, px, LDx, h, has_mid, M, NUs, *args) + cute.arch.sync_threads() + if cutlass.const_expr(sx_last > 0): + _first_stage_from_smem(sm, px, tid) + cute.arch.sync_threads() + for i in cutlass.range_constexpr(1, sx_last): + _mid_stage(sm, smtw, 0, px, i, tid) + i = sx_last + R = px.R[i] + locs = [] + xss = [] + for a in cutlass.range_constexpr(px.rounds[i]): + g = tid + NTx * a + if cutlass.const_expr(not px.exact[i]): + g = _imin(g, px.G[i] - 1) + s, k1, base = _codelet_loc(px, i, g) + locs.append((s, k1)) + xss.append(_load_codelet(sm, smtw, 0, px, i, s, k1, base)) + cute.arch.sync_threads() + # buffer b is free: refill it with this CTA's tile it + NB (overlaps the last stage) + if tid < 32: + t2 = t + NB * GX + if t2 < ntx: + if lane == 0: + fence_async_shared() + mbar_expect_tx(saddr(mbp, 8 * b), TXB) + cute.arch.sync_warp() + if lane < 2 * h: + r1, r2 = pair_rows(t2 * h + lane % h) + row = r1 + if lane >= h: + row = r2 + bulk_g2s( + saddr(sm0, (b * BUFE + lane * nf) * 16), + gaddr(xp, row * nf), + nf * 16, + saddr(mbp, 8 * b), + ) + yss = [_dft(xs) for xs in xss] + for a in cutlass.range_constexpr(px.rounds[i]): + s, k1 = locs[a] + outs = [_out(y) for y in yss[a]] + ii = s % h + ro, rp = pair_rows(t * h + ii) + row = ro + if s >= h: + row = rp + if s < h or ro != rp: + row = srow(row) + for q in cutlass.range_constexpr(R): + st2(op, row * M + px.st.out_idx(k1, q), outs[q][0], outs[q][1]) + if cutlass.const_expr(PDL): + cute.arch.griddepcontrol_launch_dependents() + _TWV[0] = False + + # ---------------- passes B (y) and C (z): plain column FFTs in place ----- + def make_inplace_col(pp, tw_off, kind): + S, NT, LD = pp.S, pp.NT, pp.LD + POL = POLZ if kind == "z" else None + gtw = pp.gtw + if gtw: + + def gload(b, i): + return ld2_hint(b, i, "L1::no_allocate") + + else: + gload = ld2 + n = pp.n + R0 = pp.R[0] + nlast = len(pp.R) - 1 + ntk = (M + S - 1) // S + + if kind == "y": + + def addr(bid, s, j): + a = bid // ntk + if cutlass.const_expr(YORD): + # planes in reverse order of the x pass's completion (most recently + # written, likely L2-resident planes first): c, c -+ 1, c +- 1, ... + i = a + d = (i + 1) // 2 + sd = 2 * d * (i % 2) + if cutlass.const_expr(nz % 2 == 0): + a = nz // 2 + d - sd + else: + a = nz // 2 - d + sd + k = _imin((bid % ntk) * S + s, M - 1) + return (a * ny + j) * M + k + + else: + + def addr(bid, s, j): + a = bid // ntk + k = _imin((bid % ntk) * S + s, M - 1) + return (j * ny + a) * M + k + + @cute.jit + def stage0(sm, op, bid, tid, a0: cutlass.Constexpr, a1: cutlass.Constexpr, sync_first: cutlass.Constexpr, pol): + locs = [] + xss = [] + for a in cutlass.range_constexpr(a0, a1): + g = tid + NT * a + if cutlass.const_expr(not pp.exact[0]): + g = _imin(g, pp.G[0] - 1) + if cutlass.const_expr(S == 1): + s = g * 0 + j1 = g + else: + s = g % S + j1 = g // S + xs = [] + rot = pp.st.stage1_rot(j1) + for r in cutlass.range_constexpr(R0): + re, im = gload(op, addr(bid, s, pp.st.stage1_src(j1, r, rot))) + xs.append(G.Cx(G.Real(re), G.Real(im))) + if cutlass.const_expr(nlast == 0): + b = j1 * R0 + else: + b = pp.st.stage1_base(j1) + locs.append((s, j1, b)) + xss.append(xs) + yss = [_dft(xs) for xs in xss] + if cutlass.const_expr(nlast == 0): + if cutlass.const_expr(sync_first): + cute.arch.sync_threads() + for j in cutlass.range_constexpr(a1 - a0): + s, j1, b = locs[j] + outs = [_out(y) for y in yss[j]] + for q in cutlass.range_constexpr(R0): + gst(op, addr(bid, s, q), outs[q][0], outs[q][1], pol) + else: + for j in cutlass.range_constexpr(a1 - a0): + a = a0 + j + s, j1, b = locs[j] + outs = [_out(y) for y in yss[j]] + if cutlass.const_expr(pp.exact[0] or (a + 1) * NT <= pp.G[0]): + for q in cutlass.range_constexpr(R0): + st2(sm, (b + q) * LD + s, outs[q][0], outs[q][1]) + else: + if tid + NT * a < pp.G[0]: + for q in cutlass.range_constexpr(R0): + st2(sm, (b + q) * LD + s, outs[q][0], outs[q][1]) + + @cute.jit + def last(sm, smtw, op, bid, tid, a0: cutlass.Constexpr, a1: cutlass.Constexpr, pol): + i = nlast + R = pp.R[i] + locs2 = [] + xss2 = [] + for a in cutlass.range_constexpr(a0, a1): + g = tid + NT * a + if cutlass.const_expr(not pp.exact[i]): + g = _imin(g, pp.G[i] - 1) + s, k1, base = _codelet_loc(pp, i, g) + locs2.append((s, k1)) + xss2.append(_load_codelet(sm, smtw, 0, pp, i, s, k1, base)) + yss2 = [_dft(xs) for xs in xss2] + if cutlass.const_expr(PDL and kind == "y"): + cute.arch.griddepcontrol_launch_dependents() + for j in cutlass.range_constexpr(a1 - a0): + s, k1 = locs2[j] + outs = [_out(y) for y in yss2[j]] + for q in cutlass.range_constexpr(R): + gst(op, addr(bid, s, pp.st.out_idx(k1, q)), outs[q][0], outs[q][1], pol) + + @cute.kernel + def kern(O: cute.Tensor, TW: cute.Tensor): + _set_soft(pp.soft) + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + alloc = cutlass.utils.SmemAllocator() + sbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * pp.smem_elems, byte_alignment=16) + sm = sbuf.llvm_ptr + op = O.iterator.llvm_ptr + tw = TW.iterator.llvm_ptr + if cutlass.const_expr(gtw): + # twiddles straight from L1 (evict_last): frees the table's shared memory + smtw = _ptr_add(tw, tw_off) + else: + tbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * n, byte_alignment=16) + smtw = tbuf.llvm_ptr + _copy_tw(smtw, tw, tw_off, n, NT, tid) + if cutlass.const_expr(PDL): + cute.arch.griddepcontrol_wait() + pol = mk_policy(POL) if POL else None + if cutlass.const_expr(nlast == 0): + # single in-place stage: every load must precede every store + stage0(sm, op, bid, tid, 0, pp.rounds[0], True, pol) + else: + full0, tail0 = _split_rounds(pp.rounds[0], pp.exact[0], NT, pp.G[0]) + if cutlass.const_expr(full0 > 0): + stage0(sm, op, bid, tid, 0, full0, False, pol) + if cutlass.const_expr(tail0): + if _warp_live(tid, NT, full0, pp.G[0]): + stage0(sm, op, bid, tid, full0, full0 + 1, False, pol) + cute.arch.sync_threads() + for i in cutlass.range_constexpr(1, nlast): + _mid_stage(sm, smtw, 0, pp, i, tid) + fl, tl = _split_rounds(pp.rounds[nlast], pp.exact[nlast], NT, pp.G[nlast]) + if cutlass.const_expr(fl > 0): + last(sm, smtw, op, bid, tid, 0, fl, pol) + if cutlass.const_expr(tl): + if _warp_live(tid, NT, fl, pp.G[nlast]): + last(sm, smtw, op, bid, tid, fl, fl + 1, pol) + + return kern, ntk + + if yz2d is None: + kyb, ntky = make_inplace_col(py, TWY, "y") + kzc, ntkz = make_inplace_col(pz, TWZ, "z") + gy = nz * ntky + gz = ny * ntkz + NTy, NTz = py.NT, pz.NT + else: + # 2D vector-radix y/z (see YZ2D): pass B = 2D DFT (A x A) per digit group times the + # merged twiddle W_N2^(b_y k1y) W_N2^(b_z k1z); pass C = the remaining digits along y + # then z on one (y, z, column) tile. Both passes run in place; pass C ends in + # natural order. + SB, NTB = yz2d["SB"], yz2d["NTB"] + SC, NTC = yz2d["SC"], yz2d["NTC"] + kyb, gy = _make_passB(yzd, ny, M, SB, NTB, yz2d.get("softB"), tw2d_offset(nz, ny, M, yzd), PDL) + kzc, gz = _make_passC( + yzd, ny, M, SC, NTC, yz2d.get("softC"), twnc_offset(nz, ny, M), PDL, yz2d.get("UC", 0), POLZ + ) + NTy, NTz = NTB, NTC + + @cute.jit + def launch(X: cute.Tensor, O: cute.Tensor, TW: cute.Tensor, stream: cuda.CUstream): + if cutlass.const_expr(XP is not None): + kxp(X, O, TW).launch( + grid=[GX, 1, 1], + block=[NTx, 1, 1], + stream=stream, + min_blocks_per_mp=XP["cps"], + preferred_smem_carveout=100, + ) + else: + kxa(X, O, TW).launch(grid=[ntx, 1, 1], block=[NTx, 1, 1], stream=stream) + kyb(O, TW).launch(grid=[gy, 1, 1], block=[NTy, 1, 1], stream=stream, use_pdl=PDL) + kzc(O, TW).launch(grid=[gz, 1, 1], block=[NTz, 1, 1], stream=stream, use_pdl=PDL) + + return launch + + +class YZ2D: + """Digit structure of the 2D (y, z) transform for ny == nz == n. + + n = N1 * N2 with N1 a coprime prime-factor part (N1 = 1: none) and N2 = A * Bc split + Cooley-Tukey: natural j -> (j1, j2) by the Stages PFA convention, j2 = Bc*a + b. + pass B: for each group (j1, b): DFT_A over a, then twiddle W_N2^(b*k1); + pass C: for each k1: DFT_Bc over b (radices rc, DIT in shared memory) and DFT_N1 over j1. + Storage between the passes: element (j1, b, a | k1) at A*(j1*Bc + b) + (a | k1), so pass B + groups are contiguous runs of A and pass C groups are the rows congruent to k1 mod A. + """ + + def __init__(self, n, A, rc, N1=1): + self.n, self.A, self.rc, self.N1 = n, A, list(rc), N1 + self.Bc = math.prod(self.rc) + self.N2 = A * self.Bc + assert N1 * self.N2 == n and math.gcd(N1, self.N2) == 1 and self.Bc >= 2 + self.nc = N1 * self.Bc + self.G = n // A # pass-B groups per dimension + if N1 > 1: + self.u1 = pow(self.N2 % N1, -1, N1) + self.u2 = pow(N1 % self.N2, -1, self.N2) + self.v1 = pow(self.Bc % N1, -1, N1) + self.v2 = pow(N1 % self.Bc, -1, self.Bc) + self.ag = (self.N2 * self.u1) % n + self.bg = (N1 * self.u2) % n + self.radC = self.rc + ([N1] if N1 > 1 else []) + + def digits(self, j): + """Natural index -> (j1, j2).""" + if self.N1 == 1: + return j * 0, j + return ((j % self.N1) * self.u1) % self.N1, ((j % self.N2) * self.u2) % self.N2 + + def p_in(self, j): + """Natural index -> storage position before pass B.""" + j1, j2 = self.digits(j) + a = j2 // self.Bc + b = j2 % self.Bc + return self.A * (j1 * self.Bc + b) + a + + def c_in(self, k1, t): + """Pass C sub-input index t (Stages natural order) of group k1 -> storage position.""" + if self.N1 == 1: + return self.A * t + k1 + j1 = ((t % self.N1) * self.v1) % self.N1 + b = ((t % self.Bc) * self.v2) % self.Bc + return self.A * (j1 * self.Bc + b) + k1 + + def c_out(self, k1, u): + """Pass C sub-output index u of group k1 -> natural output index.""" + if self.N1 == 1: + return k1 + self.A * u + q = u % self.N1 + k2 = u % self.Bc + return (q * self.ag + (k1 + self.A * k2) * self.bg) % self.n + + def twiddle_table(self): + """Merged pass-B twiddles W_N2^(b_y k1y + b_z k1z), index ((tz*G + ty)*A + k1y)*A + k1z.""" + G, A, Bc = self.G, self.A, self.Bc + w = twiddles.cis_ratio(2 * (Bc - 1) * (A - 1) + 1, self.N2) + tz, ty, k1y, k1z = torch.meshgrid( + torch.arange(G), torch.arange(G), torch.arange(A), torch.arange(A), indexing="ij" + ) + return w[((ty % Bc) * k1y + (tz % Bc) * k1z).reshape(-1)].reshape(-1) + + def nc_table(self): + return twiddles.cis(self.nc, self.nc).reshape(-1) + + +def twnc_offset(nz, ny, M): + """Offset (complex units) of the pass-C W_nc table inside the twiddle buffer.""" + return nz + ny + M + M + + +def tw2d_offset(nz, ny, M, yzd): + return twnc_offset(nz, ny, M) + yzd.nc + + +def _make_passB(yzd, ny, M, S, NT, soft, tw2off, PDL): + """Pass B: per (group ty, group tz, S columns): 2D DFT_A x DFT_A over the contiguous + runs A*ty + [0, A) (y) and A*tz + [0, A) (z), times the merged twiddle, in place.""" + A, Gd = yzd.A, yzd.G + ncol = (M + S - 1) // S + GA = A * S + rA = (GA + NT - 1) // NT + exA = GA % NT == 0 + soft = SOFT_DEFAULT if soft is None else tuple(soft) + + @cute.kernel + def kB(O: cute.Tensor, TW: cute.Tensor): + _set_soft(soft) + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + alloc = cutlass.utils.SmemAllocator() + sbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * A * A * S, byte_alignment=16) + sm = sbuf.llvm_ptr + op = O.iterator.llvm_ptr + tw = TW.iterator.llvm_ptr + if cutlass.const_expr(PDL): + cute.arch.griddepcontrol_wait() + grp = bid // ncol + c = bid % ncol + ty = grp % Gd + tz = grp // Gd + # stage A: DFT_A along y for each (z digit j2, column s) + locs = [] + xss = [] + for a in cutlass.range_constexpr(rA): + g = tid + NT * a + if cutlass.const_expr(not exA): + g = _imin(g, GA - 1) + sc = g % S + j2 = g // S + col = _imin(c * S + sc, M - 1) + xs = [] + for j1 in cutlass.range_constexpr(A): + re, im = ld2(op, ((A * tz + j2) * ny + A * ty + j1) * M + col) + xs.append(G.Cx(G.Real(re), G.Real(im))) + locs.append((sc, j2)) + xss.append(xs) + yss = [_dft(xs) for xs in xss] + for a in cutlass.range_constexpr(rA): + sc, j2 = locs[a] + outs = [_out(y) for y in yss[a]] + if cutlass.const_expr(exA or (a + 1) * NT <= GA): + for q1 in cutlass.range_constexpr(A): + st2(sm, (q1 * A + j2) * S + sc, outs[q1][0], outs[q1][1]) + else: + if tid + NT * a < GA: + for q1 in cutlass.range_constexpr(A): + st2(sm, (q1 * A + j2) * S + sc, outs[q1][0], outs[q1][1]) + cute.arch.sync_threads() + # stage B: DFT_A along z for each (y digit q1, column s), merged twiddle, store + locs = [] + xss = [] + for a in cutlass.range_constexpr(rA): + g = tid + NT * a + if cutlass.const_expr(not exA): + g = _imin(g, GA - 1) + sc = g % S + q1 = g // S + xs = [] + for j2 in cutlass.range_constexpr(A): + re, im = ld2(sm, (q1 * A + j2) * S + sc) + xs.append(G.Cx(G.Real(re), G.Real(im))) + locs.append((sc, q1)) + xss.append(xs) + yss = [_dft(xs) for xs in xss] + if cutlass.const_expr(PDL): + cute.arch.griddepcontrol_launch_dependents() + for a in cutlass.range_constexpr(rA): + sc, q1 = locs[a] + tb = tw2off + (grp * A + q1) * A + ys = [] + for q2 in cutlass.range_constexpr(A): + wr, wi = ld2(tw, tb + q2) + ys.append(G.cmul_rt(_BK, yss[a][q2], wr, wi)) + outs = [_out(y) for y in ys] + col = c * S + sc + ok = col < M + if cutlass.const_expr(not exA): + ok = ok and (tid + NT * a < GA) + if ok: + for q2 in cutlass.range_constexpr(A): + st2(op, ((A * tz + q2) * ny + A * ty + q1) * M + col, outs[q2][0], outs[q2][1]) + + return kB, Gd * Gd * ncol + + +def _u_of_storage(st, p): + """Natural output index of in-place DIT storage position p after the last stage.""" + if not st.pfa: + return p + return st.out_idx(p % st.N2, p // st.N2) + + +def _make_passC(yzd, ny, M, S, NT, soft, twnc_off, PDL, U, POL): + """Pass C: per (k1y, k1z, S columns): the remaining digits (DIT radices rc, then the + prime-factor DFT_N1) along y and then along z, on a [y][column][z] tile in shared memory. + U > 0: every stage runs as a runtime loop over groups of U codelet rounds.""" + if U > 0: + return _make_passC_loop(yzd, ny, M, S, NT, soft, twnc_off, PDL, U, POL) + A, nc = yzd.A, yzd.nc + rad = yzd.radC + pfa = yzd.N1 > 1 + nst = len(rad) + LZ = nc + 1 if nc % 2 == 0 else nc # odd z-row stride -> conflict-free z stages + pyp = PassPlan(nc, rad, S * nc, NT, False, soft, pfa) + pyp.PS = S * LZ + pyp.SQ, pyp.SS2 = nc, LZ # y sequences (sc, zp) at sc*LZ + zp + pzp = PassPlan(nc, rad, nc * S, NT, False, soft, pfa) + pzp.PS, pzp.SS = 1, LZ + st = pyp.st + R0 = rad[0] + ncol = (M + S - 1) // S + G0 = pyp.G[0] + r0 = pyp.rounds[0] + ex0 = pyp.exact[0] + twtab = nst > 1 and not (pfa and nst == 2) + + @cute.kernel + def kC(O: cute.Tensor, TW: cute.Tensor): + _set_soft(pyp.soft) + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + alloc = cutlass.utils.SmemAllocator() + sbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * nc * S * LZ, byte_alignment=16) + tbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * nc, byte_alignment=16) + sm = sbuf.llvm_ptr + smtw = tbuf.llvm_ptr + op = O.iterator.llvm_ptr + tw = TW.iterator.llvm_ptr + if cutlass.const_expr(twtab): + _copy_tw(smtw, tw, twnc_off, nc, NT, tid) + if cutlass.const_expr(PDL): + cute.arch.griddepcontrol_wait() + pol = mk_policy(POL) if POL else None + grp = bid // ncol + c = bid % ncol + k1y = grp % A + k1z = grp // A + # ---- y stage 0: global -> regs -> smem (z kept in DIT input order) + locs = [] + xss = [] + for a in cutlass.range_constexpr(r0): + g = tid + NT * a + if cutlass.const_expr(not ex0): + g = _imin(g, G0 - 1) + sig = g % (S * nc) + cy = g // (S * nc) + sc = sig % S # columns fastest: coalesced loads + tz = sig // S + col = _imin(c * S + sc, M - 1) + zrow = yzd.c_in(k1z, tz) * ny + rot = st.stage1_rot(cy) + xs = [] + for r in cutlass.range_constexpr(R0): + re, im = ld2(op, (zrow + yzd.c_in(k1y, st.stage1_src(cy, r, rot))) * M + col) + xs.append(G.Cx(G.Real(re), G.Real(im))) + if cutlass.const_expr(nst == 1): + yb = cy * R0 + zp = tz + else: + yb = st.stage1_base(cy) + zp = st.pinv(tz) + locs.append((sc, yb, zp)) + xss.append(xs) + yss = [_dft(xs) for xs in xss] + for a in cutlass.range_constexpr(r0): + sc, yb, zp = locs[a] + outs = [_out(y) for y in yss[a]] + if cutlass.const_expr(ex0 or (a + 1) * NT <= G0): + for q in cutlass.range_constexpr(R0): + st2(sm, ((yb + q) * S + sc) * LZ + zp, outs[q][0], outs[q][1]) + else: + if tid + NT * a < G0: + for q in cutlass.range_constexpr(R0): + st2(sm, ((yb + q) * S + sc) * LZ + zp, outs[q][0], outs[q][1]) + cute.arch.sync_threads() + # ---- remaining y stages, smem -> smem + for i in cutlass.range_constexpr(1, nst): + _mid_stage(sm, smtw, 0, pyp, i, tid) + # ---- z stages; the last one stores to global + if cutlass.const_expr(nst > 1): + _first_stage_from_smem(sm, pzp, tid) + cute.arch.sync_threads() + for i in cutlass.range_constexpr(1, nst - 1): + _mid_stage(sm, smtw, 0, pzp, i, tid) + i = nst - 1 + R = rad[i] + Gl = pzp.G[i] + rl = pzp.rounds[i] + exl = pzp.exact[i] + locs = [] + xss = [] + for a in cutlass.range_constexpr(rl): + g = tid + NT * a + if cutlass.const_expr(not exl): + g = _imin(g, Gl - 1) + if cutlass.const_expr(nst == 1): + sig = g % (nc * S) + k1 = g * 0 + xs = [] + for r in cutlass.range_constexpr(R): + re, im = ld2(sm, r + sig * LZ) + xs.append(G.Cx(G.Real(re), G.Real(im))) + else: + sig, k1, base = _codelet_loc(pzp, i, g) + xs = _load_codelet(sm, smtw, 0, pzp, i, sig, k1, base) + locs.append((sig, k1)) + xss.append(xs) + yss = [_dft(xs) for xs in xss] + for a in cutlass.range_constexpr(rl): + sig, k1 = locs[a] + outs = [_out(y) for y in yss[a]] + yp = sig // S + sc = sig % S + col = c * S + sc + yrow = yzd.c_out(k1y, _u_of_storage(st, yp)) + ok = col < M + if cutlass.const_expr(not exl): + ok = ok and (tid + NT * a < Gl) + if ok: + for q in cutlass.range_constexpr(R): + if cutlass.const_expr(nst == 1): + uz = q + else: + uz = st.out_idx(k1, q) + gst(op, (yzd.c_out(k1z, uz) * ny + yrow) * M + col, outs[q][0], outs[q][1], pol) + + return kC, A * A * ncol + + +def _make_passC_loop(yzd, ny, M, S, NT, soft, twnc_off, PDL, U, POL): + """_make_passC with runtime round loops (see _mid_stage_d).""" + A, nc = yzd.A, yzd.nc + rad = yzd.radC + pfa = yzd.N1 > 1 + nst = len(rad) + LZ = nc + 1 if nc % 2 == 0 else nc # odd z-row stride -> conflict-free z stages + pyp = PassPlan(nc, rad, S * nc, NT, False, soft, pfa) + pyp.PS = S * LZ + pyp.SQ, pyp.SS2 = nc, LZ # y sequences (sc, zp) at sc*LZ + zp + pzp = PassPlan(nc, rad, nc * S, NT, False, soft, pfa) + pzp.PS, pzp.SS = 1, LZ + st = pyp.st + R0 = rad[0] + ncol = (M + S - 1) // S + G0 = pyp.G[0] + twtab = nst > 1 and not (pfa and nst == 2) + iL = nst - 1 + RL = rad[iL] + GL = pzp.G[iL] + sp0 = _round_split(G0, NT, U) + spL = _round_split(GL, NT, U) + + @cute.jit + def y0_rounds(sm, op, c, k1y, k1z, tid, abase, cnt: cutlass.Constexpr, masked: cutlass.Constexpr): + locs = [] + xss = [] + oks = [] + for j in cutlass.range_constexpr(cnt): + g = tid + NT * (abase + j) + if cutlass.const_expr(masked): + oks.append(g < G0) + g = _imin(g, G0 - 1) + sig = g % (S * nc) + cy = g // (S * nc) + sc = sig % S # columns fastest: coalesced loads + tz = sig // S + col = _imin(c * S + sc, M - 1) + zrow = yzd.c_in(k1z, tz) * ny + rot = st.stage1_rot(cy) + xs = [] + for r in cutlass.range_constexpr(R0): + re, im = ld2(op, (zrow + yzd.c_in(k1y, st.stage1_src(cy, r, rot))) * M + col) + xs.append(G.Cx(G.Real(re), G.Real(im))) + if cutlass.const_expr(nst == 1): + yb = cy * R0 + zp = tz + else: + yb = st.stage1_base(cy) + zp = st.pinv(tz) + locs.append((sc, yb, zp)) + xss.append(xs) + yss = [_dft(xs) for xs in xss] + for j in cutlass.range_constexpr(cnt): + sc, yb, zp = locs[j] + outs = [_out(y) for y in yss[j]] + if cutlass.const_expr(masked): + if oks[j]: + for q in cutlass.range_constexpr(R0): + st2(sm, ((yb + q) * S + sc) * LZ + zp, outs[q][0], outs[q][1]) + else: + for q in cutlass.range_constexpr(R0): + st2(sm, ((yb + q) * S + sc) * LZ + zp, outs[q][0], outs[q][1]) + + @cute.jit + def zl_rounds(sm, smtw, op, c, k1y, k1z, tid, abase, cnt: cutlass.Constexpr, masked: cutlass.Constexpr, pol): + locs = [] + xss = [] + oks = [] + for j in cutlass.range_constexpr(cnt): + g = tid + NT * (abase + j) + if cutlass.const_expr(masked): + oks.append(g < GL) + g = _imin(g, GL - 1) + if cutlass.const_expr(nst == 1): + sig = g % (nc * S) + k1 = g * 0 + xs = [] + for r in cutlass.range_constexpr(RL): + re, im = ld2(sm, r + sig * LZ) + xs.append(G.Cx(G.Real(re), G.Real(im))) + else: + sig, k1, base = _codelet_loc(pzp, iL, g) + xs = _load_codelet(sm, smtw, 0, pzp, iL, sig, k1, base) + locs.append((sig, k1)) + xss.append(xs) + yss = [_dft(xs) for xs in xss] + for j in cutlass.range_constexpr(cnt): + sig, k1 = locs[j] + outs = [_out(y) for y in yss[j]] + yp = sig // S + sc = sig % S + col = c * S + sc + yrow = yzd.c_out(k1y, _u_of_storage(st, yp)) + ok = col < M + if cutlass.const_expr(masked): + ok = ok and oks[j] + if ok: + for q in cutlass.range_constexpr(RL): + if cutlass.const_expr(nst == 1): + uz = q + else: + uz = st.out_idx(k1, q) + gst(op, (yzd.c_out(k1z, uz) * ny + yrow) * M + col, outs[q][0], outs[q][1], pol) + + @cute.kernel + def kC(O: cute.Tensor, TW: cute.Tensor): + _set_soft(pyp.soft) + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + alloc = cutlass.utils.SmemAllocator() + sbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * nc * S * LZ, byte_alignment=16) + tbuf = alloc.allocate_array(element_type=cutlass.Float64, num_elems=2 * nc, byte_alignment=16) + sm = sbuf.llvm_ptr + smtw = tbuf.llvm_ptr + op = O.iterator.llvm_ptr + tw = TW.iterator.llvm_ptr + if cutlass.const_expr(twtab): + _copy_tw(smtw, tw, twnc_off, nc, NT, tid) + if cutlass.const_expr(PDL): + cute.arch.griddepcontrol_wait() + pol = mk_policy(POL) if POL else None + grp = bid // ncol + c = bid % ncol + k1y = grp % A + k1z = grp // A + # ---- y stage 0: global -> regs -> smem (z kept in DIT input order) + nfull, nit, nrem, tail = sp0 + if cutlass.const_expr(nit > 1): + for it in cutlass.range(nit, unroll=1): + y0_rounds(sm, op, c, k1y, k1z, tid, it * U, U, False) + elif cutlass.const_expr(nit == 1): + y0_rounds(sm, op, c, k1y, k1z, tid, 0, U, False) + if cutlass.const_expr(nrem > 0): + y0_rounds(sm, op, c, k1y, k1z, tid, nit * U, nrem, False) + if cutlass.const_expr(tail): + if _warp_live(tid, NT, nfull, G0): + y0_rounds(sm, op, c, k1y, k1z, tid, nfull, 1, True) + cute.arch.sync_threads() + # ---- remaining y stages, smem -> smem + for i in cutlass.range_constexpr(1, nst): + _mid_stage_d(sm, smtw, 0, pyp, i, tid, U) + # ---- z stages; the last one stores to global + if cutlass.const_expr(nst > 1): + _first_stage_d(sm, pzp, tid, U) + cute.arch.sync_threads() + for i in cutlass.range_constexpr(1, nst - 1): + _mid_stage_d(sm, smtw, 0, pzp, i, tid, U) + nfullL, nitL, nremL, tailL = spL + if cutlass.const_expr(nitL > 1): + for it in cutlass.range(nitL, unroll=1): + zl_rounds(sm, smtw, op, c, k1y, k1z, tid, it * U, U, False, pol) + elif cutlass.const_expr(nitL == 1): + zl_rounds(sm, smtw, op, c, k1y, k1z, tid, 0, U, False, pol) + if cutlass.const_expr(nremL > 0): + zl_rounds(sm, smtw, op, c, k1y, k1z, tid, nitL * U, nremL, False, pol) + if cutlass.const_expr(tailL): + if _warp_live(tid, NT, nfullL, GL): + zl_rounds(sm, smtw, op, c, k1y, k1z, tid, nfullL, 1, True, pol) + + return kC, A * A * ncol + + +# --------------------------------------------------------------------------- +# Generic fallback for extents without an FFT plan (prime factors > 16, or more +# than 3 radices): elementwise pairing, then three direct (O(n) per output) DFT +# passes. Correct for any size. +# --------------------------------------------------------------------------- + + +def _has_plan(n): + return bool(_radix_lists(n)) + + +def make_kernels_direct(nz, ny, M): + nf = M + 1 + total = nz * ny * M + TWZ, TWY, TWX, TWP = 0, nz, nz + ny, nz + ny + M + NT = 128 + + @cute.kernel + def kpre(X: cute.Tensor, W: cute.Tensor, TW: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + e = _imin(bid * NT + tid, total - 1) + k = e % M + p = e // M + ky = p % ny + kz = p // ny + pn = ((nz - kz) % nz) * ny + (ny - ky) % ny + xp = X.iterator.llvm_ptr + wp = W.iterator.llvm_ptr + tw = TW.iterator.llvm_ptr + a0 = ld2(xp, p * nf + k) + b0 = ld2(xp, pn * nf + (M - k)) + c0 = ld2(xp, p * nf + M) + d0 = ld2(xp, pn * nf) + w = ld2(tw, TWP + k) + # 1 <= k < M: A = a + conj(b), B = a - conj(b), Z = A + i w B + Ar = a0[0] + b0[0] + Ai = a0[1] - b0[1] + Br = a0[0] - b0[0] + Bi = a0[1] + b0[1] + zr = Ar - w[1] * Br - w[0] * Bi + zi = Ai + w[0] * Br - w[1] * Bi + # k == 0: Z = ((a+conj d) + (c+conj b))/2 + i((a+conj d) - (c+conj b))/2 + pr = a0[0] + d0[0] + pi_ = a0[1] - d0[1] + qr = c0[0] + b0[0] + qi = c0[1] - b0[1] + z0r = ((pr + qr) - (pi_ - qi)) * 0.5 + z0i = ((pi_ + qi) + (pr - qr)) * 0.5 + if k == 0: + zr = z0r + zi = z0i + if bid * NT + tid < total: + st2(wp, e, zr, zi) + + def make_dft(n, tw_off, inner): + """out[a, j, b] = sum_m in[a, m, b] w_n^(j m), element (a, j, b) at (a*n + j)*inner + b.""" + + @cute.kernel + def kd(SRC: cute.Tensor, DST: cute.Tensor, TW: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + e = _imin(bid * NT + tid, total - 1) + b = e % inner + t = e // inner + j = t % n + a = t // n + base = a * n * inner + b + ip = SRC.iterator.llvm_ptr + op = DST.iterator.llvm_ptr + tw = TW.iterator.llvm_ptr + accr = Float64(0.0) + acci = Float64(0.0) + idx = j * 0 + for m in cutlass.range(n): + xr, xi = ld2(ip, base + m * inner) + wr, wi = ld2(tw, tw_off + idx) + accr = accr + xr * wr - xi * wi + acci = acci + xr * wi + xi * wr + idx = idx + j + if idx >= n: + idx = idx - n + if bid * NT + tid < total: + st2(op, base + j * inner, accr, acci) + + return kd + + kdz = make_dft(nz, TWZ, ny * M) + kdy = make_dft(ny, TWY, M) + kdx = make_dft(M, TWX, 1) + grid = (total + NT - 1) // NT + + @cute.jit + def launch(X: cute.Tensor, O: cute.Tensor, TW: cute.Tensor, WS: cute.Tensor, stream: cuda.CUstream): + kpre(X, WS, TW).launch(grid=[grid, 1, 1], block=[NT, 1, 1], stream=stream) + kdz(WS, O, TW).launch(grid=[grid, 1, 1], block=[NT, 1, 1], stream=stream) + kdy(O, WS, TW).launch(grid=[grid, 1, 1], block=[NT, 1, 1], stream=stream) + kdx(WS, O, TW).launch(grid=[grid, 1, 1], block=[NT, 1, 1], stream=stream) + + return launch + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +def make_tables(nz, ny, M): + """Twiddle buffer: W_nz^e, W_ny^e, W_M^e and W_2M^e (cos, sin pairs) at TWZ, TWY, TWX, TWP.""" + return torch.cat([twiddles.cis(nz, nz), twiddles.cis(ny, ny), twiddles.cis(M, M), twiddles.cis(M, 2 * M)]).reshape( + -1 + ) + + +def make_tables_2d(nz, ny, M, yzd): + """make_tables() followed by the pass-C W_nc table and the merged pass-B table.""" + return torch.cat([make_tables(nz, ny, M), yzd.nc_table(), yzd.twiddle_table()]) + + +_PLANS = {} + + +def _stream(): + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + +class _Plan: + def __init__(self, nz, ny, nf, device): + M = nf - 1 + self.shape = (nz, ny, nf) + self.M = M + self.device = device + x0 = torch.empty(nz * ny * nf * 2, dtype=torch.float64, device=device) + o0 = torch.empty(nz * ny * M * 2, dtype=torch.float64, device=device) + self.direct = not (_has_plan(nz) and _has_plan(ny) and _has_plan(M)) + yz = None + if self.direct: + launch = make_kernels_direct(nz, ny, M) + elif (nz, ny, M) in _XFIRST_SHAPES: + yz = _YZ2D.get((nz, ny, M)) + launch = make_kernels_xfirst(nz, ny, M, yz2d=yz) + else: + launch = make_kernels(nz, ny, M, fused=_FUSED_YX.get((nz, ny, M))) + if yz is not None: + tables = make_tables_2d(nz, ny, M, YZ2D(ny, yz["A"], yz["rc"], yz.get("N1", 1))) + else: + tables = make_tables(nz, ny, M) + self.tables = tables.to(device) + self.tw_t = from_dlpack(self.tables, assumed_align=16) + args = [from_dlpack(x0, assumed_align=16), from_dlpack(o0, assumed_align=16), self.tw_t] + if self.direct: + args.append(from_dlpack(o0, assumed_align=16)) # workspace + self.fn = cute.compile(launch, *args, _stream()) + + def bind(self, xin, out): + xt = from_dlpack(xin.view(-1), assumed_align=16) + ot = from_dlpack(out.view(-1), assumed_align=16) + fn, tw = self.fn, self.tw_t + if self.direct: + ws = torch.empty_like(out).view(-1) + wt = from_dlpack(ws, assumed_align=16) + + def go(): + fn(xt, ot, tw, wt, _stream()) + + go.ws = ws + return go + + def go(): + fn(xt, ot, tw, _stream()) + + return go + + def run_into(self, xin, out): + self.bind(xin, out)() + + +def _get_plan(nz, ny, nf, device): + key = (nz, ny, nf, str(device)) + p = _PLANS.get(key) + if p is None: + p = _Plan(nz, ny, nf, device) + _PLANS[key] = p + return p + + +def _check(x): + if not x.is_cuda: + raise ValueError(f"c2r takes a CUDA tensor, not one on {x.device}") + if x.dtype != torch.float64: + raise TypeError(f"this c2r takes float64, not {x.dtype}") + if x.dim() != 4 or x.shape[-1] != 2 or min(x.shape[:2]) < 1 or x.shape[2] < 2: + raise ValueError(f"c2r takes a [nz, ny, nf, 2] half-spectrum with nf >= 2, not {tuple(x.shape)}") + + +def run(input): + """Inverse real 3D FFT of input (float64 [nz, ny, nf, 2]) into a new float64 [nz, ny, 2*(nf-1)].""" + x = input.contiguous() + _check(x) + nz, ny, nf, _ = x.shape + plan = _get_plan(nz, ny, nf, x.device) + out = torch.empty((nz, ny, 2 * (nf - 1)), dtype=torch.float64, device=x.device) + plan.run_into(x, out) + return out + + +def prepare(input): + """Bind a contiguous input once; launch() recomputes from its current contents and returns the output.""" + _check(input) + if not input.is_contiguous(): + raise ValueError("prepare binds the input buffer itself, so it must be contiguous") + nz, ny, nf, _ = input.shape + plan = _get_plan(nz, ny, nf, input.device) + out = torch.empty((nz, ny, 2 * (nf - 1)), dtype=torch.float64, device=input.device) + go = plan.bind(input, out) + + def launch(): + go() + return out + + return launch diff --git a/build/torch-cuda/_c2r_f64/softfp.py b/build/torch-cuda/_c2r_f64/softfp.py new file mode 100644 index 0000000000000000000000000000000000000000..38b4decd4c3e8fe43a96edd2e1e46ce8a53a823f --- /dev/null +++ b/build/torch-cuda/_c2r_f64/softfp.py @@ -0,0 +1,468 @@ +"""Software IEEE-754 binary64 addition (round to nearest even) on the integer pipes. + +On the RTX 5090 the FP64 unit runs at 1/64 of the FP32 rate while the integer ALUs sit idle +during an FFT, and the two issue concurrently, so routing part of the float64 additions +through integer instructions (inline PTX) raises throughput. + +Each operation returns the correctly rounded float64 result, bit-identical to add.rn.f64, +for finite normal or zero operands with a normal or zero result. Subnormal, Inf and NaN +operands, and results that overflow or underflow, are outside the supported domain (FFT +data of any reasonable magnitude never reaches them). A zero result is +0. + +Method: order |x| >= |y|, build 53-bit significands shifted left by 9, align y with a +jammed sticky bit, add or subtract, normalize (one CLZ) to put the leading one at bit 62, +round to nearest even, repack. + + sadd, ssub a + b, a - b + sbfly3/4/5 butterfly (a + b, a - b) sharing ordering, alignment and sticky; + three instruction schedules with identical results + stwominus 2*x - y (exact doubling, then the subtraction), == fma.rn.f64(2, x, -y) +""" + +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import Float64, T, dsl_user_op + +# One sum: $0 = $1 + $2 (all b64 bit patterns). +_PTX_ADD = r"""{ +.reg .b32 ah, al, bh, bl, aa, ba, xh, xl, yh, yl, ex, ey, d, hx, hy, mxh, myh, nd, st, msk, ng, rl, rh, v, c, lzc, lz, nl, lsb, inc, ql, qh, lzs, e2, sgn, yal, yah, t; +.reg .b64 A, B, MX, MY, MYA, L, MYP, R, RN, RR, NG, INC; +.reg .pred psw, ps, phz, pz; +mov.b64 {al, ah}, $1; +mov.b64 {bl, bh}, $2; +and.b32 aa, ah, 0x7FFFFFFF; +and.b32 ba, bh, 0x7FFFFFFF; +mov.b64 A, {al, aa}; +mov.b64 B, {bl, ba}; +setp.gt.u64 psw, B, A; +selp.b32 xh, bh, ah, psw; +selp.b32 xl, bl, al, psw; +selp.b32 yh, ah, bh, psw; +selp.b32 yl, al, bl, psw; +and.b32 ex, xh, 0x7FF00000; +and.b32 ey, yh, 0x7FF00000; +sub.u32 d, ex, ey; +shr.u32 d, d, 20; +min.u32 d, d, 63; +min.u32 hx, ex, 0x00100000; +min.u32 hy, ey, 0x00100000; +lop3.b32 mxh, xh, 0x000FFFFF, hx, 0xEA; +lop3.b32 myh, yh, 0x000FFFFF, hy, 0xEA; +mov.b64 MX, {xl, mxh}; +shl.b64 MX, MX, 9; +mov.b64 MY, {yl, myh}; +shl.b64 MY, MY, 9; +shr.b64 MYA, MY, d; +sub.u32 nd, 64, d; +shl.b64 L, MY, nd; +setp.ne.u64 ps, L, 0; +selp.b32 st, 1, 0, ps; +xor.b32 t, xh, yh; +shr.s32 msk, t, 31; +mov.b64 {yal, yah}, MYA; +lop3.b32 yal, yal, st, msk, 0x56; +xor.b32 yah, yah, msk; +mov.b64 MYP, {yal, yah}; +neg.s32 ng, msk; +cvt.u64.u32 NG, ng; +add.u64 R, MX, MYP; +add.u64 R, R, NG; +mov.b64 {rl, rh}, R; +setp.ne.u32 phz, rh, 0; +selp.b32 v, rh, rl, phz; +clz.b32 c, v; +selp.b32 lzc, -1, 31, phz; +add.u32 lz, c, lzc; +shl.b64 RN, R, lz; +cvt.u32.u64 nl, RN; +bfe.u32 lsb, nl, 10, 1; +add.u32 inc, lsb, 0x1FF; +cvt.u64.u32 INC, inc; +add.u64 RR, RN, INC; +shr.b64 RR, RR, 10; +mov.b64 {ql, qh}, RR; +shl.b32 lzs, lz, 20; +sub.u32 e2, ex, lzs; +and.b32 sgn, xh, 0x80000000; +add.u32 qh, qh, e2; +add.u32 qh, qh, sgn; +setp.eq.u32 pz, lz, 63; +selp.b32 qh, 0, qh, pz; +selp.b32 ql, 0, ql, pz; +mov.b64 $0, {ql, qh}; +}""" + +# One difference: the sum with the sign of the second operand flipped. +_PTX_SUB = _PTX_ADD.replace("mov.b64 {bl, bh}, $2;", "mov.b64 {bl, bh}, $2;\nxor.b32 bh, bh, 0x80000000;") + +# 2*x - y: x doubled by an exponent increment (exact), then the difference. +_PTX_TWOMINUS = _PTX_ADD.replace( + "mov.b64 {al, ah}, $1;", + """mov.b64 {al, ah}, $1; +{ .reg .b32 e2x; and.b32 e2x, ah, 0x7FF00000; min.u32 e2x, e2x, 0x00100000; add.u32 ah, ah, e2x; }""", +).replace("mov.b64 {bl, bh}, $2;", "mov.b64 {bl, bh}, $2;\nxor.b32 bh, bh, 0x80000000;") + +# Butterfly ($0, $1) = ($2 + $3, $2 - $3). With x the operand of larger magnitude, +# S = MX + MY and D = MX - MY are both formed; a+b takes S iff the signs agree and a-b +# takes the other one. S needs at most a 1-bit normalization (no CLZ). The exponent +# difference is taken from the operands directly (parallel to the magnitude swap), the +# sticky bit is jammed by a predicated or, S keeps sign(a), D gets sign(a) ^ swapped, and +# zero results are folded into two selects. +_PTX_BFLY3 = r"""{ +.reg .b32 ah, al, bh, bl, aa, ba, ea, eb, d, nd, xh, xl, yh, yl, hx, hy, mxh, myh, yal, yah, exf, sa, sd, sh, sl, t, sn, snl, slsb, sinc, ph, pl, dh, dl, v, c, lzb, lz, dnl, dlsb, dinc, qh, ql, zh, zl, rh, rl; +.reg .b64 A, B, MX, MY, MYA, L, S, D, SN, SR, DN, DR, I1, I2; +.reg .pred sw, st, dhz, dz, xz, eq; +mov.b64 {al, ah}, $2; +mov.b64 {bl, bh}, $3; +and.b32 aa, ah, 0x7FFFFFFF; +and.b32 ba, bh, 0x7FFFFFFF; +mov.b64 A, {al, aa}; +mov.b64 B, {bl, ba}; +setp.gt.u64 sw, B, A; +shr.u32 ea, aa, 20; +shr.u32 eb, ba, 20; +sub.s32 d, ea, eb; +abs.s32 d, d; +min.u32 d, d, 63; +selp.b32 xh, ba, aa, sw; +selp.b32 xl, bl, al, sw; +selp.b32 yh, aa, ba, sw; +selp.b32 yl, al, bl, sw; +and.b32 exf, xh, 0x7FF00000; +and.b32 hy, yh, 0x7FF00000; +min.u32 hx, exf, 0x00100000; +min.u32 hy, hy, 0x00100000; +lop3.b32 mxh, xh, 0x000FFFFF, hx, 0xEA; +lop3.b32 myh, yh, 0x000FFFFF, hy, 0xEA; +mov.b64 MX, {xl, mxh}; +shl.b64 MX, MX, 9; +mov.b64 MY, {yl, myh}; +shl.b64 MY, MY, 9; +shr.b64 MYA, MY, d; +sub.u32 nd, 64, d; +shl.b64 L, MY, nd; +setp.ne.u64 st, L, 0; +mov.b64 {yal, yah}, MYA; +@st or.b32 yal, yal, 1; +mov.b64 MYA, {yal, yah}; +add.u64 S, MX, MYA; +sub.u64 D, MX, MYA; +// ---- S: leading one at bit 62 or 61 -> to 62, round at bit 10 (RNE) +mov.b64 {sl, sh}, S; +shr.u32 t, sh, 30; +xor.b32 sn, t, 1; +shl.b64 SN, S, sn; +cvt.u32.u64 snl, SN; +bfe.u32 slsb, snl, 10, 1; +add.u32 sinc, slsb, 0x1FF; +cvt.u64.u32 I1, sinc; +add.u64 SR, SN, I1; +shr.b64 SR, SR, 10; +mov.b64 {pl, ph}, SR; +mad.lo.u32 ph, sn, 0xFFF00000, ph; +add.u32 ph, ph, exf; +// ---- D: general normalization to bit 62 +mov.b64 {dl, dh}, D; +setp.ne.u32 dhz, dh, 0; +selp.b32 v, dh, dl, dhz; +clz.b32 c, v; +selp.b32 lzb, -1, 31, dhz; +add.u32 lz, c, lzb; +shl.b64 DN, D, lz; +cvt.u32.u64 dnl, DN; +bfe.u32 dlsb, dnl, 10, 1; +add.u32 dinc, dlsb, 0x1FF; +cvt.u64.u32 I2, dinc; +add.u64 DR, DN, I2; +shr.b64 DR, DR, 10; +mov.b64 {ql, qh}, DR; +mad.lo.u32 qh, lz, 0xFFF00000, qh; +add.u32 qh, qh, exf; +// ---- signs: S takes sign(a); D takes sign(a) ^ swapped; zeros -> +0 +and.b32 sa, ah, 0x80000000; +selp.b32 sd, 0x80000000, 0, sw; +xor.b32 sd, sd, sa; +setp.eq.u32 xz, exf, 0; +setp.eq.u32 dz, lz, 63; +or.b32 ph, ph, sa; +or.b32 qh, qh, sd; +selp.b32 ph, 0, ph, xz; +selp.b32 pl, 0, pl, xz; +selp.b32 qh, 0, qh, dz; +// a+b = eq ? S : D ; a-b = eq ? D : S (eq: same signs) +xor.b32 t, ah, bh; +setp.ge.s32 eq, t, 0; +selp.b32 zh, ph, qh, eq; +selp.b32 zl, pl, ql, eq; +selp.b32 rh, qh, ph, eq; +selp.b32 rl, ql, pl, eq; +mov.b64 $0, {zl, zh}; +mov.b64 $1, {rl, rh}; +}""" + +# Butterfly, v3 with fewer instructions: sticky from a round trip of the aligned operand, +# rounding increment from a predicate (no bit extract), a single zero select per output. +_PTX_BFLY4 = r"""{ +.reg .b32 ah, al, bh, bl, aa, ba, ea, eb, d, xh, xl, yh, yl, hx, hy, exf, mxh, myh, yal, yah, sa, sd, sh, sl, t, sn, snl, sinc, ph, pl, dh, dl, v, c, lzb, lz, dnl, dinc, qh, ql, zh, zl, rh, rl; +.reg .b64 A, B, MX, MY, MYA, L, S, D, SN, SR, DN, DR, I1, I2; +.reg .pred sw, st, dhz, dz, xz, eq, sl1, dl1; +mov.b64 {al, ah}, $2; +mov.b64 {bl, bh}, $3; +and.b32 aa, ah, 0x7FFFFFFF; +and.b32 ba, bh, 0x7FFFFFFF; +mov.b64 A, {al, aa}; +mov.b64 B, {bl, ba}; +setp.gt.u64 sw, B, A; +shr.u32 ea, aa, 20; +shr.u32 eb, ba, 20; +sub.s32 d, ea, eb; +abs.s32 d, d; +min.u32 d, d, 63; +selp.b32 xh, ba, aa, sw; +selp.b32 xl, bl, al, sw; +selp.b32 yh, aa, ba, sw; +selp.b32 yl, al, bl, sw; +and.b32 exf, xh, 0x7FF00000; +and.b32 hy, yh, 0x7FF00000; +min.u32 hx, exf, 0x00100000; +min.u32 hy, hy, 0x00100000; +lop3.b32 mxh, xh, 0x000FFFFF, hx, 0xEA; +lop3.b32 myh, yh, 0x000FFFFF, hy, 0xEA; +mov.b64 MX, {xl, mxh}; +shl.b64 MX, MX, 9; +mov.b64 MY, {yl, myh}; +shl.b64 MY, MY, 9; +shr.b64 MYA, MY, d; +shl.b64 L, MYA, d; +setp.ne.u64 st, L, MY; +mov.b64 {yal, yah}, MYA; +@st or.b32 yal, yal, 1; +mov.b64 MYA, {yal, yah}; +add.u64 S, MX, MYA; +sub.u64 D, MX, MYA; +// ---- S: leading one at bit 62 or 61 -> to 62, round at bit 10 (RNE: + 0x1FF + lsb) +mov.b64 {sl, sh}, S; +shr.u32 t, sh, 30; +xor.b32 sn, t, 1; +shl.b64 SN, S, sn; +cvt.u32.u64 snl, SN; +and.b32 t, snl, 0x400; +setp.ne.u32 sl1, t, 0; +selp.b32 sinc, 0x200, 0x1FF, sl1; +cvt.u64.u32 I1, sinc; +add.u64 SR, SN, I1; +shr.b64 SR, SR, 10; +mov.b64 {pl, ph}, SR; +mad.lo.u32 ph, sn, 0xFFF00000, ph; +add.u32 ph, ph, exf; +// ---- D: general normalization to bit 62 +mov.b64 {dl, dh}, D; +setp.ne.u32 dhz, dh, 0; +selp.b32 v, dh, dl, dhz; +clz.b32 c, v; +selp.b32 lzb, -1, 31, dhz; +add.u32 lz, c, lzb; +shl.b64 DN, D, lz; +cvt.u32.u64 dnl, DN; +and.b32 t, dnl, 0x400; +setp.ne.u32 dl1, t, 0; +selp.b32 dinc, 0x200, 0x1FF, dl1; +cvt.u64.u32 I2, dinc; +add.u64 DR, DN, I2; +shr.b64 DR, DR, 10; +mov.b64 {ql, qh}, DR; +mad.lo.u32 qh, lz, 0xFFF00000, qh; +add.u32 qh, qh, exf; +// ---- signs: S takes sign(a); D takes sign(a) ^ swapped; zero results -> +0 +// (a zero S or D has a zero low word already: only the high word needs the select) +and.b32 sa, ah, 0x80000000; +selp.b32 sd, 0x80000000, 0, sw; +xor.b32 sd, sd, sa; +setp.eq.u32 xz, exf, 0; +setp.eq.u32 dz, lz, 63; +or.b32 ph, ph, sa; +or.b32 qh, qh, sd; +selp.b32 ph, 0, ph, xz; +selp.b32 qh, 0, qh, dz; +// a+b = eq ? S : D ; a-b = eq ? D : S (eq: same signs) +xor.b32 t, ah, bh; +setp.ge.s32 eq, t, 0; +selp.b32 zh, ph, qh, eq; +selp.b32 zl, pl, ql, eq; +selp.b32 rh, qh, ph, eq; +selp.b32 rl, ql, pl, eq; +mov.b64 $0, {zl, zh}; +mov.b64 $1, {rl, rh}; +}""" + +# Butterfly, v4 with the v3 sticky (computed in parallel to the alignment shift). +_PTX_BFLY5 = r"""{ +.reg .b32 ah, al, bh, bl, aa, ba, ea, eb, d, xh, xl, yh, yl, hx, hy, exf, mxh, myh, yal, yah, sa, sd, sh, sl, t, sn, snl, sinc, ph, pl, dh, dl, v, c, lzb, lz, dnl, dinc, qh, ql, zh, zl, rh, rl; +.reg .b64 A, B, MX, MY, MYA, L, S, D, SN, SR, DN, DR, I1, I2; +.reg .pred sw, st, dhz, dz, xz, eq, sl1, dl1; +mov.b64 {al, ah}, $2; +mov.b64 {bl, bh}, $3; +and.b32 aa, ah, 0x7FFFFFFF; +and.b32 ba, bh, 0x7FFFFFFF; +mov.b64 A, {al, aa}; +mov.b64 B, {bl, ba}; +setp.gt.u64 sw, B, A; +shr.u32 ea, aa, 20; +shr.u32 eb, ba, 20; +sub.s32 d, ea, eb; +abs.s32 d, d; +min.u32 d, d, 63; +selp.b32 xh, ba, aa, sw; +selp.b32 xl, bl, al, sw; +selp.b32 yh, aa, ba, sw; +selp.b32 yl, al, bl, sw; +and.b32 exf, xh, 0x7FF00000; +and.b32 hy, yh, 0x7FF00000; +min.u32 hx, exf, 0x00100000; +min.u32 hy, hy, 0x00100000; +lop3.b32 mxh, xh, 0x000FFFFF, hx, 0xEA; +lop3.b32 myh, yh, 0x000FFFFF, hy, 0xEA; +mov.b64 MX, {xl, mxh}; +shl.b64 MX, MX, 9; +mov.b64 MY, {yl, myh}; +shl.b64 MY, MY, 9; +shr.b64 MYA, MY, d; +sub.u32 t, 64, d; +shl.b64 L, MY, t; +setp.ne.u64 st, L, 0; +mov.b64 {yal, yah}, MYA; +@st or.b32 yal, yal, 1; +mov.b64 MYA, {yal, yah}; +add.u64 S, MX, MYA; +sub.u64 D, MX, MYA; +// ---- S: leading one at bit 62 or 61 -> to 62, round at bit 10 (RNE: + 0x1FF + lsb) +mov.b64 {sl, sh}, S; +shr.u32 t, sh, 30; +xor.b32 sn, t, 1; +shl.b64 SN, S, sn; +cvt.u32.u64 snl, SN; +and.b32 t, snl, 0x400; +setp.ne.u32 sl1, t, 0; +selp.b32 sinc, 0x200, 0x1FF, sl1; +cvt.u64.u32 I1, sinc; +add.u64 SR, SN, I1; +shr.b64 SR, SR, 10; +mov.b64 {pl, ph}, SR; +mad.lo.u32 ph, sn, 0xFFF00000, ph; +add.u32 ph, ph, exf; +// ---- D: general normalization to bit 62 +mov.b64 {dl, dh}, D; +setp.ne.u32 dhz, dh, 0; +selp.b32 v, dh, dl, dhz; +clz.b32 c, v; +selp.b32 lzb, -1, 31, dhz; +add.u32 lz, c, lzb; +shl.b64 DN, D, lz; +cvt.u32.u64 dnl, DN; +and.b32 t, dnl, 0x400; +setp.ne.u32 dl1, t, 0; +selp.b32 dinc, 0x200, 0x1FF, dl1; +cvt.u64.u32 I2, dinc; +add.u64 DR, DN, I2; +shr.b64 DR, DR, 10; +mov.b64 {ql, qh}, DR; +mad.lo.u32 qh, lz, 0xFFF00000, qh; +add.u32 qh, qh, exf; +// ---- signs: S takes sign(a); D takes sign(a) ^ swapped; zero results -> +0 +// (a zero S or D has a zero low word already: only the high word needs the select) +and.b32 sa, ah, 0x80000000; +selp.b32 sd, 0x80000000, 0, sw; +xor.b32 sd, sd, sa; +setp.eq.u32 xz, exf, 0; +setp.eq.u32 dz, lz, 63; +or.b32 ph, ph, sa; +or.b32 qh, qh, sd; +selp.b32 ph, 0, ph, xz; +selp.b32 qh, 0, qh, dz; +// a+b = eq ? S : D ; a-b = eq ? D : S (eq: same signs) +xor.b32 t, ah, bh; +setp.ge.s32 eq, t, 0; +selp.b32 zh, ph, qh, eq; +selp.b32 zl, pl, ql, eq; +selp.b32 rh, qh, ph, eq; +selp.b32 rl, ql, pl, eq; +mov.b64 $0, {zl, zh}; +mov.b64 $1, {rl, rh}; +}""" + + +def _bits(x, loc, ip): + return llvm.bitcast(T.i64(), Float64(x).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + + +def _f64(r, loc, ip): + return Float64(llvm.bitcast(T.f64(), r, loc=loc, ip=ip)) + + +def _one(ptx, a, b, loc, ip): + """$0 = op($1, $2) on float64 bit patterns.""" + r = llvm.inline_asm( + T.i64(), + [_bits(a, loc, ip), _bits(b, loc, ip)], + ptx, + "=l,l,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return _f64(r, loc, ip) + + +def _two(ptx, a, b, loc, ip): + """($0, $1) = op($2, $3) on float64 bit patterns.""" + res = llvm.inline_asm( + llvm.StructType.get_literal([T.i64(), T.i64()]), + [_bits(a, loc, ip), _bits(b, loc, ip)], + ptx, + "=l,=l,l,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + s = llvm.extractvalue(T.i64(), res, [0], loc=loc, ip=ip) + d = llvm.extractvalue(T.i64(), res, [1], loc=loc, ip=ip) + return _f64(s, loc, ip), _f64(d, loc, ip) + + +@dsl_user_op +def sadd(a, b, *, loc=None, ip=None): + """Software float64 a + b (bit pattern of add.rn.f64).""" + return _one(_PTX_ADD, a, b, loc, ip) + + +@dsl_user_op +def ssub(a, b, *, loc=None, ip=None): + """Software float64 a - b.""" + return _one(_PTX_SUB, a, b, loc, ip) + + +@dsl_user_op +def stwominus(x, y, *, loc=None, ip=None): + """Software float64 2*x - y.""" + return _one(_PTX_TWOMINUS, x, y, loc, ip) + + +@dsl_user_op +def sbfly3(a, b, *, loc=None, ip=None): + """Software float64 butterfly (a + b, a - b).""" + return _two(_PTX_BFLY3, a, b, loc, ip) + + +@dsl_user_op +def sbfly4(a, b, *, loc=None, ip=None): + """Software float64 butterfly (a + b, a - b); same results as sbfly3, fewer instructions.""" + return _two(_PTX_BFLY4, a, b, loc, ip) + + +@dsl_user_op +def sbfly5(a, b, *, loc=None, ip=None): + """Software float64 butterfly (a + b, a - b); sbfly4 with the sbfly3 sticky computation.""" + return _two(_PTX_BFLY5, a, b, loc, ip) diff --git a/build/torch-cuda/_c2r_f64/twiddles.py b/build/torch-cuda/_c2r_f64/twiddles.py new file mode 100644 index 0000000000000000000000000000000000000000..7c9b9152ba67d5b21bd543cfd809f7cab1bb931d --- /dev/null +++ b/build/torch-cuda/_c2r_f64/twiddles.py @@ -0,0 +1,88 @@ +"""Host-side twiddle tables: float64 (cos, sin) pairs, built once per plan. + +The values are those of an x87 long double construction (the one the kernel was tuned and +validated with): the angle 2*pi*e/d is evaluated from 2*pi rounded to float64 with a +rounding to a 64-bit significand after every operation (emulated exactly with fractions), +its cosine and sine come from the C library's cosl/sinl (called through ctypes on the +80-bit value), and the results are rounded to float64. Where the C long double is not the +x87 extended format, cos/sin of the float64-rounded angle come from the math module +instead (absolute differences of a few 1e-16, not bit-identical). +""" + +import ctypes +import functools +import math +import platform +from fractions import Fraction + +import torch + +_TWO_PI = Fraction(2 * math.pi) # 2*pi rounded to float64, exactly + + +def _log2_floor(q): + n, d = q.numerator, q.denominator + k = n.bit_length() - d.bit_length() + if (n << max(-k, 0)) < (d << max(k, 0)): + k -= 1 + return k + + +def _x87(q): + """q >= 0 rounded to the x87 extended format (64-bit significand, ties to even).""" + if q == 0: + return q + scale = Fraction(2) ** (63 - _log2_floor(q)) + t = q * scale # in [2^63, 2^64) + m, r = divmod(t.numerator, t.denominator) + if 2 * r > t.denominator or (2 * r == t.denominator and m & 1): + m += 1 + return m / scale + + +@functools.cache +def _libm(): + """The C library with cosl/sinl taking and returning x87 long doubles, or None.""" + if platform.machine().lower() not in ("x86_64", "amd64", "i386", "i686") or ctypes.sizeof(ctypes.c_longdouble) < 10: + return None + try: + lib = ctypes.CDLL("libm.so.6") + except OSError: + return None + for f in (lib.cosl, lib.sinl): + f.restype = ctypes.c_longdouble + f.argtypes = [ctypes.c_longdouble] + return lib + + +def _long_double(q): + """x87 extended encoding of q >= 0 (exactly representable, e.g. a result of _x87).""" + if q == 0: + return ctypes.c_longdouble(0.0) + k = _log2_floor(q) + m = q * Fraction(2) ** (63 - k) + assert m.denominator == 1 + raw = int(m).to_bytes(8, "little") + (k + 16383).to_bytes(2, "little") + return ctypes.c_longdouble.from_buffer_copy(raw.ljust(ctypes.sizeof(ctypes.c_longdouble), b"\0")) + + +def _cos_sin(ang): + """(cos, sin) of the extended-precision angle ang, rounded to float64.""" + lib = _libm() + if lib is None: + a = float(ang) + return math.cos(a), math.sin(a) + x = _long_double(ang) + return lib.cosl(x), lib.sinl(x) + + +def cis(n, denom): + """[n, 2] float64: (cos, sin) of 2*pi*e/denom for e = 0..n-1, evaluated as ((2*pi) * e) / denom.""" + vals = [_cos_sin(_x87(_x87(_TWO_PI * e) / denom)) for e in range(n)] + return torch.tensor(vals, dtype=torch.float64).reshape(n, 2) + + +def cis_ratio(n, denom): + """[n, 2] float64: (cos, sin) of 2*pi*e/denom for e = 0..n-1, evaluated as (2*pi) * (e / denom).""" + vals = [_cos_sin(_x87(_TWO_PI * _x87(Fraction(e, denom)))) for e in range(n)] + return torch.tensor(vals, dtype=torch.float64).reshape(n, 2) diff --git a/build/torch-cuda/_ops.py b/build/torch-cuda/_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..275435b912ae688c2e8f5ff0d75abfb6b2a7993d --- /dev/null +++ b/build/torch-cuda/_ops.py @@ -0,0 +1,44 @@ +import torch + +def get_backend() -> str: + """Detect the backend by inspecting torch.""" + import torch + + if hasattr(torch.backends, "tpu"): + # torch_tpu sets torch.backends.tpu when it is imported (via + # torch's device-backend autoload), regardless of whether TPU + # hardware is present — analogous to torch.version.cuda being + # set on CUDA builds without a GPU. + return "tpu" + elif hasattr(torch, "neuron"): + # Needs to be sorted before specific Torch builds, since Neuron + # extension can be loaded into e.g. CUDA Torch builds. + return "neuron" + elif torch.version.cuda is not None: + return "cuda" + elif torch.version.hip is not None: + return "rocm" + elif torch.backends.mps.is_available(): + return "metal" + elif hasattr(torch.version, "xpu") and torch.version.xpu is not None: + return "xpu" + else: + return "cpu" + + +def _find_ops_name() -> str: + kernel_name = "fft3d" + unique_id = "1ec826fa776c" + backend = get_backend() + return f"_{kernel_name}_{backend}_{unique_id}" + + +_OPS_NAME = _find_ops_name() + +ops = getattr(torch.ops, _OPS_NAME) + +def add_op_namespace_prefix(op_name: str) -> str: + """ + Prefix op by namespace. + """ + return f"{_OPS_NAME}::{op_name}" \ No newline at end of file diff --git a/build/torch-cuda/_r2c_f32/__init__.py b/build/torch-cuda/_r2c_f32/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..da73cecdd121ff7137ad262c0288d86adad00dd8 --- /dev/null +++ b/build/torch-cuda/_r2c_f32/__init__.py @@ -0,0 +1,5 @@ +"""fp32 forward real-to-complex 3-D FFT (torch.fft.rfftn semantics); see kernel.py.""" + +from .kernel import prepare, run + +__all__ = ["prepare", "run"] diff --git a/build/torch-cuda/_r2c_f32/fft3k.py b/build/torch-cuda/_r2c_f32/fft3k.py new file mode 100644 index 0000000000000000000000000000000000000000..5320f66bd805e0510039e299594b57c4057113d7 --- /dev/null +++ b/build/torch-cuda/_r2c_f32/fft3k.py @@ -0,0 +1,562 @@ +"""Row (x) and column (y / z) passes; data moves between passes through the +output buffer (L2 resident). + +x pass : real rows -> half spectrum. A row of nx reals is an M = nx/2-point + complex FFT; the last Stockham pass pairs butterflies (j, nb - j) so + each thread owns both Z[k] and Z[M-k] and finishes the real-to-complex + split in registers. Rows move through shared memory with coalesced + 16-byte copies. With pack=True column 0 stores X[0] + i X[nx/2] + (both real); x_task_fn is also the x stage of the plane2, ysplit and + xy kernels. +y pass : FFT along y on (z, kx-tile) tiles, in place; the tile holding the + packed column separates it into kx = 0 (column 0) and kx = nx/2 + (column M) using Hermitian symmetry in ky. +z pass : FFT along z, in place, on flat tiles of W consecutive (ky, kx) groups + (one contiguous W * 8-byte segment per z), all nx/2 + 1 columns. +""" + +import math + +import cutlass +from cutlass import Float32, Int32, Int64, cute + +from .fft_core import as_f32 +from .plan import ( + SMEM_LIMIT, + Layout, + choose_layout, + choose_line_plan, + conflict_cost, + lanes_groups, + line_plan_cost, + line_plans, +) +from .ptx import ldg2, ldg4, lds2, lds4, stg2, stg2_if, stg4, stg4_if, sts2, sts2_if, sts4, sts4_if +from .stockham import ( + TwSmem, + access_offsets, + and_pv, + b2i, + csel, + emit_fft, + fsel, + idiv, + imod, + lay_off, + pass1_indices, + pass1_static, +) + +# ----------------------------------------------------------------------------- x pass + + +def choose_x_plan(M, emax=32, t_target=8): + """Plan for the x pass: last pass uses mirror pairing (per == 2, T = nb/2). + Earlier passes may leave threads idle.""" + best = None + for pl in line_plans(M, emax=emax): + lp = pl.passes[-1] + if lp.nb % 2 or lp.nb // 2 != pl.T or lp.per != 2: + continue + c = line_plan_cost(pl) + c *= 1.0 + 0.35 * abs(math.log2(max(pl.T, 1) / t_target)) + if best is None or c < best[0]: + best = (c, pl) + if best is None: + raise ValueError(f"no x plan for M={M}") + return best[1] + + +class XCfg: + """Row pass on blocks of `rows` rows. A CTA's input rows and output rows + (each a contiguous block of global memory) move through shared memory with + 16-byte vector copies, so global traffic is fully coalesced.""" + + def __init__(self, nz, ny, nx, rows=None, plan=None, nt_target=128, t_target=8): + assert nx % 4 == 0 + self.nz, self.ny, self.nx = nz, ny, nx + self.M = nx // 2 + self.nf = self.M + 1 + self.plan = plan or choose_x_plan(self.M, t_target=t_target) + T = self.plan.T + self.kinds = ["cyclic"] * (len(self.plan.passes) - 1) + ["mirror"] + nrows = nz * ny + if rows is None: + rows = max(2, nt_target // T) + if rows % 2: + rows += 1 + self.rows = rows + self.nrows = nrows + self.nt = rows * T + self.nt_r = -(-self.nt // 32) * 32 + self.grid = -(-nrows // rows) + M = self.M + + def acc(lay): + def f(tid): + if tid >= self.nt: + return None + r, t = tid // T, tid % T + return access_offsets(self.plan, self.kinds, t, lambda g: r, lay, 1) + + return lanes_groups(self.nt, f) + + if len(self.plan.passes) > 1: + self.lay = choose_layout(M, rows, acc) + xbytes = self.lay.size * 8 + else: + self.lay = Layout(M, rows, M, M, 0) + xbytes = 16 + # input staging: row stride (complex units) chosen against pass-1 reads + best = None + for extra in range(16): + ist = M + extra + if ist % 2: + continue + + def f(tid, ist=ist): + if tid >= self.nt: + return None + r, t = tid // T, tid % T + return [r * ist + i for i in pass1_static(self.plan, self.kinds, t)] + + c = conflict_cost(lanes_groups(self.nt, f)) + if best is None or c < best[0]: + best = (c, ist) + self.ist = best[1] + self.smem = max(rows * self.ist * 8, xbytes, rows * self.nf * 8) + self.twoff = -(-self.smem // 16) * 16 + twb = sum(p.Ns * p.R * 8 for p in self.plan.passes if p.Ns > 1) + self.M * 8 + 64 + if self.twoff + twb > SMEM_LIMIT - 1024 or self.nt_r > 1024: + if rows > 2: + self.__init__( + nz, ny, nx, rows=max(2, rows // 2), plan=self.plan, nt_target=nt_target, t_target=t_target + ) + return + raise ValueError(f"x pass M={M} does not fit") + + def __repr__(self): + return ( + f"XCfg(M={self.M}, {self.plan}, rows={self.rows}, nt={self.nt}, grid={self.grid}, " + f"lay={self.lay}, ist={self.ist}, smem={self.smem})" + ) + + +def x_task_fn(cfg, r2c_off, tws, NT, copy_out=True, ysplit=1, pack=True): + """body(tid, bid, sbase, gin, gout): x work for row block `bid` with a CTA of + NT >= cfg.nt threads; twiddle tables `tws` must already be staged. With + copy_out=False the result is left in smem at (row * nf + k) * 8. + + ysplit = S > 1: block bid = z * S + p holds the rows y = p + S * r of plane + z (r < ny / S) -- strided rows for a split y transform (copy_out=False). + pack=False: X[0] and X[nx/2] go to columns 0 and nx/2 (both real) instead of + being packed into column 0.""" + plan, lay, M, nx, nf = cfg.plan, cfg.lay, cfg.M, cfg.nx, cfg.nf + T = plan.T + nb = plan.passes[-1].nb + R = plan.passes[-1].R + rows, nrows = cfg.rows, cfg.nrows + in_f4 = rows * nx // 4 # float4 chunks of a full input block + out_f4 = rows * nf // 2 # float4 chunks of a full output block + tail_rows = nrows % rows + q4 = nx // 4 # float4 chunks per input row + + def copyin_slots(tid, bid): + """(f, pv) per copy-in float4 slot of this thread for row block bid.""" + r0 = bid * rows + nrem = cutlass.min(Int32(nrows) - r0, Int32(rows)) if tail_rows else None + out = [] + for i in range(-(-in_f4 // NT)): + f = tid + i * NT if i else tid + pv = None + if (i + 1) * NT > in_f4: + pv = b2i(f < in_f4) + f = cutlass.min(f, Int32(in_f4 - 1)) + if nrem is not None: + ok = b2i(f < nrem * q4) + pv = ok if pv is None else (pv & ok) + f = cutlass.min(f, nrem * q4 - 1) + out.append((f, pv)) + return out + + def src_addr(gin, bid, f): + """Global address of copy-in float4 chunk f of block bid.""" + if ysplit == 1: + return gin + Int64(bid * rows) * (nx * 4) + Int64(f) * 16 + zz = bid // ysplit + pp = bid - zz * ysplit + rr = idiv(f, q4) + cc = imod(f, q4) + return gin + (Int64(zz * cfg.ny + pp + rr * ysplit) * nx + Int64(cc) * 4) * 4 + + def body(tid, bid, sbase, gin, gout): + r = tid // T + t = tid % T + valid = None + if cfg.nt != NT: + valid = b2i(tid < cfg.nt) + r = cutlass.min(r, Int32(rows - 1)) + t = cutlass.min(t, Int32(T - 1)) + r0 = bid * rows + row = r0 + r + if tail_rows: + ok = b2i(row < nrows) + valid = ok if valid is None else (valid & ok) + blk_out = gout + Int64(r0) * (nf * 8) + nrem = None + if tail_rows: + nrem = cutlass.min(Int32(nrows) - r0, Int32(rows)) + + # ---- coalesced copy-in: global block -> smem rows (stride ist) + slots = copyin_slots(tid, bid) + vals = [ldg4(src_addr(gin, bid, f), ".cs") for (f, _) in slots] + for (f, pv), v in zip(slots, vals): + rr = idiv(f, q4) + cc = imod(f, q4) + a = sbase + rr * (cfg.ist * 8) + cc * 16 + if pv is None: + sts4(a, v) + else: + sts4_if(pv, a, v) + cute.arch.sync_threads() + + def x_load(idx): + return [lds2(sbase + r * (cfg.ist * 8) + idx * 8)] + + def x_addr(g, idx): + return sbase + lay_off(lay, r, idx) * 8 + + def split(za, zb, w): + """(X[k], X[M-k]) from Z[k], Z[M-k] and w = exp(-2 pi i k / nx).""" + ar = (za[0] + zb[0]) * 0.5 + ai = (za[1] - zb[1]) * 0.5 + dr = za[0] - zb[0] + di = za[1] + zb[1] + br = (w[0] * di + w[1] * dr) * 0.5 + bi = (w[1] * di - w[0] * dr) * 0.5 + return (ar + br, ai + bi), (ar - br, bi - ai) + + def wk(k): + return tws.elem(r2c_off, k) + + def put(k, v, extra=None): + v = (as_f32(v[0]), as_f32(v[1])) + a = sbase + (r * nf + k) * 8 + q = valid + if extra is not None: + q = extra if q is None else (q & extra) + if q is None: + sts2(a, v) + else: + sts2_if(q, a, v) + + def x_final(outs): + # slot 0: butterfly t ; slot 1: butterfly nb - t (thread 0: nb/2). + (_, y0, _), (j1, y1, _) = outs + y0 = y0[0] + y1 = y1[0] + is0 = t == 0 + for m in range(R): + mb = R - 1 - m + za = y0[m] + zb = csel(is0, y0[(R - m) % R], y1[mb]) + ka = t + m * nb if m else t + xa, xb = split(za, zb, wk(ka)) + if m == 0 and pack: + va = (xa[0], fsel(is0, xb[0], xa[1])) # packed X[0] + i X[nx/2] + else: + va = xa + if m == 0 and not pack: + put(M, (xb[0], 0.0), b2i(is0)) # X[nx/2] (thread 0 only) + x2, _ = split(y1[mb], y1[m], wk(nb // 2 + mb * nb)) + vb = csel(is0, x2, xb) + kb = j1 + mb * nb if mb else j1 + put(ka, va) + put(kb, vb) + + def xsync(): + cute.arch.sync_threads() + + emit_fft( + plan, cfg.kinds, t, 1, x_load, None, x_addr, tws, xsync, final=x_final, first_hazard=True, last_hazard=True + ) + + if copy_out: + cute.arch.sync_threads() + # ---- coalesced copy-out: smem rows (stride nf) -> global block + for i in range(-(-out_f4 // NT)): + f = tid + i * NT if i else tid + pv = None + if (i + 1) * NT > out_f4: + pv = b2i(f < out_f4) + f = cutlass.min(f, Int32(out_f4 - 1)) + if nrem is not None: + lim = (nrem * nf) // 2 + ok = b2i(f < lim) + pv = ok if pv is None else (pv & ok) + f = cutlass.min(f, lim - 1) + v = lds4(sbase + f * 16) + a = blk_out + Int64(f) * 16 + if pv is None: + stg4(a, v) + else: + stg4_if(pv, a, v) + if nrem is not None: + # odd element count in a partial block: last complex value + last = nrem * nf - 1 + pv = b2i(tid == 0) & b2i((nrem * nf) % 2 == 1) + v = lds2(sbase + last * 8) + stg2_if(pv, blk_out + Int64(last) * 8, v) + + return body + + +def make_x_kernel(cfg, r2c_off): + """x pass kernel (3k strategy); triggers the dependent y kernel early (PDL).""" + tws = TwSmem.for_plans([cfg.plan], extra=[(r2c_off, cfg.M)]) + smem_bytes = cfg.twoff + tws.size + NT = cfg.nt_r + body = x_task_fn(cfg, r2c_off, tws, NT) + + @cute.kernel + def x_kernel(inp: cute.Tensor, out: cute.Tensor, tw: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + smem = cutlass.utils.SmemAllocator() + buf = smem.allocate_array(Float32, num_elems=smem_bytes // 4, byte_alignment=16) + sbase = buf.toint() + cute.arch.griddepcontrol_launch_dependents() + gin = inp.iterator.toint() + tws.stage(sbase + cfg.twoff, tid, NT, tw.iterator.toint()) + if cutlass.const_expr(len(cfg.plan.passes) == 1): + cute.arch.sync_threads() + body(tid, bid, sbase, gin, out.iterator.toint()) + + @cute.jit + def launch(inp: cute.Tensor, out: cute.Tensor, tw: cute.Tensor, stream): + x_kernel(inp, out, tw).launch(grid=(cfg.grid, 1, 1), block=(cfg.nt_r, 1, 1), stream=stream) + + return launch + + +# ----------------------------------------------------------------------------- column passes (y / z) + + +class ColCfg: + """FFT along a strided axis of length n on tiles of W columns. + + axis 'y': tiles (z, kx-tile) over the first M columns; the tile holding + the packed column separates it into columns 0 and M. + axis 'z': flat tiles of W consecutive (ky, kx) groups g = ky * nf + kx -- + one contiguous W * 8-byte segment per z, all nf columns. + """ + + def __init__(self, axis, nz, ny, nx, W=None, nt_target=None): + self.flat = axis == "z" + W0 = W + while True: + self._init(axis, nz, ny, nx, W0, nt_target) + if (self.smem <= SMEM_LIMIT - 1024 and self.nt_r <= 1024) or self.W == 1: + break + W0 = self.W // 2 + if self.smem > SMEM_LIMIT - 1024 or self.nt_r > 1024: + raise ValueError(f"column pass n={self.n} does not fit") + + def _init(self, axis, nz, ny, nx, W, nt_target): + self.axis = axis + self.nz, self.ny, self.nx = nz, ny, nx + self.M = nx // 2 + self.nf = self.M + 1 + self.n = ny if axis == "y" else nz + n = self.n + if W is None: + W = 32 if n >= 192 else 16 + if self.M % W and not self.flat: + # nearest divisor of M (prefer wider tiles: longer contiguous segments) + divs = [d for d in range(1, self.M + 1) if self.M % d == 0 and d <= max(W, 8) * 2] + W = min(divs, key=lambda d: (abs(d - W) if d >= W // 2 else 1000, -d)) + self.W = W + if nt_target is None: + nt_target = 512 if W >= 32 else (256 if W >= 16 else 128) + self.plan = choose_line_plan(n, t_target=max(1, nt_target // W), emax=32) + T = self.plan.T + self.kinds = ["cyclic"] * len(self.plan.passes) + self.nt = W * T + self.nt_r = -(-self.nt // 32) * 32 + if self.flat: + self.NG = ny * self.nf + self.ntiles = -(-self.NG // W) + self.grid = self.ntiles + else: + self.ntiles = self.M // W + self.grid = nz * self.ntiles + + def acc(lay): + def f(tid): + if tid >= self.nt: + return None + c, t = tid % W, tid // W + return access_offsets(self.plan, self.kinds, t, lambda g: c, lay, 1) + + return lanes_groups(self.nt, f) + + if len(self.plan.passes) > 1: + self.lay = choose_layout(n, W, acc) + size = self.lay.size * 8 + else: + self.lay = Layout(n, W, n, n, 0) + size = 16 + self.sepoff = -(-size // 16) * 16 + self.twoff = -(-(self.sepoff + (n * 8 if axis == "y" else 0)) // 16) * 16 + self.smem = self.twoff + sum(p.Ns * p.R * 8 for p in self.plan.passes if p.Ns > 1) + + def __repr__(self): + return ( + f"ColCfg({self.axis}, n={self.n}, W={self.W}, {self.plan}, nt={self.nt}, " + f"grid={self.grid}, lay={self.lay}, smem={self.smem})" + ) + + +class ColTask: + """Per-tile work of a column pass, usable from any CTA of NT >= cfg.nt + threads. Twiddle tables `tws` must already be staged in smem.""" + + def __init__(self, cfg, tws, NT): + self.cfg, self.tws, self.NT = cfg, tws, NT + if cfg.axis == "y": + self.stride = cfg.nf * 8 # bytes between consecutive y + self.outer_stride = cfg.ny * cfg.nf * 8 # per z + else: + self.stride = cfg.ny * cfg.nf * 8 # between consecutive z + + def setup(self, tid): + cfg = self.cfg + W, T = cfg.W, cfg.plan.T + c = tid % W + t = tid // W + valid = None + if cfg.nt != self.NT: + valid = b2i(tid < cfg.nt) + t = cutlass.min(t, Int32(T - 1)) + return c, t, valid + + def tile_info(self, tile, c, valid, gout): + cfg = self.cfg + W = cfg.W + v = valid + if cfg.flat: + g = tile * W + c + if cfg.NG % W: + ok = b2i(g < cfg.NG) + v = ok if v is None else (v & ok) + g = cutlass.min(g, Int32(cfg.NG - 1)) + return Int32(1), gout + Int64(g) * 8, v + tcol = tile % cfg.ntiles + outer = tile // cfg.ntiles + kx = tcol * W + c + colbase = gout + Int64(outer) * self.outer_stride + Int64(kx) * 8 + return tcol, colbase, v + + def prefetch(self, tile, c, t, valid, gout, cop=""): + cfg = self.cfg + _, colbase, _ = self.tile_info(tile, c, valid, gout) + return [ldg2(colbase + Int64(idx) * self.stride, cop) for idx in pass1_indices(cfg.plan, cfg.kinds, t)] + + def compute(self, cur, tile, sbase, c, t, valid, gout): + cfg = self.cfg + lay, stride = cfg.lay, self.stride + tcol, colbase, v = self.tile_info(tile, c, valid, gout) + sep = None + if cfg.axis == "y": + sep = b2i(tcol == 0) & b2i(c == 0) + if v is not None: + sep = sep & v + nsep = v + if sep is not None: + nsep = Int32(1) - sep if v is None else (Int32(1) - sep) & v + + def c_addr(g, idx): + return sbase + lay_off(lay, c, idx) * 8 + + def c_store(idx, vals, pv): + val = (as_f32(vals[0][0]), as_f32(vals[0][1])) + q = and_pv(pv, nsep) + if q is None: + stg2(colbase + Int64(idx) * stride, val) + else: + stg2_if(q, colbase + Int64(idx) * stride, val) + if sep is not None: + sts2_if(and_pv(pv, sep), sbase + cfg.sepoff + idx * 8, val) + + NT = self.NT + + def csync(): + cute.arch.barrier(barrier_id=1, number_of_threads=NT) + + emit_fft(cfg.plan, cfg.kinds, t, 1, None, c_store, c_addr, self.tws, csync, pre=[[x] for x in cur]) + return {"t": t, "sep": sep, "colbase": colbase, "tcol": tcol} + + def separate(self, st, sbase): + """Packed column G = P + iQ with P, Q Hermitian in ky -> col 0, col M.""" + cfg = self.cfg + p = cfg.plan.passes[-1] + T, n, M = cfg.plan.T, cfg.n, cfg.M + sepb = sbase + cfg.sepoff + for b in range(p.per): + j = st["t"] + T * b if b else st["t"] + guard = (not p.full) and b == p.per - 1 + jl = cutlass.min(j, Int32(p.nb - 1)) if guard else j + pv = b2i(j < p.nb) if guard else None + for m in range(p.R): + idx = jl + m * p.nb if m else jl + nidx = Int32(cutlass.select_(idx == 0, Int32(0), Int32(n) - idx)) + f = lds2(sepb + idx * 8) + fp = lds2(sepb + nidx * 8) + a = ((f[0] + fp[0]) * 0.5, (f[1] - fp[1]) * 0.5) + bq = ((f[1] + fp[1]) * 0.5, (fp[0] - f[0]) * 0.5) + q = and_pv(pv, st["sep"]) + addr = st["colbase"] + Int64(idx) * self.stride + stg2_if(q, addr, a) + stg2_if(q, addr + M * 8, bq) + + +def make_col_kernel(cfg, trigger=False): + """One CTA per tile. Waits for the previous kernel (PDL griddepcontrol.wait) + before touching the output; trigger=True also lets the next kernel launch + early.""" + plan = cfg.plan + axis = cfg.axis + tws = TwSmem.for_plans([plan]) + smem_bytes = cfg.twoff + max(tws.size, 16) + grid = cfg.grid + task = ColTask(cfg, tws, cfg.nt_r) + + @cute.kernel + def col_kernel(out: cute.Tensor, tw: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + smem = cutlass.utils.SmemAllocator() + buf = smem.allocate_array(Float32, num_elems=smem_bytes // 4, byte_alignment=16) + sbase = buf.toint() + gout = out.iterator.toint() + if cutlass.const_expr(trigger): + cute.arch.griddepcontrol_launch_dependents() + c, t, valid = task.setup(tid) + cute.arch.griddepcontrol_wait() + tws.stage(sbase + cfg.twoff, tid, cfg.nt_r, tw.iterator.toint()) + cur = task.prefetch(bid, c, t, valid, gout) + if cutlass.const_expr(len(plan.passes) == 1): + cute.arch.sync_threads() + st = task.compute(cur, bid, sbase, c, t, valid, gout) + # trace-time and runtime conditions must stay separate `if`s + if cutlass.const_expr(axis == "y"): # noqa: SIM102 + if st["tcol"] == 0: + cute.arch.sync_threads() + task.separate(st, sbase) + cute.arch.sync_threads() + + @cute.jit + def launch(out: cute.Tensor, tw: cute.Tensor, stream): + col_kernel(out, tw).launch(grid=(grid, 1, 1), block=(cfg.nt_r, 1, 1), stream=stream, use_pdl=True) + + return launch diff --git a/build/torch-cuda/_r2c_f32/fft_core.py b/build/torch-cuda/_r2c_f32/fft_core.py new file mode 100644 index 0000000000000000000000000000000000000000..b3b981df002aac253a86f7228eef27f4256939dc --- /dev/null +++ b/build/torch-cuda/_r2c_f32/fft_core.py @@ -0,0 +1,297 @@ +"""Trace-time FFT building blocks for CuTe-DSL kernels. + +Complex values are Python tuples ``(re, im)`` whose components are either +CuTe-DSL ``Float32`` values or Python floats (compile-time constants, folded +where possible). Every helper here runs at trace time and emits straight-line +fp32 arithmetic; all twiddle constants are computed in float64 and rounded once. +""" + +import math + +from cutlass import Float32 + +# --------------------------------------------------------------------------- +# scalar helpers with constant folding +# --------------------------------------------------------------------------- + + +def _isc(x): + return isinstance(x, (int, float)) + + +def sadd(a, b): + if _isc(a) and _isc(b): + return float(a + b) + if _isc(a) and a == 0.0: + return b + if _isc(b) and b == 0.0: + return a + return a + b + + +def ssub(a, b): + if _isc(a) and _isc(b): + return float(a - b) + if _isc(b) and b == 0.0: + return a + if _isc(a) and a == 0.0: + return -b + return a - b + + +def sneg(a): + if _isc(a): + return -float(a) + return -a + + +def smul(a, c): + """a * c where c is a python constant.""" + if _isc(a): + return float(a * c) + if c == 1.0: + return a + if c == -1.0: + return -a + if c == 0.0: + return 0.0 + return a * Float32(c) + + +def smuladd(a, c, b): + """a * c + b (c constant) -- contracted to an FMA by the backend.""" + if c == 0.0 or (_isc(a) and a == 0.0): + return b + return sadd(smul(a, c), b) + + +def as_f32(x): + if _isc(x): + return Float32(float(x)) + return x + + +# --------------------------------------------------------------------------- +# complex helpers +# --------------------------------------------------------------------------- + + +def cadd(a, b): + return (sadd(a[0], b[0]), sadd(a[1], b[1])) + + +def csub(a, b): + return (ssub(a[0], b[0]), ssub(a[1], b[1])) + + +def cneg(a): + return (sneg(a[0]), sneg(a[1])) + + +def cmul_i(a, s=1): + """Multiply by s*i (s = +1 or -1).""" + if s > 0: + return (sneg(a[1]), a[0]) + return (a[1], sneg(a[0])) + + +def cscale(a, c): + return (smul(a[0], c), smul(a[1], c)) + + +def _snap(v): + """Snap values that are exactly 0, +-1, +-0.5 up to float64 noise.""" + for t in (0.0, 1.0, -1.0, 0.5, -0.5): + if abs(v - t) < 1e-15: + return t + return v + + +def root(n, e, sign=-1): + """exp(sign * 2*pi*i*e/n) as python floats, computed exactly for the + special angles.""" + e %= n + if e == 0: + return (1.0, 0.0) + if (4 * e) % n == 0: + q = (4 * e) // n + return [(1.0, 0.0), (0.0, float(sign)), (-1.0, 0.0), (0.0, float(-sign))][q] + ang = 2.0 * math.pi * e / n + return (_snap(math.cos(ang)), _snap(sign * math.sin(ang))) + + +def cmul_const(a, w): + """a * w for a python-constant complex w.""" + c, s = w + if s == 0.0: + if c == 1.0: + return a + if c == -1.0: + return cneg(a) + return cscale(a, c) + if c == 0.0: + if s == 1.0: + return cmul_i(a, 1) + if s == -1.0: + return cmul_i(a, -1) + return cmul_i(cscale(a, abs(s)), 1 if s > 0 else -1) + if abs(abs(c) - abs(s)) < 1e-15: + # 45-degree multiples: (re*c - im*s, re*s + im*c) with |c| == |s| + m = abs(c) + if c > 0 and s < 0: # (1 - i)/sqrt2 + return (smul(sadd(a[0], a[1]), m), smul(ssub(a[1], a[0]), m)) + if c > 0 and s > 0: # (1 + i)/sqrt2 + return (smul(ssub(a[0], a[1]), m), smul(sadd(a[0], a[1]), m)) + if c < 0 and s < 0: # (-1 - i)/sqrt2 + return (smul(ssub(a[1], a[0]), m), smul(sneg(sadd(a[0], a[1])), m)) + # c < 0, s > 0: (-1 + i)/sqrt2 + return (smul(sneg(sadd(a[0], a[1])), m), smul(ssub(a[0], a[1]), m)) + re = smuladd(a[0], c, smul(a[1], -s)) + im = smuladd(a[0], s, smul(a[1], c)) + return (re, im) + + +def cmul(a, w): + """a * w for runtime complex w (components Float32).""" + re = a[0] * w[0] - a[1] * w[1] + im = a[0] * w[1] + a[1] * w[0] + return (re, im) + + +# --------------------------------------------------------------------------- +# small DFTs (forward: sign = -1) +# --------------------------------------------------------------------------- + +_PRIMES = (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61) + + +def factorize(n): + f = [] + d = 2 + while d * d <= n: + while n % d == 0: + f.append(d) + n //= d + d += 1 + if n > 1: + f.append(n) + return f + + +def _dft_prime(x, sign): + p = len(x) + h = (p - 1) // 2 + a = [cadd(x[j], x[p - j]) for j in range(1, h + 1)] + b = [csub(x[j], x[p - j]) for j in range(1, h + 1)] + out = [None] * p + s0 = x[0] + for v in a: + s0 = cadd(s0, v) + out[0] = s0 + for k in range(1, h + 1): + cre, cim = x[0] + sre, sim = 0.0, 0.0 + for j in range(1, h + 1): + ang = 2.0 * math.pi * ((j * k) % p) / p + c = _snap(math.cos(ang)) + s = _snap(math.sin(ang)) + cre = smuladd(a[j - 1][0], c, cre) + cim = smuladd(a[j - 1][1], c, cim) + sre = smuladd(b[j - 1][0], s, sre) + sim = smuladd(b[j - 1][1], s, sim) + # X[k] = C + sign*i*S, X[p-k] = C - sign*i*S ; i*S = (-sim, sre) + if sign < 0: + out[k] = (sadd(cre, sim), ssub(cim, sre)) + out[p - k] = (ssub(cre, sim), sadd(cim, sre)) + else: + out[k] = (ssub(cre, sim), sadd(cim, sre)) + out[p - k] = (sadd(cre, sim), ssub(cim, sre)) + return out + + +def _dft4(x, sign): + t0 = cadd(x[0], x[2]) + t1 = csub(x[0], x[2]) + t2 = cadd(x[1], x[3]) + t3 = csub(x[1], x[3]) + t3r = cmul_i(t3, sign) # sign*i*t3 + return [cadd(t0, t2), cadd(t1, t3r), csub(t0, t2), csub(t1, t3r)] + + +def _egcd_inv(a, m): + return pow(a, -1, m) + + +def _dft_ct(x, n1, n2, sign): + """Cooley-Tukey n = n1*n2: inner DFT_n1 over stride-n2 subsequences.""" + n = n1 * n2 + y = [] + for j2 in range(n2): + sub = [x[n2 * j1 + j2] for j1 in range(n1)] + yk = dft(sub, sign) + y.append([cmul_const(yk[k1], root(n, j2 * k1, sign)) for k1 in range(n1)]) + out = [None] * n + for k1 in range(n1): + col = [y[j2][k1] for j2 in range(n2)] + zk = dft(col, sign) + for k2 in range(n2): + out[k1 + n1 * k2] = zk[k2] + return out + + +def _dft_pfa(x, n1, n2, sign): + """Good-Thomas prime-factor algorithm, gcd(n1, n2) == 1 (no twiddles).""" + n = n1 * n2 + y = [] + for j2 in range(n2): + sub = [x[(j1 * n2 + j2 * n1) % n] for j1 in range(n1)] + y.append(dft(sub, sign)) + inv2 = _egcd_inv(n2 % n1, n1) if n1 > 1 else 0 + inv1 = _egcd_inv(n1 % n2, n2) if n2 > 1 else 0 + out = [None] * n + for k1 in range(n1): + col = [y[j2][k1] for j2 in range(n2)] + zk = dft(col, sign) + for k2 in range(n2): + k = (k1 * n2 * inv2 + k2 * n1 * inv1) % n + out[k] = zk[k2] + return out + + +def _split(n): + """Choose a factor split (n1, n2, kind) for a composite n.""" + f = factorize(n) + # coprime split preferred (no twiddles) + for p in sorted(set(f)): + pk = 1 + m = n + while m % p == 0: + pk *= p + m //= p + if m > 1: + return pk, m, "pfa" + # prime power: radix-4 first for powers of two + if f[0] == 2: + if n % 4 == 0 and n >= 16: + return 4, n // 4, "ct" + return 2, n // 2, "ct" + return f[0], n // f[0], "ct" + + +def dft(x, sign=-1): + """Unnormalized DFT of a python list of complex values (trace time).""" + n = len(x) + if n == 1: + return [x[0]] + if n == 2: + return [cadd(x[0], x[1]), csub(x[0], x[1])] + if n == 4: + return _dft4(x, sign) + if n == 8: + return _dft_ct(x, 2, 4, sign) + if n in _PRIMES: + return _dft_prime(x, sign) + n1, n2, kind = _split(n) + if kind == "pfa": + return _dft_pfa(x, n1, n2, sign) + return _dft_ct(x, n1, n2, sign) diff --git a/build/torch-cuda/_r2c_f32/fft_generic.py b/build/torch-cuda/_r2c_f32/fft_generic.py new file mode 100644 index 0000000000000000000000000000000000000000..babd451cbaea34d7980d8918097105c694e9c548 --- /dev/null +++ b/build/torch-cuda/_r2c_f32/fft_generic.py @@ -0,0 +1,116 @@ +"""Generic fallback (strategy generic) for shapes the fast paths do not cover +(nx % 4 != 0, lengths without a fitting radix plan, ...): three separable +passes of direct DFTs with float64 accumulation and float64 twiddle tables. + + x pass : real row (nx) -> half spectrum (nx//2 + 1) written to the output + y pass : complex DFT along y, in place on the output (one line per CTA) + z pass : complex DFT along z, in place on the output (one line per CTA) + +Each CTA stages its line in shared memory, so in-place updates are safe. The +twiddle index (j*k) mod n is advanced incrementally. O(n^2) per line, so this +is a correctness path, not a performance path. +""" + +import math + +import cutlass +import torch +from cutlass import Float32, Float64, Int32, cute + +NT = 128 + + +def twiddle_f64(n): + """float64 CPU tensor [cos(a_m), -sin(a_m)] for m in [0, n), a_m = (2 pi m) / n.""" + out = [] + for m in range(n): + a = 2.0 * math.pi * m / n + out += (math.cos(a), -math.sin(a)) + return torch.tensor(out, dtype=torch.float64) + + +def make_generic_x_kernel(nz, ny, nx): + nf = nx // 2 + 1 + rows = nz * ny + + @cute.kernel + def gx_kernel(inp: cute.Tensor, out: cute.Tensor, tw: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + row, _, _ = cute.arch.block_idx() + smem = cutlass.utils.SmemAllocator() + s = smem.allocate_tensor(Float32, cute.make_layout(nx), byte_alignment=16) + for i in range(tid, nx, NT): + s[i] = inp[row * nx + i] + cute.arch.sync_threads() + for k in range(tid, nf, NT): + ar = Float64(0.0) + ai = Float64(0.0) + idx = Int32(0) + for j in range(nx): + v = Float64(s[j]) + ar = ar + v * tw[2 * idx] + ai = ai + v * tw[2 * idx + 1] + idx = idx + k + idx = Int32(cutlass.select_(idx >= nx, idx - nx, idx)) + o = (row * nf + k) * 2 + out[o] = Float32(ar) + out[o + 1] = Float32(ai) + + @cute.jit + def launch(inp: cute.Tensor, out: cute.Tensor, tw: cute.Tensor, stream): + gx_kernel(inp, out, tw).launch(grid=(rows, 1, 1), block=(NT, 1, 1), stream=stream) + + return launch + + +def make_generic_col_kernel(axis, nz, ny, nx): + """Complex DFT along y or z over all nx//2 + 1 columns, in place.""" + nf = nx // 2 + 1 + if axis == "y": + n, outer = ny, nz + else: + n, outer = nz, ny + lines = outer * nf + + @cute.kernel + def gc_kernel(out: cute.Tensor, tw: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + line, _, _ = cute.arch.block_idx() + o = line // nf + kx = line - o * nf + smem = cutlass.utils.SmemAllocator() + s = smem.allocate_tensor(Float32, cute.make_layout(2 * n), byte_alignment=16) + # element j of this line lives at complex index base + j * step + if cutlass.const_expr(axis == "y"): + base = o * ny * nf + kx + step = nf + else: + base = o * nf + kx + step = ny * nf + for j in range(tid, n, NT): + e = (base + j * step) * 2 + s[2 * j] = out[e] + s[2 * j + 1] = out[e + 1] + cute.arch.sync_threads() + for k in range(tid, n, NT): + ar = Float64(0.0) + ai = Float64(0.0) + idx = Int32(0) + for j in range(n): + xr = Float64(s[2 * j]) + xi = Float64(s[2 * j + 1]) + wr = tw[2 * idx] + wi = tw[2 * idx + 1] + ar = ar + xr * wr - xi * wi + ai = ai + xr * wi + xi * wr + idx = idx + k + idx = Int32(cutlass.select_(idx >= n, idx - n, idx)) + e = (base + k * step) * 2 + out[e] = Float32(ar) + out[e + 1] = Float32(ai) + + @cute.jit + def launch(out: cute.Tensor, tw: cute.Tensor, stream): + gc_kernel(out, tw).launch(grid=(lines, 1, 1), block=(NT, 1, 1), stream=stream) + + return launch diff --git a/build/torch-cuda/_r2c_f32/fft_plane2.py b/build/torch-cuda/_r2c_f32/fft_plane2.py new file mode 100644 index 0000000000000000000000000000000000000000..3989f41a5b5737cfabed7a4059dc37fe20876c9f --- /dev/null +++ b/build/torch-cuda/_r2c_f32/fft_plane2.py @@ -0,0 +1,155 @@ +"""Plane kernel (strategy plane2): one CTA per z-plane, the whole plane resident +in shared memory. + +copy-in (16-byte coalesced, the plane is contiguous) -> x pass (real rows -> +packed half spectrum, fft3k.x_task_fn) -> y pass in place on the smem plane +(row stride nx/2+1) -> split of the packed column into kx = 0 and kx = nx/2 -> +copy-out (16-byte coalesced, the output plane is contiguous). The z pass +follows as a separate column kernel (fft3k.make_col_kernel). +""" + +import cutlass +from cutlass import Float32, Int32, Int64, cute + +from .fft3k import XCfg, x_task_fn +from .fft_core import as_f32 +from .plan import choose_line_plan +from .ptx import lds2, lds4, stg4, stg4_if, sts2, sts2_if +from .stockham import TwSmem, b2i, emit_fft + + +class Plane2Cfg: + def __init__(self, nz, ny, nx, yplan=None, xplan=None): + assert nx % 4 == 0 and ny % 2 == 0 + self.nz, self.ny, self.nx = nz, ny, nx + self.M = nx // 2 + self.nf = self.M + 1 + nt_target = 512 + self.yplan = yplan or choose_line_plan(ny, t_target=max(1, nt_target // self.M), emax=32) + self.nty = self.M * self.yplan.T + # x threads per row: match the y-phase thread count (<= 1024 threads) + tx = max(1, min(self.nty, 1024) // ny) + self.cx = XCfg(nz, ny, nx, rows=ny, nt_target=nt_target, t_target=tx, plan=xplan) + if self.cx.nt > 1024 and xplan is None: + self.cx = XCfg(nz, ny, nx, rows=ny, nt_target=nt_target, t_target=max(1, tx // 2)) + self.nt = max(self.cx.nt_r, -(-self.nty // 32) * 32) + self.data = -(-max(self.cx.smem, ny * self.nf * 8) // 16) * 16 + + def finalize(self, r2c_off): + self.tws = TwSmem.for_plans([self.cx.plan, self.yplan], extra=[(r2c_off, self.M)]) + self.smem = self.data + self.tws.size + + def __repr__(self): + return f"Plane2Cfg(nt={self.nt}, data={self.data}, y={self.yplan})\n {self.cx}" + + +def plane2_task_fn(cfg, r2c_off, tws, NT): + """body(tid, z, sbase, gin, gout): full x+y of plane z through smem + (data region [sbase, sbase + cfg.data)); tables `tws` must be staged.""" + cx, yp = cfg.cx, cfg.yplan + ny, M, nf = cfg.ny, cfg.M, cfg.nf + Ty = yp.T + xbody = x_task_fn(cx, r2c_off, tws, NT, copy_out=False) + kinds = ["cyclic"] * len(yp.passes) + out_f4 = ny * nf // 2 + + def y_phase(tid, pl): + k = tid % M + t = tid // M + valid = None + if cfg.nty != NT: + valid = b2i(tid < cfg.nty) + t = cutlass.min(t, Int32(Ty - 1)) + colb = pl + k * 8 + + def y_load(idx): + return [lds2(colb + idx * (nf * 8))] + + def y_addr(g, idx): + return colb + idx * (nf * 8) + + def y_store(idx, vals, pv): + q = pv + if valid is not None: + q = valid if q is None else (q & valid) + v = (as_f32(vals[0][0]), as_f32(vals[0][1])) + if q is None: + sts2(colb + idx * (nf * 8), v) + else: + sts2_if(q, colb + idx * (nf * 8), v) + + def ysync(): + cute.arch.sync_threads() + + # in place: sync between pass-1 loads and its Stockham stores + emit_fft(yp, kinds, t, 1, y_load, y_store, y_addr, tws, ysync, first_hazard=True) + + def separate(tid, pl): + """Column 0 holds P + iQ (P, Q Hermitian in ky): P -> col 0, Q -> col M.""" + vals = [] + for i in range(-(-ny // NT)): + ky = tid + i * NT if i else tid + ok = b2i(ky < ny) + kyc = cutlass.min(ky, Int32(ny - 1)) + nky = Int32(cutlass.select_(kyc == 0, Int32(0), Int32(ny) - kyc)) + f = lds2(pl + kyc * (nf * 8)) + fp = lds2(pl + nky * (nf * 8)) + vals.append((ok, kyc, f, fp)) + cute.arch.sync_threads() + for ok, kyc, f, fp in vals: + a = ((f[0] + fp[0]) * 0.5, (f[1] - fp[1]) * 0.5) + bq = ((f[1] + fp[1]) * 0.5, (fp[0] - f[0]) * 0.5) + sts2_if(ok, pl + kyc * (nf * 8), a) + sts2_if(ok, pl + kyc * (nf * 8) + M * 8, bq) + + def copy_out(tid, pl, gout, z): + blk = gout + Int64(z) * (ny * nf * 8) + for i in range(-(-out_f4 // NT)): + f = tid + i * NT if i else tid + pv = None + if (i + 1) * NT > out_f4: + pv = b2i(f < out_f4) + f = cutlass.min(f, Int32(out_f4 - 1)) + v = lds4(pl + f * 16) + if pv is None: + stg4(blk + Int64(f) * 16, v) + else: + stg4_if(pv, blk + Int64(f) * 16, v) + + def body(tid, z, sbase, gin, gout): + xbody(tid, z, sbase, gin, gout) + cute.arch.sync_threads() + y_phase(tid, sbase) + cute.arch.sync_threads() + separate(tid, sbase) + cute.arch.sync_threads() + copy_out(tid, sbase, gout, z) + + return body + + +def make_plane2_kernel(cfg, r2c_off): + cfg.finalize(r2c_off) + NT = cfg.nt + tws = cfg.tws + smem_bytes = cfg.smem + body = plane2_task_fn(cfg, r2c_off, tws, NT) + + @cute.kernel + def plane2_kernel(inp: cute.Tensor, out: cute.Tensor, tw: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + z, _, _ = cute.arch.block_idx() + smem = cutlass.utils.SmemAllocator() + buf = smem.allocate_array(Float32, num_elems=smem_bytes // 4, byte_alignment=16) + sbase = buf.toint() + # single-wave kernel: let the dependent (z pass) launch early onto idle SMs + cute.arch.griddepcontrol_launch_dependents() + gin = inp.iterator.toint() + tws.stage(sbase + cfg.data, tid, NT, tw.iterator.toint()) + body(tid, z, sbase, gin, out.iterator.toint()) + + @cute.jit + def launch(inp: cute.Tensor, out: cute.Tensor, tw: cute.Tensor, stream): + plane2_kernel(inp, out, tw).launch(grid=(cfg.nz, 1, 1), block=(NT, 1, 1), stream=stream) + + return launch diff --git a/build/torch-cuda/_r2c_f32/fft_xy.py b/build/torch-cuda/_r2c_f32/fft_xy.py new file mode 100644 index 0000000000000000000000000000000000000000..96f78fd04ac1d289d950ba07d9af9a722fbdb371 --- /dev/null +++ b/build/torch-cuda/_r2c_f32/fft_xy.py @@ -0,0 +1,179 @@ +"""Dataflow kernel fusing the x and y passes (strategy xy). + +One launch processes a queue of tasks: x tasks (a block of rows of one plane) +and y tasks (a column tile of one plane). Tasks are fetched with an atomic +counter, so a task only ever waits on tasks with smaller ids, which are already +held by running CTAs (no co-residency assumption). The y tasks of plane z are +queued `lag` planes after its x tasks and wait on a per-plane counter that each +x task of the plane releases after its stores. L2-bound y work therefore runs +alongside the DRAM-bound x work instead of after it. The last CTA to leave +resets the queue state for the next launch, so stream-ordered launches can reuse +one sync buffer; launches that may overlap need their own (kernel.prepare gives +each launch object one). The z pass follows as a separate column kernel +(fft3k.make_col_kernel). +""" + +import cutlass +from cutlass import Float32, Int32, Int64, cute + +from .fft3k import ColCfg, ColTask, XCfg, x_task_fn +from .ptx import atom_add_gpu, lds_s32, release_add_gpu, spin_until_ge, st_global_s32_if, sts_s32_if +from .stockham import TwSmem, b2i + +NUM_SMS = 170 # RTX 5090 + + +class XYCfg: + def __init__(self, nz, ny, nx, W): + self.nz, self.ny, self.nx = nz, ny, nx + cx = XCfg(nz, ny, nx) + # x tasks must not straddle planes + r = cx.rows + while ny % r: + r -= 2 if r > 2 else 1 + if r < 1: + r = 1 + break + if r != cx.rows: + cx = XCfg(nz, ny, nx, rows=r) + cy = ColCfg("y", nz, ny, nx, W=W) + self.cx, self.cy = cx, cy + self.xpp = ny // cx.rows + self.ypp = cy.ntiles + self.nt = max(cx.nt_r, cy.nt_r) + self.data = max(cx.smem, cy.twoff) + self.data = -(-self.data // 16) * 16 + self.total = nz * (self.xpp + self.ypp) + + def finalize(self, r2c_off): + """After twiddle offsets are assigned.""" + cx, cy = self.cx, self.cy + self.tws = TwSmem.for_plans([cx.plan, cy.plan], extra=[(r2c_off, cx.M)]) + self.slot = self.data + self.tws.size + self.smem = self.slot + 16 + occ = (100 * 1024) // (self.smem + 1024) + occ = min(occ, 1536 // self.nt, 65536 // (self.nt * 96)) + self.occ = max(occ, 1) + self.grid = min(self.total, NUM_SMS * self.occ) + B = self.xpp + self.ypp + lag = -(-self.grid // B) + 2 + self.lag = max(1, min(lag, self.nz)) + + def __repr__(self): + return ( + f"XYCfg(xpp={self.xpp}, ypp={self.ypp}, nt={self.nt}, smem={getattr(self, 'smem', None)}, " + f"grid={getattr(self, 'grid', None)}, lag={getattr(self, 'lag', None)}, total={self.total})\n" + f" {self.cx}\n {self.cy}" + ) + + +def make_xy_kernel(cfg, r2c_off): + cfg.finalize(r2c_off) + cx, cy = cfg.cx, cfg.cy + NT = cfg.nt + nz, xpp, ypp, lag = cfg.nz, cfg.xpp, cfg.ypp, cfg.lag + B = xpp + ypp + A = lag * xpp + nsteps = nz - lag # steps with both y (plane s) and x (plane s + lag) + total = cfg.total + tws = cfg.tws + xbody = x_task_fn(cx, r2c_off, tws, NT) + ytask = ColTask(cy, tws, NT) + smem_bytes = cfg.smem + G = cfg.grid + # sync buffer (int32): [0] = task head, [1] = exit count, [2 + z] = plane ready + HEAD, EXITC, READY = 0, 1, 2 + + def decode(t): + """task id -> (is_x, plane, idx), all Int32.""" + u = t - A + s = u // B + w = u - s * B + u2 = u - nsteps * B + in_a = t < A + in_tail = u >= nsteps * B + # region A: x tasks of planes [0, lag) + za, ia = t // xpp, t - (t // xpp) * xpp + # steady: [y tasks of plane s][x tasks of plane s + lag] + isy_m = w < ypp + zm = Int32(cutlass.select_(isy_m, s, s + lag)) + im = Int32(cutlass.select_(isy_m, w, w - ypp)) + # tail: y tasks of planes [nsteps, nz) + zt = nsteps + u2 // ypp + it = u2 - (u2 // ypp) * ypp + is_x = Int32(cutlass.select_(in_a, Int32(1), Int32(cutlass.select_(in_tail, Int32(0), b2i(~isy_m))))) + z = Int32(cutlass.select_(in_a, za, Int32(cutlass.select_(in_tail, zt, zm)))) + idx = Int32(cutlass.select_(in_a, ia, Int32(cutlass.select_(in_tail, it, im)))) + return is_x, z, idx + + def fetch(tid, sync_g, slot): + old = atom_add_gpu(b2i(tid == 0), sync_g + HEAD * 4, Int32(1)) + sts_s32_if(b2i(tid == 0), slot, old) + cute.arch.sync_threads() + return lds_s32(slot) + + def issue_next(tid, sync_g): + """thread 0 grabs the next task id early (published through smem after the task).""" + return atom_add_gpu(b2i(tid == 0), sync_g + HEAD * 4, Int32(1)) + + def do_x(tid, z, idx, sbase, gin, gout, sync_g): + xbody(tid, z * xpp + idx, sbase, gin, gout) + cute.arch.sync_threads() + release_add_gpu(b2i(tid == 0), sync_g + Int64(READY + z) * 4, Int32(1)) + + def do_y(tid, z, idx, sbase, gout, sync_g): + spin_until_ge(b2i(tid == 0), sync_g + Int64(READY + z) * 4, Int32(xpp)) + cute.arch.sync_threads() + c, t, valid = ytask.setup(tid) + tile = z * ypp + idx + cur = ytask.prefetch(tile, c, t, valid, gout, ".cg") + st = ytask.compute(cur, tile, sbase, c, t, valid, gout) + return st + + def finish(tid, sync_g, slot): + cute.arch.sync_threads() + old = atom_add_gpu(b2i(tid == 0), sync_g + EXITC * 4, Int32(1)) + sts_s32_if(b2i(tid == 0), slot, old) + cute.arch.sync_threads() + last = b2i(lds_s32(slot) == G - 1) + # the last CTA out resets head, exit count and the plane counters + for i in range(-(-(nz + 2) // NT)): + e = tid + i * NT if i else tid + ok = last & b2i(e < nz + 2) + e = cutlass.min(e, Int32(nz + 1)) + st_global_s32_if(ok, sync_g + Int64(e) * 4, Int32(0)) + + @cute.kernel + def xy_kernel(inp: cute.Tensor, out: cute.Tensor, tw: cute.Tensor, syncb: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + smem = cutlass.utils.SmemAllocator() + buf = smem.allocate_array(Float32, num_elems=smem_bytes // 4, byte_alignment=16) + sbase = buf.toint() + slot = sbase + cfg.slot + gin = inp.iterator.toint() + gout = out.iterator.toint() + sync_g = syncb.iterator.toint() + cute.arch.griddepcontrol_launch_dependents() + tws.stage(sbase + cfg.data, tid, NT, tw.iterator.toint()) + t = fetch(tid, sync_g, slot) + while t < total: + nxt = issue_next(tid, sync_g) + is_x, z, idx = decode(t) + if is_x != 0: + do_x(tid, z, idx, sbase, gin, gout, sync_g) + else: + st = do_y(tid, z, idx, sbase, gout, sync_g) + if st["tcol"] == 0: + cute.arch.sync_threads() + ytask.separate(st, sbase) + sts_s32_if(b2i(tid == 0), slot, nxt) + cute.arch.sync_threads() + t = lds_s32(slot) + cute.arch.sync_threads() + finish(tid, sync_g, slot) + + @cute.jit + def launch(inp: cute.Tensor, out: cute.Tensor, tw: cute.Tensor, syncb: cute.Tensor, stream): + xy_kernel(inp, out, tw, syncb).launch(grid=(G, 1, 1), block=(NT, 1, 1), stream=stream) + + return launch diff --git a/build/torch-cuda/_r2c_f32/fft_ys1w.py b/build/torch-cuda/_r2c_f32/fft_ys1w.py new file mode 100644 index 0000000000000000000000000000000000000000..f66e487a62f4c2bcdad96a58e65d2b20167cf749 --- /dev/null +++ b/build/torch-cuda/_r2c_f32/fft_ys1w.py @@ -0,0 +1,288 @@ +"""Split-y kernel 1 with a warp-local x pass. + +Same work as fft_ysplit.make_ysplit_kernel1 (x FFT of the rows y = S j + q of +plane z, then the ny/S-point y FFT of every column, stored straight to the +output), but each x row is owned by T consecutive lanes of one warp: + + * staged=True: each warp loads its own rows with coalesced 16-byte loads + and stages them in shared memory; staged=False: the first x pass loads + its inputs straight from global memory into registers (8-byte loads). + Either way the loads are issued before the twiddle tables are staged; + * all x exchanges stay inside the warp's own rows of shared memory, so the + x passes synchronise with bar.warp.sync instead of CTA barriers, and every + warp starts as soon as its own rows have arrived; + * the x output of row r goes to [r * ls, r * ls + nf) inside the same rows, + then one CTA barrier hands the columns to the y pass. +""" + +import cutlass +from cutlass import Float32, Int32, Int64, cute + +from .fft_core import as_f32 +from .plan import SMEM_LIMIT, Layout, choose_layout, lanes_groups +from .ptx import ldg2, ldg4, lds2, stg2, stg2_if, sts2, sts2_if, sts4, sts4_if +from .stockham import access_offsets, b2i, csel, emit_fft, idiv, imod, lay_off, pass1_indices + + +class YS1W: + """Thread / shared-memory geometry of the warp-local kernel 1 for a YSplitCfg.""" + + def __init__(self, cfg): + self.cfg = cfg + xp = cfg.cx.plan + self.T = T = xp.T + if T > 32: + raise ValueError("ys1w: x plan needs T <= 32") + self.rpw = 32 // T # rows per warp + nys, nf, M = cfg.nys, cfg.nf, cfg.M + self.nxw = -(-nys // self.rpw) # warps doing x + self.kinds = ["cyclic"] * (len(xp.passes) - 1) + ["mirror"] + nthr = self.nxw * 32 + + def acc(lay): + def f(tid): + lane = tid % 32 + if lane >= self.rpw * T: + return None + r = (tid // 32) * self.rpw + lane // T + if r >= nys: + return None + return access_offsets(xp, self.kinds, lane % T, lambda g: r, lay, 1) + + return lanes_groups(nthr, f) + + lay = choose_layout(M, nys, acc, min_ls=nf) + if lay.ls % 2: + lay = Layout(lay.n, lay.lines, lay.ls + 1, lay.pp, lay.pq) + self.lay = lay + self.ls = self.lay.ls + self.data = -(-nys * self.ls * 8 // 16) * 16 + self.nty = nf * cfg.yplan.T + self.nt = max(nthr, -(-self.nty // 32) * 32) + if self.nt > 1024: + raise ValueError("ys1w: too many threads") + self.smem = self.data + cfg.tws1.size + if self.smem > SMEM_LIMIT - 1024: + raise ValueError("ys1w: shared memory") + + def __repr__(self): + return f"YS1W(T={self.T}, rpw={self.rpw}, nxw={self.nxw}, nt={self.nt}, lay={self.lay}, smem={self.smem})" + + +def make_ysplit_kernel1w(cfg, staged=True): + """Raises ValueError if the warp-local geometry does not fit.""" + geo = YS1W(cfg) + xp, yp = cfg.cx.plan, cfg.yplan + nz, ny, nx, M, nf, S, nys = cfg.nz, cfg.ny, cfg.nx, cfg.M, cfg.nf, cfg.S, cfg.nys + T, rpw, lay, ls, NT = geo.T, geo.rpw, geo.lay, geo.ls, geo.nt + kinds = geo.kinds + tws = cfg.tws1 + r2c_off = cfg.r2c_off + lp = xp.passes[-1] + R, nb = lp.R, lp.nb + Ty = yp.T + ykinds = ["cyclic"] * len(yp.passes) + smem_bytes = geo.smem + data = geo.data + full_lanes = rpw * T == 32 + nxw = geo.nxw + full_rows = nxw * rpw == nys + + def x_lane(tid): + """(row r, thread t in row, valid) for the x pass; idle lanes duplicate + a valid lane of the same warp (identical values to identical addresses).""" + w = tid // 32 + lane = tid - w * 32 + rl = lane // T + t = lane - rl * T + valid = None + if not full_lanes: + ok = b2i(lane < rpw * T) + valid = ok + rl = cutlass.min(rl, Int32(rpw - 1)) + t = Int32(cutlass.select_(lane < rpw * T, t, Int32(T - 1))) + r = w * rpw + rl + if not full_rows: + ok = b2i(r < nys) + valid = ok if valid is None else (valid & ok) + r = cutlass.min(r, Int32(nys - 1)) + return r, t, valid + + q4 = nx // 4 # float4 chunks per input row + wf4 = rpw * q4 # float4 chunks of one warp's rows + + def warp_chunks(tid): + """(f, pv) per staged float4 chunk of this lane: chunk f of the warp's rows.""" + lane = tid % 32 + out = [] + for i in range(-(-wf4 // 32)): + f = lane + 32 * i if i else lane + pv = None + if (i + 1) * 32 > wf4: + pv = b2i(f < wf4) + f = cutlass.min(f, Int32(wf4 - 1)) + out.append((f, pv)) + return out + + def x_phase(tid, sbase, pre, r, t, valid): + w0 = tid // 32 + + def x_addr(g, idx): + return sbase + lay_off(lay, r, idx) * 8 + + def split(za, zb, w): + ar = (za[0] + zb[0]) * 0.5 + ai = (za[1] - zb[1]) * 0.5 + dr = za[0] - zb[0] + di = za[1] + zb[1] + br = (w[0] * di + w[1] * dr) * 0.5 + bi = (w[1] * di - w[0] * dr) * 0.5 + return (ar + br, ai + bi), (ar - br, bi - ai) + + def wk(k): + return tws.elem(r2c_off, k) + + rowo = sbase + r * (ls * 8) + + def put(k, v, extra=None): + v = (as_f32(v[0]), as_f32(v[1])) + qv = valid + if extra is not None: + qv = extra if qv is None else (qv & extra) + if qv is None: + sts2(rowo + k * 8, v) + else: + sts2_if(qv, rowo + k * 8, v) + + def x_final(outs): + (_, y0, _), (j1, y1, _) = outs + y0 = y0[0] + y1 = y1[0] + is0 = t == 0 + for m in range(R): + mb = R - 1 - m + za = y0[m] + zb = csel(is0, y0[(R - m) % R], y1[mb]) + ka = t + m * nb if m else t + xa, xb = split(za, zb, wk(ka)) + if m == 0: + put(M, (xb[0], 0.0), b2i(is0)) # X[nx/2] (thread 0 only) + x2, _ = split(y1[mb], y1[m], wk(nb // 2 + mb * nb)) + vb = csel(is0, x2, xb) + kb = j1 + mb * nb if mb else j1 + put(ka, xa) + put(kb, vb) + + def wsync(): + cute.arch.sync_warp() + + if staged: + # pre = this lane's float4 chunks of the warp's rows (coalesced loads) + for (f, pv), v in pre: + a = sbase + (w0 * rpw) * (ls * 8) + idiv(f, q4) * (ls * 8) + imod(f, q4) * 16 + if pv is None: + sts4(a, v) + else: + sts4_if(pv, a, v) + wsync() + + def x_load(idx): + return [lds2(rowo + idx * 8)] + + emit_fft( + xp, kinds, t, 1, x_load, None, x_addr, tws, wsync, final=x_final, first_hazard=True, last_hazard=True + ) + else: + emit_fft( + xp, + kinds, + t, + 1, + None, + None, + x_addr, + tws, + wsync, + final=x_final, + first_hazard=False, + last_hazard=len(xp.passes) > 1, + pre=pre, + ) + + def y_phase(tid, sbase, gcol): + k = tid % nf + t = tid // nf + valid = None + if geo.nty != NT: + valid = b2i(tid < geo.nty) + t = cutlass.min(t, Int32(Ty - 1)) + colb = sbase + k * 8 + gcol = gcol + Int64(k) * 8 + + def y_load(idx): + return [lds2(colb + idx * (ls * 8))] + + def y_addr(g, idx): + return colb + idx * (ls * 8) + + def y_store(idx, vals, pv): + qv = pv + if valid is not None: + qv = valid if qv is None else (qv & valid) + v = (as_f32(vals[0][0]), as_f32(vals[0][1])) + a = gcol + Int64(idx) * (nf * 8) + if qv is None: + stg2(a, v) + else: + stg2_if(qv, a, v) + + def ysync(): + cute.arch.sync_threads() + + emit_fft(yp, ykinds, t, 1, y_load, y_store, y_addr, tws, ysync, first_hazard=True) + + @cute.kernel + def ys1w_kernel(inp: cute.Tensor, out: cute.Tensor, tw: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + smem = cutlass.utils.SmemAllocator() + buf = smem.allocate_array(Float32, num_elems=smem_bytes // 4, byte_alignment=16) + sbase = buf.toint() + cute.arch.griddepcontrol_launch_dependents() + gin = inp.iterator.toint() + gout = out.iterator.toint() + z = bid // S + q = bid - z * S + r, t, valid = x_lane(tid) + if cutlass.const_expr(NT > nxw * 32): + r = cutlass.min(r, Int32(nys - 1)) # y-only warps: keep their (unused) loads in bounds + if cutlass.const_expr(staged): + # coalesced 16-byte loads of the warp's own rows (row y = q + S * rr of plane z) + wrow0 = cutlass.min((tid // 32) * rpw, Int32(nys - 1)) + pre = [] + for f, pv in warp_chunks(tid): + rr = cutlass.min(wrow0 + idiv(f, q4), Int32(nys - 1)) + a = gin + (Int64(z * ny + q + S * rr) * nx + Int64(imod(f, q4)) * 4) * 4 + rows_ok = b2i(wrow0 + idiv(f, q4) < nys) if not full_rows else None + pvv = pv if rows_ok is None else (rows_ok if pv is None else (pv & rows_ok)) + pre.append(((f, pvv), ldg4(a, ".cs"))) + else: + rin = gin + Int64((z * ny + q + S * r) * nx) * 4 + # first x pass inputs straight from global (complex z[i] = x[2i] + i x[2i+1]) + pre = [[ldg2(rin + Int64(idx) * 8, ".cs")] for idx in pass1_indices(xp, kinds, t)] + tws.stage(sbase + data, tid, NT, tw.iterator.toint()) + cute.arch.sync_threads() + if cutlass.const_expr(NT > nxw * 32): + # warps past the x rows only take part in the y pass + if tid < nxw * 32: + x_phase(tid, sbase, pre, r, t, valid) + else: + x_phase(tid, sbase, pre, r, t, valid) + cute.arch.sync_threads() + y_phase(tid, sbase, gout + Int64(z * ny + q * nys) * (nf * 8)) + + @cute.jit + def launch(inp: cute.Tensor, out: cute.Tensor, tw: cute.Tensor, stream): + ys1w_kernel(inp, out, tw).launch(grid=(nz * S, 1, 1), block=(NT, 1, 1), stream=stream) + + return launch diff --git a/build/torch-cuda/_r2c_f32/fft_ysplit.py b/build/torch-cuda/_r2c_f32/fft_ysplit.py new file mode 100644 index 0000000000000000000000000000000000000000..c1b14a566e71cfea458a13ae16f65d20588eda5b --- /dev/null +++ b/build/torch-cuda/_r2c_f32/fft_ysplit.py @@ -0,0 +1,268 @@ +"""Split-y two-kernel scheme (strategy ysplit). + +The y transform of length ny is split by decimation in time, ny = S * ny_s: + Y_q[k'] = FFT_{ny_s}( x[S j + q] )[k'], q < S + X[k' + k2 ny_s] = sum_q w_ny^{q k'} w_S^{q k2} Y_q[k'] + +kernel 1 (CTA per (z, q)): x pass of rows y = S j + q of plane z (staged, + unpacked: columns 0 .. nx/2 incl. the Nyquist column), then the ny_s-point + y FFT of every column in shared memory; the last y pass stores Y_q straight + to output rows [q ny_s, (q + 1) ny_s) of plane z. fft_ys1w has a variant + with a warp-local x pass. +kernel 2 (flat tiles of W consecutive (k', kx) groups, S rows per thread): + loads Y_q[k'] for all z, finishes y with the twiddled S-point DFT, then the + z FFT (`rounds` exchange rounds of S / rounds rows each), and stores rows + k' + k2 ny_s in place. +""" + +import cutlass +from cutlass import Float32, Int32, Int64, cute + +from .fft3k import XCfg, x_task_fn +from .fft_core import as_f32, cmul, dft +from .plan import SMEM_LIMIT, Layout, choose_layout, choose_line_plan, lanes_groups +from .ptx import ldg2, lds2, stg2, stg2_if +from .stockham import TwSmem, access_offsets, b2i, emit_fft, lay_off, pass1_indices + + +class YSplitCfg: + def __init__(self, nz, ny, nx, S=2, nt_target=512, yplan=None, W=None, zplan=None, znt=None, rounds=None): + W0 = W if W is not None else 16 + self.rounds = rounds if rounds else 1 + while True: + self._init(nz, ny, nx, S, nt_target, yplan, W0, zplan, znt) + est2 = self.data2 + sum(p.Ns * p.R * 8 for p in self.zplan.passes if p.Ns > 1) + ny * 8 + 64 + if est2 <= SMEM_LIMIT - 1024 or self.W <= 2: + break + W0 = max(1, self.W // 2) + + def _init(self, nz, ny, nx, S, nt_target, yplan, W, zplan, znt): + if nx % 4 or ny % S or (ny // S) % 2: + raise ValueError("ysplit: needs nx % 4 == 0 and (ny / S) even") + self.nz, self.ny, self.nx, self.S = nz, ny, nx, S + self.M = nx // 2 + self.nf = self.M + 1 + self.nys = ny // S + nf, nys = self.nf, self.nys + # ---- kernel 1 + self.yplan = yplan or choose_line_plan(nys, t_target=max(1, nt_target // nf), emax=32) + self.nty = nf * self.yplan.T + tx = max(1, min(self.nty, 1024) // nys) + self.cx = XCfg(nz, ny, nx, rows=nys, nt_target=nt_target, t_target=tx) + self.nt1 = max(self.cx.nt_r, -(-self.nty // 32) * 32) + self.data1 = -(-max(self.cx.smem, nys * nf * 8) // 16) * 16 + if self.nt1 > 1024: + raise ValueError("ysplit: kernel 1 needs too many threads") + # ---- kernel 2: flat tiles of W consecutive (k', kx) groups. Rows k' of + # one row set are adjacent in memory, so a tile is one contiguous W * 8 + # byte segment per (row set, z) -- aligned when ny_s * nf is, no idle lanes. + self.NG = nys * nf + self.W = W + if zplan is None: + target = znt if znt else 512 + zplan = choose_line_plan(nz, t_target=max(1, target // W), emax=max(8, 32 // S)) + self.zplan = zplan + T = zplan.T + self.nt2 = -(-(W * T) // 32) * 32 + if self.nt2 > 1024: + raise ValueError("ysplit: kernel 2 needs too many threads") + self.grid2 = -(-self.NG // W) + G = S // self.rounds # lines resident in the exchange buffer per round + kinds = ["cyclic"] * len(zplan.passes) + self.kinds2 = kinds + + def acc(lay): + def f(tid): + if tid >= W * T: + return None + c, t = tid % W, tid // W + return access_offsets(zplan, kinds, t, lambda g: g * W + c, lay, G) + + return lanes_groups(W * T, f) + + if len(zplan.passes) > 1: + self.lay2 = choose_layout(nz, G * W, acc) + self.data2 = -(-self.lay2.size * 8 // 16) * 16 + else: + self.lay2 = Layout(nz, G * W, nz, nz, 0) + self.data2 = 16 + self.G2 = G + + def finalize(self, r2c_off, wy_off): + self.r2c_off, self.wy_off = r2c_off, wy_off + self.tws1 = TwSmem.for_plans([self.cx.plan, self.yplan], extra=[(r2c_off, self.M)]) + self.smem1 = self.data1 + self.tws1.size + self.tws2 = TwSmem.for_plans([self.zplan], extra=[(wy_off, self.ny)]) + self.smem2 = self.data2 + self.tws2.size + if self.smem1 > SMEM_LIMIT - 1024 or self.smem2 > SMEM_LIMIT - 1024: + raise ValueError("ysplit: shared memory") + + def __repr__(self): + return ( + f"YSplitCfg(S={self.S}, nt1={self.nt1}, data1={self.data1}, y={self.yplan}, " + f"W={self.W}, nt2={self.nt2}, z={self.zplan}, lay2={self.lay2}, grid2={self.grid2})\n" + f" {self.cx}" + ) + + +def make_ysplit_kernel1(cfg): + cx, yp = cfg.cx, cfg.yplan + nz, ny, nf, S, nys = cfg.nz, cfg.ny, cfg.nf, cfg.S, cfg.nys + NT = cfg.nt1 + tws = cfg.tws1 + Ty = yp.T + xbody = x_task_fn(cx, cfg.r2c_off, tws, NT, copy_out=False, ysplit=S, pack=False) + kinds = ["cyclic"] * len(yp.passes) + smem_bytes = cfg.smem1 + + def y_phase(tid, pl, gcol): + k = tid % nf + t = tid // nf + valid = None + if cfg.nty != NT: + valid = b2i(tid < cfg.nty) + t = cutlass.min(t, Int32(Ty - 1)) + colb = pl + k * 8 + gcol = gcol + Int64(k) * 8 + + def y_load(idx): + return [lds2(colb + idx * (nf * 8))] + + def y_addr(g, idx): + return colb + idx * (nf * 8) + + def y_store(idx, vals, pv): + q = pv + if valid is not None: + q = valid if q is None else (q & valid) + v = (as_f32(vals[0][0]), as_f32(vals[0][1])) + a = gcol + Int64(idx) * (nf * 8) + if q is None: + stg2(a, v) + else: + stg2_if(q, a, v) + + def ysync(): + cute.arch.sync_threads() + + emit_fft(yp, kinds, t, 1, y_load, y_store, y_addr, tws, ysync, first_hazard=True) + + @cute.kernel + def ys1_kernel(inp: cute.Tensor, out: cute.Tensor, tw: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + smem = cutlass.utils.SmemAllocator() + buf = smem.allocate_array(Float32, num_elems=smem_bytes // 4, byte_alignment=16) + sbase = buf.toint() + cute.arch.griddepcontrol_launch_dependents() + gout = out.iterator.toint() + tws.stage(sbase + cfg.data1, tid, NT, tw.iterator.toint()) + z = bid // S + q = bid - z * S + xbody(tid, bid, sbase, inp.iterator.toint(), gout) + cute.arch.sync_threads() + y_phase(tid, sbase, gout + Int64(z * ny + q * nys) * (nf * 8)) + + @cute.jit + def launch(inp: cute.Tensor, out: cute.Tensor, tw: cute.Tensor, stream): + ys1_kernel(inp, out, tw).launch(grid=(nz * S, 1, 1), block=(NT, 1, 1), stream=stream) + + return launch + + +def make_ysplit_kernel2(cfg): + """y combine (twiddled S-point DFT over the row groups) + z FFT, in place.""" + ny, nf, S, nys = cfg.ny, cfg.nf, cfg.S, cfg.nys + W, zp, lay = cfg.W, cfg.zplan, cfg.lay2 + T = zp.T + G = S + GR = cfg.G2 # lines per exchange round + rounds = G // GR + NT = cfg.nt2 + tws = cfg.tws2 + kinds = cfg.kinds2 + smem_bytes = cfg.smem2 + zstride = ny * nf * 8 # bytes between consecutive z + nthr = W * T + + def setup(tid): + c = tid % W + t = tid // W + valid = None + if nthr != NT: + valid = b2i(tid < nthr) + t = cutlass.min(t, Int32(T - 1)) + return c, t, valid + + def tile_info(bid, c, valid): + g = bid * W + c # flat group k' * nf + kx + v = valid + if cfg.NG % W: + ok = b2i(g < cfg.NG) + v = ok if v is None else (v & ok) + g = cutlass.min(g, Int32(cfg.NG - 1)) + kp = g // nf + return kp, g, v + + def body(tid, bid, sbase, gout): + c, t, valid = setup(tid) + kp, g, v = tile_info(bid, c, valid) + # per-row-group bases: row k' + q * ny_s, column kx (g = k' * nf + kx) + bases = [gout + Int64(g + q * nys * nf) * 8 for q in range(S)] + # twiddles w_ny^{q k'} for q >= 1 + wq = [None] + [tws.elem(cfg.wy_off, kp * q) for q in range(1, S)] + + def z_load(idx): + return [ldg2(bases[q] + Int64(idx) * zstride) for q in range(S)] + + def z_finish(raw): + ys = [raw[0]] + [cmul(raw[q], wq[q]) for q in range(1, S)] + return dft(ys, -1) + + def z_addr(g, idx): + return sbase + lay_off(lay, c + g * W, idx) * 8 + + def z_store_round(r0): + def z_store(idx, vals, pv): + q_ = pv if v is None else (v if pv is None else (pv & v)) + for g in range(len(vals)): + val = (as_f32(vals[g][0]), as_f32(vals[g][1])) + a = bases[r0 + g] + Int64(idx) * zstride + if q_ is None: + stg2(a, val) + else: + stg2_if(q_, a, val) + + return z_store + + def zsync(): + cute.arch.barrier(barrier_id=1, number_of_threads=NT) + + if rounds == 1: + emit_fft(zp, kinds, t, G, z_load, z_store_round(0), z_addr, tws, zsync, finish=z_finish) + else: + # load + combine all S rows in registers, then z FFT GR rows at a time + comb = [z_finish(z_load(idx)) for idx in pass1_indices(zp, kinds, t)] + for r in range(rounds): + if r: + zsync() + pre = [cv[r * GR : (r + 1) * GR] for cv in comb] + emit_fft(zp, kinds, t, GR, None, z_store_round(r * GR), z_addr, tws, zsync, pre=pre) + + @cute.kernel + def ys2_kernel(out: cute.Tensor, tw: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + smem = cutlass.utils.SmemAllocator() + buf = smem.allocate_array(Float32, num_elems=smem_bytes // 4, byte_alignment=16) + sbase = buf.toint() + tws.stage(sbase + cfg.data2, tid, NT, tw.iterator.toint()) + cute.arch.griddepcontrol_wait() + cute.arch.sync_threads() + body(tid, bid, sbase, out.iterator.toint()) + + @cute.jit + def launch(out: cute.Tensor, tw: cute.Tensor, stream): + ys2_kernel(out, tw).launch(grid=(cfg.grid2, 1, 1), block=(NT, 1, 1), stream=stream, use_pdl=True) + + return launch diff --git a/build/torch-cuda/_r2c_f32/kernel.py b/build/torch-cuda/_r2c_f32/kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..21bf88bc47539ae7aecde71a54717638fcd588f9 --- /dev/null +++ b/build/torch-cuda/_r2c_f32/kernel.py @@ -0,0 +1,392 @@ +"""3-D forward real-to-complex FFT (fp32) for the RTX 5090, in CuTe DSL. + +run(input) -> output + input: contiguous float32 CUDA tensor [nz, ny, nx] (x fastest), nz, ny >= 8, + nx >= 16. output: new contiguous float32 [nz, ny, nx//2 + 1, 2] equal to + torch.view_as_real(torch.fft.rfftn(input, norm="backward")). +prepare(input) -> launch + Plans and compiles once per (shape, device) and allocates the output; + launch() takes no arguments, recomputes from the input buffer's current + contents into that same output tensor and returns it. +Distinct launch objects and run() calls share no mutable device state, so they +may run concurrently on different streams; the launches of one launch object +must be stream-ordered (they write the same output). + +Routing (_default_strategy, first match; the benchmark grids use offline-tuned +parameters): + generic : nx % 4 != 0 or no radix plan fits (also the fallback if a fast plan + fails to build): direct DFTs, float64 accumulation (fft_generic). + ysplit : the tuned grids, and >= 2^24-point grids whose output stays L2 + resident. The y FFT is split by decimation in time, ny = S * ny/S: + kernel 1 does x + the ny/S-point y FFT of every S-th row of a plane + (fft_ysplit, or fft_ys1w with a warp-local x pass), kernel 2 the + twiddled radix-S y combine fused with the z FFT. + plane2 : a z-plane fits in shared memory: kernel 1 does x + y of one plane + (fft_plane2), kernel 2 the z pass (fft3k column kernel). + ysplit : < 2^24 points, if the ysplit kernels fit. + xy : >= 2^24 points: one dataflow kernel interleaves x row blocks with y + column tiles (fft_xy), then the z pass. + 3k : otherwise: x, y and z passes as three kernels (fft3k). +The x pass computes a real row as an nx/2-point complex FFT and splits it into +the half spectrum in registers; in plane2 / xy / 3k, X[0] and X[nx/2] travel +packed in column 0 until the y pass separates them by Hermitian symmetry. The +kernels of one transform are chained with programmatic dependent launch. +Twiddle factors are computed in float64 on the host and rounded once to fp32. +""" + +import copy +import functools + +import cutlass +import torch +from cutlass import cute + +from .fft3k import ColCfg, XCfg, make_col_kernel, make_x_kernel +from .fft_generic import make_generic_col_kernel, make_generic_x_kernel, twiddle_f64 +from .fft_plane2 import Plane2Cfg, make_plane2_kernel +from .fft_xy import XYCfg, make_xy_kernel +from .fft_ys1w import make_ysplit_kernel1w +from .fft_ysplit import YSplitCfg, make_ysplit_kernel1, make_ysplit_kernel2 +from .plan import SMEM_LIMIT, LinePlan, TwiddleBuilder + + +def _fake1d(align=16): + return cute.runtime.make_fake_compact_tensor(cutlass.Float32, (cute.sym_int(),), assumed_align=align) + + +def _fake1d_f64(align=16): + return cute.runtime.make_fake_compact_tensor(cutlass.Float64, (cute.sym_int(),), assumed_align=align) + + +def _fake1d_i32(align=16): + return cute.runtime.make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(),), assumed_align=align) + + +def _compile(fn, *args): + return cute.compile(fn, *args, options="--enable-tvm-ffi") + + +# Offline-tuned per-shape choices for the benchmark grids; any other shape uses +# the heuristics. +# plane2: (x plan (T, radices), y plan (T, radices)); z: ColCfg kwargs +_TUNED_PLANE2 = { + (32, 32, 32): ((2, [4, 4]), (4, [8, 4])), +} +_TUNED_Z = { + (32, 32, 32): {"W": 16, "nt_target": 128}, +} +# split-y parameters; a shape listed here uses "ysplit". Two-pass z plans (one +# exchange per round) beat three-pass ones. k1: kernel-1 variant, "w" = +# warp-local x pass with coalesced staging, "wd" = warp-local x pass loading +# straight from global (fft_ys1w); absent = fft_ysplit kernel 1. +_TUNED_YS = { + (64, 64, 64): {"S": 2, "W": 8, "znt": 256, "nt_target": 512, "rounds": 2, "k1": "w"}, + (96, 96, 96): { + "S": 4, + "W": 8, + "znt": 256, + "nt_target": 256, + "rounds": 4, + "k1": "w", + "yplan": LinePlan(24, 6, [4, 6]), + "zplan": LinePlan(96, 12, [8, 12]), + }, + (128, 128, 128): { + "S": 8, + "W": 8, + "znt": 512, + "nt_target": 1024, + "rounds": 8, + "k1": "w", + "zplan": LinePlan(128, 16, [8, 16]), + }, + (80, 80, 112): { + "S": 4, + "W": 8, + "znt": 256, + "nt_target": 256, + "rounds": 2, + "k1": "wd", + "zplan": LinePlan(80, 10, [8, 10]), + }, + (96, 96, 144): { + "S": 8, + "W": 8, + "znt": 256, + "nt_target": 256, + "rounds": 4, + "k1": "w", + "zplan": LinePlan(96, 12, [8, 12]), + }, + (192, 192, 192): {"S": 4, "W": 8, "znt": 128, "nt_target": 512, "rounds": 2, "zplan": LinePlan(192, 16, [16, 12])}, + (256, 256, 256): { + "S": 4, + "W": 8, + "znt": 1024, + "nt_target": 512, + "rounds": 4, + "k1": "w", + "zplan": LinePlan(256, 16, [16, 16]), + }, +} + + +def _plane2cfg(nz, ny, nx): + t = _TUNED_PLANE2.get((nz, ny, nx)) + if t is None: + return Plane2Cfg(nz, ny, nx) + (tx, rx), (ty, ry) = t + return Plane2Cfg(nz, ny, nx, xplan=LinePlan(nx // 2, tx, rx), yplan=LinePlan(ny, ty, ry)) + + +def _zcfg(nz, ny, nx): + return ColCfg("z", nz, ny, nx, **_TUNED_Z.get((nz, ny, nx), {})) + + +def _fast_ok(nz, ny, nx): + """Can the fast (half-spectrum, mixed-radix) kernels handle this?""" + if nx % 4 or nx < 16 or ny < 2 or nz < 2: + return False + try: + XCfg(nz, ny, nx) + ColCfg("y", nz, ny, nx) + ColCfg("z", nz, ny, nx) + except Exception: # noqa: BLE001 -- any config failure means "not on the fast path" + return False + return True + + +def _ys_fits(c, ny): + tw1 = sum(p.Ns * p.R * 8 for pl in (c.cx.plan, c.yplan) for p in pl.passes if p.Ns > 1) + c.M * 8 + 64 + tw2 = sum(p.Ns * p.R * 8 for p in c.zplan.passes if p.Ns > 1) + ny * 8 + 64 + return ( + c.data1 + tw1 <= SMEM_LIMIT - 1024 + and c.data2 + tw2 <= SMEM_LIMIT - 1024 + and c.nt1 <= 1024 + and c.cx.nt <= 1024 + and c.nt2 <= 1024 + ) + + +@functools.cache +def _ys_kw_cached(nz, ny, nx): + t = _TUNED_YS.get((nz, ny, nx)) + if t is not None: + return dict(t) + if nx % 4: + return None + nf = nx // 2 + 1 + for S in (4, 2, 8): + # cheap pre-check: the kernel-1 sub-plane must fit in shared memory + if ny % (2 * S) or (ny // S) * nf * 8 > SMEM_LIMIT - 4096: + continue + kw = {"S": S, "W": 8, "znt": 256, "nt_target": 512, "rounds": 2 if S >= 4 else 1} + try: + if _ys_fits(YSplitCfg(nz, ny, nx, **kw), ny): + return kw + except (ValueError, AssertionError): + pass + return None + + +def _ys_kw(nz, ny, nx): + """Split-y parameters: tuned table, else the first S in (4, 2, 8) that + divides ny (with ny / S even) and fits; None if none does.""" + kw = _ys_kw_cached(nz, ny, nx) + return copy.deepcopy(kw) if kw is not None else None # plans get annotated per build + + +def _default_strategy(nz, ny, nx): + if not _fast_ok(nz, ny, nx): + return "generic" + if (nz, ny, nx) in _TUNED_YS: + return "ysplit" + big = nz * ny * nx >= 1 << 24 + # split-y while its intermediate stays L2 resident (beyond that xy is as fast) + if big and nz * ny * (nx // 2 + 1) * 8 <= 72 << 20 and _ys_kw(nz, ny, nx) is not None: + return "ysplit" + if nx % 4 == 0 and ny % 2 == 0 and ny * (nx // 2 + 1) * 8 <= SMEM_LIMIT - 4096: + try: + cp = Plane2Cfg(nz, ny, nx) + twb = sum(p.Ns * p.R * 8 for pl in (cp.cx.plan, cp.yplan) for p in pl.passes if p.Ns > 1) + cp.M * 8 + 64 + if cp.data + twb <= SMEM_LIMIT - 1024 and cp.nt <= 1024 and cp.cx.nt <= 1024: + return "plane2" + except (ValueError, AssertionError): + pass + if not big and _ys_kw(nz, ny, nx) is not None: + return "ysplit" + if big: + try: + c = XYCfg(nz, ny, nx, W=16) + if c.data + 16384 <= SMEM_LIMIT - 1024 and c.nt <= 1024: + return "xy" + except (ValueError, AssertionError): + pass + return "3k" + + +class _Plan: + """Configs, twiddle buffer and compiled kernels for one (shape, device); all of it + is read-only once built, so concurrent launches can share a plan. + steps: the kernels of one transform, each called as f(inp_flat, out_flat), or as + f(inp_flat, out_flat, sync) if sync_words > 0: the int32 scratch `sync` of that + many words is what the kernels synchronize through (xy). It must be zero before + its first launch, the kernels leave it zero again, and launches that may overlap + need their own (see prepare).""" + + def __init__(self, nz, ny, nx, device, strategy=None): + self.shape = (nz, ny, nx) + self.device = device + self.sync_words = 0 + if strategy is None: + strategy = _default_strategy(nz, ny, nx) + self.strategy = strategy + M = nx // 2 + tb = TwiddleBuilder() + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + if strategy == "plane2": + cp = _plane2cfg(nz, ny, nx) + cz = _zcfg(nz, ny, nx) + for pl in (cp.cx.plan, cp.yplan, cz.plan): + tb.add_pass_tables(pl) + r2c_off = tb.add_roots(range(M), nx) + self.tw = tb.array().to(device) + kp = _compile(make_plane2_kernel(cp, r2c_off), _fake1d(), _fake1d(), _fake1d(), stream) + kz = _compile(make_col_kernel(cz), _fake1d(), _fake1d(), stream) + self.steps = [lambda i, o: kp(i, o, self.tw), lambda i, o: kz(o, self.tw)] + elif strategy == "generic": + self.twx = twiddle_f64(nx).to(device) + self.twy = twiddle_f64(ny).to(device) + self.twz = twiddle_f64(nz).to(device) + kx = _compile(make_generic_x_kernel(nz, ny, nx), _fake1d(4), _fake1d(4), _fake1d_f64(8), stream) + ky = _compile(make_generic_col_kernel("y", nz, ny, nx), _fake1d(4), _fake1d_f64(8), stream) + kz = _compile(make_generic_col_kernel("z", nz, ny, nx), _fake1d(4), _fake1d_f64(8), stream) + self.steps = [ + lambda i, o: kx(i, o, self.twx), + lambda i, o: ky(o, self.twy), + lambda i, o: kz(o, self.twz), + ] + elif strategy == "ysplit": + ykw = _ys_kw(nz, ny, nx) or {} + k1v = ykw.pop("k1", None) + if (nz, ny, nx) not in _TUNED_YS and nz * ny * nx <= 1 << 22: + k1v = "w" + cys = YSplitCfg(nz, ny, nx, **ykw) + for pl in (cys.cx.plan, cys.yplan, cys.zplan): + tb.add_pass_tables(pl) + r2c_off = tb.add_roots(range(M), nx) + wy_off = tb.add_roots(range(ny), ny) + cys.finalize(r2c_off, wy_off) + self.tw = tb.array().to(device) + k1 = None + if k1v in ("w", "wd"): + try: + mk = make_ysplit_kernel1w(cys, staged=k1v == "w") + except (ValueError, AssertionError): + mk = None # warp-local geometry does not fit + if mk is not None: + k1 = _compile(mk, _fake1d(), _fake1d(), _fake1d(), stream) + if k1 is None: + k1 = _compile(make_ysplit_kernel1(cys), _fake1d(), _fake1d(), _fake1d(), stream) + k2 = _compile(make_ysplit_kernel2(cys), _fake1d(), _fake1d(), stream) + self.steps = [lambda i, o: k1(i, o, self.tw), lambda i, o: k2(o, self.tw)] + elif strategy == "xy": + cxy = XYCfg(nz, ny, nx, W=16) + cz = _zcfg(nz, ny, nx) + for c in (cxy.cx, cxy.cy, cz): + tb.add_pass_tables(c.plan) + r2c_off = tb.add_roots(range(M), nx) + self.tw = tb.array().to(device) + # kxy's task queue head, exit count and per-plane ready counters + self.sync_words = nz + 2 + kxy = _compile(make_xy_kernel(cxy, r2c_off), _fake1d(), _fake1d(), _fake1d(), _fake1d_i32(), stream) + kz = _compile(make_col_kernel(cz), _fake1d(), _fake1d(), stream) + self.steps = [lambda i, o, s: kxy(i, o, self.tw, s), lambda i, o, s: kz(o, self.tw)] + else: # "3k" + cx = XCfg(nz, ny, nx) + cy = ColCfg("y", nz, ny, nx) + cz = _zcfg(nz, ny, nx) + for c in (cx, cy, cz): + tb.add_pass_tables(c.plan) + r2c_off = tb.add_roots(range(M), nx) + self.tw = tb.array().to(device) + kx = _compile(make_x_kernel(cx, r2c_off), _fake1d(), _fake1d(), _fake1d(), stream) + ky = _compile(make_col_kernel(cy, trigger=True), _fake1d(), _fake1d(), stream) + kz = _compile(make_col_kernel(cz), _fake1d(), _fake1d(), stream) + self.steps = [ + lambda i, o: kx(i, o, self.tw), + lambda i, o: ky(o, self.tw), + lambda i, o: kz(o, self.tw), + ] + + def run(self, inp_flat, out_flat): + for f in self.steps: + f(inp_flat, out_flat) + + def run_sync(self, inp_flat, out_flat, sync): + for f in self.steps: + f(inp_flat, out_flat, sync) + + +_PLANS = {} + + +def _get_plan(shape, device): + key = (tuple(shape), str(device)) + pl = _PLANS.get(key) + if pl is None: + nz, ny, nx = shape + try: + pl = _Plan(nz, ny, nx, device) + except Exception: # noqa: BLE001 -- a fast plan that fails to build or compile falls back to generic + pl = _Plan(nz, ny, nx, device, strategy="generic") + _PLANS[key] = pl + return pl + + +def _reject(input): + """Raise the error saying why `input` is not a valid run / prepare input.""" + if input.dim() != 3: + raise ValueError(f"input must be a 3-D [nz, ny, nx] tensor, got shape {tuple(input.shape)}") + if input.dtype != torch.float32: + raise TypeError(f"input must be float32, got {input.dtype}") + if not input.is_cuda: + raise ValueError(f"input must be a CUDA tensor, got device {input.device}") + raise ValueError("input must be contiguous") + + +def prepare(input): + if not (input.dim() == 3 and input.dtype == torch.float32 and input.is_cuda and input.is_contiguous()): + _reject(input) + nz, ny, nx = input.shape + pl = _get_plan((nz, ny, nx), input.device) + out = torch.empty((nz, ny, nx // 2 + 1, 2), device=input.device, dtype=torch.float32) + inp_flat = input.view(-1) + out_flat = out.view(-1) + + if not pl.sync_words: + + def launch(): + pl.run(inp_flat, out_flat) + return out + + return launch + + # Sync scratch per launch object, so launches of one shape on different streams + # never share it. It is zeroed on the stream of the first launch (ordered before + # it); the kernels re-arm it on exit for the next launch, which writes the same + # output and so is stream-ordered after this one. + sync = None + + def launch(): + nonlocal sync + if sync is None: + sync = torch.zeros(pl.sync_words, dtype=torch.int32, device=out.device) + pl.run_sync(inp_flat, out_flat, sync) + return out + + return launch + + +def run(input): + return prepare(input)() diff --git a/build/torch-cuda/_r2c_f32/plan.py b/build/torch-cuda/_r2c_f32/plan.py new file mode 100644 index 0000000000000000000000000000000000000000..588ba3169380d5826c56abc5bfd5abfc51bc8787 --- /dev/null +++ b/build/torch-cuda/_r2c_f32/plan.py @@ -0,0 +1,244 @@ +"""Host-side planning: radix plans for the 1-D line FFTs, the twiddle table buffer +and shared-memory layouts (padding chosen by simulating warp bank conflicts).""" + +import math + +import torch + +# max dynamic shared memory per CTA on sm_120 +SMEM_LIMIT = 99 * 1024 + +# rough fp-instruction cost of one radix-R butterfly (generator output) +_DFT_OPS = { + 1: 0, + 2: 4, + 3: 12, + 4: 16, + 5: 32, + 6: 36, + 7: 60, + 8: 52, + 9: 80, + 10: 84, + 11: 110, + 12: 96, + 13: 140, + 14: 140, + 15: 150, + 16: 144, + 18: 190, + 20: 220, + 24: 260, + 25: 300, + 27: 330, + 28: 340, + 32: 376, +} + + +def dft_ops(r): + if r in _DFT_OPS: + return _DFT_OPS[r] + return 2 * r * r + + +class Pass: + """One Stockham radix-R pass over an n-point line: butterfly j in + [0, n/R) reads x[j + m*n/R] and writes y[(j//Ns)*Ns*R + j%Ns + m*Ns]. + T threads share the line, each doing `per` butterflies.""" + + def __init__(self, n, R, Ns, T): + self.n, self.R, self.Ns, self.T = n, R, Ns, T + self.nb = n // R + self.per = -(-self.nb // T) + self.full = self.nb % T == 0 + self.tw_off = None # complex offset of the [Ns][R] twiddle table (Ns > 1) + + def __repr__(self): + return f"R{self.R}(Ns={self.Ns},per={self.per}{'' if self.full else '*'})" + + +class LinePlan: + """n-point line FFT as Stockham passes with the given radices, T threads per line.""" + + def __init__(self, n, T, radices): + assert math.prod(radices) == n + self.n, self.T, self.radices = n, T, list(radices) + self.passes = [] + ns = 1 + for r in radices: + self.passes.append(Pass(n, r, ns, T)) + ns *= r + self.emax = max(p.per * p.R for p in self.passes) + + def __repr__(self): + return f"LinePlan(n={self.n}, T={self.T}, {self.passes})" + + +def _ordered_factorizations(n, maxlen, maxr): + out = [] + + def rec(rem, acc): + if rem == 1: + if acc: + out.append(list(acc)) + return + if len(acc) == maxlen: + return + for r in range(2, min(rem, maxr) + 1): + if rem % r == 0: + acc.append(r) + rec(rem // r, acc) + acc.pop() + + rec(n, []) + return out + + +def line_plan_cost(plan, sync_cost=60.0): + tot = 0.0 + for p in plan.passes: + bf = dft_ops(p.R) + 3 * p.R + (6 * (p.R - 1) if p.Ns > 1 else 0) + tot += p.per * plan.T * bf + tot /= plan.n + return tot + sync_cost * (len(plan.passes) - 1) / plan.n * plan.T + + +def line_plans(n, emax=32, maxr=16): + """All (T, radices) plans for an n-point line with per-thread elements <= emax.""" + out = [] + for rad in _ordered_factorizations(n, 3 if n <= 4096 else 4, maxr): + cand_t = set() + for r in rad: + nb = n // r + for t in range(1, nb + 1): + if nb % t == 0: + cand_t.add(t) + for t in sorted(cand_t): + pl = LinePlan(n, t, rad) + if pl.emax <= emax: + out.append(pl) + return out + + +def choose_line_plan(n, t_target=None, emax=32, maxr=16): + """Pick a plan; with t_target, prefer T close to it (thread budget).""" + best = None + for pl in line_plans(n, emax, maxr): + c = line_plan_cost(pl) + if t_target is not None: + c *= 1.0 + 0.35 * abs(math.log2(pl.T / t_target)) + if best is None or c < best[0] - 1e-9: + best = (c, pl) + if best is None: + raise ValueError(f"no line plan for n={n}") + return best[1] + + +# ----------------------------------------------------------------------------- twiddles + + +def _root(e, n): + """exp(-2 pi i e / n) as float64 (cos, sin). The angle is rounded as + (-2 pi e) * fl(1 / n), which reproduces the judged tables bit for bit.""" + if e == 0: + return 1.0, 0.0 + a = (-2.0 * math.pi * e) * (1.0 / n) + return math.cos(a), math.sin(a) + + +class TwiddleBuilder: + """Accumulates twiddle tables (computed in float64, rounded once to fp32) + into one interleaved (re, im) buffer; each table starts 16-byte aligned.""" + + def __init__(self): + self.vals = [] # complex entries as (re, im) + + def add_roots(self, exps, n): + """Append exp(-2 pi i e / n) for e in exps; returns the complex offset.""" + if len(self.vals) % 2: + self.vals.append((0.0, 0.0)) + off = len(self.vals) + self.vals.extend(_root(e, n) for e in exps) + return off + + def add_pass_tables(self, plan): + """[Ns][R] table w_L^{k m}, L = Ns R, for every pass with Ns > 1.""" + for p in plan.passes: + if p.Ns > 1: + L = p.Ns * p.R + p.tw_off = self.add_roots([(k * m) % L for k in range(p.Ns) for m in range(p.R)], L) + + def array(self): + """float32 CPU tensor [re0, im0, re1, im1, ...].""" + return torch.tensor(self.vals, dtype=torch.float64).reshape(-1).to(torch.float32) + + +# ----------------------------------------------------------------------------- smem layouts + + +class Layout: + """offset(line, i) = line*ls + i + (i // pp) * pq (complex units).""" + + def __init__(self, n, lines, ls, pp, pq): + self.n, self.lines, self.ls, self.pp, self.pq = n, lines, ls, pp, pq + + def off(self, line, i): + return line * self.ls + i + (i // self.pp) * self.pq + + @property + def size(self): + return (self.lines - 1) * self.ls + self.off(0, self.n - 1) + 1 + + def __repr__(self): + return f"Layout(n={self.n}, lines={self.lines}, ls={self.ls}, pad={self.pq}/{self.pp})" + + +def lanes_groups(nthreads, fn): + """Simulate warps of `nthreads` threads; fn(tid) -> list of offsets (one per + access slot, same length for all threads or None). Returns half-warp groups.""" + per_tid = {tid: fn(tid) for tid in range(nthreads)} + nslots = max((len(v) for v in per_tid.values() if v), default=0) + groups = [] + for w0 in range(0, nthreads, 16): + for s in range(nslots): + g = [] + for tid in range(w0, min(w0 + 16, nthreads)): + v = per_tid[tid] + if v and s < len(v) and v[s] is not None: + g.append(v[s]) + if g: + groups.append(g) + return groups + + +def conflict_cost(groups): + """groups: list of lists of complex offsets accessed by one half-warp + (8-byte accesses). Returns total wavefronts.""" + tot = 0 + for g in groups: + banks = {} + for o in set(g): + banks.setdefault(o % 16, set()).add(o) + tot += max(len(v) for v in banks.values()) if banks else 0 + return tot + + +def choose_layout(n, lines, access_fn, min_ls=None): + """access_fn(layout) -> list of half-warp offset groups. Searches padding.""" + best = None + base = n if min_ls is None else max(n, min_ls) + for pp in sorted({8, 16, 32, n if n > 0 else 1}): + for pq in (0, 1, 2, 4): + if pq == 0 and pp != n: + continue + first = Layout(n, 1, 1, pp, pq).off(0, n - 1) + 1 + lo = max(base, first) + for extra in range(16): + ls = lo + extra + lay = Layout(n, lines, ls, pp, pq) + c = conflict_cost(access_fn(lay)) + key = (c, lay.size) + if best is None or key < best[0]: + best = (key, lay) + return best[1] diff --git a/build/torch-cuda/_r2c_f32/ptx.py b/build/torch-cuda/_r2c_f32/ptx.py new file mode 100644 index 0000000000000000000000000000000000000000..3597b81b99abf26d4bdfd9c358b777e339f01366 --- /dev/null +++ b/build/torch-cuda/_r2c_f32/ptx.py @@ -0,0 +1,252 @@ +"""Thin inline-PTX wrappers: vector global / shared memory accesses (optionally +predicated) and the inter-CTA counters of the x/y dataflow kernel.""" + +from cutlass import Float32, Int32, Int64 +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + + +def _iv(x, ty, loc=None, ip=None): + if isinstance(x, ir.Value): + return x + return ty(x).ir_value(loc=loc, ip=ip) + + +def _asm(rtypes, args, ptx, cons, side=True, loc=None, ip=None): + if not rtypes: + llvm.inline_asm(None, args, ptx, cons, has_side_effects=side, asm_dialect=0, loc=loc, ip=ip) + return [] + if len(rtypes) == 1: + r = llvm.inline_asm(rtypes[0], args, ptx, cons, has_side_effects=side, asm_dialect=0, loc=loc, ip=ip) + return [r] + st = ir.Type.parse("!llvm.struct<(" + ",".join(str(t) for t in rtypes) + ")>") + r = llvm.inline_asm(st, args, ptx, cons, has_side_effects=side, asm_dialect=0, loc=loc, ip=ip) + return [llvm.extractvalue(t, r, [i], loc=loc, ip=ip) for i, t in enumerate(rtypes)] + + +# ----------------------------------------------------------------------------- global + + +@dsl_user_op +def ldg2(addr, cop="", *, loc=None, ip=None): + r = _asm( + [T.f32()] * 2, + [_iv(addr, Int64, loc, ip)], + f"ld.global{cop}.v2.f32 {{$0, $1}}, [$2];", + "=f,=f,l", + loc=loc, + ip=ip, + ) + return (Float32(r[0]), Float32(r[1])) + + +@dsl_user_op +def ldg4(addr, cop="", *, loc=None, ip=None): + r = _asm( + [T.f32()] * 4, + [_iv(addr, Int64, loc, ip)], + f"ld.global{cop}.v4.f32 {{$0, $1, $2, $3}}, [$4];", + "=f,=f,=f,=f,l", + loc=loc, + ip=ip, + ) + return tuple(Float32(v) for v in r) + + +@dsl_user_op +def stg2(addr, v, cop="", *, loc=None, ip=None): + _asm( + [], + [_iv(addr, Int64, loc, ip), _iv(v[0], Float32, loc, ip), _iv(v[1], Float32, loc, ip)], + f"st.global{cop}.v2.f32 [$0], {{$1, $2}};", + "l,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stg4(addr, v, cop="", *, loc=None, ip=None): + _asm( + [], + [_iv(addr, Int64, loc, ip)] + [_iv(x, Float32, loc, ip) for x in v], + f"st.global{cop}.v4.f32 [$0], {{$1, $2, $3, $4}};", + "l,f,f,f,f", + loc=loc, + ip=ip, + ) + + +# ----------------------------------------------------------------------------- shared + + +@dsl_user_op +def lds2(addr, *, loc=None, ip=None): + r = _asm([T.f32()] * 2, [_iv(addr, Int32, loc, ip)], "ld.shared.v2.f32 {$0, $1}, [$2];", "=f,=f,r", loc=loc, ip=ip) + return (Float32(r[0]), Float32(r[1])) + + +@dsl_user_op +def sts2(addr, v, *, loc=None, ip=None): + _asm( + [], + [_iv(addr, Int32, loc, ip), _iv(v[0], Float32, loc, ip), _iv(v[1], Float32, loc, ip)], + "st.shared.v2.f32 [$0], {$1, $2};", + "r,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def lds4(addr, *, loc=None, ip=None): + r = _asm( + [T.f32()] * 4, + [_iv(addr, Int32, loc, ip)], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,r", + loc=loc, + ip=ip, + ) + return tuple(Float32(v) for v in r) + + +@dsl_user_op +def sts4(addr, v, *, loc=None, ip=None): + _asm( + [], + [_iv(addr, Int32, loc, ip)] + [_iv(x, Float32, loc, ip) for x in v], + "st.shared.v4.f32 [$0], {$1, $2, $3, $4};", + "r,f,f,f,f", + loc=loc, + ip=ip, + ) + + +# ----------------------------------------------------------------------------- predicated stores + + +@dsl_user_op +def stg2_if(pv, addr, v, cop="", *, loc=None, ip=None): + """Store if pv != 0.""" + _asm( + [], + [_iv(pv, Int32, loc, ip), _iv(addr, Int64, loc, ip), _iv(v[0], Float32, loc, ip), _iv(v[1], Float32, loc, ip)], + "{ .reg .pred p; setp.ne.s32 p, $0, 0; @p st.global" + cop + ".v2.f32 [$1], {$2, $3}; }", + "r,l,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stg4_if(pv, addr, v, cop="", *, loc=None, ip=None): + _asm( + [], + [_iv(pv, Int32, loc, ip), _iv(addr, Int64, loc, ip)] + [_iv(x, Float32, loc, ip) for x in v], + "{ .reg .pred p; setp.ne.s32 p, $0, 0; @p st.global" + cop + ".v4.f32 [$1], {$2, $3, $4, $5}; }", + "r,l,f,f,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def sts2_if(pv, addr, v, *, loc=None, ip=None): + _asm( + [], + [_iv(pv, Int32, loc, ip), _iv(addr, Int32, loc, ip), _iv(v[0], Float32, loc, ip), _iv(v[1], Float32, loc, ip)], + "{ .reg .pred p; setp.ne.s32 p, $0, 0; @p st.shared.v2.f32 [$1], {$2, $3}; }", + "r,r,f,f", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def sts4_if(pv, addr, v, *, loc=None, ip=None): + _asm( + [], + [_iv(pv, Int32, loc, ip), _iv(addr, Int32, loc, ip)] + [_iv(x, Float32, loc, ip) for x in v], + "{ .reg .pred p; setp.ne.s32 p, $0, 0; @p st.shared.v4.f32 [$1], {$2, $3, $4, $5}; }", + "r,r,f,f,f,f", + loc=loc, + ip=ip, + ) + + +# ----------------------------------------------------------------------------- inter-CTA sync + + +@dsl_user_op +def atom_add_gpu(pv, addr, val, *, loc=None, ip=None): + """if pv: old = atom.acq_rel.gpu.add(addr, val) (else returns 0).""" + r = _asm( + [T.i32()], + [_iv(pv, Int32, loc, ip), _iv(addr, Int64, loc, ip), _iv(val, Int32, loc, ip)], + "{ .reg .pred p; setp.ne.s32 p, $1, 0; mov.s32 $0, 0; @p atom.acq_rel.gpu.global.add.s32 $0, [$2], $3; }", + "=r,r,l,r", + loc=loc, + ip=ip, + ) + return Int32(r[0]) + + +@dsl_user_op +def release_add_gpu(pv, addr, val, *, loc=None, ip=None): + """if pv: fence.acq_rel.gpu; red.release.gpu.add(addr, val).""" + _asm( + [], + [_iv(pv, Int32, loc, ip), _iv(addr, Int64, loc, ip), _iv(val, Int32, loc, ip)], + "{ .reg .pred p; setp.ne.s32 p, $0, 0; @p fence.acq_rel.gpu; @p red.release.gpu.global.add.s32 [$1], $2; }", + "r,l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def spin_until_ge(pv, addr, target, *, loc=None, ip=None): + """if pv: spin (ld.acquire.gpu) until *addr >= target.""" + _asm( + [], + [_iv(pv, Int32, loc, ip), _iv(addr, Int64, loc, ip), _iv(target, Int32, loc, ip)], + "{ .reg .pred p, q; .reg .s32 v; setp.ne.s32 p, $0, 0; @!p bra SPIN_DONE_${:uid}; " + "SPIN_LOOP_${:uid}: ld.acquire.gpu.global.s32 v, [$1]; setp.lt.s32 q, v, $2; " + "@!q bra SPIN_DONE_${:uid}; nanosleep.u32 32; bra SPIN_LOOP_${:uid}; " + "SPIN_DONE_${:uid}: }", + "r,l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def st_global_s32_if(pv, addr, val, *, loc=None, ip=None): + _asm( + [], + [_iv(pv, Int32, loc, ip), _iv(addr, Int64, loc, ip), _iv(val, Int32, loc, ip)], + "{ .reg .pred p; setp.ne.s32 p, $0, 0; @p st.relaxed.gpu.global.s32 [$1], $2; }", + "r,l,r", + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def lds_s32(addr, *, loc=None, ip=None): + r = _asm([T.i32()], [_iv(addr, Int32, loc, ip)], "ld.shared.s32 $0, [$1];", "=r,r", loc=loc, ip=ip) + return Int32(r[0]) + + +@dsl_user_op +def sts_s32_if(pv, addr, val, *, loc=None, ip=None): + _asm( + [], + [_iv(pv, Int32, loc, ip), _iv(addr, Int32, loc, ip), _iv(val, Int32, loc, ip)], + "{ .reg .pred p; setp.ne.s32 p, $0, 0; @p st.shared.s32 [$1], $2; }", + "r,r,r", + loc=loc, + ip=ip, + ) diff --git a/build/torch-cuda/_r2c_f32/stockham.py b/build/torch-cuda/_r2c_f32/stockham.py new file mode 100644 index 0000000000000000000000000000000000000000..516a3a960388cb92f73dc7e6c9014bf211e48e94 --- /dev/null +++ b/build/torch-cuda/_r2c_f32/stockham.py @@ -0,0 +1,308 @@ +"""Trace-time Stockham FFT emitter shared by all fast-path kernels. + +emit_fft runs the passes of a LinePlan over G lines per thread: pass-1 inputs +come from a caller-supplied load (or preloaded registers), intermediate passes +exchange through a padded shared-memory Layout, and the last pass hands its +outputs to a store callback (or to `final` for the real-to-complex split). +Butterflies are assigned per pass: "cyclic" (j = t + T*b) or "mirror" (the two +butterflies of a thread are j and nb - j, used by the x pass so a thread owns +both Z[k] and Z[M-k]). Pass twiddles come from TwSmem, the [Ns][R] tables of +the global twiddle buffer staged in shared memory. access_offsets replays the +same index math on Python ints for the bank-conflict layout search. +""" + +import math + +import cutlass +from cutlass import Float32, Int32, Int64 + +from .fft_core import as_f32, cmul, dft +from .ptx import ldg4, lds2, lds4, sts2, sts2_if, sts4 + +# ----------------------------------------------------------------------------- DSL helpers + + +def _pow2(n): + return n > 0 and (n & (n - 1)) == 0 + + +def idiv(x, d): + if d == 1: + return x + if _pow2(d): + return x >> int(math.log2(d)) + return x // d + + +def imod(x, d): + if d == 1: + return Int32(0) * x + if _pow2(d): + return x & (d - 1) + return x % d + + +def b2i(c): + return Int32(cutlass.select_(c, Int32(1), Int32(0))) + + +def fsel(c, a, b): + return Float32(cutlass.select_(c, as_f32(a), as_f32(b))) + + +def csel(c, a, b): + return (fsel(c, a[0], b[0]), fsel(c, a[1], b[1])) + + +def and_pv(a, b): + """AND of two optional Int32 predicates (None = always true).""" + if a is None: + return b + if b is None: + return a + return a & b + + +def lay_off(lay, line, idx): + """Runtime complex offset in a Layout (line, idx may be runtime Int32).""" + o = idx + if lay.pq: + o = o + idiv(idx, lay.pp) * lay.pq + if isinstance(line, int): + if line: + o = o + line * lay.ls + else: + o = o + line * lay.ls + return o + + +# ----------------------------------------------------------------------------- twiddles in smem + + +class TwSmem: + """Twiddle tables staged in shared memory. + + tables: list of (global complex offset, complex count), each 16B aligned. + """ + + def __init__(self, tables): + self.tables = [] + self.soff = {} + off = 0 + for goff, cnt in tables: + if goff in self.soff: + continue + self.soff[goff] = off + self.tables.append((goff, cnt, off)) + off += cnt + (cnt % 2) + self.size = off * 8 # bytes + + def stage(self, sbase, tid, nt, twb): + """Copy tables global -> smem (caller must barrier before use).""" + self.sbase = sbase + for goff, cnt, off in self.tables: + chunks = (cnt + 1) // 2 + for i in range(-(-chunks // nt)): + c = tid + i * nt if i else tid + if (i + 1) * nt > chunks: + c = cutlass.min(c, Int32(chunks - 1)) + v = ldg4(twb + Int64(c * 2 + goff) * 8, ".nc") + sts4(sbase + (c * 2 + off) * 8, v) + + def row(self, p, k): + """Twiddles w_{Ns*R}^{k*m}, m = 0..R-1, of pass p for runtime k.""" + R = p.R + base = self.sbase + (k * R + self.soff[p.tw_off]) * 8 + tw = [(1.0, 0.0)] + m = 1 + while m < R: + if R % 2 == 0 and m % 2 == 0 and m + 1 < R: + v = lds4(base + m * 8) + tw.append((v[0], v[1])) + tw.append((v[2], v[3])) + m += 2 + else: + tw.append(lds2(base + m * 8)) + m += 1 + return tw + + def elem(self, goff, k): + return lds2(self.sbase + (k + self.soff[goff]) * 8) + + @staticmethod + def for_plans(plans, extra=()): + tabs = [] + for pl in plans: + for p in pl.passes: + if p.Ns > 1: + tabs.append((p.tw_off, p.Ns * p.R)) + tabs.extend(extra) + return TwSmem(tabs) + + +# ----------------------------------------------------------------------------- butterfly assignment + + +def _sel(c, a, b): + """select usable at trace time with python ints or DSL values.""" + if isinstance(c, bool): + return a if c else b + return Int32(cutlass.select_(c, Int32(a) if isinstance(a, int) else a, Int32(b) if isinstance(b, int) else b)) + + +def bfly(p, kind, t, b, T): + """Butterfly index for thread t, slot b of pass p. + + cyclic : j = t + T*b + mirror : per == 2, T == nb/2: slot0 -> t, slot1 -> nb - t (t = 0 -> nb/2) + """ + if kind == "cyclic": + return t + T * b if b else t + if kind == "mirror": + if b == 0: + return t + if isinstance(t, int): + return p.nb // 2 if t == 0 else p.nb - t + return _sel(t == 0, p.nb // 2, Int32(p.nb) - t) + raise ValueError(kind) + + +def pass1_indices(plan, kinds, t): + """Pass-1 input element indices in emission order (b-major, then m).""" + p = plan.passes[0] + out = [] + for b in range(p.per): + j = bfly(p, kinds[0], t, b, plan.T) + guard = (not p.full) and kinds[0] == "cyclic" and b == p.per - 1 + jl = cutlass.min(j, Int32(p.nb - 1)) if guard else j + for m in range(p.R): + out.append(jl + m * p.nb if m else jl) + return out + + +# ----------------------------------------------------------------------------- emitter + + +def emit_fft( + plan, + kinds, + t, + G, + load_in, + store_out, + xaddr, + twb, + sync, + first_hazard=False, + finish=None, + final=None, + pre=None, + last_hazard=False, +): + """Stockham FFT over G lines handled by thread t (runtime index in [0, T)). + + load_in(idx) -> list of G complex: pass-1 input (unless `pre`) + pre pass-1 inputs already in registers, in pass1_indices order + finish(raw) optional transform of each pass-1 load (e.g. a twiddle) + store_out(idx, vals, pv) final outputs (natural order); pv: Int32 predicate or None + final(items) replaces store_out: called with the last pass results, + items = list of (j, [ys per line], pv), ys[m] = output j + m*nb + xaddr(g, idx) Int32 shared address of exchange element idx of line g + twb TwSmem with the pass twiddle tables staged + sync() barrier over the participating threads; first_hazard / + last_hazard add one before the first / last pass's stores + (when the exchange buffer aliases the input / output) + """ + T = plan.T + npass = len(plan.passes) + for pi, p in enumerate(plan.passes): + first = pi == 0 + last = pi == npass - 1 + R, Ns, nb = p.R, p.Ns, p.nb + kind = kinds[pi] + items = [] + for b in range(p.per): + j = bfly(p, kind, t, b, T) + guard = (not p.full) and kind == "cyclic" and b == p.per - 1 + jl = cutlass.min(j, Int32(nb - 1)) if guard else j + raws = [] + for m in range(R): + idx = jl + m * nb if m else jl + if first: + raws.append(pre[b * R + m] if pre is not None else load_in(idx)) + else: + raws.append([lds2(xaddr(g, idx)) for g in range(G)]) + items.append([j, jl, guard, raws]) + # finish loads only after all of this pass's loads have been issued + for it in items: + xs = [[None] * R for _ in range(G)] + for m in range(R): + vals = finish(it[3][m]) if (first and finish is not None) else it[3][m] + for g in range(G): + xs[g][m] = vals[g] + it[3] = xs + if (first and first_hazard) or (not first and not last) or (last and last_hazard): + sync() + outs = [] + for j, jl, guard, xs in items: + if Ns > 1: + k = imod(jl, Ns) + tw = twb.row(p, k) + for g in range(G): + for m in range(1, R): + xs[g][m] = cmul(xs[g][m], tw[m]) + ys = [dft(xs[g], -1) for g in range(G)] + pv = b2i(j < nb) if guard else None + if last: + if final is not None: + outs.append((jl, ys, pv)) + else: + for m in range(R): + store_out(jl + m * nb if m else jl, [ys[g][m] for g in range(G)], pv) + else: + ob = idiv(jl, Ns) * (Ns * R) + imod(jl, Ns) if Ns > 1 else jl * R + for m in range(R): + oi = ob + m * Ns if m else ob + for g in range(G): + a = xaddr(g, oi) + v = (as_f32(ys[g][m][0]), as_f32(ys[g][m][1])) + if pv is None: + sts2(a, v) + else: + sts2_if(pv, a, v) + if last and final is not None: + final(outs) + if not last: + sync() + + +def access_offsets(plan, kinds, t, line_of_g, lay, G): + """Static exchange-buffer offsets for thread t (layout search).""" + offs = [] + npass = len(plan.passes) + for pi, p in enumerate(plan.passes): + first = pi == 0 + last = pi == npass - 1 + for b in range(p.per): + j = min(bfly(p, kinds[pi], t, b, plan.T), p.nb - 1) + if not first: + for m in range(p.R): + for g in range(G): + offs.append(lay.off(line_of_g(g), j + m * p.nb)) + if not last: + ob = (j // p.Ns) * p.Ns * p.R + j % p.Ns + for m in range(p.R): + for g in range(G): + offs.append(lay.off(line_of_g(g), ob + m * p.Ns)) + return offs + + +def pass1_static(plan, kinds, t): + """pass1_indices on Python ints (clamped), for the staging-stride search.""" + p = plan.passes[0] + out = [] + for b in range(p.per): + j = min(bfly(p, kinds[0], t, b, plan.T), p.nb - 1) + for m in range(p.R): + out.append(j + m * p.nb) + return out diff --git a/build/torch-cuda/_r2c_f64/__init__.py b/build/torch-cuda/_r2c_f64/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..43a032e1bd70bc454aa5d41cb4c0bf561c8dae4d --- /dev/null +++ b/build/torch-cuda/_r2c_f64/__init__.py @@ -0,0 +1,3 @@ +from .kernel import run + +__all__ = ["run"] diff --git a/build/torch-cuda/_r2c_f64/dsl_ops.py b/build/torch-cuda/_r2c_f64/dsl_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..f6fe7360fc74bd46dc0ecdaf23fe8dd8469044bb --- /dev/null +++ b/build/torch-cuda/_r2c_f64/dsl_ops.py @@ -0,0 +1,306 @@ +"""CuTe-DSL building blocks used while the FFT kernels are traced. + +f64x2 vector loads/stores, global loads/stores with an L2 evict-first policy, float64 +helpers that run on the integer pipes (sign flip, exact halving and ``emu_addsub``, a +bit-exact IEEE add/sub pair), the ``DSLBackend`` through which the codelets of +fft_codelets.py emit instructions, and the counter spin-wait of the X->Y pass handoff. +""" + +import cutlass +from cutlass import Float64, cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith, scf, vector + + +def _v2t(): + return ir.VectorType.get([2], Float64.mlir_type) + + +def vload2(ptr): + v = cute.arch.load(ptr, _v2t()) + return Float64(vector.extract(v, [], [0])), Float64(vector.extract(v, [], [1])) + + +def vstore2(ptr, a, b): + v = vector.from_elements(_v2t(), [a.ir_value(), b.ir_value()]) + cute.arch.store(ptr, v) + + +_EF_POLICY = "createpolicy.fractional.L2::evict_first.b64 pol, 1.0;" + + +def load_f64_ef(ptr): + """Scalar f64 global load with an L2 evict-first policy (input that is read once).""" + v = cute.arch.inline_ptx( + "{\n\t.reg .b64 pol;\n\t" + _EF_POLICY + "\n\tld.global.L2::cache_hint.b64 {$w0}, [{$r0}], pol;\n\t}", + write_only_types=[cutlass.Int64], + read_only_args=[ptr.toint()], + ) + return Float64(arith.bitcast(Float64.mlir_type, v.ir_value())) + + +def vload2_ef(ptr): + """f64x2 global load with an L2 evict-first policy.""" + a, b = cute.arch.inline_ptx( + "{\n\t.reg .b64 pol;\n\t" + + _EF_POLICY + + "\n\tld.global.L2::cache_hint.v2.b64 {{$w0}, {$w1}}, [{$r0}], pol;\n\t}", + write_only_types=[cutlass.Int64, cutlass.Int64], + read_only_args=[ptr.toint()], + ) + return Float64(arith.bitcast(Float64.mlir_type, a.ir_value())), Float64( + arith.bitcast(Float64.mlir_type, b.ir_value()) + ) + + +def vstore2_ef(ptr, a, b): + """f64x2 global store with an L2 evict-first policy (final output that is not re-read).""" + cute.arch.inline_ptx( + "{\n\t.reg .b64 pol;\n\t" + + _EF_POLICY + + "\n\tst.global.L2::cache_hint.v2.b64 [{$r0}], {{$r1}, {$r2}}, pol;\n\t}", + read_only_args=[ + ptr.toint(), + cutlass.Int64(arith.bitcast(cutlass.Int64.mlir_type, a.ir_value())), + cutlass.Int64(arith.bitcast(cutlass.Int64.mlir_type, b.ir_value())), + ], + ) + + +def _when(cond, fn): + """Emit fn() under a runtime condition (None means always).""" + if cond is None: + fn() + else: + cutlass.if_generate(cond, fn) + + +def _and(a, b): + if a is None: + return b + if b is None: + return a + return a & b + + +class DSLBackend: + def __init__(self, kmap, mz): + self.kmap = kmap + self.mz = mz + + def fma(self, a, b, c): + return cute.math.fma(a, b, c) + + def neg(self, a): + return -a + + def const(self, v): + return self.kmap[float(v)] + + def mzero(self): + return self.mz + + def select(self, f, a, b): + return Float64(cutlass.select_(f, a, b)) + + def bneg(self, a): + return bits_neg(a) + + def addsub(self, a, b): + return emu_addsub(a, b) + + +def bits_neg(x): + """-x via an integer sign-bit flip (a standalone neg.f64 costs a DADD).""" + xi = cutlass.Int64(arith.bitcast(cutlass.Int64.mlir_type, x.ir_value())) + yi = cute.arch.inline_ptx( + "xor.b64 {$w0}, {$r0}, 0x8000000000000000;", write_only_types=[cutlass.Int64], read_only_args=[xi] + ) + return Float64(arith.bitcast(Float64.mlir_type, yi.ir_value())) + + +# Bit-exact fp64 add/sub on the integer pipes. On sm_120 the FP64 pipe issues one warp +# instruction per 64 cycles per sub-partition and ptxas pads every DFMA with NOPs; an IEEE +# round-to-nearest-even add/sub pair (shared exponent alignment, 64-bit integer ops) runs on the +# ALU/FMA pipes in those idle issue slots, so unit-coefficient butterflies can bypass the FP64 pipe. +# Results are identical to DADD/DSUB for normal operands; subnormal operands/results flush to zero +# (the sign of a zero result may differ from IEEE); Inf/NaN are not propagated. +def _emu_ptx_addsub(half=False, A="{$r0}", B="{$r1}", S="{$w0}", D="{$w1}"): + """PTX for S = A + B, D = A - B (each halved if `half`). Mantissas sit at bits [61:9] + (9 guard bits, sticky jammed into bit 0); the add path normalises by at most one bit, the + subtract path with a 64-bit count-leading-zeros; RNE rounding on 32-bit halves.""" + hdec = 1 if half else 0 + return f"""{{ + .reg .b64 aa, ab, ax, ay, mx, my, ys, t64, mp, mm, np, nm; + .reg .b32 al, ah, bl, bh, axl, axh, ayl, ayh, ex, ey, dd, mxl, mxh, myl, myh, t, u, sh, lzm, shm; + .reg .b32 pl, ph, ql, qh, rl, rh, sl, sh2, s1, swm, e1, e2; + .reg .pred sw, same, zx, zy, st, p62, zm; + mov.b64 {{al, ah}}, {A}; + mov.b64 {{bl, bh}}, {B}; + and.b32 t, ah, 0x7FFFFFFF; + mov.b64 aa, {{al, t}}; + and.b32 t, bh, 0x7FFFFFFF; + mov.b64 ab, {{bl, t}}; + setp.gt.u64 sw, ab, aa; + max.u64 ax, aa, ab; + min.u64 ay, aa, ab; + mov.b64 {{axl, axh}}, ax; + mov.b64 {{ayl, ayh}}, ay; + shr.u32 ex, axh, 20; + shr.u32 ey, ayh, 20; + setp.eq.u32 zx, ex, 0; + setp.eq.u32 zy, ey, 0; + shf.l.clamp.b32 mxh, axl, axh, 9; + shl.b32 mxl, axl, 9; + lop3.b32 mxh, mxh, 0x1FFFFFFF, 0x20000000, 0xEA; + shf.l.clamp.b32 myh, ayl, ayh, 9; + shl.b32 myl, ayl, 9; + lop3.b32 myh, myh, 0x1FFFFFFF, 0x20000000, 0xEA; + mov.b64 mx, {{mxl, mxh}}; + mov.b64 my, {{myl, myh}}; + selp.b64 mx, 0, mx, zx; + selp.b64 my, 0, my, zy; + sub.s32 dd, ex, ey; + min.u32 dd, dd, 63; + shr.b64 ys, my, dd; + shl.b64 t64, ys, dd; + setp.ne.u64 st, t64, my; + selp.b64 t64, 1, 0, st; + or.b64 ys, ys, t64; + add.s64 mp, mx, ys; + sub.s64 mm, mx, ys; + // add path: mp in [2^61, 2^63) -> MSB to bit 62 + mov.b64 {{pl, ph}}, mp; + setp.lt.u32 p62, ph, 0x40000000; + shl.b64 t64, mp, 1; + selp.b64 np, t64, mp, p62; + selp.u32 sh, 1, 0, p62; + // sub path: general normalization + clz.b64 lzm, mm; + sub.u32 shm, lzm, 1; + shl.b64 nm, mm, shm; + setp.eq.u64 zm, mm, 0; + // round both (RNE at bit 10) and pack + mov.b64 {{pl, ph}}, np; + bfe.u32 t, pl, 10, 1; + add.u32 t, t, 0x1FF; + add.cc.u32 pl, pl, t; + addc.u32 ph, ph, 0; + shf.r.clamp.b32 rl, pl, ph, 10; + shr.u32 rh, ph, 10; + sub.u32 e1, ex, sh; + sub.u32 e1, e1, {hdec}; + shl.b32 e1, e1, 20; + add.u32 rh, rh, e1; + selp.u32 rh, 0, rh, zx; + mov.b64 {{ql, qh}}, nm; + bfe.u32 t, ql, 10, 1; + add.u32 t, t, 0x1FF; + add.cc.u32 ql, ql, t; + addc.u32 qh, qh, 0; + shf.r.clamp.b32 sl, ql, qh, 10; + shr.u32 sh2, qh, 10; + sub.u32 e2, ex, shm; + sub.u32 e2, e2, {hdec}; + shl.b32 e2, e2, 20; + add.u32 sh2, sh2, e2; + selp.u32 sh2, 0, sh2, zm; + // assign magnitudes: a+b takes the add path when the signs agree + xor.b32 u, ah, bh; + setp.ge.s32 same, u, 0; + selp.b32 pl, rl, sl, same; + selp.b32 ph, rh, sh2, same; + selp.b32 ql, sl, rl, same; + selp.b32 qh, sh2, rh, same; + // signs: a+b takes the sign of the larger operand; a-b flips it when |b| > |a| + selp.b32 s1, bh, ah, sw; + selp.b32 swm, 0x80000000, 0, sw; + lop3.b32 ph, ph, s1, 0x80000000, 0xF8; + xor.b32 t, s1, swm; + lop3.b32 qh, qh, t, 0x80000000, 0xF8; + mov.b64 {S}, {{pl, ph}}; + mov.b64 {D}, {{ql, qh}}; + }}""" + + +_PTX_ADDSUB = _emu_ptx_addsub(False) +_PTX_ADDSUB_HALF = _emu_ptx_addsub(True) + + +def _as_i64(x): + return cutlass.Int64(arith.bitcast(cutlass.Int64.mlir_type, x.ir_value())) + + +def _as_f64(v): + return Float64(arith.bitcast(Float64.mlir_type, v.ir_value())) + + +def emu_addsub(x, y, half=False): + """(x + y, x - y) (halved if `half`), bit-exact, on the integer pipes.""" + s_, d_ = cute.arch.inline_ptx( + _PTX_ADDSUB_HALF if half else _PTX_ADDSUB, + write_only_types=[cutlass.Int64, cutlass.Int64], + read_only_args=[_as_i64(x), _as_i64(y)], + ) + return _as_f64(s_), _as_f64(d_) + + +def exact_half(x): + """x * 0.5 by decrementing the exponent field (integer ops, no FP64 instruction). + + Exact for normal results; |x| < 2^-1021 flushes to a signed zero (FTZ), and + Inf/NaN pass through unchanged.""" + xi = cutlass.Int64(arith.bitcast(cutlass.Int64.mlir_type, x.ir_value())) + yi = cute.arch.inline_ptx( + "{\n\t.reg .b32 lo, hi, e, h2, hs, h3, l3;\n\t.reg .pred p, q;\n\t" + "mov.b64 {lo, hi}, {$r0};\n\t" + "and.b32 e, hi, 0x7FF00000;\n\t" + "setp.gt.u32 p, e, 0x00100000;\n\t" + "setp.eq.u32 q, e, 0x7FF00000;\n\t" + "sub.s32 h2, hi, 0x00100000;\n\t" + "and.b32 hs, hi, 0x80000000;\n\t" + "selp.b32 h3, h2, hs, p;\n\t" + "selp.b32 l3, lo, 0, p;\n\t" + "selp.b32 h3, hi, h3, q;\n\t" + "selp.b32 l3, lo, l3, q;\n\t" + "mov.b64 {$w0}, {l3, h3};\n\t}", + write_only_types=[cutlass.Int64], + read_only_args=[xi], + ) + return Float64(arith.bitcast(Float64.mlir_type, yi.ir_value())) + + +def vload_i4(ptr): + v = cute.arch.load(ptr, ir.VectorType.get([4], cutlass.Int32.mlir_type)) + return [cutlass.Int32(vector.extract(v, [], [j])) for j in range(4)] + + +def _while(init, cond_fn, body_fn): + ir_vals = [v.ir_value() for v in init] + types = [v.type for v in ir_vals] + w = scf.WhileOp(types, ir_vals) + before = w.before.blocks.append(*types) + after = w.after.blocks.append(*types) + with ir.InsertionPoint(before): + args = [cutlass.Int32(a) for a in before.arguments] + scf.ConditionOp(cutlass.Boolean(cond_fn(*args)).ir_value(), list(before.arguments)) + with ir.InsertionPoint(after): + args = [cutlass.Int32(a) for a in after.arguments] + scf.YieldOp([v.ir_value() for v in body_fn(*args)]) + + +def flag_spin(ptr, target, sleep_ns=128): + """Spin (relaxed loads + nanosleep) until *ptr >= target.""" + + def cond(x): + return cute.arch.load(ptr, cutlass.Int32, sem="relaxed", scope="gpu") < target + + def body(x): + cute.arch.inline_ptx("nanosleep.u32 {$r0};", read_only_args=[cutlass.Int32(sleep_ns)]) + return [x] + + _while([cutlass.Int32(0)], cond, body) + + +def flag_acquire(ptr): + cute.arch.load(ptr, cutlass.Int32, sem="acquire", scope="gpu") diff --git a/build/torch-cuda/_r2c_f64/fft_codelets.py b/build/torch-cuda/_r2c_f64/fft_codelets.py new file mode 100644 index 0000000000000000000000000000000000000000..3608ef13da93bc74ee0237fb2b5a3cb9119a0cab --- /dev/null +++ b/build/torch-cuda/_r2c_f64/fft_codelets.py @@ -0,0 +1,581 @@ +"""FFT codelets built only from fused multiply-adds. + +Every floating-point operation goes through the backend's ``fma``; additions and +plain multiplies are FMAs against register constants (1.0 and -0.0). On +sm_120 a DFMA whose operands are all registers issues at the full FP64 rate, +while DADD/DMUL and FMAs with immediate operands issue ~16% slower. + +A complex value is ``CVal(re, im, s)`` meaning ``s * (re + i*im)``, where the +scale ``s`` is a compile-time float. Constant twiddles are applied in +Linzer-Feig form (2 FMAs plus a scale change) and the scales are absorbed by +the FMA coefficients of later additions, so a codelet's outputs end at scale 1 +whenever its first input does. + +Butterflies go through ``Ops.bfly`` (a + w b, a - w b with w in {+-1, +-i}); when an +``EmuPolicy`` is attached, the real or imaginary part (or both) of selected unit-ratio pairs +is computed by the backend's ``addsub`` -- a bit-exact IEEE add/sub pair on the integer +pipes -- instead of two FMAs. + +Backends implement ``fma(a, b, c)`` (a*b+c), ``neg(a)``, ``const(v)`` (a register holding +the positive float ``v``) and ``addsub(a, b)`` (a+b, a-b). +""" + +import math + +EPS_EXACT = 1e-15 + + +class CVal: + __slots__ = ("im", "re", "s") + + def __init__(self, re, im, s=1.0): + self.re = re + self.im = im + self.s = s + + +def _is_one(v): + return (not isinstance(v, RScale)) and abs(v - 1.0) < 1e-15 + + +class RScale: + """Runtime-dependent scale c * prod_m sigma_m**e_m (sigma_m: per-lane twiddle scales).""" + + __slots__ = ("c", "e") + + def __init__(self, c, e): + self.c = c + self.e = tuple(e) + + def __repr__(self): + return f"RScale({self.c:.6g}, {self.e})" + + +def _norm(c, e): + if all(v == 0 for v in e): + return c + return RScale(c, e) + + +def _parts(s): + if isinstance(s, RScale): + return s.c, s.e + return s, () + + +def s_mul(a, b): + ca, ea = _parts(a) + cb, eb = _parts(b) + n = max(len(ea), len(eb)) + ea = ea + (0,) * (n - len(ea)) + eb = eb + (0,) * (n - len(eb)) + return _norm(ca * cb, tuple(x + y for x, y in zip(ea, eb))) + + +def s_div(a, b): + ca, ea = _parts(a) + cb, eb = _parts(b) + n = max(len(ea), len(eb)) + ea = ea + (0,) * (n - len(ea)) + eb = eb + (0,) * (n - len(eb)) + return _norm(ca / cb, tuple(x - y for x, y in zip(ea, eb))) + + +def s_is_unit(s): + return (not isinstance(s, RScale)) and abs(abs(s) - 1.0) < 1e-15 + + +class RtRegistry: + """Runtime scale ratios a codelet needs, keyed by (|c|, exponents). + + In record mode get() remembers the key and returns a dummy value; in emit + mode it returns the backend register loaded for that key.""" + + def __init__(self, values=None, keys=None): + # keys: list of (exact |c|, exponents); matching uses a rounded |c| + self.keys = [] if keys is None else list(keys) + self.index = {self._match(c, e): i for i, (c, e) in enumerate(self.keys)} + self.values = values + + @staticmethod + def _match(c, e): + return (float(f"{abs(c):.12e}"), tuple(e)) + + def get(self, rs): + k = self._match(rs.c, rs.e) + if self.values is None: + if k not in self.index: + self.index[k] = len(self.keys) + self.keys.append((abs(rs.c), rs.e)) + return 1.0 + return self.values[self.index[k]] + + +class EmuPolicy: + """Chooses which unit-coefficient add/sub pairs run as integer-pipe emulation. + + Count mode (ratio is None): records the eligible real ops Q, emulates none. + Emit mode: emulates about `ratio` of the eligible ops, evenly spread in program order. + take() returns a part mask for a complex butterfly (1 = real part, 2 = imaginary part, + 3 = both): emulating one part of each butterfly keeps integer and FP64 work side by side + at every level of the codelet, so ptxas interleaves the two pipes.""" + + def __init__(self, ratio=None): + self.ratio = ratio + self.q = 0 + self.Q = 0 + self.done = 0 + self.flip = 0 + + def take(self, k): + self.q += k + if self.ratio is None: + self.Q = self.q + return 0 + want = self.q * self.ratio - self.done + if k == 4: + parts = int(max(0, min(2, math.floor(want / 2 + 0.5)))) + self.done += 2 * parts + if parts == 2: + return 3 + if parts == 1: + self.flip ^= 1 + return 1 if self.flip else 2 + return 0 + if want >= k / 2: + self.done += k + return 3 + return 0 + + +class Ops: + """Complex helpers over a scalar backend.""" + + def __init__(self, be, rt=None, emu=None): + self.be = be + self.rt = rt + self.emu = emu + + def _unit_ratio(self, a, b): + """+1/-1 if b.s / a.s is exactly a unit compile-time ratio, else 0.""" + if isinstance(a.s, RScale) or isinstance(b.s, RScale): + if ( + isinstance(a.s, RScale) + and isinstance(b.s, RScale) + and a.s.e == b.s.e + and abs(abs(b.s.c / a.s.c) - 1.0) < 1e-15 + ): + return 1 if b.s.c / a.s.c > 0 else -1 + return 0 + r = b.s / a.s + if abs(abs(r) - 1.0) < 1e-15: + return 1 if r > 0 else -1 + return 0 + + def bfly(self, a, b, rot=0): + """(a + i^rot b, a - i^rot b); parts of unit-ratio pairs may run as emulated add/sub pairs.""" + sig = self._unit_ratio(a, b) if self.emu is not None else 0 + mode = self.emu.take(4) if sig != 0 else 0 + if mode: + (bre, bres), (bim, bims) = self._rot(b, rot) + out = [] + for bit, x, (y, ys) in ((1, a.re, (bre, bres)), (2, a.im, (bim, bims))): + c = sig * ys + if mode & bit: + sm, df = self.be.addsub(x, y) + out.append((sm, df) if c > 0 else (df, sm)) + else: + out.append((self.axpy(x, 1, y, float(c)), self.axpy(x, 1, y, float(-c)))) + (sr, dr), (si, di) = out + return CVal(sr, si, a.s), CVal(dr, di, a.s) + return self.add(a, b, 1, rot), self.sub(a, b, rot) + + def axpy(self, x, xs, y, r): + be = self.be + xv = x if xs > 0 else be.neg(x) + if isinstance(r, RScale): + v = self.rt.get(r) + yv = be.neg(y) if r.c < 0 else y + return be.fma(yv, v, xv) + if r == 0.0: + return xv + if r < 0: + yv = be.neg(y) + r = -r + else: + yv = y + return be.fma(yv, be.const(1.0 if _is_one(r) else r), xv) + + @staticmethod + def _sneg(r): + if isinstance(r, RScale): + return RScale(-r.c, r.e) + return -r + + # rotation by i^rot of value components: returns ((re_v, re_sign), (im_v, im_sign)) + @staticmethod + def _rot(v, rot): + rot &= 3 + if rot == 0: + return (v.re, 1), (v.im, 1) + if rot == 1: + return (v.im, -1), (v.re, 1) + if rot == 2: + return (v.re, -1), (v.im, -1) + return (v.im, 1), (v.re, -1) + + def add(self, a, b, sgn=1, rot=0, prefer=None): + """a + sgn * i^rot * b, choosing the output scale.""" + (bre, bres), (bim, bims) = self._rot(b, rot) + sb = s_mul(b.s, float(sgn)) + use_b = False + if prefer == "b" or (prefer is None and not s_is_unit(a.s) and s_is_unit(sb)): + use_b = True + if not use_b: + r = s_div(sb, a.s) + re = self.axpy(a.re, 1, bre, r if bres > 0 else self._sneg(r)) + im = self.axpy(a.im, 1, bim, r if bims > 0 else self._sneg(r)) + return CVal(re, im, a.s) + # output scale sb: result = sb * (b' + (a.s/sb) * a) + r = s_div(a.s, sb) + re = self.axpy(bre, bres, a.re, r) + im = self.axpy(bim, bims, a.im, r) + return CVal(re, im, sb) + + def sub(self, a, b, rot=0, prefer=None): + return self.add(a, b, -1, rot, prefer) + + def rot(self, a, rot): + """Multiply by i^rot (free).""" + be = self.be + rot &= 3 + if rot == 0: + return a + if rot == 1: + return CVal(be.neg(a.im), a.re, a.s) + if rot == 2: + return CVal(a.re, a.im, s_mul(a.s, -1.0)) + return CVal(a.im, be.neg(a.re), a.s) + + def mul_const(self, a, wr, wi): + """a * (wr + i wi) for compile-time constants, Linzer-Feig style.""" + if abs(wi) < EPS_EXACT: + return CVal(a.re, a.im, s_mul(a.s, wr)) + if abs(wr) < EPS_EXACT: + # wi * i + v = self.rot(a, 1) + return CVal(v.re, v.im, s_mul(v.s, wi)) + if abs(wr) >= abs(wi): + t = wi / wr + if self.emu is not None and abs(abs(t) - 1.0) < 1e-15 and self.emu.take(2): + # t = +-1: (re - t im, im + t re) is an add/sub pair of (re, im) + if t > 0: + sm, df = self.be.addsub(a.re, a.im) + return CVal(df, sm, s_mul(a.s, wr)) + sm, df = self.be.addsub(a.im, a.re) + return CVal(sm, df, s_mul(a.s, wr)) + # u = a + i t a = (re - t im, im + t re); value = wr * u + re = self.axpy(a.re, 1, a.im, -t) + im = self.axpy(a.im, 1, a.re, t) + return CVal(re, im, s_mul(a.s, wr)) + t = wr / wi + # w = wi (t + i); w a = wi (t re - im, t im + re) + re = self.axpy(a.im, -1, a.re, t) + im = self.axpy(a.re, 1, a.im, t) + return CVal(re, im, s_mul(a.s, wi)) + + def lf_twiddle(self, a, m, T1, t, f, nvars): + """a * w for a runtime twiddle in Linzer-Feig form w = sigma_m (1 + i t) + (flag f = 0) or sigma_m (t + i) (f = 1); 2 FMAs plus integer selects. + T1 = -t (f = 0) or t (f = 1). The result carries scale a.s * sigma_m.""" + be = self.be + s1 = be.select(f, a.re, a.im) + s2 = be.select(f, a.im, a.re) + s2n = be.select(f, be.bneg(a.im), a.re) + re = be.fma(T1, s1, s2n) + im = be.fma(t, s2, s1) + e = [0] * nvars + e[m] = 1 + return CVal(re, im, s_mul(a.s, RScale(1.0, e))) + + def mul_rt(self, a, wr, wi): + """a * (wr + i wi) for runtime backend scalars (4 FMAs, scale kept).""" + be = self.be + mz = be.mzero() + t1 = be.fma(a.im, wi, mz) + re = be.fma(a.re, wr, be.neg(t1)) + t2 = be.fma(a.im, wr, mz) + im = be.fma(a.re, wi, t2) + return CVal(re, im, a.s) + + def finalize(self, a): + """Bring a value to scale 1.""" + if _is_one(a.s): + return a + be = self.be + if isinstance(a.s, RScale): + v = self.rt.get(a.s) + mz = be.mzero() + sg = a.s.c < 0 + re = be.fma(be.neg(a.re) if sg else a.re, v, mz) + im = be.fma(be.neg(a.im) if sg else a.im, v, mz) + return CVal(re, im, 1.0) + if _is_one(-a.s): + return CVal(be.neg(a.re), be.neg(a.im), 1.0) + mz = be.mzero() + s = a.s + k = be.const(abs(s)) + re = be.fma(a.re if s > 0 else be.neg(a.re), k, mz) + im = be.fma(a.im if s > 0 else be.neg(a.im), k, mz) + return CVal(re, im, 1.0) + + +# --------------------------------------------------------------------------- +# DFT codelets (forward, W = exp(-2 pi i / n)) +# --------------------------------------------------------------------------- + + +def _twiddle(n, e): + e %= n + ang = -2.0 * math.pi * e / n + # exact values at multiples of pi/4 multiples + if (8 * e) % n == 0: + q = (8 * e) // n + c = math.sqrt(0.5) + tab = [(1.0, 0.0), (c, -c), (0.0, -1.0), (-c, -c), (-1.0, 0.0), (-c, c), (0.0, 1.0), (c, c)] + return tab[q] + return math.cos(ang), math.sin(ang) + + +def dft2(o, x): + return list(o.bfly(x[0], x[1])) + + +def dft4(o, x): + a, b = o.bfly(x[0], x[2]) + c, d = o.bfly(x[1], x[3]) + y0, y2 = o.bfly(a, c) + y3, y1 = o.bfly(b, d, 1) + return [y0, y1, y2, y3] + + +def dft_odd_prime(o, x): + """Pair-based DFT for odd prime length (3, 5, 7, 11, 13...).""" + p = len(x) + h = (p - 1) // 2 + t = [] + d = [] + for k in range(1, h + 1): + tk, dk = o.bfly(x[k], x[p - k]) + t.append(tk) + d.append(dk) + # X0 = x0 + sum t_k + acc = x[0] + for k in range(h): + acc = o.add(acc, t[k], prefer="a") + out = [None] * p + out[0] = acc + for m in range(1, h + 1): + a = x[0] + for k in range(1, h + 1): + c = math.cos(2.0 * math.pi * k * m / p) + a = o.add(a, CVal(t[k - 1].re, t[k - 1].im, s_mul(t[k - 1].s, c)), prefer="a") + # B_m = sum_k sin(2 pi k m / p) * d_k (scale carried) + s1 = math.sin(2.0 * math.pi * m / p) + b = CVal(d[0].re, d[0].im, s_mul(d[0].s, s1)) + for k in range(2, h + 1): + s = math.sin(2.0 * math.pi * k * m / p) + b = o.add(b, CVal(d[k - 1].re, d[k - 1].im, s_mul(d[k - 1].s, s)), prefer="a") + # X_m = A - i B ; X_{p-m} = A + i B + if o._unit_ratio(a, b) != 0: + out[p - m], out[m] = o.bfly(a, b, 1) + else: + out[m] = o.add(a, b, 1, 3, prefer="a") + out[p - m] = o.add(a, b, 1, 1, prefer="a") + return out + + +def dft3(o, x): + # t = x1 + x2, d = x1 - x2 ; X0 = x0 + t ; m = x0 - t/2 ; X1 = m - i(sqrt3/2) d + t, d = o.bfly(x[1], x[2]) + X0 = o.add(x[0], t, prefer="a") + m = o.add(x[0], CVal(t.re, t.im, s_mul(t.s, -0.5)), prefer="a") + s3 = math.sqrt(3.0) / 2.0 + dd = CVal(d.re, d.im, s_mul(d.s, s3)) + if o._unit_ratio(m, dd) != 0: + X2, X1 = o.bfly(m, dd, 1) + else: + X1 = o.add(m, dd, 1, 3, prefer="a") + X2 = o.add(m, dd, 1, 1, prefer="a") + return [X0, X1, X2] + + +def dft_ct(o, x, n1, n2, sub1, sub2): + """Cooley-Tukey n = n1*n2 with constant twiddles. + + x[j1 + n1*j2]: inner DFTs of length n2 over j2 for each j1, twiddle + W_n^(j1*k2), outer DFTs of length n1 over j1. Output index k2 + n2*k1. + """ + n = n1 * n2 + y = [] + for j1 in range(n1): + y.append(sub2(o, [x[j1 + n1 * j2] for j2 in range(n2)])) + out = [None] * n + for k2 in range(n2): + col = [] + for j1 in range(n1): + v = y[j1][k2] + e = (j1 * k2) % n + if e: + wr, wi = _twiddle(n, e) + v = o.mul_const(v, wr, wi) + col.append(v) + z = sub1(o, col) + for k1 in range(n1): + out[k2 + n2 * k1] = z[k1] + return out + + +def dft_pfa(o, x, n1, n2, sub1, sub2): + """Good-Thomas n = n1*n2 with gcd(n1, n2) == 1 (no twiddles).""" + n = n1 * n2 + # input map j = (j1*n2 + j2*n1) mod n ; output map CRT + y = [] + for j1 in range(n1): + y.append(sub2(o, [x[(j1 * n2 + j2 * n1) % n] for j2 in range(n2)])) + e1 = n2 * pow(n2, -1, n1) + e2 = n1 * pow(n1, -1, n2) + out = [None] * n + for k2 in range(n2): + z = sub1(o, [y[j1][k2] for j1 in range(n1)]) + for k1 in range(n1): + out[(k1 * e1 + k2 * e2) % n] = z[k1] + return out + + +def _dft1(o, x): + return [x[0]] + + +def get_dft(n): + """Return codelet function f(o, xs) -> outputs for length n.""" + special = {1: _dft1, 2: dft2, 3: dft3, 4: dft4} + if n in special: + return special[n] + f = _factor(n) + if len(set(f)) == 1: + p = f[0] + if len(f) == 1: + return dft_odd_prime + if p == 2: + n1 = {8: 2, 16: 4, 32: 4, 64: 8}.get(n, 4) + else: + n1 = p + n2 = n // n1 + return lambda o, x: dft_ct(o, x, n1, n2, get_dft(n1), get_dft(n2)) + p0 = f[0] + n1 = 1 + while n % (n1 * p0) == 0: + n1 *= p0 + n2 = n // n1 + return lambda o, x: dft_pfa(o, x, n1, n2, get_dft(n1), get_dft(n2)) + + +def _factor(n): + f = [] + d = 2 + while d * d <= n: + while n % d == 0: + f.append(d) + n //= d + d += 1 + if n > 1: + f.append(n) + return f + + +def dft(o, xs): + out = get_dft(len(xs))(o, xs) + return out + + +# --------------------------------------------------------------------------- +# Counting backend: host-side dry runs of a codelet (cost model, constant sets) +# --------------------------------------------------------------------------- + + +class CountingBackend: + """Evaluates on Python floats and counts FMAs and the constants they use.""" + + def __init__(self): + self.nfma = 0 + self.consts = {} + + def fma(self, a, b, c): + self.nfma += 1 + return a * b + c + + def addsub(self, a, b): + return a + b, a - b + + def neg(self, a): + return -a + + def const(self, v): + key = float(v) + self.consts[key] = self.consts.get(key, 0) + 1 + return key + + def mzero(self): + return -0.0 + + def select(self, f, a, b): + return a if f else b + + def bneg(self, a): + return -a + + +# --------------------------------------------------------------------------- +# Half butterflies: two lanes share one radix-2/4 butterfly, split by output +# parity p (runtime, per lane). Used when a stage leaves half of the lanes of a +# round idle: the per-lane serial FMA chain is halved at no extra FMA cost +# (radix 4: 8 instead of 16 FMAs per lane; radix 2: 2 instead of 4). +# --------------------------------------------------------------------------- + + +def _signed_coef(ops, r, p): + """Register holding r (p false) or -r (p true); r compile-time float or RScale.""" + be = ops.be + if isinstance(r, RScale): + v = ops.rt.get(r) + if r.c < 0: + v = be.bneg(v) + else: + v = be.const(1.0 if _is_one(abs(r)) else abs(r)) + if r < 0: + v = be.bneg(v) + return be.select(p, be.bneg(v), v) + + +def add_ps(ops, a, b, p): + """a + s*b with s = +1 (p false) or -1 (p true); result keeps a's scale.""" + be = ops.be + c = _signed_coef(ops, s_div(b.s, a.s), p) + return CVal(be.fma(b.re, c, a.re), be.fma(b.im, c, a.im), a.s) + + +def rot_ps(ops, v, p): + """v (p false) or -i*v (p true): integer selects only.""" + be = ops.be + return CVal(be.select(p, v.im, v.re), be.select(p, be.bneg(v.re), v.im), v.s) + + +def dft_half(ops, xs, p): + """Outputs p, p+2 (radix 4) or p (radix 2) of the forward DFT of xs.""" + if len(xs) == 2: + return [add_ps(ops, xs[0], xs[1], p)] + assert len(xs) == 4 + u = add_ps(ops, xs[0], xs[2], p) + v = rot_ps(ops, add_ps(ops, xs[1], xs[3], p), p) + return [ops.add(u, v), ops.sub(u, v)] diff --git a/build/torch-cuda/_r2c_f64/fft_direct.py b/build/torch-cuda/_r2c_f64/fft_direct.py new file mode 100644 index 0000000000000000000000000000000000000000..18757f274b08572e73cd1c577780706891edc6d3 --- /dev/null +++ b/build/torch-cuda/_r2c_f64/fft_direct.py @@ -0,0 +1,200 @@ +"""Direct-DFT fallback kernels for axes the codelet passes do not cover. + +Used for passes the planner gives no codelet config: lengths with a prime factor above +13 (for X rows: when neither the half-length nor the packed route is available). O(n) +work per output with a runtime loop over the input and an index into the W_n table +(k * j mod n); the accumulation is FMA-based in float64. +""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.utils as cutils +from cutlass import Float64, cute + +from .dsl_ops import vload2, vstore2 +from .fft_plan import exact_twiddles +from .fft_stages import const_table, interleave, load_consts + +DIRECT_P = 128 +# The column kernel stages whole columns in shared memory (CC * n complex float64 = 16 bytes +# per element, CC = 1 above n = 2048), and one CTA can opt into at most 99 KB of dynamic shared +# memory on sm_120: the longest Y/Z axis it can transform is 101376 / 16 = 6336 (measured: +# 6336 runs, 6337 fails to launch with CUDA_ERROR_INVALID_VALUE). The row kernel (X axis) +# needs no shared memory. +MAX_SMEM_PER_CTA = 99 * 1024 +MAX_DIRECT_COL_LEN = MAX_SMEM_PER_CTA // 16 + + +def _dft_table(n): + """W_n^i for i in [0, n), interleaved (re, im) float64.""" + return interleave(exact_twiddles(n)) + + +class DirectColCfg: + def __init__(self, n, ncols, M, plane_stride, es): + self.n, self.ncols, self.M, self.plane_stride, self.es = n, ncols, M, plane_stride, es + # columns per CTA so that the smem tile stays modest + self.CC = max(1, min(8, (32 * 1024) // (n * 16))) + self.ntiles = (ncols + self.CC - 1) // self.CC + self.smem_bytes = self.CC * n * 16 + + def __repr__(self): + return f"dcol(n={self.n}, CC={self.CC}, tiles={self.ntiles})" + + +def build_direct_cols(cfg): + n, ncols, M, plane_stride, es, CC = cfg.n, cfg.ncols, cfg.M, cfg.plane_stride, cfg.es, cfg.CC + P = DIRECT_P + tab = _dft_table(n) + keys = [1.0] + ntiles, smem_bytes = cfg.ntiles, cfg.smem_bytes + nload = (CC * n + P - 1) // P + + @cute.kernel + def kern(gio: cute.Tensor, gtw: cute.Tensor, gk: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + _body(gio, gtw, gk, tid, bid) + + def _body(gio, gtw, gk, tid, bid): + lane = tid % 32 + _, mz = load_consts(gk, keys, lane) + alloc = cutils.SmemAllocator() + sm = alloc.allocate_array(Float64, CC * n * 2, byte_alignment=16) + pio = gio.iterator + pw = gtw.iterator + c0 = bid * CC + + def col_base(col): + if M >= ncols: + return col + return (col // M) * plane_stride + (col % M) + + # stage the tile's columns in shared memory (element-major, column fastest) + for i in range(nload): + e = tid + i * P + + def ld(e=e): + c = e % CC + j = e // CC + col = cutlass.min(c0 + c, ncols - 1) + re, im = vload2(pio + (col_base(col) + j * es) * 2) + vstore2(sm + e * 2, re, im) + + if (i + 1) * P <= CC * n: + ld() + else: + cutlass.if_generate(e < CC * n, ld) + cute.arch.sync_threads() + for i in range(nload): + e = tid + i * P + + def out(e=e): + c = e % CC + k = e // CC + zero = mz + + def body(j, args): + ar, ai, idx = args + xr, xi = vload2(sm + (j * CC + c) * 2) + wr, wi = vload2(pw + idx * 2) + ar = cute.math.fma(xr, wr, ar) + ar = cute.math.fma(-xi, wi, ar) + ai = cute.math.fma(xr, wi, ai) + ai = cute.math.fma(xi, wr, ai) + nidx = idx + k + nidx = cutlass.Int32(cutlass.select_(nidx >= n, nidx - n, nidx)) + return [ar, ai, nidx] + + acc = [zero, zero, cutlass.Int32(0)] + for j, args, res in cutlass.for_generate(0, n, 1, iter_args=acc): + cutlass.yield_out(body(j, list(args))) + ar, ai, _ = res + col = c0 + c + + def st(): + vstore2(pio + (col_base(col) + k * es) * 2, ar, ai) + + cutlass.if_generate(col < ncols, st) + + if (i + 1) * P <= CC * n: + out() + else: + cutlass.if_generate(e < CC * n, out) + + @cute.jit + def launch(gio: cute.Tensor, gtw: cute.Tensor, gk: cute.Tensor, stream: cuda.CUstream): + kern(gio, gtw, gk).launch(grid=[ntiles, 1, 1], block=[P, 1, 1], smem=smem_bytes, stream=stream) + + return launch, tab, const_table(keys) + + +class DirectRowCfg: + def __init__(self, nx, nrows): + self.nx, self.nrows = nx, nrows + self.nf = nx // 2 + 1 + self.RR = max(1, min(8, 1024 // self.nf)) + self.ntiles = (nrows + self.RR - 1) // self.RR + + def __repr__(self): + return f"drow(nx={self.nx}, RR={self.RR}, tiles={self.ntiles})" + + +def build_direct_rows(cfg): + nx, nrows, nf, RR = cfg.nx, cfg.nrows, cfg.nf, cfg.RR + P = DIRECT_P + tab = _dft_table(nx) + keys = [1.0] + ntiles = cfg.ntiles + nout = (RR * nf + P - 1) // P + + @cute.kernel + def kern(gin: cute.Tensor, gout: cute.Tensor, gtw: cute.Tensor, gk: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + _body(gin, gout, gtw, gk, tid, bid) + + def _body(gin, gout, gtw, gk, tid, bid): + lane = tid % 32 + _, mz = load_consts(gk, keys, lane) + pin, pout, pw = gin.iterator, gout.iterator, gtw.iterator + r0 = bid * RR + for i in range(nout): + e = tid + i * P + + def out(e=e): + rr = e // nf + k = e % nf + row = cutlass.min(r0 + rr, nrows - 1) + xbase = pin + row * nx + + def body(j, args): + ar, ai, idx = args + x = (xbase + j).load() + wr, wi = vload2(pw + idx * 2) + ar = cute.math.fma(x, wr, ar) + ai = cute.math.fma(x, wi, ai) + nidx = idx + k + nidx = cutlass.Int32(cutlass.select_(nidx >= nx, nidx - nx, nidx)) + return [ar, ai, nidx] + + acc = [mz, mz, cutlass.Int32(0)] + for j, args, res in cutlass.for_generate(0, nx, 1, iter_args=acc): + cutlass.yield_out(body(j, list(args))) + ar, ai, _ = res + + def st(): + vstore2(pout + ((r0 + rr) * nf + k) * 2, ar, ai) + + cutlass.if_generate(r0 + rr < nrows, st) + + if (i + 1) * P <= RR * nf: + out() + else: + cutlass.if_generate(e < RR * nf, out) + + @cute.jit + def launch(gin: cute.Tensor, gout: cute.Tensor, gtw: cute.Tensor, gk: cute.Tensor, stream: cuda.CUstream): + kern(gin, gout, gtw, gk).launch(grid=[ntiles, 1, 1], block=[P, 1, 1], smem=16, stream=stream) + + return launch, tab, const_table(keys) diff --git a/build/torch-cuda/_r2c_f64/fft_passes.py b/build/torch-cuda/_r2c_f64/fft_passes.py new file mode 100644 index 0000000000000000000000000000000000000000..c4c24c7590d3198b9560a35dbdf86a205b287d3f --- /dev/null +++ b/build/torch-cuda/_r2c_f64/fft_passes.py @@ -0,0 +1,802 @@ +"""The three passes as CuTe-DSL kernels, and their per-pass configuration. + +``PassCfg`` fixes a pass's plan, tile (C transforms per warp, W warps per CTA, or a tile +shared by ``split`` warps), shared-memory layout and options (see tuned.py). Builders: + + build_r2c X pass, nx even: half-length complex FFT of each row (pairs of reals + as one complex), then the split post-processing with H_k = W_nx^k / 2 + build_r2c_pack2 X pass: two real rows packed into one complex FFT of length nx; the + split costs 4 FMAs per k and halves exactly (exponent decrement), or + runs as halved integer-pipe add/sub pairs (``emu_split``) + build_colpass Y / Z pass: complex FFT along a strided axis, in place on the output + +Each builder returns (jit launcher, twiddle table, constant table, twiddle-flag table); +the tables are CPU tensors the caller copies to the device. Passes chain with programmatic +dependent launch; in mode 3 Y tiles instead wait on per-plane counters that X tiles +increment, and the last Y tile re-arms them for the next launch. +""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.utils as cutils +from cutlass import Float64, cute + +from .dsl_ops import ( + DSLBackend, + _and, + _when, + bits_neg, + emu_addsub, + exact_half, + flag_acquire, + flag_spin, + load_f64_ef, + vload2, + vload2_ef, + vstore2, + vstore2_ef, +) +from .fft_codelets import CVal, Ops +from .fft_plan import exact_twiddle +from .fft_stages import ( + Ctx, + Layout, + StageTables, + collect_consts, + const_table, + emit_fft, + emit_fft_xs, + load_consts, + preload_twiddles, +) + +NUM_SMS = 170 + + +def _nbuf(S, hold, kind): + if kind == "col": + if S == 1: + return 0 + if S == 2 or hold: + return 1 + return 2 + # r2c / pack2: last stage writes natural order into smem + if S == 1 or hold: + return 1 + return 2 + + +class PassCfg: + def __init__( + self, + kind, + n, + plan, + C, + W, + hold, + ntiles, + mode, + preload=False, + split=1, + stagger=False, + smgrid=False, + evict_first=False, + reverse=False, + half=False, + emu=0.0, + lf=True, + emu_split=None, + xs=False, + ): + self.xs = bool(xs) and kind in ("col", "pack2") # split real/imaginary smem exchange (half the smem per tile) + # target fraction of each stage's FP ops done as integer-pipe add/sub pairs (one value, or + # one per stage) + self.emu = [float(e) for e in emu] if isinstance(emu, (list, tuple)) else float(emu) + # pack2 X passes: run the two-row split as integer add/sub pairs (default: when emu is on) + self.emu_split = ( + (max(self.emu) if isinstance(self.emu, list) else self.emu) > 0 if emu_split is None else bool(emu_split) + ) + self.lf = bool(lf) # runtime twiddles in Linzer-Feig form (else plain 4-FMA products) + self.preload = preload # issue all twiddle-table loads before stage 1 + self.half = bool(half) # half butterflies in stages that leave lanes idle (see emit_fft) + # X passes: read the input with an L2 evict-first policy, so that the X output (re-read by + # the next pass) survives in L2 when the problem does not fit + self.evict_first = bool(evict_first) and kind != "col" + # column passes: write the tile results with an L2 evict-first policy (last pass only) + self.store_ef = bool(evict_first) and kind == "col" + # column passes: process tiles in reverse order (start on the planes the previous pass wrote last) + self.reverse = bool(reverse) and kind == "col" + # smgrid: one CTA of W warps per SM, tile = warp * NUM_SMS + cta (even tile count per SM) + self.smgrid = bool(smgrid) and split == 1 and W * NUM_SMS >= ntiles + # X passes: warps of a CTA issue their input loads rank by rank (rank = warp // 4, one warp + # per SM sub-partition in each rank), so every sub-partition gets one tile's data early + self.stagger = bool(stagger) and split == 1 and W >= 8 and (self.smgrid or ntiles % W == 0) + self.split = split # >1: one tile per CTA shared by `split` warps (W ignored) + if split > 1: + W = split + self.kind = kind + self.n = n + self.plan = plan + self.C = C + self.W = W + self.P = 32 * W + self.hold = hold + self.ntiles = ntiles + if split > 1: + self.ngrid = ntiles + elif self.smgrid: + self.ngrid = NUM_SMS + else: + self.ngrid = (ntiles + W - 1) // W + self.lay = Layout(mode, n, C) + self.nbuf = _nbuf(len(plan.stages), hold, kind) + self.warp_elems = self.nbuf * self.lay.size * 2 # doubles per tile + if self.xs and self.nbuf > 0: + self.nbuf = 1 + self.warp_elems = self.lay.size + self.smem_bytes = max(16, self.warp_elems * 8 * (1 if split > 1 else W)) + + def __repr__(self): + return ( + f"{self.kind}(n={self.n}, C={self.C}, W={self.W}, split={self.split}, hold={self.hold}, " + f"pre={int(self.preload)}, stg={int(self.stagger)}, smg={int(self.smgrid)}, ef={int(self.evict_first)}, " + f"rev={int(self.reverse)}, half={int(self.half)}, xs={int(self.xs)}, tiles={self.ntiles}, " + f"smem={self.smem_bytes}, mode={self.lay.mode}, emu={self.emu}, lf={int(self.lf)}, " + f"esplit={int(self.emu_split)}, plan={self.plan.radices})" + ) + + +def _warp_setup(cfg, tid, bid): + warp = tid // 32 + lane = tid % 32 + if cfg.split > 1: + return warp, lane, bid + if cfg.smgrid: + return warp, lane, warp * NUM_SMS + bid + tile = bid * cfg.W + warp + return warp, lane, tile + + +def _setup_split(ctx, cfg, warp): + if cfg.split > 1: + ctx.split = cfg.split + ctx.wsub = warp + ctx.sync = cute.arch.sync_threads + + +def _warp_bufs(cfg, warp, alloc=None): + if cfg.nbuf == 0: + return [] + if alloc is None: + alloc = cutils.SmemAllocator() + if cfg.split > 1: + base = alloc.allocate_array(Float64, cfg.warp_elems, byte_alignment=16) + wb = base + else: + base = alloc.allocate_array(Float64, cfg.warp_elems * cfg.W, byte_alignment=16) + wb = base + warp * cfg.warp_elems + bufs = [wb] + if cfg.nbuf == 2: + bufs.append(wb + cfg.lay.size * 2) + return bufs + + +FLAG_PAD = 32 # ints per counter (one 128-byte line each) + + +def x_signal(ctx, flags, gsync, lane, r0, rows): + """After an X tile: publish its rows to the per-plane counters (rows within one plane). + + The warp (or CTA) barrier orders every lane's output stores before the leader's release + reduction; unlike a full fence.acq_rel this neither invalidates the SM's L1 nor makes all + lanes wait for the write-back.""" + if flags is None: + return + getattr(ctx, "sync", cute.arch.sync_warp)() + lead = lane == 0 + if getattr(ctx, "split", 1) > 1: + lead = lead & (ctx.wsub == 0) + z = r0 // flags["ny"] + + def inc(): + cute.arch.atomic_add(gsync.iterator + z * FLAG_PAD, cutlass.Int32(rows), sem="release", scope="gpu") + + cutlass.if_generate(lead, inc) + + +def _setup_stagger(ctx, cfg, warp): + """Rank-ordered first loads (see PassCfg.stagger): rank r waits on named barrier r until + rank r-1 has issued its loads, then releases rank r+1 after issuing its own.""" + if not cfg.stagger: + return + W = cfg.W + nr = (W + 3) // 4 + rank = warp // 4 + + def nthr(r): # threads meeting at barrier r (ranks r-1 and r) + return 32 * (min(4, W - 4 * (r - 1)) + min(4, W - 4 * r)) + + def before(): + for r in range(1, nr): + cutlass.if_generate(rank == r, lambda r=r: cute.arch.barrier(barrier_id=r, number_of_threads=nthr(r))) + + def after(): + for r in range(nr - 1): + cutlass.if_generate( + rank == r, lambda r=r: cute.arch.barrier_arrive(barrier_id=r + 1, number_of_threads=nthr(r + 1)) + ) + + ctx.before_first_loads = before + ctx.after_first_loads = after + + +def _guard_tile(cfg, tile, fn, ctx=None): + if cfg.split > 1 or (cfg.ntiles % cfg.W == 0 and not cfg.smgrid): + fn() + elif cfg.stagger and ctx is not None: + # idle warps still take part in the rank barriers + def idle(): + ctx.before_first_loads() + ctx.after_first_loads() + + cutlass.if_generate(tile < cfg.ntiles, fn, idle) + else: + cutlass.if_generate(tile < cfg.ntiles, fn) + + +# --------------------------------------------------------------------------- +# Column pass (complex FFT along a strided axis, in place) +# --------------------------------------------------------------------------- + + +def build_colpass(cfg, ncols, M, plane_stride, es, pdl_wait=False, pdl_trigger=None, flags=None): + plan = cfg.plan + C, P = cfg.C, cfg.P + keys = collect_consts([st.r for st in plan.stages]) + tabs = StageTables(plan, lf=cfg.lf) + lay = cfg.lay + + @cute.kernel + def kern(gio: cute.Tensor, gtw: cute.Tensor, gk: cute.Tensor, gfl: cute.Tensor, gsync: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + _body(gio, gtw, gk, gfl, gsync, tid, bid) + + def _body(gio, gtw, gk, gfl, gsync, tid, bid): + warp, lane, tile = _warp_setup(cfg, tid, bid) + bufs = _warp_bufs(cfg, warp) + kmap, mz = load_consts(gk, keys, lane) + ctx = Ctx() + ctx.ops = Ops(DSLBackend(kmap, mz)) + ctx.gtw = gtw + ctx.gfl = gfl + ctx.tabs = tabs + ctx.half = cfg.half + ctx.emu = cfg.emu + ctx.loads_upfront = cfg.stagger and flags is None + if flags is None: + _setup_stagger(ctx, cfg, warp) + _setup_split(ctx, cfg, warp) + pio = gio.iterator + if pdl_trigger == "start": + cute.arch.griddepcontrol_launch_dependents() + if pdl_wait: + cute.arch.griddepcontrol_wait() + + def col_base(col): + if M >= ncols: + return col + return (col // M) * plane_stride + (col % M) + + def make_io(tile): + if cfg.reverse and flags is None: + tile = (cfg.ntiles - 1) - tile + c0 = tile * C + + def load_first(c, g): + col = c0 + c + if ncols % C != 0: + col = cutlass.min(col, ncols - 1) + return vload2(pio + (col_base(col) + g * es) * 2) + + def store_last(c, kidx, re, im): + col = c0 + c + ptr = pio + (col_base(col) + kidx * es) * 2 + st = vstore2_ef if cfg.store_ef else vstore2 + if ncols % C == 0: + st(ptr, re, im) + else: + cutlass.if_generate(col < ncols, lambda: st(ptr, re, im)) + + return load_first, store_last + + lf, sl = make_io(tile) + ps = gsync.iterator + leader_warp = (ctx.wsub == 0) if cfg.split > 1 else None + + def wait_inputs(): + if flags is None: + return + nf_ = flags["nf"] + c0 = tile * C + if flags["role"] == "y": + # planes touched by columns [c0, c0+C) (z-major columns z*nf + kx); + # each plane counter sits on its own 128-byte line to avoid polling hot spots + za = c0 // nf_ + zb = cutlass.min(c0 + (C - 1), ncols - 1) // nf_ + cutlass.if_generate(lane == 0, lambda: flag_spin(ps + za * FLAG_PAD, flags["ny"])) + cutlass.if_generate((lane == 1) & (zb != za), lambda: flag_spin(ps + zb * FLAG_PAD, flags["ny"])) + cute.arch.sync_warp() + flag_acquire(ps + za * FLAG_PAD) + flag_acquire(ps + zb * FLAG_PAD) + + def signal_outputs(): + if flags is None: + return + # Y pass: count finished tiles; the last one re-arms the plane counters for the next + # launch (every X->Y counter wait has completed by then). + nz_ = flags["nz"] + ex = ps + nz_ * FLAG_PAD + lead = lane == 0 if leader_warp is None else (lane == 0) & leader_warp + + def lead_fn(): + old_ = cute.arch.atomic_add(ex, cutlass.Int32(1), sem="relaxed", scope="gpu") + + def reset(): + for i in cutlass.for_generate(0, nz_ + 1, 1): + (ps + i * FLAG_PAD).store(cutlass.Int32(0)) + cutlass.yield_out() + + cutlass.if_generate(old_ == flags["ntiles"] - 1, reset) + + cutlass.if_generate(lead, lead_fn) + + def work(): + ptw = preload_twiddles(ctx, plan, lay, C, lane) if cfg.preload else None + wait_inputs() + if cfg.xs and len(plan.stages) > 1: + emit_fft_xs(ctx, plan, lay, C, lane, bufs[0], lf, sl, False, pre_tw=ptw) + else: + emit_fft(ctx, plan, lay, C, lane, bufs, lf, sl, False, cfg.hold, pre_tw=ptw) + signal_outputs() + + _guard_tile(cfg, tile, work, ctx if flags is None else None) + if pdl_trigger == "end": + cute.arch.griddepcontrol_launch_dependents() + + ngrid = cfg.ngrid + smem_bytes = cfg.smem_bytes + + use_pdl = pdl_wait or (flags is not None) + + @cute.jit + def launch( + gio: cute.Tensor, gtw: cute.Tensor, gk: cute.Tensor, gfl: cute.Tensor, gsync: cute.Tensor, stream: cuda.CUstream + ): + kern(gio, gtw, gk, gfl, gsync).launch( + grid=[ngrid, 1, 1], block=[P, 1, 1], smem=smem_bytes, stream=stream, use_pdl=use_pdl + ) + + return launch, tabs.dtab, const_table(keys), tabs.itab + + +# --------------------------------------------------------------------------- +# X pass: real-to-complex FFT of each row +# --------------------------------------------------------------------------- + + +def r2c_slots(h): + return h // 2 if h % 2 == 0 else (h + 1) // 2 + + +def r2c_post_twiddles(nx): + """H_k = W_nx^k / 2 for k in [0, h//2], as the complex product W * (0.5 + 0i).""" + h = nx // 2 + out = [] + for k in range(h // 2 + 1): + w = exact_twiddle(nx, k) + out.append(complex(w.real * 0.5 - w.imag * 0.0, w.real * 0.0 + w.imag * 0.5)) + return out + + +def build_r2c(cfg, nx, nrows, pdl_trigger=None, prefetch_bytes=(), flags=None): + plan = cfg.plan + C, P = cfg.C, cfg.P + h, nf = nx // 2, nx // 2 + 1 + keys = collect_consts([st.r for st in plan.stages], extra=(0.5,)) + tabs = StageTables(plan, lf=cfg.lf, extra_complex=r2c_post_twiddles(nx)) + post_off = tabs.extra_off // 2 # in complex elements + lay = cfg.lay + hh = h // 2 + even_h = h % 2 == 0 + nslot = r2c_slots(h) + + npf = len(prefetch_bytes) + if npf: + # also warm this pass's own twiddle/flag tables (used from stage 2 on); + # the kernel always takes 6 prefetch tensors (zero-byte entries are padding) + assert npf <= 6 + prefetch_bytes = (int(tabs.dtab.nbytes), int(tabs.itab.nbytes)) + tuple(prefetch_bytes) + (0,) * (6 - npf) + npf = len(prefetch_bytes) + + @cute.kernel + def kern( + gin: cute.Tensor, + gout: cute.Tensor, + gtw: cute.Tensor, + gk: cute.Tensor, + gfl: cute.Tensor, + gp0: cute.Tensor, + gp1: cute.Tensor, + gp2: cute.Tensor, + gp3: cute.Tensor, + gp4: cute.Tensor, + gp5: cute.Tensor, + gsync: cute.Tensor, + ): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + _body(gin, gout, gtw, gk, gfl, [gtw, gfl, gp0, gp1, gp2, gp3, gp4, gp5], gsync, tid, bid) + + def _body(gin, gout, gtw, gk, gfl, gps, gsync, tid, bid): + warp, lane, tile = _warp_setup(cfg, tid, bid) + if npf: + # warm the later passes' tables in L2 while this pass runs + def pf(): + for gi in range(npf): + nlines = (prefetch_bytes[gi] + 127) // 128 + if nlines == 0: + continue + for li in range(0, nlines, 32): + + def one(gi=gi, li=li): + addr = (gps[gi].iterator + (lane + li) * (128 // gps[gi].element_type.width * 8)).toint() + cute.arch.inline_ptx("prefetch.global.L2 [{$r0}];", read_only_args=[addr]) + + if li + 32 <= nlines: + one() + else: + cutlass.if_generate(lane + li < nlines, one) + + cutlass.if_generate((bid == 0) & (tid < 32), pf) + bufs = _warp_bufs(cfg, warp) + kmap, mz = load_consts(gk, keys, lane) + ctx = Ctx() + ops = Ops(DSLBackend(kmap, mz)) + ctx.ops = ops + ctx.gtw = gtw + ctx.gfl = gfl + ctx.tabs = tabs + ctx.half = cfg.half + ctx.emu = cfg.emu + ctx.loads_upfront = True # issue all first-stage global loads before any arithmetic + _setup_stagger(ctx, cfg, warp) + _setup_split(ctx, cfg, warp) + pin = gin.iterator + pout = gout.iterator + if pdl_trigger == "start": + cute.arch.griddepcontrol_launch_dependents() + half = kmap[0.5] + one = kmap[1.0] + + def make_load(tile): + r0 = tile * C + + def load_first(c, g): + row = r0 + c + if nrows % C != 0: + row = cutlass.min(row, nrows - 1) + if cfg.evict_first: + return vload2_ef(pin + row * nx + g * 2) + return vload2(pin + row * nx + g * 2) + + return load_first + + B = C * nslot + split = cfg.split + T = ((B + 31) // 32 + split - 1) // split + + def post_idx(t): + pred = None + if split == 1: + q = lane + t * 32 + if B % 32 != 0 and (t + 1) * 32 > B: + pred = q < B + q = cutlass.min(q, B - 1) + else: + q = lane + (warp + t * split) * 32 + if B % (32 * split) != 0: + pred = q < B + q = cutlass.min(q, B - 1) + if lay.mode == "row": + return pred, q // nslot, q % nslot + return pred, q % C, q // C + + def work(tile): + r0 = tile * C + ptw = preload_twiddles(ctx, plan, lay, C, lane) if cfg.preload else None + hpre = [vload2(gtw.iterator + (post_off + post_idx(t)[2]) * 2) for t in range(T)] if cfg.preload else None + zb = emit_fft(ctx, plan, lay, C, lane, bufs, make_load(tile), None, True, cfg.hold, pre_tw=ptw) + for t in range(T): + pred, c, k = post_idx(t) + km = cutlass.Int32(cutlass.select_(k == 0, cutlass.Int32(0), h - k)) + ar, ai = vload2(zb + lay.off(k, c) * 2) + br, bi = vload2(zb + lay.off(km, c) * 2) + hr, hi_ = hpre[t] if hpre is not None else vload2(gtw.iterator + (post_off + k) * 2) + sr = cute.math.fma(br, one, ar) + si = cute.math.fma(-bi, one, ai) + dr = cute.math.fma(-br, one, ar) + di = cute.math.fma(bi, one, ai) + v = ops.mul_rt(CVal(di, -dr), hr, hi_) + x0r = cute.math.fma(sr, half, v.re) + x0i = cute.math.fma(si, half, v.im) + x1r = cute.math.fma(sr, half, -v.re) + x1i = cute.math.fma(-si, half, v.im) + row = r0 + c + obase = pout + row * (nf * 2) + kout2 = h - k + + def _st(obase=obase, k=k, kout2=kout2, x0r=x0r, x0i=x0i, x1r=x1r, x1i=x1i, c=c): + vstore2(obase + k * 2, x0r, x0i) + vstore2(obase + kout2 * 2, x1r, x1i) + if even_h: + + def _mid(): + mr, mi = vload2(zb + lay.off(hh, c) * 2) + vstore2(obase + hh * 2, mr, bits_neg(mi)) + + cutlass.if_generate(k == 0, _mid) + + cond = pred + if nrows % C != 0: + cond = _and(cond, row < nrows) + _when(cond, _st) + x_signal(ctx, flags, gsync, lane, r0, C) + + _guard_tile(cfg, tile, lambda: work(tile), ctx) + if pdl_trigger == "end": + cute.arch.griddepcontrol_launch_dependents() + + ngrid = cfg.ngrid + smem_bytes = cfg.smem_bytes + + @cute.jit + def launch( + gin: cute.Tensor, + gout: cute.Tensor, + gtw: cute.Tensor, + gk: cute.Tensor, + gfl: cute.Tensor, + gp0: cute.Tensor, + gp1: cute.Tensor, + gp2: cute.Tensor, + gp3: cute.Tensor, + gp4: cute.Tensor, + gp5: cute.Tensor, + gsync: cute.Tensor, + stream: cuda.CUstream, + ): + kern(gin, gout, gtw, gk, gfl, gp0, gp1, gp2, gp3, gp4, gp5, gsync).launch( + grid=[ngrid, 1, 1], block=[P, 1, 1], smem=smem_bytes, stream=stream + ) + + return launch, tabs.dtab, const_table(keys), tabs.itab + + +# --------------------------------------------------------------------------- +# X pass, packed variant: two real rows per complex FFT of length nx +# --------------------------------------------------------------------------- +# +# z[n] = x_a[n] + i x_b[n]; Z = FFT_nx(z); for k in [0, nx/2]: +# X_a[k] = (Z[k] + conj Z[nx-k]) / 2, X_b[k] = -i (Z[k] - conj Z[nx-k]) / 2. +# The split costs 4 FMAs per k (2 per output); the factor 1/2 is applied with an +# exact exponent decrement. Works for odd nx as well. + + +def pack2_slots(nx): + return nx // 2 + 1 + + +def build_r2c_pack2(cfg, nx, nrows, pdl_trigger=None, prefetch_bytes=(), flags=None): + plan = cfg.plan + C, P = cfg.C, cfg.P # C = row pairs per tile + nf = nx // 2 + 1 + keys = collect_consts([st.r for st in plan.stages]) + tabs = StageTables(plan, lf=cfg.lf) + lay = cfg.lay + nslot = nf + npf = len(prefetch_bytes) + if npf: + assert npf <= 6 + prefetch_bytes = (int(tabs.dtab.nbytes), int(tabs.itab.nbytes)) + tuple(prefetch_bytes) + (0,) * (6 - npf) + npf = len(prefetch_bytes) + rows_per_tile = 2 * C + + @cute.kernel + def kern( + gin: cute.Tensor, + gout: cute.Tensor, + gtw: cute.Tensor, + gk: cute.Tensor, + gfl: cute.Tensor, + gp0: cute.Tensor, + gp1: cute.Tensor, + gp2: cute.Tensor, + gp3: cute.Tensor, + gp4: cute.Tensor, + gp5: cute.Tensor, + gsync: cute.Tensor, + ): + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + _body(gin, gout, gtw, gk, gfl, [gtw, gfl, gp0, gp1, gp2, gp3, gp4, gp5], gsync, tid, bid) + + def _body(gin, gout, gtw, gk, gfl, gps, gsync, tid, bid): + warp, lane, tile = _warp_setup(cfg, tid, bid) + if npf: + + def pf(): + for gi in range(npf): + nlines = (prefetch_bytes[gi] + 127) // 128 + if nlines == 0: + continue + for li in range(0, nlines, 32): + + def one(gi=gi, li=li): + addr = (gps[gi].iterator + (lane + li) * (128 // gps[gi].element_type.width * 8)).toint() + cute.arch.inline_ptx("prefetch.global.L2 [{$r0}];", read_only_args=[addr]) + + if li + 32 <= nlines: + one() + else: + cutlass.if_generate(lane + li < nlines, one) + + cutlass.if_generate((bid == 0) & (tid < 32), pf) + bufs = _warp_bufs(cfg, warp) + kmap, mz = load_consts(gk, keys, lane) + ctx = Ctx() + ops = Ops(DSLBackend(kmap, mz)) + ctx.ops = ops + ctx.gtw = gtw + ctx.gfl = gfl + ctx.tabs = tabs + ctx.half = cfg.half + ctx.emu = cfg.emu + ctx.loads_upfront = True # issue all first-stage global loads before any arithmetic + _setup_stagger(ctx, cfg, warp) + _setup_split(ctx, cfg, warp) + pin = gin.iterator + pout = gout.iterator + if pdl_trigger == "start": + cute.arch.griddepcontrol_launch_dependents() + one = kmap[1.0] + B = C * nslot + split = cfg.split + T = ((B + 31) // 32 + split - 1) // split + + def slot_idx(t): + pred = None + if split == 1: + q = lane + t * 32 + if B % 32 != 0 and (t + 1) * 32 > B: + pred = q < B + q = cutlass.min(q, B - 1) + else: + q = lane + (warp + t * split) * 32 + if B % (32 * split) != 0: + pred = q < B + q = cutlass.min(q, B - 1) + if lay.mode == "row": + return pred, q // nslot, q % nslot + return pred, q % C, q // C + + def work(tile): + r0 = tile * rows_per_tile + + def load_first(c, g): + ra = r0 + 2 * c + rb = ra + 1 + if nrows % rows_per_tile != 0: + ra = cutlass.min(ra, nrows - 1) + rb = cutlass.min(rb, nrows - 1) + if cfg.evict_first: + return load_f64_ef(pin + ra * nx + g), load_f64_ef(pin + rb * nx + g) + return (pin + ra * nx + g).load(), (pin + rb * nx + g).load() + + ptw = preload_twiddles(ctx, plan, lay, C, lane) if cfg.preload else None + + def body(): + if cfg.xs: + # last-stage outputs reach the split through the half buffer, real parts first + outs = emit_fft_xs(ctx, plan, lay, C, lane, bufs[0], load_first, None, True, pre_tw=ptw) + hb = bufs[0] + sync_ = getattr(ctx, "sync", cute.arch.sync_warp) + slots = [slot_idx(t) for t in range(T)] + kms = [cutlass.Int32(cutlass.select_(k == 0, cutlass.Int32(0), nx - k)) for _, _, k in slots] + if len(plan.stages) > 1: + sync_() + zparts = [] + for part in (0, 1): + for pred_, c_, items in outs: + + def st_s(items=items, c_=c_, part=part): + for kidx, y in items: + (hb + lay.off(kidx, c_)).store(y.re if part == 0 else y.im) + + _when(pred_, st_s) + sync_() + zparts.append( + [ + ((hb + lay.off(k, c)).load(), (hb + lay.off(km, c)).load()) + for (_, c, k), km in zip(slots, kms) + ] + ) + if part == 0: + sync_() + else: + zb = emit_fft(ctx, plan, lay, C, lane, bufs, load_first, None, True, cfg.hold, pre_tw=ptw) + for t in range(T): + pred, c, k = slot_idx(t) + km = cutlass.Int32(cutlass.select_(k == 0, cutlass.Int32(0), nx - k)) + if cfg.xs: + zr, wr = zparts[0][t] + zi, wi = zparts[1][t] + else: + zr, zi = vload2(zb + lay.off(k, c) * 2) + wr, wi = vload2(zb + lay.off(km, c) * 2) + if cfg.emu_split: + ar, bi = emu_addsub(wr, zr, half=True) + br, ai = emu_addsub(zi, wi, half=True) + else: + ar = exact_half(cute.math.fma(wr, one, zr)) + ai = exact_half(cute.math.fma(-wi, one, zi)) + br = exact_half(cute.math.fma(wi, one, zi)) + bi = exact_half(cute.math.fma(-zr, one, wr)) + ra = r0 + 2 * c + + def st_a(ra=ra, k=k, ar=ar, ai=ai): + vstore2(pout + (ra * nf + k) * 2, ar, ai) + + def st_b(ra=ra, k=k, br=br, bi=bi): + vstore2(pout + ((ra + 1) * nf + k) * 2, br, bi) + + ca = pred + cb = pred + if nrows % rows_per_tile != 0: + ca = _and(ca, ra < nrows) + cb = _and(cb, ra + 1 < nrows) + _when(ca, st_a) + _when(cb, st_b) + + body() + x_signal(ctx, flags, gsync, lane, r0, rows_per_tile) + + _guard_tile(cfg, tile, lambda: work(tile), ctx) + if pdl_trigger == "end": + cute.arch.griddepcontrol_launch_dependents() + + ngrid = cfg.ngrid + smem_bytes = cfg.smem_bytes + + @cute.jit + def launch( + gin: cute.Tensor, + gout: cute.Tensor, + gtw: cute.Tensor, + gk: cute.Tensor, + gfl: cute.Tensor, + gp0: cute.Tensor, + gp1: cute.Tensor, + gp2: cute.Tensor, + gp3: cute.Tensor, + gp4: cute.Tensor, + gp5: cute.Tensor, + gsync: cute.Tensor, + stream: cuda.CUstream, + ): + kern(gin, gout, gtw, gk, gfl, gp0, gp1, gp2, gp3, gp4, gp5, gsync).launch( + grid=[ngrid, 1, 1], block=[P, 1, 1], smem=smem_bytes, stream=stream + ) + + return launch, tabs.dtab, const_table(keys), tabs.itab diff --git a/build/torch-cuda/_r2c_f64/fft_plan.py b/build/torch-cuda/_r2c_f64/fft_plan.py new file mode 100644 index 0000000000000000000000000000000000000000..b900be1e7bb8e98cd91d9ce96f01710e5040c96b --- /dev/null +++ b/build/torch-cuda/_r2c_f64/fft_plan.py @@ -0,0 +1,218 @@ +"""Per-axis FFT plans: Good-Thomas split into coprime factors, mixed-radix +Stockham stages inside each factor, and the host-side twiddle tables. + +An axis of length n is viewed as a D-dimensional array (N_1, ..., N_D) with +pairwise-coprime N_d (prime-factor algorithm, no twiddles between factors). +Element (j_1..j_D) holds input x[(sum_d j_d * n/N_d) mod n]; after all stages +element (k_1..k_D) holds X[(sum_d k_d * e_d) mod n] with +e_d = (n/N_d) * ((n/N_d)^-1 mod N_d). + +Inside factor d the transform is a sequence of Stockham DIT stages with radices +r_1..r_S. Stage s (Ns = r_1...r_{s-1}) has N_d/r butterflies; butterfly b +reads positions b + m*N_d/r, multiplies input m by W_{Ns*r}^(m*k) with +k = b mod Ns, runs a radix-r DFT and writes output m to position +(b // Ns)*Ns*r + k + m*Ns. + +Twiddles are float64 values of exp(-2 pi i k / N) computed with the math module +(exact at multiples of pi/4); the Linzer-Feig tables are built as CPU torch tensors. +""" + +import math + +import torch + +from .fft_codelets import CountingBackend, CVal, Ops, RScale, RtRegistry, dft + + +def factorize(n): + f = [] + d = 2 + while d * d <= n: + while n % d == 0: + f.append(d) + n //= d + d += 1 + if n > 1: + f.append(n) + return f + + +MAX_CODELET = 16 # largest radix we unroll in registers +SUPPORTED_PRIMES = (2, 3, 5, 7, 11, 13) + + +def is_smooth(n): + return all(p in SUPPORTED_PRIMES for p in factorize(n)) + + +def coprime_split(n): + """Split n into pairwise coprime prime-power factors.""" + pw = {} + for p in factorize(n): + pw[p] = pw.get(p, 1) * p + return [pw[p] for p in sorted(pw)] + + +def radix_options(N): + """Candidate radix sequences for a prime-power factor N (product == N).""" + f = factorize(N) + p = f[0] + if N == 1: + return [[]] + out = [] + + def rec(rem, prefix): + if rem == 1: + out.append(list(prefix)) + return + r = p + while rem % r == 0 and r <= MAX_CODELET: + rec(rem // r, prefix + [r]) + r *= p + + rec(N, []) + return out + + +class Stage: + __slots__ = ("Nd", "Ns", "dim", "nb", "r") + + def __init__(self, dim, r, Ns, Nd): + self.dim = dim + self.r = r + self.Ns = Ns + self.Nd = Nd + self.nb = Nd // r # butterflies along this dim + + def __repr__(self): + return f"Stage(d={self.dim}, r={self.r}, Ns={self.Ns}, Nd={self.Nd})" + + +class AxisPlan: + def __init__(self, n, dims, radices): + self.n = n + self.dims = list(dims) + self.radices = [list(r) for r in radices] + assert math.prod(self.dims) == n + for d, rs in zip(self.dims, self.radices): + assert math.prod(rs) == d, (d, rs) + self.D = len(self.dims) + # smem layout: dim 0 slowest + self.strides = [] + s = 1 + for d in reversed(self.dims): + self.strides.append(s) + s *= d + self.strides = list(reversed(self.strides)) + self.in_coef = [n // d for d in self.dims] + self.out_coef = [(n // d) * pow(n // d, -1, d) % n if d > 1 else 0 for d in self.dims] + self.stages = [] + for di, (Nd, rs) in enumerate(zip(self.dims, self.radices)): + Ns = 1 + for r in rs: + self.stages.append(Stage(di, r, Ns, Nd)) + Ns *= r + + def __repr__(self): + return f"AxisPlan(n={self.n}, dims={self.dims}, radices={self.radices})" + + +_OCTANT = ( + (1.0, 0.0), + (math.sqrt(0.5), -math.sqrt(0.5)), + (0.0, -1.0), + (-math.sqrt(0.5), -math.sqrt(0.5)), + (-1.0, 0.0), + (-math.sqrt(0.5), math.sqrt(0.5)), + (0.0, 1.0), + (math.sqrt(0.5), math.sqrt(0.5)), +) + + +def exact_twiddle(N, k): + """W_N^k = exp(-2 pi i k / N) for k in [0, N) as a Python complex (exact at multiples of pi/4).""" + q, rem = divmod(8 * k, N) + if rem == 0: + return complex(*_OCTANT[q]) + th = 2.0 * math.pi * k / N + return complex(math.cos(th), -math.sin(th)) + + +def exact_twiddles(N): + """[W_N^i for i in range(N)] (see exact_twiddle).""" + return [exact_twiddle(N, i) for i in range(N)] + + +# --------------------------------------------------------------------------- +# Linzer-Feig runtime twiddles: per-row tables +# --------------------------------------------------------------------------- + + +def lf_params(L, e): + """Twiddle W_L^e = sigma * (1 + i t) (f=0) or sigma * (t + i) (f=1). + Returns (sigma, T1, t, f) with T1 = -t (f=0) or t (f=1).""" + w = exact_twiddle(L, e % L) + c, s = w.real, w.imag + if abs(c) >= abs(s): + t = s / c + return c, -t, t, 0 + t = c / s + return s, t, t, 1 + + +_LF_KEYS = {} + + +def lf_registry_keys(r): + """Runtime ratio keys a radix-r codelet needs when inputs m>=1 carry sigma_m.""" + if r in _LF_KEYS: + return _LF_KEYS[r] + be = CountingBackend() + rt = RtRegistry() + o = Ops(be, rt) + xs = [CVal(1.0, 0.5)] + for m in range(1, r): + e = [0] * r + e[m] = 1 + xs.append(CVal(1.0, 0.5, RScale(1.0, e))) + ys = dft(o, xs) + [o.finalize(y) for y in ys] + _LF_KEYS[r] = (list(rt.keys), be.nfma) + return _LF_KEYS[r] + + +def lf_stage_table(st): + """Per-k rows for a twiddled stage, as CPU tensors (float64 [Ns, dlen], int32 [Ns, ilen]). + + doubles row: [T1_1, t_1, ..., T1_{r-1}, t_{r-1}, ratio_0, ..., ratio_{R-1}] (+pad to even) + ints row: [f_1, ..., f_{r-1}] (+pad to multiple of 4) + """ + r, Ns = st.r, st.Ns + L = Ns * r + keys, _ = lf_registry_keys(r) + R = len(keys) + dlen = 2 * (r - 1) + R + dlen += dlen % 2 + ilen = (r - 1 + 3) // 4 * 4 + drows, irows = [], [] + for k in range(Ns): + drow = [0.0] * dlen + irow = [0] * ilen + sig = [1.0] * r + for m in range(1, r): + sigma, T1, t, f = lf_params(L, m * k) + sig[m] = sigma + drow[2 * (m - 1)] = T1 + drow[2 * (m - 1) + 1] = t + irow[m - 1] = f + for j, (c, e) in enumerate(keys): + v = c + for m, p in enumerate(e): + if p: + v *= sig[m] ** p + drow[2 * (r - 1) + j] = v + drows.append(drow) + irows.append(irow) + dtab = torch.tensor(drows, dtype=torch.float64).reshape(Ns, dlen) + itab = torch.tensor(irows, dtype=torch.int32).reshape(Ns, ilen) + return dtab, itab, R diff --git a/build/torch-cuda/_r2c_f64/fft_stages.py b/build/torch-cuda/_r2c_f64/fft_stages.py new file mode 100644 index 0000000000000000000000000000000000000000..6c4721ae7ae94968c4b25b085851f04d2a12a5c0 --- /dev/null +++ b/build/torch-cuda/_r2c_f64/fft_stages.py @@ -0,0 +1,647 @@ +"""Stage engine: emits the Stockham stages of an AxisPlan for one warp's (or CTA's) tile. + +Runs at trace time and produces straight-line code. A tile holds C transforms of one +axis in shared memory (``Layout``: columns interleaved or padded rows); the butterflies +of a stage are dealt out in 32-lane rounds, so FP64 instructions run with full lanes. +Twiddled stages read per-k rows from ``StageTables`` (Linzer-Feig form by default, or +plain complex W), and every constant a codelet multiplies by is loaded per lane from +``const_table``. ``emit_fft`` ping-pongs complex tiles between shared-memory buffers (or +keeps a stage's outputs in registers across one buffer, ``hold``); ``emit_fft_xs`` +exchanges real and imaginary parts through a half-size buffer. With a nonzero ``emu`` +fraction, part of each stage's unit-coefficient add/sub pairs runs on the integer pipes +(``EmuPolicy`` + ``dsl_ops.emu_addsub``). +""" + +import cutlass +import torch +from cutlass import cute + +from .dsl_ops import _when, vload2, vload_i4, vstore2 +from .fft_codelets import CountingBackend, CVal, EmuPolicy, Ops, RtRegistry, dft, dft_half +from .fft_plan import exact_twiddles, lf_registry_keys, lf_stage_table + + +def collect_consts(radices, extra=()): + be = CountingBackend() + o = Ops(be) + for r in radices: + xs = [CVal(1.0, 0.5) for _ in range(r)] + ys = dft(o, xs) + [o.finalize(y) for y in ys] + keys = set(be.consts.keys()) + keys.update(float(e) for e in extra) + keys.add(1.0) + return sorted(keys) + + +def const_table(keys): + """Constants replicated per lane, layout [i][lane]. + + Slot 0 holds -0.0 (the additive identity that preserves signed zeros). + Each lane loads its own copy so ptxas cannot prove the values warp-uniform: + uniform values get moved to uniform registers, and a DFMA with a UR operand + issues ~16% slower than the all-vector-register form. + """ + base = torch.tensor([-0.0] + list(keys), dtype=torch.float64) + return base.repeat_interleave(32) + + +def load_consts(gk, keys, lane): + pk = gk.iterator + lane + mz = (pk + 0).load() + kmap = {} + for i, v in enumerate(keys): + kmap[v] = (pk + (i + 1) * 32).load() + return kmap, mz + + +def interleave(values): + """Complex values -> CPU float64 tensor [re_0, im_0, re_1, im_1, ...].""" + return torch.tensor([x for v in values for x in (v.real, v.imag)], dtype=torch.float64) + + +class StageTables: + """Device tables for a plan's twiddled stages. + + Classic: complex W rows [k][m-1]. Linzer-Feig: double rows + [T1_1, t_1, ..., ratios...] and int flag rows (see fft_plan.lf_stage_table). + Offsets are in doubles / ints; dtab / itab are CPU float64 / int32 tensors. + """ + + def __init__(self, plan, lf, extra_complex=None): + self.lf = lf + dparts, iparts = [], [] + self.doff, self.ioff, self.dlen, self.ilen, self.R = [], [], [], [], [] + dpos = ipos = 0 + for st in plan.stages: + if st.Ns == 1: + for lst in (self.doff, self.ioff, self.dlen, self.ilen, self.R): + lst.append(-1) + continue + if lf: + dtab, itab, R = lf_stage_table(st) + dparts.append(dtab.reshape(-1)) + iparts.append(itab.reshape(-1)) + self.doff.append(dpos) + self.ioff.append(ipos) + self.dlen.append(dtab.shape[1]) + self.ilen.append(itab.shape[1]) + self.R.append(R) + dpos += dtab.numel() + ipos += itab.numel() + else: + L = st.Ns * st.r + w = exact_twiddles(L) + d = interleave([w[(m * k) % L] for k in range(st.Ns) for m in range(1, st.r)]) + dparts.append(d) + self.doff.append(dpos) + self.ioff.append(-1) + self.dlen.append(2 * (st.r - 1)) + self.ilen.append(0) + self.R.append(0) + dpos += d.numel() + self.extra_off = dpos + if extra_complex is not None: + e = interleave(extra_complex) + dparts.append(e) + dpos += e.numel() + self.dtab = torch.cat(dparts) if dparts else torch.zeros(2, dtype=torch.float64) + self.itab = torch.cat(iparts).to(torch.int32) if iparts else torch.zeros(4, dtype=torch.int32) + + +# --------------------------------------------------------------------------- +# Stage engine (trace-time code generation) +# --------------------------------------------------------------------------- + + +class Ctx: + pass + + +def _decompose_other(plan, di, o): + """o enumerates positions of all dims except di (row-major, dim 0 slowest).""" + others = [d for d in range(plan.D) if d != di] + res = {} + rem = o + for idx in reversed(range(len(others))): + d = others[idx] + Nd = plan.dims[d] + if idx == 0: + res[d] = rem + else: + res[d] = rem % Nd + rem = rem // Nd + return res + + +class Layout: + """Shared-memory layout of a tile of C transforms of length n. + + mode 'col': element L of column c at (L*C + c) (c fastest) + mode 'row': element L of row c at c*LD + L + L//16 (L fastest, padded) + Returns offsets in complex elements. + """ + + def __init__(self, mode, n, C): + self.mode = mode + self.n = n + self.C = C + if mode == "row": + self.LD = n + (n + 15) // 16 + 1 + self.size = self.LD * C + else: + self.size = n * C + + def off(self, L, c): + if self.mode == "row": + return c * self.LD + L + L // 16 + return L * self.C + c + + def split(self, q, nbeta): + """butterfly id q -> (c, beta).""" + if self.mode == "row": + return q // nbeta, q % nbeta + return q % self.C, q // self.C + + +def _stage_tw_load(ctx, si, st, b): + """Load the twiddle data of butterfly b of stage si (k = b mod Ns).""" + r = st.r + k = b % st.Ns + tb = ctx.tabs + twp = ctx.gtw.iterator + tb.doff[si] + k * tb.dlen[si] + if tb.lf: + flp = ctx.gfl.iterator + tb.ioff[si] + k * tb.ilen[si] + flags = [] + for j in range(tb.ilen[si] // 4): + flags += vload_i4(flp + 4 * j) + R = tb.R[si] + rv = [] + for j in range((R + 1) // 2): + a0, a1 = vload2(twp + 2 * (r - 1) + 2 * j) + rv += [a0, a1] + pairs = [vload2(twp + 2 * (m - 1)) for m in range(1, r)] + return flags, rv[:R], pairs + return [vload2(twp + (m - 1) * 2) for m in range(1, r)] + + +def _butterfly_b(ctx, plan, lay, C, lane, st, t): + n = plan.n + nbeta = n // st.r + B = C * nbeta + split = getattr(ctx, "split", 1) + if split == 1: + q = lane + t * 32 + if B % 32 != 0 and (t + 1) * 32 > B: + q = cutlass.min(q, B - 1) + else: + q = lane + (ctx.wsub + t * split) * 32 + if B % (32 * split) != 0: + q = cutlass.min(q, B - 1) + _, beta = lay.split(q, nbeta) + return beta % st.nb if nbeta != st.nb else beta + + +def preload_twiddles(ctx, plan, lay, C, lane): + """Issue every twiddled stage's table loads up front (keyed by (stage, round)).""" + out = {} + split = getattr(ctx, "split", 1) + for si, st in enumerate(plan.stages): + if st.Ns == 1: + continue + T = (C * plan.n // st.r + 31) // 32 + for t in range((T + split - 1) // split): + out[(si, t)] = _stage_tw_load(ctx, si, st, _butterfly_b(ctx, plan, lay, C, lane, st, t)) + return out + + +_EMU_PLANS = {} + + +def stage_emu_plan(r, twiddled, lf, frac): + """(ratio, Q, nfma) for one butterfly of a radix-r stage: Q real ops sit in unit-coefficient + add/sub pairs; emulating `ratio` of them moves about `frac` of the stage's floating-point ops + off the FP64 pipe (dry run of exactly the code emit_fft generates).""" + key = (r, bool(twiddled), bool(lf), float(frac)) + if key not in _EMU_PLANS: + be = CountingBackend() + pol = EmuPolicy() + if twiddled and lf: + keys, _ = lf_registry_keys(r) + o = Ops(be, RtRegistry(values=[1.0] * len(keys), keys=keys), emu=pol) + xs = [CVal(1.0, 0.5)] + [o.lf_twiddle(CVal(1.0, 0.5), m, -0.25, 0.25, False, r) for m in range(1, r)] + else: + o = Ops(be, None, emu=pol) + xs = [CVal(1.0, 0.5) for _ in range(r)] + if twiddled: + xs = [xs[0]] + [o.mul_rt(x, 0.6, 0.8) for x in xs[1:]] + ys = dft(o, xs) + [o.finalize(y) for y in ys] + nf = be.nfma + ratio = min(1.0, frac * nf / pol.Q) if pol.Q else 0.0 + _EMU_PLANS[key] = (ratio, pol.Q, nf) + return _EMU_PLANS[key] + + +def emit_fft(ctx, plan, lay, C, lane, bufs, load_first, store_last, last_natural, hold, pre_tw=None): + """Emit all stages of `plan` for one warp's tile of C transforms. + + load_first(c, gidx) -> (re, im); store_last(c, kidx, re, im). + bufs: smem pointers (this warp's region). With hold=True a single buffer is + used and smem->smem stages keep their outputs in registers across the warp + barrier; otherwise stages ping-pong between two buffers. + If last_natural, the last stage writes output kidx at smem position kidx of + the returned buffer instead of calling store_last. + """ + o_ops = ctx.ops + n = plan.n + S = len(plan.stages) + cur_src = None + final_buf = None + for si, st in enumerate(plan.stages): + first = si == 0 + last = si == S - 1 + di = st.dim + r = st.r + nb = st.nb + nbeta = n // r + B = C * nbeta + T = (B + 31) // 32 + stride_d = plan.strides[di] + src = cur_src + writes_smem = (not last) or last_natural + if writes_smem: + if hold or first: + dst = bufs[0] + else: + dst = bufs[1] if src is bufs[0] else bufs[0] + else: + dst = None + hazard = writes_smem and (src is not None) and (src is dst) + pending = [] + split = getattr(ctx, "split", 1) + # half butterflies: two lanes per radix-2/4 butterfly (outputs split by parity) when that + # does not add rounds, i.e. when the stage leaves lanes idle (latency-bound small sizes) + half = ( + getattr(ctx, "half", False) + and r in (2, 4) + and (2 * B + 32 * split - 1) // (32 * split) == (B + 32 * split - 1) // (32 * split) + ) + Bt = 2 * B if half else B + T = (Bt + 31) // 32 + nrounds = (T + split - 1) // split + + def prep( + t, + si=si, + st=st, + first=first, + di=di, + r=r, + nb=nb, + nbeta=nbeta, + B=B, + Bt=Bt, + half=half, + stride_d=stride_d, + src=src, + split=split, + ): + if split == 1: + q = lane + t * 32 + pred = None + if Bt % 32 != 0 and (t + 1) * 32 > Bt: + pred = q < Bt + q = cutlass.min(q, Bt - 1) + else: + q = lane + (ctx.wsub + t * split) * 32 + pred = None + if Bt % (32 * split) != 0: + pred = q < Bt + q = cutlass.min(q, Bt - 1) + hp = None + if half: + hp = q >= B + q = q - cutlass.Int32(cutlass.select_(hp, cutlass.Int32(B), cutlass.Int32(0))) + c, beta = lay.split(q, nbeta) + b = beta % nb if nbeta != nb else beta + o = beta // nb if nbeta != nb else None + oth = _decompose_other(plan, di, o) if plan.D > 1 else {} + L0 = None + for d, jd in oth.items(): + tt = jd * plan.strides[d] + L0 = tt if L0 is None else L0 + tt + xs = [] + for m in range(r): + pos = b + (m * nb) + if first: + if plan.D > 1: + g = pos * plan.in_coef[di] + for d, jd in oth.items(): + g = g + jd * plan.in_coef[d] + g = g % n + else: + g = pos + re, im = load_first(c, g) + else: + L = pos * stride_d if L0 is None else L0 + pos * stride_d + re, im = vload2(src + lay.off(L, c) * 2) + xs.append(CVal(re, im)) + return pred, c, b, oth, L0, xs, hp + + # with loads_upfront, every global load of the first stage is issued before any of + # its arithmetic (one memory latency per tile instead of one per batch of rounds) + if first and getattr(ctx, "loads_upfront", False): + hook = getattr(ctx, "before_first_loads", None) + if hook is not None: + hook() + preps = [prep(t) for t in range(nrounds)] + hook = getattr(ctx, "after_first_loads", None) + if hook is not None: + hook() + else: + preps = None + for t in range(nrounds): + pred, c, b, oth, L0, xs, hp = preps[t] if preps is not None else prep(t) + pol = None + emu_frac = getattr(ctx, "emu", 0.0) + if isinstance(emu_frac, (list, tuple)): + emu_frac = emu_frac[si] if si < len(emu_frac) else emu_frac[-1] + if emu_frac > 0 and not half: + ratio, _, _ = stage_emu_plan(r, st.Ns > 1, ctx.tabs.lf, emu_frac) + if ratio > 0: + pol = EmuPolicy(ratio) + b_ops = o_ops if pol is None else Ops(o_ops.be, None, emu=pol) + if st.Ns > 1: + tw = pre_tw.get((si, t)) if (pre_tw is not None and not half) else None + if tw is None: + tw = _stage_tw_load(ctx, si, st, b) + if ctx.tabs.lf: + flags, rv, pairs = tw + keys, _ = lf_registry_keys(r) + b_ops = Ops(o_ops.be, RtRegistry(values=rv, keys=keys), emu=pol) + for m in range(1, r): + T1, tt = pairs[m - 1] + xs[m] = b_ops.lf_twiddle(xs[m], m, T1, tt, flags[m - 1] != 0, r) + else: + for m in range(1, r): + wr, wi = tw[m - 1] + xs[m] = o_ops.mul_rt(xs[m], wr, wi) + if half: + ys = dft_half(b_ops, xs, hp) + hpi = cutlass.Int32(cutlass.select_(hp, cutlass.Int32(1), cutlass.Int32(0))) + ms = [hpi + 2 * j for j in range(r // 2)] + else: + ys = dft(b_ops, xs) + ms = list(range(r)) + ys = [b_ops.finalize(y) for y in ys] + if st.Ns > 1: + base_pos = (b // st.Ns) * (st.Ns * r) + (b % st.Ns) + else: + base_pos = b * r + stores = [] + for j, m in enumerate(ms): + pos = base_pos + m * st.Ns + y = ys[j] + if last: + if plan.D > 1: + kidx = pos * plan.out_coef[di] + for d, kd in oth.items(): + kidx = kidx + kd * plan.out_coef[d] + kidx = kidx % n + else: + kidx = pos + stores.append(("s" if last_natural else "g", kidx, y)) + else: + L = pos * stride_d if L0 is None else L0 + pos * stride_d + stores.append(("s", L, y)) + + def do_stores(stores=stores, c=c, dst=dst): + for kind, idx, y in stores: + if kind == "s": + vstore2(dst + lay.off(idx, c) * 2, y.re, y.im) + else: + store_last(c, idx, y.re, y.im) + + if hazard: + pending.append((pred, do_stores)) + else: + _when(pred, do_stores) + sync = getattr(ctx, "sync", cute.arch.sync_warp) + if hazard: + sync() + for pred, fn in pending: + _when(pred, fn) + if writes_smem: + sync() + cur_src = dst + final_buf = dst + return final_buf + + +# --------------------------------------------------------------------------- +# Split exchange: smem holds only real or only imaginary parts at a time +# --------------------------------------------------------------------------- +# +# Between two stages the outputs are exchanged through a buffer of n*C doubles (half of +# the complex tile): store real parts, barrier, load real parts of the next stage's inputs, +# barrier, store imaginary parts, barrier, load imaginary parts. Half the shared memory +# per warp (higher occupancy where smem limits it) for three extra barriers per boundary. + + +def _stage_spec(ctx, plan, C, si): + st = plan.stages[si] + n = plan.n + r = st.r + nbeta = n // r + B = C * nbeta + split = getattr(ctx, "split", 1) + half = ( + getattr(ctx, "half", False) + and r in (2, 4) + and (2 * B + 32 * split - 1) // (32 * split) == (B + 32 * split - 1) // (32 * split) + ) + Bt = 2 * B if half else B + T = (Bt + 31) // 32 + nrounds = (T + split - 1) // split + return st, r, nbeta, B, half, Bt, nrounds + + +def _stage_geom(ctx, plan, lay, C, lane, si, t): + """Index math of round t of stage si (as in emit_fft's prep): (pred, c, b, oth, L0, hp).""" + st, _, nbeta, B, half, Bt, _ = _stage_spec(ctx, plan, C, si) + split = getattr(ctx, "split", 1) + nb = st.nb + di = st.dim + if split == 1: + q = lane + t * 32 + pred = None + if Bt % 32 != 0 and (t + 1) * 32 > Bt: + pred = q < Bt + q = cutlass.min(q, Bt - 1) + else: + q = lane + (ctx.wsub + t * split) * 32 + pred = None + if Bt % (32 * split) != 0: + pred = q < Bt + q = cutlass.min(q, Bt - 1) + hp = None + if half: + hp = q >= B + q = q - cutlass.Int32(cutlass.select_(hp, cutlass.Int32(B), cutlass.Int32(0))) + c, beta = lay.split(q, nbeta) + b = beta % nb if nbeta != nb else beta + o = beta // nb if nbeta != nb else None + oth = _decompose_other(plan, di, o) if plan.D > 1 else {} + L0 = None + for d, jd in oth.items(): + tt = jd * plan.strides[d] + L0 = tt if L0 is None else L0 + tt + return pred, c, b, oth, L0, hp + + +def emit_fft_xs(ctx, plan, lay, C, lane, hbuf, load_first, store_last, last_natural, pre_tw=None): + """emit_fft with split real/imaginary exchanges through the half buffer hbuf (n*C doubles). + + Returns None, or for last_natural the list of (pred, c, [(kidx, y)]) of the last stage + (outputs still in registers; the caller exchanges them).""" + o_ops = ctx.ops + n = plan.n + S = len(plan.stages) + sync = getattr(ctx, "sync", cute.arch.sync_warp) + pre_in = None + for si in range(S): + st, r, _, _, half, _, nrounds = _stage_spec(ctx, plan, C, si) + first = si == 0 + last = si == S - 1 + di = st.dim + nb = st.nb + stride_d = plan.strides[di] + geoms = [_stage_geom(ctx, plan, lay, C, lane, si, t) for t in range(nrounds)] + if first: + + def load_round(t, geoms=geoms, r=r, nb=nb, di=di): + _, c, b, oth, _, _ = geoms[t] + xs = [] + for m in range(r): + pos = b + (m * nb) + if plan.D > 1: + g = pos * plan.in_coef[di] + for d, jd in oth.items(): + g = g + jd * plan.in_coef[d] + g = g % n + else: + g = pos + re, im = load_first(c, g) + xs.append(CVal(re, im)) + return xs + + if getattr(ctx, "loads_upfront", False): + hook = getattr(ctx, "before_first_loads", None) + if hook is not None: + hook() + ins = [load_round(t) for t in range(nrounds)] + hook = getattr(ctx, "after_first_loads", None) + if hook is not None: + hook() + else: + ins = None + else: + ins = pre_in + outs = [] + for t in range(nrounds): + pred, c, b, oth, L0, hp = geoms[t] + xs = list(ins[t]) if ins is not None else load_round(t) + pol = None + emu_frac = getattr(ctx, "emu", 0.0) + if isinstance(emu_frac, (list, tuple)): + emu_frac = emu_frac[si] if si < len(emu_frac) else emu_frac[-1] + if emu_frac > 0 and not half: + ratio, _, _ = stage_emu_plan(r, st.Ns > 1, ctx.tabs.lf, emu_frac) + if ratio > 0: + pol = EmuPolicy(ratio) + b_ops = o_ops if pol is None else Ops(o_ops.be, None, emu=pol) + if st.Ns > 1: + tw = pre_tw.get((si, t)) if (pre_tw is not None and not half) else None + if tw is None: + tw = _stage_tw_load(ctx, si, st, b) + if ctx.tabs.lf: + flags, rv, pairs = tw + keys, _ = lf_registry_keys(r) + b_ops = Ops(o_ops.be, RtRegistry(values=rv, keys=keys), emu=pol) + for m in range(1, r): + T1, tt = pairs[m - 1] + xs[m] = b_ops.lf_twiddle(xs[m], m, T1, tt, flags[m - 1] != 0, r) + else: + for m in range(1, r): + wr, wi = tw[m - 1] + xs[m] = o_ops.mul_rt(xs[m], wr, wi) + if half: + ys = dft_half(b_ops, xs, hp) + hpi = cutlass.Int32(cutlass.select_(hp, cutlass.Int32(1), cutlass.Int32(0))) + ms = [hpi + 2 * j for j in range(r // 2)] + else: + ys = dft(b_ops, xs) + ms = list(range(r)) + ys = [b_ops.finalize(y) for y in ys] + if st.Ns > 1: + base_pos = (b // st.Ns) * (st.Ns * r) + (b % st.Ns) + else: + base_pos = b * r + items = [] + for j, m in enumerate(ms): + pos = base_pos + m * st.Ns + if last: + if plan.D > 1: + kidx = pos * plan.out_coef[di] + for d, kd in oth.items(): + kidx = kidx + kd * plan.out_coef[d] + kidx = kidx % n + else: + kidx = pos + items.append((kidx, ys[j])) + else: + L = pos * stride_d if L0 is None else L0 + pos * stride_d + items.append((L, ys[j])) + outs.append((pred, c, items)) + if last and not last_natural: + + def st_g(items=items, c=c): + for kidx, y in items: + store_last(c, kidx, y.re, y.im) + + _when(pred, st_g) + if last: + return outs if last_natural else None + # exchange to the next stage + nst, nr, _, _, _, _, nnrounds = _stage_spec(ctx, plan, C, si + 1) + ngeoms = [_stage_geom(ctx, plan, lay, C, lane, si + 1, t) for t in range(nnrounds)] + nstride = plan.strides[nst.dim] + + def in_offsets(t, ngeoms=ngeoms, nr=nr, nst=nst, nstride=nstride): + _, c, b, _, L0, _ = ngeoms[t] + offs = [] + for m in range(nr): + pos = b + (m * nst.nb) + L = pos * nstride if L0 is None else L0 + pos * nstride + offs.append(lay.off(L, c)) + return offs + + noffs = [in_offsets(t) for t in range(nnrounds)] + if not first: + sync() # every lane has loaded the previous exchange's imaginary parts + parts = [] + for part in (0, 1): + for pred, c, items in outs: + + def st_s(items=items, c=c, part=part): + for L, y in items: + (hbuf + lay.off(L, c)).store(y.re if part == 0 else y.im) + + _when(pred, st_s) + sync() + parts.append([[(hbuf + o).load() for o in offs] for offs in noffs]) + if part == 0: + sync() + pre_in = [[CVal(parts[0][t][m], parts[1][t][m]) for m in range(nr)] for t in range(nnrounds)] + return None diff --git a/build/torch-cuda/_r2c_f64/kernel.py b/build/torch-cuda/_r2c_f64/kernel.py new file mode 100644 index 0000000000000000000000000000000000000000..66102aa577ef99bbb5ad821fb52c4ac211f3bfeb --- /dev/null +++ b/build/torch-cuda/_r2c_f64/kernel.py @@ -0,0 +1,254 @@ +"""float64 forward real-to-complex 3D FFT for the RTX 5090 (sm_120) in CuTe DSL. + +run(input): input is a float64 CUDA tensor [nz, ny, nx] on the current device (x fastest; +tuned for nz, ny >= 8, nx >= 16, any size >= 1 works). Returns a new contiguous float64 +[nz, ny, nx//2+1, 2] equal, up to rounding, to torch.view_as_real(torch.fft.rfftn(input, +norm="backward")). Other dtypes, devices and ranks raise TypeError / ValueError before any +kernel runs; non-contiguous, negative-bit or (for the vectorized X pass) not 16-byte aligned +input is first copied to a fresh contiguous buffer. Y/Z lengths that need the direct DFT +beyond its shared-memory limit raise NotImplementedError (planner.py). + +Three kernels per call, all writing the output buffer (fft_passes.py): + X real FFT of every x-row: half-length complex FFT plus split post-processing ("r2c", + even nx) or two real rows packed into one complex FFT of length nx ("pack2") + Y complex FFT along y, in place + Z complex FFT along z, in place +A warp (or a CTA of `split` warps) owns a tile of C rows/columns and runs a Good-Thomas + +mixed-radix Stockham FFT through shared memory (fft_plan.py, fft_stages.py) built only from +register-operand DFMAs (fft_codelets.py). sm_120 issues one FP64 warp instruction per 64 +cycles per SM sub-partition, so a tuned fraction of the unit-coefficient add/sub pairs runs +as a bit-exact IEEE add/sub on the integer pipes instead (dsl_ops.emu_addsub; options emu / +emu_split), keeping both pipes busy. That adder is exact for finite normal or zero operands +and results; subnormal or overflowing results, Inf and NaN come out wrong. Axes with a +prime factor above 13 use direct-DFT kernels (fft_direct.py). + +Pass configurations come from the tuned table (tuned.py) or a cost model (planner.py); each +shape's plan is built and compiled on first use and cached per (shape, device). + +Shapes chained with per-plane X->Y counters (pdl mode 3) keep one counter buffer per launch +stream: the last Y tile of a launch zeroes the counters again, and the next launch on the +same stream starts only after that launch's Z pass (which waits for Y) has finished, so +launches on one stream can share a buffer, while launches on different streams never do. +Calls made while a CUDA graph is being captured get a buffer of their own per capture and +stream, zeroed outside the graph, so concurrent graph replays never share one either. +""" + +import cuda.bindings.driver as cuda +import torch +from cutlass import cute +from cutlass.cute.runtime import from_dlpack + +from .fft_direct import DirectColCfg, DirectRowCfg, build_direct_cols, build_direct_rows +from .fft_passes import FLAG_PAD, build_colpass, build_r2c, build_r2c_pack2 +from .planner import default_cfgs +from .tuned import TUNED + +_CACHE = {} +_SIDE = {} # device -> stream that zeroes counter buffers during graph capture + +_raw_stream = torch._C._cuda_getCurrentRawStream # device index -> current stream handle +_capturing = torch._C._cuda_isCurrentStreamCapturing +_current_device = torch._C._cuda_getDevice + + +def _dev(a, device): + """Host table (1-D CPU tensor) -> device copy.""" + return a.to(device) + + +def _check(err, what): + if err != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"{what} failed: {err}") + + +def _capture_id(handle): + err, _, cid, *_ = cuda.cuStreamGetCaptureInfo(cuda.CUstream(handle)) + _check(err, "cuStreamGetCaptureInfo") + return int(cid) + + +def _zeros_outside_capture(n, device): + """A zeroed int32 [n] tensor while the current stream is capturing a CUDA graph, without + adding a memset node to the graph: the fill runs on a private side stream (relaxed capture + mode for this thread) and has completed before this returns.""" + err, prev = cuda.cuThreadExchangeStreamCaptureMode(cuda.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_RELAXED) + _check(err, "cuThreadExchangeStreamCaptureMode") + try: + side = _SIDE.get(device) + if side is None: + err, h = cuda.cuStreamCreate(cuda.CUstream_flags.CU_STREAM_NON_BLOCKING.value) + _check(err, "cuStreamCreate") + side = _SIDE[device] = torch.cuda.ExternalStream(int(h), device=device) + with torch.cuda.stream(side): + buf = torch.zeros(n, dtype=torch.int32, device=device) + side.synchronize() + finally: + cuda.cuThreadExchangeStreamCaptureMode(prev) + return buf + + +class FFT3DPlan: + """The three compiled passes of one shape, their device tables and the X->Y counters.""" + + def __init__(self, nz, ny, nx, device): + self.shape = (nz, ny, nx) + nf = nx // 2 + 1 + self.nf = nf + self.device = device + cfgs = default_cfgs(nz, ny, nx) + xcfg, ycfg, zcfg = cfgs + self.cfgs = cfgs + all_fast = all(c is not None for c in cfgs) + # programmatic dependent launch mode: 0 off, 1 trigger at CTA end, 2 at CTA start, + # 3 start trigger + per-plane X->Y counters + pdl = TUNED.get(f"{nz}x{ny}x{nx}", {}).get("pdl", 0) + if not all_fast: + pdl = 0 + if int(pdl) == 3: + xrows = 2 * xcfg.C if xcfg.kind == "pack2" else xcfg.C + if ny % xrows != 0 or ycfg.C > nf or zcfg.C > nf: + pdl = 2 + self.pdl = int(pdl) + flagged = self.pdl == 3 + self.flagged = flagged + trig = {0: None, 1: "end", 2: "start", 3: "start"}[int(pdl)] + # counters: one 128-byte line per z plane + the Y exit counter. Without mode 3 no + # kernel touches them and this one tensor fills the argument of every launch; in mode 3 + # it only serves compilation and each launch stream gets its own (_launch_args). + self.sync = torch.zeros((nz + 1) * FLAG_PAD, dtype=torch.int32, device=device) + # the "r2c" X pass reads input row pairs as f64x2 vectors (16-byte aligned); pack2 and + # the direct DFT read scalars + self.in_align = 16 if xcfg is not None and xcfg.kind == "r2c" else 8 + fy = None + fx = None + if flagged: + fy = {"role": "y", "nf": nf, "nz": nz, "ny": ny, "ntiles": ycfg.ntiles if ycfg is not None else 0} + fx = {"ny": ny} + pwait = bool(pdl) and not flagged + # Y and Z passes (in place on the output buffer) + if ycfg is not None: + self.yl, *ytabs = build_colpass( + ycfg, nz * nf, nf, ny * nf, nf, pdl_wait=pwait, pdl_trigger=("end" if flagged else trig), flags=fy + ) + else: + self.yl, *ytabs = build_direct_cols(DirectColCfg(ny, nz * nf, nf, ny * nf, nf)) + if zcfg is not None: + self.zl, *ztabs = build_colpass(zcfg, ny * nf, ny * nf, 0, ny * nf, pdl_wait=bool(pdl)) + else: + self.zl, *ztabs = build_direct_cols(DirectColCfg(nz, ny * nf, ny * nf, 0, ny * nf)) + later = tuple(ytabs) + tuple(ztabs) + if xcfg is not None: + builder = build_r2c_pack2 if xcfg.kind == "pack2" else build_r2c + self.xl, *xtabs = builder( + xcfg, nx, nz * ny, pdl_trigger=trig, prefetch_bytes=tuple(int(a.nbytes) for a in later), flags=fx + ) + self.x_prefetch = True + else: + self.xl, *xtabs = build_direct_rows(DirectRowCfg(nx, nz * ny)) + self.x_prefetch = False + self.nx_t, self.ny_t = len(xtabs), len(ytabs) + self.tabs = [_dev(t, device) for t in list(xtabs) + list(ytabs) + list(ztabs)] + self.compiled = None + self._args = {} # launch key (stream, or (stream, capture id)) -> launch arguments + + def _split(self, t): + a, b = self.nx_t, self.nx_t + self.ny_t + return t[:a], t[a:b], t[b:] + + def _pass_args(self, sync): + """Table + counter arguments of the X, Y and Z launchers for one counter tensor.""" + tx, ty, tz = self._split(self._ctabs) + later = list(ty) + list(tz) + fs = from_dlpack(sync) + xargs = list(tx) + ((later + [later[0]] * (6 - len(later)) + [fs]) if self.x_prefetch else []) + ty = list(ty) + ([fs] if self.cfgs[1] is not None else []) + tz = list(tz) + ([fs] if self.cfgs[2] is not None else []) + return xargs, ty, tz + + def compile(self, inp, out, st): + fi = from_dlpack(inp.view(-1)) + fo = from_dlpack(out.view(-1)) + self._ctabs = [from_dlpack(x) for x in self.tabs] + xargs, ty, tz = self._pass_args(self.sync) + cx = cute.compile(self.xl, fi, fo, *xargs, st) + cy = cute.compile(self.yl, fo, *ty, st) + cz = cute.compile(self.zl, fo, *tz, st) + self.compiled = (cx, cy, cz) + + def _launch_args(self, key, handle): + """(stream, X args, Y args, Z args, counters) for launches on stream `handle`.""" + if self.flagged: + n = self.sync.numel() + sync = ( + _zeros_outside_capture(n, self.device) + if isinstance(key, tuple) + # zeroed on this stream, ahead of the first launch that uses it + else torch.zeros(n, dtype=torch.int32, device=self.device) + ) + else: + sync = self.sync + args = (cuda.CUstream(handle), *self._pass_args(sync), sync) + self._args[key] = args + return args + + def launch(self, inp, out): + handle = _raw_stream(inp.get_device()) + key = (handle, _capture_id(handle)) if _capturing() else handle + args = self._args.get(key) + if args is None: + if self.compiled is None: + self.compile(inp, out, cuda.CUstream(handle)) + args = self._launch_args(key, handle) + cx, cy, cz = self.compiled + st, xargs, ty, tz, _ = args + fi = from_dlpack(inp.view(-1)) + fo = from_dlpack(out.view(-1)) + cx(fi, fo, *xargs, st) + cy(fo, *ty, st) + cz(fo, *tz, st) + + +def _get_plan(shape, device): + key = (tuple(shape), str(device)) + p = _CACHE.get(key) + if p is None: + p = FFT3DPlan(*shape, device) + _CACHE[key] = p + return p + + +def _validate(x): + """Reject what the kernels cannot take; return x as a contiguous tensor without the + negative bit (a copy when needed).""" + if not isinstance(x, torch.Tensor): + raise TypeError(f"run() expects a torch.Tensor, got {type(x).__name__}") + if x.dtype != torch.float64: + raise TypeError(f"run() expects a float64 tensor, got {x.dtype}") + if not x.is_cuda: + raise ValueError(f"run() expects a CUDA tensor, got one on {x.device}") + if x.dim() != 3: + raise ValueError(f"run() expects a 3-D tensor [nz, ny, nx], got shape {tuple(x.shape)}") + if 0 in x.shape: + raise ValueError(f"run() needs every dimension >= 1, got shape {tuple(x.shape)}") + if x.get_device() != _current_device(): + raise ValueError( + f"input is on {x.device} but the current CUDA device is cuda:{_current_device()}; " + "call run() under torch.cuda.device(input.device)" + ) + if x.requires_grad: + raise ValueError("run() does not support autograd; pass input.detach()") + if x.is_neg(): + x = x.resolve_neg() + return x.contiguous() + + +def run(input): + """rfftn of a float64 [nz, ny, nx] CUDA tensor -> new float64 [nz, ny, nx//2+1, 2] (see module docstring).""" + x = _validate(input) + nz, ny, nx = x.shape + plan = _get_plan(x.shape, x.device) + if x.data_ptr() % plan.in_align: + x = x.clone(memory_format=torch.contiguous_format) + out = torch.empty((nz, ny, nx // 2 + 1, 2), dtype=torch.float64, device=x.device) + plan.launch(x, out) + return out diff --git a/build/torch-cuda/_r2c_f64/planner.py b/build/torch-cuda/_r2c_f64/planner.py new file mode 100644 index 0000000000000000000000000000000000000000..16a7892a6524bb9caf52d8f80a3cd8ed5af474b5 --- /dev/null +++ b/build/torch-cuda/_r2c_f64/planner.py @@ -0,0 +1,222 @@ +"""Per-shape pass configurations: the tuned table, else an FP64-issue-slot cost model. + +For every pass the model enumerates plans (Good-Thomas factors x radix sequences of at +most three stages per factor), tile sizes C, register-hold and shared-memory layouts, +estimates the FP64 slots per SM including occupancy (shared memory, registers) and picks +the cheapest. +The X pass compares the half-length "r2c" route (even nx) with "pack2" (two rows per +complex FFT of length nx). A pass the codelets cannot run (a prime factor above 13, or +no candidate that fits) gets None: the direct-DFT fallback (fft_direct.py), which handles +Y/Z axes up to length MAX_DIRECT_COL_LEN (6336); longer ones raise NotImplementedError. +""" + +import math + +from .fft_codelets import CountingBackend, CVal, Ops, dft +from .fft_direct import MAX_DIRECT_COL_LEN +from .fft_passes import NUM_SMS, PassCfg, _nbuf, pack2_slots, r2c_slots +from .fft_plan import AxisPlan, coprime_split, is_smooth, radix_options +from .fft_stages import Layout, collect_consts +from .tuned import TUNED + + +def codelet_cost(r): + be = CountingBackend() + o = Ops(be) + xs = [CVal(1.0, 0.5) for _ in range(r)] + ys = dft(o, xs) + [o.finalize(y) for y in ys] + return be.nfma + + +SMEM_PER_SM = 99 * 1024 +REGS_PER_SM = 65536 +MAX_WARPS_SM = 48 + + +def _plans_for(n): + dims = coprime_split(n) + opts = [radix_options(d) for d in dims] + out = [] + + def rec(i, cur): + if i == len(dims): + out.append(AxisPlan(n, dims, cur)) + return + for o in opts[i]: + if len(o) > 3: + continue + rec(i + 1, cur + [o]) + + rec(0, []) + return out + + +def _stage_slots(plan, C): + tot = 0 + for st in plan.stages: + B = C * plan.n // st.r + T = (B + 31) // 32 + c = codelet_cost(st.r) + (4 * (st.r - 1) if st.Ns > 1 else 0) + tot += T * 32 * c + return tot + + +def _regs_est(plan, C, hold, kind): + rmax = max(st.r for st in plan.stages) + nconst = len(collect_consts([st.r for st in plan.stages])) + regs = 40 + 6 * rmax + 2 * nconst + if hold: + S = len(plan.stages) + for si, st in enumerate(plan.stages): + hz = (si > 0) and ((si < S - 1) or kind == "r2c") + if hz: + T = (C * plan.n // st.r + 31) // 32 + regs = max(regs, 40 + 4 * st.r * T + 2 * nconst) + return min(255, regs) + + +def _util(w): + # FP64 pipe utilisation vs concurrent warps per SM (latency hiding) + return min(1.0, 0.55 + 0.04 * w) + + +def _model_time(ntiles, slots, smem_warp, regs): + regs_alloc = (regs + 7) // 8 * 8 + k = MAX_WARPS_SM + if smem_warp > 0: + k = min(k, SMEM_PER_SM // smem_warp) + k = min(k, REGS_PER_SM // (32 * regs_alloc)) + if k < 1: + return None, 0 + m = math.ceil(ntiles / NUM_SMS) # tiles per SM + conc = min(m, k) + pipes = min(4, conc) + t = m * slots / (0.5 * pipes) / _util(conc) + return t, k + + +def _candidates(kind, n, nitems, extra_slot=0): + """Enumerate (time, cfg-args) for a pass.""" + out = [] + for plan in _plans_for(n): + S = len(plan.stages) + for C in (1, 2, 4, 8, 16, 32): + if C > nitems: + continue + slots = _stage_slots(plan, C) + if extra_slot: + Bp = C * extra_slot + slots += ((Bp + 31) // 32) * 32 * (4 if kind == "pack2" else 12) + ntiles = (nitems + C - 1) // C + for hold in (False, True) if (S > 2 or (kind in ("r2c", "pack2") and S > 1)) else (False,): + for mode in ("row", "col") if kind in ("r2c", "pack2") else ("col",): + lay = Layout(mode, n, C) + nbuf = _nbuf(S, hold, kind) + smem_warp = nbuf * lay.size * 16 + if smem_warp > 48 * 1024: + continue + regs = _regs_est(plan, C, hold, kind) + if regs >= 255: + continue + t, k = _model_time(ntiles, slots, smem_warp, regs) + if t is None: + continue + if kind == "col" and C < 4: + t *= 1.0 + 0.04 * (4 - C) + out.append((t, plan, C, hold, mode, k)) + out.sort(key=lambda x: x[0]) + return out + + +def _pick_W(ntiles, k): + for W in (4, 2, 1): + if (ntiles >= W * 2 * NUM_SMS or W == 1) and W <= k: + return W + return 1 + + +def _cfg_from_dict(kind, n, nitems, d): + plan = AxisPlan(n, coprime_split(n), d["radices"]) + C = d["C"] + return PassCfg( + kind, + n, + plan, + C, + d["W"], + d["hold"], + (nitems + C - 1) // C, + d["mode"], + d.get("pre", False), + d.get("split", 1), + d.get("stg", False), + d.get("smg", False), + d.get("ef", False), + d.get("rev", False), + d.get("half", False), + d.get("emu", 0.0), + d.get("lf", True), + d.get("emu_split", None), + d.get("xs", False), + ) + + +def _fast_cfg(kind, n, nitems, with_time=False): + """Best modelled config for a fast (codelet) pass, or None if the length is unsupported.""" + if not is_smooth(n): + return (None, None) if with_time else None + try: + if kind == "r2c": + cands = _candidates("r2c", n, nitems, extra_slot=r2c_slots(n)) + elif kind == "pack2": + cands = _candidates("pack2", n, nitems, extra_slot=pack2_slots(n)) + else: + cands = _candidates("col", n, nitems) + except Exception: # noqa: BLE001 - a length the planner cannot handle falls back to the direct DFT + cands = [] + if not cands: + return (None, None) if with_time else None + t, plan, C, hold, mode, k = cands[0] + ntiles = (nitems + C - 1) // C + cfg = PassCfg(kind, n, plan, C, _pick_W(ntiles, k), hold, ntiles, mode) + return (cfg, t) if with_time else cfg + + +def _x_cfg(nz, ny, nx): + nrows = nz * ny + npairs = (nrows + 1) // 2 + p2, tp2 = _fast_cfg("pack2", nx, npairs, with_time=True) + if nx % 2 == 0: + r2, tr2 = _fast_cfg("r2c", nx // 2, nrows, with_time=True) + if r2 is not None and (p2 is None or tr2 <= tp2): + return r2 + return p2 + + +def default_cfgs(nz, ny, nx): + """Per-pass configs; None marks a pass that uses the direct-DFT fallback. + + Raises NotImplementedError for a Y or Z axis that needs the direct DFT but is longer + than its shared-memory limit (fft_direct.MAX_DIRECT_COL_LEN).""" + nf = nx // 2 + 1 + t = TUNED.get(f"{nz}x{ny}x{nx}") + if t is not None: + if t["x"].get("kind") == "pack2": + xcfg = _cfg_from_dict("pack2", nx, (nz * ny + 1) // 2, t["x"]) + else: + xcfg = _cfg_from_dict("r2c", nx // 2, nz * ny, t["x"]) + return ( + xcfg, + _cfg_from_dict("col", ny, nz * nf, t["y"]), + _cfg_from_dict("col", nz, ny * nf, t["z"]), + ) + cfgs = (_x_cfg(nz, ny, nx), _fast_cfg("col", ny, nz * nf), _fast_cfg("col", nz, ny * nf)) + for axis, n, cfg in (("y", ny, cfgs[1]), ("z", nz, cfgs[2])): + if cfg is None and n > MAX_DIRECT_COL_LEN: + raise NotImplementedError( + f"shape {nz}x{ny}x{nx}: the {axis} axis (length {n}) has no codelet plan (prime factor " + f"above 13, or too long for one tile) and the direct-DFT fallback keeps a whole column " + f"in shared memory, which limits it to length {MAX_DIRECT_COL_LEN}" + ) + return cfgs diff --git a/build/torch-cuda/_r2c_f64/tuned.py b/build/torch-cuda/_r2c_f64/tuned.py new file mode 100644 index 0000000000000000000000000000000000000000..068756822d06a27c02c9599115de2957c96eb754 --- /dev/null +++ b/build/torch-cuda/_r2c_f64/tuned.py @@ -0,0 +1,345 @@ +"""Tuned pass configurations per shape "NZxNYxNX" (measured on the RTX 5090). + +Shapes not listed here use the cost model in planner.py. Per entry: "pdl" (pass chaining: +0 none, 1/2 programmatic dependent launch with the trigger at CTA end/start, 3 start +trigger + per-plane X->Y counters), "t_us" (measured time, informational) and one dict per +pass "x", "y", "z" with + kind X pass route, "pack2" or "r2c" (the default) + radices radix sequence per Good-Thomas factor of the pass length + C, W transforms per tile, warps per CTA + hold keep stage outputs in registers across one shared buffer (else ping-pong) + mode shared-memory tile layout, "col" (interleaved) or "row" (padded rows) + split several warps share one tile (CTA barrier between stages) + pre issue all twiddle loads before the first stage + smg/stg X pass with one CTA per SM and rank-ordered first loads, so the cold DRAM read + reaches one warp per SM sub-partition first + ef/rev X input read with an L2 evict-first policy, Y tiles in reverse plane order + (problems larger than L2: Y starts on what X wrote last); on the Z pass, ef + writes the final results with an L2 evict-first policy, so the Y output still + in L2 is not displaced before Z reads it + half two lanes per radix-2/4 butterfly in stages that leave lanes idle + xs exchange real and imaginary parts between stages through a half-size buffer + (half the shared memory per tile, three more barriers per exchange) + emu fraction of each stage's floating-point ops executed as bit-exact IEEE add/sub + pairs on the integer pipes (one value, or one per stage); the FP64 pipe and the + integer ALU then work side by side + emu_split pack2 X pass: run the two-row split as halved integer add/sub pairs (default: + when emu is on) + lf runtime twiddles in Linzer-Feig form (default) or as plain 4-FMA products, which + keep every codelet addition unit-coefficient and therefore emulatable +""" + +TUNED = { + "128x128x128": { + "pdl": 3, + "t_us": 64.62, + "x": { + "C": 4, + "W": 1, + "emu_split": True, + "hold": True, + "kind": "pack2", + "mode": "col", + "pre": False, + "radices": [[16, 8]], + "split": 1, + }, + "y": { + "C": 4, + "W": 2, + "emu": 0.3, + "hold": False, + "lf": False, + "mode": "col", + "pre": False, + "radices": [[16, 8]], + "split": 2, + }, + "z": { + "C": 4, + "W": 2, + "emu": 0.15, + "hold": False, + "lf": True, + "mode": "col", + "pre": False, + "radices": [[8, 2, 8]], + "split": 2, + }, + }, + "192x192x192": { + "pdl": 3, + "t_us": 193.92, + "x": { + "C": 2, + "W": 2, + "emu": [0.15, 0.15, 0.15, 0.0], + "hold": True, + "kind": "pack2", + "lf": True, + "mode": "row", + "pre": False, + "radices": [[4, 4, 4], [3]], + "split": 1, + }, + "y": { + "C": 4, + "W": 2, + "emu": 1.0, + "hold": True, + "lf": True, + "mode": "col", + "pre": False, + "radices": [[4, 4, 4], [3]], + "split": 2, + }, + "z": { + "C": 4, + "W": 2, + "emu": 1.0, + "hold": True, + "lf": True, + "mode": "col", + "pre": False, + "radices": [[4, 4, 4], [3]], + "split": 2, + }, + }, + "256x256x256": { + "pdl": 2, + "t_us": 508.93, + "x": { + "C": 2, + "W": 1, + "ef": True, + "emu": [0.0, 0.15], + "hold": True, + "kind": "pack2", + "lf": True, + "mode": "col", + "pre": True, + "radices": [[16, 16]], + "split": 1, + }, + "y": { + "C": 8, + "W": 4, + "emu": 1.0, + "hold": False, + "lf": True, + "mode": "col", + "pre": False, + "radices": [[16, 16]], + "rev": True, + "split": 4, + "xs": True, + }, + "z": { + "C": 8, + "W": 4, + "ef": True, + "emu": 0.75, + "hold": False, + "lf": True, + "mode": "col", + "pre": False, + "radices": [[16, 16]], + "split": 4, + "xs": True, + }, + }, + "32x32x32": { + "pdl": 2, + "t_us": 5.94, + "x": { + "C": 2, + "W": 4, + "emu": [0.0, 0.0, 0.6], + "half": True, + "hold": False, + "mode": "row", + "pre": False, + "radices": [[4, 2, 2]], + "split": 1, + }, + "y": { + "C": 4, + "W": 2, + "emu": 0.3, + "half": True, + "hold": False, + "lf": False, + "mode": "col", + "pre": True, + "radices": [[4, 2, 4]], + "split": 2, + }, + "z": { + "C": 2, + "W": 4, + "half": True, + "hold": True, + "mode": "col", + "pre": True, + "radices": [[4, 2, 4]], + "split": 1, + }, + }, + "64x64x64": { + "pdl": 2, + "t_us": 12.54, + "x": { + "C": 2, + "W": 8, + "hold": True, + "kind": "pack2", + "mode": "col", + "pre": False, + "radices": [[4, 4, 4]], + "smg": True, + "split": 1, + }, + "y": { + "C": 2, + "W": 4, + "emu": [0.3, 0.3, 0.3], + "hold": True, + "lf": False, + "mode": "col", + "pre": True, + "radices": [[4, 4, 4]], + "split": 1, + }, + "z": { + "C": 2, + "W": 1, + "emu": [1.0, 0.0, 1.0], + "hold": False, + "lf": True, + "mode": "col", + "pre": False, + "radices": [[4, 4, 4]], + "split": 1, + }, + }, + "80x80x112": { + "pdl": 0, + "t_us": 28.13, + "x": { + "C": 2, + "W": 10, + "emu": [0.0, 0.0, 0.6], + "emu_split": True, + "hold": False, + "kind": "pack2", + "mode": "col", + "pre": False, + "radices": [[8, 2], [7]], + "smg": True, + "split": 1, + "stg": True, + }, + "y": { + "C": 4, + "W": 1, + "emu": [0.6, 0.6, 0.6, 0.3], + "hold": True, + "lf": True, + "mode": "col", + "pre": False, + "radices": [[4, 2, 2], [5]], + "split": 1, + }, + "z": { + "C": 4, + "W": 1, + "emu": [0.75, 0.6, 0.75, 0.75], + "hold": True, + "lf": False, + "mode": "col", + "pre": False, + "radices": [[4, 2, 2], [5]], + "split": 1, + }, + }, + "96x96x144": { + "pdl": 0, + "t_us": 44.85, + "x": { + "C": 2, + "W": 14, + "emu": 0.3, + "hold": True, + "kind": "pack2", + "lf": False, + "mode": "row", + "pre": True, + "radices": [[4, 2, 2], [9]], + "smg": True, + "split": 1, + "stg": True, + }, + "y": { + "C": 4, + "W": 1, + "emu": 1.0, + "hold": True, + "lf": True, + "mode": "col", + "pre": False, + "radices": [[4, 2, 4], [3]], + "split": 1, + }, + "z": { + "C": 4, + "W": 1, + "emu": 0.6, + "hold": True, + "lf": True, + "mode": "col", + "pre": False, + "radices": [[4, 2, 4], [3]], + "split": 1, + }, + }, + "96x96x96": { + "pdl": 3, + "t_us": 30.67, + "x": { + "C": 2, + "W": 14, + "emu": 0.3, + "hold": True, + "kind": "pack2", + "lf": True, + "mode": "row", + "pre": False, + "radices": [[8, 2, 2], [3]], + "smg": True, + "split": 1, + "stg": True, + }, + "y": { + "C": 4, + "W": 2, + "emu": 0.6, + "hold": False, + "lf": False, + "mode": "col", + "pre": False, + "radices": [[8, 2, 2], [3]], + "split": 2, + }, + "z": { + "C": 4, + "W": 1, + "emu": 0.6, + "hold": True, + "lf": True, + "mode": "col", + "pre": False, + "radices": [[4, 2, 4], [3]], + "split": 1, + }, + }, +} diff --git a/build/torch-cuda/metadata.json b/build/torch-cuda/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..af9095304c9d2e2477e8980ce353231b845c697d --- /dev/null +++ b/build/torch-cuda/metadata.json @@ -0,0 +1,166 @@ +{ + "name": "fft3d", + "id": "_fft3d_cuda_1ec826fa776c", + "version": 1, + "license": "MIT", + "python-depends": [ + "nvidia-cutlass-dsl" + ], + "kernel-depends": [], + "backend": { + "type": "cuda", + "archs": [ + "10.0", + "12.0" + ] + }, + "digest": { + "algorithm": "sha256", + "files": { + "LICENSE.txt": "AIyEysQ0a1okjH0MEq1kTbdKj938zWEHbQJBfSkK23s=", + "__init__.py": "cyv3oPsvItJIjU9bnO31sjePmP2Yd5a3+Mjj01oLGUQ=", + "_b200/__init__.py": "NyP640/tMaZvJN5nlKDXoB+zzhXSm2tpc0UPN3W+uUI=", + "_b200/c2r_cuda_z_inverse_cuda_kernel.py": "PjB5WazUXdSzIjmErZAKnn4KlmVCSxTg35LiFbPUHZY=", + "_b200/c2r_cuda_z_inverse_fft_axis.cu": "iwEGEo/0rN4qsT5cwH1MDv2WEiww8cGm6AyskUlnJ6Q=", + "_b200/c2r_dense_half_inverse_wide_redirect.cu": "hbkD9Vf/GhGIBFR8kBvRSEVaIFAFirJBkqwQLZsK3dU=", + "_b200/c2r_f64_cuda_z_inverse_cuda_kernel.py": "TR8wh6H/487/0jLzZaLSKfnwUOViN11O3V9dlyYZKKc=", + "_b200/c2r_f64_cuda_z_inverse_fft_axis.cu": "A3jerp74Qllbc729k1e2tNpW0mP+g2zQJfKzBIwriTk=", + "_b200/c2r_f64_cuda_z_inverse_lto.py": "cViADUIaPPvBY542VeDv3DSV47X/A6y35I5vGyfzESc=", + "_b200/c2r_f64_cuda_z_inverse_native_dispatch.py": "j839Ez5AX6OohPVKsefY4NFQdEMLH5Fce4eKkFq0IZ0=", + "_b200/c2r_f64_cuda_z_inverse_wide_redirect.cu": "5tHUmgUk+l3OqTrFq1tMMf/ycJAxxGeivgL/jsY2bwU=", + "_b200/c2r_f64_dense_half_inverse_lto.py": "Z1EvMX+NwoK0R4hJlhGIzVV0Kn8TtN0NB6UzHR6aeiw=", + "_b200/c2r_f64_dense_half_inverse_wide_redirect.cu": "hGd+/s6c4FXFqLwqHmQcE+bNPPq+0Sd///OUL0zc2gs=", + "_b200/c2r_f64_half_vector_inverse_wide_redirect.cu": "5sX5LS8iqn++joD9AWyXodpYsWzWQPLo04diXI6MiW0=", + "_b200/c2r_f64_narrow_index_inverse_scalar_output_reconstruct.cu": "JN0WZbI2IqhB7T9HzqdZ++jUNKlxGqk4MJCAGPeZklU=", + "_b200/c2r_f64_paired_x_inverse_cuda_kernel.py": "vIMZWXJGO3mTY/IwozNmfRQEZcspeWgsLOXSU7Pm6Is=", + "_b200/c2r_f64_paired_x_inverse_fft_axis.cu": "ihwoeg9AV2rvuQCYlBMcTYoGDScZVIzSKPuvV/eNK1E=", + "_b200/c2r_f64_paired_x_inverse_wide_redirect.cu": "W82vPTTpc5iZykTmPizPKut/UUwnsfc6HCBdx+bhwaE=", + "_b200/c2r_f64_paired_y_direct_inverse_lto.py": "FwU7tKphogBA/QMk1dCHmPCetXaL/m9jRGC6prGuZRo=", + "_b200/c2r_f64_paired_y_direct_inverse_wide_redirect.cu": "C6qCRADcX2MOxDJTgewlH4/tAGhn3o31fg38ra3PyHs=", + "_b200/c2r_f64_partial_y_inverse_wide_redirect.cu": "MsQFI2fMGSYOK6S32ylFfvnrGPfFj5FrlIk3KJq6BRM=", + "_b200/c2r_f64_scalar_input_inverse_scalar_io.cu": "TzhzYvtnKNndRgBGBiqXSLE7VXTYuhGLCkN9H851/sk=", + "_b200/c2r_paired_y_direct_inverse_wide_redirect.cu": "HPMACQxYzkAQOhrIuqYXHXEX+v17s9hLYY3wc0GPTwM=", + "_b200/c2r_paired_y_inverse_wide_redirect.cu": "QuOzmgP87B9nIVzK8dfHMtzJWbicrUOKGouO+x1f8eg=", + "_b200/c2r_pitched_half_inverse_wide_redirect.cu": "1+G0zliatnWTK8xMfTqbGxr1W3mitunGEEZOL515j9k=", + "_b200/cuda_kernel.py": "pACmB/oK7X23b8R0B8oMLYqYVEccJYOwKEj224h+P3Y=", + "_b200/cuda_twiddle_forward_lto.py": "p5mtO5nL+Ygim3w+6OeEUEQEo3QHPL3iReSFDNxsdXA=", + "_b200/cuda_twiddle_forward_native_dispatch.c": "GdxR/QB2u2KqzDKEI5UY6i5jOEOrGlTwyMmyLE/Mk+c=", + "_b200/cuda_twiddle_forward_native_dispatch.py": "I1agbrlbyIBpB0CQdjGdoq4TBeQdXQ1zTuCCRPsu8rg=", + "_b200/cuda_z_inverse_cufft_support.py": "WxfLUGrbYhp1lwZCD/7iaLaekP7hmDxQv2/YfvHH2k0=", + "_b200/cuda_z_inverse_lto.py": "QczP4CY+5DzX+jZUbYmdVHvLBieHL70k+4gnAP7CNFs=", + "_b200/cuda_z_inverse_native_dispatch.c": "QJv1Yai0EDc8KuYjQpjkfKosXYsdiyOhFQDWUwYuLEQ=", + "_b200/cuda_z_inverse_native_dispatch.py": "gvHMxfC3iGDhWzOzP2FheWxFyFf9v9pY/gBtfeMKCmU=", + "_b200/cuda_z_inverse_schedule.py": "c4KtEQzNERCe0BgdeA/GI4le/KXHj4nMcYZp0julpIU=", + "_b200/dense_half_inverse_abi.py": "9hDyt0Um3/4c9Ks26dvWrEm0ngZreSCE8iIVXZlUn7Q=", + "_b200/dense_half_inverse_kernels.py": "uDHQRMXq4/6+hW6oWISNHxDLqG2EadlioyWqPwreIyE=", + "_b200/dense_half_inverse_lto.py": "XjrTUimj8hcARP/3dBy0AP/EMC3jwG/V508R9TLWveo=", + "_b200/dense_half_inverse_native_dispatch.c": "3KMzUfNL4vSWT1ucWtdTExtkXeYQ9BuRCB1XdoAlbdc=", + "_b200/dense_half_inverse_native_dispatch.py": "Jyvyd211BjREscczb5ipaXW0lBxR8C1Ice4PfTYzCLA=", + "_b200/dense_half_inverse_schedule.py": "qVuIX5W35USjUX3v5kpLmjdQZtWdFqclgznPHYBB10s=", + "_b200/direct_dispatch.c": "cSxdVlRuluTgP61xxHQwIg4us+IaZW6gsxwLKtasnXI=", + "_b200/direct_dispatch.py": "Sn2wMR0DVersp94ABoCVx7LKWIzAg/vRde1EXwhXZG0=", + "_b200/fft_axis.cu": "O5l+aBQjqBwG+IoMuE16e4GVtcS7DJfIXd9qgxc31ac=", + "_b200/forward_cuda_f32.py": "Iam+BjI7KBIhaTy6dwNdnzHZSCynspIMq3YrEUxlTEo=", + "_b200/forward_half.py": "zOUY6uZu9hDmqhIupY3S7WR1XOyTOEAJzRVvrAsxrCQ=", + "_b200/forward_native.py": "oaGPHqlaGxs0lhGWNLBsjQPMlJsmEXKv75Agzd7xuFM=", + "_b200/forward_paired_y.py": "LApxXkstJe2/imuvd6/plmAYtlzx+won5e68QnmWsqg=", + "_b200/forward_scalar.py": "ZqjkZ9RtCrbv04r5WHp4jqZcdyBBJ1owuRB8FjKgpKs=", + "_b200/forward_small.py": "Tv9NURYQtfUFRXnVctgItcoMmOV2tv6gTATP2/LdISM=", + "_b200/grouped_dif_inverse_abi.py": "zIRVdRJLuBS6kSDqKVvlHBtF/9seccLcpM/FPJwCr2I=", + "_b200/grouped_dif_inverse_kernels.py": "mdEXtoCU/goUy7iGiYu6CNT9qV2Svkxt6EgCVX7nut8=", + "_b200/grouped_dif_inverse_native_dispatch.c": "GagjZWJakPAy0hTtUQFC272GmqsvxNqFmOfDad4KBHg=", + "_b200/grouped_dif_inverse_native_dispatch.py": "zrzUFe5XThSGIliqAKobwgkI5VsidS37bm/6YYtXRvA=", + "_b200/grouped_dif_inverse_schedule.py": "o3XGdIW2pEkw8kEKJyElTgbVrpWc9ieZLCBchewkqgc=", + "_b200/half_forward_cufft_z_lto.py": "0OyQGdDqh885v9YCfN2hsi2Eq0rtGAqRjTRlHU7DMPE=", + "_b200/half_forward_cufft_z_plan_safety.py": "mjbIXvUnG2cWYAg+hZjFLy8y13k/gz8sAXaAza6qb/E=", + "_b200/half_vector_inverse_lto.py": "QFFP/F5cFhCQ+pCUL+JCp6aLETZotEP/Fp1uvVEQ5m4=", + "_b200/hermitian_projection_cufft_support.py": "qq6po9ykbrtIV3dDpvqypGPi8c64CcXYL77VXq05Aeg=", + "_b200/hermitian_projection_kernels.py": "xBvZKiL7xxm+QDnvsyKbAZo8u1QyMM/Vxd7wAWwLJQM=", + "_b200/hermitian_projection_lto.py": "QkSuQ6cjMhnvT3eVJZZv9AhtI5mZxgUA7EovZSqxRjw=", + "_b200/inverse_cuda_f32.py": "Y2tiP66WxZPORXpfWdnidhBNW4nE1LR2nYYx8gbDEPo=", + "_b200/inverse_cuda_f64.py": "Ya5U0IocJ4UwsqmTB4a+Z0rPuG1KNjicQAv0i7WeG3w=", + "_b200/inverse_dense_f32.py": "NZ+76/r/tTh2bF4PyEd0l356WVfLuejffZE3n0bVwhc=", + "_b200/inverse_dense_f64.py": "Eb6F9lHxKwv1mZ8tpZxx0GT7tc2/fERnQ+bj1NFvnYs=", + "_b200/inverse_narrow.py": "X0c8pYzG7Xl50x6M9DqtALyYZ+opN8GmiKKagQysFho=", + "_b200/inverse_paired_x_f32.py": "cf9FaOe2JJEduVeyu9tAzx7sPBaoafBAVnwdK402IiY=", + "_b200/inverse_paired_y.py": "QKmkN3aFdw2Rw3bJKcaYKRuMc5qxxDezn3OFugValVk=", + "_b200/inverse_paired_y_f32.py": "6L720c9C4sY0NpOtWtp7AQbj1LBpf9/4znAy5aJCASc=", + "_b200/inverse_paired_y_f64.py": "UKhOKwMDT2JFQcpko2UFBoEXT7iqd4HcXSms6yoXJGY=", + "_b200/inverse_partial_f32.py": "6FVPDzmWvr9GtZ9QPhLMEXK0rYNTVKGCi0QTKVzsXQU=", + "_b200/inverse_pitched.py": "/uUteJKXsDubV9c7bTzinen6aTmmaE35Z/t/b5rxU6I=", + "_b200/inverse_scalar.py": "WuZBY6JkU8DwHOOZmJM/5s48iha8JkYLX1mJm0Del+Y=", + "_b200/inverse_small_f32.py": "TJ2KIZt6MPhmBIoQsVry688RU56/in4s9NGHsbsMJ+E=", + "_b200/inverse_small_f64.py": "H9BzxYu4eDDdfxNcmXSJOtr3u2E1zsqMvRjJI6X2LEE=", + "_b200/inverse_vector.py": "aVblDFIMXi8GtVv8/2L2Dd6PsCmg8rTyObr5nFPUSoQ=", + "_b200/narrow_index_inverse_cufft_support.py": "gUhXx+1hCOHJa0vcETIM50u5HFXomJbDlm+drSI6OVE=", + "_b200/narrow_index_inverse_lto.py": "5V1C71Dd3ko3pBSbW25pisjRhqrXNqIJVFO74yKeuYI=", + "_b200/pair_load.cu": "IMHoK72vmsG2Jc36HUc8KIwWejnC8toW8G0wAPkJPFA=", + "_b200/paired_vector_store_forward_abi.py": "0YrKitrultbYyAhMICDInCREEgfIPx0zHdWY7QU2K/g=", + "_b200/paired_vector_store_forward_backend.py": "KOkHlphj5S+Ha+W6bHyTgLRDVKfp5LZBAmq4k3KF6+s=", + "_b200/paired_vector_store_forward_kernels.py": "jxqHWBHrL8zknW1+CSR9oweB2y8y08fIIJ66a+iwrVU=", + "_b200/paired_vector_store_forward_lto.py": "dfuSAkfO5W49iAGZR0guyjXGpngnKYLjOC7sTbaPGGI=", + "_b200/paired_vector_store_forward_native_dispatch.c": "An+VLfeh36eRhHH1ZUaL19w5kbqfD0YMdQWAXWVkGQU=", + "_b200/paired_vector_store_forward_native_dispatch.py": "ou9x6q/+ae1n1sHPuCS6kFNCYpweAl5iROeN4LBCy+E=", + "_b200/paired_x_inverse_lto.py": "P/8gquCooN5eqtdPGflvmrzfNhPtaRVr/EKyYaYWHOo=", + "_b200/paired_x_inverse_native_dispatch.c": "MqYLx2ElzjV3VugV8/pR2Jy4mMGdaQv+0dgB2NkB5LM=", + "_b200/paired_x_inverse_native_dispatch.py": "AI/fiv4IvRQ6tGxRSWHOkPYV8PXsQZX/3Fq8LSsQuj8=", + "_b200/paired_x_inverse_plan_safety.py": "v89YGJzYSUnhPZ102qfYcKiVwBuij9jgVvDRDxci+cs=", + "_b200/paired_y_direct_inverse_lto.py": "ZKhPSU5t8GqR89LUtaTHFJBrYqCNe1OVcZNAYZD1M58=", + "_b200/paired_y_forward_lto.py": "us/RQWeK7m3bC56g1gKz92Ztdmsynk5zsTHAEymmMmc=", + "_b200/paired_y_inverse_lto.py": "npM2AKTam/it8za3C9LG97kgS93Shu1IHCY2Uo1gSwc=", + "_b200/partial_y_inverse_lto.py": "tLSTXg19hrBg+QjI/yUgk3dVMIt+cgvniW3w3KwLQzM=", + "_b200/pitched_half_inverse_lto.py": "CQR07VHrdTfbOQvDKmZ+YE3SQfJX8a5pZiD+fNy4MMQ=", + "_b200/prepared_call.cpp": "2522MPXkipZwlsqTYhMs2GzUyAuqudFrToRCB3Kbs8M=", + "_b200/prepared_call.py": "xSEya0XmcMBlkmx+nLnvIdJrR8bWb/nCC2GvjzXGcD0=", + "_b200/projection.cu": "0DnSox0lUFN+uG+5+YnQPHlfsdeDZ6WHGiIbNGsSGLg=", + "_b200/r2c_half_forward_cufft_z_scalar_io.cu": "y3kUflxPlOkrMvSOe9gRQUm9ANSCPwFfZ46kpvjw+kk=", + "_b200/r2c_paired_y_forward_scalar_io.cu": "/3UPYGt8xaz4ko6gXiwz3z98HevXMi+gs6TZhGnNxyw=", + "_b200/reconstruct.cu": "8Jiv0szf6vlZiV6aIIT5wfmjBPb+Tus/F+rNX7JPcHE=", + "_b200/scalar_input_inverse_backend.py": "jbnVll8Ekx4z4HJtKsCeDQIiJG29zqQ3SKhuK8Y1ZYQ=", + "_b200/scalar_input_inverse_cufft_support.py": "WXruVUYvN88GEHlE7gQOSlv2ksCL/DADwOFkaHTy6Sk=", + "_b200/scalar_input_inverse_lto.py": "ecu0bAMsVFAolkOMyUy/xncXfcgYlGzlvhW+Z5CWgiw=", + "_b200/scalar_io.cu": "HPdWEM/5g2QSc8hRQYFJ9xHRJkITsOZY4lKgKGpSyXo=", + "_b200/scalar_output_backend.py": "01u/6ZRR9NdsdChAqipE/GWvPjJDpYeq1aQcyv7bRhk=", + "_b200/scalar_output_lto.py": "F1j27fJ8Jl5w1gwPtQWyQ7an0wLvxATBS6CupywwF7g=", + "_b200/wide_native_kernels.py": "mupjMNvM5dQDkeG7qOu/UCfJxQ9VvyFePW/VKp00TCM=", + "_b200/wide_native_schedule.py": "8M0FA28qeZhEVTs6e4fi308kbM/qeDGfQHyAEWjv9mg=", + "_b200/wide_redirect.cu": "Eq43OTnEoHJPclHn2A2I4kEfENwYrv5T48CI6YKlPaQ=", + "_b200/y_residue_inverse_kernels.py": "T4Hy3Tq2ROdffxtqTq1eMWS+DGgo4+5QyXhaP8TI69Q=", + "_c2r_f32/__init__.py": "XBy0pSUwgUTmssbhC7XEt/ELeTb4ayilaHy7B/V0pHw=", + "_c2r_f32/fftgen.py": "vlccjhLwsmKsoyZ+C8DPeR7fIe7LPMxEVsq1cDF5faE=", + "_c2r_f32/generic.py": "F902qd+wV64kuLP9J/ko0jIWWH/wApP3vh28C9iuiy8=", + "_c2r_f32/inplace.py": "J+tnMJWGaiz3Gug+wbJzgqgskLgHxnbqHWNb9h1GSCU=", + "_c2r_f32/kernel.py": "bP0gzcYzOJ3nFKSA0EdILiAAOfxbqRjre4kphnPoD8o=", + "_c2r_f32/tables.py": "RvVnytUIHgceogs37Z8lcjI5mxWclCA8D7kqC2YKm9w=", + "_c2r_f32/twopass.py": "zI/xziZlOyfjNByjxB9Lag16rgClQVEn9kwlcoKQXT0=", + "_c2r_f64/__init__.py": "6NiqCxU8/XZoiu7IY/nThfL5G3EovojIu+BSdpPGqEk=", + "_c2r_f64/fftgen.py": "1Fn09b4uYmDSwItnhAycTXZVDWcmvXmf28BdjtMjsEQ=", + "_c2r_f64/kernel.py": "x6MD24RFx/OPF9OQgWckUwKqPEGRhY3PRKmbq5mvrcw=", + "_c2r_f64/softfp.py": "KsaELtKI5TNfIAwqTZJc9ROp7PeKpV1sUrCcXH3jVwA=", + "_c2r_f64/twiddles.py": "WXaDgblQUWoX/1Qoofd5mD5aUpG1EkjYsdKD3YAUPW4=", + "_ops.py": "e5UmLzFiGFFZCkWwlfncv0R0JrIA5WXumPX6cWXi29w=", + "_r2c_f32/__init__.py": "b+FuhAkJolHH1YTeeRAb4NMHe+7YQId1VYOQ+W+vqsI=", + "_r2c_f32/fft3k.py": "ZoLp8dkbyHLI57MrUAhd06t3+lDM+kIavTOk7YPaVPc=", + "_r2c_f32/fft_core.py": "WThy1Zi8S9R5zSQgluOWnrhge+zXreeLR5OjysAgXNc=", + "_r2c_f32/fft_generic.py": "FCzLSrddYV38yRsIL/rPWwIgcSZDdNElpLWo30C6oCA=", + "_r2c_f32/fft_plane2.py": "jkEtDEO5n2htSNz6tTcExrwoWUuWcTAAZu+uvhJXq+M=", + "_r2c_f32/fft_xy.py": "1Y0pV8gRN+jk6Iriz6Nzpjv0SiCF7AOUn2UR9VpEL80=", + "_r2c_f32/fft_ys1w.py": "4Vske6f2ATSSTb//EkpoHS5E3s0m9xGWOLuH3KWJSC8=", + "_r2c_f32/fft_ysplit.py": "ukVz02AMlODEqFb/0qvSwAs5rA4CMFlqpI1YSRY0yqQ=", + "_r2c_f32/kernel.py": "K47e1I7yl+Uu75FwCg9qZ2aZCvHtBI08Olz9DGxmDb0=", + "_r2c_f32/plan.py": "kp8pVa1thunGU5fK5/OtqNF+7NqWYNPac/kwpfTQKTA=", + "_r2c_f32/ptx.py": "3mKM+qkQfpV6NCAXz4V4FLf6Z35LtiONgb5oSGF09fw=", + "_r2c_f32/stockham.py": "5lS6YIDy/WKT4gTUgsQsrJHBEaAWAyjELVhD6uHJysM=", + "_r2c_f64/__init__.py": "rHmKPjGcG8wbZnWqgo9OQD4lcZYRVFCcsgC3c5sFq3s=", + "_r2c_f64/dsl_ops.py": "lyXWN7Qokc+a6avqj8mtOHYzn3WEHXqf+7VAwxGD1NY=", + "_r2c_f64/fft_codelets.py": "Np96zEz++x7McIuE7+aUz6HGkIk81aNt+Zs2TyrvXKA=", + "_r2c_f64/fft_direct.py": "bImy5bLyOOWrUTeAmTJaeUsznyQzlljKDqMuKmTvLQQ=", + "_r2c_f64/fft_passes.py": "jvB79PCpXS9oHh7unWN8JNrbrHijzm28y0QVzBTNhss=", + "_r2c_f64/fft_plan.py": "OqPrkiahnt/XreodPMi90N7+kHinYgYbLuWAdVd8eqI=", + "_r2c_f64/fft_stages.py": "5ylt/t+bA+9ILsAG5P9CCJkwdJmfvvx/oB29w1LoO9A=", + "_r2c_f64/kernel.py": "DfHSUxG56B/HQ7ZsYjz7aEBbLCCDor6q+TiXlGXzZI8=", + "_r2c_f64/planner.py": "VgtohOIirLvHoK/AAciicLFmPD5nJSOCYxOomeYQZEY=", + "_r2c_f64/tuned.py": "YtEsBZO9C63haDWFVnMBtIpTXwSKuvxAF6ffpz9wrd0=" + } + } +}