From aa40ff8fc75c5d027c028ec317de31794645b3bd Mon Sep 17 00:00:00 2001 From: Shawn Lin Date: Sat, 17 Oct 2026 17:04:07 -0500 Subject: [PATCH 1/2] wip: parallelize L2P over target boxes --- sumpy/e2p.py | 64 +++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 58 insertions(+), 6 deletions(-) diff --git a/sumpy/e2p.py b/sumpy/e2p.py index aa53985d..7feaeafa 100644 --- a/sumpy/e2p.py +++ b/sumpy/e2p.py @@ -121,6 +121,21 @@ def get_kernel_scaling_assignment(self): # {{{ 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. One group handles each target box, with targets distributed + across its workitems. *None* (the default) keeps targets serial within + each box. + """ + + def __init__(self, expansion, kernels, name=None, *, work_items_per_group=None): + super().__init__(expansion, kernels, name) + self.work_items_per_group = work_items_per_group + + def get_cache_key(self): + return (*super().get_cache_key(), self.work_items_per_group) + @property def default_name(self): return "e2p_from_single_box" @@ -132,16 +147,19 @@ def get_kernel(self): loopy_knl = make_loopy_program( [ "{[itgt_box]: 0<=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} @@ -149,7 +167,8 @@ def get_kernel(self): 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( @@ -196,8 +215,41 @@ def get_kernel(self): return self.add_loopy_eval_callable(loopy_knl) def get_optimized_kernel(self): - # FIXME knl = self.get_kernel() + if self.work_items_per_group is not None: + import namedisl as nisl + + # Target-parallel L2P scheduling follows the idea in Isuru Fernando's + # https://github.com/inducer/sumpy/pull/158. + + 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)) + + # Each workitem loads its own private coefficients and box data, + # including workitems with no targets to evaluate. + knl = lp.add_inames_to_insn(knl, "iwork_item", + "id:fetch_tgt_box_* or id:fetch_center or id:fetch_coeffs " + "or id:kernel_scaling") + 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, From 09236208a2ce741ff6c4ca5129ff7e314e751eaa Mon Sep 17 00:00:00 2001 From: Shawn Lin Date: Sat, 17 Oct 2026 20:03:57 -0500 Subject: [PATCH 2/2] l2p: tgt box parallelization --- sumpy/e2p.py | 25 +++++++++++++++---------- sumpy/fmm.py | 10 +++++++++- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/sumpy/e2p.py b/sumpy/e2p.py index 7feaeafa..844ad85f 100644 --- a/sumpy/e2p.py +++ b/sumpy/e2p.py @@ -124,9 +124,8 @@ 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. One group handles each target box, with targets distributed - across its workitems. *None* (the default) keeps targets serial within - each box. + 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): @@ -215,13 +214,15 @@ def get_kernel(self): return self.add_loopy_eval_callable(loopy_knl) def get_optimized_kernel(self): + """ + 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 - # Target-parallel L2P scheduling follows the idea in Isuru Fernando's - # https://github.com/inducer/sumpy/pull/158. - knl = lp.split_iname(knl, "itgt_offset", self.work_items_per_group, outer_iname="itgt_outer", inner_iname="iwork_item", inner_tag="l.0") @@ -244,11 +245,15 @@ def get_optimized_kernel(self): ] knl = knl.with_kernel(entry.copy(domains=domains)) - # Each workitem loads its own private coefficients and box data, - # including workitems with no targets to evaluate. knl = lp.add_inames_to_insn(knl, "iwork_item", - "id:fetch_tgt_box_* or id:fetch_center or id:fetch_coeffs " - "or id:kernel_scaling") + "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") diff --git a/sumpy/fmm.py b/sumpy/fmm.py index 8780180c..c3f4c650 100644 --- a/sumpy/fmm.py +++ b/sumpy/fmm.py @@ -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, @@ -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) @@ -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__() @@ -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 @@ -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: