API 参考 - cuTENSORMg

通用

通用函数初始化库并创建张量描述符。


cutensorMgHostDevice_t

enum cutensorMgHostDevice_t

用于主机端张量的枚举设备代码。

enumerator CUTENSOR_MG_DEVICE_HOST

内存位于主机上的常规内存中。

enumerator CUTENSOR_MG_DEVICE_HOST_PINNED

内存位于主机上的已锁定内存中。


cutensorMgHandle_t

typedef struct cutensorMgHandle_s *cutensorMgHandle_t

编码参与操作的设备。

句柄包含有关参与操作的每个设备以及用于编排主机到设备操作的主机线程的信息。


cutensorMgTensorDescriptor_t

typedef struct cutensorMgTensorDescriptor_s *cutensorMgTensorDescriptor_t

表示可能分布的张量。

张量以块循环方式在设备上布局。它可以完全位于主机上,也可以分布在多个设备上。


cutensorMgCreate()

cutensorStatus_t cutensorMgCreate(cutensorMgHandle_t *handle, uint32_t numDevices, const int32_t devices[])

创建库句柄。

句柄包含有关应参与计算的设备的信息。 持有任何张量数据或参与任何 cuTENSORMg 操作的所有设备也应包含在句柄中。 每个设备在列表中只能出现一次。 建议所有设备都是相同的(即,具有相同的峰值性能)以避免负载平衡问题,并通过 NVLink 连接以避免代价高昂的设备-主机-设备传输。 如果可能,此调用将启用已传递给它的所有设备之间的对等互连。

备注

阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:
  • handle[out] 结果库句柄。

  • numDevices[in] 参与所有后续计算的设备数量。

  • devices[in] 参与所有计算的设备。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。


cutensorMgDestroy()

cutensorStatus_t cutensorMgDestroy(cutensorMgHandle_t handle)

销毁库句柄。

所有未完成的操作必须在调用此函数之前完成。释放所有相关资源。 使用句柄创建的任何描述符或计划都将变为无效,并且只能被销毁。

备注

阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:

handle[in] 要销毁的句柄。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。


cutensorMgCreateTensorDescriptor()

cutensorStatus_t cutensorMgCreateTensorDescriptor(cutensorMgHandle_t handle, cutensorMgTensorDescriptor_t *desc, uint32_t numModes, const int64_t extent[], const int64_t elementStride[], const int64_t blockSize[], const int64_t blockStride[], const int32_t deviceCount[], uint32_t numDevices, const int32_t devices[], cudaDataType_t type)

创建张量描述符。

张量描述符完全指定(可能)分布式张量的数据布局。 它主要通过五个数据片段来实现:范围、元素步幅、块大小、块步幅和设备计数。

范围描述了每个张量模式的总大小。 例如,一个 9x9 矩阵的范围为 9 和 9。

块大小描述了数据如何被分块。 例如,如果块大小为 4x2,则第一个模式有 3 个块,第二个模式有 5 个块。

4 x 2

4 x 2

1 x 2

4 x 2

4 x 2

1 x 2

4 x 2

4 x 2

1 x 2

4 x 2

4 x 2

1 x 2

4 x 1

4 x 1

1 x 1

然后,设备计数描述了块在每个模式中分布在多少个设备上。 例如,设备计数为 2x2 将意味着块在每个模式中分布在两个设备上,即总共四个设备。 设备首先沿第一个模式排列,然后沿第二个模式排列,如下所示

设备 0

设备 2

设备 1

设备 3

特别是,设备 0 将拥有第一个维度中的第一个和第三个块,第二个维度中的第一个、第三个和第五个块(总共六个块),设备 1 将拥有第一个维度中的第一个和第三个块,以及第二个维度中的第二个和第四个块(总共四个块),设备 2 将拥有第一个维度中的第二个块,以及第二个维度中的第一个、第三个和第五个块(总共三个块),最后,设备 3 将拥有第一个维度中的第二个块和第二个维度中的第二个和第四个块(总共两个块)。

