tree_d_3d.kernels

The Barnes-Hut and FMM kernels: one folder per kernel.

Every kernel <name> lives in its own folder::

tree_d_3d/kernels/<name>/__init__.py        kernel = xp.kernels.Kernel.from_folder(...)
tree_d_3d/kernels/<name>/<name>_pyccel.py   CPU, compiled with Pyccel on first use
tree_d_3d/kernels/<name>/<name>_numpy.py    CPU, when Pyccel is unavailable
tree_d_3d/kernels/<name>/<name>_cuda.cu     GPU (CuPy), compiled with NVRTC
tree_d_3d/kernels/<name>/<name>_metal.py    Apple GPU (MLX), float32

All versions take the same arguments; a kernel runs the CUDA version on CuPy arrays and the Pyccel version (or NumPy without Pyccel) on NumPy arrays. Importing this package needs neither Pyccel nor CuPy.

Attributes

Name	Description
KERNEL_NAMES	No description.
METAL_KERNEL_NAMES	No description.

Functions

Name	Description
kernels_compiled	Whether the kernels run compiled.
metal_kernel	Return the Metal version of kernel name: a function with the Pyccel arguments.

Modules

• compileCompile tree-d-3d's Pyccel kernel modules once and reuse the builds across processes.
• direct_log_sumThe direct_log_sum kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• direct_sumThe direct_sum kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• evaluate_log_quadtreeThe evaluate_log_quadtree kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• evaluate_quadtreeThe evaluate_quadtree kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• fmm_downwardThe fmm_downward kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• fmm_evaluateThe fmm_evaluate kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• fmm_m2lThe fmm_m2l kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• fmm_traverseThe fmm_traverse kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• log_quadtree_momentsThe log_quadtree_moments kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• pair_energyThe pair_energy kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• quadtree_momentsThe quadtree_moments kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• quadtree_quadrantsThe quadtree_quadrants kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• quadtree_scatterThe quadtree_scatter kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• soften_pairsThe soften_pairs kernel of the Barnes-Hut and FMM solvers: Pyccel, NumPy and CUDA versions
• tests

KERNEL_NAMESattributemodule attribute#

View source

KERNEL_NAMES = (
  'quadtree_quadrants',
  'quadtree_scatter',
  'quadtree_moments',
  'evaluate_quadtree',
  'direct_sum',
  'pair_energy',
  'log_quadtree_moments',
  'evaluate_log_quadtree',
  'direct_log_sum',
  'fmm_traverse',
  'fmm_m2l',
  'fmm_downward',
  'fmm_evaluate',
  'soften_pairs'
)

Show moreShow less

METAL_KERNEL_NAMESattributemodule attribute#

View source

METAL_KERNEL_NAMES = KERNEL_NAMES

kernels_compiledfunction#

View source

def kernels_compiled() -> bool

Whether the kernels run compiled.

On the CuPy backend they always do (CUDA); otherwise this checks once per process whether the Pyccel versions compile. Without them the tree kernels run as Python loops, so method="auto" sums directly.

metal_kernelfunction#

View source

def metal_kernel(name: str) -> Callable[..., None]

Return the Metal version of kernel name: a function with the Pyccel arguments.

Needs MLX and a Metal GPU (pip install 'tree-d-3d[metal]'); the arrays are NumPy arrays, computed in float32 and written back into the output arrays.

Raises

KeyError

If the kernel has no Metal version.

RuntimeError

If there is no usable Metal GPU.