libwholegraph API 文档#

Doxygen WholeGraph C API 文档#

有关 doxygen 文档，请参考 Doxygen 文档

WholeGraph C API 文档#

库级别 API#

enum wholememory_error_code_t#

WholeMemory 错误码定义。

定义 WholeMemory 库的错误码。

值

enumerator WHOLEMEMORY_SUCCESS#

成功

enumerator WHOLEMEMORY_UNKNOW_ERROR#

未知错误

enumerator WHOLEMEMORY_NOT_IMPLEMENTED#

方法未实现

enumerator WHOLEMEMORY_LOGIC_ERROR#

逻辑错误

enumerator WHOLEMEMORY_CUDA_ERROR#

CUDA 错误

enumerator WHOLEMEMORY_COMMUNICATION_ERROR#

通信错误

enumerator WHOLEMEMORY_INVALID_INPUT#

输入无效，例如 nullptr

enumerator WHOLEMEMORY_INVALID_VALUE#

输入值无效

enumerator WHOLEMEMORY_OUT_OF_MEMORY#

内存不足

enumerator WHOLEMEMORY_NOT_SUPPORTED#

不支持

enumerator WHOLEMEMORY_SYSTEM_ERROR#

系统错误>

wholememory_error_code_t wholememory_init(unsigned int flags, LogLevel log_level = LEVEL_INFO)#

初始化 WholeMemory 库

参数：

  • flags – : 保留，应为 0

  • log_level – : wholememory 日志级别，默认级别为“info”

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_finalize()#

终结 WholeMemory 库

返回值：

: wholememory_error_code_t

int fork_get_device_count()#

Fork 一个新进程并获取设备数量。应在其他 CUDA 调用之前调用

返回值：

: CUDA 设备数量，错误时返回 -1

WholeMemory 通信器 API#

typedef struct wholememory_comm_ *wholememory_comm_t#

通信器的不透明句柄。

通信器的不透明句柄

struct wholememory_unique_id_t#

WholeMemory 通信器的唯一 ID。

WholeMemory 通信器的不透明句柄，暴露为 char 数组。底层实现可能是 ncclUniqueId_t

wholememory_error_code_t wholememory_create_unique_id(wholememory_unique_id_t *unique_id)#

为 WholeMemory 通信器创建 UniqueID

参数：

unique_id – : 返回的 UniqueID

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_create_communicator(wholememory_comm_t *comm, wholememory_unique_id_t unique_id, int rank, int size)#

创建 WholeMemory 通信器

参数：

  • comm – : 返回的 WholeMemory 通信器

  • unique_id – : UniqueID

  • rank – : 此进程的 rank。

  • size – : 此通信器中的进程数量

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_destroy_communicator(wholememory_comm_t comm)#

销毁 WholeMemory 通信器

参数：

comm – : 要销毁的 WholeMemory 通信器

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_communicator_get_rank(int *rank, wholememory_comm_t comm)#

获取当前进程在 WholeMemory 通信器中的 rank

参数：

  • rank – : 返回的 rank

  • comm – : WholeMemory 通信器

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_communicator_get_size(int *size, wholememory_comm_t comm)#

获取 WholeMemory 通信器的大小

参数：

  • size – : 返回的大小

  • comm – : WholeMemory 通信器

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_communicator_barrier(wholememory_comm_t comm)#

WholeMemory 通信器上的屏障同步

参数：

comm – : WholeMemory 通信器

返回值：

: wholememory_error_code_t

WholeMemoryHandle API#

enum wholememory_memory_type_t#

WholeMemory 的内存类型。

内存类型是 WholeMemory 的内存地址映射类型

值

enumerator WHOLEMEMORY_MT_NONE#

未定义。

enumerator WHOLEMEMORY_MT_CONTINUOUS#

所有 ranks 的内存都映射在连续地址空间中

enumerator WHOLEMEMORY_MT_CHUNKED#

所有 ranks 的内存都映射在分块地址空间中

enumerator WHOLEMEMORY_MT_DISTRIBUTED#

其他 ranks 的内存未映射。

enumerator WHOLEMEMORY_MT_HIERARCHY#

其他 ranks 的内存映射在分层地址空间中

enum wholememory_memory_location_t#

WholeMemory 的内存位置。

WholeMemory 的内存位置可以是主机或设备。

值

enumerator WHOLEMEMORY_ML_NONE#

未定义

enumerator WHOLEMEMORY_ML_DEVICE#

设备内存

enumerator WHOLEMEMORY_ML_HOST#

主机内存

typedef struct wholememory_handle_ *wholememory_handle_t#

WholeMemory 的不透明句柄。

WholeMemory 的不透明句柄

struct wholememory_gref_t#

WholeMemory 对象的全局引用。

