张量网络收缩¶

导论¶

高级 API 背后的目标是为 cuTensorNet 库提供一个对于 Python 程序员来说感觉自然的接口。这些 API 支持来自 NumPy、CuPy 和 PyTorch 的类 ndarray 对象，并支持将张量网络指定为爱因斯坦求和表达式。

高级 API 可以进一步分为两个级别

“粗粒度”级别，用户处理诸如 contract()、contract_path()、einsum() 和 einsum_path() 等 Python 函数。粗粒度级别是一个抽象层，通常用于单个收缩操作。

“细粒度”级别，其中交互是通过对 Network 对象的操作进行的。细粒度级别允许用户投入大量资源来查找最佳收缩路径并自动调优网络，其中对同一网络对象重复收缩允许分摊成本（另请参见状态对象中的资源管理）。

API 还允许 cuTensorNet 库和外部软件包之间的互操作性。例如，用户可以指定从不同软件包（可能是研究项目）获得的收缩顺序。或者，用户可以从 cuTensorNet 获取收缩顺序和切片模式，以供在其他地方下游使用。

使用示例¶

收缩 cuTensorNet C 示例中演示的相同张量网络非常简单，如下所示

from cuquantum import contract
from numpy.random import rand

a = rand(96,64,64,96)
b = rand(96,64,64)
c = rand(64,96,64)

r = contract("mhkn,ukh,xuy->mxny", a, b, c)

如果需要，可以为收缩提供各种选项。对于 PyTorch 张量，从 cuQuantum Python v23.10 开始，contract() 函数的工作方式类似于本机 PyTorch 运算符，可以记录在 autograd 图中并生成反向模式自动微分。有关更多详细信息和示例，请参见 contract()。

从 cuQuantum v22.11 / cuTensorNet v2.0.0 开始，如果用户将 MPI 通信器绑定到库句柄，并在此处概述的其他要求中，cuTensorNet 支持自动 MPI 并行性。为了说明这一点，假设所有进程都持有相同的一组输入张量（在不同的 GPU 上）和相同的网络表达式，这应该可以开箱即用

from cupy.cuda.runtime import getDeviceCount
from mpi4py import MPI
from cuquantum import cutensornet as cutn

# bind comm to cuTensorNet handle
handle = cutn.create()
comm = MPI.COMM_WORLD
cutn.distributed_reset_configuration(
    handle, *cutn.get_mpi_comm_pointer(comm))

# make each process run on different GPU
rank = comm.Get_rank()
device_id = rank % getDeviceCount()
cp.cuda.Device(device_id).use()

# 1. assuming input tensors a, b, and c are created on the right GPU
# 2. passing handle explicitly allows reusing it to reduce the handle creation overhead
r = contract(
    "mhkn,ukh,xuy->mxny", a, b, c,
    options={'device_id' : device_id, 'handle': handle}))

可以在 https://github.com/NVIDIA/cuQuantum/blob/main/python/samples/cutensornet/coarse/example22_mpi_auto.py 找到此类自动 MPI 用法的端到端 Python 示例。

注意

截至 cuQuantum v22.11 / cuTensorNet v2.0.0，Python wheel 不包含所需的 MPI 包装器库。用户需要从源代码（包含在 wheel 中）构建它，或者改用来自 conda-forge 的 Conda 包。

最后，对于寻求完全控制张量网络操作和并行化的用户，我们提供细粒度 API，如 Network 的文档中的示例所示。下面显示了一个完整的示例，说明了使用细粒度 API 并行实现张量网络收缩

from cupy.cuda.runtime import getDeviceCount
from mpi4py import MPI
import numpy as np

from cuquantum import Network

root = 0
comm = MPI.COMM_WORLD

rank, size = comm.Get_rank(), comm.Get_size()

expr = 'ehl,gj,edhg,bif,d,c,k,iklj,cf,a->ba'
shapes = [(8, 2, 5), (5, 7), (8, 8, 2, 5), (8, 6, 3), (8,), (6,), (5,), (6, 5, 5, 7), (6, 3), (3,)]

# Set the operand data on root.
operands = [np.random.rand(*shape) for shape in shapes] if rank == root else None

# Broadcast the operand data.
operands = comm.bcast(operands, root)
    
# Assign the device for each process.
device_id = rank % getDeviceCount()

# Create network object.
network = Network(expr, *operands, options={'device_id' : device_id})

# Compute the path on all ranks with 8 samples for hyperoptimization. Force slicing to enable parallel contraction.
path, info = network.contract_path(optimize={'samples': 8, 'slicing': {'min_slices': max(16, size)}})

# Select the best path from all ranks.
opt_cost, sender = comm.allreduce(sendobj=(info.opt_cost, rank), op=MPI.MINLOC)
if rank == root:
    print(f"Process {sender} has the path with the lowest FLOP count {opt_cost}.")

# Broadcast info from the sender to all other ranks.
info = comm.bcast(info, sender)

# Set path and slices.
path, info = network.contract_path(optimize={'path': info.path, 'slicing': info.slices})

