From 69743347fae589cadcfb71b9edd450f92e0dd6c4 Mon Sep 17 00:00:00 2001
From: Weiqun Zhang <WeiqunZhang@lbl.gov>
Date: Fri, 14 Jun 2024 09:37:04 -0700
Subject: [PATCH] TableData: Add support for row-major order (amrex::Order::C)
 (#3973)

The default is still column-major (i.e., Fortran order,
amrex::Order::F).
---
 Src/Base/AMReX_TableData.H | 254 +++++++++++++++++++++++++------------
 1 file changed, 170 insertions(+), 84 deletions(-)

diff --git a/Src/Base/AMReX_TableData.H b/Src/Base/AMReX_TableData.H
index 8ff5f608334..9c2a0d06ca1 100644
--- a/Src/Base/AMReX_TableData.H
+++ b/Src/Base/AMReX_TableData.H
@@ -72,11 +72,11 @@ struct Table1D
 #endif
 };
 
-template <typename T>
+template <typename T, typename ORDER = Order::F>
 struct Table2D
 {
     T* AMREX_RESTRICT p = nullptr;
-    Long jstride = 0;
+    Long stride1 = 0;
     GpuArray<int,2> begin{{1,1}};
     GpuArray<int,2> end{{0,0}};
 
@@ -84,9 +84,9 @@ struct Table2D
 
     template <class U=T, std::enable_if_t<std::is_const_v<U>,int> = 0>
     AMREX_GPU_HOST_DEVICE
-    constexpr Table2D (Table2D<std::remove_const_t<T>> const& rhs) noexcept
+    constexpr Table2D (Table2D<std::remove_const_t<T>, ORDER> const& rhs) noexcept
         : p(rhs.p),
-          jstride(rhs.jstride),
+          stride1(rhs.stride1),
           begin(rhs.begin),
           end(rhs.end)
         {}
@@ -96,7 +96,7 @@ struct Table2D
                        GpuArray<int,2> const& a_begin,
                        GpuArray<int,2> const& a_end) noexcept
         : p(a_p),
-          jstride(a_end[0]-a_begin[0]),
+          stride1(len0(a_begin,a_end)),
           begin(a_begin),
           end(a_end)
         {}
@@ -110,7 +110,13 @@ struct Table2D
 #if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
         index_assert(i,j);
 #endif
-        return p[(i-begin[0])+(j-begin[1])*jstride];
+        static_assert(std::is_same_v<ORDER,Order::F> ||
+                      std::is_same_v<ORDER,Order::C>);
+        if constexpr (std::is_same_v<ORDER,Order::F>) {
+            return p[(i-begin[0])+(j-begin[1])*stride1];
+        } else {
+            return p[(i-begin[0])*stride1+(j-begin[1])];
+        }
     }
 
 #if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
@@ -134,14 +140,26 @@ struct Table2D
         }
     }
 #endif
+
+private:
+
+    static constexpr int len0 (GpuArray<int,2> const& a_begin,
+                               GpuArray<int,2> const& a_end) noexcept
+    {
+        if constexpr (std::is_same_v<ORDER,Order::F>) {
+            return a_end[0] - a_begin[0];
+        } else {
+            return a_end[1] - a_begin[1];
+        }
+    }
 };
 
-template <typename T>
+template <typename T, typename ORDER = Order::F>
 struct Table3D
 {
     T* AMREX_RESTRICT p = nullptr;
-    Long jstride = 0;
-    Long kstride = 0;
+    Long stride1 = 0;
+    Long stride2 = 0;
     GpuArray<int,3> begin{{1,1,1}};
     GpuArray<int,3> end{{0,0,0}};
 
@@ -149,10 +167,10 @@ struct Table3D
 
     template <class U=T, std::enable_if_t<std::is_const_v<U>,int> = 0>
     AMREX_GPU_HOST_DEVICE
-    constexpr Table3D (Table3D<std::remove_const_t<T>> const& rhs) noexcept
+    constexpr Table3D (Table3D<std::remove_const_t<T>,ORDER> const& rhs) noexcept
         : p(rhs.p),
-          jstride(rhs.jstride),
-          kstride(rhs.kstride),
+          stride1(rhs.stride1),
+          stride2(rhs.stride2),
           begin(rhs.begin),
           end(rhs.end)
         {}
@@ -162,8 +180,8 @@ struct Table3D
                        GpuArray<int,3> const& a_begin,
                        GpuArray<int,3> const& a_end) noexcept
         : p(a_p),