全局引用适用于连续或分块的 WholeMemory 类型，在这些类型中，每个 rank 都可以直接访问所有 ranks 的所有内存。全局引用用于执行这种直接访问。

wholememory_error_code_t wholememory_malloc(wholememory_handle_t *wholememory_handle_ptr, size_t total_size, wholememory_comm_t comm, wholememory_memory_type_t memory_type, wholememory_memory_location_t memory_location, size_t data_granularity, size_t *rank_entry_partition = nullptr)#

Malloc WholeMemory

参数：

  • wholememory_handle_ptr – : 返回的 WholeMemory 句柄

  • total_size – : 总分配大小（字节）。

  • comm – : WholeMemory 通信器

  • memory_type – : WholeMemory 类型

  • memory_location – : 内存位置，主机或设备

  • data_granularity – : 数据粒度大小，保证不会被分割。

  • rank_entry_partition – : 每个 rank 的条目计数（条目大小等于 data_granularity）

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_free(wholememory_handle_t wholememory_handle)#

释放已分配的 WholeMemory 句柄

参数：

wholememory_handle – : 要释放的 WholeMemory 句柄

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_get_communicator(wholememory_comm_t *comm, wholememory_handle_t wholememory_handle)#

从 WholeMemory 句柄获取底层的 WholeMemory 通信器

参数：

  • comm – : 返回的 WholeMemory 通信器

  • wholememory_handle – : WholeMemory 句柄

返回值：

: wholememory_error_code_t

wholememory_memory_type_t wholememory_get_memory_type(wholememory_handle_t wholememory_handle)#

获取 WholeMemory 类型

参数：

wholememory_handle – : WholeMemory 句柄

返回值：

: WholeMemory 类型

wholememory_memory_location_t wholememory_get_memory_location(wholememory_handle_t wholememory_handle)#

获取 WholeMemory 位置

参数：

wholememory_handle – : WholeMemory 句柄

返回值：

: WholeMemory 位置

size_t wholememory_get_total_size(wholememory_handle_t wholememory_handle)#

获取 WholeMemory 的总大小

参数：

wholememory_handle – : WholeMemory 句柄

返回值：

: 总大小

size_t wholememory_get_data_granularity(wholememory_handle_t wholememory_handle)#

获取 WholeMemory 句柄的数据粒度

参数：

wholememory_handle – : WholeMemory 句柄

返回值：

: 数据粒度大小

wholememory_error_code_t wholememory_get_local_memory(void **local_ptr, size_t *local_size, size_t *local_offset, wholememory_handle_t wholememory_handle)#

从当前 rank 的 WholeMemory 句柄获取本地内存，本地内存可以直接访问该内存。但本地内存不一定位于本地 GPU 上。

参数：

  • local_ptr – : 返回的本地内存指针

  • local_size – : 返回的本地内存大小

  • local_offset – : 返回的本地内存相对于 WholeMemory 的偏移量

  • wholememory_handle – : WholeMemory 句柄

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_get_rank_memory(void **rank_memory_ptr, size_t *rank_memory_size, size_t *rank_memory_offset, int rank, wholememory_handle_t wholememory_handle)#

从 WholeMemory 句柄获取指定 rank 的本地内存

参数：

  • rank_memory_ptr – : 返回的指定 rank 的本地内存指针

  • rank_memory_size – : 返回的指定 rank 的本地内存大小

  • rank_memory_offset – : 返回的指定 rank 的本地内存相对于 WholeMemory 的偏移量

  • rank – : 指定的 rank

  • wholememory_handle – : WholeMemory 句柄

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_get_global_pointer(void **global_ptr, wholememory_handle_t wholememory_handle)#

从 WholeMemory 句柄获取全局内存指针。只有连续内存类型或分块主机内存具有全局指针。

参数：

  • global_ptr – : 返回的 WholeMemory 指针

  • wholememory_handle – : WholeMemory 句柄

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_get_global_reference(wholememory_gref_t *wholememory_gref, wholememory_handle_t wholememory_handle)#

从 WholeMemory 句柄获取全局引用。WholeMemory 全局引用是连续和分块内存类型的通用数据结构。

参数：

  • wholememory_gref – : 返回的 WholeMemory 全局引用

  • wholememory_handle – : WholeMemory 句柄

返回值：

: wholememory_error_code_t

警告

doxygenfunction: Cannot find function “wholememory_determine_partition_plan” in doxygen xml output for project “libwholegraph” from directory: /tmp/tmp.yOXprHRhV5

警告

doxygenfunction: Cannot find function “wholememory_determine_entry_partition_plan” in doxygen xml output for project “libwholegraph” from directory: /tmp/tmp.yOXprHRhV5

警告

