[unified] refactor type mapping

Author: MarcelKoch

common/cuda_hip/base/kernel_launch_solver.hpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -27,8 +27,7 @@ __global__ __launch_bounds__(default_block_size) void generic_kernel_2d_solver(
     if (row >= rows) {
         return;
     }
-    fn(row, col,
-       device_unpack_solver_impl<KernelArgs>::unpack(args, default_stride)...);
+    fn(row, col, device_unpack(args, default_stride)...);
 }
 
 

common/unified/base/kernel_launch.hpp

@@ -175,8 +175,11 @@ using aliased_ptr = T*;
  * objects.
  *
  * @tparam ValueType  the value type of the underlying matrix.
+ * @tparam PtrWrapper  the pointer type. By default, it's just `T*`, but it may
+ *                     be set to restricted_ptr.
  */
-template <typename ValueType, template <typename> typename PtrWrapper>
+template <typename ValueType,
+          template <typename> typename PtrWrapper = aliased_ptr>
 struct matrix_accessor {
     PtrWrapper<ValueType> data;
     int64 stride;
@@ -202,9 +205,24 @@ struct matrix_accessor {
 };
 
 
+/**
+ * Tag to signal that pointers should be annotated with `__restrict`
+ */
 struct restrict_tag {};
 
 
+/**
+ * @internal
+ * Adds a restrict annotation to an object.
+ *
+ * @note Can't be used for run_kernel_solver.
+ *
+ * @tparam T  Type that should be annotated
+ *
+ * @param orig  Original object
+ *
+ * @return The original object and a restrict_tag
+ */
 template <typename T>
 auto as_restrict(T&& orig) -> std::pair<T&&, restrict_tag>
 {
@@ -220,105 +238,139 @@ auto as_restrict(T&& orig) -> std::pair<T&&, restrict_tag>
  *
  * By default, it only maps std::complex to the corresponding device
  * representation of the complex type. There are specializations for dealing
- * with gko::array and gko::matrix::Dense (both const and mutable) that map them
+ * with gko::array and gko::matrix::Dense that map them
  * to plain pointers or matrix_accessor objects.
  *
- * @tparam T  the type being mapped. It will be used based on a
- *            forwarding-reference, i.e. preserve references in the input
- *            parameter, so special care must be taken to only return types that
- *            can be passed to the device, i.e. (structs containing) device
- *            pointers or values. This means that T will be either a r-value or
- *            l-value reference.
+ * @tparam T  the underlying type being mapped. Any references or const
+ *            qualifiers have to be resolved before passing the type.
+ *            The distinction between const/mutable objects is done by
+ *            overloading the map_to_device function.
+ * @tparam PtrWrapper  the pointer type. By default, it's just `T*`, but it may
+ *                     be set to restricted_ptr.
  */
 template <typename T, template <typename> typename PtrWrapper = aliased_ptr>
 struct to_device_type_impl {
-    using type = std::decay_t<device_type<T>>;
-    static type map_to_device(T in) { return as_device_type(in); }
+    static auto map_to_device(T in) -> device_type<T>
+    {
+        return as_device_type(in);
+    }
 };
 
-template <typename ValueType, template <typename> typename PtrWrapper>
-struct to_device_type_impl<matrix::Dense<ValueType>*&, PtrWrapper> {
-    using type = matrix_accessor<device_type<ValueType>, PtrWrapper>;
-    static type map_to_device(matrix::Dense<ValueType>* mtx)
+template <typename T, template <typename> typename PtrWrapper>
+struct to_device_type_impl<T*, PtrWrapper> {
+    static auto map_to_device(T* in) -> PtrWrapper<device_type<T>>
+    {
+        return {as_device_type(in)};
+    }
+    static auto map_to_device(const T* in) -> PtrWrapper<const device_type<T>>
     {
-        return to_device_type_impl<matrix::Dense<ValueType>* const&,
-                                   PtrWrapper>::map_to_device(mtx);
+        return {as_device_type(in)};
     }
 };
 
 template <typename ValueType, template <typename> typename PtrWrapper>
-struct to_device_type_impl<matrix::Dense<ValueType>* const&, PtrWrapper> {
-    using type = matrix_accessor<device_type<ValueType>, PtrWrapper>;
-    static type map_to_device(matrix::Dense<ValueType>* mtx)
+struct to_device_type_impl<matrix::Dense<ValueType>*, PtrWrapper> {
+    static auto map_to_device(matrix::Dense<ValueType>* mtx)
+        -> matrix_accessor<device_type<ValueType>, PtrWrapper>
     {
         return {as_device_type(mtx->get_values()),
                 static_cast<int64>(mtx->get_stride())};
     }
-};
 
-template <typename ValueType, template <typename> typename PtrWrapper>
-struct to_device_type_impl<const matrix::Dense<ValueType>*&, PtrWrapper> {
-    using type = matrix_accessor<const device_type<ValueType>, PtrWrapper>;
-    static type map_to_device(const matrix::Dense<ValueType>* mtx)
+    static auto map_to_device(const matrix::Dense<ValueType>* mtx)
+        -> matrix_accessor<const device_type<ValueType>, PtrWrapper>
     {
         return {as_device_type(mtx->get_const_values()),
                 static_cast<int64>(mtx->get_stride())};
     }
 };
 
 template <typename ValueType, template <typename> typename PtrWrapper>
-struct to_device_type_impl<array<ValueType>&, PtrWrapper> {
-    using type = PtrWrapper<device_type<ValueType>>;
-    static type map_to_device(array<ValueType>& array)
+struct to_device_type_impl<array<ValueType>, PtrWrapper> {
+    static auto map_to_device(array<ValueType>& array)
+        -> PtrWrapper<device_type<ValueType>>
     {
         return {as_device_type(array.get_data())};
     }
-};
 
-template <typename ValueType, template <typename> typename PtrWrapper>
-struct to_device_type_impl<const array<ValueType>&, PtrWrapper> {
-    using type = PtrWrapper<const device_type<ValueType>>;
-    static type map_to_device(const array<ValueType>& array)
+    static auto map_to_device(const array<ValueType>& array)
+        -> PtrWrapper<const device_type<ValueType>>
     {
         return {as_device_type(array.get_const_data())};
     }
 };
 
-template <typename T, template <typename> typename PtrWrapper>
-struct to_device_type_impl<T*, PtrWrapper> {
-    using type = PtrWrapper<device_type<T>>;
-    static type map_to_device(T in) { return {as_device_type(in)}; }
+/**
+ * Specialization for handling objects annotated by as_restrict.
+ * It changes the pointer wrapper type to restricted_ptr.
+ */
+template <typename T>
+struct to_device_type_impl<std::pair<T, restrict_tag>, aliased_ptr> {
+    template <typename U>
+    static auto map_to_device(U&& in)
+    {
+        return to_device_type_impl<T, restricted_ptr>::map_to_device(in.first);
+    }
 };
 
-template <typename T, template <typename> typename PtrWrapper>
-struct to_device_type_impl<const T*, PtrWrapper> {
-    using type = PtrWrapper<const device_type<T>>;
-    static type map_to_device(T in) { return {as_device_type(in)}; }
+
+namespace detail {
+
+
+/**
+ * Similar to std::remove_cv_t except that it remove the const from pointers,
+ * i.e. `const T*` -> `T*`.
+ */
+template <typename T>
+struct aggressive_remove_const {
+    using type = std::remove_cv_t<T>;
 };
 
 template <typename T>
-struct to_device_type_impl<std::pair<T&, restrict_tag>&, aliased_ptr> {
-    using type = typename to_device_type_impl<T&, restricted_ptr>::type;
-    static type map_to_device(std::pair<T&, restrict_tag> in)
-    {
-        return to_device_type_impl<T&, restricted_ptr>::map_to_device(in.first);
-    }
+struct aggressive_remove_const<const T*> {
+    using type = T*;
 };
 
+/**
+ * Similar to std::decay, except that it also applies std::decay on the first
+ * nesting of types, i.e. `T<U&>` -> `T<U>`.
+ * This only resolves a single level of nesting.
+ */
 template <typename T>
-struct to_device_type_impl<std::pair<T, restrict_tag>&, aliased_ptr> {
-    using type = typename to_device_type_impl<T, restricted_ptr>::type;
-    static type map_to_device(std::pair<T, restrict_tag> in)
-    {
-        return to_device_type_impl<T, restricted_ptr>::map_to_device(in.first);
-    }
+struct nested_decay;
+
+/**
+ * Helper type for nested_decay.
+ * This is necessary, since the references in the top-level type have to be
+ * removed, before std::decay may be applied to the nested type.
+ */
+template <typename T>
+struct nested_decay_inner {
+    using type = typename aggressive_remove_const<std::decay_t<T>>::type;
+};
+
+template <typename T, typename Tag>
+struct nested_decay_inner<std::pair<T, Tag>> {
+    using type = std::pair<typename nested_decay_inner<T>::type, Tag>;
 };
 
+template <typename T>
+struct nested_decay {
+    using type = typename nested_decay_inner<std::decay_t<T>>::type;
+};
+
+
+}  // namespace detail
+
+
+template <typename T>
+using to_device_type =
+    to_device_type_impl<typename detail::nested_decay<T>::type>;
 
 template <typename T>
-typename to_device_type_impl<T>::type map_to_device(T&& param)
+auto map_to_device(T&& param)
 {
-    return to_device_type_impl<T>::map_to_device(param);
+    return to_device_type<T>::map_to_device(std::forward<T>(param));
 }
 
 

common/unified/base/kernel_launch_solver.hpp

@@ -43,7 +43,7 @@ struct device_unpack_solver_impl {
 
 template <typename ValueType>
 struct device_unpack_solver_impl<default_stride_dense_wrapper<ValueType>> {
-    using type = matrix_accessor<ValueType, aliased_ptr>;
+    using type = matrix_accessor<ValueType>;
     static GKO_INLINE GKO_ATTRIBUTES type
     unpack(default_stride_dense_wrapper<ValueType> param, int64 default_stride)
     {
@@ -52,6 +52,16 @@ struct device_unpack_solver_impl<default_stride_dense_wrapper<ValueType>> {
 };
 
 
+template <typename T>
+GKO_INLINE GKO_ATTRIBUTES auto device_unpack(T&& param, int64 default_stride)
+{
+    using device_type =
+        std::decay_t<decltype(map_to_device(std::forward<T>(param)))>;
+    return device_unpack_solver_impl<device_type>::unpack(
+        std::forward<T>(param), default_stride);
+}
+
+
 /**
  * @internal
  * Wraps the given matrix in a wrapper signifying that it has the default stride

dpcpp/base/kernel_launch_solver.dp.hpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -22,9 +22,7 @@ void generic_kernel_2d_solver(sycl::handler& cgh, int64 rows, int64 cols,
                      [=](sycl::id<1> idx) {
                          auto row = static_cast<int64>(idx[0] / cols);
                          auto col = static_cast<int64>(idx[0] % cols);
-                         fn(row, col,
-                            device_unpack_solver_impl<KernelArgs>::unpack(
-                                args, default_stride)...);
+                         fn(row, col, device_unpack(args, default_stride)...);
                      });
 }
 

omp/base/kernel_launch_solver.hpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -13,23 +13,14 @@ namespace kernels {
 namespace omp {
 
 
-template <typename T>
-typename device_unpack_solver_impl<typename to_device_type_impl<T>::type>::type
-map_to_device_solver(T&& param, int64 default_stride)
-{
-    return device_unpack_solver_impl<typename to_device_type_impl<T>::type>::
-        unpack(to_device_type_impl<T>::map_to_device(param), default_stride);
-}
-
-
 template <typename KernelFunction, typename... KernelArgs>
 void run_kernel_solver(std::shared_ptr<const OmpExecutor> exec,
                        KernelFunction fn, dim<2> size, size_type default_stride,
                        KernelArgs&&... args)
 {
-    run_kernel_impl(
-        exec, fn, size,
-        map_to_device_solver(args, static_cast<int64>(default_stride))...);
+    run_kernel_impl(exec, fn, size,
+                    device_unpack(map_to_device(std::forward<KernelArgs>(args)),
+                                  static_cast<int64>(default_stride))...);
 }