matmul#

nvmath.linalg.advanced.matmul(
a,
b,
/,
c=None,
*,
alpha=None,
beta=None,
epilog=None,
epilog_inputs=None,
qualifiers=None,
options=None,
preferences=None,
algorithm=None,
stream=None,
)[source]#

执行指定的矩阵乘法计算 \(F(\alpha a @ b + \beta c)\)，其中 \(F\) 是后记。此函数形式是围绕有状态的 Matmul 对象 API 的包装器，旨在用于单次使用（例如，用户只需要执行一次矩阵乘法），在这种情况下，不可能分摊准备成本。

关于此函数内部发生情况的详细信息，可以通过将 logging.Logger 对象传递给 MatmulOptions 或通过在默认使用的根 logger 对象中设置适当的选项来获得

>>> import logging
>>> logging.basicConfig(
...     level=logging.INFO,
...     format="%(asctime)s %(levelname)-8s %(message)s",
...     datefmt="%m-%d %H:%M:%S",
... )

用户可以选择所需的日志记录级别，并且通常可以利用 Python logging 模块提供的所有功能。

参数:

  • a – 表示矩阵乘法中第一个操作数的张量（参见 Semantics）。目前支持的类型包括 numpy.ndarray、 cupy.ndarray 和 torch.Tensor。

  • b – 表示矩阵乘法中第二个操作数的张量（参见 Semantics）。目前支持的类型包括 numpy.ndarray、 cupy.ndarray 和 torch.Tensor。

  • c –

    (可选) 表示要添加到矩阵乘法结果中的操作数的张量（参见 Semantics）。目前支持的类型包括 numpy.ndarray、 cupy.ndarray 和 torch.Tensor。

    注意

    1-D（向量）c 的广播行为与等效的 NumPy 表达式有所不同。使用 nvmath-python，c 在内部被提升为形状 (M, 1)，以便与 a @ b 进行广播；这与 cuBLASLt 的行为相匹配。使用 NumPy，1-D c 在表达式 a @ b + c 中表现得好像它具有形状 (1, N)。

    版本 0.2.1 中已弃用: 为了避免广播行为的歧义，从版本 0.3.0 开始，nvmath-python 将不再接受 1-D（向量）c。请使用单例维度将您的输入数组转换为 2-D。

  • alpha – 矩阵乘法项的比例因子，为实数或复数。默认值为 \(1.0\)。

  • beta – 矩阵加法项的比例因子，为实数或复数。如果指定了操作数 c，则必须提供 beta 的值。来自先前计划和自动调整的矩阵乘法。

  • epilog – 将 epilog \(F\) 指定为 MatmulEpilog 类型的对象，以应用于矩阵乘法的结果：\(F(\alpha A @ B + \beta C\))。默认情况下没有 epilog。有关可用 epilog 的列表，请参见 cuBLASLt 文档。

  • epilog_inputs – 将所选 epilog 所需的额外输入指定为字典，其中键是 epilog 输入名称，值是 epilog 输入。epilog 输入必须是与操作数具有相同包并在相同内存空间中的张量（有关操作数的更多信息，请参见构造函数）。如果未提供所需的 epilog 输入，则会引发异常，其中列出了所需的 epilog 输入。某些 epilog 输入由其他 epilog 生成。例如，MatmulEpilog.DRELU 的 epilog 输入由使用 MatmulEpilog.RELU_AUX 的相同操作数进行矩阵乘法生成。

  • qualifiers – 如果需要，将矩阵限定符指定为长度为 3 的 numpy.ndarray 类型的 matrix_qualifiers_dtype 对象，分别对应于操作数 a、 b 和 c。

  • options – 将矩阵乘法的选项指定为 MatmulOptions 对象。或者，也可以提供包含 MatmulOptions 构造函数参数的 dict。如果未指定，则该值将设置为默认构造的 MatmulOptions 对象。

  • preferences – 此参数将规划的首选项指定为 MatmulPlanPreferences 对象。或者，也可以提供包含 MatmulPlanPreferences 构造函数参数的字典。如果未指定，则该值将设置为默认构造的 MatmulPlanPreferences 对象。

  • algorithm – 如果需要绕过规划，可以直接提供 Algorithm 类型的对象。算法对象必须与矩阵乘法兼容。此选项的典型用途是提供从先前计划和自动调整的矩阵乘法中序列化（pickle 化）的算法。

  • stream – 提供 CUDA 流以用于执行操作。可接受的输入包括 cudaStream_t (作为 Python int)、 cupy.cuda.Stream 和 torch.cuda.Stream。如果未提供流，则将使用来自操作数包的当前流。

