|
| 1 | +/* |
| 2 | + * Copyright (c) 2021 Mellanox Technologies. All rights reserved. |
| 3 | + * $COPYRIGHT$ |
| 4 | + * |
| 5 | + * Additional copyrights may follow |
| 6 | + * |
| 7 | + * $HEADER$ |
| 8 | + */ |
| 9 | + |
| 10 | +#ifndef MCA_SCOLL_UCC_H |
| 11 | +#define MCA_SCOLL_UCC_H |
| 12 | + |
| 13 | +#include "oshmem_config.h" |
| 14 | + |
| 15 | +#include "shmem.h" |
| 16 | +#include "oshmem/mca/mca.h" |
| 17 | +#include "oshmem/mca/scoll/scoll.h" |
| 18 | +#include "oshmem/proc/proc.h" |
| 19 | + |
| 20 | +#include "scoll_ucc_debug.h" |
| 21 | + |
| 22 | +#include <ucc/api/ucc.h> |
| 23 | + |
| 24 | +BEGIN_C_DECLS |
| 25 | + |
| 26 | +#define SCOLL_UCC_CTS (UCC_COLL_TYPE_BARRIER | UCC_COLL_TYPE_BCAST | \ |
| 27 | + UCC_COLL_TYPE_ALLREDUCE | UCC_COLL_TYPE_ALLGATHER | \ |
| 28 | + UCC_COLL_TYPE_ALLTOALL) |
| 29 | + |
| 30 | +#define SCOLL_UCC_CTS_STR "barrier,broadcast,reduce,collect,alltoall" |
| 31 | + |
| 32 | +int mca_scoll_ucc_progress(void); |
| 33 | + |
| 34 | +/** |
| 35 | + * Globally exported structure |
| 36 | + */ |
| 37 | +struct mca_scoll_ucc_component_t { |
| 38 | + mca_scoll_base_component_1_0_0_t super; |
| 39 | + int ucc_priority; |
| 40 | + int ucc_verbose; |
| 41 | + int ucc_enable; |
| 42 | + int ucc_np; |
| 43 | + char * cls; |
| 44 | + char * cts; |
| 45 | + bool libucc_initialized; |
| 46 | + ucc_lib_h ucc_lib; |
| 47 | + ucc_lib_attr_t ucc_lib_attr; |
| 48 | + ucc_coll_type_t cts_requested; |
| 49 | + ucc_context_h ucc_context; |
| 50 | +}; |
| 51 | +typedef struct mca_scoll_ucc_component_t mca_scoll_ucc_component_t; |
| 52 | + |
| 53 | +OMPI_MODULE_DECLSPEC extern mca_scoll_ucc_component_t mca_scoll_ucc_component; |
| 54 | + |
| 55 | +/** |
| 56 | + * UCC enabled team |
| 57 | + */ |
| 58 | +struct mca_scoll_ucc_module_t { |
| 59 | + mca_scoll_base_module_t super; |
| 60 | + |
| 61 | + oshmem_group_t *group; |
| 62 | + ucc_team_h ucc_team; |
| 63 | + |
| 64 | + /* Saved handlers - for fallback */ |
| 65 | + mca_scoll_base_module_reduce_fn_t previous_reduce; |
| 66 | + mca_scoll_base_module_t *previous_reduce_module; |
| 67 | + mca_scoll_base_module_broadcast_fn_t previous_broadcast; |
| 68 | + mca_scoll_base_module_t *previous_broadcast_module; |
| 69 | + mca_scoll_base_module_barrier_fn_t previous_barrier; |
| 70 | + mca_scoll_base_module_t *previous_barrier_module; |
| 71 | + mca_scoll_base_module_collect_fn_t previous_collect; |
| 72 | + mca_scoll_base_module_t *previous_collect_module; |
| 73 | + mca_scoll_base_module_alltoall_fn_t previous_alltoall; |
| 74 | + mca_scoll_base_module_t *previous_alltoall_module; |
| 75 | +}; |
| 76 | +typedef struct mca_scoll_ucc_module_t mca_scoll_ucc_module_t; |
| 77 | + |
| 78 | +OBJ_CLASS_DECLARATION(mca_scoll_ucc_module_t); |
| 79 | + |
| 80 | +/* API functions */ |
| 81 | +int mca_scoll_ucc_init_query(bool enable_progress_threads, bool enable_mpi_threads); |
| 82 | + |
| 83 | +mca_scoll_base_module_t* mca_scoll_ucc_comm_query(oshmem_group_t *osh_group, int *priority); |
| 84 | + |
| 85 | +int mca_scoll_ucc_barrier(struct oshmem_group_t *group, long *pSync, int alg); |
| 86 | + |
| 87 | +int mca_scoll_ucc_broadcast(struct oshmem_group_t *group, |
| 88 | + int PE_root, |
| 89 | + void *target, |
| 90 | + const void *source, |
| 91 | + size_t nlong, |
| 92 | + long *pSync, |
| 93 | + bool nlong_type, |
| 94 | + int alg); |
| 95 | + |
| 96 | +int mca_scoll_ucc_collect(struct oshmem_group_t *group, |
| 97 | + void *target, |
| 98 | + const void *source, |
| 99 | + size_t nlong, |
| 100 | + long *pSync, |
| 101 | + bool nlong_type, |
| 102 | + int alg); |
| 103 | + |
| 104 | +int mca_scoll_ucc_reduce(struct oshmem_group_t *group, |
| 105 | + struct oshmem_op_t *op, |
| 106 | + void *target, |
| 107 | + const void *source, |
| 108 | + size_t nlong, |
| 109 | + long *pSync, |
| 110 | + void *pWrk, |
| 111 | + int alg); |
| 112 | + |
| 113 | +int mca_scoll_ucc_alltoall(struct oshmem_group_t *group, |
| 114 | + void *target, |
| 115 | + const void *source, |
| 116 | + ptrdiff_t dst, ptrdiff_t sst, |
| 117 | + size_t nelems, |
| 118 | + size_t element_size, |
| 119 | + long *pSync, |
| 120 | + int alg); |
| 121 | + |
| 122 | +END_C_DECLS |
| 123 | + |
| 124 | +#endif |
0 commit comments