Python API#

nvCOMP Python API 参考

这是 NVIDIA® nvCOMP 库的 Python API 参考。

BitstreamKind#

class nvidia.nvcomp.BitstreamKind#

定义 nvcomp 中缓冲区的压缩方式

成员

NVCOMP_NATIVE : 每个输入缓冲区根据管理器设置分块并并行压缩。允许计算校验和。在压缩数据的开头添加带有 nvCOMP 元数据的自定义标头。

RAW : 按原样压缩输入数据，仅使用底层压缩算法。不添加带有 nvCOMP 元数据的标头。

WITH_UNCOMPRESSED_SIZE : 类似于 RAW，但在压缩数据的开头添加带有未压缩大小的自定义标头。

CudaStream#

class nvidia.nvcomp.CudaStream#

CUDA 流的包装器。根据是通过 borrow 还是 make_new 构建，提供共享所有权或视图语义。

CudaStream 是传递给分配函数的流参数类型，可以与 set_*_allocator 一起使用。如果此类内存的释放需要访问传递给分配函数的流，则分配函数应返回一个 ExternalMemory 实例，该实例包装新构造的内存对象和 CudaStream 参数。从那时起，内存对象应仅通过 ExternalMemory 包装器访问。这确保在释放内存时流仍然处于活动状态。

不希望在分配函数之外使用 CudaStream。borrow 和 make_new 仍然为了完整性而提供。

static borrow(
cuda_stream: int,
device_idx: int = -1,
) → nvidia.nvcomp.nvcomp_impl.CudaStream#

创建流视图。

设备索引主要用于特殊的 CUDA 流（即默认流、旧版流和每线程流），这些流的设备无法从流值本身推断出来。默认情况下，它等于 -1，这是一个特殊值，其含义取决于 stream 是否特殊。如果 stream 是特殊的，则默认值将共享流与当前设备关联。否则，CudaStream 将始终与流的实际设备关联。在这种情况下，传递既不是默认值也不是流的实际设备的 device_idx 将引发异常。

参数:

  • cuda_stream – 要包装的 cudaStream_t，表示为 Python 整数。

  • device_idx – 可选的设备索引，用于关联借用的流。有关详细信息，请参阅函数描述。默认为 -1。

property device#

与流关联的设备索引。

property is_special#

底层流是否为特殊流之一（默认流、旧版流或每线程流）。

请注意，将特殊流传递给任何 CUDA API 调用实际上都会传递当前设备对应的特殊流。因此，必须确保在使用流之前选择流的关联设备（由 device 给出）。这目前完全是用户的责任。

static make_new(
device_idx: int = -1,
) → nvidia.nvcomp.nvcomp_impl.CudaStream#

创建具有共享所有权的新流。

参数:

device_idx – 可选的设备索引，用于关联新创建的流。默认情况下等于 -1，这是一个表示当前设备的特殊值。

property ptr#

底层 cudaStream_t，表示为 Python 整数。

属性名称遵循 cupy.Stream 的约定，并反映了 cudaStream_t 在内部是一个指针的事实。

Codec#

class nvidia.nvcomp.Codec#
__init__(
self: nvidia.nvcomp.nvcomp_impl.Codec,
**kwargs,
) → None#

初始化编解码器。

参数:

  • algorithm – 要使用的可选压缩算法名称。默认情况下为空，可以在解码期间推断算法。

  • device_id – 可选的设备 ID，用于在其上执行解码/编码。如果未指定，将使用默认设备。

  • cuda_stream – 可选的 cudaStream_t，表示为 Python 整数。默认情况下，为给定的设备 ID 创建内部 cuda 流。

  • uncomp_chunk_size – 可选的未压缩数据块大小。默认情况下为 65536。

  • checksum_policy –

    定义校验和的计算和验证策略。默认情况下，假定为 NO_COMPUTE_NO_VERIFY。

    LZ4 算法特定选项

    data_type: 用于默认数据类型的可选数组协议类型字符串。

    GDeflate 算法特定选项
    algorithm_type: 要使用的压缩算法类型。允许的值为

    • 0 : 最高吞吐量，仅熵压缩（用于对称压缩/解压缩性能）

    • 1 : 高吞吐量，低压缩率（默认）

    • 2 : 中等吞吐量，中等压缩率，在压缩率上击败 Zlib level 1

    • 3 : 进一步压缩级别支持的占位符，此时将归入 MEDIUM_COMPRESSION

    • 4 : 较低吞吐量，较高压缩率，在压缩率上击败 Zlib level 6

    • 5 : 最低吞吐量，最高压缩率

    Deflate 算法特定选项
    algorithm_type: 要使用的压缩算法类型。允许的值为

    • 0 : 最高吞吐量，仅熵压缩（用于对称压缩/解压缩性能）

    • 1 : 高吞吐量，低压缩率（默认）

    • 2 : 中等吞吐量，中等压缩率，在压缩率上击败 Zlib level 1

    • 3 : 进一步压缩级别支持的占位符，此时将归入 MEDIUM_COMPRESSION

    • 4 : 较低吞吐量，较高压缩率，在压缩率上击败 Zlib level 6

    • 5 : 最低吞吐量，最高压缩率

    Bitcomp 算法特定选项
    algorithm_type: 使用的 Bitcomp 算法类型。

    • 0 : 默认算法，通常给出最佳压缩率

    • 1 : “稀疏”算法，在稀疏数据（包含大量零）上效果良好，并且通常比默认算法更快。

    data_type: 用于默认数据类型的可选数组协议类型字符串。

    ANS 算法特定选项
    data_type: 用于默认数据类型的可选数组协议类型字符串。允许的值为

    • |u1 : 对于无符号 8 位整数

    • <f2 : 对于 16 位小端浮点数。要求 uncomp_chunk_size 是 2 的倍数

    Cascaded 算法特定选项

    data_type: 用于默认数据类型的可选数组协议类型字符串。

    num_rles: 要执行的行程长度编码的数量。默认情况下等于 2

    num_deltas: 要执行的 Delta 编码的数量。默认情况下等于 1

    use_bitpack: 是否对最终层进行位压缩。默认情况下为 True。

