Python的GPU编程实例——近邻表计算
技术背景
GPU加速是现代工业各种场景中非常常用的一种技术,这得益于GPU计算的高度并行化。在Python中存在有多种GPU并行优化的解决方案,包括之前的博客中提到的cupy、pycuda和numba.cuda,都是GPU加速的标志性Python库。这里我们重点推numba.cuda这一解决方案,因为cupy的优势在于实现好了的众多的函数,在算法实现的灵活性上还比较欠缺;而pycuda虽然提供了很好的灵活性和相当高的性能,但是这要求我们必须在Python的代码中插入C代码,这显然是非常不Pythonic的解决方案。因此我们可以选择numba.cuda这一解决方案,只要在Python函数前方加一个numba.cuda.jit的修饰器,就可以在Python中用最Python的编程语法,实现GPU的加速效果。
加速场景
我们需要先了解的是,GPU在什么样的计算场景下能够实现加速的效果,很显然的是,并不是所有的计算过程都能在GPU上表现出加速的效果。前面说道,GPU的加速作用,是源自于高度的并行化,所谓的并行,就要求进程之前互不干扰或者依赖。如果说一个进程的计算过程或者结果,依赖于另一个进程中的计算结果,那么就无法实现完全的并行,只能使用串行的技术。这里为了展示GPU加速的效果,我们就引入一个在分子动力学模拟领域中常见的问题:近邻表的计算。
近邻表计算的问题是这样描述的:给定一堆数量为n的原子系统,每一个原子的三维坐标都是已知的,给定一个截断常数\(d_0\),当两个原子之间的距离\(d_{i,j}<=d_0\)时,则认为这两个原子是相邻近的原子。那么最终我们需要给出一个0-1矩阵\(A_{i,j}\),当\(A_{i,j}=0\)时,表示\(i,j\)两个原子互不相邻,反之则相邻。那么对于这个问题场景,我们就可以并行化的遍历\(n\times n\)的空间,直接输出\(A_{n\times n}\)大小的近邻表。这个计算场景是一个非常适合用GPU来加速的计算,以下我们先看一下不用GPU加速时的常规实现方案:
# cuda_neighbor_list.py
from numba import jit
from numba import cuda
import numpy as np
@jit
def neighbor_list(crd, neighbors, data_length, cutoff):
"""CPU based neighbor list calculation.
"""
for i in range(data_length):
for j in range(i+1, data_length):
if np.linalg.norm(crd[i]-crd[j]) <= cutoff:
neighbors[i][j] = 1
neighbors[j][i] = 1
return neighbors
if __name__ == '__main__':
np.random.seed(1)
atoms = 2**2
cutoff = 0.5
crd = np.random.random((atoms,3))
adjacent = np.zeros((atoms, atoms))
adjacent = neighbor_list(crd, adjacent, atoms, cutoff)
print (adjacent)
这是最常规的一种CPU上的实现方案,遍历所有的原子,计算原子间距,然后填充近邻表。这里我们还使用到了numba.jit即时编译的功能,这个功能是在执行到相关函数时再对其进行编译的方法,在矢量化的计算中有可能使用到芯片厂商所提供的SIMD的一些优化。当然,这里都是CPU层面的执行和优化,执行结果如下:
$ python3 cuda_neighbor_list.py
[[0. 0. 0. 0.]
[0. 0. 1. 0.]
[0. 1. 0. 1.]
[0. 0. 1. 0.]]
这个输出的结果就是一个0-1近邻表。
基于Numba的GPU加速
对于上述的近邻表计算的场景,我们很容易的想到这个neighbor_list
函数可以用GPU的函数来进行改造。对于每一个\(d_{i,j}\)我们都可以启动一个线程去执行计算,类似于CPU上的SIMD技术,GPU中的这项优化称为SIMT。而在Python中改造成GPU函数的方法也非常简单,只需要把函数前的修饰器改一下,去掉函数内部的for循环,就基本完成了,比如下面这个改造的近邻表计算的案例:
# cuda_neighbor_list.py
from numba import jit
from numba import cuda
import numpy as np
@jit
def neighbor_list(crd, neighbors, data_length, cutoff):
"""CPU based neighbor list calculation.
"""
for i in range(data_length):
for j in range(i+1, data_length):
if np.linalg.norm(crd[i]-crd[j]) <= cutoff:
neighbors[i][j] = 1
neighbors[j][i] = 1
return neighbors
@cuda.jit
def cuda_neighbor_list(crd, neighbors, cutoff):
"""GPU based neighbor list calculation.
"""
i, j = cuda.grid(2)
dis = ((crd[i][0]-crd[j][0])**2+\
(crd[i][1]-crd[j][1])**2+\
(crd[i][2]-crd[j][2])**2)**0.5
neighbors[i][j] = dis <= cutoff[0] and dis > 0
if __name__ == '__main__':
import time
np.random.seed(1)
atoms = 2**5
cutoff = 0.5
cutoff_cuda = cuda.to_device(np.array([cutoff]).astype(np.float32))
crd = np.random.random((atoms,3)).astype(np.float32)
crd_cuda = cuda.to_device(crd)
adjacent = np.zeros((atoms, atoms)).astype(np.float32)
adjacent_cuda = cuda.to_device(adjacent)
time0 = time.time()
adjacent_c = neighbor_list(crd, adjacent, atoms, cutoff)
time1 = time.time()
cuda_neighbor_list[(atoms, atoms), (1, 1)](crd_cuda,
adjacent_cuda,
cutoff_cuda)
time2 = time.time()
adjacent_g = adjacent_cuda.copy_to_host()
print ('The time cost of CPU with numba.jit is: {}s'.format(\
time1-time0))
print ('The time cost of GPU with cuda.jit is: {}s'.format(\
time2-time1))
print ('The result error is: {}'.format(np.sum(adjacent_c-\
adjacent_g)))
需要说明的是,当前Numba并未支持所有的numpy的函数,因此有一些计算的功能需要我们自己去手动实现一下,比如计算一个Norm的值。这里我们在输出结果中不仅统计了结果的正确性,也给出了运行的时间:
$ python3 cuda_neighbor_list.py
The time cost of CPU with numba.jit is: 0.6401469707489014s
The time cost of GPU with cuda.jit is: 0.19208502769470215s
The result error is: 0.0
需要说明的是,这里仅仅运行了一次的程序,而jit即时编译的加速效果在第一次的运行中其实并不明显,甚至还有一些速度偏慢,但是在后续过程的函数调用中,就能够起到比较大的加速效果。所以这里的运行时间并没有太大的代表性,比较有代表性的时间对比可以看如下的案例:
# cuda_neighbor_list.py
from numba import jit
from numba import cuda
import numpy as np
@jit
def neighbor_list(crd, neighbors, data_length, cutoff):
"""CPU based neighbor list calculation.
"""
for i in range(data_length):
for j in range(i+1, data_length):
if np.linalg.norm(crd[i]-crd[j]) <= cutoff:
neighbors[i][j] = 1
neighbors[j][i] = 1
return neighbors
@cuda.jit
def cuda_neighbor_list(crd, neighbors, cutoff):
"""GPU based neighbor list calculation.
"""
i, j = cuda.grid(2)
dis = ((crd[i][0]-crd[j][0])**2+\
(crd[i][1]-crd[j][1])**2+\
(crd[i][2]-crd[j][2])**2)**0.5
neighbors[i][j] = dis <= cutoff[0] and dis > 0
if __name__ == '__main__':
import time
np.random.seed(1)
atoms = 2**10
cutoff = 0.5
cutoff_cuda = cuda.to_device(np.array([cutoff]).astype(np.float32))
crd = np.random.random((atoms,3)).astype(np.float32)
crd_cuda = cuda.to_device(crd)
adjacent = np.zeros((atoms, atoms)).astype(np.float32)
adjacent_cuda = cuda.to_device(adjacent)
time_c = 0.0
time_g = 0.0
for _ in range(100):
time0 = time.time()
adjacent_c = neighbor_list(crd, adjacent, atoms, cutoff)
time1 = time.time()
cuda_neighbor_list[(atoms, atoms), (1, 1)](crd_cuda,
adjacent_cuda,
cutoff_cuda)
time2 = time.time()
if _ != 0:
time_c += time1 - time0
time_g += time2 - time1
print ('The total time cost of CPU with numba.jit is: {}s'.format(\
time_c))
print ('The total time cost of GPU with cuda.jit is: {}s'.format(\
time_g))
这个案例中也没有修改较多的地方,只是把一次计算的时间调整为多次计算的时间,并且忽略第一次计算过程中的即时编译,最终输出结果如下:
$ python3 cuda_neighbor_list.py
The total time cost of CPU with numba.jit is: 14.955506563186646s
The total time cost of GPU with cuda.jit is: 0.018685102462768555s
可以看到,在GPU加速后,相比于CPU的高性能运算,能够提速达将近1000倍!
总结概要
对于Pythoner而言,苦其性能已久。如果能够用一种非常Pythonic的方法来实现GPU的加速效果,对于Pythoner而言无疑是巨大的好消息,Numba就为我们提供了这样的一个基础功能。本文通过一个近邻表计算的案例,给出了适用于GPU加速的计算场景。这种计算场景可并行化的程度较高,而且函数会被多次用到(在分子动力学模拟的过程中,每一个step都会调用到这个函数),因此这是一种最典型的、最适用于GPU加速场景的案例。
版权声明
本文首发链接为:https://www.cnblogs.com/dechinphy/p/cuda-neighbor.html
作者ID:DechinPhy
更多原著文章请参考:https://www.cnblogs.com/dechinphy/
打赏专用链接:https://www.cnblogs.com/dechinphy/gallery/image/379634.html
腾讯云专栏同步:https://cloud.tencent.com/developer/column/91958
Python的GPU编程实例——近邻表计算的更多相关文章
- JAX-MD在近邻表的计算中,使用了什么奇技淫巧?(一)
技术背景 JAX-MD是一款基于JAX的纯Python高性能分子动力学模拟软件,应该说在纯Python的软件中很难超越其性能.当然,比一部分直接基于CUDA的分子动力学模拟软件性能还是有些差距.而在计 ...
- Python进阶:函数式编程实例(附代码)
Python进阶:函数式编程实例(附代码) 上篇文章"几个小例子告诉你, 一行Python代码能干哪些事 -- 知乎专栏"中用到了一些列表解析.生成器.map.filter.lam ...
- Serverless助力AI计算:阿里云ACK Serverless/ECI发布GPU容器实例
ACK Serverless(Serverless Kubernetes)近期基于ECI(弹性容器实例)正式推出GPU容器实例支持,让用户以serverless的方式快速运行AI计算任务,极大降低AI ...
- Python Socket,How to Create Socket Cilent? - 网络编程实例
文章出自:Python socket – network programming tutorial by Silver Moon 原创译文,如有版权问题请联系删除. Network programin ...
- python学习_数据处理编程实例(二)
在上一节python学习_数据处理编程实例(二)的基础上数据发生了变化,文件中除了学生的成绩外,新增了学生姓名和出生年月的信息,因此将要成变成:分别根据姓名输出每个学生的无重复的前三个最好成绩和出生年 ...
- Python并发编程实例教程
有关Python中的并发编程实例,主要是对Threading模块的应用,文中自定义了一个Threading类库. 一.简介 我们将一个正在运行的程序称为进程.每个进程都有它自己的系统状态,包含内存状态 ...
- Tableau学习Step5一表计算、详细级别表达式、动作、外接python
Tableau学习Step5一表计算.详细级别表达式.动作.外接python 本文首发于博客冰山一树Sankey,去博客浏览效果更好. ) Tableau学习Step4一数据解释.异常值监测.参数使用 ...
- Python黑帽编程2.2 数值类型
Python黑帽编程2.2 数值类型 数值类型,说白了就是处理各种各样的数字,Python中的数值类型包括整型.长整型.布尔.双精度浮点.十进制浮点和复数,这些类型在很多方面与传统的C类型有很大的区 ...
- Python黑帽编程2.3 字符串、列表、元组、字典和集合
Python黑帽编程2.3 字符串.列表.元组.字典和集合 本节要介绍的是Python里面常用的几种数据结构.通常情况下,声明一个变量只保存一个值是远远不够的,我们需要将一组或多组数据进行存储.查询 ...
随机推荐
- ThinkPHP3.2.3 语言包切换中英文切换
今天要用ThinkPHP3.2.3做一个小网站,其中涉及到切换中文与英文,通过查询手册和百度实现了该操作,现在将我具体的操作步骤记录下来,作为笔记和大家分享. php开发框架:ThinkPHP3.2. ...
- 记一次lombok踩坑记
引言 今天中午正在带着耳机遨游在代码的世界里,被运营在群里@了,气冲冲的反问我最近有删生产的用户数据的吗?我肯定客气的回答道没有呀?生产的数据我怎么能随随便便可以删除,这可是公司的红线,再说了我也没有 ...
- javascript里面的document.getElementById
一.getElementById:获取对 ID 标签属性为指定值的第一个对象的引用,它有 value 和 length 等属性 1.获取当前页面的值input标签值:var attr1=documen ...
- Windows10 64位解决无法使用Microsoft.Jet.OLEDB.4.0的方法
本机软件环境:Windows10 64位+Office2003 (32位) ============================================= 1.下载 ACE2010的驱动, ...
- Python基础之tabview
以前写过界面,但是没有记录下来,以至于现在得从头学习一次,论做好笔记的重要性. 现在学习的是怎么写一个tabview出来,也就是用tkinter做一个界面切换的效果.参考链接:https://blog ...
- Linux 查看内存命令
Linux 查看内存命令 top命令, Linux的top命令提供Linux资源使用情况的实时更新信息.不仅可以查看Linux内存,也可以查看CPU以及各个进程之间的对资源的占用情况.使用方式如下: ...
- YsoSerial 工具常用Payload分析之Common-Collections2、4(五)
前言 Common-Collections <= 3.2.1 对应与YsoSerial为CC1.3.5.6.7 ,Commno-collections4.0对应与CC2.4. 这篇文章结束官方原 ...
- openssl查看证书命令
openssl x509部分命令打印出证书的内容:openssl x509 -in cert.pem -noout -text打印出证书的系列号openssl x509 -in cert.pem -n ...
- [CISCN2019 华北赛区 Day2 Web1]Hack World(二分法写布尔注入脚本)
记一道布尔注入的题,存在过滤字符. 从题目看应该是一道注入题.提示存在flag表flag列. 输入1和2的返回结果不一样,可能是布尔注入. 简单用万能密码尝试了一下.提示SQL Injection C ...
- Bugku-misc 1-8题总结
1.签到题 略过 2.这是一张单纯的图片 拉入winhex,在最后面有一段Uniocde编码,解码得到flag. 3.隐写 题目是隐写,binwalk打开分析 得到两个Zlib(提供数据压缩用的函式库 ...