Tensor descriptions and handles#

Defines

WHOLEMEMORY_MAX_TENSOR_DIM#

Enums

enum wholememory_dtype_t#

defines WholeMemory data type for tensors

Values:

enumerator WHOLEMEMORY_DT_UNKNOWN#

Unknown type

enumerator WHOLEMEMORY_DT_FLOAT#

32-bit float type

enumerator WHOLEMEMORY_DT_HALF#

16-bit half float type

enumerator WHOLEMEMORY_DT_DOUBLE#

64-bit double type

enumerator WHOLEMEMORY_DT_BF16#

16-bit bfloat type

enumerator WHOLEMEMORY_DT_INT#

32-bit signed integer type

enumerator WHOLEMEMORY_DT_INT64#

64-bit signed integer type

enumerator WHOLEMEMORY_DT_INT16#

16-bit signed integer type

enumerator WHOLEMEMORY_DT_INT8#

8-bit signed integer type

enumerator WHOLEMEMORY_DT_COUNT#

total count if types

Functions

size_t wholememory_dtype_get_element_size(wholememory_dtype_t dtype)#

Get element size of wholememory_dtype_t

Parameters:

dtype – : wholememory_dtype_t

Returns:

: element size of dtype, -1 on invalid dtype.

bool wholememory_dtype_is_floating_number(wholememory_dtype_t dtype)#

Check if dtype is floating number

Parameters:

dtype – : wholememory_dtype_t

Returns:

: True if dtype is WHOLEMEMORY_DT_FLOAT, WHOLEMEMORY_DT_HALF, WHOLEMEMORY_DT_DOUBLE or WHOLEMEMORY_DT_BF16. False otherwise.

bool wholememory_dtype_is_integer_number(wholememory_dtype_t dtype)#

Check if dtype is integer number

Parameters:

dtype – : wholememory_dtype_t

Returns:

: True if dtype is WHOLEMEMORY_DT_INT, WHOLEMEMORY_DT_INT64, WHOLEMEMORY_DT_INT16 or WHOLEMEMORY_DT_INT8, False otherwise.

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

Create wholememory_array_description_t object

Parameters:

  • size – : array size in number of elements

  • storage_offset – : storage offset in number of elements

  • dtype – : data type of array elements

Returns:

created 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
)#

Create wholememory_matrix_description_t object

Parameters:

  • sizes – : matrix sizes array, counted in number of elements, sizes[1] changes fastest.

  • stride – : stride of first dimension(slower changed dimension), stride is counted in number of elements

  • storage_offset – : storage offset in number of elements

  • dtype – : data type of matrix elements

Returns:

created wholememory_matrix_description_t

void wholememory_initialize_tensor_desc(
wholememory_tensor_description_t *p_tensor_description
)#

Initialize wholememory_tensor_description_t, set sizes and strides to all ones, and set storage_offset to 0, set dtype to WHOLEMEMORY_DT_UNKNOWN, set dim to 0.

Parameters:

p_tensor_description – : pointer to 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
)#

Copy array description to tensor description

Parameters:
void wholememory_copy_array_desc_to_tensor(
wholememory_tensor_description_t *p_tensor_description,
wholememory_array_description_t *p_array_description
)#

Copy array description to tensor description

Parameters:
void wholememory_copy_matrix_desc_to_tensor(
wholememory_tensor_description_t *p_tensor_description,
wholememory_matrix_description_t *p_matrix_description
)#

Copy matrix description to tensor description

Parameters:
bool wholememory_convert_tensor_desc_to_array(
wholememory_array_description_t *p_array_description,
wholememory_tensor_description_t *p_tensor_description
)#

Convert tensor description to array description

Parameters:
Returns:

: Return true if convertible else false.

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

Convert tensor description to matrix description

Parameters:
Returns:

: Return true if convertible else false.

int64_t wholememory_get_memory_element_count_from_array(
wholememory_array_description_t *p_array_description
)#

Get total element count from array description.

Parameters:

p_array_description – : pointer to wholememory_array_description_t.

Returns:

: Return element count.

int64_t wholememory_get_memory_size_from_array(
wholememory_array_description_t *p_array_description
)#

Get total memory size from array description.

Parameters:

p_array_description – : pointer to wholememory_array_description_t.

Returns:

: Return memory size.

int64_t wholememory_get_memory_element_count_from_matrix(
wholememory_matrix_description_t *p_matrix_description
)#

Get total element count from matrix description.

Parameters:

p_matrix_description – : pointer to wholememory_matrix_description_t.

Returns:

: Return element count.

int64_t wholememory_get_memory_size_from_matrix(
wholememory_matrix_description_t *p_matrix_description
)#