decode(*args, **kwargs)#

重载函数。

  1. decode(self: nvidia.nvcomp.nvcomp_impl.Codec, src: nvidia.nvcomp.nvcomp_impl.Array, data_type: str = ‘’) -> object

    执行从 Array 句柄解码数据。

    Args

    src: 解码源对象。

    data_type: 用于输出数据类型的可选数组协议类型字符串。默认情况下等于 |u1。

    Returns

    nvcomp.Array

  2. decode(self: nvidia.nvcomp.nvcomp_impl.Codec, srcs: list[nvidia.nvcomp.nvcomp_impl.Array], data_type: str = ‘’) -> list[object]

    执行从一批 Array 句柄解码。

    Args

    srcs: Array 对象列表

    data_type: 用于输出数据类型的可选数组协议类型字符串。

    Returns

    解码后的 nvcomp.Array 列表

encode(*args, **kwargs)#

重载函数。

  1. encode(self: nvidia.nvcomp.nvcomp_impl.Codec, array_s: nvidia.nvcomp.nvcomp_impl.Array) -> object

    编码数组。

    Args

    array: 要编码的数组

    Returns

    编码后的 nvcomp.Array

  2. encode(self: nvidia.nvcomp.nvcomp_impl.Codec, srcs: list[nvidia.nvcomp.nvcomp_impl.Array]) -> list[object]

    执行从一批 Array 句柄编码。

    Args

    srcs: Array 对象列表

    Returns

    编码后的 nvcomp.Array 列表

ArrayBufferKind#

class nvidia.nvcomp.ArrayBufferKind#

定义存储数组数据的缓冲区类型。

成员

STRIDED_DEVICE : GPU 可访问，采用 pitch-linear 布局。

STRIDED_HOST : 主机可访问，采用 pitch-linear 布局。

Array#

class nvidia.nvcomp.Array#

包装数组的类。它可以是解码的数据或要编码的数据。

property __cuda_array_interface__#

与 Numba v0.39.0 或更高版本兼容的 CUDA 数组交换接口（有关详细信息，请参阅 CUDA Array Interface）

__dlpack__(
self: nvidia.nvcomp.nvcomp_impl.Array,
stream: object = None,
) → capsule#

将数组导出为 DLPack 张量

__dlpack_device__(
self: nvidia.nvcomp.nvcomp_impl.Array,
) → tuple#

获取与缓冲区关联的设备

property buffer_kind#

存储数组数据的缓冲区类型。

property buffer_size#

存储数组的总字节数。

cpu(self: nvidia.nvcomp.nvcomp_impl.Array) → object#

返回此数组在 CPU 内存中的副本。如果此数组已在 CPU 内存中，则不执行复制，并返回原始对象。

Returns:

内容位于 CPU 内存中的 Array 对象，如果无法完成复制，则返回 None。

cuda(
self: nvidia.nvcomp.nvcomp_impl.Array,
synchronize: bool = True,
cuda_stream: int = 0,
) → object#

返回此数组在设备内存中的副本。如果此数组已在设备内存中，则不执行复制，并返回原始对象。

参数:

  • synchronize – 如果为 True（默认），则它会阻塞并等待从主机到设备的复制完成，否则不执行同步，并且需要使用例如 __cuda_array_interface__ 提供的 cuda 流进行进一步同步。

  • cuda_stream – 可选的 cudaStream_t，表示为 Python 整数，用于将主机缓冲区复制到其中。

Returns:

内容位于设备内存中的 Array 对象，如果无法完成复制，则返回 None。

property dtype#
property item_size#

每个元素的大小（以字节为单位）。

property ndim#
property precision#

数据类型中有效位的最大数量。值 0 表示精度等于数据类型位深度

property shape#
property size#

此数组容纳的元素数量。

property strides#

轴的步幅（以字节为单位）

to_dlpack(
self: nvidia.nvcomp.nvcomp_impl.Array,
cuda_stream: object = None,
) → capsule#

使用零拷贝转换将数组导出为 DLPack 张量。

