diff --git a/devito/core/operator.py b/devito/core/operator.py index 6cf20767b2..39a0d3a2d1 100644 --- a/devito/core/operator.py +++ b/devito/core/operator.py @@ -263,6 +263,15 @@ def _check_kwargs(cls, **kwargs): "`npthreads` must be a positive integer" ) + async_degree = oo['buf-async-degree'] + if async_degree is not None and ( + isinstance(async_degree, bool) or + not is_integer(async_degree) or async_degree < 0 + ): + raise InvalidOperator( + "`buf-async-degree` must be a non-negative integer" + ) + if oo['cire-maxpar'] not in (False, 'basic', 'compact'): raise InvalidOperator("Illegal `cire-maxpar` value") diff --git a/devito/passes/clusters/buffering.py b/devito/passes/clusters/buffering.py index 9c627284ec..e9d90517a3 100644 --- a/devito/passes/clusters/buffering.py +++ b/devito/passes/clusters/buffering.py @@ -44,11 +44,12 @@ def buffering(clusters, key, sregistry, options, **kwargs): Accepted: ['buf-async-degree', 'buf-reuse', 'npthreads']. * 'buf-async-degree': Specify the size of the buffer. By default, the buffer size is the minimal one, inferred from the memory accesses in - the ``clusters`` themselves. An asynchronous degree equals to `k` - means that the buffer will be enforced to size=`k` along the introduced - ModuloDimensions. This might help relieving the synchronization - overhead when asynchronous operations are used (these are however - implemented by other passes). + the ``clusters`` themselves. A positive asynchronous degree `k` + requests `k` slots; values below the inferred minimum are ignored. + Zero disables buffering. Read-buffer initialization remains limited to + the minimum number of slots required by the memory accesses. A larger + buffer might relieve synchronization overhead in asynchronous operations + introduced by other passes. * 'buf-reuse': If True, the pass will try to reuse existing Buffers for different buffered Functions. By default, False. * 'npthreads': Number of pthreads for asynchronous tasks. The tasks are @@ -640,10 +641,12 @@ def write_to(self): # might be accessed through a stencil ispace = ispace.promote(lambda d: d.is_AbstractSub, mode='total') - # Analogous to the above, we need to include the halo region as well + # Include the spatial halo without widening the temporal interval, + # which may already be restricted by an earlier buffering round ihalo = IntervalGroup([ Interval(i.dim, -h.left, h.right, i.stamp) for i, h in zip(ispace, self.b._size_halo, strict=False) + if i.dim is not self.xd ]) ispace = IterationSpace.union(ispace, IterationSpace(ihalo)) @@ -793,6 +796,7 @@ def init_buffers(descriptors, options): Create the initializing Clusters for the given buffers. """ init_onwrite = options['buf-init-onwrite'] + async_degree = options['buf-async-degree'] init = [] for b, v in descriptors.flat_items(): @@ -803,6 +807,7 @@ def init_buffers(descriptors, options): # multiple) buffering because it's completely unnecessary if v.is_double_buffering: continue + lhs = b.indexify()._subs(v.xd, v.first_idx.b) rhs = f.indexify()._subs(v.dim, v.first_idx.f) @@ -817,6 +822,16 @@ def init_buffers(descriptors, options): expr = lower_exprs(expr) ispace = v.write_to + if v.is_read and async_degree is not None: + # The allocated capacity (`v.size`) may exceed the time-window width + # that must be loaded before computation starts (`size` below). E.g., + # reads at u[t-1], u[t] and u[t+1] make `infer_buffer_size` return 3, + # even if `buf-async-degree` gives us 4 slots (`v.size == 4`). Seed + # only db0=0..2; the spare slot is filled as computation advances. + # This preserves the stencil's data space and iteration bounds, + # without requiring extra input time levels to fill the ring. + size = infer_buffer_size(f, v.dim, v.clusters) + ispace = ispace.translate(v.xd, 0, size - v.size) guards = {} guards[None] = GuardBound(v.dim.root.symbolic_min, v.dim.root.symbolic_max) diff --git a/tests/test_buffering.py b/tests/test_buffering.py index eb2ac804ab..b28e83aac8 100644 --- a/tests/test_buffering.py +++ b/tests/test_buffering.py @@ -113,7 +113,8 @@ def test_read_only_w_offset(): assert np.all(v.data == v1.data) -def test_read_only_backwards(): +@pytest.mark.parametrize('async_degree', [None, 4]) +def test_read_only_backwards(async_degree): nt = 10 grid = Grid(shape=(2, 2)) @@ -127,7 +128,8 @@ def test_read_only_backwards(): eqns = [Eq(v.backward, v + u.backward + u + u.forward + 1.)] op0 = Operator(eqns, opt='noop') - op1 = Operator(eqns, opt='buffering') + op1 = Operator(eqns, opt=('buffering', + {'buf-async-degree': async_degree})) # Check generated code assert len(retrieve_iteration_tree(op1)) == 4 @@ -171,32 +173,83 @@ def test_read_only_backwards_unstructured(): assert np.all(v.data == v1.data) -@pytest.mark.parametrize('async_degree', [2, 4]) -def test_async_degree(async_degree): +@pytest.mark.parametrize('async_degree', [1, 2, 4]) +@pytest.mark.parametrize('backward', [False, True], + ids=['forward', 'backward']) +def test_async_degree(async_degree, backward): nt = 10 grid = Grid(shape=(4, 4)) u = TimeFunction(name='u', grid=grid, save=nt) u1 = TimeFunction(name='u', grid=grid, save=nt) - eqn = Eq(u.forward, u + 1) + lhs = u.backward if backward else u.forward + eqn = Eq(lhs, u + 1) op0 = Operator(eqn, opt='noop') op1 = Operator(eqn, opt=('buffering', {'buf-async-degree': async_degree})) # Check generated code assert len(retrieve_iteration_tree(op1)) == 3 - buffers = [i for i in FindSymbols().visit(op1) if i.is_Array and i._mem_heap] + buffers = [i for i in FindSymbols().visit(op1) + if i.is_Array and i._mem_heap] assert len(buffers) == 1 - assert buffers.pop().symbolic_shape[0] == async_degree + assert buffers.pop().symbolic_shape[0] == max(2, async_degree) - op0.apply(time_M=nt-2) - op1.apply(time_M=nt-2, u=u1) + kwargs = {'time_m': 1} if backward else {'time_M': nt - 2} + op0.apply(**kwargs) + op1.apply(u=u1, **kwargs) assert np.all(u.data == u1.data) -def test_two_homogeneous_buffers(): +@pytest.mark.parametrize('backward,expected_bounds', [ + pytest.param(False, (0, 8), id='forward'), + pytest.param(True, (1, 9), id='backward') +]) +@pytest.mark.parametrize('async_degree', [0, 1, 4, 16]) +def test_async_degree_read_only(backward, expected_bounds, async_degree): + nt = 10 + grid = Grid(shape=(4, 4)) + + u = TimeFunction(name='u', grid=grid, save=nt) + v = TimeFunction(name='v', grid=grid) + v1 = TimeFunction(name='v', grid=grid) + + u.data[:] = np.arange(nt).reshape(nt, 1, 1) + + lhs = v.backward if backward else v.forward + eqn = Eq(lhs, v + u) + + op0 = Operator(eqn, opt='noop', name='op0') + op1 = Operator(eqn, opt=('buffering', + {'buf-async-degree': async_degree}), name='op1') + + buffers = [i for i in FindSymbols().visit(op1) + if i.is_Array and i._mem_heap] + assert len(buffers) == int(async_degree != 0) + if async_degree: + assert buffers[0].symbolic_shape[0] == async_degree + + for op in [op0, op1]: + args = op.arguments() + assert (args['time_m'], args['time_M']) == expected_bounds + + # Default bounds, either endpoint, a partial ring, and an empty interval + time_m, time_M = expected_bounds + for kwargs in [{}, {'time_m': time_m, 'time_M': time_m}, + {'time_m': time_M, 'time_M': time_M}, + {'time_m': 3, 'time_M': 4}, {'time_m': 1, 'time_M': 0}]: + v.data[:] = 0 + v1.data[:] = 0 + op0.apply(**kwargs) + op1.apply(v=v1, **kwargs) + + assert np.all(v.data == v1.data) + + +@pytest.mark.parametrize('async_degree', [None, 4]) +def test_two_homogeneous_buffers(async_degree): nt = 10 grid = Grid(shape=(4, 4)) @@ -209,8 +262,10 @@ def test_two_homogeneous_buffers(): Eq(v.forward, u + v + u.backward + v.backward + 1.)] op0 = Operator(eqns, opt='noop') - op1 = Operator(eqns, opt='buffering') - op2 = Operator(eqns, opt=('buffering', 'fuse')) + op1 = Operator(eqns, opt=('buffering', + {'buf-async-degree': async_degree})) + op2 = Operator(eqns, opt=('buffering', 'fuse', + {'buf-async-degree': async_degree})) # Check generated code assert len(retrieve_iteration_tree(op1)) == 5 @@ -224,8 +279,16 @@ def test_two_homogeneous_buffers(): assert np.all(u.data == u1.data) assert np.all(v.data == v1.data) + u1.data[:] = 0 + v1.data[:] = 0 + op2.apply(time_M=nt-2, u=u1, v=v1) + + assert np.all(u.data == u1.data) + assert np.all(v.data == v1.data) + -def test_two_heterogeneous_buffers(): +@pytest.mark.parametrize('async_degree', [None, 4]) +def test_two_heterogeneous_buffers(async_degree): nt = 10 grid = Grid(shape=(4, 4)) @@ -242,7 +305,8 @@ def test_two_heterogeneous_buffers(): Eq(v.forward, u + v + v.backward)] op0 = Operator(eqns, opt='noop') - op1 = Operator(eqns, opt='buffering') + op1 = Operator(eqns, opt=('buffering', + {'buf-async-degree': async_degree})) # Check generated code assert len(retrieve_iteration_tree(op1)) == 5 diff --git a/tests/test_gpu_common.py b/tests/test_gpu_common.py index c0e4d00ade..a9f0632121 100644 --- a/tests/test_gpu_common.py +++ b/tests/test_gpu_common.py @@ -814,6 +814,7 @@ def test_streaming_conddim_forward(self, opt): @pytest.mark.parametrize('opt', [ ('buffering', 'streaming', 'orchestrate'), + ('buffering', 'streaming', 'orchestrate', {'buf-async-degree': 4}), ]) def test_streaming_conddim_backward(self, opt): nt = 10 @@ -849,6 +850,50 @@ def test_streaming_conddim_backward(self, opt): # 3rd time u[1] = u[0]+u[1]+usave[2] = 0+7+2 = 9 assert np.all(u.data[1] == 9) + def run_streaming_async_degree(self, backward, expected_bounds, async_degree): + nt = 10 + grid = Grid(shape=(4, 4)) + + usave = TimeFunction(name='usave', grid=grid, save=nt) + v = TimeFunction(name='v', grid=grid) + v1 = TimeFunction(name='v', grid=grid) + + usave.data._local[:] = np.arange(nt).reshape(nt, 1, 1) + + lhs = v.backward if backward else v.forward + eqn = Eq(lhs, v + usave) + + op0 = Operator(eqn, opt=('noop', {'gpu-fit': usave}), name='op0') + op1 = Operator(eqn, opt=('buffering', 'streaming', 'orchestrate', + {'buf-async-degree': async_degree}), name='op1') + + for op in [op0, op1]: + args = op.arguments() + assert (args['time_m'], args['time_M']) == expected_bounds + + op0.apply() + op1.apply(v=v1) + + assert np.all(v.data == v1.data) + + @pytest.mark.parametrize('backward,expected_bounds', [ + pytest.param(False, (0, 8), id='forward'), + pytest.param(True, (1, 9), id='backward') + ]) + @pytest.mark.parametrize('async_degree', [4, 16]) + def test_streaming_async_degree(self, backward, expected_bounds, async_degree): + self.run_streaming_async_degree(backward, expected_bounds, async_degree) + + @pytest.mark.parallel(mode=2) + @pytest.mark.parametrize('backward,expected_bounds', [ + pytest.param(False, (0, 8), id='forward'), + pytest.param(True, (1, 9), id='backward') + ]) + @pytest.mark.parametrize('async_degree', [4, 16]) + def test_streaming_async_degree_mpi(self, backward, expected_bounds, + async_degree, mode): + self.run_streaming_async_degree(backward, expected_bounds, async_degree) + @pytest.mark.parametrize('opt,ntmps', [ (('buffering', 'streaming', 'orchestrate'), 3), ]) @@ -910,7 +955,8 @@ def test_streaming_multi_input_conddim_foward(self): assert np.all(v.data == v1.data) - def test_streaming_multi_input_conddim_backward(self): + @pytest.mark.parametrize('async_degree', [None, 5]) + def test_streaming_multi_input_conddim_backward(self, async_degree): nt = 10 grid = Grid(shape=(4, 4)) time_dim = grid.time_dim @@ -931,7 +977,8 @@ def test_streaming_multi_input_conddim_backward(self): eqns = [Eq(v.backward, v + expr + 1.)] op0 = Operator(eqns, opt=('noop', {'gpu-fit': u})) - op1 = Operator(eqns, opt=('buffering', 'streaming', 'orchestrate')) + op1 = Operator(eqns, opt=('buffering', 'streaming', 'orchestrate', + {'buf-async-degree': async_degree})) op0.apply(time_M=nt, dt=.01) op1.apply(time_M=nt, dt=.01, v=v1) diff --git a/tests/test_operator.py b/tests/test_operator.py index 218c5e6710..59cf0013c7 100644 --- a/tests/test_operator.py +++ b/tests/test_operator.py @@ -141,6 +141,11 @@ def test_opt_options(self): Operator(Eq(u, u + 1), opt=('advanced', {'npthreads': npthreads})) + for async_degree in (False, True, -1, 1.5): + with pytest.raises(InvalidOperator, match='non-negative integer'): + Operator(Eq(u, u + 1), + opt=('advanced', {'buf-async-degree': async_degree})) + def test_compiler_uniqueness(self): grid = Grid(shape=(3, 3, 3))