From a8b17a85ea41e44120a967980d4a788439f27bbc Mon Sep 17 00:00:00 2001
From: William Hicks <whicks@nvidia.com>
Date: Mon, 2 Oct 2023 11:29:23 -0400
Subject: [PATCH] Add static asserts for mdspan_copyable

---
 cpp/test/core/mdspan_copy.cpp | 56 +++++++++++++++++++++++++----------
 cpp/test/core/mdspan_copy.cu  | 54 +++++++++++++++++++++++++++++++++
 2 files changed, 95 insertions(+), 15 deletions(-)

diff --git a/cpp/test/core/mdspan_copy.cpp b/cpp/test/core/mdspan_copy.cpp
index b64ad0355b..2f938e3035 100644
--- a/cpp/test/core/mdspan_copy.cpp
+++ b/cpp/test/core/mdspan_copy.cpp
@@ -38,7 +38,10 @@ TEST(MDSpanCopy, Mdspan1DHostHost)
   }
 
   auto out_right = make_host_vector<double, std::uint32_t, layout_f_contiguous>(res, cols);
-  // std::copy
+  static_assert(detail::mdspan_copyable<true,
+                                        decltype(out_right.view()),
+                                        decltype(in_left.view())>::can_use_std_copy,
+                "Current implementation should use std::copy for this copy");
   copy(res, out_right.view(), in_left.view());
   for (auto i = std::uint32_t{}; i < cols; ++i) {
     ASSERT_TRUE(match(out_right(i), double(gen_unique_entry(i)), CompareApprox<double>{0.0001}));
@@ -57,8 +60,11 @@ TEST(MDSpanCopy, Mdspan1DHostDevice)
     in_left(i) = gen_unique_entry(i);
   }
 
-  // raft::copy
   auto out_right = make_device_vector<float, std::uint32_t, layout_f_contiguous>(res, cols);
+  static_assert(detail::mdspan_copyable<true,
+                                        decltype(out_right.view()),
+                                        decltype(in_left.view())>::can_use_raft_copy,
+                "Current implementation should use raft::copy for this copy");
   copy(res, out_right.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < cols; ++i) {
@@ -78,8 +84,11 @@ TEST(MDSpanCopy, Mdspan1DDeviceHost)
     in_left(i) = gen_unique_entry(i);
   }
 
-  // raft::copy
   auto out_right = make_host_vector<float, std::uint32_t, layout_f_contiguous>(res, cols);
