diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 59f7237fcaa..ec6381698e3 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -40,7 +40,10 @@ def iet_build(stree): nsections += 1 elif i.is_Halo: - body = HaloSpot(queues.pop(i), i.halo_scheme) + try: + body = HaloSpot(queues.pop(i), i.halo_scheme) + except KeyError: + body = HaloSpot(None, i.halo_scheme) elif i.is_Sync: body = SyncSpot(i.sync_ops, body=queues.pop(i, None)) diff --git a/devito/ir/stree/algorithms.py b/devito/ir/stree/algorithms.py index b679ef0d3a1..4f7dae168ef 100644 --- a/devito/ir/stree/algorithms.py +++ b/devito/ir/stree/algorithms.py @@ -81,10 +81,13 @@ def stree_build(clusters, profiler=None, **kwargs): tip = augment_whole_subtree(c, tip, mapper, it) # Attach NodeHalo if necessary - for it, v in mapper.items(): + for it, v in reversed(mapper.items()): if needs_nodehalo(it.dim, c.halo_scheme): v.bottom.parent = NodeHalo(c.halo_scheme, v.bottom.parent) break + if needs_nodehalo_dim(it.dim, c.halo_scheme): + v.bottom.children = [NodeHalo(c.halo_scheme, v.bottom.parent)] + break # Add in NodeExprs exprs = [] @@ -182,7 +185,19 @@ def preprocess(clusters, options=None, **kwargs): processed.append(c.rebuild(exprs=[], ispace=ispace, syncs=syncs)) halo_scheme = HaloScheme.union([c1.halo_scheme for c1 in found]) - processed.append(c.rebuild(halo_scheme=halo_scheme)) + if halo_scheme: + itdims = set(c.ispace.dimensions) + hispace = c.ispace.project(itdims - halo_scheme.distributed_aindices) + else: + hispace = None + + if hispace: + processed.append(c.rebuild(exprs=[], + ispace=hispace, + halo_scheme=halo_scheme)) + processed.append(c.rebuild(halo_scheme=None)) + else: + processed.append(c.rebuild(halo_scheme=halo_scheme)) # Sanity check! try: @@ -229,6 +244,10 @@ def needs_nodehalo(d, hs): return d and hs and d._defines.intersection(hs.distributed_aindices) +def needs_nodehalo_dim(d, hs): + return d and hs and d._defines.intersection(hs.loc_indices) + + def reuse_section(candidate, section): try: if not section or candidate.siblings[-1] is not section: diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py index bee8bc0de48..98e0105e8aa 100644 --- a/devito/operations/interpolators.py +++ b/devito/operations/interpolators.py @@ -195,7 +195,7 @@ def _augment_implicit_dims(self, implicit_dims, extras=None): idims = self.sfunction.dimensions + as_tuple(implicit_dims) + extra else: idims = extra + as_tuple(implicit_dims) + self.sfunction.dimensions - return tuple(filter_ordered(idims)) + return tuple(idims) def _coeff_temps(self, implicit_dims): return [] diff --git a/devito/passes/iet/mpi.py b/devito/passes/iet/mpi.py index 00d96213aa4..dfea62089c2 100644 --- a/devito/passes/iet/mpi.py +++ b/devito/passes/iet/mpi.py @@ -21,7 +21,7 @@ def optimize_halospots(iet, **kwargs): around in order to improve the halo exchange performance. """ iet = _drop_halospots(iet) - iet = _hoist_halospots(iet) + # iet = _hoist_halospots(iet) iet = _merge_halospots(iet) iet = _drop_if_unwritten(iet, **kwargs) iet = _mark_overlappable(iet) diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index efde8674024..506a3b88def 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -703,8 +703,8 @@ class SparseFirst(SparseFunction): ds = DefaultDimension("ps", default_value=3) grid = Grid((11, 11)) dims = grid.dimensions - s = SparseFirst(name="s", grid=grid, npoint=2, dimensions=(dr, ds), shape=(2, 3)) - s.coordinates.data[:] = [[.5, .5], [.2, .2]] + s = SparseFirst(name="s", grid=grid, npoint=2, dimensions=(dr, ds), shape=(2, 3), + coordinates=[[.5, .5], [.2, .2]]) # Check dimensions and shape are correctly initialized assert s.indices[s._sparse_position] == dr diff --git a/tests/test_mpi.py b/tests/test_mpi.py index 25f85d29229..2b62774e51a 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -601,6 +601,41 @@ def test_precomputed_sparse(self, r): Operator(sf1.interpolate(u))() assert np.all(sf1.data == 4) + @pytest.mark.parallel(mode=4) + def test_sparse_first(self): + """ + Tests custom sprase function with sparse dimension as first index. + """ + + class SparseFirst(SparseFunction): + """ Custom sparse class with the sparse dimension as the first one""" + _sparse_position = 0 + + dr = Dimension("cd") + ds = DefaultDimension("ps", default_value=3) + grid = Grid((11, 11)) + dims = grid.dimensions + s = SparseFirst(name="s", grid=grid, npoint=2, dimensions=(dr, ds), shape=(2, 3), + coordinates=[[.5, .5], [.2, .2]]) + + # Check dimensions and shape are correctly initialized + assert s.indices[s._sparse_position] == dr + assert s.shape == (2, 3) + assert s.coordinates.indices[0] == dr + + # Operator + u = TimeFunction(name="u", grid=grid, time_order=1) + fs = Function(name="fs", grid=grid, dimensions=(*dims, ds), shape=(11, 11, 3)) + + eqs = [Eq(u.forward, u+1), Eq(fs, u)] + # No time dependence so need the implicit dim + rec = s.interpolate(expr=s+fs, implicit_dims=grid.stepping_dim) + op = Operator(eqs + rec) + + op(time_M=10) + expected = 10*11/2 # n (n+1)/2 + assert np.allclose(s.data, expected) + @pytest.mark.parallel(mode=4) def test_no_grid_dim_slow(self): shape = (12, 13, 14) @@ -624,7 +659,6 @@ class CoordSlowSparseFunction(SparseFunction): rec_eq = s.interpolate(expr=u) op = Operator([Eq(u, 1)] + rec_eq) - print(op) op.apply() assert np.all(s.data == 1)