From 6f9beba75b27d2a5a03d07b54e0c03f482aeaf41 Mon Sep 17 00:00:00 2001 From: mloubout Date: Fri, 15 Sep 2023 09:19:22 -0400 Subject: [PATCH] api: process injected expression dimensions in case it's not the sparse function --- devito/builtins/initializers.py | 7 ++++++- devito/core/autotuning.py | 4 +--- devito/operations/interpolators.py | 30 ++++++++++++++++++++---------- devito/operator/operator.py | 8 ++++++++ devito/tools/data_structures.py | 3 +++ devito/types/dimension.py | 15 ++++++--------- tests/test_buffering.py | 2 +- tests/test_dimension.py | 3 +-- tests/test_dle.py | 11 ++++++++--- tests/test_dse.py | 28 +++++++++++++++------------- tests/test_interpolation.py | 26 ++++++++++++++++++++++++++ tests/test_mpi.py | 4 ++-- 12 files changed, 97 insertions(+), 44 deletions(-) diff --git a/devito/builtins/initializers.py b/devito/builtins/initializers.py index f338e194e1e..83bad735fae 100644 --- a/devito/builtins/initializers.py +++ b/devito/builtins/initializers.py @@ -77,7 +77,12 @@ def assign(f, rhs=0, options=None, name='assign', assign_halo=False, **kwargs): symbolic_max=d.symbolic_max + h.right) eqs = [eq.xreplace(subs) for eq in eqs] - dv.Operator(eqs, name=name, **kwargs)() + op = dv.Operator(eqs, name=name, **kwargs) + try: + op() + except ValueError: + # Corner case such as assign(u, v) with v a Buffered TimeFunction + op(time_M=f._time_size) def smooth(f, g, axis=None): diff --git a/devito/core/autotuning.py b/devito/core/autotuning.py index 16b146721ae..f30c020ef7e 100644 --- a/devito/core/autotuning.py +++ b/devito/core/autotuning.py @@ -209,9 +209,7 @@ 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: - if isinstance(args[dim.size_name], range): - pass - else: + if not isinstance(args[dim.size_name], range): # 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]: diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py index 92bc3923926..dae96d8dfe5 100644 --- a/devito/operations/interpolators.py +++ b/devito/operations/interpolators.py @@ -169,11 +169,17 @@ def _rdim(self): return DimensionTuple(*rdims, getters=self._gdims) - def _augment_implicit_dims(self, implicit_dims): + def _augment_implicit_dims(self, implicit_dims, extras=None): + if extras is not None: + extra = set([i for v in extras for i in v.dimensions]) - set(self._gdims) + extra = tuple(extra) + else: + extra = tuple() + if self.sfunction._sparse_position == -1: - return self.sfunction.dimensions + as_tuple(implicit_dims) + return self.sfunction.dimensions + as_tuple(implicit_dims) + extra else: - return as_tuple(implicit_dims) + self.sfunction.dimensions + return as_tuple(implicit_dims) + self.sfunction.dimensions + extra def _coeff_temps(self, implicit_dims): return [] @@ -252,8 +258,6 @@ def _interpolate(self, expr, increment=False, self_subs={}, implicit_dims=None): interpolation expression, but that should be honored when constructing the operator. """ - implicit_dims = self._augment_implicit_dims(implicit_dims) - # Derivatives must be evaluated before the introduction of indirect accesses try: _expr = expr.evaluate @@ -263,6 +267,9 @@ def _interpolate(self, expr, increment=False, self_subs={}, implicit_dims=None): variables = list(retrieve_function_carriers(_expr)) + # Implicit dimensions + implicit_dims = self._augment_implicit_dims(implicit_dims) + # List of indirection indices for all adjacent grid points idx_subs, temps = self._interp_idx(variables, implicit_dims=implicit_dims) @@ -295,8 +302,6 @@ 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) - # Make iterable to support inject((u, v), expr=expr) # or inject((u, v), expr=(expr1, expr2)) fields, exprs = as_tuple(field), as_tuple(expr) @@ -315,6 +320,10 @@ def _inject(self, field, expr, implicit_dims=None): _exprs = exprs variables = list(v for e in _exprs for v in retrieve_function_carriers(e)) + + # Implicit dimensions + implicit_dims = self._augment_implicit_dims(implicit_dims, variables) + variables = variables + list(fields) # List of indirection indices for all adjacent grid points @@ -385,6 +394,7 @@ 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.parent.symbolic_min}) - for (ri, rd) in enumerate(self._rdim)]) + mappers = [{ddim: ri, cdim: rd-rd.parent.symbolic_min} + for (ri, rd) in enumerate(self._rdim)] + return Mul(*[self.interpolation_coeffs.subs(mapper) + for mapper in mappers]) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index b91d51727fa..609c69295f6 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -566,6 +566,7 @@ def _prepare_arguments(self, autotune=None, **kwargs): "`%s=%s`, while `%s=%s` is expected. Perhaps you " "forgot to override `%s`?" % (p, k, v, k, args[k], p)) + args = kwargs['args'] = args.reduce_all() # DiscreteFunctions may be created from CartesianDiscretizations, which in @@ -573,6 +574,10 @@ def _prepare_arguments(self, autotune=None, **kwargs): discretizations = {getattr(kwargs[p.name], 'grid', None) for p in overrides} discretizations.update({getattr(p, 'grid', None) for p in defaults}) discretizations.discard(None) + # Remove subgrids if multiple grids + if len(discretizations) > 1: + discretizations = {g for g in discretizations + if not any(d.is_Derived for d in g.dimensions)} for i in discretizations: args.update(i._arg_values(**kwargs)) @@ -585,6 +590,9 @@ def _prepare_arguments(self, autotune=None, **kwargs): if configuration['mpi']: raise ValueError("Multiple Grids found") try: + # Take biggest grid, i.e discard grids with subdimensions + grids = {g for g in grids if not any(d.is_Sub for d in g.dimensions)} + # First grid as there is no heuristic on how to choose from the leftovers grid = grids.pop() except KeyError: grid = None diff --git a/devito/tools/data_structures.py b/devito/tools/data_structures.py index 3afe7197eb7..01a9a3f4bd8 100644 --- a/devito/tools/data_structures.py +++ b/devito/tools/data_structures.py @@ -137,6 +137,9 @@ def compare_to_first(v): return candidates[0] elif all(map(compare_to_first, candidates)): # return first non-range + for c in candidates: + if not isinstance(c, range): + return c return candidates[0] else: raise ValueError("Unable to find unique value for key %s, candidates: %s" diff --git a/devito/types/dimension.py b/devito/types/dimension.py index 6044f014690..2d11ccb2206 100644 --- a/devito/types/dimension.py +++ b/devito/types/dimension.py @@ -298,14 +298,14 @@ 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 = loc_minv.start + loc_minv = loc_minv.stop except AttributeError: try: loc_minv = sorted(loc_minv).pop(0) except TypeError: pass try: - loc_maxv = loc_maxv.start + loc_maxv = loc_maxv.stop except AttributeError: try: loc_maxv = sorted(loc_maxv).pop(0) @@ -859,7 +859,8 @@ 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] = range(1, factor*(size)) + + defaults[dim.parent.max_name] = range(1, factor*size - 1) return defaults @@ -983,8 +984,7 @@ def bound_symbols(self): return set(self.parent.bound_symbols) 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)} + return {} def _arg_values(self, *args, **kwargs): return {} @@ -1466,10 +1466,7 @@ def _arg_defaults(self, _min=None, size=None, **kwargs): A SteppingDimension does not know its max point and therefore does not have a size argument. """ - args = {self.parent.min_name: _min} - if size: - args[self.parent.size_name] = range(size-1, np.iinfo(np.int32).max) - return args + return {self.parent.min_name: _min} def _arg_values(self, *args, **kwargs): """ diff --git a/tests/test_buffering.py b/tests/test_buffering.py index 16f98b4f940..ba200d220c7 100644 --- a/tests/test_buffering.py +++ b/tests/test_buffering.py @@ -272,7 +272,7 @@ def test_over_injection(): # Check generated code assert len(retrieve_iteration_tree(op1)) == \ - 7 + int(configuration['language'] != 'C') + 8 + int(configuration['language'] != 'C') buffers = [i for i in FindSymbols().visit(op1) if i.is_Array] assert len(buffers) == 1 diff --git a/tests/test_dimension.py b/tests/test_dimension.py index 32da3b22e30..9d41dddf484 100644 --- a/tests/test_dimension.py +++ b/tests/test_dimension.py @@ -1515,7 +1515,7 @@ def test_issue_1927(self, factor): op = Operator(Eq(f, 1)) - assert op.arguments()['time_M'] == 4*(save-1) # == min legal endpoint + assert op.arguments()['time_M'] == 4*save-1 # == min legal endpoint # Also no issues when supplying an override assert op.arguments(time_M=10)['time_M'] == 10 @@ -1530,7 +1530,6 @@ def test_issue_1927_v2(self): i = Dimension(name='i') ci = ConditionalDimension(name='ci', parent=i, factor=factor) - g = Function(name='g', shape=(size,), dimensions=(i,)) f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,)) diff --git a/tests/test_dle.py b/tests/test_dle.py index 86a288ac008..3b9883e6652 100644 --- a/tests/test_dle.py +++ b/tests/test_dle.py @@ -187,9 +187,14 @@ def test_cache_blocking_structure_optrelax(): op = Operator(eqns, opt=('advanced', {'blockrelax': True})) - bns, _ = assert_blocking(op, {'x0_blk0', 'p_src0_blk0'}) + bns, _ = assert_blocking(op, {'x0_blk0', 'p_src0_blk0', 'p_src1_blk0'}) iters = FindNodes(Iteration).visit(bns['p_src0_blk0']) + assert len(iters) == 2 + assert iters[0].dim.is_Block + assert iters[1].dim.is_Block + + iters = FindNodes(Iteration).visit(bns['p_src1_blk0']) assert len(iters) == 5 assert iters[0].dim.is_Block assert iters[1].dim.is_Block @@ -286,7 +291,7 @@ def test_cache_blocking_structure_optrelax_prec_inject(): 'openmp': True, 'par-collapse-ncores': 1})) - assert_structure(op, ['t', 't,p_s0_blk0,p_s,rsx,rsy'], + assert_structure(op, ['t', 't,p_s0_blk0,p_s', 't,p_s0_blk0,p_s,rsx,rsy'], 't,p_s0_blk0,p_s,rsx,rsy') @@ -958,7 +963,7 @@ def test_parallel_prec_inject(self): iterations = FindNodes(Iteration).visit(op0) assert not iterations[0].pragmas - assert 'omp for collapse(2)' in iterations[1].pragmas[0].value + assert 'omp for' in iterations[1].pragmas[0].value class TestNestedParallelism(object): diff --git a/tests/test_dse.py b/tests/test_dse.py index 2aefe69ed49..1e18157c77e 100644 --- a/tests/test_dse.py +++ b/tests/test_dse.py @@ -48,9 +48,9 @@ def test_scheduling_after_rewrite(): trees = retrieve_iteration_tree(op) # Check loop nest structure - assert all(i.dim is j for i, j in zip(trees[0], grid.dimensions)) # time invariant - assert trees[1].root.dim is grid.time_dim - assert all(trees[1].root.dim is tree.root.dim for tree in trees[1:]) + assert all(i.dim is j for i, j in zip(trees[1], grid.dimensions)) # time invariant + assert trees[2].root.dim is grid.time_dim + assert all(trees[2].root.dim is tree.root.dim for tree in trees[2:]) @pytest.mark.parametrize('exprs,expected,min_cost', [ @@ -1687,7 +1687,7 @@ def test_drop_redundants_after_fusion(self, rotate): op = Operator(eqns, opt=('advanced', {'cire-rotate': rotate})) arrays = [i for i in FindSymbols().visit(op) if i.is_Array] - assert len(arrays) == 2 + assert len(arrays) == 4 assert all(i._mem_heap and not i._mem_external for i in arrays) def test_full_shape_big_temporaries(self): @@ -2711,9 +2711,11 @@ def test_fullopt(self): assert np.isclose(summary0[('section0', None)].oi, 2.851, atol=0.001) assert summary1[('section0', None)].ops == 9 - assert summary1[('section1', None)].ops == 31 - assert summary1[('section2', None)].ops == 88 - assert np.isclose(summary1[('section1', None)].oi, 1.767, atol=0.001) + assert summary1[('section1', None)].ops == 9 + assert summary1[('section2', None)].ops == 31 + assert summary1[('section3', None)].ops == 26 + assert summary1[('section4', None)].ops == 22 + assert np.isclose(summary1[('section2', None)].oi, 1.767, atol=0.001) assert np.allclose(u0.data, u1.data, atol=10e-5) assert np.allclose(rec0.data, rec1.data, atol=10e-5) @@ -2773,8 +2775,8 @@ def test_fullopt(self): assert np.allclose(self.tti_noopt[1].data, rec.data, atol=10e-1) # Check expected opcount/oi - assert summary[('section2', None)].ops == 92 - assert np.isclose(summary[('section2', None)].oi, 2.074, atol=0.001) + assert summary[('section3', None)].ops == 92 + assert np.isclose(summary[('section3', None)].oi, 2.074, atol=0.001) # With optimizations enabled, there should be exactly four BlockDimensions op = wavesolver.op_fwd() @@ -2792,7 +2794,7 @@ def test_fullopt(self): # 3 Arrays are defined globally for the sparse positions temporaries # and two additional bock-sized Arrays are defined locally arrays = [i for i in FindSymbols().visit(op) if i.is_Array] - extra_arrays = 2+3 + extra_arrays = 2+3+3 assert len(arrays) == 4 + extra_arrays assert all(i._mem_heap and not i._mem_external for i in arrays) bns, pbs = assert_blocking(op, {'x0_blk0'}) @@ -2828,7 +2830,7 @@ def test_fullopt_w_mpi(self): def test_opcounts(self, space_order, expected): op = self.tti_operator(opt='advanced', space_order=space_order) sections = list(op.op_fwd()._profiler._sections.values()) - assert sections[2].sops == expected + assert sections[3].sops == expected @switchconfig(profiling='advanced') @pytest.mark.parametrize('space_order,expected', [ @@ -2838,8 +2840,8 @@ def test_opcounts_adjoint(self, space_order, expected): wavesolver = self.tti_operator(opt=('advanced', {'openmp': False})) op = wavesolver.op_adj() - assert op._profiler._sections['section2'].sops == expected - assert len([i for i in FindSymbols().visit(op) if i.is_Array]) == 7+3 + assert op._profiler._sections['section3'].sops == expected + assert len([i for i in FindSymbols().visit(op) if i.is_Array]) == 7+3+3 class TestTTIv2(object): diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 3a22ca1db73..c7a15665a48 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -734,3 +734,29 @@ class SparseFirst(SparseFunction): op(time_M=10) expected = 10*11/2 # n (n+1)/2 assert np.allclose(s.data, expected) + + +def test_inject_function(): + nt = 11 + + grid = Grid(shape=(5, 5)) + u = TimeFunction(name="u", grid=grid, time_order=2) + src = SparseTimeFunction(name="src", grid=grid, nt=nt, npoint=1, + coordinates=[[0.5, 0.5]]) + + nfreq = 5 + freq_dim = DefaultDimension(name="freq", default_value=nfreq) + omega = Function(name="omega", dimensions=(freq_dim,), shape=(nfreq,), grid=grid) + omega.data.fill(1.) + + inj = src.inject(field=u.forward, expr=omega) + + op = Operator([inj]) + + op(time_M=0) + assert u.data[1, 2, 2] == nfreq + assert np.all(u.data[0] == 0) + assert np.all(u.data[2] == 0) + for i in [0, 1, 3, 4]: + for j in [0, 1, 3, 4]: + assert u.data[1, i, j] == 0 diff --git a/tests/test_mpi.py b/tests/test_mpi.py index 14ddbec249b..2860fc726e4 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -2499,8 +2499,8 @@ def test_adjoint_codegen(self, shape, kernel, space_order, save): op_adj = solver.op_adj() adj_calls = FindNodes(Call).visit(op_adj) - # one halo, ndim memalign and free (pos temp rec) - sf_calls = 2 * len(shape) + # one halo, ndim memalign and free (pos temp rec/src) + sf_calls = 2 * len(shape) * 2 assert len(fwd_calls) == 1 + sf_calls assert len(adj_calls) == 1 + sf_calls