Skip to content

M2P: target boxes parallelization and shared coefficients #313

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 5 commits into
base: main
Choose a base branch
from
Draft

M2P: target boxes parallelization and shared coefficients #313

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
5ca71c7
wip: parallelize E2P over targets within each box
ShawnL00 Oct 16, 2026
5d170a3
Share M2P coefficients within workgroups
ShawnL00 Oct 17, 2026
8105bec
add FMM support
ShawnL00 Oct 17, 2026
3d805d7
Merge branch 'inducer:main' into m2p
ShawnL00 Oct 18, 2026
598fd61
fix docstring
ShawnL00 Oct 19, 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
239 changes: 230 additions & 9 deletions sumpy/e2p.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,25 @@
# {{{ E2PFromCSR: E2P from CSR-like interaction list

class E2PFromCSR(E2PBase):
"""Evaluate CSR-listed expansions at each target box's particles.

:arg work_items_per_group: OpenCL workgroup size, with one group per
target box and target particles distributed across its workitems.
*None* uses 128 workitems when ``max_ntargets_in_one_box``
is supplied to ``__call__``, and selects the serial schedule
otherwise.
"""

# M2P (List 3) target-box parallelization:
# Each (leaf) target box is a workgroup with W workitems.
# Workitem j handles target offsets j, j + W, j + 2*W, ... within its box.
def __init__(self, expansion, kernels, name=None, *, work_items_per_group=None):

Check warning on line 246 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 246 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 246 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 246 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 246 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 246 in sumpy/e2p.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Type of parameter "kernels" is unknown (reportUnknownParameterType)