# Calculate this process's share of the slices.
num_slices = info.num_slices
chunk, extra = num_slices // size, num_slices % size
slice_begin = rank * chunk + min(rank, extra)
slice_end = num_slices if rank == size - 1 else (rank + 1) * chunk + min(rank + 1, extra)
slices = range(slice_begin, slice_end)

print(f"Process {rank} is processing slice range: {slices}.")

# Contract the group of slices the process is responsible for.
result = network.contract(slices=slices)

# Sum the partial contribution from each process on root.
result = comm.reduce(sendobj=result, op=MPI.SUM, root=root)

# Check correctness.
if rank == root:
   result_np = np.einsum(expr, *operands, optimize=True)
   print("Does the cuQuantum parallel contraction result match the numpy.einsum result?", np.allclose(result, result_np))

可以在 NVIDIA/cuQuantum 存储库中找到此“手动”MPI Python 示例（此处）。

调用阻塞行为¶

默认情况下，对执行 API 的调用（Network.autotune() 和 Network.contract() 在 Network 对象上以及函数 contract()）会阻塞，并且在操作完成之前不会返回。可以通过设置 NetworkOptions.blocking 并将选项传递给 Network 来更改此行为。当 NetworkOptions.blocking 设置为 'auto' 时，如果输入张量位于设备上，则对执行 API 的调用将在操作在 GPU 上启动后立即返回，而无需等待其完成。如果输入张量位于主机上，则执行 API 调用将始终阻塞，因为收缩的结果也是将驻留在主机上的张量。

在主机上执行的 API（例如 Network.contract_path() 在 Network 对象上，以及 contract_path() 和 einsum_path() 函数）始终阻塞。

流语义¶

流语义取决于执行 API 的行为是选择为阻塞还是非阻塞（请参见调用阻塞行为）。

对于阻塞行为，cuQuantum Python 高级 API 会自动处理流排序，对于在包内执行的操作。可以提供流有两个原因

1. 当准备输入张量的计算在调用执行 API 时尚未完成时。这是用户提供数据的正确性要求。 2. 如果设备有足够的资源并且当前流（默认流）具有伴随操作，则可以实现跨多个流的并行计算。这样做可能是出于性能原因。

对于非阻塞行为，用户有责任确保执行 API 调用之间正确的流排序。

在任何情况下，执行 API 都在提供的流上启动。

状态对象中的资源管理¶

细粒度、状态对象 API（例如 Network）的一个重要方面是资源管理。我们需要确保在对象的整个生命周期内安全地正确管理内部资源，包括库资源、内存资源和用户提供的输入操作数。因此，用户应注意一些注意事项，因为它们会影响内存水印。

状态对象 API 允许用户准备一个对象并将其重用于多个收缩或梯度计算，以分摊准备成本。根据具体问题，投入于缩短执行时间的准备工作可能是理想的解决方案。在准备步骤中，对象不可避免地会保留对设备内存的引用，以供以后重用。但是，此类问题通常意味着高内存使用率，使得同时持有多个对象成为不可能。交错多个大型张量网络的收缩是这种情况的一个示例。

为了解决此用例，从 cuQuantum Python v24.03 开始，添加了两个新功能

现在，每个执行方法都接受 release_workspace 选项。当此选项设置为 True（默认为 False）时，执行操作所需的内存将在方法返回之前释放，从而使此内存可用于其他任务。下次调用相同（或不同）方法时，将按需分配内存。因此，与 release_workspace=True 关联的开销很小，因为分配/释放内存可能需要时间，具体取决于底层内存分配器的实现（请参见下一节）；但是，多个 Network 对象共存成为可能，例如，请参见 example6_resource_mgmt_contraction.py 示例。
reset_operands() 方法现在接受设置 operands=None 以在执行后释放对输入操作数的内部引用。这减少了潜在的内存争用，从而允许以交错方式收缩具有大型输入张量的多个网络。在这种情况下，在对同一 Network 对象进行后续执行之前，应再次使用新操作数调用 reset_operands() 方法，例如，请参见 example8_reset_operand_none.py 示例。

当设备可用内存不足以一次容纳所有问题时，单独或联合使用这两种功能可以准备和使用大量 Network 对象。

外部内存管理¶

从 cuQuantum Python v22.03 开始，我们支持 EMM 类接口，如 Numba 提出和支持的那样，供用户设置其 Python 内存池。用户将选项 NetworkOptions.allocator 设置为符合 cuquantum.BaseCUDAMemoryManager 协议的 Python 对象，并将选项传递给高级 API，例如 contract() 或 Network。然后，临时内存分配将通过此接口完成。（在内部，我们使用相同的接口来使用 CuPy 或 PyTorch 的内存池，具体取决于输入张量操作数。）

注意

cuQuantum 的 BaseCUDAMemoryManager 协议与 Numba 的 EMM 接口（numba.cuda.BaseCUDAMemoryManager）略有不同，但在运行时使用现有 EMM 实例（非类型！）进行鸭子类型化应该是可能的。