doxygenfunction: Cannot find function “wholememory_get_partition_plan” in doxygen xml output for project “libwholegraph” from directory: /tmp/tmp.yOXprHRhV5

wholememory_error_code_t wholememory_load_from_file(wholememory_handle_t wholememory_handle, size_t memory_offset, size_t memory_entry_size, size_t file_entry_size, const char **file_names, int file_count, int round_robin_size)#

从二进制文件加载 WholeMemory，所有 rank 应一起调用

参数：

  • wholememory_handle – : WholeMemory 句柄

  • memory_offset – : 加载到内存偏移量

  • memory_entry_size – : WholeMemory 的条目大小

  • file_entry_size – : 文件中的条目大小，应小于或等于 memory_entry_size

  • file_names – : 文件名，所有二进制文件将被逻辑连接并加载。

  • file_count – : 文件数量。

  • round_robin_size – : 在 round-robin 分片模式下，一个 rank 的连续嵌入数量

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_store_to_file(wholememory_handle_t wholememory_handle, size_t memory_offset, size_t memory_entry_stride, size_t file_entry_size, const char *local_file_name)#

将本地 WholeMemory 存储到文件，所有 rank 都应调用此函数，使用不同的 local_file_name。

参数：

  • wholememory_handle – : WholeMemory 句柄

  • memory_offset – : 要存储的内存偏移量

  • memory_entry_stride – : WholeMemory 的条目步长

  • file_entry_size – : 文件中的条目大小，应小于或等于 memory_entry_size

  • local_file_name – : 要存储到的本地文件

返回值：

: wholememory_error_code_t

WholeMemoryTensor API#

enum wholememory_dtype_t#

定义 WholeMemory 张量的数据类型

值

enumerator WHOLEMEMORY_DT_UNKNOWN#

未知类型

enumerator WHOLEMEMORY_DT_FLOAT#

32 位浮点类型

enumerator WHOLEMEMORY_DT_HALF#

16 位半精度浮点类型

enumerator WHOLEMEMORY_DT_DOUBLE#

64 位双精度浮点类型

enumerator WHOLEMEMORY_DT_BF16#

16 位 bfloat 类型

enumerator WHOLEMEMORY_DT_INT#

32 位有符号整数类型

enumerator WHOLEMEMORY_DT_INT64#

64 位有符号整数类型

enumerator WHOLEMEMORY_DT_INT16#

16 位有符号整数类型

enumerator WHOLEMEMORY_DT_INT8#

8 位有符号整数类型

enumerator WHOLEMEMORY_DT_COUNT#

类型总数

struct wholememory_array_description_t#

WholeMemory 中数组的包装器

struct wholememory_matrix_description_t#

WholeMemory 中矩阵的包装器

struct wholememory_tensor_description_t#

WholeMemory 中的张量描述，维度 0 是变化最慢的维度。

typedef struct wholememory_tensor_ *wholememory_tensor_t#

WholeMemoryTensor 的不透明句柄。

WholeMemoryTensor 的不透明句柄

size_t wholememory_dtype_get_element_size(wholememory_dtype_t dtype)#

获取 wholememory_dtype_t 的元素大小

参数：

dtype – : wholememory_dtype_t

返回值：

: dtype 的元素大小，无效 dtype 时为 -1。

bool wholememory_dtype_is_floating_number(wholememory_dtype_t dtype)#

检查 dtype 是否为浮点数

参数：

dtype – : wholememory_dtype_t

返回值：

: 如果 dtype 是 WHOLEMEMORY_DT_FLOAT, WHOLEMEMORY_DT_HALF, WHOLEMEMORY_DT_DOUBLE 或 WHOLEMEMORY_DT_BF16 则为 True，否则为 False。

bool wholememory_dtype_is_integer_number(wholememory_dtype_t dtype)#

检查 dtype 是否为整数

参数：

dtype – : wholememory_dtype_t

返回值：

: 如果 dtype 是 WHOLEMEMORY_DT_INT, WHOLEMEMORY_DT_INT64, WHOLEMEMORY_DT_INT16 或 WHOLEMEMORY_DT_INT8 则为 True，否则为 False。

wholememory_array_description_t wholememory_create_array_desc(int64_t size, int64_t storage_offset, wholememory_dtype_t dtype)#

创建 wholememory_array_description_t 对象

参数：

  • size – : 数组大小（元素数量）

  • storage_offset – : 存储偏移量（元素数量）

  • dtype – : 数组元素的 数据类型

返回值：

创建的 wholememory_array_description_t

wholememory_matrix_description_t wholememory_create_matrix_desc(int64_t sizes[2], int64_t stride, int64_t storage_offset, wholememory_dtype_t dtype)#

创建 wholememory_matrix_description_t 对象

