ucc.h

 
#ifndef UCC_H_
#define UCC_H_
 
#include <ucc/api/ucc_def.h>
#include <ucc/api/ucc_version.h>
#include <ucc/api/ucc_status.h>
#include <stdio.h>
 
BEGIN_C_DECLS
 
typedef enum {
    UCC_COLL_TYPE_ALLGATHER          = UCC_BIT(0),
    UCC_COLL_TYPE_ALLGATHERV         = UCC_BIT(1),
    UCC_COLL_TYPE_ALLREDUCE          = UCC_BIT(2),
    UCC_COLL_TYPE_ALLTOALL           = UCC_BIT(3),
    UCC_COLL_TYPE_ALLTOALLV          = UCC_BIT(4),
    UCC_COLL_TYPE_BARRIER            = UCC_BIT(5),
    UCC_COLL_TYPE_BCAST              = UCC_BIT(6),
    UCC_COLL_TYPE_FANIN              = UCC_BIT(7),
    UCC_COLL_TYPE_FANOUT             = UCC_BIT(8),
    UCC_COLL_TYPE_GATHER             = UCC_BIT(9),
    UCC_COLL_TYPE_GATHERV            = UCC_BIT(10),
    UCC_COLL_TYPE_REDUCE             = UCC_BIT(11),
    UCC_COLL_TYPE_REDUCE_SCATTER     = UCC_BIT(12),
    UCC_COLL_TYPE_REDUCE_SCATTERV    = UCC_BIT(13),
    UCC_COLL_TYPE_SCATTER            = UCC_BIT(14),
    UCC_COLL_TYPE_SCATTERV           = UCC_BIT(15),
    UCC_COLL_TYPE_LAST
} ucc_coll_type_t;
 
typedef enum ucc_memory_type {
    UCC_MEMORY_TYPE_HOST,         
    UCC_MEMORY_TYPE_CUDA,         
    UCC_MEMORY_TYPE_CUDA_MANAGED, 
    UCC_MEMORY_TYPE_ROCM,         
    UCC_MEMORY_TYPE_ROCM_MANAGED, 
    UCC_MEMORY_TYPE_LAST,
    UCC_MEMORY_TYPE_UNKNOWN = UCC_MEMORY_TYPE_LAST
} ucc_memory_type_t;
 
typedef uint64_t ucc_datatype_t;
 
#define   UCC_DT_INT8      UCC_PREDEFINED_DT(0)
#define   UCC_DT_INT16     UCC_PREDEFINED_DT(1)
#define   UCC_DT_INT32     UCC_PREDEFINED_DT(2)
#define   UCC_DT_INT64     UCC_PREDEFINED_DT(3)
#define   UCC_DT_INT128    UCC_PREDEFINED_DT(4)
#define   UCC_DT_UINT8     UCC_PREDEFINED_DT(5)
#define   UCC_DT_UINT16    UCC_PREDEFINED_DT(6)
#define   UCC_DT_UINT32    UCC_PREDEFINED_DT(7)
#define   UCC_DT_UINT64    UCC_PREDEFINED_DT(8)
#define   UCC_DT_UINT128   UCC_PREDEFINED_DT(9)
#define   UCC_DT_FLOAT16   UCC_PREDEFINED_DT(10)
#define   UCC_DT_FLOAT32   UCC_PREDEFINED_DT(11)
#define   UCC_DT_FLOAT64   UCC_PREDEFINED_DT(12)
#define   UCC_DT_BFLOAT16  UCC_PREDEFINED_DT(13)
#define   UCC_DT_PREDEFINED_LAST  14
 
enum ucc_generic_dt_ops_field {
    UCC_GENERIC_DT_OPS_FIELD_FLAGS             = UCC_BIT(0),
};
 
typedef enum {
    UCC_GENERIC_DT_OPS_FLAG_CONTIG             = UCC_BIT(0), 
    UCC_GENERIC_DT_OPS_FLAG_REDUCE             = UCC_BIT(1), 
} ucc_generic_dt_ops_flags_t;
 
