| 168 | //------------------------------------------------------------------------------ |
| 169 | |
| 170 | int rmm_wrap_initialize // returns -1 on error, 0 on success |
| 171 | ( |
| 172 | RMM_MODE mode, // TODO: describe |
| 173 | std::size_t init_pool_size, // TODO: describe |
| 174 | std::size_t max_pool_size, // TODO: describe |
| 175 | std::size_t stream_pool_size |
| 176 | ) |
| 177 | { |
| 178 | |
| 179 | //-------------------------------------------------------------------------- |
| 180 | // check inputs |
| 181 | //-------------------------------------------------------------------------- |
| 182 | |
| 183 | if (rmm_wrap_context != NULL) |
| 184 | { |
| 185 | // rmm_wrap_initialize cannot be called twice |
| 186 | return (-1) ; |
| 187 | } |
| 188 | |
| 189 | if(stream_pool_size <= 0) |
| 190 | { |
| 191 | std::cout << "Stream pool size must be >=0" << std::endl; |
| 192 | // failed to create the alloc_map |
| 193 | return (-1) ; |
| 194 | } |
| 195 | |
| 196 | // create the RMM wrap handle and save it as a global pointer. |
| 197 | rmm_wrap_context = new RMM_Wrap_Handle() ; |
| 198 | |
| 199 | // std::cout<< " init called with mode "<<mode<<" init_size " |
| 200 | // <<init_pool_size<<" max_size "<<max_pool_size<<"\n"; |
| 201 | |
| 202 | //-------------------------------------------------------------------------- |
| 203 | // Construct a resource that uses a coalescing best-fit pool allocator |
| 204 | //-------------------------------------------------------------------------- |
| 205 | |
| 206 | // Set CUDA stream pool |
| 207 | rmm_wrap_context->stream_pool = make_and_set_cuda_stream_pool(stream_pool_size); |
| 208 | rmm_wrap_context->main_stream = make_cuda_stream(); |
| 209 | |
| 210 | if (mode == rmm_wrap_host ) |
| 211 | { |
| 212 | // rmm_wrap_context->host_resource = |
| 213 | // std::pmr::synchronized_pool_resource() ; |
| 214 | // // (init_pool_size, max_pool_size) ; |
| 215 | // rmm_wrap_context->host_resource = make_and_set_host_pool() ; |
| 216 | // // (init_pool_size, max_pool_size) ; |
| 217 | } |
| 218 | else if (mode == rmm_wrap_host_pinned ) |
| 219 | { |
| 220 | // rmm_wrap_context->host_resource = |
| 221 | // std::pmr::synchronized_pool_resource() ; |
| 222 | // // (init_pool_size, max_pool_size) ; |
| 223 | } |
| 224 | else if (mode == rmm_wrap_device ) |
| 225 | { |
| 226 | rmm_wrap_context->resource = |
| 227 | make_and_set_device_pool( init_pool_size, max_pool_size) ; |