Skip to content

Commit

Permalink
chore(numba): rename functions (#22153)
Browse files Browse the repository at this point in the history
  • Loading branch information
hongbo-miao authored Dec 30, 2024
1 parent 444e70c commit 724b87f
Showing 1 changed file with 22 additions and 22 deletions.
44 changes: 22 additions & 22 deletions data-analytics/hm-numba/src/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,15 @@
from numba import cuda


def sum_of_squares_python(n: int) -> float:
total = 0.0
for i in range(n):
total += i * i
return total


@nb.jit(nopython=True)
def sum_of_squares_numba(n: int) -> float:
def sum_of_squares_cpu_numba(n: int) -> float:
total: float = 0.0
for i in range(n):
total += i * i
Expand All @@ -21,7 +28,7 @@ def sum_of_squares_cuda_kernel(n: int, result: cuda.devicearray.DeviceNDArray) -
cuda.atomic.add(result, 0, idx * idx)


def sum_of_squares_gpu(n: int) -> float:
def sum_of_squares_gpu_numba(n: int) -> float:
# Allocate memory on GPU
result = cuda.device_array(1, dtype=np.float64)
cuda.to_device(np.array([0.0], dtype=np.float64), to=result)
Expand All @@ -37,39 +44,32 @@ def sum_of_squares_gpu(n: int) -> float:
return result.copy_to_host()[0]


def sum_of_squares_python(n: int) -> float:
total = 0.0
for i in range(n):
total += i * i
return total


def main() -> None:
n = 10_000_000

# Python version
start_time = time.time()
result_python = sum_of_squares_python(n)
python_result = sum_of_squares_python(n)
python_time = time.time() - start_time
logging.info(f"Python result: {result_python}")
logging.info(f"Python result: {python_result}")
logging.info(f"Python time: {python_time} seconds")

# CPU Numba version
start_time = time.time()
result_numba = sum_of_squares_numba(n)
numba_time = time.time() - start_time
logging.info(f"CPU Numba result: {result_numba}")
logging.info(f"CPU Numba time: {numba_time} seconds")
logging.info(f"CPU Numba speedup vs Python: {python_time / numba_time}x")
cpu_numba_result = sum_of_squares_cpu_numba(n)
cpu_numba_time = time.time() - start_time
logging.info(f"CPU Numba result: {cpu_numba_result}")
logging.info(f"CPU Numba time: {cpu_numba_time} seconds")
logging.info(f"CPU Numba speedup vs Python: {python_time / cpu_numba_time}x")

# GPU version
# GPU Numba version
try:
start_time = time.time()
result_gpu = sum_of_squares_gpu(n)
gpu_time = time.time() - start_time
logging.info(f"GPU result: {result_gpu}")
logging.info(f"GPU time: {gpu_time} seconds")
logging.info(f"GPU speedup vs Python: {python_time / gpu_time}x")
gpu_numba_result = sum_of_squares_gpu_numba(n)
gpu_numba_time = time.time() - start_time
logging.info(f"GPU Numba result: {gpu_numba_result}")
logging.info(f"GPU Numba time: {gpu_numba_time} seconds")
logging.info(f"GPU Numba speedup vs Python: {python_time / gpu_numba_time}x")
except cuda.CudaSupportError:
logging.warning("CUDA GPU is not available")

Expand Down

0 comments on commit 724b87f

Please sign in to comment.