返回值:

指定矩阵乘法（应用 epilog）的结果，该结果保留在与输入操作数相同的设备上并属于同一包。如果使用导致额外输出的 epilog（例如 nvmath.linalg.advanced.MatmulEpilog.RELU_AUX），则返回一个元组，其中第一个元素是矩阵乘法结果（应用 epilog），第二个元素是由所选 epilog 提供的辅助输出，类型为 dict。

语义

矩阵乘法的语义遵循 numpy.matmul() 语义，但对广播有一些限制。此外，矩阵融合加法的语义描述如下

  • 如果参数 a 和 b 是矩阵，则它们根据矩阵乘法规则相乘。

  • 如果参数 a 是 1-D 的，则通过在其维度前缀 1 将其提升为矩阵。矩阵乘法后，前缀的 1 将从结果的维度中移除。

  • 如果参数 b 是 1-D 的，则通过在其维度后缀 1 将其提升为矩阵。矩阵乘法后，后缀的 1 将从结果的维度中移除。

  • 如果 a 或 b 是 N-D (N > 2)，则该操作数被视为一批矩阵。如果 a 和 b 都是 N-D，则它们的批次维度必须匹配。如果 a 或 b 中只有一个是 N-D，则另一个操作数将被广播。

  • 矩阵加法的操作数 c 可以是长度为 M 的向量、形状为 (M, 1) 或 (M, N) 的矩阵，或后者的批次版本（…, M, 1）或（…, M, N）。这里 M 和 N 是矩阵乘法结果的维度。如果提供向量或 N = 1，则 c 的列将为加法进行广播。如果不存在批次维度，则 c 将根据需要跨批次进行广播。

  • 类似地，当对批次进行操作时，所有 epilog 的辅助输出均为 3-D。因此，在非批次模式下返回长度为 N 的 1-D 向量的 epilog 在批次模式下返回大小为 (batch, N, 1) 的 3-D 矩阵。

另请参阅

Matmul、 MatmulOptions、 MatmulEpilog、 MatmulPlanPreferences

示例

>>> import cupy as cp
>>> import nvmath

在 GPU 上创建三个 float32 ndarray

>>> M, N, K = 128, 64, 256
>>> a = cp.random.rand(M, K, dtype=cp.float32)
>>> b = cp.random.rand(K, N, dtype=cp.float32)
>>> c = cp.random.rand(M, N, dtype=cp.float32)

使用 matmul() 执行操作 \(\alpha A @ B + \beta C\)。结果 r 也是一个 CuPy float64 ndarray

>>> r = nvmath.linalg.advanced.matmul(a, b, c, alpha=1.23, beta=0.74)

也可以使用 epilog。这里我们执行 \(RELU(\alpha A @ B + \beta C)\)

>>> epilog = nvmath.linalg.advanced.MatmulEpilog.RELU
>>> r = nvmath.linalg.advanced.matmul(a, b, c, alpha=1.23, beta=0.74, epilog=epilog)

可以提供选项来自定义操作

>>> compute_type = nvmath.linalg.advanced.MatmulComputeType.COMPUTE_32F_FAST_TF32
>>> o = nvmath.linalg.advanced.MatmulOptions(compute_type=compute_type)
>>> r = nvmath.linalg.advanced.matmul(a, b, options=o)

有关可用选项的完整列表，请参见 MatmulOptions。

包的当前流默认使用，但可以为 Matmul 操作显式提供流。如果操作数在不同的流上计算，则可以这样做，例如

>>> s = cp.cuda.Stream()
>>> with s:
...     a = cp.random.rand(M, K)
...     b = cp.random.rand(K, N)
>>> r = nvmath.linalg.advanced.matmul(a, b, stream=s)

上面的操作在流 s 上运行，并相对于输入计算进行排序。

在 CPU 上创建 NumPy ndarray。

>>> import numpy as np
>>> a = np.random.rand(M, K)
>>> b = np.random.rand(K, N)

将 NumPy ndarray 提供给 matmul()，结果也是 NumPy ndarray

>>> r = nvmath.linalg.advanced.matmul(a, b)

注释

  • 此函数是 Matmul 的便捷包装器，专门用于单次使用。

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

在本页中