typedef struct ucc_reduce_cb_params {
    uint64_t          mask;      /*< for backward compatibility. currently ignored. */
    void             *src1;      /*< input buffer */
    void             *src2;      /*< input buffer: represents n_vectors buffers with
                                   offset "stride" between them */
    void             *dst;       /*< destination buffer */
    size_t            n_vectors; /*< number of vectors from src2 to reduce */
    size_t            count;     /*< number of elements in one vector */
    size_t            stride;    /*< stride in bytes between the vectors in src2 */
    ucc_dt_generic_t *dt;        /*< pointer to user-defined datatype used for
                                     reduction */
    void             *cb_ctx;    /*< user-defined context as defined
                                   by @ref ucc_generic_dt_ops::reduce.cb_ctx */
} ucc_reduce_cb_params_t;
typedef struct ucc_generic_dt_ops {
    uint64_t mask;
    uint64_t flags;
    size_t   contig_size; 
    void* (*start_pack)(void *context, const void *buffer, size_t count);
 
    void* (*start_unpack)(void *context, void *buffer, size_t count);
 
    size_t (*packed_size)(void *state);
 
    size_t (*pack) (void *state, size_t offset, void *dest, size_t max_length);
 
    ucc_status_t (*unpack)(void *state, size_t offset, const void *src, size_t length);
 
    void (*finish)(void *state);
 
    struct {
        ucc_status_t (*cb)(const ucc_reduce_cb_params_t *params);
        void *cb_ctx;
    } reduce;
} ucc_generic_dt_ops_t;
 
 
ucc_status_t ucc_dt_create_generic(const ucc_generic_dt_ops_t *ops, void *context,
                                   ucc_datatype_t *datatype_p);
 
void ucc_dt_destroy(ucc_datatype_t datatype);
 
typedef enum {
    UCC_OP_SUM,
    UCC_OP_PROD,
    UCC_OP_MAX,
    UCC_OP_MIN,
    UCC_OP_LAND,
    UCC_OP_LOR,
    UCC_OP_LXOR,
    UCC_OP_BAND,
    UCC_OP_BOR,
    UCC_OP_BXOR,
    UCC_OP_MAXLOC,
    UCC_OP_MINLOC,
    UCC_OP_AVG,
    UCC_OP_LAST
} ucc_reduction_op_t;
 
typedef enum {
    UCC_THREAD_SINGLE       = 0, 
    UCC_THREAD_FUNNELED     = 1, 
    UCC_THREAD_MULTIPLE     = 2  
} ucc_thread_mode_t;
 
typedef enum {
    UCC_NO_SYNC_COLLECTIVES = 0, 
    UCC_SYNC_COLLECTIVES    = 1  
} ucc_coll_sync_type_t;
 
 
enum ucc_lib_params_field{
    UCC_LIB_PARAM_FIELD_THREAD_MODE         = UCC_BIT(0),
    UCC_LIB_PARAM_FIELD_COLL_TYPES          = UCC_BIT(1),
    UCC_LIB_PARAM_FIELD_REDUCTION_TYPES     = UCC_BIT(2),
    UCC_LIB_PARAM_FIELD_SYNC_TYPE           = UCC_BIT(3)
};
 
enum ucc_lib_attr_field{
    UCC_LIB_ATTR_FIELD_THREAD_MODE       = UCC_BIT(0),
    UCC_LIB_ATTR_FIELD_COLL_TYPES        = UCC_BIT(1),
    UCC_LIB_ATTR_FIELD_REDUCTION_TYPES   = UCC_BIT(2),
    UCC_LIB_ATTR_FIELD_SYNC_TYPE         = UCC_BIT(3)
};
 
typedef struct ucc_lib_params {
    uint64_t                mask;
    ucc_thread_mode_t       thread_mode;
    uint64_t                coll_types;
    uint64_t                reduction_types;
    ucc_coll_sync_type_t    sync_type;
} ucc_lib_params_t;
 
typedef struct ucc_lib_attr {
    uint64_t                mask;
    ucc_thread_mode_t       thread_mode;
    uint64_t                coll_types;
    uint64_t                reduction_types;
    ucc_coll_sync_type_t    sync_type;
} ucc_lib_attr_t;
 
 
ucc_status_t ucc_lib_config_read(const char *env_prefix, const char *filename,
                                 ucc_lib_config_h *config);
 
void ucc_lib_config_release(ucc_lib_config_h config);
 
