Skip to content

Commit

Permalink
Non const iterator neighbor list (#994)
Browse files Browse the repository at this point in the history 
* build non-const iterator for neighbor lists

* build non-const iterator for neighbor lists

* implement non-const iterators for neighbor list

* remove second access operator for Neighbors<ParticleType>::iterator

* const

Co-authored-by: Johannes Blaschke <jpblaschke@lbl.gov>
  • Loading branch information
@atmyers @JBlaschke
atmyers and JBlaschke authored Jun 10, 2020
1 parent 399917f commit 37224d3
Show file tree
Hide file tree
Showing 2 changed files with 99 additions and 52 deletions.
145 changes: 96 additions & 49 deletions Src/Particle/AMReX_NeighborList.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,77 +10,124 @@ namespace amrex

template <class ParticleType>
struct Neighbors
{
{
struct iterator
{
AMREX_GPU_HOST_DEVICE
iterator (int start, int stop, const unsigned int * nbor_list_ptr, const ParticleType* pstruct)
: m_index(start), m_stop(stop), m_nbor_list_ptr(nbor_list_ptr), m_pstruct(pstruct)
iterator (int start, int stop, const unsigned int * nbor_list_ptr, ParticleType* pstruct)
: m_index(start), m_stop(stop), m_nbor_list_ptr(nbor_list_ptr), m_pstruct(pstruct)
{}

AMREX_GPU_HOST_DEVICE
void operator++ () { ++m_index;; }

AMREX_GPU_HOST_DEVICE
bool operator!= (iterator const& rhs) const { return m_index < m_stop; }


AMREX_GPU_HOST_DEVICE
ParticleType& operator* () const { return m_pstruct[m_nbor_list_ptr[m_index]]; }

private:
int m_index;
int m_stop;
const unsigned int* m_nbor_list_ptr;
ParticleType* m_pstruct;
};

struct const_iterator
{
AMREX_GPU_HOST_DEVICE
const_iterator (int start, int stop, const unsigned int * nbor_list_ptr, const ParticleType* pstruct)
: m_index(start), m_stop(stop), m_nbor_list_ptr(nbor_list_ptr), m_pstruct(pstruct)
{}

AMREX_GPU_HOST_DEVICE
void operator++ () { ++m_index;; }

AMREX_GPU_HOST_DEVICE
bool operator!= (const_iterator const& rhs) const { return m_index < m_stop; }

AMREX_GPU_HOST_DEVICE
ParticleType operator* () const { return m_pstruct[m_nbor_list_ptr[m_index]]; }
const ParticleType& operator* () const { return m_pstruct[m_nbor_list_ptr[m_index]]; }

private:
int m_index;
int m_stop;
const unsigned int* m_nbor_list_ptr;
const ParticleType* m_pstruct;
};

AMREX_GPU_HOST_DEVICE
iterator begin () const {
return iterator(m_nbor_offsets_ptr[m_i], m_nbor_offsets_ptr[m_i+1],
iterator begin () noexcept {
return iterator(m_nbor_offsets_ptr[m_i], m_nbor_offsets_ptr[m_i+1],
m_nbor_list_ptr, m_pstruct);
}

AMREX_GPU_HOST_DEVICE
iterator end () const {
return iterator(m_nbor_offsets_ptr[m_i+1], m_nbor_offsets_ptr[m_i+1],
iterator end () noexcept {
return iterator(m_nbor_offsets_ptr[m_i+1], m_nbor_offsets_ptr[m_i+1],
m_nbor_list_ptr, m_pstruct);
}

AMREX_GPU_HOST_DEVICE
AMREX_GPU_HOST_DEVICE
const_iterator begin () const noexcept {
return const_iterator(m_nbor_offsets_ptr[m_i], m_nbor_offsets_ptr[m_i+1],
m_nbor_list_ptr, m_pstruct);
}

AMREX_GPU_HOST_DEVICE
const_iterator end () const noexcept {
return const_iterator(m_nbor_offsets_ptr[m_i+1], m_nbor_offsets_ptr[m_i+1],
m_nbor_list_ptr, m_pstruct);
}

AMREX_GPU_HOST_DEVICE
const_iterator cbegin () const noexcept {
return const_iterator(m_nbor_offsets_ptr[m_i], m_nbor_offsets_ptr[m_i+1],
m_nbor_list_ptr, m_pstruct);
}

AMREX_GPU_HOST_DEVICE
const_iterator cend () const noexcept {
return const_iterator(m_nbor_offsets_ptr[m_i+1], m_nbor_offsets_ptr[m_i+1],
m_nbor_list_ptr, m_pstruct);
}

AMREX_GPU_HOST_DEVICE
Neighbors (int i, const unsigned int *nbor_offsets_ptr, const unsigned int *nbor_list_ptr,
const ParticleType* pstruct)
ParticleType* pstruct)
: m_i(i),
m_nbor_offsets_ptr(nbor_offsets_ptr),
m_nbor_list_ptr(nbor_list_ptr),
m_pstruct(pstruct)
{}

private:

const int m_i;
const unsigned int * m_nbor_offsets_ptr;
const unsigned int * m_nbor_list_ptr;
const ParticleType * m_pstruct;
ParticleType * m_pstruct;
};

template <class ParticleType>
struct NeighborData
{
NeighborData (const Gpu::DeviceVector<unsigned int>& offsets,
NeighborData (const Gpu::DeviceVector<unsigned int>& offsets,
const Gpu::DeviceVector<unsigned int>& list,
const ParticleType* pstruct)
ParticleType* pstruct)
: m_nbor_offsets_ptr(offsets.dataPtr()),
m_nbor_list_ptr(list.dataPtr()),
m_pstruct(pstruct)
{}

AMREX_GPU_HOST_DEVICE
amrex::Neighbors<ParticleType> getNeighbors (int i) const
{ return Neighbors<ParticleType>(i, m_nbor_offsets_ptr, m_nbor_list_ptr, m_pstruct); }

const unsigned int * m_nbor_offsets_ptr;
const unsigned int * m_nbor_list_ptr;
const ParticleType * m_pstruct;
ParticleType * m_pstruct;
};

template <class ParticleType>
Expand All @@ -89,20 +136,20 @@ class NeighborList
public:

template <class PTile, class CheckPair>
void build (const PTile& ptile,
void build (PTile& ptile,
const amrex::Box& bx, const amrex::Geometry& geom,
CheckPair check_pair, int num_cells=1)
{
const auto& vec = ptile.GetArrayOfStructs()();
auto& vec = ptile.GetArrayOfStructs()();
m_pstruct = vec.dataPtr();

const auto dxi = geom.InvCellSizeArray();
const auto plo = geom.ProbLoArray();

const size_t np_real = ptile.numRealParticles();
const size_t np_total = vec.size();
const ParticleType* pstruct_ptr = vec.dataPtr();

const auto lo = lbound(bx);
const auto hi = ubound(bx);
m_bins.build(np_total, pstruct_ptr, bx,
Expand All @@ -111,7 +158,7 @@ public:
AMREX_D_TERM(AMREX_ASSERT((p.pos(0)-plo[0])*dxi[0] - lo.x >= 0.0);,
AMREX_ASSERT((p.pos(1)-plo[1])*dxi[1] - lo.y >= 0.0);,
AMREX_ASSERT((p.pos(2)-plo[2])*dxi[2] - lo.z >= 0.0));

return IntVect(AMREX_D_DECL(static_cast<int>((p.pos(0)-plo[0])*dxi[0] - lo.x),
static_cast<int>((p.pos(1)-plo[1])*dxi[1] - lo.y),
static_cast<int>((p.pos(2)-plo[2])*dxi[2] - lo.z)));
Expand All @@ -131,13 +178,13 @@ public:
int ix = (pstruct_ptr[i].pos(0)-plo[0])*dxi[0] - lo.x;
int iy = (pstruct_ptr[i].pos(1)-plo[1])*dxi[1] - lo.y;
int iz = (pstruct_ptr[i].pos(2)-plo[2])*dxi[2] - lo.z;

int nx = hi.x-lo.x+1;
int ny = hi.y-lo.y+1;
int nz = hi.z-lo.z+1;
int nz = hi.z-lo.z+1;

int count = 0;

for (int ii = amrex::max(ix-num_cells, 0); ii <= amrex::min(ix+num_cells, nx-1); ++ii) {
for (int jj = amrex::max(iy-num_cells, 0); jj <= amrex::min(iy+num_cells, ny-1); ++jj) {
for (int kk = amrex::max(iz-num_cells, 0); kk <= amrex::min(iz+num_cells, nz-1); ++kk) {
Expand Down Expand Up @@ -173,20 +220,20 @@ public:
int ix = (pstruct_ptr[i].pos(0)-plo[0])*dxi[0] - lo.x;
int iy = (pstruct_ptr[i].pos(1)-plo[1])*dxi[1] - lo.y;
int iz = (pstruct_ptr[i].pos(2)-plo[2])*dxi[2] - lo.z;

int nx = hi.x-lo.x+1;
int ny = hi.y-lo.y+1;
int nz = hi.z-lo.z+1;
int n = 0;
int nz = hi.z-lo.z+1;

int n = 0;
for (int ii = amrex::max(ix-num_cells, 0); ii <= amrex::min(ix+num_cells, nx-1); ++ii) {
for (int jj = amrex::max(iy-num_cells, 0); jj <= amrex::min(iy+num_cells, ny-1); ++jj) {
for (int kk = amrex::max(iz-num_cells, 0); kk <= amrex::min(iz+num_cells, nz-1); ++kk) {
int index = (ii * ny + jj) * nz + kk;
for (auto p = poffset[index]; p < poffset[index+1]; ++p) {
if (pperm[p] == i) continue;
if (check_pair(pstruct_ptr[i], pstruct_ptr[pperm[p]])) {
pm_nbor_list[pnbor_offset[i] + n] = pperm[p];
pm_nbor_list[pnbor_offset[i] + n] = pperm[p];
++n;
}
}
Expand All @@ -195,10 +242,10 @@ public:
}
});
}
NeighborData<ParticleType> data ()
{
return NeighborData<ParticleType>(m_nbor_offsets, m_nbor_list, m_pstruct);

NeighborData<ParticleType> data ()
{
return NeighborData<ParticleType>(m_nbor_offsets, m_nbor_list, m_pstruct);
}

int numParticles () { return m_nbor_offsets.size() - 1; }
Expand All @@ -211,17 +258,17 @@ public:

Gpu::DeviceVector<unsigned int>& GetList () { return m_nbor_list; }
const Gpu::DeviceVector<unsigned int>& GetList () const { return m_nbor_list; }

void print ()
{
BL_PROFILE("NeighborList::print");

Gpu::HostVector<unsigned int> host_nbor_offsets(m_nbor_offsets.size());
Gpu::HostVector<unsigned int> host_nbor_list(m_nbor_list.size());
Gpu::copy(Gpu::deviceToHost, m_nbor_offsets.begin(), m_nbor_offsets.end(), host_nbor_offsets.begin());

Gpu::copy(Gpu::deviceToHost, m_nbor_offsets.begin(), m_nbor_offsets.end(), host_nbor_offsets.begin());
Gpu::copy(Gpu::deviceToHost, m_nbor_list.begin(), m_nbor_list.end(), host_nbor_list.begin());

for (int i = 0; i < numParticles(); ++i) {
amrex::Print() << "Particle " << i << " could collide with: ";
for (int j = host_nbor_offsets[i]; j < host_nbor_offsets[i+1]; ++j) {
Expand All @@ -232,14 +279,14 @@ public:
}

protected:
const ParticleType* m_pstruct;

ParticleType* m_pstruct;

// This is the neighbor list data structure
Gpu::DeviceVector<unsigned int> m_nbor_offsets;
Gpu::DeviceVector<unsigned int> m_nbor_list;
Gpu::DeviceVector<unsigned int> m_nbor_counts;

DenseBins<ParticleType> m_bins;
};

Expand Down
6 changes: 3 additions & 3 deletions Tests/Particles/NeighborParticles/MDParticleContainer.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -141,15 +141,15 @@ std::pair<Real, Real> MDParticleContainer::minAndMaxDistance()
ParticleType& p1 = pstruct[i];

for (const auto& p2 : nbor_data.getNeighbors(i))
{
{
Real dx = p1.pos(0) - p2.pos(0);
Real dy = p1.pos(1) - p2.pos(1);
Real dz = p1.pos(2) - p2.pos(2);

Real r2 = dx*dx + dy*dy + dz*dz;
r2 = amrex::max(r2, Params::min_r*Params::min_r);
Real r = sqrt(r2);

Gpu::Atomic::Min(pmin_d, r);
Gpu::Atomic::Max(pmax_d, r);
}
Expand Down

0 comments on commit 37224d3

Please sign in to comment.