参数：

  • sizes – : 矩阵大小数组，按元素数量计数，sizes[1] 变化最快。

  • stride – : 第一维度（变化较慢的维度）的步长，步长按元素数量计数

  • storage_offset – : 存储偏移量（元素数量）

  • dtype – : 矩阵元素的数据类型

返回值：

创建的 wholememory_matrix_description_t

void wholememory_initialize_tensor_desc(wholememory_tensor_description_t *p_tensor_description)#

初始化 wholememory_tensor_description_t，将 sizes 和 strides 设置为全 1，将 storage_offset 设置为 0，将 dtype 设置为 WHOLEMEMORY_DT_UNKNOWN，将 dim 设置为 0。

参数：

p_tensor_description – : 指向 wholememory_tensor_description_t 的指针。

void wholememory_copy_array_desc_to_matrix(wholememory_matrix_description_t *p_matrix_description, wholememory_array_description_t *p_array_description)#

将数组描述复制到张量描述

参数：
void wholememory_copy_array_desc_to_tensor(wholememory_tensor_description_t *p_tensor_description, wholememory_array_description_t *p_array_description)#

将数组描述复制到张量描述

参数：
void wholememory_copy_matrix_desc_to_tensor(wholememory_tensor_description_t *p_tensor_description, wholememory_matrix_description_t *p_matrix_description)#

将矩阵描述复制到张量描述

参数：
bool wholememory_convert_tensor_desc_to_array(wholememory_array_description_t *p_array_description, wholememory_tensor_description_t *p_tensor_description)#

将张量描述转换为数组描述

参数：
返回值：

: 如果可转换则返回 true，否则返回 false。

bool wholememory_convert_tensor_desc_to_matrix(wholememory_matrix_description_t *p_matrix_description, wholememory_tensor_description_t *p_tensor_description)#

将张量描述转换为矩阵描述

参数：
返回值：

: 如果可转换则返回 true，否则返回 false。

int64_t wholememory_get_memory_element_count_from_array(wholememory_array_description_t *p_array_description)#

从数组描述获取总元素数量。

参数：

p_array_description – : 指向 wholememory_array_description_t 的指针。

返回值：

: 返回元素数量。

int64_t wholememory_get_memory_size_from_array(wholememory_array_description_t *p_array_description)#

从数组描述获取总内存大小。

参数：

p_array_description – : 指向 wholememory_array_description_t 的指针。

返回值：

: 返回内存大小。

int64_t wholememory_get_memory_element_count_from_matrix(wholememory_matrix_description_t *p_matrix_description)#

从矩阵描述获取总元素数量。

参数：

p_matrix_description – : 指向 wholememory_matrix_description_t 的指针。

返回值：

: 返回元素数量。

int64_t wholememory_get_memory_size_from_matrix(wholememory_matrix_description_t *p_matrix_description)#

从矩阵描述获取总内存大小。

参数：

p_matrix_description – : 指向 wholememory_matrix_description_t 的指针。

返回值：

: 返回内存大小。

int64_t wholememory_get_memory_element_count_from_tensor(wholememory_tensor_description_t *p_tensor_description)#

从张量描述获取总元素数量。

参数：

p_tensor_description – : 指向 wholememory_tensor_description_t 的指针。

返回值：

: 返回元素数量。

int64_t wholememory_get_memory_size_from_tensor(wholememory_tensor_description_t *p_tensor_description)#

从张量描述中获取总内存大小。

参数：

p_tensor_description – : 指向 wholememory_tensor_description_t 的指针。

返回值：

: 返回内存大小。

bool wholememory_unsqueeze_tensor(wholememory_tensor_description_t *p_tensor_description, int dim)#

扩展张量维度

参数：

  • p_tensor_description – : wholememory_tensor_description_t 指针

  • dim – : 在哪个维度扩展

返回值：

: 成功返回 true，否则返回 false

wholememory_error_code_t wholememory_create_tensor(wholememory_tensor_t *wholememory_tensor, wholememory_tensor_description_t *tensor_description, wholememory_comm_t comm, wholememory_memory_type_t memory_type, wholememory_memory_location_t memory_location, size_t *tensor_entry_partition = nullptr)#

创建 WholeMemory 张量

参数：

  • wholememory_tensor – : 返回的 WholeMemory 张量句柄

  • tensor_description – : WholeMemory 张量的描述，应为不带偏移的 1 维或 2 维连续张量。

  • comm – : WholeMemory 通信器

  • memory_type – : 底层 WholeMemory 的内存类型

  • memory_location – : 底层 WholeMemory 的内存位置

  • tensor_entry_partition – : 每个 rank 的张量条目计数，长度必须等于 world_size。

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_destroy_tensor(wholememory_tensor_t wholememory_tensor)#

销毁 WholeMemory 张量

参数：