void ucc_lib_config_print(const ucc_lib_config_h config, FILE *stream,
                          const char *title, ucc_config_print_flags_t print_flags);
 
ucc_status_t ucc_lib_config_modify(ucc_lib_config_h config, const char *name,
                                   const char *value);
 
 
ucc_status_t ucc_init_version(unsigned api_major_version,
                              unsigned api_minor_version,
                              const ucc_lib_params_t *params,
                              const ucc_lib_config_h config,
                              ucc_lib_h *lib_p);
 
 
static inline ucc_status_t ucc_init(const ucc_lib_params_t *params,
                                    const ucc_lib_config_h config,
                                    ucc_lib_h *lib_p)
{
    return ucc_init_version(UCC_API_MAJOR, UCC_API_MINOR, params, config,
                            lib_p);
}
 
 
ucc_status_t ucc_finalize(ucc_lib_h lib_p);
 
 
ucc_status_t ucc_lib_get_attr(ucc_lib_h lib_p, ucc_lib_attr_t *lib_attr);
 
/*
 * *************************************************************
 *                   Context Section
 * *************************************************************
 */
 
typedef enum {
   UCC_CONTEXT_EXCLUSIVE = 0 ,
   UCC_CONTEXT_SHARED
} ucc_context_type_t;
 
enum ucc_context_params_field {
    UCC_CONTEXT_PARAM_FIELD_TYPE              = UCC_BIT(0),
    UCC_CONTEXT_PARAM_FIELD_SYNC_TYPE         = UCC_BIT(1),
    UCC_CONTEXT_PARAM_FIELD_OOB               = UCC_BIT(2),
    UCC_CONTEXT_PARAM_FIELD_ID                = UCC_BIT(3),
    UCC_CONTEXT_PARAM_FIELD_MEM_PARAMS        = UCC_BIT(4)
};
 
enum ucc_context_attr_field {
    UCC_CONTEXT_ATTR_FIELD_TYPE               = UCC_BIT(0),
    UCC_CONTEXT_ATTR_FIELD_SYNC_TYPE          = UCC_BIT(1),
    UCC_CONTEXT_ATTR_FIELD_CTX_ADDR           = UCC_BIT(2),
    UCC_CONTEXT_ATTR_FIELD_CTX_ADDR_LEN       = UCC_BIT(3),
    UCC_CONTEXT_ATTR_FIELD_WORK_BUFFER_SIZE   = UCC_BIT(4)
};
 
typedef struct ucc_oob_coll {
    ucc_status_t    (*allgather)(void *src_buf, void *recv_buf, size_t size,
                                 void *allgather_info,  void **request);
    ucc_status_t    (*req_test)(void *request);
    ucc_status_t    (*req_free)(void *request);
    void            *coll_info;
    uint32_t        n_oob_eps; 
    uint32_t        oob_ep; 
}  ucc_oob_coll_t;
 
typedef ucc_oob_coll_t ucc_context_oob_coll_t;
typedef ucc_oob_coll_t ucc_team_oob_coll_t;
 
typedef struct ucc_mem_map {
    void *   address; 
    size_t   len;     
} ucc_mem_map_t;
 
typedef struct ucc_mem_map_params {
    ucc_mem_map_t *segments;   
    uint64_t       n_segments; 
} ucc_mem_map_params_t;
 
typedef struct ucc_context_params {
    uint64_t                mask;
    ucc_context_type_t      type;
    ucc_coll_sync_type_t    sync_type;
    ucc_context_oob_coll_t  oob;
    uint64_t                ctx_id;
    ucc_mem_map_params_t    mem_params;
} ucc_context_params_t;
 
typedef struct ucc_context_attr {
    uint64_t                mask;
    ucc_context_type_t      type;
    ucc_coll_sync_type_t    sync_type;
    ucc_context_addr_h      ctx_addr;
    ucc_context_addr_len_t  ctx_addr_len;
    uint64_t                global_work_buffer_size;
} ucc_context_attr_t;
 
ucc_status_t ucc_context_config_read(ucc_lib_h lib_handle,
                                     const char *filename,
                                     ucc_context_config_h *config);
 
void ucc_context_config_release(ucc_context_config_h config);
 
void ucc_context_config_print(const ucc_context_config_h config, FILE *stream,
                              const char *title, ucc_config_print_flags_t print_flags);
 
