Support max_dulation keyword argument
+import math
import time
import numpy
@@ -8,7 +9,9 @@ from cupy import util
class _PerfCaseResult(object):
def __init__(self, name, ts):
- assert ts.ndim == 2 and ts.shape[0] == 2 and ts.shape[1] > 0
+ assert ts.ndim == 2
+ assert ts.shape[0] == 2
+ assert ts.shape[1] > 0
self.name = name
self._ts = ts
@@ -42,7 +45,9 @@ class _PerfCaseResult(object):
return self.to_str(show_gpu=True)
-def repeat(func, args=(), kwargs={}, n_times=10000, *, name=None, n_warmup=10):
+def repeat(
+ func, args=(), kwargs={}, n_repeat=10000, *,
+ name=None, n_warmup=10, max_duration=math.inf):
util.experimental('cupyx.time.repeat')
if name is None:
name = func.__name__
@@ -53,14 +58,13 @@ def repeat(func, args=(), kwargs={}, n_times=10000, *, name=None, n_warmup=10):
raise ValueError('`args` should be of tuple type.')
if not isinstance(kwargs, dict):
raise ValueError('`kwargs` should be of dict type.')
- if not isinstance(n_times, int):
- raise ValueError('`n_times` should be an integer.')
+ if not isinstance(n_repeat, int):
+ raise ValueError('`n_repeat` should be an integer.')
if not isinstance(name, str):
raise ValueError('`str` should be a string.')
if not isinstance(n_warmup, int):
raise ValueError('`n_warmup` should be an integer.')
- ts = numpy.empty((2, n_times,), dtype=numpy.float64)
ev1 = cupy.cuda.stream.Event()
ev2 = cupy.cuda.stream.Event()
@@ -70,7 +74,10 @@ def repeat(func, args=(), kwargs={}, n_times=10000, *, name=None, n_warmup=10):
ev1.record()
ev1.synchronize()
- for i in range(n_times):
+ cpu_times = []
+ gpu_times = []
+ duration = 0
+ for i in range(n_repeat):
ev1.record()
t1 = time.perf_counter()
@@ -81,7 +88,12 @@ def repeat(func, args=(), kwargs={}, n_times=10000, *, name=None, n_warmup=10):
ev2.synchronize()
cpu_time = t2 - t1
gpu_time = cupy.cuda.get_elapsed_time(ev1, ev2) * 1e-3
- ts[0, i] = cpu_time
- ts[1, i] = gpu_time
+ cpu_times.append(cpu_time)
+ gpu_times.append(gpu_time)
+ duration += cpu_time
+ if duration > max_duration:
+ break
+
+ ts = numpy.asarray([cpu_times, gpu_times], dtype=numpy.float64)
return _PerfCaseResult(name, ts)
