底层 Python 绑定¶

从 cuQuantum-python v24.11 版本开始，所有 cuStateVec 和 cuTensorNet C API 都暴露在 cuquantum.custatevec 和 cuquantum.cutensornet 模块下。cuDensityMat 的 C API 目前暴露在 cuquantum.bindings.cudensitymat 模块下。

命名与调用约定¶

当 C API 暴露时，我们遵循 PEP 8 风格指南并采用以下更改

所有库名称前缀都被剥离
函数名称按单词分隔并遵循驼峰命名法
枚举名称中每个单词的首字母都大写
每个枚举的名称前缀都从其值的名称中剥离
可在所有子模块中使用的通用枚举放在父模块 cuquantum 中
在适用情况下，输出从函数参数中剥离，并直接作为 Python 对象返回
指针作为 Python int 传递

以下是非详尽的 C 到 Python 映射示例列表

函数：custatevecGetDefaultWorkspaceSize() -> custatevec.get_default_workspace_size()。
函数：cutensornetCreateNetworkDescriptor() -> cutensornet.create_network_descriptor()。
枚举类型：custatevecMatrixLayout_t -> custatevec.MatrixLayout。
枚举类型：cutensornetContractionOptimizerConfigAttributes_t -> cutensornet.ContractionOptimizerConfigAttribute。
枚举值名称：CUSTATEVEC_MATRIX_LAYOUT_COL -> custatevec.MatrixLayout.COL。
枚举值名称：CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_HYPER_NUM_SAMPLES -> cutensornet.ContractionOptimizerConfigAttribute.HYPER_NUM_SAMPLES。
返回：custatevecSamplerCreate() 的输出是采样器描述符和所需的工作区大小，它们在相应的 custatevec.sampler_create() Python API 中包装为 2 元组。
全局枚举：custatevecComputeType_t 和 cutensornetComputeType_t -> cuquantum.ComputeType。

上述规则可能存在例外，但它们将是不言自明的并妥善记录在案。在下一节中，我们将讨论 Python 中的指针传递。

内存管理¶

指针和数据生命周期¶

与 C/C++ 不同，Python 不提供分配/释放主机内存的底层原语，更不用说设备内存了。为了使 C API 与 Python 一起工作，重要的是通过 Python 代理对象正确完成内存管理。

注意

也可以使用 array.array（以及根据需要使用 memoryview）来管理主机内存，但是与使用 numpy.ndarray 相比，它更加繁琐，尤其是在数组操作和计算方面。

注意

也可以使用 CUDA Python 来管理设备内存，但截至 CUDA 11，没有简单、Pythonic 的方法来修改存储在 GPU 上的内容，这需要自定义内核。CuPy 是一个轻量级、NumPy 兼容的数组库，旨在解决此需求。

要将数据从 Python 传递到 C，需要使用各种对象的指针地址（作为 Python int）。使用 NumPy/CuPy 数组作为代理，操作非常简单，如下所示

# create a host buffer to hold 5 int
buf = numpy.empty((5,), dtype=numpy.int32)
# pass buf's pointer to the wrapper
# buf could get modified in-place if the function writes to it
my_func(..., buf.ctypes.data, ...)
# examine/use buf's data
print(buf)

# create a device buffer to hold 10 double
buf = cupy.empty((10,), dtype=cupy.float64)
# pass buf's pointer to the wrapper
# buf could get modified in-place if the function writes to it
my_func(..., buf.data.ptr, ...)
# examine/use buf's data
print(buf)

# create an untyped device buffer of 128 bytes
buf = cupy.cuda.alloc(128)
# pass buf's pointer to the wrapper
# buf could get modified in-place if the function writes to it
my_func(..., buf.ptr, ...)
# buf is automatically destroyed when going out of scope

请注意，基本假设是数组在内存中必须是连续的（除非 C 接口允许指定数组步幅）。

因此，例如，从 cuQuantum Python v0.1.0 版本开始，所有 C 结构体（包括句柄和描述符）都未作为 Python 类公开；也就是说，它们没有自己的类型，只是简单地强制转换为普通的 Python int 以便传递。任何下游消费者都应创建包装类来保存指针地址（如果需要）。换句话说，用户完全控制（并负责）管理指针生命周期。

然而，在某些情况下，我们能够为用户转换 Python 对象（如果需要只读主机数组），以减轻用户的负担。例如，在需要序列或嵌套序列的函数中，以下操作是等效的：