设备 0

设备 2

设备 0

设备 1

设备 3

设备 1

设备 0

设备 2

设备 0

设备 1

设备 3

设备 1

设备 0

设备 2

设备 0

然后,元素步幅和块步幅描述了块如何在各个设备上布局,即该模式下元素和块之间的距离。 最后,devices 数组描述了块映射到的设备。 在这里,允许指定 CUTENSOR_MG_DEVICE_HOST 以表示这些块位于主机上。 张量必须完全位于设备上或完全位于主机上。

张量也可以被复制,其中相同的张量数据分布在多个设备上,或者复制和分布的混合。 复制通过将 numDevices 设置为 deviceCounts 乘积的倍数来表示。 此时,devices 张量被假定为具有一个额外的最终模式,张量在该模式上被复制。 复制张量可以在任何地方使用,除了作为收缩的输出。 特别是,传递复制张量可以解锁新的优化,这对于受益于减少通信的问题是有利的。

 For instance, with devices = {0,1,2,3, 4,5,6,7} the tensor would be replicated
. 如下所示

设备 0 & 4

设备 2 & 6

设备 0 & 4

设备 1 & 5

设备 3 & 7

设备 1 & 5

设备 0 & 4

设备 2 & 6

设备 0 & 4

设备 1 & 5

设备 3 & 7

设备 1 & 5

设备 0 & 4

设备 2 & 6

设备 0 & 4

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:
  • handle[in] 库句柄。

  • desc[out] 结果张量描述符。

  • numModes[in] 模式数量。

  • extent[in] 每个模式中张量的范围(大小为 numModes 的数组)。

  • elementStride[in] 每个模式中两个相邻元素之间的偏移量(以线性内存为单位)(大小为 numModes 的数组),对于密集张量可以为 NULL

  • blockSize[in] 每个模式中块的大小(大小为 numModes 的数组),对于未分块的张量(即,每个模式只有一个块,该块等于其范围)可以为 NULL

  • blockStride[in] 每个模式中两个相邻块之间的偏移量(以线性内存为单位)(大小为 numModes 的数组),对于密集块交错布局可以为 NULL

  • deviceCount[in] 每个模式以块循环方式分布的设备数量(大小为 numModes 的数组),对于未分布的张量可以为 NULL

  • numDevices[in] 张量分布的总设备数量(即,deviceCount 中所有元素的乘积乘以它复制到的设备数量)。

  • devices[in] 块分布到的设备,以列优先顺序,即首先步幅为 1(大小为 numDevices 的数组)。

  • type[in] 张量的数据类型。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。

  • CUTENSOR_STATUS_NOT_SUPPORTED – 不支持此布局或数据类型。


cutensorMgDestroyTensorDescriptor()

cutensorStatus_t cutensorMgDestroyTensorDescriptor(cutensorMgTensorDescriptor_t desc)

销毁张量描述符。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:

desc[in] 要销毁的描述符。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。


复制操作

以下类型和函数执行张量之间的复制。


cutensorMgCopyDescriptor_t

typedef struct cutensorMgCopyDescriptor_s *cutensorMgCopyDescriptor_t

描述从一种数据布局到另一种数据布局的张量复制。

它可以描述从主机、单设备和多设备复制的完整笛卡尔积,以及置换和布局更改。


cutensorMgCopyPlan_t

typedef struct cutensorMgCopyPlan_s *cutensorMgCopyPlan_t

描述实现复制操作的特定方法。

它编码了分块和其他实现细节,并且可以重复使用以减少计划开销。


cutensorMgCreateCopyDescriptor()

cutensorStatus_t cutensorMgCreateCopyDescriptor(const cutensorMgHandle_t handle, cutensorMgCopyDescriptor_t *desc, const cutensorMgTensorDescriptor_t descDst, const int32_t modesDst[], const cutensorMgTensorDescriptor_t descSrc, const int32_t modesSrc[])

创建复制描述符。

