class Timer:
    """클래스 생성 시점부터 소멸 시점까지의 시간을 출력한다."""

    def __init__(self, func_name: str='this func'):
        self.func_name: str = func_name
        self.time_start: float = 0.0

    def __enter__(self):
        import sys
        import time
        print('{} ==>'.format(self.func_name), end=' ')
        sys.stdout.flush()
        self.time_start = time.process_time()
        return self

    def __exit__(self, *args):
        import time
        time_end = time.process_time()
        interval = time_end - self.time_start
        print('Elapsed time: {:.8f} sec'.format(interval))

import pycuda.driver as drv
drv.init()

print('Detected {} CUDA Capable device(s) \n'.format(drv.Device.count()))

for i in range(drv.Device.count()):

    dev = drv.Device(i)
    print(f'Device {i}: {dev.name()}')
    compute_capability = float('%d.%d' % dev.compute_capability())
    print(f'  Compute Capability: {compute_capability}')
    print(f'  Total Memory: {dev.total_memory() // (1024**2)} MB')

    dev_attr_tuples = dev.get_attributes().items()
    dev_attributes = {}

    for k, v in dev_attr_tuples:
        dev_attributes[str(k)] = v

    num_mp = dev_attributes['MULTIPROCESSOR_COUNT']

    cuda_cores_per_mp = {5.0 : 128,
                         5.1 : 128,
                         5.2 : 128,
                         6.0 : 64,
                         6.1 : 128,
                         6.2 : 128}[compute_capability]
    print(f'  Multiprocessors: {num_mp}')
    print(f'  CUDA Cores Per Multiprocessor: {cuda_cores_per_mp}')
    print(f'  Total CUDA Cores: {num_mp * cuda_cores_per_mp}')

    dev_attributes.pop('MULTIPROCESSOR_COUNT')

    for k in dev_attributes.keys():
        print(f'  {k}: {dev_attributes[k]}')

from numba import cuda
from numba.cuda import to_device
import numpy as np

@cuda.reduce
def sum_reduce(a, b):
    return a + b

import math
import numpy as np
from numpy import float32
from mkl_random import standard_normal

n_sample = 2 ** 26

data = float32(standard_normal(n_sample))
d_data = to_device(data.copy())

with Timer('Sum By CPU'):
    sum_cpu = np.sum(data)

_ = sum_reduce(data)
    
with Timer('Sum By GPU'):
    sum_gpu = sum_reduce(d_data)
    
print(f'CPU Result: {sum_cpu}')
print(f'GPU Result: {sum_gpu}')

# Kernel Code
from numba import cuda
from numba.cuda import to_device
import numpy as np

@cuda.jit("void(float32[:], float32[:], int64)")
def sum_kernel(out, data, n):
    tid = cuda.threadIdx.x
    idx = cuda.grid(1)
    idx_block_start = cuda.blockIdx.x * cuda.blockDim.x

    if idx >= n:
        return

    stride = cuda.blockDim.x // 2

    while stride > 0:
        if tid < stride and (idx_block_start + tid + stride) < n:
            data[idx_block_start + tid] += data[idx_block_start + tid + stride]
        cuda.syncthreads()
        stride //= 2

    if (tid == 0):
        out[cuda.blockIdx.x] = data[idx_block_start]

# Host Code
import math
import numpy as np
from numpy import float32
from mkl_random import standard_normal

n_sample = 2 ** 26

data = float32(standard_normal(n_sample))
d_data = to_device(data.copy())

with Timer('Sum By CPU'):
    sum_cpu = np.sum(data)

threadsperblock = 128
blockspergrid = math.ceil(n_sample / threadsperblock)
out_of_gpu = cuda.device_array(blockspergrid, float32)

sum_kernel[blockspergrid, threadsperblock](out_of_gpu, data, n_sample) # For Compile

with Timer('Sum By GPU'):
    sum_kernel[blockspergrid, threadsperblock](out_of_gpu, d_data, n_sample)
    out_from_gpu = out_of_gpu.copy_to_host()
    sum_gpu = np.sum(out_from_gpu)

print(f'CPU Result: {sum_cpu}')
print(f'GPU Result: {sum_gpu}')