# passing a host buffer of int type can be done like this
buf = numpy.array([0, 1, 3, 5, 6], dtype=numpy.int32)
my_func(..., buf.ctypes.data, ...)

# or just this
buf = [0, 1, 3, 5, 6]
my_func(..., buf, ...)  # the underlying data type is determined by the C API

当用户需要将大量张量元数据传递给 C 时，这尤其有用（例如： cutensornet.create_network_descriptor()）。

用户提供的内存池¶

从 cuQuantum v22.03 开始，我们提供了一个接口，供用户为其 cuStateVec/cuTensorNet 库引入内存池以供使用。设置后，用户不再需要在调用 API 之前管理任何临时工作区；库将从用户的池中提取内存（并在完成后返回）。内存池的唯一要求是它必须是流有序的。有关介绍，请参见内存管理 API。目前我们只支持设备内存池。

在 cuQuantum Python 中，此接口通过底层 API custatevec.set_device_mem_handler() 和 custatevec.get_device_mem_handler() 公开（同样适用于 cutensornet）。目前我们提供三种不同的方式来设置 handler 参数：

如果给定 int，则假定它是指向完全初始化的 custatevecDeviceMemHandler_t 结构体的指针地址。

如果长度为 4 的 Python 序列，则假定为 (ctx, device_alloc, device_free, name)

如果长度为 3 的 Python 序列，则假定为 (malloc, free, name)

有关更多详细信息，请参见 API 参考。设置后，使用调用约定：

将工作区（或工作区描述符）指针地址设置为 0

将工作区大小设置为 0

无论 API 在哪里需要工作区，都会通知库它应该使用用户内存池。此示例演示了此 API 的用法。

用法示例¶

下面的代码是用 C 编写的相应 cuStateVec 示例的 Python 翻译。

import numpy as np
import cupy as cp
from cuquantum import custatevec as cusv
from cuquantum import cudaDataType as cudtype
from cuquantum import ComputeType as ctype


nIndexBits = 3
nSvSize = (1 << nIndexBits)
nTargets = 1
nControls = 2
adjoint = 0

targets = (2,)
controls = (0, 1)

d_sv = cp.asarray([[0.0, 0.0], [0.0, 0.1], [0.1, 0.1], [0.1, 0.2],
                   [0.2, 0.2], [0.3, 0.3], [0.3, 0.4], [0.4, 0.5]], dtype=np.float64)
d_sv = d_sv.view(np.complex128).reshape(-1)

d_sv_result = cp.asarray([[0.0, 0.0], [0.0, 0.1], [0.1, 0.1], [0.4, 0.5],
                          [0.2, 0.2], [0.3, 0.3], [0.3, 0.4], [0.1, 0.2]], dtype=np.float64)
d_sv_result = d_sv_result.view(np.complex128).reshape(-1)

d_matrix = cp.asarray([[0.0, 0.0], [1.0, 0.0], [1.0, 0.0], [0.0, 0.0]], dtype=np.float64)
d_matrix = d_matrix.view(np.complex128).reshape(-1)

# cuStateVec handle initialization
handle = cusv.create()

# check the size of external workspace
extraWorkspaceSizeInBytes = cusv.apply_matrix_get_workspace_size(
    handle, cudtype.CUDA_C_64F, nIndexBits, d_matrix.data.ptr, cudtype.CUDA_C_64F,
    cusv.MatrixLayout.ROW, adjoint, nTargets, nControls, ctype.COMPUTE_64F)

# allocate external workspace if necessary
if extraWorkspaceSizeInBytes > 0:
    workspace = cp.cuda.alloc(extraWorkspaceSizeInBytes)
    workspace_ptr = workspace.ptr
else:
    workspace_ptr = 0

# apply gate
cusv.apply_matrix(
    handle, d_sv.data.ptr, cudtype.CUDA_C_64F, nIndexBits,
    d_matrix.data.ptr, cudtype.CUDA_C_64F, cusv.MatrixLayout.ROW, adjoint,
    targets, len(targets), controls, 0, len(controls), ctype.COMPUTE_64F,
    workspace_ptr, extraWorkspaceSizeInBytes)

# destroy handle
cusv.destroy(handle)

# --------------------------------------------------------------------------

# check if d_sv holds the updated statevector
correct = cp.allclose(d_sv, d_sv_result)
if not correct:
    raise RuntimeError("example FAILED: wrong result")

# if this is a standalone script, everything is cleaned up properly at exit