[DO NOT MERGE] Experiment with using cuda::memory_resource interface

CMakeLists.txt

@@ -58,13 +58,16 @@ rapids_find_package(
 rapids_cpm_init()
 include(cmake/thirdparty/get_spdlog.cmake)
 include(cmake/thirdparty/get_thrust.cmake)
+include(cmake/thirdparty/get_libcudacxx.cmake)
 
 # library targets
 add_library(rmm INTERFACE)
 add_library(rmm::rmm ALIAS rmm)
 
-target_include_directories(rmm INTERFACE "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>"
-                                         "$<INSTALL_INTERFACE:include>")
+target_include_directories(
+  rmm
+  INTERFACE "$<BUILD_INTERFACE:${LIBCUDACXX_INCLUDE_DIR}>"
+            "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>" "$<INSTALL_INTERFACE:include>")
 
 if(CUDA_STATIC_RUNTIME)
   message(STATUS "RMM: Enabling static linking of cudart")
@@ -108,6 +111,10 @@ include(CPack)
 # install export targets
 install(TARGETS rmm EXPORT rmm-exports)
 install(DIRECTORY include/rmm/ DESTINATION include/rmm)
+install(
+  DIRECTORY ${RMM_GENERATED_INCLUDE_DIR}/include/libcxx
+            ${RMM_GENERATED_INCLUDE_DIR}/include/libcudacxx
+  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/rmm)
 install(FILES ${RMM_BINARY_DIR}/include/rmm/version_config.hpp DESTINATION include/rmm)
 
 set(doc_string

cmake/thirdparty/get_libcudacxx.cmake

@@ -0,0 +1,33 @@
+# =============================================================================
+# Copyright (c) 2020, 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.
+# =============================================================================
+
+# Use CPM to find or clone libcudacxx
+function(find_and_configure_libcudacxx VERSION)
+  rapids_cpm_find(
+    libcudacxx ${VERSION}
+    GIT_REPOSITORY https://github.com/mzient/libcudacxx.git
+    GIT_TAG memres_view # ${VERSION}
+    GIT_SHALLOW TRUE DOWNLOAD_ONLY TRUE)
+
+  set(LIBCUDACXX_INCLUDE_DIR
+      "${libcudacxx_SOURCE_DIR}/include"
+      PARENT_SCOPE)
+  set(LIBCXX_INCLUDE_DIR
+      "${libcudacxx_SOURCE_DIR}/libcxx/include"
+      PARENT_SCOPE)
+endfunction()
+
+set(RMM_MIN_VERSION_libcudacxx 1.5.0)
+
+find_and_configure_libcudacxx(${RMM_MIN_VERSION_libcudacxx})

include/rmm/cuda_stream_view.hpp

@@ -20,6 +20,8 @@
 
 #include <cuda_runtime_api.h>
 
+#include <cuda/stream_view>
+
 #include <atomic>
 #include <cstddef>
 #include <cstdint>
@@ -61,6 +63,11 @@ class cuda_stream_view {
    */
   constexpr operator cudaStream_t() const noexcept { return value(); }
 
+  /**
+   * @brief Implicit conversion to cuda::stream_view.
+   */
+  constexpr operator cuda::stream_view() const noexcept { return value(); }
+
   /**
    * @brief Return true if the wrapped stream is the CUDA per-thread default stream.
    */

include/rmm/mr/libcudacxx/device/cuda_memory_resource.hpp

@@ -0,0 +1,110 @@
+/*
+ * Copyright (c) 2019, 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.
+ */
+#pragma once
+
+#include <rmm/detail/aligned.hpp>
+#include <rmm/detail/error.hpp>
+
+#include <cuda/memory_resource>
+
+#include <cassert>
+
+namespace rmm {
+namespace mr {
+
+/**
+ * @brief `device_memory_resource` derived class that uses cudaMalloc/Free for
+ * allocation/deallocation.
+ */
+class cuda_memory_resource final
+  : public cuda::stream_ordered_memory_resource<cuda::memory_kind::device> {
+ public:
+  cuda_memory_resource()                            = default;
+  ~cuda_memory_resource()                           = default;
+  cuda_memory_resource(cuda_memory_resource const&) = default;
+  cuda_memory_resource(cuda_memory_resource&&)      = default;
+  cuda_memory_resource& operator=(cuda_memory_resource const&) = default;
+  cuda_memory_resource& operator=(cuda_memory_resource&&) = default;
+
+ private:
+  /**
+   * @brief Allocates memory of size at least `bytes` using cudaMalloc.
+   *
+   * The returned pointer has at least 256B alignment. Alignments greater than
+   * this are not supported.
+   *
+   * @throws `rmm::bad_alloc` if the requested allocation could not be fulfilled
+   * @throws `rmm::logic_error` if the requested alignment is greater than 256B
+   *
+   * @param bytes The size, in bytes, of the allocation
+   * @return void* Pointer to the newly allocated memory
+   */
+  void* do_allocate(std::size_t bytes, std::size_t alignment) override
+  {
+    RMM_EXPECTS(detail::is_aligned(detail::CUDA_ALLOCATION_ALIGNMENT, alignment),
+                "Unsupported alignment");
+    void* p{nullptr};
+    RMM_CUDA_TRY(cudaMalloc(&p, bytes), rmm::bad_alloc);
+    return p;
+  }
+
+  /**
+   * @brief Deallocate memory pointed to by \p p.
+   *
+   * Alignments greater than 256B are not supported and behavior is undefined.
+   *
+   * @throws Nothing.
+   *
+   * @param p Pointer to be deallocated
+   */
+  void do_deallocate(void* p, std::size_t, std::size_t alignment) override
+  {
+    assert(detail::is_aligned(detail::CUDA_ALLOCATION_ALIGNMENT, alignment));
+    RMM_ASSERT_CUDA_SUCCESS(cudaFree(p));
+  }
+
+  void* do_allocate_async(std::size_t bytes, std::size_t alignment, cuda::stream_view) override
+  {
+    return do_allocate(bytes, alignment);
+  }
+
+  void do_deallocate_async(void* p,
+                           std::size_t bytes,
+                           std::size_t alignment,
+                           cuda::stream_view) override
+  {
+    return do_deallocate(p, bytes, alignment);
+  };
+
+  /**
+   * @brief Compare this resource to another.
+   *
+   * Two cuda_memory_resources always compare equal, because they can each
+   * deallocate memory allocated by the other.
+   *
+   * @throws Nothing.
+   *
+   * @param other The other resource to compare to
+   * @return true If the two resources are equivalent
+   * @return false If the two resources are not equal
+   */
+  bool do_is_equal(memory_resource const& other) const noexcept override
+  {
+    return dynamic_cast<cuda_memory_resource const*>(&other) != nullptr;
+  }
+};
+}  // namespace mr
+}  // namespace rmm