From d715f3f310229a9427de21785d41f29614b3a8a0 Mon Sep 17 00:00:00 2001 From: mloubout Date: Mon, 11 Sep 2023 09:35:15 -0400 Subject: [PATCH] api: always use conditional dimension for interpolation radius dim --- devito/core/autotuning.py | 7 +++++-- devito/operations/interpolators.py | 28 +++++++++++++++---------- devito/operator/operator.py | 5 +++-- devito/tools/data_structures.py | 9 ++++++++ devito/tools/utils.py | 9 +++++++- devito/types/basic.py | 4 +++- devito/types/dense.py | 2 +- devito/types/dimension.py | 33 +++++++++++++++++++----------- tests/test_operator.py | 6 ++++-- 9 files changed, 71 insertions(+), 32 deletions(-) diff --git a/devito/core/autotuning.py b/devito/core/autotuning.py index 603a78efd66..16b146721ae 100644 --- a/devito/core/autotuning.py +++ b/devito/core/autotuning.py @@ -209,8 +209,11 @@ def init_time_bounds(stepper, at_args, args): else: at_args[dim.max_name] = at_args[dim.min_name] + options['squeezer'] if dim.size_name in args: - # May need to shrink to avoid OOB accesses - at_args[dim.max_name] = min(at_args[dim.max_name], args[dim.max_name]) + if isinstance(args[dim.size_name], range): + pass + else: + # May need to shrink to avoid OOB accesses + at_args[dim.max_name] = min(at_args[dim.max_name], args[dim.max_name]) if at_args[dim.min_name] > at_args[dim.max_name]: warning("too few time iterations; skipping") return False diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py index ab758ba9820..92bc3923926 100644 --- a/devito/operations/interpolators.py +++ b/devito/operations/interpolators.py @@ -155,7 +155,19 @@ def _rdim(self): -self.r+1, self.r, 2*self.r, parent) for d in self._gdims] - return DimensionTuple(*dims, getters=self._gdims) + # Make radius dimension conditional to avoid OOB + rdims = [] + pos = self.sfunction._position_map.values() + + for (d, rd, p) in zip(self._gdims, dims, pos): + # Add conditional to avoid OOB + lb = sympy.And(rd + p >= d.symbolic_min - self.r, evaluate=False) + ub = sympy.And(rd + p <= d.symbolic_max + self.r, evaluate=False) + cond = sympy.And(lb, ub, evaluate=False) + rdims.append(ConditionalDimension(rd.name, rd, condition=cond, + indirect=True)) + + return DimensionTuple(*rdims, getters=self._gdims) def _augment_implicit_dims(self, implicit_dims): if self.sfunction._sparse_position == -1: @@ -177,13 +189,6 @@ def _interp_idx(self, variables, implicit_dims=None): mapper = {} pos = self.sfunction._position_map.values() - for ((di, d), rd, p) in zip(enumerate(self._gdims), self._rdim, pos): - # Add conditional to avoid OOB - lb = sympy.And(rd + p >= d.symbolic_min - self.r, evaluate=False) - ub = sympy.And(rd + p <= d.symbolic_max + self.r, evaluate=False) - cond = sympy.And(lb, ub, evaluate=False) - mapper[d] = ConditionalDimension(rd.name, rd, condition=cond, indirect=True) - # Temporaries for the position temps = self._positions(implicit_dims) @@ -191,10 +196,10 @@ def _interp_idx(self, variables, implicit_dims=None): temps.extend(self._coeff_temps(implicit_dims)) # Substitution mapper for variables + mapper = self._rdim._getters idx_subs = {v: v.subs({k: c + p for ((k, c), p) in zip(mapper.items(), pos)}) for v in variables} - idx_subs.update(dict(zip(self._rdim, mapper.values()))) return idx_subs, temps @@ -290,7 +295,7 @@ def _inject(self, field, expr, implicit_dims=None): injection expression, but that should be honored when constructing the operator. """ - implicit_dims = self._augment_implicit_dims(implicit_dims) + self._rdim + implicit_dims = self._augment_implicit_dims(implicit_dims) # Make iterable to support inject((u, v), expr=expr) # or inject((u, v), expr=(expr1, expr2)) @@ -380,5 +385,6 @@ def interpolation_coeffs(self): @property def _weights(self): ddim, cdim = self.interpolation_coeffs.dimensions[1:] - return Mul(*[self.interpolation_coeffs.subs({ddim: ri, cdim: rd-rd.symbolic_min}) + return Mul(*[self.interpolation_coeffs.subs({ddim: ri, + cdim: rd-rd.parent.symbolic_min}) for (ri, rd) in enumerate(self._rdim)]) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index d1bee9fa665..b91d51727fa 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -24,7 +24,7 @@ from devito.symbolics import estimate_cost from devito.tools import (DAG, OrderedSet, Signer, ReducerMap, as_tuple, flatten, filter_sorted, frozendict, is_integer, split, timed_pass, - timed_region) + timed_region, contains_val) from devito.types import Grid, Evaluable __all__ = ['Operator'] @@ -526,6 +526,7 @@ def _prepare_arguments(self, autotune=None, **kwargs): edges = [(i, i.parent) for i in self.dimensions if i.is_Derived and i.parent in set(nodes)] toposort = DAG(nodes, edges).topological_sort() + futures = {} for d in reversed(toposort): if set(d._arg_names).intersection(kwargs): @@ -560,7 +561,7 @@ def _prepare_arguments(self, autotune=None, **kwargs): # a TimeFunction `usave(t_sub, x, y)`, an override for `fact` is # supplied w/o overriding `usave`; that's legal pass - elif is_integer(args[k]) and args[k] not in as_tuple(v): + elif is_integer(args[k]) and not contains_val(args[k], v): raise ValueError("Default `%s` is incompatible with other args as " "`%s=%s`, while `%s=%s` is expected. Perhaps you " "forgot to override `%s`?" % diff --git a/devito/tools/data_structures.py b/devito/tools/data_structures.py index 539f75d5934..3afe7197eb7 100644 --- a/devito/tools/data_structures.py +++ b/devito/tools/data_structures.py @@ -110,6 +110,7 @@ def unique(self, key): Key for which to retrieve a unique value. """ candidates = self.getall(key) + candidates = [c for c in candidates if c is not None] def compare_to_first(v): first = candidates[0] @@ -122,12 +123,20 @@ def compare_to_first(v): return first in v elif isinstance(first, Set): return v in first + elif isinstance(v, range): + if isinstance(first, range): + return first.stop > v.start or v.stop > first.start + else: + return first >= v.start and first < v.stop + elif isinstance(first, range): + return v >= first.start and v < first.stop else: return first == v if len(candidates) == 1: return candidates[0] elif all(map(compare_to_first, candidates)): + # return first non-range return candidates[0] else: raise ValueError("Unable to find unique value for key %s, candidates: %s" diff --git a/devito/tools/utils.py b/devito/tools/utils.py index 16f09879304..d99bf34b25e 100644 --- a/devito/tools/utils.py +++ b/devito/tools/utils.py @@ -12,7 +12,7 @@ 'roundm', 'powerset', 'invert', 'flatten', 'single_or', 'filter_ordered', 'as_mapper', 'filter_sorted', 'pprint', 'sweep', 'all_equal', 'as_list', 'indices_to_slices', 'indices_to_sections', 'transitive_closure', - 'humanbytes'] + 'humanbytes', 'contains_val'] def prod(iterable, initial=1): @@ -75,6 +75,13 @@ def is_integer(value): return isinstance(value, (int, np.integer, sympy.Integer)) +def contains_val(val, items): + try: + return val in items + except TypeError: + return val == items + + def generator(): """ Return a function ``f`` that generates integer numbers starting at 0 diff --git a/devito/types/basic.py b/devito/types/basic.py index c36496d195d..400232faed2 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -838,13 +838,15 @@ def __new__(cls, *args, **kwargs): newobj._dimensions = dimensions newobj._shape = cls.__shape_setup__(**kwargs) newobj._dtype = cls.__dtype_setup__(**kwargs) - newobj.__init_finalize__(*args, **kwargs) # All objects created off an existing AbstractFunction `f` (e.g., # via .func, or .subs, such as `f(x + 1)`) keep a reference to `f` # through the `function` field newobj.function = function or newobj + # Initialization + newobj.__init_finalize__(*args, **kwargs) + return newobj def __init__(self, *args, **kwargs): diff --git a/devito/types/dense.py b/devito/types/dense.py index ec371b662c9..2a13fa91af4 100644 --- a/devito/types/dense.py +++ b/devito/types/dense.py @@ -94,7 +94,7 @@ def __init_finalize__(self, *args, function=None, **kwargs): # a reference to the user-provided buffer self._initializer = None if len(initializer) > 0: - self.data_with_halo[:] = initializer + self.data_with_halo[:] = initializer[:] else: # This is a corner case -- we might get here, for example, when # running with MPI and some processes get 0-size arrays after diff --git a/devito/types/dimension.py b/devito/types/dimension.py index 76d9d9e60a1..6044f014690 100644 --- a/devito/types/dimension.py +++ b/devito/types/dimension.py @@ -298,13 +298,19 @@ def _arg_values(self, interval, grid=None, args=None, **kwargs): # may represent sets of legal values. If that's the case, here we just # pick one. Note that we sort for determinism try: - loc_minv = sorted(loc_minv).pop(0) - except TypeError: - pass + loc_minv = loc_minv.start + except AttributeError: + try: + loc_minv = sorted(loc_minv).pop(0) + except TypeError: + pass try: - loc_maxv = sorted(loc_maxv).pop(0) - except TypeError: - pass + loc_maxv = loc_maxv.start + except AttributeError: + try: + loc_maxv = sorted(loc_maxv).pop(0) + except TypeError: + pass return {self.min_name: loc_minv, self.max_name: loc_maxv} @@ -853,8 +859,7 @@ def _arg_defaults(self, _min=None, size=None, alias=None): factor = defaults[dim._factor.name] = dim._factor.data except AttributeError: factor = dim._factor - defaults[dim.parent.max_name] = \ - frozenset(range(factor*(size - 1), factor*(size))) + defaults[dim.parent.max_name] = range(1, factor*(size)) return defaults @@ -977,8 +982,9 @@ def symbolic_incr(self): def bound_symbols(self): return set(self.parent.bound_symbols) - def _arg_defaults(self, **kwargs): - return {} + def _arg_defaults(self, alias=None, **kwargs): + dim = alias or self + return {dim.parent.size_name: range(self.symbolic_size, np.iinfo(np.int64).max)} def _arg_values(self, *args, **kwargs): return {} @@ -1446,7 +1452,7 @@ def symbolic_max(self): def _arg_names(self): return (self.min_name, self.max_name, self.name) + self.parent._arg_names - def _arg_defaults(self, _min=None, **kwargs): + def _arg_defaults(self, _min=None, size=None, **kwargs): """ A map of default argument values defined by this dimension. @@ -1460,7 +1466,10 @@ def _arg_defaults(self, _min=None, **kwargs): A SteppingDimension does not know its max point and therefore does not have a size argument. """ - return {self.parent.min_name: _min} + args = {self.parent.min_name: _min} + if size: + args[self.parent.size_name] = range(size-1, np.iinfo(np.int32).max) + return args def _arg_values(self, *args, **kwargs): """ diff --git a/tests/test_operator.py b/tests/test_operator.py index f38ac01942a..4f8228bc24a 100644 --- a/tests/test_operator.py +++ b/tests/test_operator.py @@ -707,6 +707,8 @@ def verify_arguments(self, arguments, expected): if isinstance(v, (Function, SparseFunction)): condition = v._C_as_ndarray(arguments[name])[v._mask_domain] == v.data condition = condition.all() + elif isinstance(arguments[name], range): + condition = arguments[name].start <= v < arguments[name].stop else: condition = arguments[name] == v @@ -1803,7 +1805,7 @@ def test_scheduling_sparse_functions(self): # `trees` than 6 op = Operator([eqn1] + eqn2 + [eqn3] + eqn4, opt=('noop', {'openmp': False})) trees = retrieve_iteration_tree(op) - assert len(trees) == 5 + assert len(trees) == 6 # Time loop not shared due to the WAR assert trees[0][0].dim is time and trees[0][0] is trees[1][0] # this IS shared assert trees[1][0] is not trees[3][0] @@ -1813,7 +1815,7 @@ def test_scheduling_sparse_functions(self): eqn2 = sf1.inject(u1.forward, expr=sf1) op = Operator([eqn1] + eqn2 + [eqn3] + eqn4, opt=('noop', {'openmp': False})) trees = retrieve_iteration_tree(op) - assert len(trees) == 5 + assert len(trees) == 6 assert all(trees[0][0] is i[0] for i in trees) def test_scheduling_with_free_dims(self):