Matmul#

class nvmath.linalg.advanced.Matmul(
a,
b,
/,
c=None,
*,
alpha=None,
beta=None,
qualifiers=None,
options=None,
stream=None,
)[source]#

创建封装指定矩阵乘法计算 \(\alpha a @ b + \beta c\) 和执行该操作所需资源的状态对象。状态对象可用于分摊跨多次执行的准备成本（矩阵乘法情况下的规划）（另请参阅有状态 API 部分）。

函数形式 API matmul() 是使用状态对象进行单次使用的便捷替代方法（例如，用户只需要执行一次矩阵乘法），在这种情况下，没有可能分摊准备成本。函数形式 API 只是状态对象 API 的便捷包装器。

使用状态对象通常涉及以下步骤

  1. 问题规范：使用定义的操作和选项初始化对象。

  2. 准备：使用 plan() 确定此特定矩阵乘法操作的最佳算法实现。

  3. 执行：使用 execute() 执行矩阵乘法计算。

  4. 资源管理：通过显式调用 free() 或通过在上下文管理器中管理状态对象，确保释放所有资源。

通过将 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 – 表示矩阵乘法的第一个操作数的张量（参见 语义）。当前支持的类型为 numpy.ndarray、cupy.ndarray 和 torch.Tensor。

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

  • c –

    （可选）表示要添加到矩阵乘法结果的操作数的张量（参见 语义）。当前支持的类型为 numpy.ndarray、cupy.ndarray 和 torch.Tensor。

    注意

    1-D（向量）c 的广播行为与等效的 NumPy 表达式有所不同。使用 nvmath-python，为了与 a @ b 进行广播，c 会在内部提升为形状 (M, 1)；这与 cuBLASLt 的行为相匹配。使用 NumPy，在表达式 a @ b + c 中，1-D c 的行为就像其形状为 (1, N) 一样。

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

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

  • beta – 矩阵加法项的比例因子，为实数或复数。如果指定了操作数 c，则必须提供 beta 的值。

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

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

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

语义

矩阵乘法的语义遵循 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 将根据需要跨批次广播。

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

另请参阅

autotune(), plan(), reset_operands(), execute()

示例

>>> import numpy as np
>>> import nvmath

在 CPU 上创建两个 2-D float64 ndarray

>>> M, N, K = 1024, 1024, 1024
>>> a = np.random.rand(M, K)
>>> b = np.random.rand(K, N)

我们将使用专门的矩阵乘法接口定义一个矩阵乘法运算，后跟一个 RELU 后记函数。

创建一个封装上述问题规范的 Matmul 对象

>>> mm = nvmath.linalg.advanced.Matmul(a, b)

可以在上面提供选项，以使用 options 参数控制操作的行为（请参阅 MatmulOptions）。

接下来，规划操作。指定了后记，并且可以选择性地指定规划的首选项

>>> epilog = nvmath.linalg.advanced.MatmulEpilog.RELU
>>> algorithms = mm.plan(epilog=epilog)

某些后记选择（如 nvmath.linalg.advanced.MatmulEpilog.BIAS）需要使用 plan() 的 epilog_inputs 参数提供额外的输入。

现在执行矩阵乘法，并获取 NumPy ndarray 形式的结果 r1。

>>> r1 = mm.execute()

最后，释放对象的资源。为了避免必须显式调用此方法，建议尽可能使用如下所示的上下文管理器来使用 Matmul 对象。

>>> mm.free()

请注意，默认情况下，所有 Matmul 方法都在当前流上执行。或者，可以使用 stream 参数在指定的流上运行方法。

现在让我们看看在 GPU 上使用 CuPy ndarray 的相同问题。

在 GPU 上创建 3-D complex128 CuPy ndarray

>>> import cupy as cp
>>> a = cp.random.rand(M, K)
>>> b = cp.random.rand(K, N)

创建一个封装前面描述的问题规范的 Matmul 对象，并将其用作上下文管理器。

>>> with nvmath.linalg.advanced.Matmul(a, b) as mm:
...     algorithms = mm.plan(epilog=epilog)
...
...     # Execute the operation to get the first result.
...     r1 = mm.execute()
...
...     # Update operands A and B in-place (see reset_operands() for an
...     # alternative).
...     a[:] = cp.random.rand(M, K)
...     b[:] = cp.random.rand(K, N)
...
...     # Execute the operation to get the new result.
...     r2 = mm.execute()