参数:

cuda_stream – 可选的 cudaStream_t，表示为 Python 整数，必须在其上进行同步才能在创建的 Array 中进行。

Returns:

DLPack 张量，封装在 PyCapsule 对象中。

as_array#

nvidia.nvcomp.as_array(
source: object,
cuda_stream: int = 0,
) → nvidia.nvcomp.nvcomp_impl.Array#

将外部缓冲区包装为数组，并将缓冲区生命周期与数组绑定

参数:

  • source – 输入 DLPack 张量，封装在 PyCapsule 对象中，或其他具有 __cuda_array_interface__、__array_interface__ 或 __dlpack__ 和 __dlpack_device__ 方法的对象。

  • cuda_stream – 可选的 cudaStream_t，表示为 Python 整数，必须在其上进行同步才能在创建的 Array 中进行。

Returns:

nvcomp.Array

as_arrays#

nvidia.nvcomp.as_arrays(sources: list[object], cuda_stream: int = 0) → list[object]#

将所有外部缓冲区包装为数组，并将缓冲区生命周期与数组绑定

参数:

  • sources – 输入 DLPack 张量列表，封装在 PyCapsule 对象中，或其他具有 __cuda_array_interface__、__array_interface__ 或 __dlpack__ 和 __dlpack_device__ 方法的对象。

  • cuda_stream – 可选的 cudaStream_t，表示为 Python 整数，必须在其上进行同步才能在创建的 Array 中进行。

Returns:

nvcomp.Array 列表

from_dlpack#

nvidia.nvcomp.from_dlpack(
source: object,
cuda_stream: int = 0,
) → nvidia.nvcomp.nvcomp_impl.Array#

从 DLPack 张量到数组的零拷贝转换。

参数:

  • source – 输入 DLPack 张量，封装在 PyCapsule 对象中，或其他（数组）对象，具有 __dlpack__ 和 __dlpack_device__ 方法。

  • cuda_stream – 可选的 cudaStream_t，表示为 Python 整数，必须在其上进行同步才能在创建的 Array 中进行。

Returns:

nvcomp.Array

set_device_allocator#

nvidia.nvcomp.set_device_allocator(allocator: object = None) → None#

设置一个新的分配器，用于未来的设备分配。

分配器的签名应如以下示例所示

def my_allocator(nbytes: int, stream: nvcomp.Stream) -> PtrProtocol:
    return MyBuffer(nbytes, stream)

PtrProtocol 表示任何具有整数类型的 ptr 属性的对象。这应该是指向已分配缓冲区的指针（表示为整数）。

在签名中，nbytes 是请求缓冲区中的字节数。stream 是 CUDA 流，用于执行分配和/或释放（如果分配器是流排序的）。非流排序的分配器可以忽略 stream，或者可以在释放之前与其同步，具体取决于所需的行为。目前不支持单独的分配和释放流。

返回的对象应使得，当它被删除时，要么是因为没有更多有效的 Python 引用指向它，要么是因为它被垃圾回收，内存将被释放。在自定义 Python 类中，可以通过 __del__ 方法来实现。这被认为是高级用法模式，因此推荐的方法是从其他库（例如 cupy 的 Memory 类和 rmm 的 DeviceBuffer）组合预先存在的解决方案。

通常允许在先前分配器分配的一个或多个缓冲区仍处于活动状态时设置新的分配器。但是，各个分配器实现可以选择禁止这样做。

如果释放需要访问 stream，则分配器应返回一个 ExternalMemory 实例，该实例包装新构造的内存对象和 CudaStream 参数。从那时起，内存对象应仅通过 ExternalMemory 包装器访问。这确保在释放内存时流仍然处于活动状态。

自定义分配器的简单但通用的示例由 rmm_nvcomp_allocator 给出。

分配的内存必须是设备可访问的。

参数:

allocator – 满足上述条件的可调用对象。

set_pinned_allocator#

nvidia.nvcomp.set_pinned_allocator(allocator: object = None) → None#

设置一个新的分配器，用于未来的固定主机分配。

请注意，他应该分配固定的主机内存。对于非固定主机内存，请使用 set_host_allocator。使用此分配器分配非固定主机内存不是错误，但可能导致性能下降。

分配器必须分配主机可访问的内存。除此之外，关于 allocator 的条件与 set_device_allocator 中的条件相同，包括流语义。

参数:

allocator – 满足上述条件的可调用对象。

set_host_allocator#

nvidia.nvcomp.set_host_allocator(allocator: object = None) → None#

设置一个新的分配器，用于未来的非固定主机分配。

这主要用于可能很大的分配，例如支持 CPU Array 实例的分配。中等大小的内部主机分配可能仍使用系统分配内存。

这应该分配非固定主机内存。对于固定主机内存，请使用 set_pinned_allocator。使用此分配器分配固定主机内存不是错误，但可能导致性能下降。

分配器必须分配主机可访问的内存。除此之外，关于 allocator 的条件与 set_device_allocator 中的条件相同，包括流语义。

参数:

allocator – 满足上述条件的可调用对象。