matmul#

nvmath.device.matmul(*, compiler=None, **kwargs)[源代码]#

创建一个 BlasOptions 对象，该对象封装了一个已编译且可直接使用的用于矩阵乘法的设备函数。

参数:

  • size – 一个整数序列，表示矩阵乘法问题的三个维度 (m, n, k)。

  • precision – 计算精度，指定为 numpy 浮点 dtype，目前支持 numpy.float16、numpy.float32 和 numpy.float64。

  • data_type – 输入矩阵的数据类型，可以是 'real' 或 'complex'。

  • compiler – 一个字符串，用于指定设备代码的编译器，目前支持 None (默认) 和 'Numba'

  • code_type (CodeType) – 目标 GPU 代码和计算能力。

  • block_size (int) – 总块大小，可选。如果未提供或设置为 'suggested'，则将设置为 1D 块维度的建议值。

  • block_dim (Dim3) – 用于启动 CUDA 内核的块维度，可选。如果未提供或设置为 'suggested'，则将设置为建议值。当显式指定 block_size 时，不能使用此参数。

  • leading_dimension (LeadingDimension) – 输入矩阵的前导维度，可选。如果未提供，将设置为匹配矩阵的行/列维度。或者，如果作为 'suggested' 提供，则将设置为最佳性能的建议值。

  • transpose_mode (TransposeMode) – 所有输入矩阵的转置模式。如果未提供，则默认不进行转置。

  • function (str) – 一个字符串，用于指定函数名称。目前支持 'MM' (默认) 用于矩阵乘法。

  • execution (str) – 一个字符串，用于指定执行方法，可以是 'Block' 或 'Thread'。

另请参阅

BlasOptions 的属性与 CUDA C++ cuBLASDx API 提供了 1:1 的映射。更多详细信息，请参阅 cuBLASDx 文档。

示例

>>> from numba import cuda
>>> from nvmath.device import matmul
>>> import numpy as np
>>> m, n, k = 32, 16, 64
>>> block_size = 256

使用 nvmath.device.matmul() 创建编译后的矩阵乘法对象

>>> MM = matmul(
...     size=(m, n, k),
...     precision=np.float32,
...     data_type="real",
...     transpose_mode=("non_transposed", "transposed"),
...     execution="Block",
...     block_size=block_size,
...     compiler="numba",
... )

在定义内核时，将 link=MM.files 传递给 numba.cuda.jit() 装饰器，以便与编译后的代码链接。

cuBLASDx 适用于共享内存数组。它需要列优先（F 顺序）数组，但 cuda.shared.array 仅创建行优先（C 顺序）数组。您可以通过翻转维度来模拟列优先数组。在您的共享内存数组准备就绪并填充实际数据后，您可以通过调用 MM 来运行矩阵乘法

>>> a_dim, b_dim, c_dim = MM.a_dim, MM.b_dim, MM.c_dim
>>> @cuda.jit(link=MM.files)
... def f():
...     a = cuda.shared.array(shape=(a_dim[1], a_dim[0]), dtype=np.float32)
...     b = cuda.shared.array(shape=(b_dim[1], b_dim[0]), dtype=np.float32)
...     c = cuda.shared.array(shape=(c_dim[1], c_dim[0]), dtype=np.float32)
...     # TODO: Populate the arrays with actual data.
...     alpha, beta = 1.0, 0.0
...     MM(alpha, a, b, beta, c)
...     cuda.syncthreads()
...     # TODO: Copy the result (c) from the shared memory
>>> f[1, block_size]()

更多示例可以在 nvmath/examples/device 目录中找到。

本页内容