Tensor and graph operations

Functions

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 Op

Parameters:

• wholememory_tensor – : WholeMemory Tensor of embedding table.

• indices_tensor – : indices to gather from, should NOT be WholeMemory Tensor

• output_tensor – : output tensor to gather to, should NOT be WholeMemoryTensor

• p_env_fns – : pointers to environment functions.

• stream – : cudaStream_t to use.

• gather_sms – : the number of stream multiprocessor used in gather kernel

Returns:

: 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 Op

Parameters:

• input_tensor – : input tensor tor scatter from, should NOT be WholeMemory Tensor

• indices_tensor – : indices to scatter to, should NOT be WholeMemory Tensor

• wholememory_tensor – : WholeMemory Tensor of embedding table.

• p_env_fns – : pointers to environment functions.

• stream – : cudaStream_t to use.

• scatter_sms – : the number of stream multiprocessor used in scatter kernel

Returns:

: wholememory_error_code_t

wholememory_error_code_t wholememory_env_test_op(

wholememory_tensor_t input_tensor,

wholememory_tensor_t output_fixed_tensor,

void *output_variable_device_tensor_handle,

void *output_variable_pinned_tensor_handle,

void *output_variable_host_tensor_handle,

int64_t output_variable_entry_count,

wholememory_env_func_t *p_env_fns,

void *stream

)

Just a test function,

Parameters:

• input_tensor – : input tensor

• output_fixed_tensor – : fixed size tensor of output

• output_variable_device_tensor_handle – : device version variable tensor

• output_variable_pinned_tensor_handle – : pinned version variable tensor

• output_variable_host_tensor_handle – : host version variable tensor

• output_variable_entry_count – : output entry count

• p_env_fns – : pointers to environment functions.

• stream – : cudaStream_t to use.

Returns:

: wholememory_error_code_t

Functions

wholememory_error_code_t graph_append_unique(

wholememory_tensor_t target_nodes_tensor,

wholememory_tensor_t 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 op

Parameters:

• target_nodes_tensor – : Wholememory Tensor of graph csr_row_ptr

• neighbor_nodes_tensor – : Wholememory Tensor of graph csr_col_ptr

• output_unique_node_memory_context – : memory context to output dest nodes

• output_neighbor_raw_to_unique_mapping_tensor – : pointer to output sample offset, optional output

• p_env_fns – : pointers to environment functions.

• stream – : CUDA stream to use

Returns:

: 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 Add Self Loop Op

Parameters:

• csr_row_ptr_tensor – : Wholememory Tensor of local graph csr_row_ptr

• csr_col_ptr_tensor – : Wholememory Tensor of csr_col_ptr

• output_csr_row_ptr_tensor – : Wholememory Tensor of output_csr_row_ptr

• output_csr_col_ptr_tensor – : Wholememory Tensor of output_csr_col_ptr

• stream – : CUDA stream to use

Returns:

: wholememory_error_code_t

Functions

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

)

Unweighted sample without replacement kernel op

Parameters:

• wm_csr_row_ptr_tensor – : Wholememory Tensor of graph csr_row_ptr

• wm_csr_col_ptr_tensor – : Wholememory Tensor of graph csr_col_ptr

• center_nodes_tensor – : None Wholememory Tensor of center node to sample

• max_sample_count – : maximum sample count

• output_sample_offset_tensor – : pointer to output sample offset

• output_dest_memory_context – : memory context to output dest nodes

• output_center_localid_memory_context – : memory context to output center local id

• output_edge_gid_memory_context – : memory context to output edge global id

• random_seed – random number generator seed

• p_env_fns – : pointers to environment functions.

• stream – : CUDA stream to use

Returns:

: 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

)

Unweighted sample without replacement kernel op

Parameters:

• wm_csr_row_ptr_tensor – : Wholememory Tensor of graph csr_row_ptr

• wm_csr_col_ptr_tensor – : Wholememory Tensor of graph csr_col_ptr

• wm_csr_weight_ptr_tensor – : Wholememory Tensor of graph edge weight

• center_nodes_tensor – : None Wholememory Tensor of center node to sample

• max_sample_count – : maximum sample count

• output_sample_offset_tensor – : pointer to output sample offset

• output_dest_memory_context – : memory context to output dest nodes

• output_center_localid_memory_context – : memory context to output center local id

• output_edge_gid_memory_context – : memory context to output edge global id

• random_seed – random number generator seed

• p_env_fns – : pointers to environment functions.

• stream – : CUDA stream to use

Returns:

: wholememory_error_code_t

wholememory_error_code_t generate_random_positive_int_cpu(

int64_t random_seed,

int64_t subsequence,

wholememory_tensor_t output

)

raft_pcg_generator_random_int cpu op

Parameters:

• random_seed – : random seed

• subsequence – : subsequence for generating random value

• output – : Wholememory Tensor of output

Returns:

: wholememory_error_code_t

wholememory_error_code_t generate_exponential_distribution_negative_float_cpu(

int64_t random_seed,

int64_t subsequence,

wholememory_tensor_t output

)

raft_pcg_generator_random_float cpu op

Parameters:

• random_seed – : random seed

• subsequence – : subsequence for generating random value

• output – : Wholememory Tensor of output

Returns:

: wholememory_error_code_t