From cbb20aceb2329a5fce61cb7244f9f3be7ff743f6 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 18 Nov 2025 17:58:49 -0600 Subject: [PATCH 1/6] VkFFT: use public API --- sumpy/tools.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/sumpy/tools.py b/sumpy/tools.py index ef9bde38..2cf06436 100644 --- a/sumpy/tools.py +++ b/sumpy/tools.py @@ -40,7 +40,6 @@ import loopy as lp import pytools.obj_array as obj_array from pymbolic.mapper.dependency import DependencyMapper -from pyopencl import MemoryObjectHolder from pytools import memoize_method from pytools.tag import Tag, tag_dataclass @@ -1033,17 +1032,9 @@ def run_opencl_fft( else: output_vec = cla.empty_like(input_vec, queue) - # FIXME: use the public API once - # https://github.com/vincefn/pyvkfft/pull/17 is in - from pyvkfft.opencl import _vkfft_opencl - if inverse: # noqa: SIM108 - meth = _vkfft_opencl.ifft - else: - meth = _vkfft_opencl.fft + meth = app.ifft if inverse else app.fft - assert isinstance(output_vec.data, MemoryObjectHolder) - meth(app.app, int(input_vec.data.int_ptr), - int(output_vec.data.int_ptr), int(queue.int_ptr)) + meth(input_vec, output_vec, queue=queue) if queue.device.platform.name == "NVIDIA CUDA": end_evt = cl.enqueue_marker(queue) From f1563a2917df03511be89402d8ba1d7a851b8b93 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 18 Nov 2025 17:59:12 -0600 Subject: [PATCH 2/6] VkFFT: Fall back to loopy FFT on pocl 7 --- sumpy/tools.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/sumpy/tools.py b/sumpy/tools.py index 2cf06436..c7f76d64 100644 --- a/sumpy/tools.py +++ b/sumpy/tools.py @@ -40,6 +40,7 @@ import loopy as lp import pytools.obj_array as obj_array from pymbolic.mapper.dependency import DependencyMapper +from pyopencl.characterize import get_pocl_version from pytools import memoize_method from pytools.tag import Tag, tag_dataclass @@ -955,15 +956,22 @@ def _get_fft_backend(queue: pyopencl.CommandQueue) -> FFTBackend: import platform import sys - if (sys.platform == "darwin" - and platform.machine() == "x86_64" - and queue.context.devices[0].platform.name - == "Portable Computing Language"): - warnings.warn( - "PoCL miscompiles some VkFFT kernels. " - "See https://github.com/inducer/sumpy/issues/129. " - "Falling back to slower implementation.", stacklevel=3) - return FFTBackend.loopy + pocl_ver = get_pocl_version(queue.device.platform) + if pocl_ver is not None: + if pocl_ver >= (7,): + warnings.warn( + "PoCL>=7 miscompiles VkFFT. " + "See https://github.com/pocl/pocl/issues/2069 for details. " + "Falling back to slower implementation.", stacklevel=3) + return FFTBackend.loopy + + if (sys.platform == "darwin" + and platform.machine() == "x86_64"): + warnings.warn( + "PoCL crashes on some VkFFT kernels on MacOS. " + "See https://github.com/inducer/sumpy/issues/129. " + "Falling back to slower implementation.", stacklevel=3) + return FFTBackend.loopy return FFTBackend.pyvkfft From 2409892a2502060ab4c2588b3043620b8a284cdf Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 18 Nov 2025 17:59:28 -0600 Subject: [PATCH 3/6] Expansion toys: make compatible with Loopy FFT --- sumpy/toys.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sumpy/toys.py b/sumpy/toys.py index 18460f98..e095a40a 100644 --- a/sumpy/toys.py +++ b/sumpy/toys.py @@ -427,9 +427,9 @@ def _m2l(psource, to_center, to_rscale, to_order: int, e2e, expn_class, expn_kwa from sumpy.tools import get_native_event, get_opencl_fft_app, run_opencl_fft if toy_ctx.use_fft: - fft_app = get_opencl_fft_app(queue, (expn_size,), + fft_app = get_opencl_fft_app(queue, (1, expn_size,), dtype=preprocessed_src_expansions.dtype, inverse=False) - ifft_app = get_opencl_fft_app(queue, (expn_size,), + ifft_app = get_opencl_fft_app(queue, (1, expn_size,), dtype=preprocessed_src_expansions.dtype, inverse=True) evt, preprocessed_src_expansions = run_opencl_fft(fft_app, queue, From 755d9f1a707ae8213fe7dee28e5c5cdab972f3d2 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 18 Nov 2025 18:12:02 -0600 Subject: [PATCH 4/6] Add CI job with Loopy FFT --- .github/workflows/ci.yml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ae476b17..6ffb1318 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -71,6 +71,17 @@ jobs: curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh . ./build-and-test-py-project-within-miniconda.sh + pytest_symengine_loopy_fft: + name: Conda Pytest Symengine with Loopy FFT + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + - name: "Main Script" + run: | + curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh + export SUMPY_FFT_BACKEND=loopy + . ./build-and-test-py-project-within-miniconda.sh + pytest_symengine: name: Conda Pytest Symengine runs-on: ubuntu-latest From 90ad7a2f7beb3ef460619bff5822288a90925f30 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 18 Nov 2025 18:41:01 -0600 Subject: [PATCH 5/6] Loopy FFT: don't bake batch dimensions lengths into kernel --- sumpy/test/test_tools.py | 6 +++++- sumpy/tools.py | 41 +++++++++++++++++++++++----------------- 2 files changed, 29 insertions(+), 18 deletions(-) diff --git a/sumpy/test/test_tools.py b/sumpy/test/test_tools.py index e68cff4c..89ec231d 100644 --- a/sumpy/test/test_tools.py +++ b/sumpy/test/test_tools.py @@ -93,7 +93,11 @@ def test_fft(actx_factory: ArrayContextFactory, size: int): inp_dev = actx.from_numpy(inp) out = fft(inp) - fft_func = loopy_fft(inp.shape, inverse=False, complex_dtype=inp.dtype.type) + fft_func = loopy_fft( + inp.shape[-1], + n_batch_dims=len(inp.shape) - 1, + inverse=False, + complex_dtype=inp.dtype.type) _evt, (out_dev,) = fft_func(actx.queue, y=inp_dev) assert np.allclose(actx.to_numpy(out_dev), out) diff --git a/sumpy/tools.py b/sumpy/tools.py index c7f76d64..f644a304 100644 --- a/sumpy/tools.py +++ b/sumpy/tools.py @@ -752,7 +752,8 @@ def wait(self): def loopy_fft( - shape: tuple[int, ...], + n: int, + *, n_batch_dims: int, inverse: bool, complex_dtype: DTypeLike, index_dtype: DTypeLike | None = None, @@ -766,7 +767,6 @@ def loopy_fft( complex_dtype = np.dtype(complex_dtype) sign = 1 if not inverse else -1 - n = shape[-1] m = n factors = [] @@ -776,13 +776,13 @@ def loopy_fft( nfft = n - broadcast_dims = tuple(var(f"j{d}") for d in range(len(shape) - 1)) + batch_dims = tuple(var(f"j{d}") for d in range(n_batch_dims)) domains = [ "{[i]: 0<=i Date: Tue, 18 Nov 2025 18:12:17 -0600 Subject: [PATCH 6/6] Update baseline --- .basedpyright/baseline.json | 36 ++++++++++-------------------------- 1 file changed, 10 insertions(+), 26 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index c362b12b..8a15fab4 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -36153,8 +36153,8 @@ "code": "reportAny", "range": { "startColumn": 15, - "endColumn": 83, - "lineCount": 1 + "endColumn": 46, + "lineCount": 5 } }, { @@ -36229,43 +36229,27 @@ "lineCount": 1 } }, - { - "code": "reportMissingTypeStubs", - "range": { - "startColumn": 13, - "endColumn": 27, - "lineCount": 1 - } - }, - { - "code": "reportAny", - "range": { - "startColumn": 13, - "endColumn": 20, - "lineCount": 1 - } - }, { "code": "reportAny", "range": { - "startColumn": 26, - "endColumn": 40, + "startColumn": 8, + "endColumn": 12, "lineCount": 1 } }, { "code": "reportAny", "range": { - "startColumn": 26, - "endColumn": 48, + "startColumn": 15, + "endColumn": 23, "lineCount": 1 } }, { "code": "reportAny", "range": { - "startColumn": 26, - "endColumn": 48, + "startColumn": 40, + "endColumn": 47, "lineCount": 1 } }, @@ -37331,7 +37315,7 @@ "code": "reportUnknownArgumentType", "range": { "startColumn": 48, - "endColumn": 60, + "endColumn": 63, "lineCount": 1 } }, @@ -37355,7 +37339,7 @@ "code": "reportUnknownArgumentType", "range": { "startColumn": 49, - "endColumn": 61, + "endColumn": 64, "lineCount": 1 } },