Check warning on line 246 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 246 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 247 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 247 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 250 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_csr"
Expand All @@ -242,7 +261,7 @@
loopy_knl = make_loopy_program(
[
"{[itgt_box]: 0<=itgt_box<ntgt_boxes}",
"{[itgt]: itgt_start<=itgt<itgt_end}",
"{[itgt_offset]: 0<=itgt_offset<itgt_count}",
"{[isrc_box]: isrc_box_start<=isrc_box<isrc_box_end }",
"{[idim]: 0<=idim<dim}",
"{[icoeff]: 0<=icoeff<ncoeffs}",
Expand All @@ -251,11 +270,13 @@
*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}

for itgt
for itgt_offset
<> itgt = itgt_start + itgt_offset
<> tgt[idim] = targets[idim,itgt] {id=fetch_tgt,dup=idim}

<> isrc_box_start = source_box_starts[itgt_box]
Expand Down Expand Up @@ -315,15 +336,210 @@

loopy_knl = lp.tag_inames(loopy_knl, "idim*:unr")
loopy_knl = lp.tag_inames(loopy_knl, "iknl*:unr")
loopy_knl = lp.prioritize_loops(loopy_knl, "itgt_box,itgt,isrc_box")
loopy_knl = lp.prioritize_loops(loopy_knl, "itgt_box,itgt_offset,isrc_box")
loopy_knl = self.add_loopy_eval_callable(loopy_knl)
return lp.tag_array_axes(loopy_knl, "targets", "sep,C")

def get_optimized_kernel(self):
# FIXME
def get_shared_coefficient_kernel(self, max_ntargets_in_one_box):

@inducer inducer Oct 19, 2026

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd only like to maintain one of these kernels.

"""
- One workgroup handles one target box with W workitems.
For each source box, the workgroup loads the coefficients
into local memory once and shares them across its workitems. Loopy
inserts barriers to ensure the load finishes before evaluation and
all workitems finish using the coefficients before they are
overwritten for the next source box.

- Each round covers W target slots: in round ``itgt_outer``, workitem
``iwork_item`` handles offset ``W*itgt_outer + iwork_item`` within the
box. Each workitem keeps separate sums for each round, accumulating
contributions from all source boxes. After the last source box, it
adds those sums to the global result.

- All boxes run ``ntgt_outer`` = ceil(*max_ntargets_in_one_box* / W)
rounds. Slots beyond a box's actual target count are padding: guards
skip their target reads, evaluations, and result writes.
A zero maximum uses one guarded round for temporary-array sizing.

*max_ntargets_in_one_box* must be at least the target count of every box,
or the targets beyond it are skipped. The summation order per target is
that of :meth:`get_kernel`, so the results agree exactly.

Adapted from Isuru Fernando's https://github.com/inducer/sumpy/pull/180.
"""
ncoeffs = len(self.expansion)
work_items_per_group = self.work_items_per_group
if work_items_per_group is None:
work_items_per_group = 128
loopy_args = self.get_loopy_args()

loopy_knl = make_loopy_program(

@inducer inducer Oct 19, 2026

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See if you can make Loopy do the split?

[
"{[itgt_box]: 0<=itgt_box<ntgt_boxes}",
"{[iwork_item]: 0<=iwork_item<work_items_per_group}",
"{[itgt_outer]: 0<=itgt_outer<ntgt_outer}",
"{[isrc_box]: isrc_box_start<=isrc_box<isrc_box_end }",
"{[idim]: 0<=idim<dim}",
"{[icoeff]: 0<=icoeff<ncoeffs}",
"{[iknl]: 0<=iknl<nresults}",
], [
*self.get_kernel_scaling_assignment(),
"""
for itgt_box
for iwork_item
<> tgt_ibox = target_boxes[itgt_box] {id=fetch_tgt_box_0}
<> itgt_start = box_target_starts[tgt_ibox] \
{id=fetch_tgt_box_1}
<> itgt_end = itgt_start \
+ box_target_counts_nonchild[tgt_ibox] \
{id=fetch_tgt_box_2}
<> isrc_box_start = source_box_starts[itgt_box] \
{id=fetch_src_list_start}
<> isrc_box_end = source_box_starts[itgt_box+1] \
{id=fetch_src_list_end}
<> result_temp[itgt_outer, iknl] = 0 \
{id=init_result,dup=itgt_outer:iknl}
end

for isrc_box
for iwork_item
<> src_ibox = source_box_lists[isrc_box] \
{id=fetch_src_box}
<> center[idim] = centers[idim, src_ibox] \
{dup=idim,id=fetch_center}
end
# Loaded the coefficients into local memory once
# for all workitems in the group.
<> coeffs[icoeff] = \
src_expansions[src_ibox - src_base_ibox, icoeff] \
{id=fetch_coeffs}

for iwork_item, itgt_outer
<> itgt = itgt_start + work_items_per_group*itgt_outer \
+ iwork_item
<> cond_itgt = itgt < itgt_end
<> tgt[idim] = targets[idim,itgt] \
{id=fetch_tgt,dup=idim,if=cond_itgt}

[iknl]: result_temp[itgt_outer, iknl] = e2p(
[iknl]: result_temp[itgt_outer, iknl],
[icoeff]: coeffs[icoeff],
[idim]: center[idim],
[idim]: tgt[idim],
rscale,
itgt,
ntargets,
targets,
""" + ",".join(arg.name for arg in loopy_args) + """
) {id=update_result, \
dep=fetch_coeffs:fetch_center:fetch_tgt:init_result, \
if=cond_itgt}
end
end

# Workitems in the group write their results to global memory.
for iwork_item, itgt_outer
<> itgt_write = itgt_start \
+ work_items_per_group*itgt_outer + iwork_item \
{id=init_itgt_for_write}
<> cond_itgt_write = itgt_write < itgt_end \
{id=init_cond_for_write}
result[iknl, itgt_write] = result[iknl, itgt_write] \
+ result_temp[itgt_outer, iknl] * kernel_scaling \
{dep=update_result:init_result,id=write_result, \
dup=iknl,if=cond_itgt_write}
end
end
"""],
[
lp.GlobalArg("targets", None, shape=(self.dim, "ntargets")),
lp.GlobalArg("box_target_starts,box_target_counts_nonchild",
None, shape=None),
lp.GlobalArg("centers", None, shape="dim, aligned_nboxes"),
lp.GlobalArg("src_expansions", None,
shape=("nsrc_level_boxes", ncoeffs), offset=lp.auto),
lp.ValueArg("src_base_ibox", np.int32),
lp.ValueArg("nsrc_level_boxes,aligned_nboxes", np.int32),
lp.ValueArg("ntargets", np.int32),
lp.GlobalArg("result", None, shape="nresults,ntargets",
dim_tags="sep,C"),
lp.GlobalArg("source_box_starts, source_box_lists,",
None, shape=None, offset=lp.auto),
*loopy_args,
"..."
],
name=self.name,
assumptions="ntgt_boxes>=1",
silenced_warnings="write_race(*_result)",
fixed_parameters={
"ncoeffs": ncoeffs,
"dim": self.dim,
"work_items_per_group": work_items_per_group,
"ntgt_outer": max(1, -(
-max_ntargets_in_one_box // work_items_per_group)),
"nresults": len(self.kernels)})

loopy_knl = lp.tag_inames(loopy_knl, "idim*:unr")
loopy_knl = lp.tag_inames(loopy_knl, "iknl*:unr")
knl = self.add_loopy_eval_callable(loopy_knl)

# When writing to the global result, these three instructions must run
# in one loop over the same iname. Duplicate itgt_outer for all three
# together to keep this loop separate from the evaluation loop.
knl = lp.duplicate_inames(knl, "itgt_outer",
"id:write_result or id:init_itgt_for_write or id:init_cond_for_write")
knl = lp.tag_inames(knl, {"itgt_box": "g.0", "iwork_item": "l.0"})

# The coefficient load is spread over the workitems.
knl = lp.split_iname(knl, "icoeff", 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.add_inames_to_insn(knl, "itgt_box,iwork_item",
"id:kernel_scaling")

return lp.set_options(knl, enforce_variable_access_ordered="no_check")

def get_optimized_kernel(self, max_ntargets_in_one_box=None):
if max_ntargets_in_one_box is not None:
return self.get_shared_coefficient_kernel(max_ntargets_in_one_box)

knl = self.get_kernel()
if self.work_items_per_group is not None:
import namedisl as nisl

# Target-parallel E2P scheduling follows the idea in Isuru Fernando's
# https://github.com/inducer/sumpy/pull/180.
# Unlike #180, the coefficients are not loaded into local memory:
# each workitem reads every source box's coefficients itself and
# accumulates the sum over source boxes.

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

# Loopy does not allow a parallel iname (iwork_item) in a domain
# whose bound depends on data (itgt_count), so we split the domain.
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]: "
f"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.tag_inames(knl, {"itgt_box": "g.0"})
knl = lp.add_inames_to_insn(knl, "itgt_box", "id:kernel_scaling")
if self.work_items_per_group is not None:
knl = lp.add_inames_to_insn(knl, "iwork_item",
"id:fetch_tgt_box_* or id:kernel_scaling")
return lp.set_options(knl,
enforce_variable_access_ordered="no_check")

Expand All @@ -332,8 +548,13 @@
# "1" may be passed for rscale, which won't have its type
# meaningfully inferred. Make the type of rscale explicit.
rscale = centers.dtype.type(kwargs.pop("rscale"))
max_ntargets_in_one_box = kwargs.pop("max_ntargets_in_one_box", None)

knl = self.get_cached_kernel()
if max_ntargets_in_one_box is None:
knl = self.get_cached_kernel()
else:
knl = self.get_cached_kernel(
max_ntargets_in_one_box=int(max_ntargets_in_one_box))
result = actx.call_loopy(
knl,
centers=centers,
Expand Down
9 changes: 8 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 @@
use_rscale: bool | None
strength_usage: Sequence[int] | None
fft_backend: FFTBackend | None
m2p_work_items_per_group: int | None

def __init__(self,
array_context: ArrayContext,
Expand All @@ -141,6 +142,7 @@
strength_usage: Sequence[int] | None = None,
source_kernels: Sequence[ScalarKernel] | None = None,
fft_backend: FFTBackend | None = None,
m2p_work_items_per_group: int | None = None,
) -> None:
"""
:arg multipole_expansion_factory: a callable of a single argument (order)
Expand All @@ -154,6 +156,8 @@
: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 m2p_work_items_per_group: the OpenCL workgroup size of m2p,
passed unchanged to :class:`~sumpy.e2p.E2PFromCSR`.
"""
super().__init__()

Expand All @@ -170,6 +174,7 @@
self.use_rscale = use_rscale
self.strength_usage = strength_usage
self.fft_backend = fft_backend
self.m2p_work_items_per_group = m2p_work_items_per_group

@memoize_method
def get_base_kernel(self) -> ScalarKernel:
Expand Down Expand Up @@ -257,7 +262,8 @@
def m2p(self, src_order: int) -> E2PBase:
return E2PFromCSR(
self.multipole_expansion(src_order),
self.target_kernels, name="m2p")
self.target_kernels, name="m2p",
work_items_per_group=self.m2p_work_items_per_group)

@memoize_method
def l2p(self, src_order: int) -> E2PBase:
Expand Down Expand Up @@ -1036,7 +1042,7 @@
source_level_start_ibox, source_mpoles_view = \
self.multipole_expansions_view(mpole_exps, isrc_level)

pot_res = m2p(

Check failure on line 1045 in sumpy/fmm.py

View workflow job for this annotation

GitHub Actions / basedpyright 

Object of type "E2PBase" is not callable
  Attribute "__call__" is unknown (reportCallIssue)
actx,

src_expansions=source_mpoles_view,
Expand All @@ -1047,6 +1053,7 @@
source_box_lists=ssn.lists,
centers=self.tree.box_centers,
result=pot,
max_ntargets_in_one_box=self.max_ntargets_in_one_box,

rscale=self.level_to_rscale(isrc_level),

Expand Down
Loading