Skip to content

L2P: target boxes parallelization and shared coefficients #314

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ShawnL00 wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

L2P: target boxes parallelization and shared coefficients #314

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
aa40ff8
wip: parallelize L2P over target boxes
ShawnL00 Oct 17, 2026
0923620
l2p: tgt box parallelization
ShawnL00 Oct 18, 2026
cb67c37
Merge branch 'inducer:main' into l2p
ShawnL00 Oct 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 63 additions & 6 deletions sumpy/e2p.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,20 @@
# {{{ E2PFromSingleBox: E2P to single box (L2P, likely)

class E2PFromSingleBox(E2PBase):
"""Evaluate one expansion at each target in its box.

:arg work_items_per_group: a positive OpenCL workgroup size supported by
the device. Each workgroup handles one target box, with targets
distributed across its workitems.
"""

def __init__(self, expansion, kernels, name=None, *, work_items_per_group=None):

Check warning on line 131 in sumpy/e2p.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Type annotation is missing for parameter "work_items_per_group" (reportMissingParameterType)

Check warning on line 131 in sumpy/e2p.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Type of parameter "work_items_per_group" is partially unknown
  Parameter type is "Unknown | None" (reportUnknownParameterType)

Check warning on line 131 in sumpy/e2p.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Type annotation is missing for parameter "name" (reportMissingParameterType)

Check warning on line 131 in sumpy/e2p.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Type of parameter "name" is partially unknown
  Parameter type is "Unknown | None" (reportUnknownParameterType)

Check warning on line 131 in sumpy/e2p.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Type annotation is missing for parameter "kernels" (reportMissingParameterType)

Check warning on line 131 in sumpy/e2p.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Type of parameter "kernels" is unknown (reportUnknownParameterType)

Check warning on line 131 in sumpy/e2p.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Type annotation is missing for parameter "expansion" (reportMissingParameterType)

Check warning on line 131 in sumpy/e2p.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Type of parameter "expansion" is unknown (reportUnknownParameterType)
super().__init__(expansion, kernels, name)

Check warning on line 132 in sumpy/e2p.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Argument type is unknown
  Argument corresponds to parameter "expansion" in function "__init__" (reportUnknownArgumentType)

Check warning on line 132 in sumpy/e2p.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Type of "__init__" is partially unknown
  Type of "__init__" is "(expansion: Unknown, kernels: Unknown, name: Unknown | None = None) -> None" (reportUnknownMemberType)
self.work_items_per_group = work_items_per_group

def get_cache_key(self):

Check failure on line 135 in sumpy/e2p.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Method "get_cache_key" overrides class "E2PBase" in an incompatible manner
  Return type mismatch: base method returns type "tuple[str, Unknown, tuple[ScalarKernel, ...]]", override returns type "tuple[str, Unknown, tuple[ScalarKernel, ...], Unknown | None]"
    "tuple[str, Unknown, tuple[ScalarKernel, ...], Unknown | None]" is not assignable to "tuple[str, Unknown, tuple[ScalarKernel, ...]]"
      Tuple size mismatch; expected 3 but received 4 (reportIncompatibleMethodOverride)
return (*super().get_cache_key(), self.work_items_per_group)

