diff --git a/include/rmm/cuda_stream_view.hpp b/include/rmm/cuda_stream_view.hpp index 1641a80ec..c80d4de2f 100644 --- a/include/rmm/cuda_stream_view.hpp +++ b/include/rmm/cuda_stream_view.hpp @@ -57,9 +57,9 @@ class cuda_stream_view { constexpr cudaStream_t value() const noexcept { return stream_; } /** - * @brief Explicit conversion to cudaStream_t. + * @brief Implicit conversion to cudaStream_t. */ - explicit constexpr operator cudaStream_t() const noexcept { return value(); } + constexpr operator cudaStream_t() const noexcept { return value(); } /** * @brief Return true if the wrapped stream is the CUDA per-thread default stream. diff --git a/python/rmm/_lib/cuda_stream_pool.pxd b/python/rmm/_lib/cuda_stream_pool.pxd new file mode 100644 index 000000000..e2db041d3 --- /dev/null +++ b/python/rmm/_lib/cuda_stream_pool.pxd @@ -0,0 +1,24 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +cimport cython + +from rmm._lib.cuda_stream_view cimport cuda_stream_view + +cdef extern from "rmm/cuda_stream_pool.hpp" namespace "rmm" nogil: + cdef cppclass cuda_stream_pool: + cuda_stream_pool(size_t pool_size) + cuda_stream_view get_stream() + cuda_stream_view get_stream(size_t stream_id) except + + size_t get_pool_size()