对象使用的所有资源都在块的末尾释放。

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

方法

__init__(
a,
b,
/,
c=None,
*,
alpha=None,
beta=None,
qualifiers=None,
options=None,
stream=None,
)[source]#
applicable_algorithm_ids(limit=8)[source]#

获取适用于此矩阵乘法的算法 ID。

参数:

limit – 所需的适用算法 ID 的最大数量

返回:

适用于此矩阵乘法问题规范的算法 ID 序列，顺序随机。

autotune(
iterations=3,
prune=None,
release_workspace=False,
stream=None,
)[source]#

自动调整矩阵乘法，以按测量的执行时间从最快到最慢的顺序排列算法。自动调整后，可以使用 algorithms 访问最佳排序的算法序列。

参数:

  • iterations – 要执行的自动调整迭代次数。

  • prune – 整数 N，指定自动调整后要保留的前 N 个最快算法。默认值是保留所有算法。

  • release_workspace – True 值指定状态对象应在函数返回时将工作区内存释放回包内存池，而 False 值指定对象应保留内存。如果应用程序在连续调用（相同或不同的）execute() API 之间执行其他消耗大量内存的操作，则可以将此选项设置为 True，但这会因每次调用时从包内存池获取和释放工作区内存而产生少量开销。默认值为 False。

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

execute(
*,
algorithm=None,
release_workspace=False,
stream=None,
)[source]#

执行准备好的（已规划且可能已自动调整的）矩阵乘法。

参数:

  • algorithm – （实验性）从 plan() 或 algorithms 返回的序列中选择的算法。默认情况下，使用序列中的第一个算法。

  • release_workspace – True 值指定状态对象应在函数返回时将工作区内存释放回包内存池，而 False 值指定对象应保留内存。如果应用程序在连续调用（相同或不同的）execute() API 之间执行其他消耗大量内存的操作，则可以将此选项设置为 True，但这会因每次调用时从包内存池获取和释放工作区内存而产生少量开销。默认值为 False。

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

返回:

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

free()[source]#

释放 Matmul 资源。

建议在上下文中使用 Matmul 对象，但如果不可能，则必须显式调用此方法，以确保矩阵乘法资源（尤其是内部库对象）得到正确清理。

plan(
*,
preferences=None,
algorithms=None,
epilog=None,
epilog_inputs=None,
stream=None,
)[source]#

规划矩阵乘法运算，考虑后记（如果提供）。

参数:

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

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

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

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

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

返回:

适用于此矩阵乘法问题规范的 nvmath.linalg.advanced.Algorithm 对象序列，按从最快到最慢的启发式顺序排列。

注释

名称中包含 BIAS 的后记需要一个键为 'bias' 的后记输入。包含 DRELU 的后记需要一个键为 'relu_aux' 的后记输入，该输入在“前向传递”后记（如 RELU_AUX 或 RELU_AUX_BIAS）中生成。类似地，名称中包含 DGELU 的后记需要一个键为 'gelu_aux' 的后记输入，该输入在相应的正向传递操作中生成。

示例

>>> import numpy as np
>>> import nvmath

在 CPU 上创建两个 3-D float64 ndarray，表示批处理矩阵，以及一个偏置向量

>>> batch = 32
>>> M, N, K = 1024, 1024, 1024
>>> a = np.random.rand(batch, M, K)
>>> b = np.random.rand(batch, K, N)
>>> # The bias vector will be broadcast along the columns, as well as along the
>>> # batch dimension.
>>> bias = np.random.rand(M)

我们将定义一个矩阵乘法运算，后跟一个 nvmath.linalg.advanced.MatmulEpilog.RELU_BIAS 后记函数。