-          jstride(a_end[0]-a_begin[0]),
-          kstride(jstride*(a_end[1]-a_begin[1])),
+          stride1(        len0(a_begin,a_end)),
+          stride2(stride1*len1(a_begin,a_end)),
           begin(a_begin),
           end(a_end)
         {}
@@ -177,7 +195,13 @@ struct Table3D
 #if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
         index_assert(i,j,k);
 #endif
-        return p[(i-begin[0])+(j-begin[1])*jstride+(k-begin[2])*kstride];
+        static_assert(std::is_same_v<ORDER,Order::F> ||
+                      std::is_same_v<ORDER,Order::C>);
+        if constexpr (std::is_same_v<ORDER,Order::F>) {
+            return p[(i-begin[0])+(j-begin[1])*stride1+(k-begin[2])*stride2];
+        } else {
+            return p[(i-begin[0])*stride2+(j-begin[1])*stride1+(k-begin[2])];
+        }
     }
 
 #if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
@@ -204,15 +228,33 @@ struct Table3D
         }
     }
 #endif
+
+private:
+
+    static constexpr int len0 (GpuArray<int,3> const& a_begin,
+                               GpuArray<int,3> const& a_end) noexcept
+    {
+        if constexpr (std::is_same_v<ORDER,Order::F>) {
+            return a_end[0] - a_begin[0];
+        } else {
+            return a_end[2] - a_begin[2];
+        }
+    }
+
+    static constexpr int len1 (GpuArray<int,3> const& a_begin,
+                               GpuArray<int,3> const& a_end) noexcept
+    {
+        return a_end[1] - a_begin[1];
+    }
 };
 
-template <typename T>
+template <typename T, typename ORDER = Order::F>
 struct Table4D
 {
     T* AMREX_RESTRICT p = nullptr;
-    Long jstride = 0;
-    Long kstride = 0;
-    Long nstride = 0;
+    Long stride1 = 0;
+    Long stride2 = 0;
+    Long stride3 = 0;
     GpuArray<int,4> begin{{1,1,1,1}};
     GpuArray<int,4> end{{0,0,0,0}};
 
@@ -220,11 +262,11 @@ struct Table4D
 
     template <class U=T, std::enable_if_t<std::is_const_v<U>,int> = 0>
     AMREX_GPU_HOST_DEVICE
-    constexpr Table4D (Table4D<std::remove_const_t<T>> const& rhs) noexcept
+    constexpr Table4D (Table4D<std::remove_const_t<T>,ORDER> const& rhs) noexcept
         : p(rhs.p),
-          jstride(rhs.jstride),
-          kstride(rhs.kstride),
-          nstride(rhs.nstride),
+          stride1(rhs.stride1),
+          stride2(rhs.stride2),
+          stride3(rhs.stride3),
           begin(rhs.begin),
           end(rhs.end)
         {}
@@ -234,9 +276,9 @@ struct Table4D
                        GpuArray<int,4> const& a_begin,
                        GpuArray<int,4> const& a_end) noexcept
         : p(a_p),
-          jstride(a_end[0]-a_begin[0]),
-          kstride(jstride*(a_end[1]-a_begin[1])),
-          nstride(kstride*(a_end[2]-a_begin[2])),
+          stride1(        len0(a_begin,a_end)),
+          stride2(stride1*len1(a_begin,a_end)),
+          stride3(stride2*len2(a_begin,a_end)),
           begin(a_begin),
           end(a_end)
         {}
@@ -250,7 +292,13 @@ struct Table4D
 #if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
         index_assert(i,j,k,n);
 #endif
-        return p[(i-begin[0])+(j-begin[1])*jstride+(k-begin[2])*kstride+(n-begin[3])*nstride];
+        static_assert(std::is_same_v<ORDER,Order::F> ||
+                      std::is_same_v<ORDER,Order::C>);
+        if constexpr (std::is_same_v<ORDER,Order::F>) {
+            return p[(i-begin[0])+(j-begin[1])*stride1+(k-begin[2])*stride2+(n-begin[3])*stride3];
+        } else {
+            return p[(i-begin[0])*stride3+(j-begin[1])*stride2+(k-begin[2])*stride1+(n-begin[3])];
+        }
     }
 
 #if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
