ucx1-sys 0.1.0

Rust FFI bindings to UCX.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

/**
 * Copyright (C) Mellanox Technologies Ltd. 2018.  ALL RIGHTS RESERVED.
 * See file LICENSE for terms.
 */

#ifndef UCT_CUDA_IFACE_H
#define UCT_CUDA_IFACE_H

#include <uct/base/uct_iface.h>
#include <ucs/sys/preprocessor.h>
#include <ucs/profile/profile.h>
#include <cuda_runtime.h>
#include <cuda.h>
#include <nvml.h>


#define UCT_CUDA_DEV_NAME       "cuda"

#define UCT_CUDAR_CALL(_log_level, _func, ...) \
    ({ \
        cudaError_t _result = UCS_PROFILE_CALL(_func, __VA_ARGS__); \
        ucs_status_t _status; \
        \
        if (cudaSuccess != _result) { \
            ucs_log((_log_level), "%s() failed: %s", \
                    UCS_PP_MAKE_STRING(_func), cudaGetErrorString(_result)); \
            _status = UCS_ERR_IO_ERROR; \
        } else { \
            _status = UCS_OK; \
        } \
        _status; \
    })


#define UCT_CUDAR_CALL_LOG_ERR(_func, ...) \
    UCT_CUDAR_CALL(UCS_LOG_LEVEL_ERROR, _func, __VA_ARGS__)


#define UCT_CUDA_FUNC(_func, _log_level)                        \
    ({                                                          \
        ucs_status_t _status = UCS_OK;                          \
        do {                                                    \
            cudaError_t _result = (_func);                      \
            if (cudaSuccess != _result) {                       \
                ucs_log((_log_level), "%s() failed: %s",        \
                        UCS_PP_MAKE_STRING(_func),              \
                        cudaGetErrorString(_result));           \
                _status = UCS_ERR_IO_ERROR;                     \
            }                                                   \
        } while (0);                                            \
        _status;                                                \
    })


#define UCT_CUDA_FUNC_LOG_ERR(_func) \
    UCT_CUDA_FUNC(_func, UCS_LOG_LEVEL_ERROR)


#define UCT_NVML_FUNC(_func, _log_level)                        \
    ({                                                          \
        ucs_status_t _status = UCS_OK;                          \
        do {                                                    \
            nvmlReturn_t _err = (_func);                        \
            if (NVML_SUCCESS != _err) {                         \
                ucs_log((_log_level), "%s failed: %s",          \
                        UCS_PP_MAKE_STRING(_func),              \
                        nvmlErrorString(_err));                 \
                _status = UCS_ERR_IO_ERROR;                     \
            }                                                   \
        } while (0);                                            \
        _status;                                                \
    })


#define UCT_NVML_FUNC_LOG_ERR(_func) \
    UCT_NVML_FUNC(_func, UCS_LOG_LEVEL_ERROR)


#define UCT_CUDADRV_FUNC(_func, _log_level)                     \
    ({                                                          \
        ucs_status_t _status = UCS_OK;                          \
        do {                                                    \
            CUresult _result = (_func);                         \
            const char *cu_err_str;                             \
            if (CUDA_ERROR_NOT_READY == _result) {              \
                _status = UCS_INPROGRESS;                       \
            } else if (CUDA_SUCCESS != _result) {               \
                cuGetErrorString(_result, &cu_err_str);         \
                ucs_log((_log_level), "%s() failed: %s",        \
                        UCS_PP_MAKE_STRING(_func), cu_err_str); \
                _status = UCS_ERR_IO_ERROR;                     \
            }                                                   \
        } while (0);                                            \
        _status;                                                \
    })


#define UCT_CUDADRV_FUNC_LOG_ERR(_func) \
    UCT_CUDADRV_FUNC(_func, UCS_LOG_LEVEL_ERROR)


#define UCT_CUDADRV_CTX_ACTIVE(_state) \
    { \
        CUdevice _dev; \
        CUcontext _ctx; \
        int _flags; \
        if (CUDA_SUCCESS == cuCtxGetDevice(&_dev)) { \
            cuDevicePrimaryCtxGetState(_dev, &_flags, &_state); \
            if (_state == 0) { \
                /* need to retain for malloc purposes */ \
                if (CUDA_SUCCESS != cuDevicePrimaryCtxRetain(&_ctx, _dev)) { \
                    ucs_fatal("unable to retain ctx after detecting device"); \
                } \
            } \
            _state = 1; \
        } else { \
            _state = 0; \
        } \
    }


typedef enum uct_cuda_base_gen {
    UCT_CUDA_BASE_GEN_P100 = 6,
    UCT_CUDA_BASE_GEN_V100 = 7,
    UCT_CUDA_BASE_GEN_A100 = 8
} uct_cuda_base_gen_t;


ucs_status_t
uct_cuda_base_query_devices_common(
        uct_md_h md, uct_device_type_t dev_type,
        uct_tl_device_resource_t **tl_devices_p, unsigned *num_tl_devices_p);

ucs_status_t
uct_cuda_base_query_devices(
        uct_md_h md, uct_tl_device_resource_t **tl_devices_p,
        unsigned *num_tl_devices_p);

ucs_status_t
uct_cuda_base_get_sys_dev(CUdevice cuda_device,
                          ucs_sys_device_t *sys_dev_p);

#endif