Allow the reordering of a DiscreteDomain (#420)

* Add a reordered assignment operator --------- Co-authored-by: Emily Bourne <[email protected]> Co-authored-by: Thomas Padioleau <[email protected]>
CExA-project · Apr 22, 2024 · 9ef820e · 9ef820e
1 parent c13d10e
commit 9ef820e
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 0 deletions.
diff --git a/include/ddc/discrete_domain.hpp b/include/ddc/discrete_domain.hpp
@@ -100,6 +100,25 @@ class DiscreteDomain
 
     KOKKOS_DEFAULTED_FUNCTION DiscreteDomain& operator=(DiscreteDomain&& x) = default;
 
+    /**
+     * @brief Copy a DiscreteDomain by reordering and slicing.
+     *
+     * An assign operator to build a DiscreteDomain from another compatible domain.
+     * A domain is compatible if it either contains the same dimensions as this
+     * domain (even if they are in a different order) or if it contains at
+     * the dimensions of this domain plus some additional dimensions which will
+     * be unused here.
+     *
+     * @param domain A compatible domain.
+     */
+    template <class... ODDims>
+    DiscreteDomain& KOKKOS_FUNCTION operator=(DiscreteDomain<ODDims...> const& domain)
+    {
+        m_element_begin = DiscreteElement<DDims...>(domain.front());
+        m_element_end = m_element_begin + DiscreteVector<DDims...>(domain.extents());
+        return *this;
+    }
+
     template <class... ODims>
     KOKKOS_FUNCTION constexpr bool operator==(DiscreteDomain<ODims...> const& other) const
     {

diff --git a/tests/discrete_domain.cpp b/tests/discrete_domain.cpp
@@ -54,6 +54,9 @@ using DElemXYZ = ddc::DiscreteElement<DDimX, DDimY, DDimZ>;
 using DVectXYZ = ddc::DiscreteVector<DDimX, DDimY, DDimZ>;
 using DDomXYZ = ddc::DiscreteDomain<DDimX, DDimY, DDimZ>;
 
+using DElemZYX = ddc::DiscreteElement<DDimZ, DDimY, DDimX>;
+using DVectZYX = ddc::DiscreteVector<DDimZ, DDimY, DDimX>;
+using DDomZYX = ddc::DiscreteDomain<DDimZ, DDimY, DDimX>;
 
 static DElemX constexpr lbound_x(50);
 static DVectX constexpr nelems_x(3);
@@ -234,3 +237,34 @@ TEST(ProductMDomainTest, SliceDomainXToolate)
             R"rgx([Aa]ssert.*uid<ODDims>\(m_element_end\).*uid<ODDims>\(odomain\.m_element_end\).*)rgx");
 #endif
 }
+
+TEST(ProductMDomainTest, Transpose3DConstructor)
+{
+    DDomX const dom_x(lbound_x, nelems_x);
+    DDomY const dom_y(lbound_y, nelems_y);
+    DDomZ const dom_z(lbound_z, nelems_z);
+    DDomXYZ const dom_x_y_z(dom_x, dom_y, dom_z);
+    DDomZYX const dom_z_y_x(dom_x_y_z);
+    EXPECT_EQ(ddc::select<DDimX>(dom_x_y_z.front()), ddc::select<DDimX>(dom_z_y_x.front()));
+    EXPECT_EQ(ddc::select<DDimY>(dom_x_y_z.front()), ddc::select<DDimY>(dom_z_y_x.front()));
+    EXPECT_EQ(ddc::select<DDimZ>(dom_x_y_z.front()), ddc::select<DDimZ>(dom_z_y_x.front()));
+    EXPECT_EQ(ddc::select<DDimX>(dom_x_y_z.back()), ddc::select<DDimX>(dom_z_y_x.back()));
+    EXPECT_EQ(ddc::select<DDimY>(dom_x_y_z.back()), ddc::select<DDimY>(dom_z_y_x.back()));
+    EXPECT_EQ(ddc::select<DDimZ>(dom_x_y_z.back()), ddc::select<DDimZ>(dom_z_y_x.back()));
+}
+
+TEST(ProductMDomainTest, Transpose3DAssign)
+{
+    DDomX const dom_x(lbound_x, nelems_x);
+    DDomY const dom_y(lbound_y, nelems_y);
+    DDomZ const dom_z(lbound_z, nelems_z);
+    DDomXYZ const dom_x_y_z(dom_x, dom_y, dom_z);
+    DDomZYX dom_z_y_x;
+    dom_z_y_x = dom_x_y_z;
+    EXPECT_EQ(ddc::select<DDimX>(dom_x_y_z.front()), ddc::select<DDimX>(dom_z_y_x.front()));
+    EXPECT_EQ(ddc::select<DDimY>(dom_x_y_z.front()), ddc::select<DDimY>(dom_z_y_x.front()));
+    EXPECT_EQ(ddc::select<DDimZ>(dom_x_y_z.front()), ddc::select<DDimZ>(dom_z_y_x.front()));
+    EXPECT_EQ(ddc::select<DDimX>(dom_x_y_z.back()), ddc::select<DDimX>(dom_z_y_x.back()));
+    EXPECT_EQ(ddc::select<DDimY>(dom_x_y_z.back()), ddc::select<DDimY>(dom_z_y_x.back()));
+    EXPECT_EQ(ddc::select<DDimZ>(dom_x_y_z.back()), ddc::select<DDimZ>(dom_z_y_x.back()));
+}