ucc_status_t ucc_context_config_modify(ucc_context_config_h config,
                                       const char *component, const char *name,
                                       const char *value);
 
ucc_status_t ucc_context_create(ucc_lib_h lib_handle,
                                const ucc_context_params_t *params,
                                const ucc_context_config_h  config,
                                ucc_context_h *context);
 
ucc_status_t ucc_context_progress(ucc_context_h context);
 
ucc_status_t ucc_context_destroy(ucc_context_h context);
 
ucc_status_t ucc_context_get_attr(ucc_context_h context,
                                  ucc_context_attr_t *context_attr);
 
/*
 * *************************************************************
 *                   Teams Section
 * *************************************************************
 */
 
enum ucc_team_params_field {
    UCC_TEAM_PARAM_FIELD_ORDERING               = UCC_BIT(0),
    UCC_TEAM_PARAM_FIELD_OUTSTANDING_COLLS      = UCC_BIT(1),
    UCC_TEAM_PARAM_FIELD_EP                     = UCC_BIT(2),
    UCC_TEAM_PARAM_FIELD_EP_LIST                = UCC_BIT(3),
    UCC_TEAM_PARAM_FIELD_EP_RANGE               = UCC_BIT(4),
    UCC_TEAM_PARAM_FIELD_TEAM_SIZE              = UCC_BIT(5),
    UCC_TEAM_PARAM_FIELD_SYNC_TYPE              = UCC_BIT(6),
    UCC_TEAM_PARAM_FIELD_OOB                    = UCC_BIT(7),
    UCC_TEAM_PARAM_FIELD_P2P_CONN               = UCC_BIT(8),
    UCC_TEAM_PARAM_FIELD_MEM_PARAMS             = UCC_BIT(9),
    UCC_TEAM_PARAM_FIELD_EP_MAP                 = UCC_BIT(10),
    UCC_TEAM_PARAM_FIELD_ID                     = UCC_BIT(11),
    UCC_TEAM_PARAM_FIELD_FLAGS                  = UCC_BIT(12)
};
 
enum ucc_team_attr_field {
    UCC_TEAM_ATTR_FIELD_POST_ORDERING          = UCC_BIT(0),
    UCC_TEAM_ATTR_FIELD_OUTSTANDING_CALLS      = UCC_BIT(1),
    UCC_TEAM_ATTR_FIELD_EP                     = UCC_BIT(2),
    UCC_TEAM_ATTR_FIELD_EP_RANGE               = UCC_BIT(3),
    UCC_TEAM_ATTR_FIELD_SYNC_TYPE              = UCC_BIT(4),
    UCC_TEAM_ATTR_FIELD_MEM_PARAMS             = UCC_BIT(5)
};
 
enum ucc_team_flags {
    UCC_TEAM_FLAG_COLL_WORK_BUFFER             = UCC_BIT(0) /*< If set, this indicates
                                                                the user will provide 
                                                                a scratchpad buffer for 
                                                                use in one-sided 
                                                                collectives. Otherwise, 
                                                                an internal buffer will
                                                                used. */
};
 
typedef struct ucc_team_p2p_conn {
    int             (*conn_info_lookup)(void *conn_ctx, uint64_t ep, ucc_p2p_conn_t **conn_info, void *request);
    int             (*conn_info_release)(ucc_p2p_conn_t *conn_info);
    void            *conn_ctx;
    ucc_status_t    (*req_test)(void *request);
    ucc_status_t    (*req_free)(void *request);
} ucc_team_p2p_conn_t;
 
typedef enum {
 
    UCC_COLLECTIVE_POST_ORDERED             = 0,
 
    UCC_COLLECTIVE_POST_UNORDERED           = 1,
 
    UCC_COLLECTIVE_INIT_ORDERED             = 2,
 
    UCC_COLLECTIVE_INIT_UNORDERED           = 3,
 
    UCC_COLLECTIVE_INIT_AND_POST_ORDERED    = 4,
 
    UCC_COLLECTIVE_INIT_AND_POST_UNORDERED  = 5
} ucc_post_ordering_t;
 
typedef enum {
    UCC_COLLECTIVE_EP_RANGE_CONTIG      = 0,
    UCC_COLLECTIVE_EP_RANGE_NONCONTIG   = 1
} ucc_ep_range_type_t;
 
