numbapig.py 1.21 KB
Newer Older
Neil Gershenfeld's avatar
wip  
Neil Gershenfeld committed
1
2
#
# numbapig.py
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
3
4
# Neil Gershenfeld 3/1/20
# calculation of pi by a Numba CUDA sum
Neil Gershenfeld's avatar
wip  
Neil Gershenfeld committed
5
6
7
8
9
# pi = 3.14159265358979323846 
#
from numba import cuda
import numpy as np
import time
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
10
11
12
13
block_size = 1024
grid_size = 1024
nloop = 1000000
npts = grid_size*block_size
Neil Gershenfeld's avatar
wip  
Neil Gershenfeld committed
14
@cuda.jit
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
15
16
17
18
19
20
21
def init(arr,nloop):
   i = cuda.blockIdx.x*cuda.blockDim.x+cuda.threadIdx.x;
   start = nloop*i+1;
   end = nloop*(i+1)+1;
   arr[i] = 0;
   for j in range(start,end):
      arr[i] += 0.5/((j-0.75)*(j-0.25));
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
22
@cuda.jit
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
23
24
25
26
27
28
def reduce_sum(arr,len):
   i = cuda.blockIdx.x*cuda.blockDim.x+cuda.threadIdx.x;
   if (i < len):
      arr[i] += arr[i+len]
def reduce(arr,npts):
   len = npts >> 1
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
29
   while (1):
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
30
      reduce_sum[grid_size,block_size](arr,len)
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
31
32
33
      len = len >> 1
      if (len == 0):
         return
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
34
35
36
darr = cuda.device_array(npts,np.float64)
init[grid_size,block_size](darr,nloop) # compile kernel
reduce(darr,npts) # compile kernel
Neil Gershenfeld's avatar
wip  
Neil Gershenfeld committed
37
start_time = time.time()
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
38
39
init[grid_size,block_size](darr,nloop)
reduce(darr,npts)
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
40
cuda.synchronize()
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
41
end_time = time.time()
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
42
43
44
45
mflops = npts*nloop*5.0/(1.0e6*(end_time-start_time))
harr = darr.copy_to_host()
print("npts = %d, nloop = %d, pi = %f"%(npts,nloop,harr[0]))
print("time = %f, estimated MFlops = %f"%(end_time-start_time,mflops))
Neil Gershenfeld's avatar
wip    
Neil Gershenfeld committed
46