wholememory_tensor – : 要销毁的 WholeMemory 张量

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_make_tensor_from_pointer(wholememory_tensor_t *wholememory_tensor, void *storage_ptr, wholememory_tensor_description_t *tensor_description)#

从本地内存创建 WholeMemory 张量

参数：

  • wholememory_tensor – : 返回的 WholeMemory 张量句柄

  • storage_ptr – : 指向底层存储内存的指针。注意：存储指针可能与数据指针不同。

  • tensor_description – : WholeMemory 张量的描述，应为 1 维或 2 维

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_make_tensor_from_handle(wholememory_tensor_t *wholememory_tensor, wholememory_handle_t wholememory_handle, wholememory_tensor_description_t *tensor_description)#

从本地内存创建 WholeMemory 张量

参数：

  • wholememory_tensor – : 返回的 WholeMemory 张量句柄

  • wholememory_handle – : WholeMemory 句柄

  • tensor_description – : WholeMemory 张量的描述，应为 1 维或 2 维

返回值：

: wholememory_error_code_t

bool wholememory_tensor_has_handle(wholememory_tensor_t wholememory_tensor)#

检查是否具有 WholeMemory 句柄，wholmemory_make_tensor 创建的 WholeMemory 张量没有句柄

参数：

wholememory_tensor – : WholeMemory 张量

返回值：

: 是否具有 WholeMemory 句柄

wholememory_handle_t wholememory_tensor_get_memory_handle(wholememory_tensor_t wholememory_tensor)#

从 WholeMemory 张量获取 WholeMemory 句柄

参数：

wholememory_tensor – : WholeMemory 张量

返回值：

: WholeMemory 句柄

wholememory_tensor_description_t *wholememory_tensor_get_tensor_description(wholememory_tensor_t wholememory_tensor)#

从 WholeMemory 张量获取张量描述

参数：

wholememory_tensor – : WholeMemory 张量

返回值：

: 指向底层 wholememory_tensor_description_t 的指针

wholememory_error_code_t wholememory_tensor_get_global_reference(wholememory_tensor_t wholememory_tensor, wholememory_gref_t *wholememory_gref)#

从 WholeMemory 张量获取全局引用

参数：

  • wholememory_tensor – : WholeMemory 张量

  • wholememory_gref – : 全局引用

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_tensor_map_local_tensor(wholememory_tensor_t wholememory_tensor, wholememory_tensor_t *local_tensor)#

映射 WholeMemory 张量的局部张量。仅支持具有 WholeMemory 句柄的 1 维和 2 维张量。对于 1 维张量，storage_offset 应为 0。对于 2 维张量，storage_offset + size[1] 应 <= stride[0]

参数：

  • wholememory_tensor – : WholeMemory 张量。

  • local_tensor – : 返回的局部张量，需要被销毁。

返回值：

: wholememory_error_code_t

void *wholememory_tensor_get_data_pointer(wholememory_tensor_t wholememory_tensor)#

从 WholeMemory 张量获取数据指针

参数：

wholememory_tensor – : WholeMemory 张量

返回值：

: 指向 CONTINUOUS WholeMemory 或非 WholeMemory 的第一个数据点的指针。

警告

doxygenfunction: Cannot find function “wholememory_tensor_get_entry_per_partition” in doxygen xml output for project “libwholegraph” from directory: /tmp/tmp.yOXprHRhV5

wholememory_error_code_t wholememory_tensor_get_subtensor(wholememory_tensor_t wholememory_tensor, int64_t *starts, int64_t *ends, wholememory_tensor_t *sub_wholememory_tensor)#

获取 WholeMemory 张量的子张量

参数：

  • wholememory_tensor – : WholeMemory 张量

  • starts – : 各维度的起始索引，长度应与 wholememory_tensor 的维度相同。

  • ends – : 各维度的结束索引，长度应与 wholememory_tensor 的维度相同

  • sub_wholememory_tensor – : 指向返回的子张量的指针

返回值：

: wholememory_error_code_t

wholememory_tensor_t wholememory_tensor_get_root(wholememory_tensor_t wholememory_tensor)#

获取 WholeMemory 张量的根张量。根张量表示它不是任何 WholeMemory 张量的子张量。

参数：

wholememory_tensor – : WholeMemory 张量

返回值：

: 当前 WholeMemory 张量的根张量，可能与 wholememory_tensor 相同。

WholeMemory 张量操作#

wholememory_error_code_t wholememory_gather(wholememory_tensor_t wholememory_tensor, wholememory_tensor_t indices_tensor, wholememory_tensor_t output_tensor, wholememory_env_func_t *p_env_fns, void *stream, int gather_sms = -1)#

Gather 操作

