diff --git a/benchmarks/README.md b/benchmarks/README.md index ab878247..eb8b68a3 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -12,6 +12,18 @@ Performance benchmarks for [mkl_fft](https://github.com/IntelPython/mkl_fft) usi | `bench_interfaces.py` | `mkl_fft.interfaces.{numpy_fft, scipy_fft}` | All exported functions; selected by a `module` parameter. Hermitian 2-D/N-D (`hfft2`, `hfftn`) are scipy-only. | float32, float64, complex64, complex128 | power-of-two and cubic | | `bench_memory.py` | `mkl_fft` | Peak RSS for 1-D, 2-D, and 3-D transforms | float32, float64, complex128 | power-of-two | +The suites above vary transform size, shape, and dtype while calling with +default arguments on a freshly allocated C-contiguous array. The suites below +hold the transform fixed and vary how it is *requested*, which selects between +backend dispatch paths that differ by orders of magnitude in cost: + +| File | Varies | Why it matters | +|------|--------|----------------| +| `bench_axes.py` | `axes=` subsets; `axis=` position | A strict subset of axes is dispatched slice by slice over the complementary axes, so cost scales with slice count rather than transform size. A 1-D transform of a rank > 2 array is batched into one `DftiCompute` only when the axis is first or last. | +| `bench_args.py` | `norm=`, `s=`, `n=` | `norm` other than `None`/`"backward"` computes a scale factor in Python; an `s` that pads or truncates leaves the direct N-D path for `_iter_fftnd`. | +| `bench_layout.py` | input layout (C/F/strided), `out=` layout | A non-contiguous input is not one segment, so for rank > 2 the backend iterates one vector at a time. OneMKL requires `out` to have the same element strides as the input; otherwise `_pydfti` allocates a temporary and copies. | +| `bench_descriptor.py` | call sequence; per-call floor | The thread-local descriptor cache holds a single entry, so interleaving lengths, dtypes, or domains rebuilds it on every call. | + ## Threading Set `MKL_NUM_THREADS` in the environment before running ASV to control the @@ -32,12 +44,37 @@ recommendation. ### DFTI descriptor warmup -MKL creates a DFTI descriptor on the first FFT call for a given (size, dtype, -strides) combination and reuses it on subsequent calls. To avoid charging -that one-time cost to the first measured iteration, each benchmark's `setup` -performs an explicit warmup call after preparing the input array. ASV's -default `warmup_time` (0.1s) already amortizes this for sub-millisecond -transforms, but the explicit warmup makes the intent visible. +For **1-D** transforms, `_pydfti` keeps one DFTI descriptor in thread-local +storage and reuses it when the next call matches on rank, precision, domain, +and length. To avoid charging that one-time build to the first measured +iteration, each benchmark's `setup` performs an explicit warmup call after +preparing the input array. ASV's default `warmup_time` (0.1s) already +amortizes this for sub-millisecond transforms, but the explicit warmup makes +the intent visible. + +Two limits of that cache are worth knowing when reading results: + +- It holds a **single** descriptor, so a call sequence that alternates length, + dtype, or forward domain rebuilds it every time rather than reusing it. Only + `bench_descriptor.py` exercises that; every other suite repeats one shape and + so always hits. +- The **N-D** entry points in `mkl_fft/src/mklfft.c.src` take no cache + argument at all and build, commit, and free a descriptor on every call. The + warmup in the `fft2`/`fftn` suites therefore warms MKL's own internal state + but does not prime a descriptor for reuse. + +### Paired controls + +The `bench_descriptor.py` suites are written as pairs: an alternating call +sequence alongside a control that issues the same number of transforms with +unchanging parameters. Neither number means much alone — the quantity of +interest is the difference. `time_switch_direction` is a deliberate negative +control: forward and backward scales live on the same descriptor, so it should +sit close to `time_repeat`. + +Several suites here also carry an in-suite control among their parameters: +the full-dimensional `axes` tuple in `bench_axes.py`, `norm=None` in +`bench_args.py`, and `out_kind="none"` in `bench_layout.py`. ## Running Benchmarks diff --git a/benchmarks/benchmarks/bench_args.py b/benchmarks/benchmarks/bench_args.py new file mode 100644 index 00000000..a29fb209 --- /dev/null +++ b/benchmarks/benchmarks/bench_args.py @@ -0,0 +1,181 @@ +"""Benchmarks for argument handling overhead. + +These suites hold the transform itself fixed and vary only the keyword +arguments, isolating the cost of the Python-level argument processing that +runs before MKL is reached: + +* ``norm`` other than ``None``/``"backward"`` goes through + ``_fft_utils._compute_fwd_scale``. On small transforms the scale computation + is a measurable fraction of the whole call. +* An explicit ``s`` that matches the input shape still reaches the direct N-D + path, but padding or truncating dispatches through ``_fft_utils._iter_fftnd``, + which re-normalizes shape and axes on every call and performs the pad/trim. + +Every other suite in this directory calls with default arguments only. +""" + +import mkl_fft + +from ._utils import _DTYPES_REAL, _DTYPES_REDUCED, BenchC2C, BenchR2C + +_NORMS = [None, "backward", "forward", "ortho"] + +# "match" is the in-suite control: it should stay on the direct N-D path. +_S_MODES = ["match", "trunc", "pad"] + + +def _shape_for_mode(shape, mode): + """Return an ``s`` argument that matches, truncates, or pads *shape*.""" + if mode == "trunc": + return tuple(d // 2 for d in shape) + if mode == "pad": + return tuple(d + d // 2 for d in shape) + return tuple(shape) + + +# --------------------------------------------------------------------------- +# norm +# --------------------------------------------------------------------------- + + +class BenchNorm1D(BenchC2C): + """fft with each ``norm`` mode. + + ``n=64`` is small enough that the transform costs about a microsecond, so + scale-factor computation shows up directly. + """ + + params = [[64, 1024, 16384], _DTYPES_REDUCED, _NORMS] + param_names = ["n", "dtype", "norm"] + + def setup(self, n, dtype, norm): + super().setup(n, dtype) + mkl_fft.fft(self.x, norm=norm) + + def time_fft(self, n, dtype, norm): + mkl_fft.fft(self.x, norm=norm) + + def time_ifft(self, n, dtype, norm): + mkl_fft.ifft(self.x, norm=norm) + + +class BenchNormND(BenchC2C): + """fftn with each ``norm`` mode. + + For N-D transforms the forward scale is a product over the whole shape, + so this path does strictly more work than the 1-D one. + """ + + params = [[(64, 64), (256, 256), (32, 32, 32)], ["complex128"], _NORMS] + param_names = ["shape", "dtype", "norm"] + + def setup(self, shape, dtype, norm): + super().setup(shape, dtype) + mkl_fft.fftn(self.x, norm=norm) + + def time_fftn(self, shape, dtype, norm): + mkl_fft.fftn(self.x, norm=norm) + + def time_ifftn(self, shape, dtype, norm): + mkl_fft.ifftn(self.x, norm=norm) + + +class BenchNormR2C1D(BenchR2C): + """rfft / irfft with each ``norm`` mode.""" + + params = [[64, 16384], _DTYPES_REAL, _NORMS] + param_names = ["n", "dtype", "norm"] + + def setup(self, n, dtype, norm): + super().setup(n, dtype) + mkl_fft.rfft(self.x_real, norm=norm) + mkl_fft.irfft(self.x_complex, n=n, norm=norm) + + def time_rfft(self, n, dtype, norm): + mkl_fft.rfft(self.x_real, norm=norm) + + def time_irfft(self, n, dtype, norm): + mkl_fft.irfft(self.x_complex, n=n, norm=norm) + + +# --------------------------------------------------------------------------- +# explicit output shape (s) +# --------------------------------------------------------------------------- + + +class BenchShapeArg2D(BenchC2C): + """fftn with an explicit ``s`` that matches, truncates, or pads.""" + + params = [[(256, 256)], _DTYPES_REDUCED, _S_MODES] + param_names = ["shape", "dtype", "mode"] + + def setup(self, shape, dtype, mode): + super().setup(shape, dtype) + self.s = _shape_for_mode(shape, mode) + mkl_fft.fftn(self.x, s=self.s) + + def time_fftn(self, shape, dtype, mode): + mkl_fft.fftn(self.x, s=self.s) + + def time_ifftn(self, shape, dtype, mode): + mkl_fft.ifftn(self.x, s=self.s) + + +class BenchShapeArg3D(BenchC2C): + """fftn with an explicit ``s`` on a 3-D array.""" + + params = [[(32, 32, 32)], _DTYPES_REDUCED, _S_MODES] + param_names = ["shape", "dtype", "mode"] + + def setup(self, shape, dtype, mode): + super().setup(shape, dtype) + self.s = _shape_for_mode(shape, mode) + mkl_fft.fftn(self.x, s=self.s) + + def time_fftn(self, shape, dtype, mode): + mkl_fft.fftn(self.x, s=self.s) + + def time_ifftn(self, shape, dtype, mode): + mkl_fft.ifftn(self.x, s=self.s) + + +class BenchShapeArgR2C2D(BenchR2C): + """rfftn with an explicit ``s`` that matches, truncates, or pads.""" + + params = [[(256, 256)], _DTYPES_REAL, _S_MODES] + param_names = ["shape", "dtype", "mode"] + + def setup(self, shape, dtype, mode): + super().setup(shape, dtype) + self.s = _shape_for_mode(shape, mode) + mkl_fft.rfftn(self.x_real, s=self.s) + + def time_rfftn(self, shape, dtype, mode): + mkl_fft.rfftn(self.x_real, s=self.s) + + +# --------------------------------------------------------------------------- +# 1-D explicit length (n) +# --------------------------------------------------------------------------- + + +class BenchLengthArg1D(BenchC2C): + """fft with an explicit ``n`` that matches, truncates, or pads. + + Padding forces a copy into a larger buffer inside ``_pydfti._pad_array``; + truncating returns a view of a longer in-place result. + """ + + params = [[16384], _DTYPES_REDUCED, _S_MODES] + param_names = ["n", "dtype", "mode"] + + def setup(self, n, dtype, mode): + super().setup(n, dtype) + self.n = _shape_for_mode((n,), mode)[0] + mkl_fft.fft(self.x, n=self.n) + + def time_fft(self, n, dtype, mode): + mkl_fft.fft(self.x, n=self.n) + + def time_ifft(self, n, dtype, mode): + mkl_fft.ifft(self.x, n=self.n) diff --git a/benchmarks/benchmarks/bench_axes.py b/benchmarks/benchmarks/bench_axes.py new file mode 100644 index 00000000..941c7184 --- /dev/null +++ b/benchmarks/benchmarks/bench_axes.py @@ -0,0 +1,164 @@ +"""Benchmarks for axis and axes selection. + +Which axes a transform is asked for changes the dispatch path, not just the +amount of arithmetic: + +* ``axes=None`` (or a tuple covering every axis) reaches the batched N-D MKL + descriptor in one call. +* A strict *subset* of axes is dispatched slice by slice over the + complementary axes (``_fft_utils._iter_complementary``), so cost scales with + the number of complementary slices rather than with transform size. +* For a 1-D transform of an array of rank > 2, the C backend can issue a + single batched ``DftiCompute`` only when the axis is the first or the last + one (``mklfft.c.src``, ``compute_strides_and_distances``). Any interior axis + falls back to iterating one vector at a time. + +The other suites in this directory always request every axis, so none of these +paths were previously visible to the dashboard. +""" + +import mkl_fft + +from ._utils import _DTYPES_REAL, _DTYPES_REDUCED, BenchC2C, BenchR2C + +# axes subsets, plus the full-dimensional tuple as an in-suite control +_AXES_2D = [(0,), (1,), (0, 1)] +_AXES_3D = [(0,), (1,), (2,), (0, 1), (1, 2), (0, 1, 2)] + +# For c2r the half-spectrum input is laid out along the array's last axis, so +# only axes tuples ending on that axis are meaningful here. +_AXES_3D_R2C = [(2,), (1, 2), (0, 2), (0, 1, 2)] + + +# --------------------------------------------------------------------------- +# 2-D complex-to-complex over a subset of axes +# --------------------------------------------------------------------------- + + +class BenchAxes2D(BenchC2C): + """fftn / ifftn over a subset of the axes of a 2-D array.""" + + params = [[(128, 128), (512, 512)], _DTYPES_REDUCED, _AXES_2D] + param_names = ["shape", "dtype", "axes"] + + def setup(self, shape, dtype, axes): + super().setup(shape, dtype) + mkl_fft.fftn(self.x, axes=axes) + + def time_fftn(self, shape, dtype, axes): + mkl_fft.fftn(self.x, axes=axes) + + def time_ifftn(self, shape, dtype, axes): + mkl_fft.ifftn(self.x, axes=axes) + + +# --------------------------------------------------------------------------- +# 3-D complex-to-complex over a subset of axes +# --------------------------------------------------------------------------- + + +class BenchAxes3D(BenchC2C): + """fftn / ifftn over a subset of the axes of a 3-D array. + + The number of complementary slices spans three orders of magnitude across + this parameter set: ``(0, 1, 2)`` is a single batched call, ``(1, 2)`` + iterates 32 slices, and ``(2,)`` iterates 1024. The slice count is the + product of the untransformed axes, so larger shapes amplify the effect; + this one is kept modest to bound suite runtime. + """ + + params = [[(32, 32, 32)], _DTYPES_REDUCED, _AXES_3D] + param_names = ["shape", "dtype", "axes"] + + def setup(self, shape, dtype, axes): + super().setup(shape, dtype) + mkl_fft.fftn(self.x, axes=axes) + + def time_fftn(self, shape, dtype, axes): + mkl_fft.fftn(self.x, axes=axes) + + def time_ifftn(self, shape, dtype, axes): + mkl_fft.ifftn(self.x, axes=axes) + + +# --------------------------------------------------------------------------- +# 3-D real-to-complex / complex-to-real over a subset of axes +# --------------------------------------------------------------------------- + + +class BenchAxesR2C3D(BenchR2C): + """rfftn / irfftn over a subset of the axes of a 3-D array.""" + + params = [[(32, 32, 32)], _DTYPES_REAL, _AXES_3D_R2C] + param_names = ["shape", "dtype", "axes"] + + def setup(self, shape, dtype, axes): + super().setup(shape, dtype) + # shape of the result along the requested axes + self.s = tuple(shape[ax] for ax in axes) + mkl_fft.rfftn(self.x_real, axes=axes) + mkl_fft.irfftn(self.x_complex, s=self.s, axes=axes) + + def time_rfftn(self, shape, dtype, axes): + mkl_fft.rfftn(self.x_real, axes=axes) + + def time_irfftn(self, shape, dtype, axes): + mkl_fft.irfftn(self.x_complex, s=self.s, axes=axes) + + +# --------------------------------------------------------------------------- +# 1-D transform along each axis of a higher-rank array +# --------------------------------------------------------------------------- + + +class BenchAxis3D(BenchC2C): + """fft / ifft along each individual axis of a 3-D array. + + The backend batches the whole transform into one ``DftiCompute`` call only + for ``axis=0`` and ``axis=2``; ``axis=1`` iterates vector by vector. + """ + + params = [[(32, 32, 32), (64, 64, 64)], _DTYPES_REDUCED, [0, 1, 2]] + param_names = ["shape", "dtype", "axis"] + + def setup(self, shape, dtype, axis): + super().setup(shape, dtype) + mkl_fft.fft(self.x, axis=axis) + + def time_fft(self, shape, dtype, axis): + mkl_fft.fft(self.x, axis=axis) + + def time_ifft(self, shape, dtype, axis): + mkl_fft.ifft(self.x, axis=axis) + + +class BenchAxis4D(BenchC2C): + """fft along each individual axis of a 4-D array. + + A rank-4 array has two interior axes, so the batched and per-vector paths + are exercised twice each within one parameter sweep. + """ + + params = [[(16, 16, 16, 16)], ["complex128"], [0, 1, 2, 3]] + param_names = ["shape", "dtype", "axis"] + + def setup(self, shape, dtype, axis): + super().setup(shape, dtype) + mkl_fft.fft(self.x, axis=axis) + + def time_fft(self, shape, dtype, axis): + mkl_fft.fft(self.x, axis=axis) + + +class BenchAxisR2C3D(BenchR2C): + """rfft / irfft along each individual axis of a 3-D array.""" + + params = [[(64, 64, 64)], _DTYPES_REAL, [0, 1, 2]] + param_names = ["shape", "dtype", "axis"] + + def setup(self, shape, dtype, axis): + super().setup(shape, dtype) + mkl_fft.rfft(self.x_real, axis=axis) + + def time_rfft(self, shape, dtype, axis): + mkl_fft.rfft(self.x_real, axis=axis) diff --git a/benchmarks/benchmarks/bench_descriptor.py b/benchmarks/benchmarks/bench_descriptor.py new file mode 100644 index 00000000..ad5fcdfd --- /dev/null +++ b/benchmarks/benchmarks/bench_descriptor.py @@ -0,0 +1,163 @@ +"""Benchmarks for DFTI descriptor lifecycle and call patterns. + +Every other suite in this directory calls one transform shape repeatedly, so +the thread-local DFTI descriptor cache in ``_pydfti`` always hits after +warmup. Real callers interleave transforms of different length, dtype, or +domain. The cache holds a single descriptor, so each switch frees the cached +one and builds and commits a replacement. + +Each suite here pairs an alternating call sequence with a same-parameter +control that issues the same number of transforms, so the switch cost is +readable as the difference between the two. + +``BenchFixedCost`` measures the per-call floor at sizes where the transform +itself is negligible. The N-D entry points in ``mklfft.c.src`` take no cache +argument and build a descriptor on every call, so the N-D floor is expected to +sit above the 1-D one. +""" + +import numpy as np + +import mkl_fft + +from ._utils import _RNG_SEED, _make_input + +# Offset used to build a second, differently sized input. Small enough that +# both lengths factor similarly, so the difference measured is descriptor +# rebuild rather than a change of MKL algorithm class. +_SIZE_DELTA = 24 + + +# --------------------------------------------------------------------------- +# 1-D descriptor switching +# --------------------------------------------------------------------------- + + +class BenchDescriptorSwitch1D: + """Alternating vs repeated transform parameters, 1-D. + + ``time_repeat`` is the control for ``time_switch_size`` and + ``time_switch_dtype``. ``time_switch_domain`` is compared against + ``time_repeat_domain``; those two are not an exact control pair, because + rfft and fft do not cost the same, but at the smaller ``n`` the descriptor + rebuild dominates that difference. + """ + + params = [[1024, 65536]] + param_names = ["n"] + + def setup(self, n): + rng = np.random.default_rng(_RNG_SEED) + self.a = _make_input(rng, n, "complex128") + self.b = _make_input(rng, n - _SIZE_DELTA, "complex128") + self.a32 = self.a.astype("complex64") + self.real = _make_input(rng, n, "float64") + mkl_fft.fft(self.a) + mkl_fft.fft(self.b) + mkl_fft.fft(self.a32) + mkl_fft.rfft(self.real) + + def time_repeat(self, n): + mkl_fft.fft(self.a) + mkl_fft.fft(self.a) + + def time_switch_size(self, n): + mkl_fft.fft(self.a) + mkl_fft.fft(self.b) + + def time_switch_dtype(self, n): + mkl_fft.fft(self.a) + mkl_fft.fft(self.a32) + + def time_repeat_domain(self, n): + mkl_fft.rfft(self.real) + mkl_fft.rfft(self.real) + + def time_switch_domain(self, n): + mkl_fft.rfft(self.real) + mkl_fft.fft(self.a) + + def time_switch_direction(self, n): + mkl_fft.fft(self.a) + mkl_fft.ifft(self.a) + + +# --------------------------------------------------------------------------- +# Descriptor switching driven by axis choice +# --------------------------------------------------------------------------- + + +class BenchDescriptorSwitchAxis: + """Alternating the transformed axis of a 2-D array. + + This is the call pattern ``_fft_utils._iter_fftnd`` generates internally, + and that 2-D user code generates directly. + + The two axes of a C-contiguous array do not cost the same — axis 0 walks + the long stride — so ``time_switch_axis`` cannot be read against a single + control. Both per-axis controls are therefore reported, and the descriptor + contribution is ``switch - (repeat_axis0 + repeat_axis1) / 2``. On a + non-square array the transform length changes with the axis, so the cached + descriptor is discarded on every call; on a square array it survives. + """ + + params = [[(512, 512), (512, 256)]] + param_names = ["shape"] + + def setup(self, shape): + rng = np.random.default_rng(_RNG_SEED) + self.x = _make_input(rng, shape, "complex128") + mkl_fft.fft(self.x, axis=0) + mkl_fft.fft(self.x, axis=1) + + def time_repeat_axis0(self, shape): + mkl_fft.fft(self.x, axis=0) + mkl_fft.fft(self.x, axis=0) + + def time_repeat_axis1(self, shape): + mkl_fft.fft(self.x, axis=1) + mkl_fft.fft(self.x, axis=1) + + def time_switch_axis(self, shape): + mkl_fft.fft(self.x, axis=1) + mkl_fft.fft(self.x, axis=0) + + +# --------------------------------------------------------------------------- +# Per-call fixed cost +# --------------------------------------------------------------------------- + + +class BenchFixedCost: + """Per-call cost at sizes where the transform itself is negligible. + + Anything these report above a few hundred nanoseconds is argument + processing, dispatch, allocation, and descriptor setup rather than + arithmetic. + """ + + params = [["float64", "complex128"]] + param_names = ["dtype"] + + def setup(self, dtype): + rng = np.random.default_rng(_RNG_SEED) + self.x1 = _make_input(rng, 8, dtype) + self.x2 = _make_input(rng, (4, 4), dtype) + self.x3 = _make_input(rng, (4, 4, 4), dtype) + self.real1 = _make_input(rng, 8, "float64") + mkl_fft.fft(self.x1) + mkl_fft.fft2(self.x2) + mkl_fft.fftn(self.x3) + mkl_fft.rfft(self.real1) + + def time_fft_1d(self, dtype): + mkl_fft.fft(self.x1) + + def time_fft2_2d(self, dtype): + mkl_fft.fft2(self.x2) + + def time_fftn_3d(self, dtype): + mkl_fft.fftn(self.x3) + + def time_rfft_1d(self, dtype): + mkl_fft.rfft(self.real1) diff --git a/benchmarks/benchmarks/bench_layout.py b/benchmarks/benchmarks/bench_layout.py new file mode 100644 index 00000000..a6fbcfd5 --- /dev/null +++ b/benchmarks/benchmarks/bench_layout.py @@ -0,0 +1,153 @@ +"""Benchmarks for input and output memory layout. + +Every other suite in this directory transforms a freshly allocated +C-contiguous array. Callers routinely pass something else, and layout decides +which backend path runs: + +* A C- or F-contiguous array is one segment, so the C backend can describe the + whole batch to MKL with a single stride/distance pair. +* A non-contiguous view is not one segment, so for rank > 2 the backend falls + back to iterating one vector at a time + (``mklfft.c.src``, ``compute_strides_and_distances``). +* OneMKL requires ``out`` to have the same element strides as the input. When + they differ, ``_pydfti`` allocates a temporary and copies the result into + ``out``, costing an extra full pass. +""" + +import numpy as np + +import mkl_fft + +from ._utils import _DTYPES_REDUCED, _RNG_SEED, _make_input + +_LAYOUTS = ["C", "F", "strided"] + +_OUT_KINDS = ["none", "contig", "strided"] + + +def _layout_input(shape, dtype, layout): + """Return an array of *shape* and *dtype* in the requested *layout*.""" + rng = np.random.default_rng(_RNG_SEED) + if layout == "strided": + # A view with a gap between elements on every axis; not one segment. + big = _make_input(rng, tuple(2 * d for d in shape), dtype) + return big[(slice(None, None, 2),) * len(shape)] + x = _make_input(rng, shape, dtype) + if layout == "F": + return np.asfortranarray(x) + return x + + +# --------------------------------------------------------------------------- +# 2-D input layout +# --------------------------------------------------------------------------- + + +class BenchLayout2D: + """fft2 / ifft2 over C-contiguous, F-ordered, and strided input.""" + + params = [[(256, 256), (512, 512)], _DTYPES_REDUCED, _LAYOUTS] + param_names = ["shape", "dtype", "layout"] + + def setup(self, shape, dtype, layout): + self.x = _layout_input(shape, dtype, layout) + mkl_fft.fft2(self.x) + + def time_fft2(self, shape, dtype, layout): + mkl_fft.fft2(self.x) + + def time_ifft2(self, shape, dtype, layout): + mkl_fft.ifft2(self.x) + + +# --------------------------------------------------------------------------- +# 3-D input layout +# --------------------------------------------------------------------------- + + +class BenchLayout3D: + """fftn / ifftn over C-contiguous, F-ordered, and strided input.""" + + params = [[(64, 64, 64)], _DTYPES_REDUCED, _LAYOUTS] + param_names = ["shape", "dtype", "layout"] + + def setup(self, shape, dtype, layout): + self.x = _layout_input(shape, dtype, layout) + mkl_fft.fftn(self.x) + + def time_fftn(self, shape, dtype, layout): + mkl_fft.fftn(self.x) + + def time_ifftn(self, shape, dtype, layout): + mkl_fft.ifftn(self.x) + + +class BenchLayoutAxis3D: + """1-D fft along the last axis of a 3-D array of each layout. + + A strided rank-3 input cannot be handled as one batched call even along + the last axis, so this pairs with ``bench_axes.BenchAxis3D`` to separate + the layout effect from the axis-position effect. + """ + + params = [[(64, 64, 64)], ["complex128"], _LAYOUTS] + param_names = ["shape", "dtype", "layout"] + + def setup(self, shape, dtype, layout): + self.x = _layout_input(shape, dtype, layout) + mkl_fft.fft(self.x, axis=-1) + + def time_fft(self, shape, dtype, layout): + mkl_fft.fft(self.x, axis=-1) + + +# --------------------------------------------------------------------------- +# out= layout +# --------------------------------------------------------------------------- + + +class BenchOut2D: + """fft2 writing into a caller-supplied ``out`` of each layout. + + ``none`` is the control; ``contig`` should hit the in-place-into-out fast + path; ``strided`` forces a temporary allocation plus a copy. + """ + + params = [[(512, 512)], ["complex128"], _OUT_KINDS] + param_names = ["shape", "dtype", "out_kind"] + + def setup(self, shape, dtype, out_kind): + rng = np.random.default_rng(_RNG_SEED) + self.x = _make_input(rng, shape, dtype) + if out_kind == "contig": + self.out = np.empty(shape, dtype=dtype) + elif out_kind == "strided": + wide = shape[:-1] + (2 * shape[-1],) + self.out = np.empty(wide, dtype=dtype)[..., ::2] + else: + self.out = None + mkl_fft.fft2(self.x, out=self.out) + + def time_fft2(self, shape, dtype, out_kind): + mkl_fft.fft2(self.x, out=self.out) + + +class BenchOut1D: + """fft writing into a caller-supplied ``out`` of each layout.""" + + params = [[65536], ["complex128"], _OUT_KINDS] + param_names = ["n", "dtype", "out_kind"] + + def setup(self, n, dtype, out_kind): + rng = np.random.default_rng(_RNG_SEED) + self.x = _make_input(rng, n, dtype) + if out_kind == "contig": + self.out = np.empty(n, dtype=dtype) + elif out_kind == "strided": + self.out = np.empty(2 * n, dtype=dtype)[::2] + else: + self.out = None + mkl_fft.fft(self.x, out=self.out) + + def time_fft(self, n, dtype, out_kind): + mkl_fft.fft(self.x, out=self.out)