@@ -279,6 +327,38 @@ struct Table4D
         }
     }
 #endif
+
+private:
+
+    static constexpr int len0 (GpuArray<int,4> const& a_begin,
+                               GpuArray<int,4> const& a_end) noexcept
+    {
+        if constexpr (std::is_same_v<ORDER,Order::F>) {
+            return a_end[0] - a_begin[0];
+        } else {
+            return a_end[3] - a_begin[3];
+        }
+    }
+
+    static constexpr int len1 (GpuArray<int,4> const& a_begin,
+                               GpuArray<int,4> const& a_end) noexcept
+    {
+        if constexpr (std::is_same_v<ORDER,Order::F>) {
+            return a_end[1] - a_begin[1];
+        } else {
+            return a_end[2] - a_begin[2];
+        }
+    }
+
+    static constexpr int len2 (GpuArray<int,4> const& a_begin,
+                               GpuArray<int,4> const& a_end) noexcept
+    {
+        if constexpr (std::is_same_v<ORDER,Order::F>) {
+            return a_end[2] - a_begin[2];
+        } else {
+            return a_end[1] - a_begin[1];
+        }
+    }
 };
 
 /**
@@ -287,9 +367,12 @@ struct Table4D
  * This class is somewhat similar to FArrayBox/BaseFab. The main difference
  * is the dimension of the array in this class can be 1, 2, 3, or 4, whereas
  * the dimension of FArrayBox/BaseFab is the spatial dimension
- * (AMREX_SPACEDIM) plus a component dimension.  Below is an example of
- * using it to store a 3D table of data that is initialized on CPU and is
- * read-only by all GPU threads on the device.
+ * (AMREX_SPACEDIM) plus a component dimension. Another difference is that
+ * this class supports both column-major order (i.e., Fortran order) and
+ * row-major order (i.e., C order), whereas FArrayBox/BaseFab is always
+ * column-major. Below is an example of using it to store a 3D table of data
+ * that is initialized on CPU and is read-only by all GPU threads on the
+ * device.
  *
  * \code
  *      Array<int,3> tlo{0,0,0}; // lower bounds
@@ -316,22 +399,22 @@ struct Table4D
  *      // We can now use table in device lambda.
  * \endcode
  */
-template <typename T, int N>
+template <typename T, int N, typename ORDER = Order::F>
 class TableData
     : public DataAllocator
 {
 public:
 
-    template <class U, int M> friend class TableData;
+    template <class U, int M, class O> friend class TableData;
     using value_type = T;
     using table_type = std::conditional_t<N==1, Table1D<T>,
-                       std::conditional_t<N==2, Table2D<T>,
-                       std::conditional_t<N==3, Table3D<T>,
-                                                Table4D<T> > > >;
+                       std::conditional_t<N==2, Table2D<T, ORDER>,
+                       std::conditional_t<N==3, Table3D<T, ORDER>,
+                                                Table4D<T, ORDER> > > >;
     using const_table_type = std::conditional_t<N==1, Table1D<T const>,
-                             std::conditional_t<N==2, Table2D<T const>,
-                             std::conditional_t<N==3, Table3D<T const>,
-                                                      Table4D<T const> > > >;
+                             std::conditional_t<N==2, Table2D<T const, ORDER>,
+                             std::conditional_t<N==3, Table3D<T const, ORDER>,
+                                                      Table4D<T const, ORDER> > > >;
 
     TableData () noexcept = default;
 
@@ -339,11 +422,11 @@ public:
 
     TableData (Array<int,N> const& lo, Array<int,N> const& hi, Arena* ar = nullptr);
 
-    TableData (TableData<T,N> const&) = delete;
-    TableData<T,N>& operator= (TableData<T,N> const&) = delete;
+    TableData (TableData<T,N,ORDER> const&) = delete;
+    TableData<T,N,ORDER>& operator= (TableData<T,N,ORDER> const&) = delete;
 
-    TableData (TableData<T,N>&& rhs) noexcept;
-    TableData<T,N>& operator= (TableData<T,N> && rhs) noexcept;
+    TableData (TableData<T,N,ORDER>&& rhs) noexcept;
+    TableData<T,N,ORDER>& operator= (TableData<T,N,ORDER> && rhs) noexcept;
 
     ~TableData () noexcept;
 
@@ -359,7 +442,7 @@ public:
 
     void clear () noexcept;
 
-    void copy (TableData<T,N> const& rhs) noexcept;
+    void copy (TableData<T,N,ORDER> const& rhs) noexcept;
 
     table_type table () noexcept;
     const_table_type table () const noexcept;
@@ -376,16 +459,16 @@ private:
     bool m_ptr_owner = false;
 };
 