>>> with nvmath.linalg.advanced.Matmul(a, b) as mm:
...     # Plan the operation with RELU_BIAS epilog and corresponding epilog
...     # input.
...     p = nvmath.linalg.advanced.MatmulPlanPreferences(limit=8)
...     epilog = nvmath.linalg.advanced.MatmulEpilog.RELU_BIAS
...     epilog_inputs = {"bias": bias}
...     # The preferences can also be provided as a dict: {'limit': 8}
...     algorithms = mm.plan(
...         preferences=p,
...         epilog=epilog,
...         epilog_inputs=epilog_inputs,
...     )
...
...     # Execute the matrix multiplication, and obtain the result `r` as a
...     # NumPy ndarray.
...     r = mm.execute()

某些后记（如 nvmath.linalg.advanced.MatmulEpilog.RELU_AUX）会生成辅助输出。

>>> with nvmath.linalg.advanced.Matmul(a, b) as mm:
...     # Plan the operation with RELU_AUX epilog>
...     epilog = nvmath.linalg.advanced.MatmulEpilog.RELU_AUX
...     algorithms = mm.plan(epilog=epilog)
...
...     # Execute the matrix multiplication, and obtain the result `r` along
...     # with the auxiliary output.
...     r, auxiliary = mm.execute()

辅助输出是一个 Python dict，其中每个辅助输出的名称作为键。

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

reset_operands(
a=None,
b=None,
c=None,
*,
alpha=None,
beta=None,
epilog_inputs=None,
stream=None,
)[source]#

重置此 Matmul 实例持有的操作数。

此方法有两个用例

  1. 当原始操作数在 CPU 上时，它可用于为执行提供新的操作数

  2. 它可用于释放对先前操作数的内部引用，并通过为所有参数传递 None 使其内存可用于其他用途。在这种情况下，必须再次调用此方法以提供所需的操作数，然后才能再次调用执行 API（如 autotune() 或 execute()）。

当操作数驻留在 GPU 上并且使用就地操作来更新操作数值时，不需要此方法。

此方法将对新操作数执行各种检查，以确保

  • 形状、步幅、数据类型与旧操作数的数据类型匹配。

  • 操作数所属的包与旧操作数的包匹配。

  • 如果输入张量在 GPU 上，则设备必须匹配。

参数:

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

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

  • c –

    （可选）表示要添加到矩阵乘法结果的操作数的张量（参见 语义）。当前支持的类型为 numpy.ndarray、cupy.ndarray 和 torch.Tensor。

    注意

    1-D（向量）c 的广播行为与等效的 NumPy 表达式有所不同。使用 nvmath-python，为了与 a @ b 进行广播，c 会在内部提升为形状 (M, 1)；这与 cuBLASLt 的行为相匹配。使用 NumPy，在表达式 a @ b + c 中，1-D c 的行为就像其形状为 (1, N) 一样。

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

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

  • beta – 矩阵加法项的比例因子，为实数或复数。如果指定了操作数 c，则必须提供 beta 的值。

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

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

示例

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

在 GPU 上创建两个 3-D float64 ndarray

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

创建一个矩阵乘法对象作为上下文管理器

>>> with nvmath.linalg.advanced.Matmul(a, b) as mm:
...     # Plan the operation.
...     algorithms = mm.plan()
...
...     # Execute the MM to get the first result.
...     r1 = mm.execute()
...
...     # Reset the operands to new CuPy ndarrays.
...     c = cp.random.rand(M, K)
...     d = cp.random.rand(K, N)
...     mm.reset_operands(c, d)
...
...     # Execute to get the new result corresponding to the updated operands.
...     r2 = mm.execute()

请注意，如果仅重置操作数的子集，则未重置的操作数将保留其原始值。

使用 reset_operands()，可以实现最小的开销，因为问题规范和规划仅执行一次。

对于上面的特定示例，显式调用 reset_operands() 等效于就地更新操作数，即，将 mm.reset_operand(c, d) 替换为 a[:]=c 和 b[:]=d。请注意，应谨慎采用就地更新操作数，因为它只有在以下附加约束条件下才能产生预期的结果

  • 操作数在 GPU 上（更准确地说，应可以从执行空间访问操作数内存空间）。

有关更多详细信息，请参阅 就地更新示例。

属性

algorithms#

在使用 plan() 进行规划后，获取算法对象序列以查询其功能、配置它们或序列化它们以供以后使用。

返回:

适用于此矩阵乘法问题规范的 nvmath.linalg.advanced.Algorithm 对象序列。