参数：

  • wholememory_tensor – : 嵌入表的 WholeMemory 张量。

  • indices_tensor – : 要 Gather 的索引，不应是 WholeMemory 张量

  • output_tensor – : Gather 的输出张量，不应是 WholeMemory 张量

  • p_env_fns – : 指向环境函数的指针。

  • stream – : 要使用的 cudaStream_t。

  • gather_sms – : Gather 内核中使用的流多处理器数量

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_scatter(wholememory_tensor_t input_tensor, wholememory_tensor_t indices_tensor, wholememory_tensor_t wholememory_tensor, wholememory_env_func_t *p_env_fns, void *stream, int scatter_sms = -1)#

Scatter 操作

参数：

  • input_tensor – : 要 Scatter 的输入张量，不应是 WholeMemory 张量

  • indices_tensor – : 要 Scatter 的索引，不应是 WholeMemory 张量

  • wholememory_tensor – : 嵌入表的 WholeMemory 张量。

  • p_env_fns – : 指向环境函数的指针。

  • stream – : 要使用的 cudaStream_t。

  • scatter_sms – : Scatter 内核中使用的流多处理器数量

返回值：

: wholememory_error_code_t

WholeTensorEmbedding API#

typedef struct wholememory_embedding_cache_policy_ *wholememory_embedding_cache_policy_t#

WholeMemory 嵌入缓存策略的不透明句柄。

WholeMemory 嵌入缓存策略的不透明句柄

typedef struct wholememory_embedding_optimizer_ *wholememory_embedding_optimizer_t#

WholeMemory 嵌入优化器的不透明句柄。

WholeMemory 嵌入优化器的不透明句柄

typedef struct wholememory_embedding_ *wholememory_embedding_t#

WholeMemory 嵌入的不透明句柄。

WholeMemory 嵌入的不透明句柄

enum wholememory_access_type_t#

定义 WholeMemory 嵌入的访问类型

值

enumerator WHOLEMEMORY_AT_NONE#

未定义

enumerator WHOLEMEMORY_AT_READONLY#

仅对 WholeMemory 具有只读访问权限

enumerator WHOLEMEMORY_AT_READWRITE#

可能对 WholeMemory 具有写访问权限

enum wholememory_optimizer_type_t#

定义 WholeMemory 嵌入的优化器类型

值

enumerator WHOLEMEMORY_OPT_NONE#

无需优化器

enumerator WHOLEMEMORY_OPT_SGD#

使用 SGD 优化器

enumerator WHOLEMEMORY_OPT_LAZY_ADAM#

使用 Lazy Adam 优化器

enumerator WHOLEMEMORY_OPT_RMSPROP#

使用 RMSProp 优化器

enumerator WHOLEMEMORY_OPT_ADAGRAD#

使用 AdaGrad 优化器

wholememory_error_code_t wholememory_create_embedding_optimizer(wholememory_embedding_optimizer_t *optimizer, wholememory_optimizer_type_t optimizer_type)#

创建优化器

参数：

  • optimizer – : 返回的 wholememory_embedding_optimizer_t

  • optimizer_type – : 优化器类型

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_optimizer_set_parameter(wholememory_embedding_optimizer_t optimizer, const char *parameter_name, void *value)#

为优化器设置参数。

参数：

  • optimizer – : 要设置参数的优化器

  • parameter_name – : 参数名称

  • value – : 参数值

返回值：

: wholememory_error_code_t

void wholememory_destroy_embedding_optimizer(wholememory_embedding_optimizer_t optimizer)#

销毁优化器

参数：

optimizer – : 要销毁的优化器。

wholememory_error_code_t wholememory_create_embedding_cache_policy(wholememory_embedding_cache_policy_t *cache_policy, wholememory_comm_t cache_level_comm, wholememory_memory_type_t memory_type, wholememory_memory_location_t memory_location, wholememory_access_type_t access_type, float cache_ratio)#

创建 WholeMemory 嵌入缓存策略

参数：

  • cache_policy – : 返回的 wholememory_embedding_cache_policy_t

  • cache_level_comm – : 在哪个级别缓存完整嵌入。大多数情况下应与 wholememory_embedding_t 的 comm 相同。如果 access_type 是 WHOLEMEMORY_AT_READONLY，多个只读缓存可以有不同的级别。例如，对于一个多节点的 WHOLEMEMORY_MT_DISTRIBUTED WHOLEMEMORY_AT_READONLY 嵌入，它可以有一个节点内的 WHOLEMEMORY_MT_CHUNKED 缓存，或者一个多节点的 WHOLEMEMORY_MT_DISTRIBUTED 缓存。

  • memory_type – : 缓存的底层 WholeMemory 的内存类型

  • memory_location – : 缓存的底层 WholeMemory 的内存位置

  • access_type – : 只读或读写

  • cache_ratio – : 建议的缓存比例，取值范围应在 [1.0 / 512, 1.0] 内

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_destroy_embedding_cache_policy(wholememory_embedding_cache_policy_t cache_policy)#

