Uploaded using `kernel-builder`.

Browse files

Files changed (8) hide show

README.md +21 -1
build/torch-cuda/__init__.py +18 -6
build/torch-cuda/_c2r_f32/kernel.py +19 -1
build/torch-cuda/_c2r_f64/kernel.py +24 -19
build/torch-cuda/_ops.py +1 -1
build/torch-cuda/_r2c_f32/kernel.py +10 -1
build/torch-cuda/_r2c_f64/kernel.py +82 -56
build/torch-cuda/metadata.json +8 -8

README.md CHANGED Viewed

@@ -65,6 +65,10 @@ spectrum = launch()  # ...then transform its current contents
   fp64 r2c has no two-phase form; its `launch()` calls `r2c(x)`.
 - Calls on different CUDA streams may run concurrently, including calls of the
   same shape.
 - The first call for a new shape compiles its kernels (CuTe DSL JIT), typically
   in 1 to 15 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;
@@ -168,6 +172,17 @@ Every sampled grid is faster than cuFFT's 3D plan and than `torch.fft`. Large
 prime extents gain most: 251x251x254 takes 0.25 ms in fp32 r2c (cuFFT 0.59 ms)
 and 2.6 ms in fp64 c2r (cuFFT 4.9 ms).
 ## Changes
 Updated on 2026-10-24 with fixes found by sweeping every planner signature over
@@ -213,6 +228,11 @@ two-stream checks):
 - fp64 r2c on data with a large mean (randn + 1e3) missed atol 1e-6 on 21 of 100
   sample grids of 1M points and more, where cuFFT passed; now at most 1.07x
   cuFFT's error.
 The eight grids keep their plans and their outputs are bitwise unchanged.
@@ -224,7 +244,7 @@ built from the requests in [NVlabs/kda#7](https://github.com/NVlabs/kda/pull/7),
 and accepted by the KDA judge. Published from commit `e539498` of
 `humanfia/kda-kernels-internal`, directory `rtx5090/fft3d/` (now
 `sm120/fft3d/`), which also holds its tests and the fixes made after the
-judge's checks; updated from commit `f8b3c93` (see [Changes](#changes)).
 MIT-licensed.
 ## Available functions

   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.
+- A call costs about 13 to 17 us of host time (TVM-FFI launches on the current
+  stream). Timed per call as the requests' benchmark does (flashinfer-bench:
+  synchronize, event, call, event), every kernel is faster than `torch.fft` on
+  every grid tried, the smallest included (see Performance).
 - The first call for a new shape compiles its kernels (CuTe DSL JIT), typically
   in 1 to 15 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;
 prime extents gain most: 251x251x254 takes 0.25 ms in fp32 r2c (cuFFT 0.59 ms)
 and 2.6 ms in fp64 c2r (cuFFT 4.9 ms).
+Timed per call the way the requests' benchmark (flashinfer-bench 0.1.2) does it,
+on fresh copies of the input with a cleared L2 and the host's launch work
+included, against the requests' baseline (`torch.fft`):
+| | eight grids (geomean) | 8x8x16 to 48x40x36 (slowest) |
+|---|---:|---:|
+| r2c fp32 | 1.91x | 1.54x |
+| c2r fp32 | 1.72x | 1.11x |
+| r2c fp64 | 2.20x | 1.29x |
+| c2r fp64 | 2.23x | 1.12x |
 ## Changes
 Updated on 2026-10-24 with fixes found by sweeping every planner signature over
 - fp64 r2c on data with a large mean (randn + 1e3) missed atol 1e-6 on 21 of 100
   sample grids of 1M points and more, where cuFFT passed; now at most 1.07x
   cuFFT's error.