Get total memory size from matrix description.

Parameters:

p_matrix_description – : pointer to wholememory_matrix_description_t.

Returns:

: Return memory size.

int64_t wholememory_get_memory_element_count_from_tensor(
wholememory_tensor_description_t *p_tensor_description
)#

Get total element count from tensor description.

Parameters:

p_tensor_description – : pointer to wholememory_tensor_description_t.

Returns:

: Return element count.

int64_t wholememory_get_memory_size_from_tensor(
wholememory_tensor_description_t *p_tensor_description
)#

Get total memory size from tensor description.

Parameters:

p_tensor_description – : pointer to wholememory_tensor_description_t.

Returns:

: Return memory size.

bool wholememory_squeeze_tensor(
wholememory_tensor_description_t *p_tensor_description,
int dim
)#

Squeeze tensor

Parameters:
Returns:

: true if success else false

bool wholememory_unsqueeze_tensor(
wholememory_tensor_description_t *p_tensor_description,
int dim
)#

Unsqueeze tensor

Parameters:
Returns:

: true if success else false

struct wholememory_array_description_t#
#include <tensor_description.h>

wrapper for array in WholeMemory

Public Members

int64_t size#

size of the array in elements.

int64_t storage_offset#

offset in number of elements, NOT in bytes.

wholememory_dtype_t dtype#

data type of the array

struct wholememory_matrix_description_t#
#include <tensor_description.h>

wrapper for matrix in WholeMemory

Public Members

int64_t sizes[2]#

sizes[0] is row of the matrix, sizes[1] is column of the matrix

int64_t stride#

stride of first dimension, in number of elements

int64_t storage_offset#

offset in number of elements, NOT in bytes.

wholememory_dtype_t dtype#

data type of the matrix

struct wholememory_tensor_description_t#
#include <tensor_description.h>

Tensor description in WholeMemory, dimension 0 is the slowest changed dimension.

Public Members

int64_t sizes[WHOLEMEMORY_MAX_TENSOR_DIM]#

size of each dimension of the tensor, in number of elements

int64_t strides[WHOLEMEMORY_MAX_TENSOR_DIM]#

stride of the tensor, in number of elements

int64_t storage_offset#

offset in number of elements, NOT in bytes.

int dim#

dim of the tensor

wholememory_dtype_t dtype#

data type of the tensor

Defines

WM_TENSOR_COUNT_DEBUG#

Typedefs

typedef struct wholememory_tensor_ *wholememory_tensor_t#

Opaque handle to WholeMemoryTensor.

An Opaque handle to WholeMemoryTensor

Functions

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
)#

Create WholeMemory Tensor

Parameters:

  • wholememory_tensor – : returned WholeMemory Tensor handle

  • tensor_description – : description of the WholeMemory Tensor, should be 1-D or 2-D continuous tensor without offset.

  • comm – : WholeMemory Communicator

  • memory_type – : Memory Type of the underlying WholeMemory

  • memory_location – : Memory Location of the underlying WholeMemory

  • tensor_entry_partition – : Tensor entry count of each rank, the length must be world_size.

Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_destroy_tensor(
wholememory_tensor_t wholememory_tensor
)#

Destroy WholeMemory Tensor

Parameters:

wholememory_tensor – : WholeMemory Tensor to destroy

Returns:

: 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
)#

Make WholeMemory Tensor from local memory

Parameters:

  • wholememory_tensor – : returned WholeMemory Tensor handle

  • storage_ptr – : pointer to underlying storage memory. Note: storage pointer may be not same as data pointer.

  • tensor_description – : description of the WholeMemory Tensor, should be 1-D or 2-D

Returns:

: 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
)#

Make WholeMemory Tensor from local memory

Parameters:

  • wholememory_tensor – : returned WholeMemory Tensor handle

  • wholememory_handle – : WholeMemory Handle

  • tensor_description – : description of the WholeMemory Tensor, should be 1-D or 2-D

Returns:

: wholememory_error_code_t

bool wholememory_tensor_has_handle(
wholememory_tensor_t wholememory_tensor
)#

Check if has WholeMemory Handle, WholeMemory Tensor created by wholememory_make_tensor has no Handle

Parameters:

wholememory_tensor – : WholeMemory Tensor

Returns:

: if has WholeMemory Handle

wholememory_handle_t wholememory_tensor_get_memory_handle(
wholememory_tensor_t wholememory_tensor
)#

Get WholeMemory handle from WholeMemory Tensor

Parameters:

wholememory_tensor – : WholeMemory Tensor

Returns:

: WholeMemory handle