销毁 WholeMemory 嵌入缓存策略

参数：

cache_policy – : 要销毁的 WholeMemory 嵌入缓存策略。

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_create_embedding(wholememory_embedding_t *wholememory_embedding, wholememory_tensor_description_t *embedding_tensor_description, wholememory_comm_t comm, wholememory_memory_type_t memory_type, wholememory_memory_location_t memory_location, wholememory_embedding_cache_policy_t cache_policy, size_t *embedding_entry_partition = nullptr, int user_defined_sms = -1, int round_robin_size = 0)#

创建 WholeMemory 嵌入

参数：

  • wholememory_embedding – : 返回的 wholememory_embedding_t

  • embedding_tensor_description – : 嵌入的描述，使用的大小和 dtype，stride 和 storage_offset 被忽略。必须是矩阵。

  • comm – : WholeMemory 通信器

  • memory_type – : 底层 WholeMemory 的内存类型

  • memory_location – : 底层 WholeMemory 的内存位置

  • cache_policy – : 此嵌入的缓存策略，如果不需要缓存，请使用 nullptr

  • embedding_entry_partition – 每个 rank 的嵌入条目计数，长度必须等于 world_size

  • user_defined_sms – : 用户定义的用于原始嵌入 gather/scatter 的 sms 数量

  • round_robin_size – : 轮询分片模式下每个 rank 的连续嵌入大小

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_destroy_embedding(wholememory_embedding_t wholememory_embedding)#

销毁 WholeMemory 嵌入

参数：

wholememory_embedding – : 要销毁的 WholeMemory 嵌入

返回值：

: wholememory_error_code_t

wholememory_tensor_t wholememory_embedding_get_embedding_tensor(wholememory_embedding_t wholememory_embedding)#

从 WholeMemory 嵌入获取 WholeMemory 张量。

参数：

wholememory_embedding – : WholeMemory 嵌入

返回值：

: WholeMemory 张量

wholememory_error_code_t wholememory_embedding_gather(wholememory_embedding_t wholememory_embedding, wholememory_tensor_t indices, wholememory_tensor_t output, bool adjust_cache, wholememory_env_func_t *p_env_fns, int64_t stream_int)#

从 WholeMemory 嵌入 Gather

参数：

  • wholememory_embedding – : WholeMemory 嵌入

  • indices – : 要 Gather 的索引

  • output – : 输出张量

  • adjust_cache – : 是否应在此 Gather 操作中调整缓存

  • p_env_fns – : 环境函数

  • stream_int – : 要使用的 CUDA 流

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_embedding_gather_gradient_apply(wholememory_embedding_t wholememory_embedding, wholememory_tensor_t indices, wholememory_tensor_t grads, bool adjust_cache, float lr, wholememory_env_func_t *p_env_fns, int64_t stream_int)#

WholeMemory 嵌入的 Gather 反向传播

参数：

  • wholememory_embedding – : WholeMemory 嵌入

  • indices – : 要 Gather 的索引

  • grads – : 输出张量的梯度

  • adjust_cache – : 是否应在此 Gather 操作中调整缓存

  • lr – : 当前步骤的学习率。

  • p_env_fns – : 环境函数

  • stream_int – : 要使用的 CUDA 流

返回值：

: wholememory_error_code_t

const char *const *wholememory_embedding_get_optimizer_state_names(wholememory_embedding_t wholememory_embedding)#

获取优化器内部状态名称

参数：

wholememory_embedding – : WholeMemory 嵌入

返回值：

: 以 nullptr 结尾的名称列表。

wholememory_tensor_t wholememory_embedding_get_optimizer_state(wholememory_embedding_t wholememory_embedding, const char *name)#

获取优化器内部状态

参数：

  • wholememory_embedding – : WholeMemory 嵌入

  • name – : 状态名称

返回值：

: 内部状态，如果不存在则为 nullptr。

wholememory_error_code_t wholememory_embedding_writeback_cache(wholememory_embedding_t wholememory_embedding, int64_t stream_int)#

回写所有 WholeMemory 嵌入的缓存

参数：

  • wholememory_embedding – : WholeMemory 嵌入

  • stream_int – : 要使用的 CUDA 流。

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholememory_embedding_drop_all_cache(wholememory_embedding_t wholememory_embedding, int64_t stream_int)#

丢弃 WholeMemory 嵌入中的所有缓存

参数：

  • wholememory_embedding – : WholeMemory 嵌入

  • stream_int – : 要使用的 CUDA 流。

返回值：

: wholememory_error_code_t

存储在 WholeMemory 中的图操作#