+- Per-call host time: the fp64 kernels spent about 40 to 50 us of host time per
+  call (0.48x to 0.69x of `torch.fft` on small grids when timed per call, as the
+  requests' benchmark does). All kernels now launch through TVM-FFI, at about
+  13 to 17 us; outputs are bitwise unchanged. fp64 c2r now transforms an input
+  that is not 16-byte aligned from an aligned copy instead of raising.
 The eight grids keep their plans and their outputs are bitwise unchanged.
 and accepted by the KDA judge. Published from commit `e539498` of
 `humanfia/kda-kernels-internal`, directory `rtx5090/fft3d/` (now
 `sm120/fft3d/`), which also holds its tests and the fixes made after the
+judge's checks; updated from commit `b0811f8` (see [Changes](#changes)).
 MIT-licensed.
 ## Available functions

build/torch-cuda/__init__.py CHANGED Viewed

@@ -32,16 +32,28 @@ _PACKAGES = {
 }
 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}")
-    return importlib.import_module(f".{package}", __name__)
 def _on(tensor):
-    """The input's device, current while a call runs (the kernels launch on the current device)."""
-    return torch.cuda.device(tensor.device) if tensor.is_cuda else contextlib.nullcontext()
 def r2c(x):

 }
+_MODULES = {}
 def _kernel(direction, tensor):
+    module = _MODULES.get((direction, tensor.dtype))
+    if module is None:
+        package = _PACKAGES.get((direction, tensor.dtype))
+        if package is None:
+            raise TypeError(f"{direction} takes float32 or float64, not {tensor.dtype}")
+        module = _MODULES[direction, tensor.dtype] = importlib.import_module(f".{package}", __name__)
+    return module
+_NULL = contextlib.nullcontext()
 def _on(tensor):
+    """The input's device, current while a call runs (the kernels launch on the current device). A no-op
+    when it already is, which is the common case; switching devices costs a few microseconds per call."""
+    if not tensor.is_cuda or tensor.device.index == torch.cuda.current_device():
+        return _NULL
+    return torch.cuda.device(tensor.device)
 def r2c(x):

build/torch-cuda/_c2r_f32/kernel.py CHANGED Viewed

@@ -316,4 +316,22 @@ def prepare(input):
 def run(input):
     """Return the unnormalized 3D C2R inverse FFT of ``input`` as a new [nz, ny, nx] tensor."""
-    return prepare(input)()

 def run(input):
     """Return the unnormalized 3D C2R inverse FFT of ``input`` as a new [nz, ny, nx] tensor."""
+    dev = input.device
+    if (
+        dev.type != "cuda"
+        or dev.index != torch.cuda.current_device()
+        or not input.is_contiguous()
+        or input.data_ptr() % 16
+    ):
+        return prepare(input)()
+    # the common case without building a launch closure: per-call overhead counts on small grids
+    _check(input)
+    nz, ny, nf, _ = input.shape
+    nx = 2 * (nf - 1)
+    plan = _get_plan(nz, ny, nx, dev)
+    out = torch.empty((nz, ny, nx), dtype=torch.float32, device=dev)
+    if hasattr(plan, "nyq_numel"):
+        plan(input.view(-1), out.view(-1), torch.empty(plan.nyq_numel, dtype=torch.float32, device=dev))
+    else:
+        plan(input.view(-1), out.view(-1))
+    return out

build/torch-cuda/_c2r_f64/kernel.py CHANGED Viewed

@@ -42,7 +42,6 @@ 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 bigprime as BP
@@ -2841,8 +2840,9 @@ def make_tables_2d(nz, ny, M, yzd):
 _PLANS = {}
-def _stream():
-    return cuda.CUstream(torch.cuda.current_stream().cuda_stream)
 class _Plan:
@@ -2851,8 +2851,6 @@ class _Plan:
         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)
         axes = {"z": (nz, 1, ny * M), "y": (ny, nz, M), "x": (M, nz * ny, 1)}  # (n, A, inner)
         planned = {ax: _has_plan(n) for ax, (n, _, _) in axes.items()}
         self.direct = not all(planned[ax] or BP.supported(n) for ax, (n, _, _) in axes.items())
@@ -2877,38 +2875,41 @@ class _Plan:
             tables = make_tables_2d(nz, ny, M, YZ2D(ny, yz["A"], yz["rc"], yz.get("N1", 1)))
         else:
             tables = torch.cat([make_tables(nz, ny, M)] + [cfg.table() for cfg in big.values()])
-        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)
@@ -2931,6 +2932,8 @@ def run(input):
     _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
