diff options
Diffstat (limited to 'threadtest.py')
| -rw-r--r-- | threadtest.py | 57 |
1 files changed, 33 insertions, 24 deletions
diff --git a/threadtest.py b/threadtest.py index 6d4e228..11bdaa4 100644 --- a/threadtest.py +++ b/threadtest.py @@ -11,16 +11,6 @@ from histogram import Histogram jobs = Queue() output = Queue() -lbne = LBNE() -lbne.build(bits=8) - -cuda.init() - -gputhreads = [] -for i in range(cuda.Device.count()): - gputhreads.append(GPUThread(i, lbne, jobs, output)) - gputhreads[-1].start() - def generate_event(pos=(0,0,0), nphotons=1000): origins_float3 = np.tile(gpuarray.vec.make_float3(*pos), (nphotons,1)) directions = uniform_sphere(nphotons) @@ -67,18 +57,37 @@ def likelihood(event_bincount, pos=(0,0,0), nphotons=1000, neval=1000): return -log_likelihood -def stop(): - for gputhread in gputhreads: - gputhread.stop() - - for gputhread in gputhreads: - gputhread.join() - if __name__ == '__main__': - event_bincount = generate_event() - - for z in np.linspace(-1.0, 1.0, 100): - log_likelihood = likelihood(event_bincount, (z,0,0)) - print 'z = %f, likelihood = %s' % (z, log_likelihood) - - stop() + import sys + import optparse + import time + + parser = optparse.OptionParser('%prog') + parser.add_option('-b', type='int', dest='nbits', default=8) + parser.add_option('-j', type='int', dest='ndevices', default=1) + options, args = parser.parse_args() + + lbne = LBNE() + lbne.build(bits=options.nbits) + + cuda.init() + + gputhreads = [] + for i in range(options.ndevices): + gputhreads.append(GPUThread(i, lbne, jobs, output)) + gputhreads[-1].start() + + try: + event_bincount = generate_event() + + for z in np.linspace(-1.0, 1.0, 100): + t0 = time.time() + log_likelihood = likelihood(event_bincount, (z,0,0)) + elapsed = time.time() - t0 + print 'z = %f, likelihood = %s, elapsed %f sec' % (z, log_likelihood, elapsed) + finally: + for gputhread in gputhreads: + gputhread.stop() + + for gputhread in gputhreads: + gputhread.join() |