Implement & test transposed(layout_{left,right}_padded)

kokkos · Apr 4, 2024 · 4791e71 · 4791e71
1 parent 89ca9ad
commit 4791e71
Show file tree

Hide file tree

Showing 2 changed files with 122 additions and 1 deletion.
diff --git a/include/experimental/__p1673_bits/transposed.hpp b/include/experimental/__p1673_bits/transposed.hpp
@@ -99,7 +99,7 @@ template<class Layout>
 class layout_transpose {
 public:
   using nested_layout_type = Layout;
-  
+
   template<class Extents>
   struct mapping {
   private:
@@ -265,6 +265,27 @@ namespace impl {
     }
   };
 
+  template<size_t PaddingValue>
+  struct transposed_layout<layout_left_padded<PaddingValue>> {
+    using layout_type = layout_right_padded<PaddingValue>;
+
+    template<class OriginalExtents>
+    static auto mapping(const typename layout_left_padded<PaddingValue>::template mapping<OriginalExtents>& orig_map) {
+      using input_mapping_type =
+        typename layout_left_padded<PaddingValue>::template mapping<OriginalExtents>;
+      using output_extents_type =
+        transpose_extents_t<typename input_mapping_type::extents_type>;
+      using output_mapping_type =
+        typename layout_type::template mapping<output_extents_type>;
+
+      const auto padding_value = orig_map.stride(1);
+      return output_mapping_type{
+	transpose_extents(orig_map.extents()),
+        padding_value
+      };
+    }
+  };
+
   template<class StorageOrder>
   using opposite_storage_t = std::conditional_t<
     std::is_same_v<StorageOrder, column_major_t>,

diff --git a/tests/native/transposed.cpp b/tests/native/transposed.cpp
@@ -199,6 +199,106 @@ namespace {
     }
   }
 
+  template<size_t PaddingValue>
+  void test_transposed_layout_left_padded(auto runtime_padding_value)
+  {
+    auto test_one = [=] (auto in_exts, auto out_exts,
+                         std::vector<char>& fake_storage)
+    {
+      using in_extents_type = decltype(in_exts);
+      using out_extents_type = decltype(out_exts);
+
+      typename layout_left_padded<PaddingValue>::template mapping<in_extents_type> in_map{
+        in_exts, runtime_padding_value
+      };
+      typename layout_right_padded<PaddingValue>::template mapping<out_extents_type> out_map{
+        out_exts, runtime_padding_value
+      };
+      test_transposed_layout(in_map, out_map, fake_storage);
+    };
+
+    std::vector<char> storage;
+    {
+      using in_extents_type = extents<int, 3, 4>;
+      using out_extents_type = extents<int, 4, 3>;
+      test_one(in_extents_type{}, out_extents_type{}, storage);
+    }
+    {
+      using in_extents_type = extents<int, 3, dynamic_extent>;
+      using out_extents_type = extents<int, dynamic_extent, 3>;
+      test_one(in_extents_type{4}, out_extents_type{4}, storage);
+    }
+    {
+      using in_extents_type = extents<int, dynamic_extent, dynamic_extent>;
+      using out_extents_type = extents<int, dynamic_extent, dynamic_extent>;
+      test_one(in_extents_type{3, 4}, out_extents_type{4, 3}, storage);
+    }
+  }
+
+  TEST(transposed_layout, layout_left_padded)
+  {
+    {
+      constexpr size_t padding_value = dynamic_extent;
+      constexpr size_t runtime_padding_value = 5u;
+      test_transposed_layout_left_padded<padding_value>(runtime_padding_value);
+    }
+    {
+      constexpr size_t padding_value = 5u;
+      constexpr size_t runtime_padding_value = padding_value;
+      test_transposed_layout_left_padded<padding_value>(runtime_padding_value);
+    }
+  }
+
+  template<size_t PaddingValue>
+  void test_transposed_layout_right_padded(auto runtime_padding_value)
+  {
+    auto test_one = [=] (auto in_exts, auto out_exts,
+                         std::vector<char>& fake_storage)
+    {
+      using in_extents_type = decltype(in_exts);
+      using out_extents_type = decltype(out_exts);
+
+      typename layout_left_padded<PaddingValue>::template mapping<in_extents_type> in_map{
+        in_exts, runtime_padding_value
+      };
+      typename layout_right_padded<PaddingValue>::template mapping<out_extents_type> out_map{
+        out_exts, runtime_padding_value
+      };
+      test_transposed_layout(in_map, out_map, fake_storage);
+    };
+
+    std::vector<char> storage;
+    {
+      using in_extents_type = extents<int, 3, 4>;
+      using out_extents_type = extents<int, 4, 3>;
+      test_one(in_extents_type{}, out_extents_type{}, storage);
+    }
+    {
+      using in_extents_type = extents<int, 3, dynamic_extent>;
+      using out_extents_type = extents<int, dynamic_extent, 3>;
+      test_one(in_extents_type{4}, out_extents_type{4}, storage);
+    }
+    {
+      using in_extents_type = extents<int, dynamic_extent, dynamic_extent>;
+      using out_extents_type = extents<int, dynamic_extent, dynamic_extent>;
+      test_one(in_extents_type{3, 4}, out_extents_type{4, 3}, storage);
+    }
+  }
+
+  TEST(transposed_layout, layout_right_padded)
+  {
+    {
+      constexpr size_t padding_value = dynamic_extent;
+      constexpr size_t runtime_padding_value = 5u;
+      test_transposed_layout_right_padded<padding_value>(runtime_padding_value);
+    }
+    {
+      constexpr size_t padding_value = 5u;
+      constexpr size_t runtime_padding_value = padding_value;
+      test_transposed_layout_right_padded<padding_value>(runtime_padding_value);
+    }
+  }
+
   TEST(transposed, mdspan_double)
   {
     using real_t = double;