Line data    Source code

       1             : #ifndef AMREX_GPU_RANGE_H_
       2             : #define AMREX_GPU_RANGE_H_
       3             : #include <AMReX_Config.H>
       4             : 
       5             : #include <AMReX_GpuQualifiers.H>
       6             : #include <AMReX_GpuControl.H>
       7             : #include <AMReX_GpuError.H>
       8             : #include <AMReX_Box.H>
       9             : #include <AMReX_TypeTraits.H>
      10             : 
      11             : namespace amrex {
      12             : 
      13             : template <typename T, std::enable_if_t<std::is_integral_v<T>,int> = 0>
      14             : bool isEmpty (T n) noexcept { return n <= 0; }
      15             : 
      16             : AMREX_FORCE_INLINE bool isEmpty (Box const& b) noexcept { return b.isEmpty(); }
      17             : 
      18             : namespace Gpu {
      19             : 
      20             : namespace range_detail {
      21             : 
      22             : //! integer version
      23             : template <typename T, std::enable_if_t<std::is_integral_v<T>,int> = 0>
      24             : AMREX_GPU_HOST_DEVICE
      25             : Long size (T const& b) noexcept { return static_cast<Long>(b); }
      26             : 
      27             : template <typename T, std::enable_if_t<std::is_integral_v<T>,int> = 0>
      28             : AMREX_GPU_HOST_DEVICE
      29             : Long at (T const& /*b*/, Long offset) noexcept { return offset; }
      30             : 
      31             : //! Box version
      32             : AMREX_GPU_HOST_DEVICE
      33             : AMREX_FORCE_INLINE Long size (Box const& b) noexcept
      34             : {
      35             :     AMREX_IF_ON_DEVICE((return b.numPts();))
      36           0 :     AMREX_IF_ON_HOST((
      37             :         amrex::ignore_unused(b);
      38             :         return 1;
      39             :     ))
      40             : }
      41             : 
      42             : AMREX_GPU_HOST_DEVICE
      43             : AMREX_FORCE_INLINE Box at (Box const& b, Long offset) noexcept
      44             : {
      45             :     AMREX_IF_ON_DEVICE((
      46             :         auto len = b.length3d();
      47             :         Long k = offset / (len[0]*len[1]);
      48             :         Long j = (offset - k*(len[0]*len[1])) / len[0];
      49             :         Long i = (offset - k*(len[0]*len[1])) - j*len[0];
      50             :         IntVect iv{AMREX_D_DECL(static_cast<int>(i),
      51             :                                 static_cast<int>(j),
      52             :                                 static_cast<int>(k))};
      53             :         iv += b.smallEnd();
      54             :         return Box(iv,iv,b.ixType());
      55             :     ))
      56           0 :     AMREX_IF_ON_HOST((
      57             :         amrex::ignore_unused(offset);
      58             :         return b;
      59             :     ))
      60             : }
      61             : 
      62             : template <typename T>
      63             : struct range_impl
      64             : {
      65             :     AMREX_GPU_HOST_DEVICE
      66           0 :     explicit range_impl (T const& b) noexcept : m_b(b), m_n(range_detail::size(b)) {}
      67             : 
      68             : #ifdef AMREX_USE_SYCL
      69             :     range_impl (T const& b, Long gid, Long grange) noexcept
      70             :         : m_b(b), m_n(range_detail::size(b)), m_gid(gid), m_grange(grange) {}
      71             : #endif
      72             : 
      73             :     struct iterator
      74             :     {
      75             :         AMREX_GPU_HOST_DEVICE
      76           0 :         iterator (T const& b, Long i, Long s) noexcept : mi_b(&b), mi_i(i), mi_s(s)  {}
      77             : 
      78             :         AMREX_GPU_HOST_DEVICE
      79           0 :         void operator++ () noexcept { mi_i += mi_s; }
      80             : 
      81             :         AMREX_GPU_HOST_DEVICE
      82           0 :         bool operator!= (iterator const& rhs) const noexcept { return mi_i < rhs.mi_i; }
      83             : 
      84             :         AMREX_GPU_HOST_DEVICE
      85           0 :         T operator* () const noexcept { return range_detail::at(*mi_b,mi_i); }
      86             : 
      87             :     private:
      88             :         T const* mi_b;
      89             :         Long mi_i;
      90             :         Long mi_s;
      91             :     };
      92             : 
      93             :     [[nodiscard]] AMREX_GPU_HOST_DEVICE
      94           0 :     iterator begin () const noexcept {
      95             : #if defined (__SYCL_DEVICE_ONLY__)
      96             :         return iterator(m_b, m_gid, m_grange);
      97             : #else
      98             :         AMREX_IF_ON_DEVICE((
      99             :             return iterator(m_b, blockDim.x*blockIdx.x+threadIdx.x, blockDim.x*gridDim.x);
     100             :         ))
     101           0 :         AMREX_IF_ON_HOST((
     102             :             return iterator(m_b,0,1);
     103             :         ))
     104             : #endif
     105             :     }
     106             : 
     107             :     [[nodiscard]] AMREX_GPU_HOST_DEVICE
     108           0 :     iterator end () const noexcept { return iterator(m_b,m_n,0); }
     109             : 
     110             : private:
     111             :     T m_b;
     112             :     Long m_n;
     113             : #ifdef AMREX_USE_SYCL
     114             :     Long m_gid;
     115             :     Long m_grange;
     116             : #endif
     117             : };
     118             : }
     119             : 
     120             : #ifdef AMREX_USE_SYCL
     121             : template <typename T>
     122             : range_detail::range_impl<T> Range (T const& b, Long gid, Long grange) noexcept {
     123             :     return range_detail::range_impl<T>(b,gid,grange);
     124             : }
     125             : #endif
     126             : 
     127             : template <typename T>
     128             : AMREX_GPU_HOST_DEVICE
     129           0 : range_detail::range_impl<T> Range (T const& b) noexcept { return range_detail::range_impl<T>(b); }
     130             : 
     131             : }}
     132             : 
     133             : #endif