diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d8042a..acac314 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,13 @@ format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ## [Unreleased] +### Changed + +- The kernels are built with cunumpy's `KernelPackage` (`mpifield.kernels.KERNEL_PACKAGE`): + `kernels/_build.py`, `kernels/_metal.py` and `kernels/compile.py` are gone, and the build + cache is still `MPIFIELD_KERNEL_CACHE`. `shared_memory_fits` is + `xp.cuda.shared_memory_fits`; the catalog no longer exports it. + ### Added - `Field`, moved out of picnic (`picnic.field`): a field on the nodes of a uniform 1D, 2D or diff --git a/src/mpifield/kernels/__init__.py b/src/mpifield/kernels/__init__.py index e169518..eb0122d 100644 --- a/src/mpifield/kernels/__init__.py +++ b/src/mpifield/kernels/__init__.py @@ -1 +1,23 @@ -"""Compiled linear gather and scatter kernels (see :mod:`mpifield.kernels.catalog`).""" +"""mpifield's compiled kernels, one folder per kernel (see :mod:`mpifield.kernels.catalog`). + +Every kernel folder builds its kernel with :data:`KERNEL_PACKAGE`:: + + kernel = KERNEL_PACKAGE.from_folder(__name__) + +with the defaults of :class:`cunumpy.kernels.KernelPackage`: the Pyccel version is +compiled on its first host call and kept under ``MPIFIELD_KERNEL_CACHE`` (default +``~/.cache/mpifield/kernels``; ``MPIFIELD_OPENMP=0`` builds serial), the CUDA version runs +for CuPy arguments and the host version for NumPy ones, one CUDA thread per entry of +the last axis of the first array. Importing a kernel needs neither Pyccel nor CuPy. + +The Metal versions (``/_metal.py``) wrap a kernel built by +``KERNEL_PACKAGE.metal_kernel`` in a function with the arguments of the Pyccel +version. The Apple GPU has no float64, so they compute in float32: expect a relative +error of about 1e-5 against the CPU versions. Import one only where +:func:`cunumpy.kernels.metal_available` is true. +""" + +import cunumpy as xp + +#: The build cache and launch options of mpifield's kernels. +KERNEL_PACKAGE = xp.kernels.KernelPackage("mpifield") diff --git a/src/mpifield/kernels/_build.py b/src/mpifield/kernels/_build.py deleted file mode 100644 index c3f1ffb..0000000 --- a/src/mpifield/kernels/_build.py +++ /dev/null @@ -1,51 +0,0 @@ -"""The options every kernel folder passes to :meth:`cunumpy.kernels.Kernel.from_folder`. - -The ``__init__.py`` of a kernel folder ``mpifield/kernels//`` declares its -kernel with these:: - - kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) - -which finds ``_pyccel.py``, ``_numpy.py`` and ``_cuda.cu`` in -the folder. -Nothing is compiled here: the Pyccel version is compiled on its first host call -(cached on disk by :func:`~mpifield.kernels.compile.compile_cached`), the CUDA -version on its first launch. Importing this module needs neither Pyccel nor CuPy. -""" - -from __future__ import annotations - -from typing import Any - -#: CUDA threads per block of every kernel. -BLOCK_SIZE = 256 - -#: Default dynamic shared memory per block of a CUDA device, in bytes. -SHARED_MEMORY_PER_BLOCK = 48 * 1024 - - -def shared_memory_fits(num_values: int, itemsize: int = 8) -> bool: - """Whether ``num_values`` values fit the default shared memory of a block.""" - return num_values * itemsize <= SHARED_MEMORY_PER_BLOCK - - -def _compile(module: Any) -> Any: - from mpifield.kernels.compile import compile_cached - - return compile_cached(module) - - -#: ``_pyccel.py`` compiled with Pyccel on first use; CUDA for CuPy -#: arguments, the host version for NumPy ones (also while CuPy is active); -#: without ``n_threads`` one CUDA thread per entry of the last axis of the -#: first array (``"last_axis"``: the marker count of a component-major -#: ``(ncomp, N)`` or ``(N,)`` marker array). A kernel whose first array is not -#: a per-marker array passes ``n_threads`` at its call sites. CUDA sources -#: include mpifield headers relative to the source root (cunumpy's default), e.g. -#: ``#include "mpifield/kernels/linear_cell.cuh"``. -KERNEL_OPTIONS: dict[str, Any] = { - "host_suffix": "_pyccel", - "compile_host": _compile, - "dispatch": "arrays", - "block_size": BLOCK_SIZE, - "n_threads_from": "last_axis", -} diff --git a/src/mpifield/kernels/_metal.py b/src/mpifield/kernels/_metal.py deleted file mode 100644 index 10360db..0000000 --- a/src/mpifield/kernels/_metal.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Helpers of the Metal versions (``/_metal.py``) of the kernels. - -A Metal version is the body of a Metal Shading Language kernel run by -:class:`cunumpy.kernels.MetalKernel` on the GPU of an Apple-silicon Mac, wrapped in -a function with the arguments of the Pyccel version: NumPy arrays in, results -written into the output arrays. The GPU has no float64, so the kernels compute in -float32 (``float64="cast"``): expect a relative error of about 1e-5 against the -CPU versions. Import a version only where :func:`metal_available` is true. -""" - -from __future__ import annotations - -import cunumpy as xp - -#: Threads per threadgroup of every Metal kernel. -THREADGROUP = 256 - - -def metal_available() -> bool: - """Whether the Metal versions can run here (MLX installed, Metal GPU present).""" - return bool(xp.kernels.metal_available()) - - -def make_kernel( - name: str, - source: str, - inputs: list[str], - outputs: list[str], - header: str = "", -) -> xp.kernels.MetalKernel: - """A :class:`~cunumpy.kernels.MetalKernel` computing in float32, outputs zeroed.""" - return xp.kernels.MetalKernel( - source, - inputs=inputs, - outputs=outputs, - name=name, - header=header, - threadgroup=THREADGROUP, - float64="cast", - init_value=0, - ) diff --git a/src/mpifield/kernels/_metal_grid.py b/src/mpifield/kernels/_metal_grid.py index 9271f87..ee8a93f 100644 --- a/src/mpifield/kernels/_metal_grid.py +++ b/src/mpifield/kernels/_metal_grid.py @@ -6,7 +6,7 @@ ``quadratic_cell.cuh`` for the quadratic ones). A scatter accumulates into a zeroed float32 grid with float atomics (no 64-bit atomics on the Apple GPU), which the host adds to the caller's grid in float64: the summation order, hence the -last bits, differ from the CPU. See :mod:`mpifield.kernels._metal`. +last bits, differ from the CPU. See :mod:`mpifield.kernels`. """ from __future__ import annotations @@ -15,12 +15,11 @@ import numpy as np from mpifield._types import Array -from mpifield.kernels._metal import THREADGROUP, make_kernel +from mpifield.kernels import KERNEL_PACKAGE __all__ = [ "CELL_HEADER", "QUADRATIC_HEADER", - "make_kernel", "make_scatter_kernel", "relative", "scatter", @@ -94,16 +93,8 @@ def make_scatter_kernel( name: str, source: str, inputs: list[str], header: str = "" ) -> xp.kernels.MetalKernel: """A kernel adding into the zeroed float32 output ``delta`` with float atomics.""" - return xp.kernels.MetalKernel( - source, - inputs=inputs, - outputs=["delta"], - name=name, - header=header, - threadgroup=THREADGROUP, - float64="cast", - atomic_outputs=True, - init_value=0, + return KERNEL_PACKAGE.metal_kernel( + name, source, inputs, ["delta"], header=header, atomic_outputs=True ) diff --git a/src/mpifield/kernels/catalog.py b/src/mpifield/kernels/catalog.py index 2df4e07..180d011 100644 --- a/src/mpifield/kernels/catalog.py +++ b/src/mpifield/kernels/catalog.py @@ -2,7 +2,7 @@ Every kernel ```` lives in its own folder:: - mpifield/kernels//__init__.py kernel = xp.kernels.Kernel.from_folder(...) + mpifield/kernels//__init__.py kernel = KERNEL_PACKAGE.from_folder(__name__) mpifield/kernels//_pyccel.py def (...) CPU, compiled with Pyccel mpifield/kernels//_cuda.cu __global__ void (...) GPU, compiled with NVRTC mpifield/kernels//_numpy.py def (...) CPU, when Pyccel is unavailable @@ -17,10 +17,10 @@ runs the CUDA version when the arrays live on the GPU (CuPy) and the compiled Pyccel version otherwise (``dispatch="arrays"``). The kernel object is a :class:`cunumpy.kernels.Kernel`, built in the folder's ``__init__.py`` by -:meth:`cunumpy.kernels.Kernel.from_folder` with the options of -:mod:`mpifield.kernels._build`. Both versions are compiled on first use: importing -this module needs neither Pyccel nor CuPy. CUDA sources include mpifield's header -relative to the source root: ``#include "mpifield/kernels/linear_cell.cuh"``. +``KERNEL_PACKAGE.from_folder`` (see :mod:`mpifield.kernels`). Both versions are +compiled on first use: importing this module needs neither Pyccel nor CuPy. CUDA +sources include mpifield's header relative to the source root: +``#include "mpifield/kernels/linear_cell.cuh"``. The kernels are the linear (``*_linear_*``) and quadratic (``*_quadratic_*``) deposits and gathers of :mod:`mpifield.weighting`. The quadratic ones, and the 1D and 2D linear @@ -39,23 +39,9 @@ from collections.abc import Callable import cunumpy as xp +from cunumpy.kernels import metal_available -from mpifield.kernels._build import ( - BLOCK_SIZE, - SHARED_MEMORY_PER_BLOCK, - shared_memory_fits, -) -from mpifield.kernels._metal import metal_available - -__all__ = [ - "BLOCK_SIZE", - "KERNELS", - "KERNEL_NAMES", - "METAL_KERNEL_NAMES", - "SHARED_MEMORY_PER_BLOCK", - "metal_kernel", - "shared_memory_fits", -] +__all__ = ["KERNELS", "KERNEL_NAMES", "METAL_KERNEL_NAMES", "metal_kernel"] #: The kernels, by name. KERNEL_NAMES = ( diff --git a/src/mpifield/kernels/compile.py b/src/mpifield/kernels/compile.py deleted file mode 100644 index 5527af2..0000000 --- a/src/mpifield/kernels/compile.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Compile mpifield's Pyccel kernel modules once and reuse the builds across processes. - -The cache is :class:`cunumpy.kernels.PyccelCache`: a kernel module is built with -``pyccel compile`` on its first use, keyed on its source, the Pyccel and Python -versions and the platform, and stored under ``MPIFIELD_KERNEL_CACHE`` (default -``~/.cache/mpifield/kernels``, or ``/mpifield-kernels`` if that is not writable). -Later processes import the stored extension directly. - -Kernels are built with OpenMP, so the loops marked ``#$ omp parallel for`` run on -all cores (limit them with ``OMP_NUM_THREADS``; under MPI set it to the cores per -rank). ``MPIFIELD_OPENMP=0`` builds them serial, except the kernels that call the OpenMP -runtime; the setting is part of the cache key. -""" - -from __future__ import annotations - -import cunumpy as xp - -#: The on-disk build cache of mpifield's kernels. -KERNEL_CACHE = xp.kernels.PyccelCache("mpifield") - -#: Return the compiled form of a kernel module, building it on first use -#: (:meth:`cunumpy.kernels.PyccelCache.compile`). -compile_cached = KERNEL_CACHE.compile diff --git a/src/mpifield/kernels/gather_linear_1d/__init__.py b/src/mpifield/kernels/gather_linear_1d/__init__.py index 1a65084..ef8b8a5 100644 --- a/src/mpifield/kernels/gather_linear_1d/__init__.py +++ b/src/mpifield/kernels/gather_linear_1d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from mpifield.kernels import KERNEL_PACKAGE -from mpifield.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/mpifield/kernels/gather_linear_1d/gather_linear_1d_metal.py b/src/mpifield/kernels/gather_linear_1d/gather_linear_1d_metal.py index d33d4ab..7fc2b69 100644 --- a/src/mpifield/kernels/gather_linear_1d/gather_linear_1d_metal.py +++ b/src/mpifield/kernels/gather_linear_1d/gather_linear_1d_metal.py @@ -9,7 +9,8 @@ import numpy as np from mpifield._types import Array -from mpifield.kernels._metal_grid import CELL_HEADER, make_kernel, relative +from mpifield.kernels import KERNEL_PACKAGE +from mpifield.kernels._metal_grid import CELL_HEADER, relative _SOURCE = """ uint marker = thread_position_in_grid.x; @@ -25,7 +26,7 @@ } """ -_KERNEL = make_kernel( +_KERNEL = KERNEL_PACKAGE.metal_kernel( "mpifield_gather_linear_1d", _SOURCE, inputs=["rel", "field", "grid"], diff --git a/src/mpifield/kernels/gather_linear_2d/__init__.py b/src/mpifield/kernels/gather_linear_2d/__init__.py index fd7a22c..ee3822d 100644 --- a/src/mpifield/kernels/gather_linear_2d/__init__.py +++ b/src/mpifield/kernels/gather_linear_2d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from mpifield.kernels import KERNEL_PACKAGE -from mpifield.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/mpifield/kernels/gather_linear_2d/gather_linear_2d_metal.py b/src/mpifield/kernels/gather_linear_2d/gather_linear_2d_metal.py index 64f8d75..aa959bf 100644 --- a/src/mpifield/kernels/gather_linear_2d/gather_linear_2d_metal.py +++ b/src/mpifield/kernels/gather_linear_2d/gather_linear_2d_metal.py @@ -7,7 +7,8 @@ from __future__ import annotations from mpifield._types import Array -from mpifield.kernels._metal_grid import CELL_HEADER, make_kernel, relative +from mpifield.kernels import KERNEL_PACKAGE +from mpifield.kernels._metal_grid import CELL_HEADER, relative _SOURCE = """ uint marker = thread_position_in_grid.x; @@ -31,7 +32,7 @@ } """ -_KERNEL = make_kernel( +_KERNEL = KERNEL_PACKAGE.metal_kernel( "mpifield_gather_linear_2d", _SOURCE, inputs=["rel", "field", "grid"], diff --git a/src/mpifield/kernels/gather_linear_3d/__init__.py b/src/mpifield/kernels/gather_linear_3d/__init__.py index d9d7321..506fcd5 100644 --- a/src/mpifield/kernels/gather_linear_3d/__init__.py +++ b/src/mpifield/kernels/gather_linear_3d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from mpifield.kernels import KERNEL_PACKAGE -from mpifield.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/mpifield/kernels/gather_linear_3d/gather_linear_3d_metal.py b/src/mpifield/kernels/gather_linear_3d/gather_linear_3d_metal.py index 1ac5c09..50470c8 100644 --- a/src/mpifield/kernels/gather_linear_3d/gather_linear_3d_metal.py +++ b/src/mpifield/kernels/gather_linear_3d/gather_linear_3d_metal.py @@ -7,7 +7,8 @@ from __future__ import annotations from mpifield._types import Array -from mpifield.kernels._metal_grid import CELL_HEADER, make_kernel, relative +from mpifield.kernels import KERNEL_PACKAGE +from mpifield.kernels._metal_grid import CELL_HEADER, relative _SOURCE = """ uint marker = thread_position_in_grid.x; @@ -34,7 +35,7 @@ } """ -_KERNEL = make_kernel( +_KERNEL = KERNEL_PACKAGE.metal_kernel( "mpifield_gather_linear_3d", _SOURCE, inputs=["rel", "field", "grid"], diff --git a/src/mpifield/kernels/gather_quadratic_1d/__init__.py b/src/mpifield/kernels/gather_quadratic_1d/__init__.py index 5d3b43f..c4bfc4a 100644 --- a/src/mpifield/kernels/gather_quadratic_1d/__init__.py +++ b/src/mpifield/kernels/gather_quadratic_1d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from mpifield.kernels import KERNEL_PACKAGE -from mpifield.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_metal.py b/src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_metal.py index e823185..8da1ee4 100644 --- a/src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_metal.py +++ b/src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_metal.py @@ -8,7 +8,8 @@ from __future__ import annotations from mpifield._types import Array -from mpifield.kernels._metal_grid import QUADRATIC_HEADER, make_kernel, relative +from mpifield.kernels import KERNEL_PACKAGE +from mpifield.kernels._metal_grid import QUADRATIC_HEADER, relative _SOURCE = """ uint marker = thread_position_in_grid.x; @@ -26,7 +27,7 @@ } """ -_KERNEL = make_kernel( +_KERNEL = KERNEL_PACKAGE.metal_kernel( "mpifield_gather_quadratic_1d", _SOURCE, inputs=["rel", "field", "grid"], diff --git a/src/mpifield/kernels/gather_quadratic_2d/__init__.py b/src/mpifield/kernels/gather_quadratic_2d/__init__.py index e08d6bf..2f30c5f 100644 --- a/src/mpifield/kernels/gather_quadratic_2d/__init__.py +++ b/src/mpifield/kernels/gather_quadratic_2d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from mpifield.kernels import KERNEL_PACKAGE -from mpifield.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_metal.py b/src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_metal.py index ce8f443..c59c834 100644 --- a/src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_metal.py +++ b/src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_metal.py @@ -8,7 +8,8 @@ from __future__ import annotations from mpifield._types import Array -from mpifield.kernels._metal_grid import QUADRATIC_HEADER, make_kernel, relative +from mpifield.kernels import KERNEL_PACKAGE +from mpifield.kernels._metal_grid import QUADRATIC_HEADER, relative _SOURCE = """ uint marker = thread_position_in_grid.x; @@ -31,7 +32,7 @@ } """ -_KERNEL = make_kernel( +_KERNEL = KERNEL_PACKAGE.metal_kernel( "mpifield_gather_quadratic_2d", _SOURCE, inputs=["rel", "field", "grid"], diff --git a/src/mpifield/kernels/gather_quadratic_3d/__init__.py b/src/mpifield/kernels/gather_quadratic_3d/__init__.py index e69285a..f5e8cab 100644 --- a/src/mpifield/kernels/gather_quadratic_3d/__init__.py +++ b/src/mpifield/kernels/gather_quadratic_3d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from mpifield.kernels import KERNEL_PACKAGE -from mpifield.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/mpifield/kernels/gather_quadratic_3d/gather_quadratic_3d_metal.py b/src/mpifield/kernels/gather_quadratic_3d/gather_quadratic_3d_metal.py index d86389a..bd67f3c 100644 --- a/src/mpifield/kernels/gather_quadratic_3d/gather_quadratic_3d_metal.py +++ b/src/mpifield/kernels/gather_quadratic_3d/gather_quadratic_3d_metal.py @@ -8,7 +8,8 @@ from __future__ import annotations from mpifield._types import Array -from mpifield.kernels._metal_grid import QUADRATIC_HEADER, make_kernel, relative +from mpifield.kernels import KERNEL_PACKAGE +from mpifield.kernels._metal_grid import QUADRATIC_HEADER, relative _SOURCE = """ uint marker = thread_position_in_grid.x; @@ -35,7 +36,7 @@ } """ -_KERNEL = make_kernel( +_KERNEL = KERNEL_PACKAGE.metal_kernel( "mpifield_gather_quadratic_3d", _SOURCE, inputs=["rel", "field", "grid"], diff --git a/src/mpifield/kernels/scatter_linear_1d/__init__.py b/src/mpifield/kernels/scatter_linear_1d/__init__.py index 8457341..4dbe25c 100644 --- a/src/mpifield/kernels/scatter_linear_1d/__init__.py +++ b/src/mpifield/kernels/scatter_linear_1d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from mpifield.kernels import KERNEL_PACKAGE -from mpifield.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/mpifield/kernels/scatter_linear_2d/__init__.py b/src/mpifield/kernels/scatter_linear_2d/__init__.py index 29db8f9..1f5f6c1 100644 --- a/src/mpifield/kernels/scatter_linear_2d/__init__.py +++ b/src/mpifield/kernels/scatter_linear_2d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from mpifield.kernels import KERNEL_PACKAGE -from mpifield.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/mpifield/kernels/scatter_linear_3d/__init__.py b/src/mpifield/kernels/scatter_linear_3d/__init__.py index 6ece41e..6d0e547 100644 --- a/src/mpifield/kernels/scatter_linear_3d/__init__.py +++ b/src/mpifield/kernels/scatter_linear_3d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from mpifield.kernels import KERNEL_PACKAGE -from mpifield.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/mpifield/kernels/scatter_quadratic_1d/__init__.py b/src/mpifield/kernels/scatter_quadratic_1d/__init__.py index 1e76370..18a4146 100644 --- a/src/mpifield/kernels/scatter_quadratic_1d/__init__.py +++ b/src/mpifield/kernels/scatter_quadratic_1d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from mpifield.kernels import KERNEL_PACKAGE -from mpifield.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/mpifield/kernels/scatter_quadratic_2d/__init__.py b/src/mpifield/kernels/scatter_quadratic_2d/__init__.py index 5c7dc95..263a85e 100644 --- a/src/mpifield/kernels/scatter_quadratic_2d/__init__.py +++ b/src/mpifield/kernels/scatter_quadratic_2d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from mpifield.kernels import KERNEL_PACKAGE -from mpifield.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/mpifield/kernels/scatter_quadratic_3d/__init__.py b/src/mpifield/kernels/scatter_quadratic_3d/__init__.py index fa08416..e53e47a 100644 --- a/src/mpifield/kernels/scatter_quadratic_3d/__init__.py +++ b/src/mpifield/kernels/scatter_quadratic_3d/__init__.py @@ -5,8 +5,6 @@ (see ``cunumpy.kernels.set_host_kernel_implementation`` to choose). """ -import cunumpy as xp +from mpifield.kernels import KERNEL_PACKAGE -from mpifield.kernels._build import KERNEL_OPTIONS - -kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS) +kernel = KERNEL_PACKAGE.from_folder(__name__) diff --git a/src/mpifield/tests/kernels/test_metal_grid_kernels.py b/src/mpifield/tests/kernels/test_metal_grid_kernels.py index 2f98d18..e539dba 100644 --- a/src/mpifield/tests/kernels/test_metal_grid_kernels.py +++ b/src/mpifield/tests/kernels/test_metal_grid_kernels.py @@ -11,8 +11,8 @@ import numpy as np import pytest +from cunumpy.kernels import metal_available -from mpifield.kernels._metal import metal_available from mpifield.kernels.catalog import KERNELS pytestmark = pytest.mark.skipif( diff --git a/src/mpifield/weighting.py b/src/mpifield/weighting.py index 42f60d6..7331606 100644 --- a/src/mpifield/weighting.py +++ b/src/mpifield/weighting.py @@ -44,7 +44,6 @@ import cunumpy as xp from mpifield._types import Array, BlockGrid -from mpifield.kernels.catalog import shared_memory_fits from mpifield.kernels.gather_linear_1d import kernel as gather_linear_1d from mpifield.kernels.gather_linear_2d import kernel as gather_linear_2d from mpifield.kernels.gather_linear_3d import kernel as gather_linear_3d @@ -183,7 +182,7 @@ def _scatter_linear_1d( """ args = (lower_x, cell_size_x, num_x, periodic_x) gpu = xp.is_gpu(field) - if gpu and not shared_memory_fits(num_x): + if gpu and not xp.cuda.shared_memory_fits(num_x): for component in range(field.shape[1]): scatter_linear_1d_numpy( positions[0], values[component], field[:, component], *args diff --git a/uv.lock b/uv.lock index 9fbd325..85a2a02 100644 --- a/uv.lock +++ b/uv.lock @@ -920,7 +920,7 @@ wheels = [ [[package]] name = "cunumpy" version = "0.7.1" -source = { git = "https://github.com/max-models/cunumpy?branch=devel#18c7c2fbef68a73a32029981662213b73f1d2a21" } +source = { git = "https://github.com/max-models/cunumpy?branch=devel#92b38a5a0e00d7ccf03aa39849d6c7eaee946859" } dependencies = [ { name = "array-api-compat" }, { name = "maybempi" },