Merge pull request #3023 from devitocodes/revamp-buf-async-degree

compiler: Revamp bug-async-degree opt-option

Author: FabioLuporini

.github/workflows/lint.yaml

@@ -40,7 +40,7 @@ jobs:
 
     - name: Lint codebase with flake8
       run: |
-        flake8 --builtins=ArgumentError .
+        flake8 .
 
   spellcheck:
     name: "Spellcheck everything"

devito/core/operator.py

@@ -263,6 +263,13 @@ def _check_kwargs(cls, **kwargs):
                 "`npthreads` must be a positive integer"
             )
 
+        async_degree = oo['buf-async-degree']
+        if async_degree is not None and \
+           (type(async_degree) is not int 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")
 

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
@@ -559,11 +560,15 @@ def expand_halo_transfers(clusters, mapper):
     return processed
 
 
-def _include_halo(ispace, f):
-    """Extend `ispace` to include `f`'s HALO."""
+def _include_halo(ispace, f, dims=None):
+    """
+    Extend `ispace` to include `f`'s HALO along `dims`.
+    """
+    dims = dims or f.dimensions
+
     ihalo = [
         Interval(i.dim, -f._size_halo[i.dim].left, f._size_halo[i.dim].right, i.stamp)
-        for i in ispace if i.dim in f.dimensions
+        for i in ispace if i.dim in dims
     ]
 
     return IterationSpace.union(ispace, IterationSpace(ihalo))
@@ -724,8 +729,9 @@ 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
-        ispace = _include_halo(ispace, self.b)
+        # Include the spatial halo without widening the temporal interval,
+        # which may already be restricted by an earlier buffering round
+        ispace = _include_halo(ispace, self.b, self.bdims)
 
         return ispace
 
@@ -872,6 +878,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():
@@ -882,6 +889,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)
 
@@ -895,7 +903,17 @@ def init_buffers(descriptors, options):
         expr = Eq(lhs, rhs)
         expr = lower_exprs(expr)
 
-        ispace = v.write_to
+        ispace = v.write_to.concrete
+        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)

tests/test_buffering.py

@@ -212,7 +212,8 @@ def test_read_only_w_offset():
     assert np.all(v.data == v1.data)
 
 
-def test_read_only_backwards():
+@pytest.mark.parametrize('async_degree,expected_size', [(None, 3), (4, 4)])
+def test_read_only_backwards(async_degree, expected_size):
     nt = 10
     grid = Grid(shape=(2, 2))
 
@@ -226,12 +227,14 @@ 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
     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] == expected_size
 
     op0.apply(time_m=1)
     op1.apply(time_m=1, v=v1)
@@ -270,32 +273,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))
 
@@ -308,8 +362,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
@@ -323,8 +379,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))
 
@@ -341,7 +405,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

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,42 @@ 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)
 
+    @pytest.mark.parametrize('mode', [
+        pytest.param(None, id='serial'),
+        pytest.param(2, marks=pytest.mark.parallel, id='basic')
+    ])
+    @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, mode, 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('opt,ntmps', [
         (('buffering', 'streaming', 'orchestrate'), 3),
     ])
@@ -910,7 +947,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 +969,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)

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))