+  static_assert(detail::mdspan_copyable<true,
+                                        decltype(out_right.view()),
+                                        decltype(in_left.view())>::can_use_raft_copy,
+                "Current implementation should use raft::copy for this copy");
   copy(res, out_right.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < cols; ++i) {
@@ -95,9 +104,9 @@ TEST(MDSpanCopy, Mdspan3DHostHost)
   auto constexpr depth = std::uint32_t{500};
   auto constexpr rows  = std::uint32_t{300};
   auto constexpr cols  = std::uint32_t{200};
-  auto in_left = make_host_mdarray<float, std::uint32_t, layout_c_contiguous, depth, rows, cols>(
+  auto in_left = make_host_mdarray<float, std::uint32_t, layout_f_contiguous, depth, rows, cols>(
     res, extents<std::uint32_t, depth, rows, cols>{});
-  auto in_right = make_host_mdarray<float, std::uint32_t, layout_f_contiguous, depth, rows, cols>(
+  auto in_right = make_host_mdarray<float, std::uint32_t, layout_c_contiguous, depth, rows, cols>(
     res, extents<std::uint32_t, depth, rows, cols>{});
   auto gen_unique_entry = [](auto&& x, auto&& y, auto&& z) { return x * 7 + y * 11 + z * 13; };
 
@@ -112,10 +121,13 @@ TEST(MDSpanCopy, Mdspan3DHostHost)
 
   auto out_left = make_host_mdarray<double, std::uint32_t, layout_f_contiguous, depth, rows, cols>(
     res, extents<std::uint32_t, depth, rows, cols>{});
-  auto out_right = make_host_mdarray<double, std::uint32_t, layout_f_contiguous, depth, rows, cols>(
+  auto out_right = make_host_mdarray<double, std::uint32_t, layout_c_contiguous, depth, rows, cols>(
     res, extents<std::uint32_t, depth, rows, cols>{});
 
-  // std::copy
+  static_assert(detail::mdspan_copyable<true,
+                                        decltype(out_right.view()),
+                                        decltype(in_right.view())>::can_use_std_copy,
+                "Current implementation should use std::copy for this copy");
   copy(res, out_right.view(), in_right.view());
   for (auto i = std::uint32_t{}; i < depth; ++i) {
     for (auto j = std::uint32_t{}; j < rows; ++j) {
@@ -126,7 +138,6 @@ TEST(MDSpanCopy, Mdspan3DHostHost)
     }
   }
 
-  // simd or custom logic
   copy(res, out_right.view(), in_left.view());
   for (auto i = std::uint32_t{}; i < depth; ++i) {
     for (auto j = std::uint32_t{}; j < rows; ++j) {
@@ -137,7 +148,6 @@ TEST(MDSpanCopy, Mdspan3DHostHost)
     }
   }
 
-  // simd or custom logic
   copy(res, out_left.view(), in_right.view());
   for (auto i = std::uint32_t{}; i < depth; ++i) {
     for (auto j = std::uint32_t{}; j < rows; ++j) {
@@ -148,7 +158,9 @@ TEST(MDSpanCopy, Mdspan3DHostHost)
     }
   }
 
-  // std::copy
+  static_assert(detail::mdspan_copyable<true, decltype(out_left.view()), decltype(in_left.view())>::
+                  can_use_std_copy,
+                "Current implementation should use std::copy for this copy");
   copy(res, out_left.view(), in_left.view());
   for (auto i = std::uint32_t{}; i < depth; ++i) {
     for (auto j = std::uint32_t{}; j < rows; ++j) {
@@ -190,7 +202,10 @@ TEST(MDSpanCopy, Mdspan3DHostDevice)
     make_device_mdarray<float, std::uint32_t, layout_f_contiguous, depth, rows, cols>(
       res, extents<std::uint32_t, depth, rows, cols>{});
 
-  // raft::copy
+  static_assert(detail::mdspan_copyable<true,
+                                        decltype(out_right.view()),
+                                        decltype(in_right.view())>::can_use_raft_copy,
+                "Current implementation should use raft::copy for this copy");
   copy(res, out_right.view(), in_right.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < depth; ++i) {
@@ -203,7 +218,9 @@ TEST(MDSpanCopy, Mdspan3DHostDevice)
     }
   }
 
-  // raft::copy
+  static_assert(detail::mdspan_copyable<true, decltype(out_left.view()), decltype(in_left.view())>::
+                  can_use_raft_copy,
+                "Current implementation should use raft::copy for this copy");
   copy(res, out_left.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < depth; ++i) {
@@ -240,7 +257,10 @@ TEST(MDSpanCopy, Mdspan2DDeviceDevice)
   auto out_right = make_device_mdarray<float, std::uint32_t, layout_f_contiguous, rows, cols>(
     res, extents<std::uint32_t, rows, cols>{});
 
-  // raft::copy
+  static_assert(detail::mdspan_copyable<true,
+                                        decltype(out_right.view()),
+                                        decltype(in_right.view())>::can_use_raft_copy,
+                "Current implementation should use raft::copy for this copy");
   copy(res, out_right.view(), in_right.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < rows; ++i) {
@@ -250,7 +270,10 @@ TEST(MDSpanCopy, Mdspan2DDeviceDevice)
     }
   }
 
-  // cublas
+  static_assert(detail::mdspan_copyable<true,
+                                        decltype(out_right.view()),
+                                        decltype(in_left.view())>::can_use_cublas,
+                "Current implementation should use cuBLAS for this copy");
   copy(res, out_right.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < rows; ++i) {
@@ -260,7 +283,10 @@ TEST(MDSpanCopy, Mdspan2DDeviceDevice)
     }
   }
 
-  // cublas
+  static_assert(detail::mdspan_copyable<true,
+                                        decltype(out_left.view()),
+                                        decltype(in_right.view())>::can_use_cublas,
+                "Current implementation should use cuBLAS for this copy");
   copy(res, out_left.view(), in_right.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < rows; ++i) {
diff --git a/cpp/test/core/mdspan_copy.cu b/cpp/test/core/mdspan_copy.cu
index f0a22eabe8..95d7d3befd 100644
--- a/cpp/test/core/mdspan_copy.cu
+++ b/cpp/test/core/mdspan_copy.cu
@@ -50,6 +50,9 @@ TEST(MDSpanCopy, Mdspan3DDeviceDeviceCuda)
   auto out_long =
     make_device_mdarray<std::int64_t, std::uint32_t, layout_c_contiguous, depth, rows, cols>(
       res, extents<std::uint32_t, depth, rows, cols>{});
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_long.view()), decltype(in_left.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_long.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < depth; ++i) {
@@ -66,6 +69,9 @@ TEST(MDSpanCopy, Mdspan3DDeviceDeviceCuda)
   auto out_right = make_device_mdarray<int, std::uint32_t, layout_f_contiguous, depth, rows, cols>(
     res, extents<std::uint32_t, depth, rows, cols>{});
 
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_right.view()), decltype(in_left.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_right.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < depth; ++i) {
@@ -76,6 +82,9 @@ TEST(MDSpanCopy, Mdspan3DDeviceDeviceCuda)
     }
   }
 
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_left.view()), decltype(in_right.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_left.view(), in_right.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < depth; ++i) {
@@ -113,6 +122,9 @@ TEST(MDSpanCopy, Mdspan2DDeviceDeviceCuda)
   res.sync_stream();
 
   // Test dtype conversion without transpose
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_right.view()), decltype(in_right.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_right.view(), in_right.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < rows; ++i) {
@@ -123,6 +135,9 @@ TEST(MDSpanCopy, Mdspan2DDeviceDeviceCuda)
   }
 
   // Test dtype conversion with transpose
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_right.view()), decltype(in_left.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_right.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < rows; ++i) {
@@ -131,6 +146,9 @@ TEST(MDSpanCopy, Mdspan2DDeviceDeviceCuda)
         double(out_right(i, j)), double(gen_unique_entry(i, j)), CompareApprox<double>{0.0001}));
     }
   }
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_left.view()), decltype(in_right.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_left.view(), in_right.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < rows; ++i) {
@@ -167,6 +185,9 @@ TEST(MDSpanCopy, Mdspan3DDeviceHostCuda)
   auto out_long =
     make_host_mdarray<std::int64_t, std::uint32_t, layout_c_contiguous, depth, rows, cols>(
       res, extents<std::uint32_t, depth, rows, cols>{});
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_long.view()), decltype(in_left.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_long.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < depth; ++i) {
@@ -183,6 +204,9 @@ TEST(MDSpanCopy, Mdspan3DDeviceHostCuda)
   auto out_right = make_host_mdarray<int, std::uint32_t, layout_f_contiguous, depth, rows, cols>(
     res, extents<std::uint32_t, depth, rows, cols>{});
 
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_right.view()), decltype(in_left.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_right.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < depth; ++i) {
@@ -193,6 +217,9 @@ TEST(MDSpanCopy, Mdspan3DDeviceHostCuda)
     }
   }
 
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_left.view()), decltype(in_right.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_left.view(), in_right.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < depth; ++i) {
@@ -230,6 +257,9 @@ TEST(MDSpanCopy, Mdspan2DDeviceHostCuda)
   res.sync_stream();
 
   // Test dtype conversion without transpose
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_right.view()), decltype(in_right.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_right.view(), in_right.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < rows; ++i) {
@@ -240,6 +270,9 @@ TEST(MDSpanCopy, Mdspan2DDeviceHostCuda)
   }
 
   // Test dtype conversion with transpose
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_right.view()), decltype(in_left.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_right.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < rows; ++i) {
@@ -248,6 +281,9 @@ TEST(MDSpanCopy, Mdspan2DDeviceHostCuda)
         double(out_right(i, j)), double(gen_unique_entry(i, j)), CompareApprox<double>{0.0001}));
     }
   }
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_left.view()), decltype(in_right.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_left.view(), in_right.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < rows; ++i) {
@@ -285,6 +321,9 @@ TEST(MDSpanCopy, Mdspan3DHostDeviceCuda)
   auto out_long =
     make_device_mdarray<std::int64_t, std::uint32_t, layout_c_contiguous, depth, rows, cols>(
       res, extents<std::uint32_t, depth, rows, cols>{});
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_long.view()), decltype(in_left.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_long.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < depth; ++i) {
@@ -301,6 +340,9 @@ TEST(MDSpanCopy, Mdspan3DHostDeviceCuda)
   auto out_right = make_device_mdarray<int, std::uint32_t, layout_f_contiguous, depth, rows, cols>(
     res, extents<std::uint32_t, depth, rows, cols>{});
 
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_right.view()), decltype(in_left.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_right.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < depth; ++i) {
@@ -311,6 +353,9 @@ TEST(MDSpanCopy, Mdspan3DHostDeviceCuda)
     }
   }
 
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_left.view()), decltype(in_right.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_left.view(), in_right.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < depth; ++i) {
@@ -348,6 +393,9 @@ TEST(MDSpanCopy, Mdspan2DHostDeviceCuda)
   res.sync_stream();
 
   // Test dtype conversion without transpose
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_right.view()), decltype(in_right.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_right.view(), in_right.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < rows; ++i) {
@@ -358,6 +406,9 @@ TEST(MDSpanCopy, Mdspan2DHostDeviceCuda)
   }
 
   // Test dtype conversion with transpose
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_right.view()), decltype(in_left.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_right.view(), in_left.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < rows; ++i) {
@@ -366,6 +417,9 @@ TEST(MDSpanCopy, Mdspan2DHostDeviceCuda)
         double(out_right(i, j)), double(gen_unique_entry(i, j)), CompareApprox<double>{0.0001}));
     }
   }
+  static_assert(
+    detail::mdspan_copyable_with_kernel_v<decltype(out_left.view()), decltype(in_right.view())>,
+    "Current implementation should use kernel for this copy");
   copy(res, out_left.view(), in_right.view());
   res.sync_stream();
   for (auto i = std::uint32_t{}; i < rows; ++i) {