-template <typename T, int N>
-TableData<T,N>::TableData (Array<int,N> const& lo, Array<int,N> const& hi, Arena* ar)
+template <typename T, int N, typename ORDER>
+TableData<T,N,ORDER>::TableData (Array<int,N> const& lo, Array<int,N> const& hi, Arena* ar)
     : DataAllocator{ar}, m_lo(lo), m_hi(hi)
 {
     define();
 }
 
 
-template <typename T, int N>
-TableData<T,N>::TableData (TableData<T,N>&& rhs) noexcept
+template <typename T, int N, typename ORDER>
+TableData<T,N,ORDER>::TableData (TableData<T,N,ORDER>&& rhs) noexcept
     : DataAllocator{rhs.arena()},
       m_dptr(rhs.m_dptr),
       m_lo(rhs.m_lo),
@@ -397,9 +480,9 @@ TableData<T,N>::TableData (TableData<T,N>&& rhs) noexcept
     rhs.m_ptr_owner = false;
 }
 
-template <typename T, int N>
-TableData<T,N>&
-TableData<T,N>::operator= (TableData<T,N> && rhs) noexcept
+template <typename T, int N, typename ORDER>
+TableData<T,N,ORDER>&
+TableData<T,N,ORDER>::operator= (TableData<T,N,ORDER> && rhs) noexcept
 {
     if (this != &rhs) {
         clear();
@@ -415,19 +498,22 @@ TableData<T,N>::operator= (TableData<T,N> && rhs) noexcept
     return *this;
 }
 
-template <typename T, int N>
-TableData<T,N>::~TableData () noexcept
+template <typename T, int N, typename ORDER>
+TableData<T,N,ORDER>::~TableData () noexcept
 {
     static_assert(std::is_trivially_copyable<T>() &&
                   std::is_trivially_destructible<T>(),
-                  "TableData<T,N>: T must be trivially copyable and trivially destructible");
-    static_assert(N>=1 && N <=4, "TableData<T,N>: N must be in the range of [1,4]");
+                  "TableData<T,N,ORDER>: T must be trivially copyable and trivially destructible");
+    static_assert(N>=1 && N <=4, "TableData<T,N,ORDER>: N must be in the range of [1,4]");
+    static_assert(std::is_same_v<ORDER,Order::F> ||
+                  std::is_same_v<ORDER,Order::C>,
+                  "TableDat<T,N,ORDER>: ORDER must be either Order::F or Order::C");
     clear();
 }
 
-template <typename T, int N>
+template <typename T, int N, typename ORDER>
 void
-TableData<T,N>::resize (Array<int,N> const& lo, Array<int,N> const& hi, Arena* ar)
+TableData<T,N,ORDER>::resize (Array<int,N> const& lo, Array<int,N> const& hi, Arena* ar)
 {
     m_lo = lo;
     m_hi = hi;
@@ -449,9 +535,9 @@ TableData<T,N>::resize (Array<int,N> const& lo, Array<int,N> const& hi, Arena* a
     }
 }
 
-template <typename T, int N>
+template <typename T, int N, typename ORDER>
 Long
