Line data    Source code

       1             : #ifndef AMREX_ARRAY4_H_
       2             : #define AMREX_ARRAY4_H_
       3             : #include <AMReX_Config.H>
       4             : 
       5             : #include <AMReX.H>
       6             : #include <AMReX_IntVect.H>
       7             : #include <AMReX_GpuPrint.H>
       8             : 
       9             : #include <iostream>
      10             : #include <sstream>
      11             : 
      12             : namespace amrex {
      13             : 
      14             :     template <typename T>
      15             :     struct CellData // Data in a single cell
      16             :     {
      17             :         T* AMREX_RESTRICT p = nullptr;
      18             :         Long stride = 0;
      19             :         int ncomp = 0;
      20             : 
      21             :         AMREX_GPU_HOST_DEVICE
      22             :         constexpr CellData (T* a_p, Long a_stride, int a_ncomp)
      23             :             : p(a_p), stride(a_stride), ncomp(a_ncomp)
      24             :             {}
      25             : 
      26             :         template <class U=T,
      27             :                   std::enable_if_t<std::is_const_v<U>,int> = 0>
      28             :         AMREX_GPU_HOST_DEVICE
      29             :         constexpr CellData (CellData<std::remove_const_t<T>> const& rhs) noexcept
      30             :             : p(rhs.p), stride(rhs.stride), ncomp(rhs.ncomp)
      31             :             {}
      32             : 
      33             :         AMREX_GPU_HOST_DEVICE
      34             :         explicit operator bool() const noexcept { return p != nullptr; }
      35             : 
      36             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE
      37             :         int nComp() const noexcept { return ncomp; }
      38             : 
      39             :         template <class U=T,
      40             :                   std::enable_if_t<!std::is_void_v<U>,int> = 0>
      41             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
      42             :         U& operator[] (int n) const noexcept {
      43             : #if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
      44             :             if (n < 0 || n >= ncomp) {
      45             :                 AMREX_IF_ON_DEVICE((
      46             :                     AMREX_DEVICE_PRINTF(" %d is out of bound (0:%d)", n, ncomp-1);
      47             :                 ))
      48             :                 AMREX_IF_ON_HOST((
      49             :                     std::stringstream ss;
      50             :                     ss << " " << n << " is out of bound: (0:" << ncomp-1 << ")";
      51             :                     amrex::Abort(ss.str());
      52             :                 ))
      53             :             }
      54             : #endif
      55             :             return p[n*stride];
      56             :         }
      57             :     };
      58             : 
      59             :     template <typename T>
      60             :     struct Array4
      61             :     {
      62             :         T* AMREX_RESTRICT p;
      63             :         Long jstride = 0;
      64             :         Long kstride = 0;
      65             :         Long nstride = 0;
      66             :         Dim3 begin{1,1,1};
      67             :         Dim3 end{0,0,0};  // end is hi + 1
      68             :         int  ncomp=0;
      69             : 
      70             :         AMREX_GPU_HOST_DEVICE
      71        1844 :         constexpr Array4 () noexcept : p(nullptr) {}
      72             : 
      73             :         template <class U=T, std::enable_if_t<std::is_const_v<U>,int> = 0>
      74             :         AMREX_GPU_HOST_DEVICE
      75     5337436 :         constexpr Array4 (Array4<std::remove_const_t<T>> const& rhs) noexcept
      76     5337436 :             : p(rhs.p),
      77     5337436 :               jstride(rhs.jstride),
      78     5337436 :               kstride(rhs.kstride),
      79     5337436 :               nstride(rhs.nstride),
      80             :               begin(rhs.begin),
      81             :               end(rhs.end),
      82     5337436 :               ncomp(rhs.ncomp)
      83     5337436 :             {}
      84             : 
      85             :         AMREX_GPU_HOST_DEVICE
      86     4644016 :         constexpr Array4 (T* a_p, Dim3 const& a_begin, Dim3 const& a_end, int a_ncomp) noexcept
      87             :             : p(a_p),
      88     4644016 :               jstride(a_end.x-a_begin.x),
      89     4644016 :               kstride(jstride*(a_end.y-a_begin.y)),
      90     4644016 :               nstride(kstride*(a_end.z-a_begin.z)),
      91             :               begin(a_begin),
      92             :               end(a_end),
      93     4644016 :               ncomp(a_ncomp)
      94     4644016 :             {}
      95             : 
      96             :         template <class U,
      97             :                   std::enable_if_t
      98             :                   <std::is_same_v<std::remove_const_t<T>,
      99             :                                 std::remove_const_t<U>>,int> = 0>
     100             :         AMREX_GPU_HOST_DEVICE
     101             :         constexpr Array4 (Array4<U> const& rhs, int start_comp) noexcept
     102             :             : p((T*)(rhs.p+start_comp*rhs.nstride)),
     103             :               jstride(rhs.jstride),
     104             :               kstride(rhs.kstride),
     105             :               nstride(rhs.nstride),
     106             :               begin(rhs.begin),
     107             :               end(rhs.end),
     108             :               ncomp(rhs.ncomp-start_comp)
     109             :         {}
     110             : 
     111             :         template <class U,
     112             :                   std::enable_if_t
     113             :                   <std::is_same_v<std::remove_const_t<T>,
     114             :                                 std::remove_const_t<U>>,int> = 0>
     115             :         AMREX_GPU_HOST_DEVICE
     116             :         constexpr Array4 (Array4<U> const& rhs, int start_comp, int num_comps) noexcept
     117             :             : p((T*)(rhs.p+start_comp*rhs.nstride)),
     118             :               jstride(rhs.jstride),
     119             :               kstride(rhs.kstride),
     120             :               nstride(rhs.nstride),
     121             :               begin(rhs.begin),
     122             :               end(rhs.end),
     123             :               ncomp(num_comps)
     124             :         {}
     125             : 
     126             :         AMREX_GPU_HOST_DEVICE
     127             :         explicit operator bool() const noexcept { return p != nullptr; }
     128             : 
     129             :         template <class U=T, std::enable_if_t<!std::is_void_v<U>,int> = 0>
     130             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     131             :         U& operator() (int i, int j, int k) const noexcept {
     132             : #if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
     133             :             index_assert(i,j,k,0);
     134             : #endif
     135    43024975 :             return p[(i-begin.x)+(j-begin.y)*jstride+(k-begin.z)*kstride];
     136             :         }
     137             : 
     138             :         template <class U=T, std::enable_if_t<!std::is_void_v<U>,int> = 0>
     139             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     140             :         U& operator() (int i, int j, int k, int n) const noexcept {
     141             : #if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
     142             :             index_assert(i,j,k,n);
     143             : #endif
     144   901356884 :             return p[(i-begin.x)+(j-begin.y)*jstride+(k-begin.z)*kstride+n*nstride];
     145             :         }
     146             : 
     147             :         template <class U=T, std::enable_if_t<!std::is_void_v<U>,int> = 0>
     148             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     149             :         T* ptr (int i, int j, int k) const noexcept {
     150             : #if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
     151             :             index_assert(i,j,k,0);
     152             : #endif
     153             :             return p + ((i-begin.x)+(j-begin.y)*jstride+(k-begin.z)*kstride);
     154             :         }
     155             : 
     156             :         template <class U=T, std::enable_if_t<!std::is_void_v<U>,int> = 0>
     157             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     158             :         T* ptr (int i, int j, int k, int n) const noexcept {
     159             : #if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
     160             :             index_assert(i,j,k,0);
     161             : #endif
     162             :             return p + ((i-begin.x)+(j-begin.y)*jstride+(k-begin.z)*kstride+n*nstride);
     163             :         }
     164             : 
     165             :         template <class U=T, std::enable_if_t<!std::is_void_v<U>,int> = 0>
     166             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     167             :         U& operator() (IntVect const& iv) const noexcept {
     168             : #if (AMREX_SPACEDIM == 1)
     169             :             return this->operator()(iv[0],0,0);
     170             : #elif (AMREX_SPACEDIM == 2)
     171             :             return this->operator()(iv[0],iv[1],0);
     172             : #else
     173             :             return this->operator()(iv[0],iv[1],iv[2]);
     174             : #endif
     175             :         }
     176             : 
     177             :         template <class U=T, std::enable_if_t<!std::is_void_v<U>,int> = 0>
     178             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     179             :         U& operator() (IntVect const& iv, int n) const noexcept {
     180             : #if (AMREX_SPACEDIM == 1)
     181             :             return this->operator()(iv[0],0,0,n);
     182             : #elif (AMREX_SPACEDIM == 2)
     183             :             return this->operator()(iv[0],iv[1],0,n);
     184             : #else
     185             :             return this->operator()(iv[0],iv[1],iv[2],n);
     186             : #endif
     187             :         }
     188             : 
     189             :         template <class U=T, std::enable_if_t<!std::is_void_v<U>,int> = 0>
     190             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     191             :         T* ptr (IntVect const& iv) const noexcept {
     192             : #if (AMREX_SPACEDIM == 1)
     193             :             return this->ptr(iv[0],0,0);
     194             : #elif (AMREX_SPACEDIM == 2)
     195             :             return this->ptr(iv[0],iv[1],0);
     196             : #else
     197             :             return this->ptr(iv[0],iv[1],iv[2]);
     198             : #endif
     199             :         }
     200             : 
     201             :         template <class U=T, std::enable_if_t<!std::is_void_v<U>,int> = 0>
     202             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     203             :         T* ptr (IntVect const& iv, int n) const noexcept {
     204             : #if (AMREX_SPACEDIM == 1)
     205             :             return this->ptr(iv[0],0,0,n);
     206             : #elif (AMREX_SPACEDIM == 2)
     207             :             return this->ptr(iv[0],iv[1],0,n);
     208             : #else
     209             :             return this->ptr(iv[0],iv[1],iv[2],n);
     210             : #endif
     211             :         }
     212             : 
     213             :         template <class U=T, std::enable_if_t<!std::is_void_v<U>,int> = 0>
     214             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     215             :         U& operator() (Dim3 const& cell) const noexcept {
     216             :             return this->operator()(cell.x,cell.y,cell.z);
     217             :         }
     218             : 
     219             :         template <class U=T, std::enable_if_t<!std::is_void_v<U>,int> = 0>
     220             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     221             :         U& operator() (Dim3 const& cell, int n) const noexcept {
     222             :             return this->operator()(cell.x,cell.y,cell.z,n);
     223             :         }
     224             : 
     225             :         template <class U=T, std::enable_if_t<!std::is_void_v<U>,int> = 0>
     226             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     227             :         T* ptr (Dim3 const& cell) const noexcept {
     228             :             return this->ptr(cell.x,cell.y,cell.z);
     229             :         }
     230             : 
     231             :         template <class U=T, std::enable_if_t<!std::is_void_v<U>,int> = 0>
     232             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     233             :         T* ptr (Dim3 const& cell, int n) const noexcept {
     234             :             return this->ptr(cell.x,cell.y,cell.z,n);
     235             :         }
     236             : 
     237             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     238             :         T* dataPtr () const noexcept {
     239             :             return this->p;
     240             :         }
     241             : 
     242             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     243             :         std::size_t size () const noexcept {
     244           0 :             return this->nstride * this->ncomp;
     245             :         }
     246             : 
     247             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     248           0 :         int nComp () const noexcept { return ncomp; }
     249             : 
     250             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     251             :         bool contains (int i, int j, int k) const noexcept {
     252             :             return (i>=begin.x && i<end.x && j>=begin.y && j<end.y && k>=begin.z && k<end.z);
     253             :         }
     254             : 
     255             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     256             :         bool contains (IntVect const& iv) const noexcept {
     257             :             return AMREX_D_TERM(   iv[0]>=begin.x && iv[0]<end.x,
     258             :                                 && iv[1]>=begin.y && iv[1]<end.y,
     259             :                                 && iv[2]>=begin.z && iv[2]<end.z);
     260             :         }
     261             : 
     262             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     263             :         bool contains (Dim3 const& cell) const noexcept {
     264             :             return this->contains(cell.x,cell.y,cell.z);
     265             :         }
     266             : 
     267             : #if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
     268             :         AMREX_GPU_HOST_DEVICE inline
     269             :         void index_assert (int i, int j, int k, int n) const
     270             :         {
     271             :             if (i<begin.x || i>=end.x || j<begin.y || j>=end.y || k<begin.z || k>=end.z
     272             :                 || n < 0 || n >= ncomp) {
     273             :                 AMREX_IF_ON_DEVICE((
     274             :                     AMREX_DEVICE_PRINTF(" (%d,%d,%d,%d) is out of bound (%d:%d,%d:%d,%d:%d,0:%d)\n",
     275             :                                         i, j, k, n, begin.x, end.x-1, begin.y, end.y-1,
     276             :                                         begin.z, end.z-1, ncomp-1);
     277             :                     amrex::Abort();
     278             :                 ))
     279             :                 AMREX_IF_ON_HOST((
     280             :                     std::stringstream ss;
     281             :                     ss << " (" << i << "," << j << "," << k << "," <<  n
     282             :                     << ") is out of bound ("
     283             :                     << begin.x << ":" << end.x-1 << ","
     284             :                     << begin.y << ":" << end.y-1 << ","
     285             :                     << begin.z << ":" << end.z-1 << ","
     286             :                     << "0:" << ncomp-1 << ")";
     287             :                     amrex::Abort(ss.str());
     288             :                 ))
     289             :             }
     290             :         }
     291             : #endif
     292             : 
     293             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     294             :         CellData<T> cellData (int i, int j, int k) const noexcept {
     295             :             return CellData<T>{this->ptr(i,j,k), nstride, ncomp};
     296             :         }
     297             :     };
     298             : 
     299             :     template <class Tto, class Tfrom>
     300             :     [[nodiscard]] AMREX_GPU_HOST_DEVICE
     301             :     Array4<Tto> ToArray4 (Array4<Tfrom> const& a_in) noexcept
     302             :     {
     303             :         return Array4<Tto>((Tto*)(a_in.p), a_in.begin, a_in.end, a_in.ncomp);
     304             :     }
     305             : 
     306             :     template <class T>
     307             :     [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     308             :     Dim3 lbound (Array4<T> const& a) noexcept
     309             :     {
     310             :         return a.begin;
     311             :     }
     312             : 
     313             :     template <class T>
     314             :     [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     315             :     Dim3 ubound (Array4<T> const& a) noexcept
     316             :     {
     317           0 :         return Dim3{a.end.x-1,a.end.y-1,a.end.z-1};
     318             :     }
     319             : 
     320             :     template <class T>
     321             :     [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     322             :     Dim3 length (Array4<T> const& a) noexcept
     323             :     {
     324             :         return Dim3{a.end.x-a.begin.x,a.end.y-a.begin.y,a.end.z-a.begin.z};
     325             :     }
     326             : 
     327             :     template <typename T>
     328             :     std::ostream& operator<< (std::ostream& os, const Array4<T>& a) {
     329             :         os << "((" << lbound(a) << ',' << ubound(a) << ")," << a.ncomp << ')';
     330             :         return os;
     331             :     }
     332             : 
     333             :     //
     334             :     // Type traits for detecting if a class has a size() constexpr function.
     335             :     //
     336             :     template <class A, class Enable = void> struct HasMultiComp : std::false_type {};
     337             :     //
     338             :     template <class B>
     339             :     struct HasMultiComp<B, std::enable_if_t<B().size() >= 1>>
     340             :         : std::true_type {};
     341             : 
     342             :     //
     343             :     // PolymorphicArray4 can be used to access both AoS and SoA with
     344             :     // (i,j,k,n).  Here SoA refers multi-component BaseFab, and AoS refers
     345             :     // to single-component BaseFab of multi-component GpuArray.
     346             :     //
     347             :     template <typename T>
     348             :     struct PolymorphicArray4
     349             :         : public Array4<T>
     350             :     {
     351             :         AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     352             :         PolymorphicArray4 (Array4<T> const& a)
     353             :             : Array4<T>{a} {}
     354             : 
     355             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     356             :         T& operator() (int i, int j, int k) const noexcept {
     357             :             return this->Array4<T>::operator()(i,j,k);
     358             :         }
     359             : 
     360             :         template <class U=T, std::enable_if_t< amrex::HasMultiComp<U>::value,int> = 0>
     361             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     362             :         typename U::reference_type
     363             :         operator() (int i, int j, int k, int n) const noexcept {
     364             :             return this->Array4<T>::operator()(i,j,k,0)[n];
     365             :         }
     366             : 
     367             :         template <class U=T, std::enable_if_t<!amrex::HasMultiComp<U>::value,int> = 0>
     368             :         [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     369             :         U& operator() (int i, int j, int k, int n) const noexcept {
     370             :             return this->Array4<T>::operator()(i,j,k,n);
     371             :         }
     372             :     };
     373             : 
     374             :     template <typename T>
     375             :     [[nodiscard]] PolymorphicArray4<T>
     376             :     makePolymorphic (Array4<T> const& a)
     377             :     {
     378             :         return PolymorphicArray4<T>(a);
     379             :     }
     380             : }
     381             : 
     382             : #endif