wholememory_error_code_t wholegraph_csr_unweighted_sample_without_replacement(wholememory_tensor_t wm_csr_row_ptr_tensor, wholememory_tensor_t wm_csr_col_ptr_tensor, wholememory_tensor_t center_nodes_tensor, int max_sample_count, wholememory_tensor_t output_sample_offset_tensor, void *output_dest_memory_context, void *output_center_localid_memory_context, void *output_edge_gid_memory_context, unsigned long long random_seed, wholememory_env_func_t *p_env_fns, void *stream)#

无权无放回采样内核操作

参数：

  • wm_csr_row_ptr_tensor – : 图的 csr_row_ptr 的 Wholememory 张量

  • wm_csr_col_ptr_tensor – : 图的 csr_col_ptr 的 Wholememory 张量

  • center_nodes_tensor – : 要采样的中心节点的非 Wholememory 张量

  • max_sample_count – : 最大采样计数

  • output_sample_offset_tensor – : 指向输出采样偏移量的指针

  • output_dest_memory_context – : 输出目标节点的内存上下文

  • output_center_localid_memory_context – : 输出中心本地 ID 的内存上下文

  • output_edge_gid_memory_context – : 输出边全局 ID 的内存上下文

  • random_seed – 随机数生成器种子

  • p_env_fns – : 指向环境函数的指针。

  • stream – : 要使用的 CUDA 流

返回值：

: wholememory_error_code_t

wholememory_error_code_t wholegraph_csr_weighted_sample_without_replacement(wholememory_tensor_t wm_csr_row_ptr_tensor, wholememory_tensor_t wm_csr_col_ptr_tensor, wholememory_tensor_t wm_csr_weight_ptr_tensor, wholememory_tensor_t center_nodes_tensor, int max_sample_count, wholememory_tensor_t output_sample_offset_tensor, void *output_dest_memory_context, void *output_center_localid_memory_context, void *output_edge_gid_memory_context, unsigned long long random_seed, wholememory_env_func_t *p_env_fns, void *stream)#

无权无放回采样内核操作

参数：

  • wm_csr_row_ptr_tensor – : 图的 csr_row_ptr 的 Wholememory 张量

  • wm_csr_col_ptr_tensor – : 图的 csr_col_ptr 的 Wholememory 张量

  • wm_csr_weight_ptr_tensor – : 图的边权重的 Wholememory 张量

  • center_nodes_tensor – : 要采样的中心节点的非 Wholememory 张量

  • max_sample_count – : 最大采样计数

  • output_sample_offset_tensor – : 指向输出采样偏移量的指针

  • output_dest_memory_context – : 输出目标节点的内存上下文

  • output_center_localid_memory_context – : 输出中心本地 ID 的内存上下文

  • output_edge_gid_memory_context – : 输出边全局 ID 的内存上下文

  • random_seed – 随机数生成器种子

  • p_env_fns – : 指向环境函数的指针。

  • stream – : 要使用的 CUDA 流

返回值：

: wholememory_error_code_t

图的杂项操作#

wholememory_error_code_t graph_append_unique(wholememory_tensor_t target_nodes_tensor, neighbor_nodes_tensor neighbor_nodes_tensor, void *output_unique_node_memory_context, wholememory_tensor_t output_neighbor_raw_to_unique_mapping_tensor, wholememory_env_func_t *p_env_fns, void *stream)#

Append Unique 操作

参数：

  • target_nodes_tensor – : 图的 csr_row_ptr 的 Wholememory 张量

  • neighbor_nodes_tensor – : 图的 csr_col_ptr 的 Wholememory 张量

  • output_unique_node_memory_context – : 输出目标节点的内存上下文

  • output_neighbor_raw_to_unique_mapping_tensor – : 指向输出采样偏移量的指针，可选输出

  • p_env_fns – : 指向环境函数的指针。

  • stream – : 要使用的 CUDA 流

返回值：

: wholememory_error_code_t

wholememory_error_code_t csr_add_self_loop(wholememory_tensor_t csr_row_ptr_tensor, wholememory_tensor_t csr_col_ptr_tensor, wholememory_tensor_t output_csr_row_ptr_tensor, wholememory_tensor_t output_csr_col_ptr_tensor, void *stream)#

Csr 添加自环操作

参数：

  • csr_row_ptr_tensor – : 局部图的 csr_row_ptr 的 Wholememory 张量

  • csr_col_ptr_tensor – : csr_col_ptr 的 Wholememory 张量

  • output_csr_row_ptr_tensor – : output_csr_row_ptr 的 Wholememory 张量

  • output_csr_col_ptr_tensor – : output_csr_col_ptr 的 Wholememory 张量

  • stream – : 要使用的 CUDA 流

返回值：

: wholememory_error_code_t

本页内容