实现CUDA核函数——二维矩阵相乘

bachelor2001

已于 2022-02-20 00:36:02 修改

阅读量1k

点赞数

分类专栏： AI 文章标签：算法 cuda 矩阵核函数 python

于 2022-02-19 23:55:29 首次发布

本文链接：https://blog.youkuaiyun.com/bachelor2001/article/details/123025739

版权

本文介绍了作者使用CUDA实现二维矩阵相乘的过程，特别强调了并发线程数与矩阵计算数据量的自适应性。通过与numpy和torch库计算结果的对比验证，确保了CUDA实现的正确性。实验环境为win10，使用PyCUDA库，Python版本为3.9.7。代码实现中，注意grid和block参数不能超过GPU限制，需预先获取相关信息。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

初学CUDA，自己练手实现二维矩阵乘法，且并发线程数与矩阵计算数据量自适应。计算结果，与调用numpy、torch库的计算方法进行比较验证通过。

环境：
win10，PyCUDA: (2021, 1)，Python: 3.9.7

说明：调用时传入的grid、block参数，不能超过GPU的限制，所以需要先调用pycuda.driver库进行相关操作获取，网上有很多资料。

代码如下：

# -*- coding: utf-8 -*-
"""
Created on Sat Feb 19 21:33:17 2022

@author: xuning
"""

import pycuda.autoinit
from pycuda.compiler import SourceModule
import pycuda.gpuarray as gpuarray
import numpy as np
import torch
import sys
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity='all'

print("python ready")
print(f'The version of PyCUDA: {pycuda.VERSION}')
print(f'The version of Python: {sys.version}')

mod = SourceModule(r"""

/* gpuMatMultKernel_ex：GPU下矩阵乘法核函数,扩展版本
*  a:第一个矩阵指针，表示a[M][N]
*  b:第二个矩阵指针，表示b[N][S]
*  result:结果矩阵&