@property
def default_name(self):
return "e2p_from_single_box"
Expand All @@ -132,24 +146,28 @@
loopy_knl = make_loopy_program(
[
"{[itgt_box]: 0<=itgt_box<ntgt_boxes}",
"{[itgt,idim]: itgt_start<=itgt<itgt_end and 0<=idim<dim}",
"{[itgt_offset]: 0<=itgt_offset<itgt_count}",
"{[idim]: 0<=idim<dim}",
"{[icoeff]: 0<=icoeff<ncoeffs}",
"{[iknl]: 0<=iknl<nresults}",
], [
*self.get_kernel_scaling_assignment(),
"""
for itgt_box
<> tgt_ibox = target_boxes[itgt_box]
<> itgt_start = box_target_starts[tgt_ibox]
<> itgt_end = itgt_start+box_target_counts_nonchild[tgt_ibox]
<> tgt_ibox = target_boxes[itgt_box] {id=fetch_tgt_box_0}
<> itgt_start = box_target_starts[tgt_ibox] \
{id=fetch_tgt_box_1}
<> itgt_count = box_target_counts_nonchild[tgt_ibox] \
{id=fetch_tgt_box_2}

<> center[idim] = centers[idim, tgt_ibox] {id=fetch_center}

<> coeffs[icoeff] = \
src_expansions[tgt_ibox - src_base_ibox, icoeff] \
{id=fetch_coeffs}

for itgt
for itgt_offset
<> itgt = itgt_start + itgt_offset
<> tgt[idim] = targets[idim, itgt] {id=fetch_tgt,dup=idim}
<> result_temp[iknl] = 0 {id=init_result,dup=iknl}
[iknl]: result_temp[iknl] = e2p(
Expand Down Expand Up @@ -196,8 +214,47 @@
return self.add_loopy_eval_callable(loopy_knl)

def get_optimized_kernel(self):
# FIXME
"""
Adapted from Isuru Fernando's https://github.com/inducer/sumpy/pull/158.
Unlike #158, the kernel does not use `max_ntargets_in_one_box` to determine
the workgroup size.
"""
knl = self.get_kernel()
if self.work_items_per_group is not None:
import namedisl as nisl

knl = lp.split_iname(knl, "itgt_offset", self.work_items_per_group,
outer_iname="itgt_outer", inner_iname="iwork_item", inner_tag="l.0")

# [itgt_count] -> {[itgt_outer, iwork_item]: 0 <= itgt_outer and
# 0 <= iwork_item < W and W*itgt_outer + iwork_item < itgt_count}
#
# Loopy does not allow a parallel iname in such a domain, so we give
# iwork_item its own domain and make it a parameter of itgt_outer's.
entry = knl.default_entrypoint
idom = entry.get_home_domain_index("iwork_item")
domains = list(entry.domains)
domains[idom:idom + 1] = [
nisl.make_set(
f"{{[iwork_item]: 0<=iwork_item<{self.work_items_per_group}}}"),
nisl.make_set(
"[iwork_item, itgt_count] -> {[itgt_outer]: "
"0 <= itgt_outer and "
f"{self.work_items_per_group}*itgt_outer"
" + iwork_item < itgt_count}"),
]
knl = knl.with_kernel(entry.copy(domains=domains))

knl = lp.add_inames_to_insn(knl, "iwork_item",
"id:fetch_tgt_box_* or id:fetch_center or id:kernel_scaling")

# The coefficient load is spread over the workitems.
knl = lp.split_iname(knl, "icoeff", self.work_items_per_group,
outer_iname="icoeff_outer", inner_iname="icoeff_inner",
inner_tag="l.0")
knl = lp.set_temporary_address_space(
knl, "coeffs", lp.AddressSpace.LOCAL)

knl = lp.tag_inames(knl, {"itgt_box": "g.0"})
knl = lp.add_inames_to_insn(knl, "itgt_box", "id:kernel_scaling")
return lp.set_options(knl,
Expand Down
10 changes: 9 additions & 1 deletion sumpy/fmm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,7 @@ class SumpyTreeIndependentDataForWrangler(TreeIndependentDataForWrangler):
use_rscale: bool | None
strength_usage: Sequence[int] | None
fft_backend: FFTBackend | None
l2p_work_items_per_group: int | None

def __init__(self,
array_context: ArrayContext,
Expand All @@ -141,6 +142,7 @@ def __init__(self,
strength_usage: Sequence[int] | None = None,
source_kernels: Sequence[ScalarKernel] | None = None,
fft_backend: FFTBackend | None = None,
l2p_work_items_per_group: int | None = 64,
) -> None:
"""
:arg multipole_expansion_factory: a callable of a single argument (order)
Expand All @@ -154,6 +156,9 @@ def __init__(self,
:arg fft_backend: the FFT backend used for compressing M2L interactions.
If not provided, it will be determined at runtime for the given
array context.
:arg l2p_work_items_per_group: L2P workgroup size, with one group per
target box. The default workgroup size is 64. Pass *None* to disable
parallelization and keep target processing serial within the box.
"""
super().__init__()

Expand All @@ -171,6 +176,8 @@ def __init__(self,
self.strength_usage = strength_usage
self.fft_backend = fft_backend

self.l2p_work_items_per_group = l2p_work_items_per_group

@memoize_method
def get_base_kernel(self) -> ScalarKernel:
from pytools import single_valued
Expand Down Expand Up @@ -263,7 +270,8 @@ def m2p(self, src_order: int) -> E2PBase:
def l2p(self, src_order: int) -> E2PBase:
return E2PFromSingleBox(
self.local_expansion(src_order),
self.target_kernels, name="l2p")
self.target_kernels, name="l2p",
work_items_per_group=self.l2p_work_items_per_group)

@memoize_method
def p2p(self) -> P2PBase:
Expand Down
Loading