From 02073109ceb9fb6178822a47ea3b06c967ee0197 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Tue, 12 Jun 2018 17:00:01 -0500 Subject: [PATCH 1/8] [ci skip] Add @record_time decorators to the pyfmmlib FMM. Depends on inducer/pytools!14. --- boxtree/pyfmmlib_integration.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/boxtree/pyfmmlib_integration.py b/boxtree/pyfmmlib_integration.py index ddf7e20..f6143f5 100644 --- a/boxtree/pyfmmlib_integration.py +++ b/boxtree/pyfmmlib_integration.py @@ -26,7 +26,7 @@ THE SOFTWARE. import numpy as np -from pytools import memoize_method, log_process +from pytools import memoize_method, log_process, record_time import logging logger = logging.getLogger(__name__) @@ -394,10 +394,12 @@ class FMMLibExpansionWrangler(object): # }}} @log_process(logger) + @record_time("timing_data") def reorder_sources(self, source_array): return source_array[..., self.tree.user_source_ids] @log_process(logger) + @record_time("timing_data") def reorder_potentials(self, potentials): return potentials[self.tree.sorted_target_ids] @@ -420,6 +422,7 @@ class FMMLibExpansionWrangler(object): } @log_process(logger) + @record_time("timing_data") def form_multipoles(self, level_start_source_box_nrs, source_boxes, src_weights): formmp = self.get_routine("%ddformmp" + self.dp_suffix) @@ -459,6 +462,7 @@ class FMMLibExpansionWrangler(object): return mpoles @log_process(logger) + @record_time("timing_data") def coarsen_multipoles(self, level_start_source_parent_box_nrs, source_parent_boxes, mpoles): tree = self.tree @@ -512,6 +516,7 @@ class FMMLibExpansionWrangler(object): ibox - target_level_start_ibox] += new_mp[..., 0].T @log_process(logger) + @record_time("timing_data") def eval_direct(self, target_boxes, neighbor_sources_starts, neighbor_sources_lists, src_weights): output = self.output_zeros() @@ -553,6 +558,7 @@ class FMMLibExpansionWrangler(object): return output @log_process(logger) + @record_time("timing_data") def multipole_to_local(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, @@ -638,6 +644,7 @@ class FMMLibExpansionWrangler(object): return local_exps @log_process(logger) + @record_time("timing_data") def eval_multipoles(self, target_boxes_by_source_level, sep_smaller_nonsiblings_by_level, mpole_exps): @@ -680,6 +687,7 @@ class FMMLibExpansionWrangler(object): return output @log_process(logger) + @record_time("timing_data") def form_locals(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, starts, lists, src_weights): @@ -731,6 +739,7 @@ class FMMLibExpansionWrangler(object): return local_exps @log_process(logger) + @record_time("timing_data") def refine_locals(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, local_exps): @@ -777,6 +786,7 @@ class FMMLibExpansionWrangler(object): return local_exps @log_process(logger) + @record_time("timing_data") def eval_locals(self, level_start_target_box_nrs, target_boxes, local_exps): output = self.output_zeros() taeval = self.get_expn_eval_routine("ta") @@ -812,6 +822,7 @@ class FMMLibExpansionWrangler(object): return output @log_process(logger) + @record_time("timing_data") def finalize_potentials(self, potential): if self.eqn_letter == "l" and self.dim == 2: scale_factor = -1/(2*np.pi) -- GitLab From c8395c4e6b2a431bf0ddb7d916388a5c94cf91f4 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Thu, 21 Jun 2018 16:56:24 -0500 Subject: [PATCH 2/8] [ci skip] Change FMM API to allow for recording timing data. Depends on: pytools!15 --- boxtree/fmm.py | 96 ++++++++++++++++++++++++++------- boxtree/pyfmmlib_integration.py | 6 +-- boxtree/tools.py | 64 ++++++++++++++++++++++ test/test_fmm.py | 23 +++++++- 4 files changed, 165 insertions(+), 24 deletions(-) diff --git a/boxtree/fmm.py b/boxtree/fmm.py index b2dd276..357c432 100644 --- a/boxtree/fmm.py +++ b/boxtree/fmm.py @@ -28,7 +28,7 @@ logger = logging.getLogger(__name__) from pytools import ProcessLogger -def drive_fmm(traversal, expansion_wrangler, src_weights): +def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): """Top-level driver routine for a fast multipole calculation. In part, this is intended as a template for custom FMMs, in the sense that @@ -44,6 +44,8 @@ def drive_fmm(traversal, expansion_wrangler, src_weights): :class:`ExpansionWranglerInterface`. :arg src_weights: Source 'density/weights/charges'. Passed unmodified to *expansion_wrangler*. + :arg timing_data: Either *None*, or a :class:`dict` inside which timing + information, if available, is returned Returns the potentials computed by *expansion_wrangler*. """ @@ -53,6 +55,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights): # to the expansion wrangler and should not be passed. fmm_proc = ProcessLogger(logger, "qbx fmm") + recorder = TimingDataRecorder() src_weights = wrangler.reorder_sources(src_weights) @@ -61,7 +64,8 @@ def drive_fmm(traversal, expansion_wrangler, src_weights): mpole_exps = wrangler.form_multipoles( traversal.level_start_source_box_nrs, traversal.source_boxes, - src_weights) + src_weights, + timing_data=recorder.next_record()) # }}} @@ -70,7 +74,8 @@ def drive_fmm(traversal, expansion_wrangler, src_weights): wrangler.coarsen_multipoles( traversal.level_start_source_parent_box_nrs, traversal.source_parent_boxes, - mpole_exps) + mpole_exps, + timing_data=recorder.next_record()) # mpole_exps is called Phi in [1] @@ -82,7 +87,8 @@ def drive_fmm(traversal, expansion_wrangler, src_weights): traversal.target_boxes, traversal.neighbor_source_boxes_starts, traversal.neighbor_source_boxes_lists, - src_weights) + src_weights, + timing_data=recorder.next_record()) # these potentials are called alpha in [1] @@ -95,7 +101,8 @@ def drive_fmm(traversal, expansion_wrangler, src_weights): traversal.target_or_target_parent_boxes, traversal.from_sep_siblings_starts, traversal.from_sep_siblings_lists, - mpole_exps) + mpole_exps, + timing_data=recorder.next_record()) # local_exps represents both Gamma and Delta in [1] @@ -109,7 +116,8 @@ def drive_fmm(traversal, expansion_wrangler, src_weights): potentials = potentials + wrangler.eval_multipoles( traversal.target_boxes_sep_smaller_by_source_level, traversal.from_sep_smaller_by_level, - mpole_exps) + mpole_exps, + timing_data=recorder.next_record()) # these potentials are called beta in [1] @@ -121,7 +129,8 @@ def drive_fmm(traversal, expansion_wrangler, src_weights): traversal.target_boxes, traversal.from_sep_close_smaller_starts, traversal.from_sep_close_smaller_lists, - src_weights) + src_weights, + timing_data=recorder.next_record()) # }}} @@ -132,14 +141,16 @@ def drive_fmm(traversal, expansion_wrangler, src_weights): traversal.target_or_target_parent_boxes, traversal.from_sep_bigger_starts, traversal.from_sep_bigger_lists, - src_weights) + src_weights, + timing_data=recorder.next_record()) if traversal.from_sep_close_bigger_starts is not None: potentials = potentials + wrangler.eval_direct( traversal.target_or_target_parent_boxes, traversal.from_sep_close_bigger_starts, traversal.from_sep_close_bigger_lists, - src_weights) + src_weights, + timing_data=recorder.next_record()) # }}} @@ -148,7 +159,8 @@ def drive_fmm(traversal, expansion_wrangler, src_weights): wrangler.refine_locals( traversal.level_start_target_or_target_parent_box_nrs, traversal.target_or_target_parent_boxes, - local_exps) + local_exps, + timing_data=recorder.next_record()) # }}} @@ -157,7 +169,8 @@ def drive_fmm(traversal, expansion_wrangler, src_weights): potentials = potentials + wrangler.eval_locals( traversal.level_start_target_box_nrs, traversal.target_boxes, - local_exps) + local_exps, + timing_data=recorder.next_record()) # }}} @@ -167,6 +180,9 @@ def drive_fmm(traversal, expansion_wrangler, src_weights): fmm_proc.done() + if timing_data is not None: + timing_data.update(recorder.summarize()) + return result @@ -181,6 +197,14 @@ class ExpansionWranglerInterface: Will usually hold a reference (and thereby be specific to) a :class:`boxtree.Tree` instance. + + This interface supports collecting timing data. If timing data is requested, + the dictionary argument *timing_data* filled with the following two entries: + + - *description* + - *callback* + + *callback* may block until the operation is finished. """ def multipole_expansion_zeros(self): @@ -215,7 +239,8 @@ class ExpansionWranglerInterface: *source_weights* is in tree target order. """ - def form_multipoles(self, level_start_source_box_nrs, source_boxes, src_weights): + def form_multipoles(self, level_start_source_box_nrs, source_boxes, src_weights, + timing_data=None): """Return an expansions array (compatible with :meth:`multipole_expansion_zeros`) containing multipole expansions in *source_boxes* due to sources @@ -224,7 +249,7 @@ class ExpansionWranglerInterface: """ def coarsen_multipoles(self, level_start_source_parent_box_nrs, - source_parent_boxes, mpoles): + source_parent_boxes, mpoles, timing_data=None): """For each box in *source_parent_boxes*, gather (and translate) the box's children's multipole expansions in *mpole* and add the resulting expansion into the box's multipole @@ -234,7 +259,7 @@ class ExpansionWranglerInterface: """ def eval_direct(self, target_boxes, neighbor_sources_starts, - neighbor_sources_lists, src_weights): + neighbor_sources_lists, src_weights, timing_data=None): """For each box in *target_boxes*, evaluate the influence of the neighbor sources due to *src_weights*, which use :ref:`csr` and are indexed like *target_boxes*. @@ -245,7 +270,7 @@ class ExpansionWranglerInterface: def multipole_to_local(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, - starts, lists, mpole_exps): + starts, lists, mpole_exps, timing_data=None): """For each box in *target_or_target_parent_boxes*, translate and add the influence of the multipole expansion in *mpole_exps* into a new array of local expansions. *starts* and *lists* use :ref:`csr`, and @@ -256,7 +281,8 @@ class ExpansionWranglerInterface: """ def eval_multipoles(self, - target_boxes_by_source_level, from_sep_smaller_by_level, mpole_exps): + target_boxes_by_source_level, from_sep_smaller_by_level, mpole_exps, + timing_data=None): """For a level *i*, each box in *target_boxes_by_source_level[i]*, evaluate the multipole expansion in *mpole_exps* in the nearby boxes given in *from_sep_smaller_by_level*, and return a new potential array. @@ -268,7 +294,8 @@ class ExpansionWranglerInterface: def form_locals(self, level_start_target_or_target_parent_box_nrs, - target_or_target_parent_boxes, starts, lists, src_weights): + target_or_target_parent_boxes, starts, lists, src_weights, + timing_data=None): """For each box in *target_or_target_parent_boxes*, form local expansions due to the sources in the nearby boxes given in *starts* and *lists*, and return a new local expansion array. *starts* and *lists* @@ -281,7 +308,7 @@ class ExpansionWranglerInterface: pass def refine_locals(self, level_start_target_or_target_parent_box_nrs, - target_or_target_parent_boxes, local_exps): + target_or_target_parent_boxes, local_exps, timing_data=None): """For each box in *child_boxes*, translate the box's parent's local expansion in *local_exps* and add the resulting expansion into the box's local expansion in *local_exps*. @@ -289,7 +316,8 @@ class ExpansionWranglerInterface: :returns: *local_exps* """ - def eval_locals(self, level_start_target_box_nrs, target_boxes, local_exps): + def eval_locals(self, level_start_target_box_nrs, target_boxes, local_exps, + timing_data=None): """For each box in *target_boxes*, evaluate the local expansion in *local_exps* and return a new potential array. @@ -306,4 +334,34 @@ class ExpansionWranglerInterface: # }}} +# {{{ timing data recorder + +class TimingDataRecorder(object): + + def __init__(self): + self.records = [] + + def next_record(self): + self.records.append(dict()) + return self.records[-1] + + def summarize(self): + result = {} + + for record in self.records: + if "description" not in record: + continue + + description = record["description"] + + if description in result: + result[description] += record["callback"]() + else: + result[description] = record["callback"]() + + return result + +# }}} + + # vim: filetype=pyopencl:fdm=marker diff --git a/boxtree/pyfmmlib_integration.py b/boxtree/pyfmmlib_integration.py index f6143f5..32737d5 100644 --- a/boxtree/pyfmmlib_integration.py +++ b/boxtree/pyfmmlib_integration.py @@ -26,7 +26,8 @@ THE SOFTWARE. import numpy as np -from pytools import memoize_method, log_process, record_time +from pytools import memoize_method, log_process +from boxtree.tools import record_time import logging logger = logging.getLogger(__name__) @@ -394,12 +395,10 @@ class FMMLibExpansionWrangler(object): # }}} @log_process(logger) - @record_time("timing_data") def reorder_sources(self, source_array): return source_array[..., self.tree.user_source_ids] @log_process(logger) - @record_time("timing_data") def reorder_potentials(self, potentials): return potentials[self.tree.sorted_target_ids] @@ -822,7 +821,6 @@ class FMMLibExpansionWrangler(object): return output @log_process(logger) - @record_time("timing_data") def finalize_potentials(self, potential): if self.eqn_letter == "l" and self.dim == 2: scale_factor = -1/(2*np.pi) diff --git a/boxtree/tools.py b/boxtree/tools.py index 68ac5f4..b96abbb 100644 --- a/boxtree/tools.py +++ b/boxtree/tools.py @@ -501,6 +501,70 @@ class MapValuesKernel(object): # }}} +# {{{ time recording tool + +class record_time(object): # noqa: N801 + """A decorator for recording timing data for a function call. + + Timing data is saved to a :class:`dict` passed as an optional keyword + argument. The following entries are written: + + - *description* + - *callback*. + + *callback* may be called once to obtain the elapsed wall time in + seconds. + + Example usage:: + + >>> from time import sleep + >>> @record_time("timing_data") + ... def slow_function(n): + ... sleep(n) + ... + >>> timing_result = {} + >>> slow_function(1, timing_data=timing_result) + >>> elapsed_time = timing_result["callback"]() + """ + + def __init__(self, arg=None, description=None): + self.arg = arg + self.description = description + + def __call__(self, wrapped): + description = self.description or wrapped.__name__ + + from pytools import ProcessTimer + from contextlib import contextmanager + + @contextmanager + def time_process(output): + timer = ProcessTimer() + yield + timer.done() + + def callback(): + return timer.wall_elapsed + + output["description"] = description + output["callback"] = callback + + def wrapper(*args, **kwargs): + output = kwargs.pop(self.arg, None) + if output is None: + return wrapped(*args, **kwargs) + + with time_process(output): + return wrapped(*args, **kwargs) + + from functools import update_wrapper + new_wrapper = update_wrapper(wrapper, wrapped) + + return new_wrapper + +# }}} + + # {{{ binary search from mako.template import Template diff --git a/test/test_fmm.py b/test/test_fmm.py index 2eb9eb3..6133251 100644 --- a/test/test_fmm.py +++ b/test/test_fmm.py @@ -44,6 +44,16 @@ logger = logging.getLogger(__name__) # {{{ fmm interaction completeness test +def ignore_timing_data(f): + from functools import wraps + + @wraps(f) + def wrapper(timing_data=None, *args, **kwargs): + return f(*args, **kwargs) + + return wrapper + + class ConstantOneExpansionWrangler(object): """This implements the 'analytical routines' for a Green's function that is constant 1 everywhere. For 'charges' of 'ones', this should get every particle @@ -77,6 +87,7 @@ class ConstantOneExpansionWrangler(object): def reorder_potentials(self, potentials): return potentials[self.tree.sorted_target_ids] + @ignore_timing_data def form_multipoles(self, level_start_source_box_nrs, source_boxes, src_weights): mpoles = self.multipole_expansion_zeros() for ibox in source_boxes: @@ -85,6 +96,7 @@ class ConstantOneExpansionWrangler(object): return mpoles + @ignore_timing_data def coarsen_multipoles(self, level_start_source_parent_box_nrs, source_parent_boxes, mpoles): tree = self.tree @@ -104,6 +116,7 @@ class ConstantOneExpansionWrangler(object): if child: mpoles[ibox] += mpoles[child] + @ignore_timing_data def eval_direct(self, target_boxes, neighbor_sources_starts, neighbor_sources_lists, src_weights): pot = self.potential_zeros() @@ -123,6 +136,7 @@ class ConstantOneExpansionWrangler(object): return pot + @ignore_timing_data def multipole_to_local(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, @@ -141,6 +155,7 @@ class ConstantOneExpansionWrangler(object): return local_exps + @ignore_timing_data def eval_multipoles(self, target_boxes_by_source_level, from_sep_smaller_nonsiblings_by_level, mpole_exps): @@ -161,6 +176,7 @@ class ConstantOneExpansionWrangler(object): return pot + @ignore_timing_data def form_locals(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, starts, lists, src_weights): @@ -180,6 +196,7 @@ class ConstantOneExpansionWrangler(object): return local_exps + @ignore_timing_data def refine_locals(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, local_exps): @@ -191,6 +208,7 @@ class ConstantOneExpansionWrangler(object): return local_exps + @ignore_timing_data def eval_locals(self, level_start_target_box_nrs, target_boxes, local_exps): pot = self.potential_zeros() @@ -564,7 +582,10 @@ def test_pyfmmlib_fmm(ctx_getter, dims, use_dipoles, helmholtz_k): dipole_vec=dipole_vec) from boxtree.fmm import drive_fmm - pot = drive_fmm(trav, wrangler, weights) + + timing_data = {} + pot = drive_fmm(trav, wrangler, weights, timing_data=timing_data) + assert timing_data # {{{ ref fmmlib computation -- GitLab From 598fec77976f06bd51d99769e39f087ce3f2275c Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 22 Jun 2018 01:45:01 -0500 Subject: [PATCH 3/8] Avoid making timing_data the first positional argument. --- test/test_fmm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/test_fmm.py b/test/test_fmm.py index 6133251..f1b35ec 100644 --- a/test/test_fmm.py +++ b/test/test_fmm.py @@ -48,7 +48,8 @@ def ignore_timing_data(f): from functools import wraps @wraps(f) - def wrapper(timing_data=None, *args, **kwargs): + def wrapper(*args, **kwargs): + kwargs.pop("timing_data") return f(*args, **kwargs) return wrapper -- GitLab From a3e13b280e6463dec95e1c03afd0cffa6d4c1dc5 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 6 Jul 2018 01:23:11 -0500 Subject: [PATCH 4/8] Pass an object as timing_data to the FMM wrangler. --- boxtree/fmm.py | 121 +++++++++++++++++++++++++++++++++++------------ boxtree/tools.py | 30 +++--------- test/test_fmm.py | 1 + 3 files changed, 100 insertions(+), 52 deletions(-) diff --git a/boxtree/fmm.py b/boxtree/fmm.py index 357c432..3133141 100644 --- a/boxtree/fmm.py +++ b/boxtree/fmm.py @@ -25,7 +25,7 @@ THE SOFTWARE. import logging logger = logging.getLogger(__name__) -from pytools import ProcessLogger +from pytools import ProcessLogger, Record def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): @@ -44,10 +44,12 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): :class:`ExpansionWranglerInterface`. :arg src_weights: Source 'density/weights/charges'. Passed unmodified to *expansion_wrangler*. - :arg timing_data: Either *None*, or a :class:`dict` inside which timing - information, if available, is returned + :arg timing_data: Either *None*, or a :class:`dict` that is populated with + timing information for the stages of the algorithm (in the form of + instances of :class:`TimingResult`), if such information is available. Returns the potentials computed by *expansion_wrangler*. + """ wrangler = expansion_wrangler @@ -55,7 +57,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): # to the expansion wrangler and should not be passed. fmm_proc = ProcessLogger(logger, "qbx fmm") - recorder = TimingDataRecorder() + recorder = TimingRecorder() src_weights = wrangler.reorder_sources(src_weights) @@ -65,7 +67,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): traversal.level_start_source_box_nrs, traversal.source_boxes, src_weights, - timing_data=recorder.next_record()) + timing_data=recorder.next()) # }}} @@ -75,7 +77,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): traversal.level_start_source_parent_box_nrs, traversal.source_parent_boxes, mpole_exps, - timing_data=recorder.next_record()) + timing_data=recorder.next()) # mpole_exps is called Phi in [1] @@ -88,7 +90,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): traversal.neighbor_source_boxes_starts, traversal.neighbor_source_boxes_lists, src_weights, - timing_data=recorder.next_record()) + timing_data=recorder.next()) # these potentials are called alpha in [1] @@ -102,7 +104,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): traversal.from_sep_siblings_starts, traversal.from_sep_siblings_lists, mpole_exps, - timing_data=recorder.next_record()) + timing_data=recorder.next()) # local_exps represents both Gamma and Delta in [1] @@ -117,7 +119,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): traversal.target_boxes_sep_smaller_by_source_level, traversal.from_sep_smaller_by_level, mpole_exps, - timing_data=recorder.next_record()) + timing_data=recorder.next()) # these potentials are called beta in [1] @@ -130,7 +132,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): traversal.from_sep_close_smaller_starts, traversal.from_sep_close_smaller_lists, src_weights, - timing_data=recorder.next_record()) + timing_data=recorder.next()) # }}} @@ -142,7 +144,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): traversal.from_sep_bigger_starts, traversal.from_sep_bigger_lists, src_weights, - timing_data=recorder.next_record()) + timing_data=recorder.next()) if traversal.from_sep_close_bigger_starts is not None: potentials = potentials + wrangler.eval_direct( @@ -150,7 +152,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): traversal.from_sep_close_bigger_starts, traversal.from_sep_close_bigger_lists, src_weights, - timing_data=recorder.next_record()) + timing_data=recorder.next()) # }}} @@ -160,7 +162,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): traversal.level_start_target_or_target_parent_box_nrs, traversal.target_or_target_parent_boxes, local_exps, - timing_data=recorder.next_record()) + timing_data=recorder.next()) # }}} @@ -170,7 +172,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): traversal.level_start_target_box_nrs, traversal.target_boxes, local_exps, - timing_data=recorder.next_record()) + timing_data=recorder.next()) # }}} @@ -199,12 +201,7 @@ class ExpansionWranglerInterface: :class:`boxtree.Tree` instance. This interface supports collecting timing data. If timing data is requested, - the dictionary argument *timing_data* filled with the following two entries: - - - *description* - - *callback* - - *callback* may block until the operation is finished. + the *timing_data* argument is a :class:`TimingDataWaiter` whose fields can """ def multipole_expansion_zeros(self): @@ -305,7 +302,6 @@ class ExpansionWranglerInterface: :returns: a new local expansion array, see :meth:`local_expansion_zeros`. """ - pass def refine_locals(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, local_exps, timing_data=None): @@ -334,30 +330,97 @@ class ExpansionWranglerInterface: # }}} -# {{{ timing data recorder +# {{{ timing result + +class TimingResult(Record): + """ + .. automethod:: __add__ + + .. attribute:: wall_elapsed + .. attribute:: process_elapsed + """ + + def __init__(self, wall_elapsed, process_elapsed): + Record.__init__(self, + wall_elapsed=wall_elapsed, + process_elapsed=process_elapsed) + + def __add__(self, other): + wall_elapsed = self.wall_elapsed + other.wall_elapsed + process_elapsed = self.process_elapsed + other.process_elapsed + return TimingResult(wall_elapsed, process_elapsed) + +# }}} + + +# {{{ timing waiter + +class TimingWaiter(object): + """Obtains timing data through a supplied callback function. + + Attributes that can be set:: + + .. attribute:: description + + A string, the description of the timing data. + + .. attribute:: callback + + Returns a :class:`TimingResult`. + """ + + def __init__(self): + self.description = None + self.callback = None + self._result = None + + @property + def empty(self): + return not self.callback + + @property + def result(self): + if self._result is None: + self.wait() + + return self._result + + def wait(self): + if self.empty: + return + + callback_result = self.callback() + self._result = TimingResult( + callback_result.wall_elapsed, + callback_result.process_elapsed) + +# }}} + + +# {{{ timing recorder -class TimingDataRecorder(object): +class TimingRecorder(object): def __init__(self): self.records = [] - def next_record(self): - self.records.append(dict()) + def next(self): + self.records.append(TimingWaiter()) return self.records[-1] def summarize(self): result = {} for record in self.records: - if "description" not in record: + if record.empty: continue - description = record["description"] + description = record.description if description in result: - result[description] += record["callback"]() + result[description] += record.result else: - result[description] = record["callback"]() + result[description] = record.result return result diff --git a/boxtree/tools.py b/boxtree/tools.py index b96abbb..cea84a9 100644 --- a/boxtree/tools.py +++ b/boxtree/tools.py @@ -506,25 +506,12 @@ class MapValuesKernel(object): class record_time(object): # noqa: N801 """A decorator for recording timing data for a function call. - Timing data is saved to a :class:`dict` passed as an optional keyword - argument. The following entries are written: + This introduces an extra keyword argument to the decorated function. For the + newly added argument, the caller should pass either *None* or an instance of + :class:`boxtree.fmm.TimingWaiter`. If the latter gets passed, the fields + *callback* and *description* are populated with timing data for the function + call. - - *description* - - *callback*. - - *callback* may be called once to obtain the elapsed wall time in - seconds. - - Example usage:: - - >>> from time import sleep - >>> @record_time("timing_data") - ... def slow_function(n): - ... sleep(n) - ... - >>> timing_result = {} - >>> slow_function(1, timing_data=timing_result) - >>> elapsed_time = timing_result["callback"]() """ def __init__(self, arg=None, description=None): @@ -543,11 +530,8 @@ class record_time(object): # noqa: N801 yield timer.done() - def callback(): - return timer.wall_elapsed - - output["description"] = description - output["callback"] = callback + output.description = description + output.callback = lambda: timer def wrapper(*args, **kwargs): output = kwargs.pop(self.arg, None) diff --git a/test/test_fmm.py b/test/test_fmm.py index f1b35ec..5e1b6c0 100644 --- a/test/test_fmm.py +++ b/test/test_fmm.py @@ -586,6 +586,7 @@ def test_pyfmmlib_fmm(ctx_getter, dims, use_dipoles, helmholtz_k): timing_data = {} pot = drive_fmm(trav, wrangler, weights, timing_data=timing_data) + print(timing_data) assert timing_data # {{{ ref fmmlib computation -- GitLab From e63182ac75f3bae987c724b2d0e64e039a77044e Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 6 Jul 2018 01:35:29 -0500 Subject: [PATCH 5/8] Document. --- boxtree/fmm.py | 6 ++++-- doc/fmm.rst | 4 ++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/boxtree/fmm.py b/boxtree/fmm.py index 3133141..77969ae 100644 --- a/boxtree/fmm.py +++ b/boxtree/fmm.py @@ -201,7 +201,8 @@ class ExpansionWranglerInterface: :class:`boxtree.Tree` instance. This interface supports collecting timing data. If timing data is requested, - the *timing_data* argument is a :class:`TimingDataWaiter` whose fields can + the *timing_data* argument expects a :class:`TimingWaiter`. Otherwise + *timing_data* should be *None*. """ def multipole_expansion_zeros(self): @@ -358,7 +359,7 @@ class TimingResult(Record): class TimingWaiter(object): """Obtains timing data through a supplied callback function. - Attributes that can be set:: + Attributes that should be set: .. attribute:: description @@ -380,6 +381,7 @@ class TimingWaiter(object): @property def result(self): + """The timing data result obtained from *callback*.""" if self._result is None: self.wait() diff --git a/doc/fmm.rst b/doc/fmm.rst index cca8157..38f400d 100644 --- a/doc/fmm.rst +++ b/doc/fmm.rst @@ -10,6 +10,10 @@ FMM driver :undoc-members: :member-order: bysource +.. autoclass:: TimingResult + +.. autoclass:: TimingWaiter + Integration with PyFMMLib ------------------------- -- GitLab From d8821901ddff23b96e48e68a96cae4589866259c Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 6 Jul 2018 01:46:35 -0500 Subject: [PATCH 6/8] Bump version. --- boxtree/version.py | 2 +- doc/misc.rst | 7 ++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/boxtree/version.py b/boxtree/version.py index 9bf5997..aac0098 100644 --- a/boxtree/version.py +++ b/boxtree/version.py @@ -1,2 +1,2 @@ -VERSION = (2013, 1) +VERSION = (2018, 1) VERSION_TEXT = ".".join(str(i) for i in VERSION) diff --git a/doc/misc.rst b/doc/misc.rst index 29226b5..1086d8b 100644 --- a/doc/misc.rst +++ b/doc/misc.rst @@ -24,13 +24,18 @@ for instructions. User-visible Changes ==================== -Version 2013.1 +Version 2018.1 -------------- .. note:: This version is currently under development. You can get snapshots from boxtree's `git repository `_ +* Added *timing_data* parameter to FMM driver. + +Version 2013.1 +-------------- + * Initial release. .. _license: -- GitLab From 29e6f56d389e5455c9c3793730c6f28b5f4569f7 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Fri, 6 Jul 2018 17:45:34 -0500 Subject: [PATCH 7/8] Allow None attributes --- boxtree/fmm.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/boxtree/fmm.py b/boxtree/fmm.py index 77969ae..11c2178 100644 --- a/boxtree/fmm.py +++ b/boxtree/fmm.py @@ -347,8 +347,14 @@ class TimingResult(Record): process_elapsed=process_elapsed) def __add__(self, other): - wall_elapsed = self.wall_elapsed + other.wall_elapsed - process_elapsed = self.process_elapsed + other.process_elapsed + wall_elapsed = ( + self.wall_elapsed + other.wall_elapsed + if self.wall_elapsed is not None else None) + + process_elapsed = ( + self.process_elapsed + other.process_elapsed + if self.process_elapsed is not None else None) + return TimingResult(wall_elapsed, process_elapsed) # }}} -- GitLab From de2fe76b3a26856e3b82ab73dca94602ea31839a Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 9 Jul 2018 18:34:23 -0500 Subject: [PATCH 8/8] Change expansion wrangler interface to return futures in a tuple. --- boxtree/fmm.py | 225 ++++++++++++++++---------------- boxtree/pyfmmlib_integration.py | 20 +-- boxtree/tools.py | 60 ++++----- doc/fmm.rst | 2 +- test/test_fmm.py | 28 ++-- 5 files changed, 162 insertions(+), 173 deletions(-) diff --git a/boxtree/fmm.py b/boxtree/fmm.py index 11c2178..54a1649 100644 --- a/boxtree/fmm.py +++ b/boxtree/fmm.py @@ -46,7 +46,7 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): Passed unmodified to *expansion_wrangler*. :arg timing_data: Either *None*, or a :class:`dict` that is populated with timing information for the stages of the algorithm (in the form of - instances of :class:`TimingResult`), if such information is available. + :class:`TimingResult`), if such information is available. Returns the potentials computed by *expansion_wrangler*. @@ -63,21 +63,23 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): # {{{ "Step 2.1:" Construct local multipoles - mpole_exps = wrangler.form_multipoles( + mpole_exps, timing_future = wrangler.form_multipoles( traversal.level_start_source_box_nrs, traversal.source_boxes, - src_weights, - timing_data=recorder.next()) + src_weights) + + recorder.add("form_multipoles", timing_future) # }}} # {{{ "Step 2.2:" Propagate multipoles upward - wrangler.coarsen_multipoles( + mpole_exps, timing_future = wrangler.coarsen_multipoles( traversal.level_start_source_parent_box_nrs, traversal.source_parent_boxes, - mpole_exps, - timing_data=recorder.next()) + mpole_exps) + + recorder.add("coarsen_multipoles", timing_future) # mpole_exps is called Phi in [1] @@ -85,12 +87,13 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): # {{{ "Stage 3:" Direct evaluation from neighbor source boxes ("list 1") - potentials = wrangler.eval_direct( + potentials, timing_future = wrangler.eval_direct( traversal.target_boxes, traversal.neighbor_source_boxes_starts, traversal.neighbor_source_boxes_lists, - src_weights, - timing_data=recorder.next()) + src_weights) + + recorder.add("eval_direct", timing_future) # these potentials are called alpha in [1] @@ -98,13 +101,14 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): # {{{ "Stage 4:" translate separated siblings' ("list 2") mpoles to local - local_exps = wrangler.multipole_to_local( + local_exps, timing_future = wrangler.multipole_to_local( traversal.level_start_target_or_target_parent_box_nrs, traversal.target_or_target_parent_boxes, traversal.from_sep_siblings_starts, traversal.from_sep_siblings_lists, - mpole_exps, - timing_data=recorder.next()) + mpole_exps) + + recorder.add("multipole_to_local", timing_future) # local_exps represents both Gamma and Delta in [1] @@ -115,11 +119,14 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): # (the point of aiming this stage at particles is specifically to keep its # contribution *out* of the downward-propagating local expansions) - potentials = potentials + wrangler.eval_multipoles( + mpole_result, timing_future = wrangler.eval_multipoles( traversal.target_boxes_sep_smaller_by_source_level, traversal.from_sep_smaller_by_level, - mpole_exps, - timing_data=recorder.next()) + mpole_exps) + + recorder.add("eval_multipoles", timing_future) + + potentials = potentials + mpole_result # these potentials are called beta in [1] @@ -127,52 +134,65 @@ def drive_fmm(traversal, expansion_wrangler, src_weights, timing_data=None): logger.debug("evaluate separated close smaller interactions directly " "('list 3 close')") - potentials = potentials + wrangler.eval_direct( + direct_result, timing_future = wrangler.eval_direct( traversal.target_boxes, traversal.from_sep_close_smaller_starts, traversal.from_sep_close_smaller_lists, - src_weights, - timing_data=recorder.next()) + src_weights) + + recorder.add("eval_direct", timing_future) + + potentials = potentials + direct_result # }}} # {{{ "Stage 6:" form locals for separated bigger source boxes ("list 4") - local_exps = local_exps + wrangler.form_locals( + local_result, timing_future = wrangler.form_locals( traversal.level_start_target_or_target_parent_box_nrs, traversal.target_or_target_parent_boxes, traversal.from_sep_bigger_starts, traversal.from_sep_bigger_lists, - src_weights, - timing_data=recorder.next()) + src_weights) + + recorder.add("form_locals", timing_future) + + local_exps = local_exps + local_result if traversal.from_sep_close_bigger_starts is not None: - potentials = potentials + wrangler.eval_direct( + direct_result, timing_future = wrangler.eval_direct( traversal.target_or_target_parent_boxes, traversal.from_sep_close_bigger_starts, traversal.from_sep_close_bigger_lists, - src_weights, - timing_data=recorder.next()) + src_weights) + + recorder.add("eval_direct", timing_future) + + potentials = potentials + direct_result # }}} # {{{ "Stage 7:" propagate local_exps downward - wrangler.refine_locals( + local_exps, timing_future = wrangler.refine_locals( traversal.level_start_target_or_target_parent_box_nrs, traversal.target_or_target_parent_boxes, - local_exps, - timing_data=recorder.next()) + local_exps) + + recorder.add("refine_locals", timing_future) # }}} # {{{ "Stage 8:" evaluate locals - potentials = potentials + wrangler.eval_locals( + local_result, timing_future = wrangler.eval_locals( traversal.level_start_target_box_nrs, traversal.target_boxes, - local_exps, - timing_data=recorder.next()) + local_exps) + + recorder.add("eval_locals", timing_future) + + potentials = potentials + local_result # }}} @@ -200,9 +220,12 @@ class ExpansionWranglerInterface: Will usually hold a reference (and thereby be specific to) a :class:`boxtree.Tree` instance. - This interface supports collecting timing data. If timing data is requested, - the *timing_data* argument expects a :class:`TimingWaiter`. Otherwise - *timing_data* should be *None*. + Functions that support returning timing data return a value supporting the + :class:`TimingFuture` interface. + + .. versionchanged:: 2018.1 + + Changed (a subset of) functions to return timing data. """ def multipole_expansion_zeros(self): @@ -237,88 +260,89 @@ class ExpansionWranglerInterface: *source_weights* is in tree target order. """ - def form_multipoles(self, level_start_source_box_nrs, source_boxes, src_weights, - timing_data=None): + def form_multipoles(self, level_start_source_box_nrs, source_boxes, src_weights): """Return an expansions array (compatible with :meth:`multipole_expansion_zeros`) containing multipole expansions in *source_boxes* due to sources with *src_weights*. All other expansions must be zero. + + :return: A pair (*mpoles*, *timing_future*). """ def coarsen_multipoles(self, level_start_source_parent_box_nrs, - source_parent_boxes, mpoles, timing_data=None): + source_parent_boxes, mpoles): """For each box in *source_parent_boxes*, gather (and translate) the box's children's multipole expansions in *mpole* and add the resulting expansion into the box's multipole expansion in *mpole*. - :returns: *mpoles* + :returns: A pair (*mpoles*, *timing_future*). """ def eval_direct(self, target_boxes, neighbor_sources_starts, - neighbor_sources_lists, src_weights, timing_data=None): + neighbor_sources_lists, src_weights): """For each box in *target_boxes*, evaluate the influence of the neighbor sources due to *src_weights*, which use :ref:`csr` and are indexed like *target_boxes*. - :returns: a new potential array, see :meth:`output_zeros`. + :returns: A pair (*pot*, *timing_future*), where *pot* is a + a new potential array, see :meth:`output_zeros`. """ def multipole_to_local(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, - starts, lists, mpole_exps, timing_data=None): + starts, lists, mpole_exps): """For each box in *target_or_target_parent_boxes*, translate and add the influence of the multipole expansion in *mpole_exps* into a new array of local expansions. *starts* and *lists* use :ref:`csr`, and *starts* is indexed like *target_or_target_parent_boxes*. - :returns: a new (local) expansion array, see - :meth:`local_expansion_zeros`. + :returns: A pair (*pot*, *timing_future*) where *pot* is + a new (local) expansion array, see :meth:`local_expansion_zeros`. """ def eval_multipoles(self, - target_boxes_by_source_level, from_sep_smaller_by_level, mpole_exps, - timing_data=None): + target_boxes_by_source_level, from_sep_smaller_by_level, mpole_exps): """For a level *i*, each box in *target_boxes_by_source_level[i]*, evaluate the multipole expansion in *mpole_exps* in the nearby boxes given in *from_sep_smaller_by_level*, and return a new potential array. *starts* and *lists* in *from_sep_smaller_by_level[i]* use :ref:`csr` and *starts* is indexed like *target_boxes_by_source_level[i]*. - :returns: a new potential array, see :meth:`output_zeros`. + :returns: A pair (*pot*, *timing_future*) where *pot* is a new potential + array, see :meth:`output_zeros`. """ def form_locals(self, level_start_target_or_target_parent_box_nrs, - target_or_target_parent_boxes, starts, lists, src_weights, - timing_data=None): + target_or_target_parent_boxes, starts, lists, src_weights): """For each box in *target_or_target_parent_boxes*, form local expansions due to the sources in the nearby boxes given in *starts* and *lists*, and return a new local expansion array. *starts* and *lists* use :ref:`csr` and *starts* is indexed like *target_or_target_parent_boxes*. - :returns: a new local expansion array, see - :meth:`local_expansion_zeros`. + :returns: A pair (*pot*, *timing_future*) where *pot* is a new + local expansion array, see :meth:`local_expansion_zeros`. """ def refine_locals(self, level_start_target_or_target_parent_box_nrs, - target_or_target_parent_boxes, local_exps, timing_data=None): + target_or_target_parent_boxes, local_exps): """For each box in *child_boxes*, translate the box's parent's local expansion in *local_exps* and add the resulting expansion into the box's local expansion in *local_exps*. - :returns: *local_exps* + :returns: A pair (*local_exps*, *timing_future*). """ - def eval_locals(self, level_start_target_box_nrs, target_boxes, local_exps, - timing_data=None): + def eval_locals(self, level_start_target_box_nrs, target_boxes, local_exps): """For each box in *target_boxes*, evaluate the local expansion in *local_exps* and return a new potential array. - :returns: a new potential array, see :meth:`output_zeros`. + :returns: A pair (*pot*, *timing_future*) where *pot* is a new potential + array, see :meth:`output_zeros`. """ def finalize_potentials(self, potentials): @@ -335,8 +359,6 @@ class ExpansionWranglerInterface: class TimingResult(Record): """ - .. automethod:: __add__ - .. attribute:: wall_elapsed .. attribute:: process_elapsed """ @@ -346,61 +368,25 @@ class TimingResult(Record): wall_elapsed=wall_elapsed, process_elapsed=process_elapsed) - def __add__(self, other): - wall_elapsed = ( - self.wall_elapsed + other.wall_elapsed - if self.wall_elapsed is not None else None) - - process_elapsed = ( - self.process_elapsed + other.process_elapsed - if self.process_elapsed is not None else None) - - return TimingResult(wall_elapsed, process_elapsed) - # }}} -# {{{ timing waiter - -class TimingWaiter(object): - """Obtains timing data through a supplied callback function. - - Attributes that should be set: - - .. attribute:: description - - A string, the description of the timing data. +# {{{ timing future - .. attribute:: callback +class TimingFuture(object): + """Returns timing data for a potentially asynchronous operation. - Returns a :class:`TimingResult`. + .. automethod:: result + .. automethod:: done """ - def __init__(self): - self.description = None - self.callback = None - self._result = None - - @property - def empty(self): - return not self.callback - - @property def result(self): - """The timing data result obtained from *callback*.""" - if self._result is None: - self.wait() - - return self._result - - def wait(self): - if self.empty: - return + """Return a :class:`TimingResult`. May block.""" + raise NotImplementedError - callback_result = self.callback() - self._result = TimingResult( - callback_result.wall_elapsed, - callback_result.process_elapsed) + def done(self): + """Return *True* if the operation is complete.""" + raise NotImplementedError # }}} @@ -410,25 +396,34 @@ class TimingWaiter(object): class TimingRecorder(object): def __init__(self): - self.records = [] + from collections import defaultdict + self.futures = defaultdict(list) - def next(self): - self.records.append(TimingWaiter()) - return self.records[-1] + def add(self, description, future): + self.futures[description].append(future) + + def merge(self, result1, result2): + wall_elapsed = None + process_elapsed = None + + if None not in (result1.wall_elapsed, result2.wall_elapsed): + wall_elapsed = result1.wall_elapsed + result2.wall_elapsed + if None not in (result1.process_elapsed, result2.process_elapsed): + process_elapsed = result1.process_elapsed + result2.process_elapsed + + return TimingResult(wall_elapsed, process_elapsed) def summarize(self): result = {} - for record in self.records: - if record.empty: - continue + for description, futures_list in self.futures.items(): + futures = iter(futures_list) - description = record.description + timing_result = next(futures).result() + for future in futures: + timing_result = self.merge(timing_result, future.result()) - if description in result: - result[description] += record.result - else: - result[description] = record.result + result[description] = timing_result return result diff --git a/boxtree/pyfmmlib_integration.py b/boxtree/pyfmmlib_integration.py index 32737d5..c076e56 100644 --- a/boxtree/pyfmmlib_integration.py +++ b/boxtree/pyfmmlib_integration.py @@ -27,7 +27,7 @@ THE SOFTWARE. import numpy as np from pytools import memoize_method, log_process -from boxtree.tools import record_time +from boxtree.tools import return_timing_data import logging logger = logging.getLogger(__name__) @@ -421,7 +421,7 @@ class FMMLibExpansionWrangler(object): } @log_process(logger) - @record_time("timing_data") + @return_timing_data def form_multipoles(self, level_start_source_box_nrs, source_boxes, src_weights): formmp = self.get_routine("%ddformmp" + self.dp_suffix) @@ -461,7 +461,7 @@ class FMMLibExpansionWrangler(object): return mpoles @log_process(logger) - @record_time("timing_data") + @return_timing_data def coarsen_multipoles(self, level_start_source_parent_box_nrs, source_parent_boxes, mpoles): tree = self.tree @@ -514,8 +514,10 @@ class FMMLibExpansionWrangler(object): target_mpoles_view[ ibox - target_level_start_ibox] += new_mp[..., 0].T + return mpoles + @log_process(logger) - @record_time("timing_data") + @return_timing_data def eval_direct(self, target_boxes, neighbor_sources_starts, neighbor_sources_lists, src_weights): output = self.output_zeros() @@ -557,7 +559,7 @@ class FMMLibExpansionWrangler(object): return output @log_process(logger) - @record_time("timing_data") + @return_timing_data def multipole_to_local(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, @@ -643,7 +645,7 @@ class FMMLibExpansionWrangler(object): return local_exps @log_process(logger) - @record_time("timing_data") + @return_timing_data def eval_multipoles(self, target_boxes_by_source_level, sep_smaller_nonsiblings_by_level, mpole_exps): @@ -686,7 +688,7 @@ class FMMLibExpansionWrangler(object): return output @log_process(logger) - @record_time("timing_data") + @return_timing_data def form_locals(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, starts, lists, src_weights): @@ -738,7 +740,7 @@ class FMMLibExpansionWrangler(object): return local_exps @log_process(logger) - @record_time("timing_data") + @return_timing_data def refine_locals(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, local_exps): @@ -785,7 +787,7 @@ class FMMLibExpansionWrangler(object): return local_exps @log_process(logger) - @record_time("timing_data") + @return_timing_data def eval_locals(self, level_start_target_box_nrs, target_boxes, local_exps): output = self.output_zeros() taeval = self.get_expn_eval_routine("ta") diff --git a/boxtree/tools.py b/boxtree/tools.py index a89d7da..97cd43d 100644 --- a/boxtree/tools.py +++ b/boxtree/tools.py @@ -30,6 +30,7 @@ import pyopencl.array # noqa from pyopencl.tools import dtype_to_c_struct from mako.template import Template from pytools.obj_array import make_obj_array +from boxtree.fmm import TimingFuture, TimingResult AXIS_NAMES = ("x", "y", "z", "w") @@ -512,48 +513,45 @@ class MapValuesKernel(object): # {{{ time recording tool -class record_time(object): # noqa: N801 - """A decorator for recording timing data for a function call. +class DummyTimingFuture(TimingFuture): - This introduces an extra keyword argument to the decorated function. For the - newly added argument, the caller should pass either *None* or an instance of - :class:`boxtree.fmm.TimingWaiter`. If the latter gets passed, the fields - *callback* and *description* are populated with timing data for the function - call. + @classmethod + def from_timer(cls, timer): + return cls(timer.wall_elapsed, timer.process_elapsed) - """ + def __init__(self, wall_elapsed, process_elapsed): + self.wall_elapsed = wall_elapsed + self.process_elapsed = process_elapsed - def __init__(self, arg=None, description=None): - self.arg = arg - self.description = description + def result(self): + return TimingResult(self.wall_elapsed, self.process_elapsed) - def __call__(self, wrapped): - description = self.description or wrapped.__name__ + def done(self): + return True - from pytools import ProcessTimer - from contextlib import contextmanager - @contextmanager - def time_process(output): - timer = ProcessTimer() - yield - timer.done() +def return_timing_data(wrapped): + """A decorator for recording timing data for a function call. + + The decorated function returns a tuple (*retval*, *timing_future*) + where *retval* is the original return value and *timing_future* + supports the timing data future interface in :mod:`boxtree.fmm`. + """ - output.description = description - output.callback = lambda: timer + from pytools import ProcessTimer - def wrapper(*args, **kwargs): - output = kwargs.pop(self.arg, None) - if output is None: - return wrapped(*args, **kwargs) + def wrapper(*args, **kwargs): + timer = ProcessTimer() + retval = wrapped(*args, **kwargs) + timer.done() - with time_process(output): - return wrapped(*args, **kwargs) + future = DummyTimingFuture.from_timer(timer) + return (retval, future) - from functools import update_wrapper - new_wrapper = update_wrapper(wrapper, wrapped) + from functools import update_wrapper + new_wrapper = update_wrapper(wrapper, wrapped) - return new_wrapper + return new_wrapper # }}} diff --git a/doc/fmm.rst b/doc/fmm.rst index 446ef8b..97ea4e0 100644 --- a/doc/fmm.rst +++ b/doc/fmm.rst @@ -12,7 +12,7 @@ FMM driver .. autoclass:: TimingResult -.. autoclass:: TimingWaiter +.. autoclass:: TimingFuture Integration with PyFMMLib ------------------------- diff --git a/test/test_fmm.py b/test/test_fmm.py index 5e1b6c0..cb0b061 100644 --- a/test/test_fmm.py +++ b/test/test_fmm.py @@ -44,15 +44,7 @@ logger = logging.getLogger(__name__) # {{{ fmm interaction completeness test -def ignore_timing_data(f): - from functools import wraps - - @wraps(f) - def wrapper(*args, **kwargs): - kwargs.pop("timing_data") - return f(*args, **kwargs) - - return wrapper +from boxtree.tools import return_timing_data class ConstantOneExpansionWrangler(object): @@ -88,7 +80,7 @@ class ConstantOneExpansionWrangler(object): def reorder_potentials(self, potentials): return potentials[self.tree.sorted_target_ids] - @ignore_timing_data + @return_timing_data def form_multipoles(self, level_start_source_box_nrs, source_boxes, src_weights): mpoles = self.multipole_expansion_zeros() for ibox in source_boxes: @@ -97,7 +89,7 @@ class ConstantOneExpansionWrangler(object): return mpoles - @ignore_timing_data + @return_timing_data def coarsen_multipoles(self, level_start_source_parent_box_nrs, source_parent_boxes, mpoles): tree = self.tree @@ -117,7 +109,9 @@ class ConstantOneExpansionWrangler(object): if child: mpoles[ibox] += mpoles[child] - @ignore_timing_data + return mpoles + + @return_timing_data def eval_direct(self, target_boxes, neighbor_sources_starts, neighbor_sources_lists, src_weights): pot = self.potential_zeros() @@ -137,7 +131,7 @@ class ConstantOneExpansionWrangler(object): return pot - @ignore_timing_data + @return_timing_data def multipole_to_local(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, @@ -156,7 +150,7 @@ class ConstantOneExpansionWrangler(object): return local_exps - @ignore_timing_data + @return_timing_data def eval_multipoles(self, target_boxes_by_source_level, from_sep_smaller_nonsiblings_by_level, mpole_exps): @@ -177,7 +171,7 @@ class ConstantOneExpansionWrangler(object): return pot - @ignore_timing_data + @return_timing_data def form_locals(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, starts, lists, src_weights): @@ -197,7 +191,7 @@ class ConstantOneExpansionWrangler(object): return local_exps - @ignore_timing_data + @return_timing_data def refine_locals(self, level_start_target_or_target_parent_box_nrs, target_or_target_parent_boxes, local_exps): @@ -209,7 +203,7 @@ class ConstantOneExpansionWrangler(object): return local_exps - @ignore_timing_data + @return_timing_data def eval_locals(self, level_start_target_box_nrs, target_boxes, local_exps): pot = self.potential_zeros() -- GitLab