Repository: cupy/cupy

Issue:
cupyx.time.repeat does not work for multi-gpu benchmark
Currently, events are created and synchronized only for the current device in `cupyx.time.repeat`.


Discussion:

Resolving patch:
diff --git a/cupyx/time.py b/cupyx/time.py
--- a/cupyx/time.py
+++ b/cupyx/time.py
@@ -8,12 +8,13 @@
 
 
 class _PerfCaseResult(object):
-    def __init__(self, name, ts):
+    def __init__(self, name, ts, devices):
         assert ts.ndim == 2
-        assert ts.shape[0] == 2
+        assert ts.shape[0] == len(devices) + 1
         assert ts.shape[1] > 0
         self.name = name
         self._ts = ts
+        self._devices = devices
 
     @property
     def cpu_times(self):
@@ -21,7 +22,7 @@ def cpu_times(self):
 
     @property
     def gpu_times(self):
-        return self._ts[1]
+        return self._ts[1:]
 
     @staticmethod
     def _to_str_per_item(device_name, t):
@@ -38,7 +39,10 @@ def _to_str_per_item(device_name, t):
     def to_str(self, show_gpu=False):
         results = [self._to_str_per_item('CPU', self._ts[0])]
         if show_gpu:
-            results.append(self._to_str_per_item('GPU', self._ts[1]))
+            for i, d in enumerate(self._devices):
+                results.append(
+                    self._to_str_per_item('GPU-{}'.format(d),
+                                          self._ts[1 + i]))
         return '{:<20s}:{}'.format(self.name, ' '.join(results))
 
     def __str__(self):
@@ -47,11 +51,15 @@ def __str__(self):
 
 def repeat(
         func, args=(), kwargs={}, n_repeat=10000, *,
-        name=None, n_warmup=10, max_duration=math.inf):
+        name=None, n_warmup=10, max_duration=math.inf, devices=None):
+
     util.experimental('cupyx.time.repeat')
     if name is None:
         name = func.__name__
 
+    if devices is None:
+        devices = (cupy.cuda.get_device_id(),)
+
     if not callable(func):
         raise ValueError('`func` should be a callable object.')
     if not isinstance(args, tuple):
@@ -64,6 +72,23 @@ def repeat(
         raise ValueError('`str` should be a string.')
     if not isinstance(n_warmup, int):
         raise ValueError('`n_warmup` should be an integer.')
+    if not isinstance(devices, tuple):
+        raise ValueError('`devices` should be of tuple type')
+
+    return _repeat(
+        func, args, kwargs, n_repeat, name, n_warmup, max_duration, devices)
+
+
+def _repeat(
+        func, args, kwargs, n_repeat, name, n_warmup, max_duration, devices):
+
+    events_1 = []
+    events_2 = []
+
+    for i in devices:
+        with cupy.cuda.Device(i):
+            events_1.append(cupy.cuda.stream.Event())
+            events_2.append(cupy.cuda.stream.Event())
 
     ev1 = cupy.cuda.stream.Event()
     ev2 = cupy.cuda.stream.Event()
@@ -71,29 +96,40 @@ def repeat(
     for i in range(n_warmup):
         func(*args, **kwargs)
 
-    ev1.record()
-    ev1.synchronize()
+    for event, device in zip(events_1, devices):
+        with cupy.cuda.Device(device):
+            event.record()
+        event.synchronize()
 
     cpu_times = []
-    gpu_times = []
+    gpu_times = [[] for i in events_1]
     duration = 0
     for i in range(n_repeat):
-        ev1.record()
+        for event, device in zip(events_1, devices):
+            with cupy.cuda.Device(device):
+                event.record()
+
         t1 = time.perf_counter()
 
         func(*args, **kwargs)
 
         t2 = time.perf_counter()
-        ev2.record()
-        ev2.synchronize()
         cpu_time = t2 - t1
-        gpu_time = cupy.cuda.get_elapsed_time(ev1, ev2) * 1e-3
         cpu_times.append(cpu_time)
-        gpu_times.append(gpu_time)
+
+        for event, device in zip(events_2, devices):
+            with cupy.cuda.Device(device):
+                event.record()
+        for event, device in zip(events_2, devices):
+            with cupy.cuda.Device(device):
+                event.synchronize()
+        for i, (ev1, ev2) in enumerate(zip(events_1, events_2)):
+            gpu_time = cupy.cuda.get_elapsed_time(ev1, ev2) * 1e-3
+            gpu_times[i].append(gpu_time)
 
         duration += time.perf_counter() - t1
         if duration > max_duration:
             break
 
-    ts = numpy.asarray([cpu_times, gpu_times], dtype=numpy.float64)
-    return _PerfCaseResult(name, ts)
+    ts = numpy.asarray([cpu_times] + gpu_times, dtype=numpy.float64)
+    return _PerfCaseResult(name, ts, devices=devices)


Test patch:
diff --git a/tests/cupyx_tests/test_time.py b/tests/cupyx_tests/test_time.py
--- a/tests/cupyx_tests/test_time.py
+++ b/tests/cupyx_tests/test_time.py
@@ -4,6 +4,7 @@
 import numpy
 
 import cupy
+from cupy import testing
 import cupyx
 
 
@@ -26,7 +27,29 @@ def test_cpu_routine(self):
                 assert perf.name == 'test_name_xxx'
                 assert mock_func.call_count == 13
                 assert perf.cpu_times.shape == (10,)
-                assert perf.gpu_times.shape == (10,)
+                assert perf.gpu_times.shape == (1, 10,)
+                assert (perf.cpu_times == 1.4).all()
+                assert (perf.gpu_times == 2.5).all()
+
+    @testing.multi_gpu(2)
+    def test_multigpu_routine(self):
+        with mock.patch('time.perf_counter',
+                        mock.Mock(side_effect=[2.4, 3.8, 3.8] * 10)):
+            with mock.patch('cupy.cuda.get_elapsed_time',
+                            mock.Mock(return_value=2500)):
+                mock_func = mock.Mock()
+                mock_func.__name__ = 'test_name_xxx'
+                x = cupy.testing.shaped_random((2, 3), cupy, 'int32')
+                y = cupy.testing.shaped_random((2, 3), cupy, 'int32')
+                assert mock_func.call_count == 0
+
+                perf = cupyx.time.repeat(
+                    mock_func, (x, y), n_repeat=10, n_warmup=3, devices=(0, 1))
+
+                assert perf.name == 'test_name_xxx'
+                assert mock_func.call_count == 13
+                assert perf.cpu_times.shape == (10,)
+                assert perf.gpu_times.shape == (2, 10,)
                 assert (perf.cpu_times == 1.4).all()
                 assert (perf.gpu_times == 2.5).all()
 
@@ -47,7 +70,7 @@ def test_repeat_max_duration(self):
                 assert perf.name == 'test_name_xxx'
                 assert mock_func.call_count == 6
                 assert perf.cpu_times.shape == (3,)
-                assert perf.gpu_times.shape == (3,)
+                assert perf.gpu_times.shape == (1, 3)
                 assert (perf.cpu_times == 1.).all()
                 assert (perf.gpu_times == 2.5).all()
 
@@ -63,12 +86,12 @@ def test_show_gpu(self):
             [5.4, 7.1, 6.0, 5.4, 4.2],
             [6.4, 4.3, 8.9, 9.6, 3.8],
         ]) * 1e-6