复制描述符编码复制操作的源和目标。 复制操作支持主机、单设备或多设备上的张量。 它还支持布局更改和模式置换。 唯一的限制是输入和输出张量中相应模式的范围必须匹配。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:
  • handle[in] 库句柄。

  • desc[out] 结果复制描述符。

  • descDst[in] 目标张量描述符。

  • modesDst[in] 目标张量模式。

  • descSrc[in] 源张量描述符。

  • modesSrc[in] 源张量模式。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。

  • CUTENSOR_STATUS_NOT_SUPPORTED – 不支持此张量布局或精度组合。


cutensorMgDestroyCopyDescriptor()

cutensorStatus_t cutensorMgDestroyCopyDescriptor(cutensorMgCopyDescriptor_t desc)

销毁复制描述符并释放其所有先前分配的资源。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:

desc[in] 要销毁的描述符。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。


cutensorMgCopyGetWorkspace()

cutensorStatus_t cutensorMgCopyGetWorkspace(const cutensorMgHandle_t handle, const cutensorMgCopyDescriptor_t desc, int64_t deviceWorkspaceSize[], int64_t *hostWorkspaceSize)

计算复制所需的工作区。

该函数计算复制操作成功完成所需的最小工作区。 它以设备传递到库句柄的相同顺序返回设备工作区大小。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:
  • handle[in] 库句柄。

  • desc[in] 复制描述符。

  • deviceWorkspaceSize[out] 工作区大小(以字节为单位),用于句柄中的每个设备。

  • hostWorkspaceSize[out] 用于锁页主机内存的工作区大小(以字节为单位)。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。


cutensorMgCreateCopyPlan()

cutensorStatus_t cutensorMgCreateCopyPlan(const cutensorMgHandle_t handle, cutensorMgCopyPlan_t *plan, const cutensorMgCopyDescriptor_t desc, const int64_t deviceWorkspaceSize[], int64_t hostWorkspaceSize)

创建复制计划。

复制计划实现了通过复制描述符表达的复制操作。 它包含执行复制操作所需的所有信息。 如果提供的工作区不足,则计划可能会失败。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:
  • handle[in] 库句柄。

  • plan[out] 结果复制计划。

  • desc[in] 复制描述符。

  • deviceWorkspaceSize[in] 将提供的工作区量,用于句柄中的每个设备。

  • hostWorkspaceSize[in] 将提供的锁页主机工作区量。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。


cutensorMgDestroyCopyPlan()

cutensorStatus_t cutensorMgDestroyCopyPlan(cutensorMgCopyPlan_t plan)

销毁复制计划。

调用时,所有未完成的操作必须完成。 释放所有相关资源。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:

plan[in] 要销毁的计划。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。


cutensorMgCopy()

cutensorStatus_t cutensorMgCopy(const cutensorMgHandle_t handle, const cutensorMgCopyPlan_t plan, void *ptrDst[], const void *ptrSrc[], void *deviceWorkspace[], void *hostWorkspace, cudaStream_t streams[])

执行复制操作。

根据给定的计划执行复制操作。它按照相应张量描述符的 devices 参数指定的顺序接收源指针和目标指针,并按照句柄的 devices 参数指定的顺序接收设备工作区和流。如果执行过程中涉及到主机传输,该函数将阻塞,直到这些主机传输完成。只要并发线程使用不同的库句柄,该函数就是线程安全的。

备注

调用异步函数,有条件阻塞,不可重入,且有条件线程安全

返回值

指示操作成功或失败的状态代码

参数:
  • handle[in] 库句柄。

  • plan[输入] 复制计划。

  • ptrDst[输出] 目标张量指针。

  • ptrSrc[输入] 源张量指针。

  • deviceWorkspace[输出] 设备工作区。

  • hostWorkspace[输出] 主机页锁定内存工作区。

  • streams[输入] 执行流。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。

  • CUTENSOR_STATUS_CUDA_ERROR – 与 CUDA 运行时交互时出现问题。


收缩运算

以下类型和函数实现收缩运算。


