diff --git a/pytential/qbx/__init__.py b/pytential/qbx/__init__.py index adb858a119921812cf0c9b7084a2cedfb0397856..50fb8e5887266235ac7864fee91986217e06be96 100644 --- a/pytential/qbx/__init__.py +++ b/pytential/qbx/__init__.py @@ -441,16 +441,16 @@ class QBXLayerPotentialSource(LayerPotentialSourceBase): # {{{ fmm-based execution @memoize_method - def expansion_wrangler_code_container(self, base_kernel, out_kernels): + def expansion_wrangler_code_container(self, fmm_kernel, out_kernels): mpole_expn_class = \ - self.expansion_factory.get_multipole_expansion_class(base_kernel) + self.expansion_factory.get_multipole_expansion_class(fmm_kernel) local_expn_class = \ - self.expansion_factory.get_local_expansion_class(base_kernel) + self.expansion_factory.get_local_expansion_class(fmm_kernel) from functools import partial - fmm_mpole_factory = partial(mpole_expn_class, base_kernel) - fmm_local_factory = partial(local_expn_class, base_kernel) - qbx_local_factory = partial(local_expn_class, base_kernel) + fmm_mpole_factory = partial(mpole_expn_class, fmm_kernel) + fmm_local_factory = partial(local_expn_class, fmm_kernel) + qbx_local_factory = partial(local_expn_class, fmm_kernel) if self.fmm_backend == "sumpy": from pytential.qbx.fmm import \ @@ -514,67 +514,23 @@ class QBXLayerPotentialSource(LayerPotentialSourceBase): strengths = (evaluate(insn.density).with_queue(queue) * self.weights_and_area_elements()) - # {{{ get expansion wrangler - - base_kernel = None - out_kernels = [] - - from sumpy.kernel import AxisTargetDerivativeRemover - for knl in insn.kernels: - candidate_base_kernel = AxisTargetDerivativeRemover()(knl) - - if base_kernel is None: - base_kernel = candidate_base_kernel - else: - assert base_kernel == candidate_base_kernel - out_kernels = tuple(knl for knl in insn.kernels) - - if base_kernel.is_complex_valued or strengths.dtype.kind == "c": - value_dtype = self.complex_dtype - else: - value_dtype = self.real_dtype - - # {{{ build extra_kwargs dictionaries - - # This contains things like the Helmholtz parameter k or - # the normal directions for double layers. - - def reorder_sources(source_array): - if isinstance(source_array, cl.array.Array): - return (source_array - .with_queue(queue) - [geo_data.tree().user_source_ids] - .with_queue(None)) - else: - return source_array - - kernel_extra_kwargs = {} - source_extra_kwargs = {} - - from sumpy.tools import gather_arguments, gather_source_arguments - from pytools.obj_array import with_object_array_or_scalar - for func, var_dict in [ - (gather_arguments, kernel_extra_kwargs), - (gather_source_arguments, source_extra_kwargs), - ]: - for arg in func(out_kernels): - var_dict[arg.name] = with_object_array_or_scalar( - reorder_sources, - evaluate(insn.kernel_arguments[arg.name])) - - # }}} + fmm_kernel = self.get_fmm_kernel(out_kernels) + output_and_expansion_dtype = ( + self.get_fmm_output_and_expansion_dtype(fmm_kernel, strengths)) + kernel_extra_kwargs, source_extra_kwargs = ( + self.get_fmm_expansion_wrangler_extra_kwargs( + queue, out_kernels, geo_data.tree().user_source_ids, + insn.kernel_arguments, evaluate)) wrangler = self.expansion_wrangler_code_container( - base_kernel, out_kernels).get_wrangler( - queue, geo_data, value_dtype, + fmm_kernel, out_kernels).get_wrangler( + queue, geo_data, output_and_expansion_dtype, self.qbx_order, self.fmm_level_to_order, source_extra_kwargs=source_extra_kwargs, kernel_extra_kwargs=kernel_extra_kwargs) - # }}} - from pytential.qbx.geometry import target_state if (geo_data.user_target_to_center().with_queue(queue) == target_state.FAILED).any().get(): diff --git a/pytential/source.py b/pytential/source.py index 6d78da5723b10dc30ae39569aba9628c422c4b7a..75d0db70ee6babe762f9589f67fc8da877fd2379 100644 --- a/pytential/source.py +++ b/pytential/source.py @@ -206,6 +206,60 @@ class LayerPotentialSourceBase(PotentialSource): return p2p + # {{{ fmm setup helpers + + def get_fmm_kernel(self, kernels): + fmm_kernel = None + + from sumpy.kernel import AxisTargetDerivativeRemover + for knl in kernels: + candidate_fmm_kernel = AxisTargetDerivativeRemover()(knl) + + if fmm_kernel is None: + fmm_kernel = candidate_fmm_kernel + else: + assert fmm_kernel == candidate_fmm_kernel + + return fmm_kernel + + def get_fmm_output_and_expansion_dtype(self, base_kernel, strengths): + if base_kernel.is_complex_valued or strengths.dtype.kind == "c": + return self.complex_dtype + else: + return self.real_dtype + + def get_fmm_expansion_wrangler_extra_kwargs( + self, queue, out_kernels, tree_user_source_ids, arguments, evaluator): + # This contains things like the Helmholtz parameter k or + # the normal directions for double layers. + + def reorder_sources(source_array): + if isinstance(source_array, cl.array.Array): + return (source_array + .with_queue(queue) + [tree_user_source_ids] + .with_queue(None)) + else: + return source_array + + kernel_extra_kwargs = {} + source_extra_kwargs = {} + + from sumpy.tools import gather_arguments, gather_source_arguments + from pytools.obj_array import with_object_array_or_scalar + for func, var_dict in [ + (gather_arguments, kernel_extra_kwargs), + (gather_source_arguments, source_extra_kwargs), + ]: + for arg in func(out_kernels): + var_dict[arg.name] = with_object_array_or_scalar( + reorder_sources, + evaluator(arguments[arg.name])) + + return kernel_extra_kwargs, source_extra_kwargs + + # }}} + # {{{ weights and area elements @memoize_method diff --git a/pytential/unregularized.py b/pytential/unregularized.py index 4ea94bf1094fd3c820b786e3f814d43e67a7e347..fcd5298e798c41bc9853c2d64491d0e84fce2c3c 100644 --- a/pytential/unregularized.py +++ b/pytential/unregularized.py @@ -27,7 +27,15 @@ THE SOFTWARE. import six +import numpy as np +import loopy as lp + +from boxtree.tools import DeviceDataRecord from pytential.source import LayerPotentialSourceBase +from pytools import memoize_method + +import pyopencl as cl +import pyopencl.array # noqa import logging logger = logging.getLogger(__name__) @@ -43,17 +51,43 @@ __doc__ = """ class UnregularizedLayerPotentialSource(LayerPotentialSourceBase): """A source discretization for a layer potential discretized with a Nyström method that uses panel-based quadrature and does not modify the kernel. + + .. attribute:: fmm_level_to_order """ def __init__(self, density_discr, + fmm_order=False, + fmm_level_to_order=None, + expansion_factory=None, # begin undocumented arguments # FIXME default debug=False once everything works debug=True): """ + :arg fmm_order: `False` for direct calculation. """ self.density_discr = density_discr self.debug = debug + if fmm_order is not False and fmm_level_to_order is not None: + raise TypeError("may not specify both fmm_order and fmm_level_to_order") + + if fmm_level_to_order is None: + if fmm_order is not False: + def fmm_level_to_order(level): + return fmm_order + else: + fmm_level_to_order = False + + self.density_discr = density_discr + self.fmm_level_to_order = fmm_level_to_order + + if expansion_factory is None: + from sumpy.expansion import DefaultExpansionFactory + expansion_factory = DefaultExpansionFactory() + self.expansion_factory = expansion_factory + + self.debug = debug + @property def fine_density_discr(self): return self.density_discr @@ -67,9 +101,12 @@ class UnregularizedLayerPotentialSource(LayerPotentialSourceBase): def copy( self, density_discr=None, - debug=None + fmm_level_to_order=None, + debug=None, ): return type(self)( + fmm_level_to_order=( + fmm_level_to_order or self.fmm_level_to_order), density_discr=density_discr or self.density_discr, debug=debug if debug is not None else self.debug) @@ -80,7 +117,11 @@ class UnregularizedLayerPotentialSource(LayerPotentialSourceBase): value = evaluate(expr) return with_object_array_or_scalar(lambda x: x, value) - func = self.exec_compute_potential_insn_direct + if self.fmm_level_to_order is False: + func = self.exec_compute_potential_insn_direct + else: + func = self.exec_compute_potential_insn_fmm + return func(queue, insn, bound_expr, evaluate_wrapper) def op_group_features(self, expr): @@ -126,6 +167,259 @@ class UnregularizedLayerPotentialSource(LayerPotentialSourceBase): return result, [] + # {{{ fmm-based execution + + @memoize_method + def expansion_wrangler_code_container(self, fmm_kernel, out_kernels): + mpole_expn_class = \ + self.expansion_factory.get_multipole_expansion_class(fmm_kernel) + local_expn_class = \ + self.expansion_factory.get_local_expansion_class(fmm_kernel) + + from functools import partial + fmm_mpole_factory = partial(mpole_expn_class, fmm_kernel) + fmm_local_factory = partial(local_expn_class, fmm_kernel) + + from sumpy.fmm import SumpyExpansionWranglerCodeContainer + return SumpyExpansionWranglerCodeContainer( + self.cl_context, + fmm_mpole_factory, + fmm_local_factory, + out_kernels) + + @property + @memoize_method + def fmm_geometry_code_container(self): + return _FMMGeometryCodeContainer( + self.cl_context, self.ambient_dim, self.debug) + + def fmm_geometry_data(self, targets): + return _FMMGeometryData( + self, + self.fmm_geometry_code_container, + targets, + self.debug) + + def exec_compute_potential_insn_fmm(self, queue, insn, bound_expr, evaluate): + + # {{{ gather unique target discretizations used + + target_name_to_index = {} + targets = [] + + for o in insn.outputs: + assert o.qbx_forced_limit not in (-1, 1) + + if o.target_name in target_name_to_index: + continue + + target_name_to_index[o.target_name] = len(targets) + targets.append(bound_expr.places[o.target_name]) + + targets = tuple(targets) + + # }}} + + # {{{ get wrangler + + geo_data = self.fmm_geometry_data(targets) + + strengths = (evaluate(insn.density).with_queue(queue) + * self.weights_and_area_elements()) + + out_kernels = tuple(knl for knl in insn.kernels) + fmm_kernel = self.get_fmm_kernel(out_kernels) + output_and_expansion_dtype = ( + self.get_fmm_output_and_expansion_dtype(fmm_kernel, strengths)) + kernel_extra_kwargs, source_extra_kwargs = ( + self.get_fmm_expansion_wrangler_extra_kwargs( + queue, out_kernels, geo_data.tree().user_source_ids, + insn.kernel_arguments, evaluate)) + + wrangler = self.expansion_wrangler_code_container( + fmm_kernel, out_kernels).get_wrangler( + queue, + geo_data.tree(), + output_and_expansion_dtype, + self.fmm_level_to_order, + source_extra_kwargs=source_extra_kwargs, + kernel_extra_kwargs=kernel_extra_kwargs) + + # }}} + + from boxtree.fmm import drive_fmm + all_potentials_on_every_tgt = drive_fmm( + geo_data.traversal(), wrangler, strengths) + + # {{{ postprocess fmm + + result = [] + + for o in insn.outputs: + target_index = target_name_to_index[o.target_name] + target_slice = slice(*geo_data.target_info().target_discr_starts[ + target_index:target_index+2]) + + result.append( + (o.name, + all_potentials_on_every_tgt[o.kernel_index][target_slice])) + + # }}} + + return result, [] + + # }}} + +# }}} + + +# {{{ fmm tools + +class _FMMGeometryCodeContainer(object): + + def __init__(self, cl_context, ambient_dim, debug): + self.cl_context = cl_context + self.ambient_dim = ambient_dim + self.debug = debug + + @memoize_method + def copy_targets_kernel(self): + knl = lp.make_kernel( + """{[dim,i]: + 0<=dim