-        perf = cupyx.time._PerfCaseResult('test_name_xxx', times)
+        perf = cupyx.time._PerfCaseResult('test_name_xxx', times, (0,))
         expected = (
             'test_name_xxx       :'
             '    CPU:    5.620 us   +/- 0.943 '
             '(min:    4.200 / max:    7.100) us '
-            '    GPU:    6.600 us   +/- 2.344 '
+            '    GPU-0:    6.600 us   +/- 2.344 '
             '(min:    3.800 / max:    9.600) us'
         )
         assert str(perf) == expected
@@ -78,7 +101,7 @@ def test_no_show_gpu(self):
             [5.4, 7.1, 6.0, 5.4, 4.2],
             [6.4, 4.3, 8.9, 9.6, 3.8],
         ]) * 1e-6
-        perf = cupyx.time._PerfCaseResult('test_name_xxx', times)
+        perf = cupyx.time._PerfCaseResult('test_name_xxx', times, (0,))
         expected = (
             'test_name_xxx       :'
             '    CPU:    5.620 us   +/- 0.943 '
@@ -90,11 +113,19 @@ def test_no_show_gpu(self):
 
     def test_single_show_gpu(self):
         times = numpy.array([[5.4], [6.4]]) * 1e-6
-        perf = cupyx.time._PerfCaseResult('test_name_xxx', times)
+        perf = cupyx.time._PerfCaseResult('test_name_xxx', times, (0,))
         assert str(perf) == ('test_name_xxx       :    CPU:    5.400 us '
-                             '    GPU:    6.400 us')
+                             '    GPU-0:    6.400 us')
 
     def test_single_no_show_gpu(self):
         times = numpy.array([[5.4], [6.4]]) * 1e-6
-        perf = cupyx.time._PerfCaseResult('test_name_xxx', times)
+        perf = cupyx.time._PerfCaseResult('test_name_xxx', times, (0,))
         assert perf.to_str() == 'test_name_xxx       :    CPU:    5.400 us'
+
+    def test_show_multigpu(self):
+        times = numpy.array([[5.4], [6.4], [7.0], [8.1]]) * 1e-6
+        perf = cupyx.time._PerfCaseResult('test_name_xxx', times, (0, 1, 2))
+        assert str(perf) == ('test_name_xxx       :    CPU:    5.400 us '
+                             '    GPU-0:    6.400 us '
+                             '    GPU-1:    7.000 us '
+                             '    GPU-2:    8.100 us')