适用于CUDA GPU的Numba例子
矩阵乘法
这是使用CUDA内核的矩阵乘法的简单实现:
@cuda.jit
def matmul(A, B, C):
"""Perform square matrix multiplication of C = A * B
"""
i, j = cuda.grid(2)
if i < C.shape[0] and j < C.shape[1]:
tmp = 0.
for k in range(A.shape[1]):
tmp += A[i, k] * B[k, j]
C[i, j] = tmp
这种实现方式简单直观,但性能不佳,因为相同的矩阵元素将从设备内存中多次加载,这很慢(某些设备可能具有透明的数据缓存,但它们可能不足以一次容纳整个输入)。
如果使用阻塞算法来减少对设备内存的访问,它将更快。CUDA为 块中的线程提供快速共享内存,以协作执行任务。以下使用共享内存实现了方阵乘法的更快版本:
from numba import cuda, float32
# Controls threads per block and shared memory usage.
# The computation will be done on blocks of TPBxTPB elements.
TPB = 16
@cuda.jit
def fast_matmul(A, B, C):
# Define an array in the shared memory
# The size and type of the arrays must be known at compile time
sA = cuda.shared.array(shape=(TPB, TPB), dtype=float32)
sB = cuda.shared.array(shape=(TPB, TPB), dtype=float32)
x, y = cuda.grid(2)
tx = cuda.threadIdx.x
ty = cuda.threadIdx.y
bpg = cuda.gridDim.x # blocks per grid
if x >= C.shape[0] and y >= C.shape[1]:
# Quit if (x, y) is outside of valid C boundary
return
# Each thread computes one element in the result matrix.
# The dot product is chunked into dot products of TPB-long vectors.
tmp = 0.
for i in range(bpg):
# Preload data into shared memory
sA[tx, ty] = A[x, ty + i * TPB]
sB[tx, ty] = B[tx + i * TPB, y]
# Wait until all threads finish preloading
cuda.syncthreads()
# Computes partial product on the shared memory
for j in range(TPB):
tmp += sA[tx, j] * sB[j, ty]
# Wait until all threads finish computing
cuda.syncthreads()
C[x, y] = tmp
因为共享内存是有限的资源,所以代码一次从输入数组中预加载小块。然后,调用 syncthreads()以等待所有线程完成预加载,再对共享内存进行计算。计算之后,再次同步,以确保所有线程在共享内存中的数据均已完成之后,在下一个循环迭代中将其覆盖。
最新文章
- PKCS#1规范阅读笔记1--------基本概念
- 开启Apache mod_rewrite模块完全解答
- android中基于HTML模板的方式嵌入SWF
- 在Windows Azure虚拟机上开发Windows 8 应用
- K近邻分类算法实现 in Python
- 为Form中的控件增加自适应功能 转
- C++ STL stack和queue
- popupWindow使用详解
- HDOJ(HDU) 2212 DFS(阶乘相关、)
- JAE京东云引擎Git上传管理代码教程和京东云数据库导入导出管理
- java super关键字
- 省市区三级联动(jquery+ajax)(封装和不封装两种方式)-----2017-05-15
- asp.net core 开发的https证书服务-agilelabs.net
- C#中&与&&的区别
- JavaScript数据类型之null和undeined
- Kong配置参考
- 20)django-session使用
- Linux-server-sshd
- spring boot +mybatis 整合 连接数据库测试(从0到1)
- table滑块