Benchmark photon load times

1 files changed, 32 insertions, 1 deletions

diff --git a/benchmark.py b/benchmark.py
index 441b85f..d0819cb 100644..100755
--- a/benchmark.py
+++ b/benchmark.py
@@ -1,3 +1,4 @@
+#/usr/bin/env python
 import numpy as np
 from pycuda import gpuarray as ga
 import time
@@ -62,6 +63,35 @@ def ray_trace(gpu, number=1000):
 
     return pixels_gpu.size/ufloat((np.mean(run_times),np.std(run_times)))
 
+def load_photons(gpu, number=10, nphotons=500000):
+    """
+    Return the mean and standard deviation of the number of photons loaded
+    onto `gpu` per second.
+
+    Args:
+        - gpu, chroma.gpu.GPU
+            The GPU object with a geometry already loaded.
+        - number, int
+            The number of loads to average
+        - nphotons, int
+            The number of photons to load per trial
+    """
+    gpu.setup_propagate()
+    photons = Photons(np.zeros((nphotons,3)), uniform_sphere(nphotons), np.random.uniform(400,800,size=nphotons))
+
+    run_times = []
+    for i in progress(range(number)):
+        t0 = time.time()
+        gpu.load_photons(photons)
+        gpu.context.synchronize()
+        elapsed = time.time() - t0
+
+        if i > 0:
+            # first kernel call incurs some driver overhead
+            run_times.append(elapsed)
+
+    return nphotons/ufloat((np.mean(run_times),np.std(run_times)))
+
 def propagate(gpu, number=10, nphotons=500000):
     """
     Return the mean and standard deviation of the number of photons propagated
@@ -102,4 +132,5 @@ if __name__ == '__main__':
     gpu.load_geometry(lbne, print_usage=False)
 
     print '%s track steps/s' % ray_trace(gpu)
-    print '%s photons/s' % propagate(gpu)
+    print '%s loaded photons/s' % load_photons(gpu)
+    print '%s propagated photons/s' % propagate(gpu)