aboutsummaryrefslogtreecommitdiff
path: root/Source/Particles/Sorting/SortingUtils.H
blob: 1133ccab54cb9df07156b5e65b9c23e24b03ae77 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
#ifndef WARPX_PARTICLES_SORTING_SORTINGUTILS_H_
#define WARPX_PARTICLES_SORTING_SORTINGUTILS_H_

#include <WarpXParticleContainer.H>
#include <AMReX_CudaContainers.H>
#include <AMReX_Gpu.H>

// TODO: Add documentation
void fillWithConsecutiveIntegers( amrex::Gpu::ManagedDeviceVector<long>& v ) {
#ifdef AMREX_USE_GPU
    // On GPU: Use thrust
    thrust::sequence( v.begin(), v.end() );
#else
    // On CPU: Use std library
    std::iota( v.begin(), v.end(), 0L );
#endif
}

// TODO: Add documentation
class fillBufferFlag
{
    public:
        fillBufferFlag( WarpXParIter& pti, const amrex::iMultiFab* bmasks,
                        amrex::Gpu::ManagedDeviceVector<int>& inexflag,
                        const amrex::Geometry& geom ) {

            // Extract simple structure that can be used directly on the GPU
            m_particles = &(pti.GetArrayOfStructs()[0]);
            m_buffer_mask = (*bmasks)[pti].array();
            m_inexflag_ptr = inexflag.dataPtr();
            m_domain_small_end = geom.Domain().smallEnd();
            for (int idim=0; idim<AMREX_SPACEDIM; idim++) {
                m_prob_lo[idim] = geom.ProbLo(idim);
                m_inv_cell_size[idim] = geom.InvCellSize(idim);
            }
        };


        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        void operator()( const long i ) const {
            // Select a particle
            auto& p = m_particles[i];
            // Find the index of the cell where this particle is located
            amrex::IntVect iv;
            for (int idim=0; idim<AMREX_SPACEDIM; idim++) {
                iv[idim] = static_cast<int>(floor(
                    (p.m_rdata.pos[idim]-m_prob_lo[idim])*m_inv_cell_size[idim]
                ));
            }
            // Find the value of the buffer flag in this cell and
            // store it at the corresponding particle position in the array `inexflag`
            m_inexflag_ptr[i] = m_buffer_mask(iv);
        };

    private:
        amrex::Real m_prob_lo[AMREX_SPACEDIM];
        amrex::Real m_inv_cell_size[AMREX_SPACEDIM];
        amrex::IntVect m_domain_small_end;
        int* m_inexflag_ptr;
        WarpXParticleContainer::ParticleType* m_particles;
        amrex::Array4<const int> m_buffer_mask;
};

#endif // WARPX_PARTICLES_SORTING_SORTINGUTILS_H_