From 0d0ea02d1ad98150606d710586841498d4017fc9 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Tue, 24 Sep 2024 12:47:54 -0700 Subject: [PATCH] address reviews --- python/rmm/rmm/_lib/cuda_stream.pxd | 1 + python/rmm/rmm/_lib/cuda_stream_pool.pxd | 15 +++++ python/rmm/rmm/_lib/cuda_stream_view.pxd | 15 +++++ python/rmm/rmm/_lib/device_buffer.pxd | 6 ++ python/rmm/rmm/_lib/device_uvector.pxd | 15 +++++ python/rmm/rmm/_lib/logger.pxd | 75 +++++++++++++++++++++ python/rmm/rmm/_lib/memory_resource.pxd | 27 ++++++++ python/rmm/rmm/_lib/per_device_resource.pxd | 21 ++++++ python/rmm/rmm/allocators/torch.py | 2 +- python/rmm/rmm/librmm/_logger.pyx | 2 +- python/rmm/rmm/librmm/memory_resource.pxd | 2 - 11 files changed, 177 insertions(+), 4 deletions(-) create mode 100644 python/rmm/rmm/_lib/cuda_stream_pool.pxd create mode 100644 python/rmm/rmm/_lib/cuda_stream_view.pxd create mode 100644 python/rmm/rmm/_lib/device_uvector.pxd create mode 100644 python/rmm/rmm/_lib/logger.pxd create mode 100644 python/rmm/rmm/_lib/per_device_resource.pxd diff --git a/python/rmm/rmm/_lib/cuda_stream.pxd b/python/rmm/rmm/_lib/cuda_stream.pxd index f7b5759bd..afc365fbb 100644 --- a/python/rmm/rmm/_lib/cuda_stream.pxd +++ b/python/rmm/rmm/_lib/cuda_stream.pxd @@ -12,4 +12,5 @@ # See the License for the specific language governing permissions and # limitations under the License. +from rmm.librmm.cuda_stream cimport cuda_stream from rmm.pylibrmm.cuda_stream cimport CudaStream diff --git a/python/rmm/rmm/_lib/cuda_stream_pool.pxd b/python/rmm/rmm/_lib/cuda_stream_pool.pxd new file mode 100644 index 000000000..4da59cc68 --- /dev/null +++ b/python/rmm/rmm/_lib/cuda_stream_pool.pxd @@ -0,0 +1,15 @@ +# Copyright (c) 2021-2024, 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. + +from rmm.librmm.cuda_stream_pool cimport cuda_stream_pool diff --git a/python/rmm/rmm/_lib/cuda_stream_view.pxd b/python/rmm/rmm/_lib/cuda_stream_view.pxd new file mode 100644 index 000000000..78480420f --- /dev/null +++ b/python/rmm/rmm/_lib/cuda_stream_view.pxd @@ -0,0 +1,15 @@ +# Copyright (c) 2020-2024, 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. + +from rmm.librmm.cuda_stream_view cimport cuda_stream_view diff --git a/python/rmm/rmm/_lib/device_buffer.pxd b/python/rmm/rmm/_lib/device_buffer.pxd index 66c19d990..22833b1b8 100644 --- a/python/rmm/rmm/_lib/device_buffer.pxd +++ b/python/rmm/rmm/_lib/device_buffer.pxd @@ -12,6 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from rmm.librmm.device_buffer cimport ( + cuda_device_id, + device_buffer, + get_current_cuda_device, + prefetch, +) from rmm.pylibrmm.device_buffer cimport ( DeviceBuffer, copy_device_to_ptr, diff --git a/python/rmm/rmm/_lib/device_uvector.pxd b/python/rmm/rmm/_lib/device_uvector.pxd new file mode 100644 index 000000000..230b0afb3 --- /dev/null +++ b/python/rmm/rmm/_lib/device_uvector.pxd @@ -0,0 +1,15 @@ +# Copyright (c) 2021-2024, 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. + +from rmm.librmm.device_uvector cimport device_uvector diff --git a/python/rmm/rmm/_lib/logger.pxd b/python/rmm/rmm/_lib/logger.pxd new file mode 100644 index 000000000..27936d09c --- /dev/null +++ b/python/rmm/rmm/_lib/logger.pxd @@ -0,0 +1,75 @@ +# Copyright (c) 2023-2024, 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. + +from rmm.librmm.logger cimport logger, logging_level, spdlog_logger +from rmm.pylibrmm.logger cimport ( + _validate_level_type, + flush_logger, + get_flush_level, + get_logging_level, + set_flush_level, + set_logging_level, + should_log, +) + + +cdef extern from "spdlog/common.h" namespace "spdlog::level" nogil: + cpdef enum logging_level "spdlog::level::level_enum": + """ + The debug logging level for RMM. + + Debug logging prints messages to a log file. See + `Debug Logging `_ + for more information. + + Valid levels, in decreasing order of verbosity, are TRACE, DEBUG, + INFO, WARN, ERR, CRITICAL, and OFF. Default is INFO. + + Examples + -------- + >>> import rmm + >>> rmm.logging_level.DEBUG + + >>> rmm.logging_level.DEBUG.value + 1 + >>> rmm.logging_level.DEBUG.name + 'DEBUG' + + See Also + -------- + set_logging_level : Set the debug logging level + get_logging_level : Get the current debug logging level + """ + TRACE "spdlog::level::trace" + DEBUG "spdlog::level::debug" + INFO "spdlog::level::info" + WARN "spdlog::level::warn" + ERR "spdlog::level::err" + CRITICAL "spdlog::level::critical" + OFF "spdlog::level::off" + + +cdef extern from "spdlog/spdlog.h" namespace "spdlog" nogil: + cdef cppclass spdlog_logger "spdlog::logger": + spdlog_logger() except + + void set_level(logging_level level) + logging_level level() + void flush() except + + void flush_on(logging_level level) + logging_level flush_level() + bool should_log(logging_level msg_level) + + +cdef extern from "rmm/logger.hpp" namespace "rmm" nogil: + cdef spdlog_logger& logger() except + diff --git a/python/rmm/rmm/_lib/memory_resource.pxd b/python/rmm/rmm/_lib/memory_resource.pxd index 79a28cac2..983063914 100644 --- a/python/rmm/rmm/_lib/memory_resource.pxd +++ b/python/rmm/rmm/_lib/memory_resource.pxd @@ -12,6 +12,33 @@ # See the License for the specific language governing permissions and # limitations under the License. +from rmm.librmm.memory_resource cimport ( + CppExcept, + allocate_callback_t, + allocation_handle_type, + available_device_memory, + binning_memory_resource, + callback_memory_resource, + cuda_async_memory_resource, + cuda_memory_resource, + deallocate_callback_t, + device_memory_resource, + failure_callback_resource_adaptor, + failure_callback_t, + fixed_size_memory_resource, + limiting_resource_adaptor, + logging_resource_adaptor, + managed_memory_resource, + percent_of_free_device_memory, + pool_memory_resource, + prefetch_resource_adaptor, + sam_headroom_memory_resource, + statistics_resource_adaptor, + system_memory_resource, + throw_cpp_except, + tracking_resource_adaptor, + translate_python_except_to_cpp, +) from rmm.pylibrmm.memory_resource cimport ( BinningMemoryResource, CallbackMemoryResource, diff --git a/python/rmm/rmm/_lib/per_device_resource.pxd b/python/rmm/rmm/_lib/per_device_resource.pxd new file mode 100644 index 000000000..29487f503 --- /dev/null +++ b/python/rmm/rmm/_lib/per_device_resource.pxd @@ -0,0 +1,21 @@ +# Copyright (c) 2019-2024, 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. + +from rmm.librmm.per_device_resource cimport ( + cuda_device_id, + get_current_device_resource, + get_per_device_resource, + set_current_device_resource, + set_per_device_resource, +) diff --git a/python/rmm/rmm/allocators/torch.py b/python/rmm/rmm/allocators/torch.py index 462a06e80..eee0e9df9 100644 --- a/python/rmm/rmm/allocators/torch.py +++ b/python/rmm/rmm/allocators/torch.py @@ -28,7 +28,7 @@ # allocator .so relative to the current file because the current file # is pure Python and will therefore be in the source directory. # Instead, we search relative to an arbitrary file in the compiled - # package. We use the cpp._logger module because it is small. + # package. We use the librmm._logger module because it is small. from rmm.librmm import _logger sofile = pathlib.Path(_logger.__file__).parent / "_torch_allocator.so" diff --git a/python/rmm/rmm/librmm/_logger.pyx b/python/rmm/rmm/librmm/_logger.pyx index e848c6dda..4392cb106 100644 --- a/python/rmm/rmm/librmm/_logger.pyx +++ b/python/rmm/rmm/librmm/_logger.pyx @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ._logger cimport logging_level # no-cython-lint +from rmm.librmm._logger cimport logging_level # no-cython-lint diff --git a/python/rmm/rmm/librmm/memory_resource.pxd b/python/rmm/rmm/librmm/memory_resource.pxd index a8e932e6c..9ddaf04b9 100644 --- a/python/rmm/rmm/librmm/memory_resource.pxd +++ b/python/rmm/rmm/librmm/memory_resource.pxd @@ -79,8 +79,6 @@ cdef extern from *: void throw_cpp_except(CppExcept) nogil -# NOTE: Keep extern declarations in .pyx file as much as possible to avoid -# leaking dependencies when importing RMM Cython .pxd files cdef extern from "rmm/mr/device/cuda_memory_resource.hpp" \ namespace "rmm::mr" nogil: cdef cppclass cuda_memory_resource(device_memory_resource):