struct ucc_ep_map_strided {
    uint64_t    start;
    int64_t     stride;
};
 
struct ucc_ep_map_array {
    void    *map;
    size_t  elem_size; 
};
 
struct ucc_ep_map_cb {
    uint64_t   (*cb)(uint64_t ep, void *cb_ctx);
    void       *cb_ctx;
};
 
typedef enum {
    UCC_EP_MAP_FULL     = 1, 
    UCC_EP_MAP_STRIDED  = 2, 
    UCC_EP_MAP_ARRAY    = 3, 
    UCC_EP_MAP_CB       = 4, 
} ucc_ep_map_type_t;
 
typedef struct ucc_ep_map_t {
    ucc_ep_map_type_t type;
    uint64_t          ep_num; 
    union {
        struct ucc_ep_map_strided strided;
        struct ucc_ep_map_array   array;
        struct ucc_ep_map_cb      cb;
    };
} ucc_ep_map_t;
 
typedef struct ucc_team_params {
 
    uint64_t                mask;
    uint64_t                flags;
    ucc_post_ordering_t     ordering;
 
    uint64_t                outstanding_colls;
 
    uint64_t                ep;
 
    uint64_t                *ep_list;
 
    ucc_ep_range_type_t     ep_range;
 
    uint64_t                team_size;
 
    ucc_coll_sync_type_t    sync_type;
 
    ucc_team_oob_coll_t     oob;
 
    ucc_team_p2p_conn_t     p2p_conn;
 
    ucc_mem_map_params_t    mem_params;
 
    ucc_ep_map_t            ep_map;
 
    uint64_t                id;
} ucc_team_params_t;
 
typedef struct ucc_team_attr {
    uint64_t               mask;
    ucc_post_ordering_t    ordering;
    uint64_t               outstanding_colls;
    uint64_t               ep;
    ucc_ep_range_type_t    ep_range;
    ucc_coll_sync_type_t   sync_type;
    ucc_mem_map_params_t   mem_params;
} ucc_team_attr_t;
 
 
ucc_status_t ucc_team_create_post(ucc_context_h *contexts,
                                  uint32_t num_contexts,
                                  const ucc_team_params_t *team_params,
                                  ucc_team_h *new_team);
 
 
ucc_status_t ucc_team_create_test(ucc_team_h team);
 
 
ucc_status_t ucc_team_destroy(ucc_team_h team);
 
ucc_status_t ucc_team_get_attr(ucc_team_h team,
                               ucc_team_attr_t *team_attr);
 
ucc_status_t ucc_team_create_from_parent(uint64_t my_ep, uint32_t included,
                                         ucc_team_h parent_team,
                                         ucc_team_h *new_team);
 
ucc_status_t ucc_team_get_size(ucc_team_h team, uint32_t *size);
 
 
ucc_status_t ucc_team_get_my_ep(ucc_team_h team, uint64_t *ep);
 
ucc_status_t ucc_team_get_all_eps(ucc_team_h team, uint64_t **ep,
                                  uint64_t *num_eps);
 
 
/*
 * *************************************************************
 *                   Collectives Section
 * *************************************************************
 */
 
typedef enum {
    UCC_COLL_ARGS_FLAG_IN_PLACE             = UCC_BIT(0), 
    UCC_COLL_ARGS_FLAG_PERSISTENT           = UCC_BIT(1), 
    UCC_COLL_ARGS_FLAG_COUNT_64BIT          = UCC_BIT(2), 
    UCC_COLL_ARGS_FLAG_DISPLACEMENTS_64BIT  = UCC_BIT(3), 
    UCC_COLL_ARGS_FLAG_CONTIG_SRC_BUFFER    = UCC_BIT(4), 
    UCC_COLL_ARGS_FLAG_CONTIG_DST_BUFFER    = UCC_BIT(5), 
    UCC_COLL_ARGS_FLAG_TIMEOUT              = UCC_BIT(6), 
    UCC_COLL_ARGS_FLAG_MEM_MAPPED_BUFFERS   = UCC_BIT(7)  
} ucc_coll_args_flags_t;
 
