From 2a3cc05dac916961b1a5ae4c18b21bacd889e7fc Mon Sep 17 00:00:00 2001
From: Weiqun Zhang <WeiqunZhang@lbl.gov>
Date: Fri, 23 Sep 2022 12:24:05 -0700
Subject: [PATCH] CellData: data in a single cell (#2959)

This adds struct CellData that allows for accessing data in a single cell in
Array4.  This is convenient sometimes because one can omit the i, j and k
indices.  It might also be faster sometimes because it can skip the repeated
index calculation involving i,j,k.
---
 Src/Base/AMReX_Array4.H | 49 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/Src/Base/AMReX_Array4.H b/Src/Base/AMReX_Array4.H
index 0fc4c049437..296762614d3 100644
--- a/Src/Base/AMReX_Array4.H
+++ b/Src/Base/AMReX_Array4.H
@@ -11,6 +11,50 @@
 
 namespace amrex {
 
+    template <typename T>
+    struct CellData // Data in a single cell
+    {
+        T* AMREX_RESTRICT p = nullptr;
+        Long stride = 0;
+        int ncomp = 0;
+
+        AMREX_GPU_HOST_DEVICE
+        constexpr CellData (T* a_p, Long a_stride, int a_ncomp)
+            : p(a_p), stride(a_stride), ncomp(a_ncomp)
+            {}
+
+        template <class U=T,
+                  std::enable_if_t<std::is_const<U>::value,int> = 0>
+        AMREX_GPU_HOST_DEVICE
+        constexpr CellData (CellData<typename std::remove_const<T>::type> const& rhs) noexcept
+            : p(rhs.p), stride(rhs.stride), ncomp(rhs.ncomp)
+            {}
+
+        AMREX_GPU_HOST_DEVICE
+        explicit operator bool() const noexcept { return p != nullptr; }
+
+        AMREX_GPU_HOST_DEVICE
+        int nComp() const noexcept { return ncomp; }
+
+        template <class U=T,
+                  std::enable_if_t<!std::is_void<U>::value,int> = 0>
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        U& operator[] (int n) const noexcept {
+#if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
+            if (n < 0 || n >= ncomp) {
+#if AMREX_DEVICE_COMPILE
+                AMREX_DEVICE_PRINTF(" %d is out of bound (0:%d)", n, ncomp-1);
+#else
+                std::stringstream ss;
+                ss << " " << n << " is out of bound: (0:" << ncomp-1 << ")";
+                amrex::Abort(ss.str());
+#endif
+            }
+#endif
+            return p[n*stride];
+        }
+    };
+
     template <typename T>
     struct Array4
     {
@@ -207,6 +251,11 @@ namespace amrex {
             }
         }
 #endif
+
+        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+        CellData<T> cellData (int i, int j, int k) const noexcept {
+            return CellData<T>{this->ptr(i,j,k), nstride, ncomp};
+        }
     };
 
     template <class Tto, class Tfrom>