@@ -2941,6 +2944,8 @@ def prepare(input):
     _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)

 from cutlass import cute
 from cutlass._mlir import ir
 from cutlass._mlir.dialects import llvm
 from cutlass.cutlass_dsl import Float64, Int32, T, dsl_user_op
 from . import bigprime as BP
 _PLANS = {}
+def _fake(n):
+    """Compile-time stand-in for a flat float64 buffer of n elements, 16-byte aligned (the shape stays static)."""
+    return cute.runtime.make_fake_compact_tensor(cutlass.Float64, (n,), assumed_align=16)
 class _Plan:
         self.shape = (nz, ny, nf)
         self.M = M
         self.device = device
         axes = {"z": (nz, 1, ny * M), "y": (ny, nz, M), "x": (M, nz * ny, 1)}  # (n, A, inner)
         planned = {ax: _has_plan(n) for ax, (n, _, _) in axes.items()}
         self.direct = not all(planned[ax] or BP.supported(n) for ax, (n, _, _) in axes.items())
             tables = make_tables_2d(nz, ny, M, YZ2D(ny, yz["A"], yz["rc"], yz.get("N1", 1)))
         else:
             tables = torch.cat([make_tables(nz, ny, M)] + [cfg.table() for cfg in big.values()])
+        self.tables = tables.to(device).contiguous()
+        nx_, no_ = nz * ny * nf * 2, nz * ny * M * 2
+        args = [_fake(nx_), _fake(no_), _fake(self.tables.numel())]
         if self.direct:
+            args.append(_fake(no_))  # workspace
+        # TVM-FFI: called with torch tensors directly, on the caller's current stream, so a call costs a
+        # few microseconds of host time (the shapes and the 16-byte alignment are checked at the call)
+        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+        self.fn = cute.compile(launch, *args, stream, options="--enable-tvm-ffi")
     def bind(self, xin, out):
+        xf, of, fn, tw = xin.view(-1), out.view(-1), self.fn, self.tables
         if self.direct:
+            ws = torch.empty_like(of)
             def go():
+                fn(xf, of, tw, ws)
             go.ws = ws
             return go
         def go():
+            fn(xf, of, tw)
         return go
     def run_into(self, xin, out):
+        if self.direct:
+            self.fn(xin.view(-1), out.view(-1), self.tables, torch.empty_like(out).view(-1))
+        else:
+            self.fn(xin.view(-1), out.view(-1), self.tables)
 def _get_plan(nz, ny, nf, device):
+    key = (nz, ny, nf, device)
     p = _PLANS.get(key)
     if p is None:
         p = _Plan(nz, ny, nf, device)
     _check(x)
     nz, ny, nf, _ = x.shape
     plan = _get_plan(nz, ny, nf, x.device)
+    if x.data_ptr() % 16:
+        x = x.clone()  # the kernels read 16-byte vectors; an aligned copy leaves the input as it is
     out = torch.empty((nz, ny, 2 * (nf - 1)), dtype=torch.float64, device=x.device)
     plan.run_into(x, out)
     return out
     _check(input)
     if not input.is_contiguous():
         raise ValueError("prepare binds the input buffer itself, so it must be contiguous")
+    if input.data_ptr() % 16:
+        raise ValueError("prepare binds the input buffer itself, so it must be 16-byte aligned")
     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)

build/torch-cuda/_ops.py CHANGED Viewed

@@ -28,7 +28,7 @@ def get_backend() -> str:
 def _find_ops_name() -> str:
     kernel_name = "fft3d"
-    unique_id = "f2cede9"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

 def _find_ops_name() -> str:
     kernel_name = "fft3d"
+    unique_id = "9e7519b"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

build/torch-cuda/_r2c_f32/kernel.py CHANGED Viewed

@@ -408,4 +408,13 @@ def prepare(input):
 def run(input):
-    return prepare(input)()

 def run(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)