typedef struct ucc_coll_buffer_info_v {
    void             *buffer; 
    ucc_count_t      *counts; 
    ucc_aint_t       *displacements; 
    ucc_datatype_t    datatype; 
    ucc_memory_type_t mem_type; 
} ucc_coll_buffer_info_v_t;
 
typedef struct ucc_coll_buffer_info {
    void             *buffer;   
    ucc_count_t       count;    
    ucc_datatype_t    datatype; 
    ucc_memory_type_t mem_type; 
} ucc_coll_buffer_info_t;
 
typedef enum {
    UCC_ERR_TYPE_LOCAL      = 0,
    UCC_ERR_TYPE_GLOBAL     = 1
} ucc_error_type_t;
 
enum ucc_coll_args_field {
    UCC_COLL_ARGS_FIELD_FLAGS                           = UCC_BIT(0),
    UCC_COLL_ARGS_FIELD_TAG                             = UCC_BIT(1),
    UCC_COLL_ARGS_FIELD_CB                              = UCC_BIT(2),
    UCC_COLL_ARGS_FIELD_GLOBAL_WORK_BUFFER              = UCC_BIT(3)
};
 
typedef struct ucc_coll_args {
    uint64_t                        mask;
    ucc_coll_type_t                 coll_type; 
    union {
        ucc_coll_buffer_info_t      info;   
        ucc_coll_buffer_info_v_t    info_v; 
    } src;
    union {
        ucc_coll_buffer_info_t      info;   
        ucc_coll_buffer_info_v_t    info_v; 
    } dst;
    ucc_reduction_op_t              op; 
    uint64_t                        flags;
    uint64_t                        root; 
    ucc_error_type_t                error_type; 
    ucc_coll_id_t                   tag; 
    void                           *global_work_buffer; 
    ucc_coll_callback_t             cb;
    double                          timeout; 
} ucc_coll_args_t;
 
ucc_status_t ucc_collective_init(ucc_coll_args_t *coll_args,
                                 ucc_coll_req_h *request, ucc_team_h team);
 
ucc_status_t ucc_collective_post(ucc_coll_req_h request);
 
 
ucc_status_t ucc_collective_init_and_post(ucc_coll_args_t *coll_args,
                                          ucc_coll_req_h *request,
                                          ucc_team_h team);
 
static inline ucc_status_t ucc_collective_test(ucc_coll_req_h request)
{
    return request->status;
}
 
ucc_status_t ucc_collective_finalize(ucc_coll_req_h request);
 
typedef enum ucc_event_type {
    UCC_EVENT_COLLECTIVE_POST     = UCC_BIT(0),
    UCC_EVENT_COLLECTIVE_COMPLETE = UCC_BIT(1),
    UCC_EVENT_COMPUTE_COMPLETE    = UCC_BIT(2),
    UCC_EVENT_OVERFLOW            = UCC_BIT(3)
} ucc_event_type_t;
 
typedef enum ucc_ee_type {
    UCC_EE_CUDA_STREAM = 0,
    UCC_EE_CPU_THREAD,
    UCC_EE_LAST,
    UCC_EE_UNKNOWN = UCC_EE_LAST
} ucc_ee_type_t;
 
typedef struct ucc_event {
    ucc_event_type_t ev_type;
    void *           ev_context;
    size_t           ev_context_size;
    ucc_coll_req_h   req;
} ucc_ev_t;
 
typedef struct ucc_ee_params {
    ucc_ee_type_t ee_type;
    void *        ee_context;
    size_t        ee_context_size;
} ucc_ee_params_t;
 
ucc_status_t ucc_ee_create(ucc_team_h team, const ucc_ee_params_t *params,
                           ucc_ee_h *ee);
 
ucc_status_t ucc_ee_destroy(ucc_ee_h ee);
 
ucc_status_t ucc_ee_get_event(ucc_ee_h ee, ucc_ev_t **ev);
 
ucc_status_t ucc_ee_ack_event(ucc_ee_h ee, ucc_ev_t *ev);
 
ucc_status_t ucc_ee_set_event(ucc_ee_h ee, ucc_ev_t *ev);
 
ucc_status_t ucc_ee_wait(ucc_ee_h ee, ucc_ev_t *ev);
 
ucc_status_t ucc_collective_triggered_post(ucc_ee_h ee, ucc_ev_t *ee_event);
 
END_C_DECLS
#endif