wholememory_tensor_description_t *wholememory_tensor_get_tensor_description(
wholememory_tensor_t wholememory_tensor
)#

Get tensor description from WholeMemory Tensor

Parameters:

wholememory_tensor – : WholeMemory Tensor

Returns:

: pointer to the underlying wholememory_tensor_description_t

wholememory_error_code_t wholememory_tensor_get_global_reference(
wholememory_tensor_t wholememory_tensor,
wholememory_gref_t *wholememory_gref
)#

Get global reference from WholeMemory Tensor

Parameters:

  • wholememory_tensor – : WholeMemory Tensor

  • wholememory_gref – : global reference

Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_tensor_map_local_tensor(
wholememory_tensor_t wholememory_tensor,
wholememory_tensor_t *local_tensor
)#

Map local tensor of WholeMemory Tensor. Only support 1D and 2D tensor with WholeMemory Handle. For 1D tensor, storage_offset should be 0 For 2D tensor, storage_offset + size[1] should <= stride[0]

Parameters:

  • wholememory_tensor – : WholeMemory Tensor.

  • local_tensor – : returned local tensor, need to be destroyed.

Returns:

: wholememory_error_code_t

void *wholememory_tensor_get_data_pointer(
wholememory_tensor_t wholememory_tensor
)#

Get data pointer from WholeMemory Tensor

Parameters:

wholememory_tensor – : WholeMemory Tensor

Returns:

: Pointer to first data for CONTINUOUS WholeMemory or not WholeMemory.

wholememory_error_code_t wholememory_tensor_get_entry_offsets(
size_t *entry_offsets,
wholememory_tensor_t wholememory_tensor
)#

Get entry offset of each rank from WholeMemory Tensor

Parameters:

  • entry_offsets – : returned entry offset of each rank

  • wholememory_tensor – : WholeMemory Tensor

Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_tensor_get_entry_partition_sizes(
size_t *entry_partition,
wholememory_tensor_t wholememory_tensor
)#

Get entry count of each rank from WholeMemory Tensor

Parameters:

  • entry_partition – : returned entry count of each rank

  • wholememory_tensor – : WholeMemory Tensor

Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_tensor_get_local_entry_count(
size_t *local_entry_count,
wholememory_tensor_t wholememory_tensor
)#

Get entry count of current rank from WholeMemory Tensor

Parameters:

  • local_entry_count – : returned entry count of current rank

  • wholememory_tensor – : WholeMemory Tensor

Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_tensor_get_local_entry_start(
size_t *local_entry_start,
wholememory_tensor_t wholememory_tensor
)#

Get entry start of current rank from WholeMemory Tensor

Parameters:

  • local_entry_start – : returned entry start id of current rank

  • wholememory_tensor – : WholeMemory Tensor

Returns:

: wholememory_error_code_t

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
)#

Get sub tensor of a WholeMemory Tensor

Parameters:

  • wholememory_tensor – : WholeMemory Tensor

  • starts – : starts of each dim, length should be the dim of wholememory_tensor.

  • ends – : ends of each dim, length should be the dim of wholememory_tensor

  • sub_wholememory_tensor – : pointer to returned sub tensor

Returns:

: wholememory_error_code_t

wholememory_tensor_t wholememory_tensor_get_root(
wholememory_tensor_t wholememory_tensor
)#

Get root tensor of a WholeMemory Tensor, root means it is not a sub tensor of any WholeMemory Tensor.

Parameters:

wholememory_tensor – : WholeMemory Tensor

Returns:

: the root of current WholeMemory tensor, maybe same as wholememory_tensor.

int64_t get_wholememory_tensor_count()#

Functions

wholememory_gref_t wholememory_create_continuous_global_reference(
void *ptr
)#

Create global reference for continuous memory.

Parameters:

ptr – : pointer to the memory

Returns:

: wholememory_gref_t

struct wholememory_gref_t#
#include <global_reference.h>

Global reference of a WholeMemory object.

A global reference is for Continuous or Chunked WholeMemory Type, in these types, each rank can directly access all memory from all ranks. The global reference is used to do this direct access.

Public Members

void *pointer#

pointer to data for CONTINUOUS WholeMemory or pointer to data pointer array for CHUNKED WholeMemory

size_t *rank_memory_offsets#

memory offset of each rank, and the length must be world_size+1

int world_size#
size_t stride#

must be 0 for CONTINUOUS WholeMemory or memory size in byte for each pointer

bool same_chunk#

if true, rank can be got by offset/stride

struct wholememory_nvshmem_ref_t#
#include <global_reference.h>

Public Members

void *pointer#
size_t *rank_memory_offsets#
size_t stride#
int world_rank#
int world_size#
bool same_chunk#
On this page