+    if pl.sync_words:
+        return prepare(input)()
+    # the common case without building a launch closure: per-call overhead counts on small grids
+    out = torch.empty((nz, ny, nx // 2 + 1, 2), device=input.device, dtype=torch.float32)
+    pl.run(input.view(-1), out.view(-1))
+    return out

build/torch-cuda/_r2c_f64/kernel.py CHANGED Viewed

@@ -27,7 +27,10 @@ the X pass packs two real rows per complex FFT); lengths it cannot split (none u
 the 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
@@ -56,6 +59,31 @@ _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)
@@ -121,6 +149,7 @@ class FFT3DPlan:
         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
@@ -192,35 +221,33 @@ class FFT3DPlan:
         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)
@@ -228,35 +255,20 @@ class FFT3DPlan:
                 # 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
@@ -285,13 +297,27 @@ def _validate(x):
     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

 the 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).  The three
+launches are compiled as one host function for TVM-FFI calls: a call hands the input, output,
+tables and counters (each once) straight to it, and it launches the X, Y and Z kernels on the
+device's current torch stream.
 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
 _current_device = torch._C._cuda_getDevice
+def _ffi(t):
+    """Compile-time stand-in for a tensor argument: its static shape and default alignment."""
+    return from_dlpack(t, enable_tvm_ffi=True)
+def _compile(fn, *args):
+    return cute.compile(fn, *args, cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi")
+def _fused_launcher(xl, yl, zl, ntabs, xmap, ymap, zmap):
+    """One host function launching the X, Y and Z passes, so that a call crosses into compiled
+    code once.  Its parameters are gin, gout, the tables t0.. and gsync; *map name each pass's
+    table / counter arguments.  Straight-line launches only, so the AST preprocessor is off."""
+    names = ["gin", "gout"] + [f"t{i}" for i in range(ntabs)] + ["gsync"]
+    src = (
+        f"def launch({', '.join(n + ': cute.Tensor' for n in names)}, stream: cuda.CUstream):\n"
+        f"    xl(gin, gout, {''.join(a + ', ' for a in xmap)}stream)\n"
+        f"    yl(gout, {''.join(a + ', ' for a in ymap)}stream)\n"
+        f"    zl(gout, {''.join(a + ', ' for a in zmap)}stream)\n"
+    )
+    ns = {"xl": xl, "yl": yl, "zl": zl, "cute": cute, "cuda": cuda}
+    exec(src, ns)
+    return cute.jit(preprocess=False)(ns["launch"])
 def _dev(a, device):
     """Host table (1-D CPU tensor) -> device copy."""
     return a.to(device)
         self.shape = (nz, ny, nx)
         nf = nx // 2 + 1
         self.nf = nf
+        self.out_shape = (nz, ny, nf, 2)
         self.device = device
         cfgs = default_cfgs(nz, ny, nx)
         xcfg, ycfg, zcfg = cfgs
         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 = None  # table + counter arguments, when they do not depend on the stream
+        self._stream_args = {}  # mode 3: launch key (stream, or (stream, capture id)) -> arguments
+    def _pass_maps(self):
+        """Names of the X, Y and Z launchers' table + counter arguments (see _fused_launcher)."""
+        t = [f"t{i}" for i in range(len(self.tabs))]
         a, b = self.nx_t, self.nx_t + self.ny_t
+        tx, ty, tz = t[:a], t[a:b], t[b:]
+        later = ty + tz
+        xmap = tx + ((later + [later[0]] * (6 - len(later)) + ["gsync"]) if self.x_prefetch else [])
+        ymap = ty + (["gsync"] if self.cfgs[1] is not None else [])
+        zmap = tz + (["gsync"] if self.cfgs[2] is not None else [])
+        return xmap, ymap, zmap
+    def compile(self, inp, out):
+        launch = _fused_launcher(self.xl, self.yl, self.zl, len(self.tabs), *self._pass_maps())
+        args = [_ffi(inp.view(-1)), _ffi(out.view(-1))] + [_ffi(x) for x in self.tabs] + [_ffi(self.sync)]
+        self.compiled = _compile(launch, *args)
+        if not self.flagged:
+            self._args = (*self.tabs, self.sync)
+    def _launch_args(self, inp):
+        """Mode 3: the arguments for launches on the current stream, with that stream's counters."""
+        handle = _raw_stream(inp.get_device())
+        key = (handle, _capture_id(handle)) if _capturing() else handle
+        args = self._stream_args.get(key)
+        if args is None:
             n = self.sync.numel()
             sync = (
                 _zeros_outside_capture(n, self.device)
                 # zeroed on this stream, ahead of the first launch that uses it
                 else torch.zeros(n, dtype=torch.int32, device=self.device)
             )
+            args = self._stream_args[key] = (*self.tabs, sync)
         return args
     def launch(self, inp, out):
+        if self.compiled is None:
+            self.compile(inp, out)
+        self.compiled(inp.view(-1), out.view(-1), *(self._args or self._launch_args(inp)))
+def _get_plan(x):
+    key = (x.shape, x.get_device())
     p = _CACHE.get(key)
     if p is None:
+        p = _CACHE[key] = FFT3DPlan(*x.shape, x.device)
     return p
     return x.contiguous()
+def _ready(x):
+    """True for the common input that needs no checks beyond these and no copy."""
+    return (
+        type(x) is torch.Tensor
+        and x.dtype is torch.float64
+        and x.is_cuda
+        and x.dim() == 3
+        and x.is_contiguous()
+        and not x.requires_grad
+        and not x.is_neg()
+        and 0 not in x.shape
+        and x.get_device() == _current_device()
+    )
 def run(input):
     """rfftn of a float64 [nz, ny, nx] CUDA tensor -> new float64 [nz, ny, nx//2+1, 2] (see module docstring)."""
+    x = input if _ready(input) else _validate(input)
+    plan = _get_plan(x)
     if x.data_ptr() % plan.in_align:
         x = x.clone(memory_format=torch.contiguous_format)
+    out = x.new_empty(plan.out_shape)
     plan.launch(x, out)
     return out

build/torch-cuda/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "fft3d",
-  "id": "_fft3d_cuda_f2cede9",
   "version": 1,
   "license": "MIT",
   "python-depends": [
@@ -17,22 +17,22 @@
     "algorithm": "sha256",
     "files": {
       "LICENSE.txt": "AIyEysQ0a1okjH0MEq1kTbdKj938zWEHbQJBfSkK23s=",
-      "__init__.py": "FpeC/pysTUfgktdzoT9ywiuLxGtBIQ2QqqfsEDvegvA=",
       "_c2r_f32/__init__.py": "XBy0pSUwgUTmssbhC7XEt/ELeTb4ayilaHy7B/V0pHw=",
       "_c2r_f32/anyfft.py": "ehbLwanScRoMLeAvopBDiFfJlVhlpimFcRuoqKgpFy0=",
       "_c2r_f32/fftgen.py": "vlccjhLwsmKsoyZ+C8DPeR7fIe7LPMxEVsq1cDF5faE=",
       "_c2r_f32/generic.py": "F286vgjce5gqBY+L9ZzqY8kAauvLc0tvxUTAYBJBfrA=",
       "_c2r_f32/inplace.py": "J+tnMJWGaiz3Gug+wbJzgqgskLgHxnbqHWNb9h1GSCU=",
-      "_c2r_f32/kernel.py": "ebURfjnZ7meANsrBG3yY7iO1waL0TSONIbB8RX+TxG8=",
       "_c2r_f32/tables.py": "RvVnytUIHgceogs37Z8lcjI5mxWclCA8D7kqC2YKm9w=",
       "_c2r_f32/twopass.py": "zI/xziZlOyfjNByjxB9Lag16rgClQVEn9kwlcoKQXT0=",
       "_c2r_f64/__init__.py": "6NiqCxU8/XZoiu7IY/nThfL5G3EovojIu+BSdpPGqEk=",
       "_c2r_f64/bigprime.py": "02gw/YOCiA32nxWpmrF1spuAExE4M5f3JgYj/zJaI5U=",
       "_c2r_f64/fftgen.py": "1Fn09b4uYmDSwItnhAycTXZVDWcmvXmf28BdjtMjsEQ=",
-      "_c2r_f64/kernel.py": "0uvhv7AdWZHYMoQqpS9A2yqJtEM/mFrwmfRmfIgOQc8=",
       "_c2r_f64/softfp.py": "B4zzVG3AA7dtfqYCCORehavfX7Xz8ypzqiML/W5hgBE=",
       "_c2r_f64/twiddles.py": "WXaDgblQUWoX/1Qoofd5mD5aUpG1EkjYsdKD3YAUPW4=",
-      "_ops.py": "Bc1c8UdnMmHRPewDW5TdU2LnIFAJFkFkES2ME9IJPsM=",
       "_r2c_f32/__init__.py": "b+FuhAkJolHH1YTeeRAb4NMHe+7YQId1VYOQ+W+vqsI=",
       "_r2c_f32/fft3k.py": "ZoLp8dkbyHLI57MrUAhd06t3+lDM+kIavTOk7YPaVPc=",
       "_r2c_f32/fft_core.py": "WThy1Zi8S9R5zSQgluOWnrhge+zXreeLR5OjysAgXNc=",
@@ -42,7 +42,7 @@
       "_r2c_f32/fft_xy.py": "Z4nZN7GM9AbNVr9JYgckfPifxfybumN5tMcmqgm/ePo=",
       "_r2c_f32/fft_ys1w.py": "4Vske6f2ATSSTb//EkpoHS5E3s0m9xGWOLuH3KWJSC8=",
       "_r2c_f32/fft_ysplit.py": "ukVz02AMlODEqFb/0qvSwAs5rA4CMFlqpI1YSRY0yqQ=",
-      "_r2c_f32/kernel.py": "GojR3ze9eI+yfP8Yqw/fjBVLZo+BNvo2oNcdklqYa7k=",
       "_r2c_f32/plan.py": "eSTB1bK8zgHLcjpMf04C5F5LVHjqot2JdFXnm8/VRes=",
       "_r2c_f32/ptx.py": "3mKM+qkQfpV6NCAXz4V4FLf6Z35LtiONgb5oSGF09fw=",
       "_r2c_f32/stockham.py": "5lS6YIDy/WKT4gTUgsQsrJHBEaAWAyjELVhD6uHJysM=",
@@ -55,7 +55,7 @@
       "_r2c_f64/fft_plan.py": "OqPrkiahnt/XreodPMi90N7+kHinYgYbLuWAdVd8eqI=",
       "_r2c_f64/fft_stages.py": "5ylt/t+bA+9ILsAG5P9CCJkwdJmfvvx/oB29w1LoO9A=",
       "_r2c_f64/fftgen.py": "1Fn09b4uYmDSwItnhAycTXZVDWcmvXmf28BdjtMjsEQ=",
-      "_r2c_f64/kernel.py": "aPX6PcigYrN9GMVKOF8DnMtRj1I+HAx6a2MtjXjI048=",
       "_r2c_f64/planner.py": "VgtohOIirLvHoK/AAciicLFmPD5nJSOCYxOomeYQZEY=",
       "_r2c_f64/tuned.py": "YtEsBZO9C63haDWFVnMBtIpTXwSKuvxAF6ffpz9wrd0="
     }
@@ -67,7 +67,7 @@
       "dirty": false
     },
     "kernel": {
-      "commit": "f2cede95bff053e1a955986b718572d24161fa35",
       "dirty": false
     }
   }

 {
   "name": "fft3d",
+  "id": "_fft3d_cuda_9e7519b",
   "version": 1,
   "license": "MIT",
   "python-depends": [
     "algorithm": "sha256",
     "files": {
       "LICENSE.txt": "AIyEysQ0a1okjH0MEq1kTbdKj938zWEHbQJBfSkK23s=",
+      "__init__.py": "q3wlQ/Ew0+FfsBtKwNvoOH9sRHaZbLGqgCBortEP+XU=",
       "_c2r_f32/__init__.py": "XBy0pSUwgUTmssbhC7XEt/ELeTb4ayilaHy7B/V0pHw=",
       "_c2r_f32/anyfft.py": "ehbLwanScRoMLeAvopBDiFfJlVhlpimFcRuoqKgpFy0=",
       "_c2r_f32/fftgen.py": "vlccjhLwsmKsoyZ+C8DPeR7fIe7LPMxEVsq1cDF5faE=",
       "_c2r_f32/generic.py": "F286vgjce5gqBY+L9ZzqY8kAauvLc0tvxUTAYBJBfrA=",
       "_c2r_f32/inplace.py": "J+tnMJWGaiz3Gug+wbJzgqgskLgHxnbqHWNb9h1GSCU=",
+      "_c2r_f32/kernel.py": "hX4XP44Cd2BE06j95sIF8RJd8nwq7sik86DIM+6uLc0=",
       "_c2r_f32/tables.py": "RvVnytUIHgceogs37Z8lcjI5mxWclCA8D7kqC2YKm9w=",
       "_c2r_f32/twopass.py": "zI/xziZlOyfjNByjxB9Lag16rgClQVEn9kwlcoKQXT0=",
       "_c2r_f64/__init__.py": "6NiqCxU8/XZoiu7IY/nThfL5G3EovojIu+BSdpPGqEk=",
       "_c2r_f64/bigprime.py": "02gw/YOCiA32nxWpmrF1spuAExE4M5f3JgYj/zJaI5U=",
       "_c2r_f64/fftgen.py": "1Fn09b4uYmDSwItnhAycTXZVDWcmvXmf28BdjtMjsEQ=",
+      "_c2r_f64/kernel.py": "F/KBXDpll4RdODgvHEsdqqU+vMxeNrJumfmiIMDOnY0=",
       "_c2r_f64/softfp.py": "B4zzVG3AA7dtfqYCCORehavfX7Xz8ypzqiML/W5hgBE=",
       "_c2r_f64/twiddles.py": "WXaDgblQUWoX/1Qoofd5mD5aUpG1EkjYsdKD3YAUPW4=",
+      "_ops.py": "wpRc/vzJsXg3CcTlM6GxXCVZFLCb6nJ0ocSnhY/mvK8=",
       "_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_xy.py": "Z4nZN7GM9AbNVr9JYgckfPifxfybumN5tMcmqgm/ePo=",
       "_r2c_f32/fft_ys1w.py": "4Vske6f2ATSSTb//EkpoHS5E3s0m9xGWOLuH3KWJSC8=",
       "_r2c_f32/fft_ysplit.py": "ukVz02AMlODEqFb/0qvSwAs5rA4CMFlqpI1YSRY0yqQ=",
+      "_r2c_f32/kernel.py": "k/jj+b9XOwBzf/Uug9hwXg0JUqOTjVJ6mhalZs1PNi8=",
       "_r2c_f32/plan.py": "eSTB1bK8zgHLcjpMf04C5F5LVHjqot2JdFXnm8/VRes=",
       "_r2c_f32/ptx.py": "3mKM+qkQfpV6NCAXz4V4FLf6Z35LtiONgb5oSGF09fw=",
       "_r2c_f32/stockham.py": "5lS6YIDy/WKT4gTUgsQsrJHBEaAWAyjELVhD6uHJysM=",
       "_r2c_f64/fft_plan.py": "OqPrkiahnt/XreodPMi90N7+kHinYgYbLuWAdVd8eqI=",
       "_r2c_f64/fft_stages.py": "5ylt/t+bA+9ILsAG5P9CCJkwdJmfvvx/oB29w1LoO9A=",
       "_r2c_f64/fftgen.py": "1Fn09b4uYmDSwItnhAycTXZVDWcmvXmf28BdjtMjsEQ=",
+      "_r2c_f64/kernel.py": "3DZZyoLqdJnlS9WRS+ujv0M+aMAhAN1/B1zvpSSbdBY=",
       "_r2c_f64/planner.py": "VgtohOIirLvHoK/AAciicLFmPD5nJSOCYxOomeYQZEY=",
       "_r2c_f64/tuned.py": "YtEsBZO9C63haDWFVnMBtIpTXwSKuvxAF6ffpz9wrd0="
     }
       "dirty": false
     },
     "kernel": {
+      "commit": "9e7519bbed9eaea4ef8f23df5cdc6654543f8f72",
       "dirty": false
     }
   }