Adding padded layout 'layout_padded_general' (#725)

This is a different approach / followup PR of #663 for issue #497.

I implemented a `layout_padded_general` within raft to statically enforce padding on mdpsan accesses.
* The layout has template parameters for `ValueType`, `StorageOrder `(default `row_major_t`), and `ByteAlignment `(default 128)
* in order to *not* require changes upstream I skipped `submdspan `functionality right now. I have a branch on a mdspan fork where I tested this though (https://github.com/mfoerste4/mdspan/tree/layout_padded).

Authors:
  - Malte Förster (https://github.com/mfoerste4)

Approvers:
  - Artem M. Chirkin (https://github.com/achirkin)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #725

cpp/include/raft/core/device_mdspan.hpp

@@ -161,6 +161,51 @@ template <typename ElementType,
           typename LayoutPolicy = layout_c_contiguous>
 using device_matrix_view = device_mdspan<ElementType, matrix_extent<IndexType>, LayoutPolicy>;
 
+/**
+ * @brief Shorthand for 128 byte aligned device matrix view.
+ * @tparam ElementType the data type of the matrix elements
+ * @tparam IndexType the index type of the extents
+ * @tparam LayoutPolicy must be of type layout_{left/right}_padded
+ */
+template <typename ElementType,
+          typename IndexType    = std::uint32_t,
+          typename LayoutPolicy = layout_right_padded<ElementType>,
+          typename              = enable_if_layout_padded<ElementType, LayoutPolicy>>
+using device_aligned_matrix_view =
+  device_mdspan<ElementType,
+                matrix_extent<IndexType>,
+                LayoutPolicy,
+                std::experimental::aligned_accessor<ElementType, detail::alignment::value>>;
+
+/**
+ * @brief Create a 2-dim 128 byte aligned mdspan instance for device pointer. It's
+ *        expected that the given layout policy match the layout of the underlying
+ *        pointer.
+ * @tparam ElementType the data type of the matrix elements
+ * @tparam LayoutPolicy must be of type layout_{left/right}_padded
+ * @tparam IndexType the index type of the extents
+ * @param[in] ptr on device to wrap
+ * @param[in] n_rows number of rows in pointer
+ * @param[in] n_cols number of columns in pointer
+ */
+template <typename ElementType,
+          typename IndexType    = std::uint32_t,
+          typename LayoutPolicy = layout_right_padded<ElementType>>
+auto make_device_aligned_matrix_view(ElementType* ptr, IndexType n_rows, IndexType n_cols)
+{
+  using data_handle_type =
+    typename std::experimental::aligned_accessor<ElementType,
+                                                 detail::alignment::value>::data_handle_type;
+  static_assert(std::is_same<LayoutPolicy, layout_left_padded<ElementType>>::value ||
+                std::is_same<LayoutPolicy, layout_right_padded<ElementType>>::value);
+  assert(ptr == alignTo(ptr, detail::alignment::value));
+
+  data_handle_type aligned_pointer = ptr;
+
+  matrix_extent<IndexType> extents{n_rows, n_cols};
+  return device_aligned_matrix_view<ElementType, IndexType, LayoutPolicy>{aligned_pointer, extents};
+}
+
 /**
  * @brief Create a raft::managed_mdspan
  * @tparam ElementType the data type of the matrix elements

cpp/include/raft/core/host_mdspan.hpp

@@ -110,6 +110,51 @@ template <typename ElementType,
           typename LayoutPolicy = layout_c_contiguous>
 using host_matrix_view = host_mdspan<ElementType, matrix_extent<IndexType>, LayoutPolicy>;
 
+/**
+ * @brief Shorthand for 128 byte aligned host matrix view.
+ * @tparam ElementType the data type of the matrix elements
+ * @tparam IndexType the index type of the extents
+ * @tparam LayoutPolicy must be of type layout_{left/right}_padded
+ */
+template <typename ElementType,
+          typename IndexType    = std::uint32_t,
+          typename LayoutPolicy = layout_right_padded<ElementType>,
+          typename              = enable_if_layout_padded<ElementType, LayoutPolicy>>
+using host_aligned_matrix_view =
+  host_mdspan<ElementType,
+              matrix_extent<IndexType>,
+              LayoutPolicy,
+              std::experimental::aligned_accessor<ElementType, detail::alignment::value>>;
+
+/**
+ * @brief Create a 2-dim 128 byte aligned mdspan instance for host pointer. It's
+ *        expected that the given layout policy match the layout of the underlying
+ *        pointer.
+ * @tparam ElementType the data type of the matrix elements
+ * @tparam LayoutPolicy must be of type layout_{left/right}_padded
+ * @tparam IndexType the index type of the extents
+ * @param[in] ptr on host to wrap
+ * @param[in] n_rows number of rows in pointer
+ * @param[in] n_cols number of columns in pointer
+ */
+template <typename ElementType,
+          typename IndexType    = std::uint32_t,
+          typename LayoutPolicy = layout_right_padded<ElementType>>
+auto make_host_aligned_matrix_view(ElementType* ptr, IndexType n_rows, IndexType n_cols)
+{
+  using data_handle_type =
+    typename std::experimental::aligned_accessor<ElementType,
+                                                 detail::alignment::value>::data_handle_type;
+
+  static_assert(std::is_same<LayoutPolicy, layout_left_padded<ElementType>>::value ||
+                std::is_same<LayoutPolicy, layout_right_padded<ElementType>>::value);
+  assert(ptr == alignTo(ptr, detail::alignment::value));
+  data_handle_type aligned_pointer = ptr;
+
+  matrix_extent<IndexType> extents{n_rows, n_cols};
+  return host_aligned_matrix_view<ElementType, IndexType, LayoutPolicy>{aligned_pointer, extents};
+}
+
 /**
  * @brief Create a 0-dim (scalar) mdspan instance for host value.
  *

cpp/include/raft/core/mdspan.hpp

@@ -32,6 +32,40 @@ template <typename ElementType,
           typename AccessorPolicy = std::experimental::default_accessor<ElementType>>
 using mdspan = std::experimental::mdspan<ElementType, Extents, LayoutPolicy, AccessorPolicy>;
 
+namespace detail {
+
+// keeping ByteAlignment as optional to allow testing
+template <class ValueType, size_t ByteAlignment = 128>
+struct padding {
+  static_assert(std::is_same<std::remove_cv_t<ValueType>, ValueType>::value,
+                "std::experimental::padding ValueType has to be provided without "
+                "const or volatile specifiers.");
+  static_assert(ByteAlignment % sizeof(ValueType) == 0 || sizeof(ValueType) % ByteAlignment == 0,
+                "std::experimental::padding sizeof(ValueType) has to be multiple or "
+                "divider of ByteAlignment.");
+  static constexpr size_t value = std::max(ByteAlignment / sizeof(ValueType), 1ul);
+};
+
+// alignment fixed to 128 bytes
+struct alignment {
+  static constexpr size_t value = 128;
+};
+
+}  // namespace detail
+
+template <typename ElementType>
+using layout_right_padded = std::experimental::layout_right_padded<
+  detail::padding<std::remove_cv_t<std::remove_reference_t<ElementType>>>::value>;
+
+template <typename ElementType>
+using layout_left_padded = std::experimental::layout_left_padded<
+  detail::padding<std::remove_cv_t<std::remove_reference_t<ElementType>>>::value>;
+
+template <typename ElementType, typename LayoutPolicy>
+using enable_if_layout_padded =
+  std::enable_if_t<std::is_same<LayoutPolicy, layout_left_padded<ElementType>>::value ||
+                   std::is_same<LayoutPolicy, layout_right_padded<ElementType>>::value>;
+
 /**
  * Ensure all types listed in the parameter pack `Extents` are integral types.
  * Usage:
@@ -254,4 +288,5 @@ RAFT_INLINE_FUNCTION auto unravel_index(Idx idx,
     return unravel_index_impl<uint32_t>(static_cast<uint32_t>(idx), shape);
   }
 }
+
 }  // namespace raft