Provide memory_type enum

cpp/include/raft/core/device_mdspan.hpp

@@ -19,14 +19,15 @@
 #include <cstdint>
 #include <raft/core/host_device_accessor.hpp>
 #include <raft/core/mdspan.hpp>
+#include <raft/core/memory_type.hpp>
 
 namespace raft {
 
 template <typename AccessorPolicy>
-using device_accessor = host_device_accessor<AccessorPolicy, false, true>;
+using device_accessor = host_device_accessor<AccessorPolicy, memory_type::device>;
 
 template <typename AccessorPolicy>
-using managed_accessor = host_device_accessor<AccessorPolicy, true, true>;
+using managed_accessor = host_device_accessor<AccessorPolicy, memory_type::managed>;
 
 /**
  * @brief std::experimental::mdspan with device tag to avoid accessing incorrect memory location.
@@ -276,4 +277,4 @@ auto make_device_vector_view(ElementType* ptr, IndexType n)
   return device_vector_view<ElementType, IndexType, LayoutPolicy>{ptr, n};
 }
 
-}  // end namespace raft
+}  // end namespace raft

cpp/include/raft/core/host_device_accessor.hpp

@@ -15,6 +15,8 @@
  */
 
 #pragma once
+#include <raft/core/memory_type.hpp>
+#include <type_traits>
 
 namespace raft {
 
@@ -23,15 +25,19 @@ namespace raft {
  * accessor used throught RAFT's APIs to denote whether an underlying pointer
  * is accessible from device, host, or both.
  */
-template <typename AccessorPolicy, bool is_host, bool is_device>
+template <typename AccessorPolicy, memory_type MemType>
 struct host_device_accessor : public AccessorPolicy {
-  using accessor_type   = AccessorPolicy;
-  using is_host_type    = std::conditional_t<is_host, std::true_type, std::false_type>;
-  using is_device_type  = std::conditional_t<is_device, std::true_type, std::false_type>;
-  using is_managed_type = std::conditional_t<is_device && is_host, std::true_type, std::false_type>;
-  static constexpr bool is_host_accessible    = is_host;
-  static constexpr bool is_device_accessible  = is_device;
-  static constexpr bool is_managed_accessible = is_device && is_host;
+  using accessor_type                  = AccessorPolicy;
+  auto static constexpr const mem_type = MemType;
+  using is_host_type =
+    std::conditional_t<raft::is_host_accessible(mem_type), std::true_type, std::false_type>;
+  using is_device_type =
+    std::conditional_t<raft::is_device_accessible(mem_type), std::true_type, std::false_type>;
+  using is_managed_type =
+    std::conditional_t<raft::is_host_device_accessible(mem_type), std::true_type, std::false_type>;
+  static constexpr bool is_host_accessible    = raft::is_host_accessible(mem_type);
+  static constexpr bool is_device_accessible  = raft::is_device_accessible(mem_type);
+  static constexpr bool is_managed_accessible = raft::is_host_device_accessible(mem_type);
   // make sure the explicit ctor can fall through
   using AccessorPolicy::AccessorPolicy;
   using offset_policy = host_device_accessor;

cpp/include/raft/core/host_mdspan.hpp

@@ -18,13 +18,14 @@
 
 #include <cstdint>
 #include <raft/core/mdspan.hpp>
+#include <raft/core/memory_type.hpp>
 
 #include <raft/core/host_device_accessor.hpp>
 
 namespace raft {
 
 template <typename AccessorPolicy>
-using host_accessor = host_device_accessor<AccessorPolicy, true, false>;
+using host_accessor = host_device_accessor<AccessorPolicy, memory_type::host>;
 
 /**
  * @brief std::experimental::mdspan with host tag to avoid accessing incorrect memory location.
@@ -205,4 +206,4 @@ auto make_host_vector_view(ElementType* ptr, IndexType n)
 {
   return host_vector_view<ElementType, IndexType, LayoutPolicy>{ptr, n};
 }
-}  // end namespace raft
+}  // end namespace raft

cpp/include/raft/core/mdarray.hpp

@@ -28,6 +28,7 @@
 #include <raft/core/host_device_accessor.hpp>
 #include <raft/core/mdspan.hpp>
 #include <raft/core/mdspan_types.hpp>
+#include <raft/core/memory_type.hpp>
 #include <rmm/cuda_stream_view.hpp>
 
 namespace raft {
@@ -154,12 +155,11 @@ class mdarray
               std::conditional_t<std::is_const_v<E>,
                                  typename container_policy_type::const_accessor_policy,
                                  typename container_policy_type::accessor_policy>>
-  using view_type_impl = mdspan<E,
-                                extents_type,
-                                layout_type,
-                                host_device_accessor<ViewAccessorPolicy,
-                                                     container_policy_type::is_host_accessible,
-                                                     container_policy_type::is_device_accessible>>;
+  using view_type_impl =
+    mdspan<E,
+           extents_type,
+           layout_type,
+           host_device_accessor<ViewAccessorPolicy, container_policy_type::mem_type>>;
 
  public:
   /**

cpp/include/raft/core/mdspan.hpp

@@ -17,6 +17,7 @@
 
 #include <raft/core/error.hpp>
 #include <raft/core/mdspan_types.hpp>
+#include <raft/core/memory_type.hpp>
 
 #include <raft/core/detail/macros.hpp>
 #include <raft/core/detail/mdspan_util.cuh>
@@ -183,9 +184,34 @@ template <typename ElementType,
           size_t... Extents>
 auto make_mdspan(ElementType* ptr, extents<IndexType, Extents...> exts)
 {
-  using accessor_type = host_device_accessor<std::experimental::default_accessor<ElementType>,
-                                             is_host_accessible,
-                                             is_device_accessible>;
+  using accessor_type = host_device_accessor<
+    std::experimental::default_accessor<ElementType>,
+    detail::memory_type_from_access<is_host_accessible, is_device_accessible>()>;
+  /*using accessor_type = host_device_accessor<std::experimental::default_accessor<ElementType>,
+                                             mem_type>; */
+
+  return mdspan<ElementType, decltype(exts), LayoutPolicy, accessor_type>{ptr, exts};
+}
+
+/**
+ * @brief Create a raft::mdspan
+ * @tparam ElementType the data type of the matrix elements
+ * @tparam IndexType the index type of the extents
+ * @tparam LayoutPolicy policy for strides and layout ordering
+ * @tparam MemType the raft::memory_type for where the data are stored
+ * @param ptr Pointer to the data
+ * @param exts dimensionality of the array (series of integers)
+ * @return raft::mdspan
+ */
+template <typename ElementType,
+          typename IndexType    = std::uint32_t,
+          typename LayoutPolicy = layout_c_contiguous,
+          memory_type MemType   = memory_type::device,
+          size_t... Extents>
+auto make_mdspan(ElementType* ptr, extents<IndexType, Extents...> exts)
+{
+  using accessor_type =
+    host_device_accessor<std::experimental::default_accessor<ElementType>, MemType>;
 
   return mdspan<ElementType, decltype(exts), LayoutPolicy, accessor_type>{ptr, exts};
 }

cpp/include/raft/core/memory_type.hpp

@@ -0,0 +1,52 @@
+/*
+ * Copyright (c) 2022, 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
+
+namespace raft {
+enum class memory_type { host, device, managed, pinned };
+
+auto constexpr is_device_accessible(memory_type mem_type)
+{
+  return (mem_type == memory_type::device || mem_type == memory_type::managed);
+}
+auto constexpr is_host_accessible(memory_type mem_type)
+{
+  return (mem_type == memory_type::host || mem_type == memory_type::managed ||
+          mem_type == memory_type::pinned);
+}
+auto constexpr is_host_device_accessible(memory_type mem_type)
+{
+  return is_device_accessible(mem_type) && is_host_accessible(mem_type);
+}
+
+namespace detail {
+
+template <bool is_host_accessible, bool is_device_accessible>
+auto constexpr memory_type_from_access()
+{
+  if constexpr (is_host_accessible && is_device_accessible) {
+    return memory_type::managed;
+  } else if constexpr (is_host_accessible) {
+    return memory_type::host;
+  } else if constexpr (is_device_accessible) {
+    return memory_type::device;
+  }
+  static_assert(is_host_accessible || is_device_accessible,
+                "Must be either host or device accessible to return a valid memory type");
+}
+
+}  // end namespace detail
+}  // end namespace raft

cpp/test/CMakeLists.txt

@@ -98,6 +98,7 @@ if(BUILD_TESTS)
             test/nvtx.cpp
             test/mdarray.cu
             test/mdspan_utils.cu
+            test/memory_type.cpp
             test/span.cpp
             test/span.cu
             test/test.cpp

cpp/test/memory_type.cpp

@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) 2022, 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 <gtest/gtest.h>
+#include <raft/core/memory_type.hpp>
+
+namespace raft {
+TEST(MemoryType, IsDeviceAccessible)
+{
+  static_assert(!is_device_accessible(memory_type::host));
+  static_assert(is_device_accessible(memory_type::device));
+  static_assert(is_device_accessible(memory_type::managed));
+  static_assert(!is_device_accessible(memory_type::pinned));
+}
+
+TEST(MemoryType, IsHostAccessible)
+{
+  static_assert(is_host_accessible(memory_type::host));
+  static_assert(!is_host_accessible(memory_type::device));
+  static_assert(is_host_accessible(memory_type::managed));
+  static_assert(is_host_accessible(memory_type::pinned));
+}
+
+TEST(MemoryType, IsHostDeviceAccessible)
+{
+  static_assert(!is_host_device_accessible(memory_type::host));
+  static_assert(!is_host_device_accessible(memory_type::device));
+  static_assert(is_host_device_accessible(memory_type::managed));
+  static_assert(!is_host_device_accessible(memory_type::pinned));
+}
+}  // namespace raft