cublas_sgemm_matrix_multiplication_with_cuda_python_bindings.py

1import numpy as np
2from cuda import cuda, cublas
3
4# Helper function to check for CUDA errors
5def checkCudaErrors(status):
6    if isinstance(status, cuda.CUresult):
7        if status != cuda.CUresult.CUDA_SUCCESS:
8            raise RuntimeError(f"CUDA Error: {status}")
9    elif isinstance(status, cublas.cublasStatus_t):
10        if status != cublas.cublasStatus_t.CUBLAS_STATUS_SUCCESS:
11            raise RuntimeError(f"cuBLAS Error: {status}")
12
13# 1. Initialize data
14m, n, k = 4, 4, 4
15alpha, beta = 1.0, 0.0
16A = np.ones((m, k), dtype=np.float32)
17B = np.ones((k, n), dtype=np.float32)
18C = np.zeros((m, n), dtype=np.float32)
19
20# 2. Initialize CUDA and cuBLAS
21checkCudaErrors(cuda.cuInit(0))
22res, dev = cuda.cuDeviceGet(0)
23res, ctx = cuda.cuCtxCreate(0, dev)
24handle = cublas.cublasCreate()
25
26# 3. Allocate and copy memory to device
27size_A = A.nbytes
28size_B = B.nbytes
29size_C = C.nbytes
30
31res, d_A = cuda.cuMemAlloc(size_A)
32res, d_B = cuda.cuMemAlloc(size_B)
33res, d_C = cuda.cuMemAlloc(size_C)
34
35checkCudaErrors(cuda.cuMemcpyHtoD(d_A, A, size_A))
36checkCudaErrors(cuda.cuMemcpyHtoD(d_B, B, size_B))
37
38# 4. Execute Matrix Multiplication (SGEMM)
39# Note: cuBLAS uses column-major order
40checkCudaErrors(cublas.cublasSgemm(
41    handle, 
42    cublas.cublasOperation_t.CUBLAS_OP_N, 
43    cublas.cublasOperation_t.CUBLAS_OP_N, 
44    m, n, k, 
45    alpha, d_A, m, 
46    d_B, k, 
47    beta, d_C, m
48))
49
50# 5. Copy result back to host
51checkCudaErrors(cuda.cuMemcpyDtoH(C, d_C, size_C))
52
53print("Result Matrix C (first element):", C[0,0])
54
55# 6. Cleanup
56cublas.cublasDestroy(handle)
57cuda.cuMemFree(d_A)
58cuda.cuMemFree(d_B)
59cuda.cuMemFree(d_C)
60cuda.cuCtxDestroy(ctx)