-TableData<T,N>::size () const noexcept
+TableData<T,N,ORDER>::size () const noexcept
 {
     Long r = 1;
     for (int i = 0; i < N; ++i) {
@@ -460,9 +546,9 @@ TableData<T,N>::size () const noexcept
     return r;
 }
 
-template <typename T, int N>
+template <typename T, int N, typename ORDER>
 void
-TableData<T,N>::clear () noexcept
+TableData<T,N,ORDER>::clear () noexcept
 {
     if (m_dptr) {
         if (m_ptr_owner) {
@@ -473,9 +559,9 @@ TableData<T,N>::clear () noexcept
     }
 }
 
-template <typename T, int N>
+template <typename T, int N, typename ORDER>
 void
-TableData<T,N>::define ()
+TableData<T,N,ORDER>::define ()
 {
     m_truesize = size();
     AMREX_ASSERT(m_truesize >= 0);
@@ -488,48 +574,48 @@ TableData<T,N>::define ()
 }
 
 namespace detail {
-    template <typename T>
+    template <typename T, typename>
     Table1D<T> make_table (T* p, Array<int,1> const& lo, Array<int,1> const& hi) {
         return Table1D<T>(p, lo[0], hi[0]+1);
     }
-    template <typename T>
-    Table2D<T> make_table (T* p, Array<int,2> const& lo, Array<int,2> const& hi) {
-        return Table2D<T>(p, {lo[0],lo[1]}, {hi[0]+1,hi[1]+1});
+    template <typename T, typename ORDER>
+    Table2D<T,ORDER> make_table (T* p, Array<int,2> const& lo, Array<int,2> const& hi) {
+        return Table2D<T,ORDER>(p, {lo[0],lo[1]}, {hi[0]+1,hi[1]+1});
     }
-    template <typename T>
+    template <typename T, typename ORDER>
     Table3D<T> make_table (T* p, Array<int,3> const& lo, Array<int,3> const& hi) {
-        return Table3D<T>(p, {lo[0],lo[1],lo[2]}, {hi[0]+1,hi[1]+1,hi[2]+1});
+        return Table3D<T,ORDER>(p, {lo[0],lo[1],lo[2]}, {hi[0]+1,hi[1]+1,hi[2]+1});
     }
-    template <typename T>
+    template <typename T, typename ORDER>
     Table4D<T> make_table (T* p, Array<int,4> const& lo, Array<int,4> const& hi) {
-        return Table4D<T>(p, {lo[0],lo[1],lo[2],lo[3]}, {hi[0]+1,hi[1]+1,hi[2]+1,hi[3]+1});
+        return Table4D<T,ORDER>(p, {lo[0],lo[1],lo[2],lo[3]}, {hi[0]+1,hi[1]+1,hi[2]+1,hi[3]+1});
     }
 }
 
-template <typename T, int N>
-typename TableData<T,N>::table_type
-TableData<T,N>::table () noexcept
+template <typename T, int N, typename ORDER>
+typename TableData<T,N,ORDER>::table_type
+TableData<T,N,ORDER>::table () noexcept
 {
-    return detail::make_table<T>(m_dptr, m_lo, m_hi);
+    return detail::make_table<T,ORDER>(m_dptr, m_lo, m_hi);
 }
 
-template <typename T, int N>
-typename TableData<T,N>::const_table_type
-TableData<T,N>::table () const noexcept
+template <typename T, int N, typename ORDER>
+typename TableData<T,N,ORDER>::const_table_type
+TableData<T,N,ORDER>::table () const noexcept
 {
-    return detail::make_table<T const>(m_dptr, m_lo, m_hi);
+    return detail::make_table<T const, ORDER>(m_dptr, m_lo, m_hi);
 }
 
-template <typename T, int N>
-typename TableData<T,N>::const_table_type
-TableData<T,N>::const_table () const noexcept
+template <typename T, int N, typename ORDER>
+typename TableData<T,N,ORDER>::const_table_type
+TableData<T,N,ORDER>::const_table () const noexcept
 {
-    return detail::make_table<T const>(m_dptr, m_lo, m_hi);
+    return detail::make_table<T const, ORDER>(m_dptr, m_lo, m_hi);
 }
 
-template <typename T, int N>
+template <typename T, int N, typename ORDER>
 void
-TableData<T,N>::copy (TableData<T,N> const& rhs) noexcept
+TableData<T,N,ORDER>::copy (TableData<T,N,ORDER> const& rhs) noexcept
 {
     std::size_t count = sizeof(T)*size();
 #ifdef AMREX_USE_GPU