cutensorMgContractionDescriptor_t

typedef struct cutensorMgContractionDescriptor_s *cutensorMgContractionDescriptor_t

描述了两个张量收缩成第三个张量的运算,可以带有一个可选的源张量。

仅支持设备端张量。


cutensorMgContractionFind_t

typedef struct cutensorMgContractionFind_s *cutensorMgContractionFind_t

描述了实现张量收缩的算法细节。


cutensorMgContractionPlan_t

typedef struct cutensorMgContractionPlan_s *cutensorMgContractionPlan_t

描述了实现收缩运算的一种特定方法。

它编码了分块、置换和其他实现细节,并且可以重复使用以减少计划开销。


cutensorMgAlgo_t

enum cutensorMgAlgo_t

表示在规划收缩运算时选择的算法。

enumerator CUTENSORMG_ALGO_DEFAULT

让内部启发式算法选择。


cutensorMgCreateContractionDescriptor()

cutensorStatus_t cutensorMgCreateContractionDescriptor(const cutensorMgHandle_t handle, cutensorMgContractionDescriptor_t *desc, const cutensorMgTensorDescriptor_t descA, const int32_t modesA[], const cutensorMgTensorDescriptor_t descB, const int32_t modesB[], const cutensorMgTensorDescriptor_t descC, const int32_t modesC[], const cutensorMgTensorDescriptor_t descD, const int32_t modesD[], cutensorComputeType_t compute)

创建一个收缩描述符。

收缩描述符编码了形如以下形式的收缩运算的操作数

\[ D = \alpha \mathcal{A} \mathcal{B} + \beta \mathcal{C} \]
. 目前,收缩运算支持位于一个或多个设备上的张量,但暂不支持存储在主机上的张量。它使用爱因斯坦求和约定,即仅在 modesA 和 modesB 之间共享的模式会被收缩。目前,descC 和 descD 以及 modesC 和 modesD 必须相同。计算类型表示计算过程中可能使用的最低精度。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:
  • handle[in] 库句柄。

  • desc[输出] 生成的张量收缩描述符。

  • descA[输入] 操作数 A 的张量描述符。

  • modesA[输入] 操作数 A 的模式。

  • descB[输入] 操作数 B 的张量描述符。

  • modesB[输入] 操作数 B 的模式。

  • descC[输入] 操作数 C 的张量描述符。

  • modesC[输入] 操作数 C 的模式。

  • descD[输入] 操作数 D 的张量描述符。

  • modesD[输入] 操作数 D 的模式。

  • compute[输入] 运算的计算类型。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。

  • CUTENSOR_STATUS_NOT_SUPPORTED – 不支持此张量布局或精度组合。


cutensorMgDestroyContractionDescriptor()

cutensorStatus_t cutensorMgDestroyContractionDescriptor(cutensorMgContractionDescriptor_t desc)

销毁一个收缩描述符。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:

desc[in] 要销毁的描述符。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。


cutensorMgCreateContractionFind()

cutensorStatus_t cutensorMgCreateContractionFind(const cutensorMgHandle_t handle, cutensorMgContractionFind_t *find, const cutensorMgAlgo_t algo)

创建一个收缩查找。

收缩查找包含执行张量收缩的所有算法选项。目前,它唯一的参数是一个算法,而该算法当前只有一个默认值。未来可能会增加其他选项。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:
  • handle[in] 库句柄。

  • find[输出] 生成的查找。

  • algo[输入] 所需的算法。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。


cutensorMgDestroyContractionFind()

cutensorStatus_t cutensorMgDestroyContractionFind(cutensorMgContractionFind_t find)

销毁一个收缩查找。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:

find[输入] 要销毁的查找。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。


cutensorMgContractionGetWorkspace()

cutensorStatus_t cutensorMgContractionGetWorkspace(const cutensorMgHandle_t handle, const cutensorMgContractionDescriptor_t desc, const cutensorMgContractionFind_t find, cutensorWorksizePreference_t preference, int64_t deviceWorkspaceSize[], int64_t *hostWorkspaceSize)

