Merge remote-tracking branch 'upstream/branch-24.02' into cccl-2.2.0

.github/CODEOWNERS

@@ -14,5 +14,4 @@ python/            @rapidsai/rmm-python-codeowners
 .github/           @rapidsai/ops-codeowners
 ci/                @rapidsai/ops-codeowners
 conda/             @rapidsai/ops-codeowners
-**/Dockerfile      @rapidsai/ops-codeowners
-**/.dockerignore   @rapidsai/ops-codeowners
+dependencies.yaml  @rapidsai/ops-codeowners

.gitignore

@@ -22,6 +22,7 @@ rmm.egg-info/
 python/build
 python/*/build
 python/rmm/**/_lib/**/*.cpp
+!python/rmm/_lib/_torch_allocator.cpp
 python/rmm/**/_lib/**/*.h
 python/rmm/**/_lib/.nfs*
 python/rmm/_cuda/*.cpp

README.md

@@ -67,7 +67,7 @@ CUDA/GPU requirements:
 You can obtain CUDA from [https://developer.nvidia.com/cuda-downloads](https://developer.nvidia.com/cuda-downloads)
 
 Python requirements:
-* `scikit-build`
+* `scikit-build-core`
 * `cuda-python`
 * `cython`
 

build.sh

@@ -176,5 +176,5 @@ fi
 # Build and install the rmm Python package
 if (( NUMARGS == 0 )) || hasArg rmm; then
     echo "building and installing rmm..."
-    SKBUILD_CONFIGURE_OPTIONS="${SKBUILD_EXTRA_CMAKE_ARGS}" python -m pip install --no-build-isolation --no-deps ${REPODIR}/python
+    SKBUILD_CMAKE_ARGS="${SKBUILD_EXTRA_CMAKE_ARGS}" python -m pip install --no-build-isolation --no-deps ${REPODIR}/python
 fi

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -33,12 +33,11 @@ dependencies:
 - pytest
 - pytest-cov
 - python>=3.9,<3.11
-- scikit-build>=0.13.1
+- scikit-build-core>=0.7.0
 - spdlog>=1.12.0,<1.13
 - sphinx
 - sphinx-copybutton
 - sphinx-markdown-tables
 - sphinx_rtd_theme
 - sysroot_linux-64==2.17
-- tomli
 name: all_cuda-118_arch-x86_64

conda/environments/all_cuda-120_arch-x86_64.yaml

@@ -31,12 +31,11 @@ dependencies:
 - pytest
 - pytest-cov
 - python>=3.9,<3.11
-- scikit-build>=0.13.1
+- scikit-build-core>=0.7.0
 - spdlog>=1.12.0,<1.13
 - sphinx
 - sphinx-copybutton
 - sphinx-markdown-tables
 - sphinx_rtd_theme
 - sysroot_linux-64==2.17
-- tomli
 name: all_cuda-120_arch-x86_64

conda/recipes/rmm/meta.yaml

@@ -61,9 +61,7 @@ requirements:
     - cython >=3.0.0
     - librmm ={{ version }}
     - python
-    - scikit-build >=0.13.1
-    - setuptools >=61.0.0
-    - tomli  # [py<311]
+    - scikit-build-core >=0.7.0
   run:
     {% if cuda_major == "11" %}
     - cudatoolkit

dependencies.yaml

@@ -11,6 +11,7 @@ files:
       - cudatoolkit
       - develop
       - docs
+      - py_version
       - run
       - test_python
   test_python:
@@ -65,19 +66,16 @@ dependencies:
           - &cmake_ver cmake>=3.26.4
           - cython>=3.0.0
           - ninja
-          - scikit-build>=0.13.1
-          - tomli
       - output_types: conda
         packages:
           - c-compiler
           - cxx-compiler
           - fmt>=10.1.1,<11
+          - scikit-build-core>=0.7.0
           - spdlog>=1.12.0,<1.13
-          - python>=3.9,<3.11
-      - output_types: pyproject
+      - output_types: [requirements, pyproject]
         packages:
-          - wheel
-          - setuptools>=61.0.0
+          - scikit-build-core[pyproject]>=0.7.0
     specific:
       - output_types: conda
         matrices:
@@ -197,6 +195,9 @@ dependencies:
               py: "3.10"
             packages:
               - python=3.10
+          - matrix:
+            packages:
+              - python>=3.9,<3.11
   run:
     common:
       - output_types: [conda, requirements, pyproject]

python/CMakeLists.txt

@@ -21,11 +21,7 @@ include(../fetch_rapids.cmake)
 project(
   rmm-python
   VERSION ${rmm_version}
-  LANGUAGES # TODO: Building Python extension modules via the python_extension_module requires the C
-            # language to be enabled here. The test project that is built in scikit-build to verify
-            # various linking options for the python library is hardcoded to build with C, so until
-            # that is fixed we need to keep C.
-            C CXX)
+  LANGUAGES CXX)
 
 option(FIND_RMM_CPP "Search for existing RMM C++ installations before defaulting to local files"
        OFF)
@@ -45,7 +41,7 @@ if(NOT rmm_FOUND)
   add_subdirectory(../ rmm-cpp EXCLUDE_FROM_ALL)
 endif()
 
-include(rapids-cython)
+include(rapids-cython-core)
 rapids_cython_init()
 
 add_compile_definitions("SPDLOG_ACTIVE_LEVEL=SPDLOG_LEVEL_${RMM_LOGGING_LEVEL}")

python/pyproject.toml

@@ -13,16 +13,13 @@
 # limitations under the License.
 
 [build-system]
-build-backend = "setuptools.build_meta"
+build-backend = "scikit_build_core.build"
 requires = [
     "cmake>=3.26.4",
     "cuda-python>=11.7.1,<12.0a0",
     "cython>=3.0.0",
     "ninja",
-    "scikit-build>=0.13.1",
-    "setuptools>=61.0.0",
-    "tomli",
-    "wheel",
+    "scikit-build-core[pyproject]>=0.7.0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../dependencies.yaml and run `rapids-dependency-file-generator`.
 
 [project]
@@ -112,8 +109,15 @@ skip = [
     "__init__.py",
 ]
 
-[tool.setuptools]
-license-files = ["LICENSE"]
+[tool.scikit-build]
+build-dir = "build/{wheel_tag}"
+cmake.build-type = "Release"
+cmake.minimum-version = "3.26.4"
+ninja.make-fallback = true
+sdist.reproducible = true
+wheel.packages = ["rmm"]
 
-[tool.setuptools.dynamic]
-version = {file = "rmm/VERSION"}
+[tool.scikit-build.metadata.version]
+provider = "scikit_build_core.metadata.regex"
+input = "rmm/VERSION"
+regex = "(?P<value>.*)"

python/rmm/_lib/CMakeLists.txt

@@ -12,12 +12,17 @@
 # the License.
 # =============================================================================
 
-set(cython_sources device_buffer.pyx lib.pyx logger.pyx memory_resource.pyx cuda_stream.pyx
-                   torch_allocator.pyx)
+set(cython_sources device_buffer.pyx lib.pyx logger.pyx memory_resource.pyx cuda_stream.pyx)
 set(linked_libraries rmm::rmm)
 
 # Build all of the Cython targets
 rapids_cython_create_modules(SOURCE_FILES "${cython_sources}" LINKED_LIBRARIES "${linked_libraries}"
                                                                                CXX)
-# The cdef public functions in this file need to have a C ABI
-target_compile_definitions(torch_allocator PRIVATE CYTHON_EXTERN_C=extern\ "C")
+
+add_library(_torch_allocator SHARED _torch_allocator.cpp)
+# Want the output to be called _torch_allocator.so
+set_target_properties(_torch_allocator PROPERTIES PREFIX "" SUFFIX ".so")
+target_link_libraries(_torch_allocator PRIVATE rmm::rmm)
+cmake_path(RELATIVE_PATH CMAKE_CURRENT_SOURCE_DIR BASE_DIRECTORY "${PROJECT_SOURCE_DIR}"
+           OUTPUT_VARIABLE _torch_allocator_location)
+install(TARGETS _torch_allocator DESTINATION "${_torch_allocator_location}")

python/rmm/_lib/_torch_allocator.cpp

@@ -0,0 +1,60 @@
+/*
+ * Copyright (c) 2023, 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.
+ */
+
+#include <cuda_runtime_api.h>
+
+#include <rmm/cuda_device.hpp>
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/mr/device/per_device_resource.hpp>
+
+// These signatures must match those required by CUDAPluggableAllocator in
+// github.com/pytorch/pytorch/blob/main/torch/csrc/cuda/CUDAPluggableAllocator.h
+// Since the loading is done at runtime via dlopen, no error checking
+// can be performed for mismatching signatures.
+
+/**
+ * @brief Allocate memory of at least \p size bytes.
+ *
+ * @throws rmm::bad_alloc When the requested allocation cannot be satisfied.
+ *
+ * @param size The number of bytes to allocate
+ * @param device The device whose memory resource one should use
+ * @param stream CUDA stream to perform allocation on
+ * @return Pointer to the newly allocated memory
+ */
+extern "C" void* allocate(std::size_t size, int device, void* stream)
+{
+  rmm::cuda_device_id const device_id{device};
+  rmm::cuda_set_device_raii with_device{device_id};
+  auto mr = rmm::mr::get_per_device_resource(device_id);
+  return mr->allocate(size, rmm::cuda_stream_view{static_cast<cudaStream_t>(stream)});
+}
+
+/**
+ * @brief Deallocate memory pointed to by \p ptr.
+ *
+ * @param ptr Pointer to be deallocated
+ * @param size The number of bytes in the allocation
+ * @param device The device whose memory resource one should use
+ * @param stream CUDA stream to perform deallocation on
+ */
+extern "C" void deallocate(void* ptr, std::size_t size, int device, void* stream)
+{
+  rmm::cuda_device_id const device_id{device};
+  rmm::cuda_set_device_raii with_device{device_id};
+  auto mr = rmm::mr::get_per_device_resource(device_id);
+  mr->deallocate(ptr, size, rmm::cuda_stream_view{static_cast<cudaStream_t>(stream)});
+}

python/rmm/allocators/torch.py

@@ -16,11 +16,14 @@
 except ImportError:
     rmm_torch_allocator = None
 else:
-    import rmm._lib.torch_allocator
+    import pathlib
 
-    _alloc_free_lib_path = rmm._lib.torch_allocator.__file__
+    sofile = (
+        pathlib.Path(__file__).parent.parent / "_lib" / "_torch_allocator.so"
+    )
     rmm_torch_allocator = CUDAPluggableAllocator(
-        _alloc_free_lib_path,
+        str(sofile.absolute()),
         alloc_fn_name="allocate",
         free_fn_name="deallocate",
     )
+    del pathlib, sofile