计算收缩所需的workspace大小。

该函数计算成功执行收缩运算所需的工作区大小。它接受一个工作区偏好设置,可以调整所需的工作区大小。它返回设备工作区大小,顺序与设备传递给库句柄的顺序相同。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:
  • handle[in] 库句柄。

  • desc[输入] 收缩描述符。

  • find[输入] 收缩查找。

  • preference[输入] 工作区偏好设置。

  • deviceWorkspaceSize[输出] 工作区大小,以字节为单位,对应句柄中的每个设备。

  • hostWorkspaceSize[输出] 页锁定主机内存大小,以字节为单位。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。

  • CUTENSOR_STATUS_NOT_SUPPORTED – 不支持此张量布局或精度组合。


cutensorMgCreateContractionPlan()

cutensorStatus_t cutensorMgCreateContractionPlan(const cutensorMgHandle_t handle, cutensorMgContractionPlan_t *plan, const cutensorMgContractionDescriptor_t desc, const cutensorMgContractionFind_t find, const int64_t deviceWorkspaceSize[], int64_t hostWorkspaceSize)

创建一个收缩计划。

收缩计划根据收缩查找中指定的选项,实现通过收缩描述符表达的收缩运算。它包含执行收缩运算所需的所有信息。如果提供的工作区不足,计划可能会失败。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:
  • handle[in] 库句柄。

  • plan[输出] 生成的收缩计划。

  • desc[输入] 收缩描述符。

  • find[输入] 收缩查找。

  • deviceWorkspaceSize[输入] 工作区大小,以字节为单位,对应句柄中的每个设备。

  • hostWorkspaceSize[输入] 页锁定主机内存大小,以字节为单位。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。

  • CUTENSOR_STATUS_NOT_SUPPORTED – 不支持此张量布局或精度组合。


cutensorMgDestroyContractionPlan()

cutensorStatus_t cutensorMgDestroyContractionPlan(cutensorMgContractionPlan_t plan)

销毁一个收缩计划。

调用时,所有未完成的操作必须完成。 释放所有相关资源。

备注

非阻塞,不可重入且线程安全

返回值

指示操作成功或失败的状态代码

参数:

plan[in] 要销毁的计划。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。


cutensorMgContraction()

cutensorStatus_t cutensorMgContraction(const cutensorMgHandle_t handle, const cutensorMgContractionPlan_t plan, const void *alpha, const void *ptrA[], const void *ptrB[], const void *beta, const void *ptrC[], void *ptrD[], void *deviceWorkspace[], void *hostWorkspace, cudaStream_t streams[])

执行收缩运算。

根据提供的计划执行收缩运算。它接收所有操作数作为指针数组,这些指针数组根据其张量描述符的 devices 参数排序。设备工作区和流根据库句柄的 devices 参数排序。只要并发线程使用不同的库句柄,该函数就是线程安全的。

备注

调用异步函数,非阻塞,不可重入,且有条件线程安全

返回值

指示操作成功或失败的状态代码

参数:
  • handle[in] 库句柄。

  • plan[输入] 复制计划。

  • alpha[输入] alpha 缩放因子(主机指针)。

  • ptrA[输入] 操作数 A 的张量指针。

  • ptrB[输入] 操作数 B 的张量指针。

  • beta[输入] beta 缩放因子(主机指针)。

  • ptrC[输入] 操作数 C 的张量指针。

  • ptrD[输出] 操作数 D 的张量指针。

  • deviceWorkspace[输出] 设备工作区。

  • hostWorkspace[输出] 主机页锁定内存工作区。

  • streams[输入] 执行流。

返回值:
  • CUTENSOR_STATUS_SUCCESS – 操作成功完成。

  • CUTENSOR_STATUS_INVALID_VALUE – 某些输入参数无效。

  • CUTENSOR_STATUS_CUDA_ERROR – 与 CUDA 运行时交互时出现问题。