Skip to content

Commit ac3dc54

Browse files
streevestreeve
authored
Merge pull request #803 from aprokop/arborx-2.0
Update ArborX neighbor list generation to work with ArborX 2.0
2 parents 1187808 + b3a0a46 commit ac3dc54

File tree

3 files changed

+139
-28
lines changed

3 files changed

+139
-28
lines changed

‎.github/workflows/CI.yml‎

Lines changed: 26 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -153,17 +153,17 @@ jobs:
153153
cxx: 'g++'
154154
backend: 'OPENMP'
155155
cmake_build_type: 'Debug'
156-
kokkos_ver: '4.1.00'
157-
arborx: 'ArborX'
156+
kokkos_ver: '4.2.00'
157+
arborx: 'v1.7'
158158
heffte: 'OFF'
159159
hypre: 'OFF'
160160
coverage: 'OFF'
161161
- distro: 'ubuntu:latest'
162162
cxx: 'g++'
163163
backend: 'OPENMP'
164164
cmake_build_type: 'Release'
165-
kokkos_ver: '4.1.00'
166-
arborx: 'ArborX'
165+
kokkos_ver: '4.3.00'
166+
arborx: 'fee3d777c0325a90ab50671fca69f2dc98f2a882'
167167
heffte: 'OFF'
168168
hypre: 'OFF'
169169
coverage: 'OFF'
@@ -219,8 +219,8 @@ jobs:
219219
cxx: 'g++'
220220
backend: 'OPENMP'
221221
cmake_build_type: 'Debug'
222-
kokkos_ver: '4.1.00'
223-
arborx: 'ArborX'
222+
kokkos_ver: '4.2.00'
223+
arborx: 'v1.7'
224224
heffte: 'FFTW'
225225
hypre: 'HYPRE'
226226
liball: 'libALL'
@@ -232,14 +232,26 @@ jobs:
232232
cxx: 'g++'
233233
backend: 'OPENMP'
234234
cmake_build_type: 'Debug'
235-
kokkos_ver: '4.1.00'
236-
arborx: 'ArborX'
235+
kokkos_ver: '4.2.00'
236+
arborx: 'v1.7'
237237
heffte: 'FFTW'
238238
hypre: 'HYPRE'
239239
liball: 'libALL'
240240
silo: 'Silo'
241241
hdf5: 'HDF5'
242242
coverage: 'ON'
243+
- distro: 'ubuntu:latest'
244+
cxx: 'g++'
245+
backend: 'OPENMP'
246+
cmake_build_type: 'Debug'
247+
kokkos_ver: '4.3.00'
248+
arborx: 'fee3d777c0325a90ab50671fca69f2dc98f2a882'
249+
heffte: 'FFTW'
250+
hypre: 'HYPRE'
251+
liball: 'libALL'
252+
silo: 'Silo'
253+
hdf5: 'HDF5'
254+
coverage: 'OFF'
243255
runs-on: ubuntu-20.04
244256
timeout-minutes: 30
245257
container:
@@ -281,14 +293,14 @@ jobs:
281293
cmake --build build --parallel 2
282294
cmake --install build
283295
- name: Checkout arborx
284-
if: ${{ matrix.arborx == 'ArborX' }}
296+
if: ${{ matrix.arborx != 'OFF' }}
285297
uses: actions/checkout@v3
286298
with:
287299
repository: arborx/ArborX
288-
ref: v1.2
300+
ref: ${{ matrix.arborx }}
289301
path: arborx
290302
- name: Build arborx
291-
if: ${{ matrix.arborx == 'ArborX' }}
303+
if: ${{ matrix.arborx != 'OFF' }}
292304
working-directory: arborx
293305
run: |
294306
cmake -B build \
@@ -507,7 +519,7 @@ jobs:
507519
uses: actions/checkout@v3
508520
with:
509521
repository: arborx/ArborX
510-
ref: v1.2
522+
ref: v1.7
511523
path: arborx
512524
- name: Build arborx
513525
working-directory: arborx
@@ -566,7 +578,7 @@ jobs:
566578
matrix:
567579
cxx: ['nvcc']
568580
cmake_build_type: ['Release']
569-
kokkos_ver: ['4.1.00']
581+
kokkos_ver: ['4.2.00']
570582
runs-on: ubuntu-20.04
571583
container: ghcr.io/ecp-copa/ci-containers/cuda:12.2.0
572584
steps:
@@ -591,7 +603,7 @@ jobs:
591603
uses: actions/checkout@v3
592604
with:
593605
repository: arborx/ArborX
594-
ref: v1.2
606+
ref: v1.7
595607
path: arborx
596608
- name: Build arborx
597609
working-directory: arborx

‎.github/workflows/Nightly.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ jobs:
3636
uses: actions/checkout@v3
3737
with:
3838
repository: arborx/ArborX
39-
ref: v1.7
39+
ref: master
4040
path: arborx
4141
- name: Build arborx
4242
working-directory: arborx

‎core/src/Cabana_Experimental_NeighborList.hpp‎

Lines changed: 112 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,10 @@ namespace ArborX
106106
{
107107
//! Neighbor access trait for Cabana slice and/or Kokkos View.
108108
template <typename Positions>
109-
struct AccessTraits<Positions, PrimitivesTag,
109+
struct AccessTraits<Positions,
110+
#if ARBORX_VERSION < 10799
111+
PrimitivesTag,
112+
#endif
110113
std::enable_if_t<Cabana::is_slice<Positions>{} ||
111114
Kokkos::is_view<Positions>{}>>
112115
{
@@ -120,17 +123,21 @@ struct AccessTraits<Positions, PrimitivesTag,
120123
return Cabana::size( x );
121124
}
122125
//! Get the particle at the index.
123-
static KOKKOS_FUNCTION Point get( Positions const& x, size_type i )
126+
static KOKKOS_FUNCTION auto get( Positions const& x, size_type i )
124127
{
125-
return { static_cast<float>( x( i, 0 ) ),
126-
static_cast<float>( x( i, 1 ) ),
127-
static_cast<float>( x( i, 2 ) ) };
128+
return Point{ static_cast<float>( x( i, 0 ) ),
129+
static_cast<float>( x( i, 1 ) ),
130+
static_cast<float>( x( i, 2 ) ) };
128131
}
129132
};
130133
//! Neighbor access trait.
131134
template <typename Positions>
132-
struct AccessTraits<
133-
Cabana::Experimental::Impl::SubPositionsAndRadius<Positions>, PredicatesTag>
135+
struct AccessTraits<Cabana::Experimental::Impl::SubPositionsAndRadius<Positions>
136+
#if ARBORX_VERSION < 10799
137+
,
138+
PredicatesTag
139+
#endif
140+
>
134141
{
135142
//! Position wrapper with partial range and radius information.
136143
using PositionLike =
@@ -148,10 +155,15 @@ struct AccessTraits<
148155
static KOKKOS_FUNCTION auto get( PositionLike const& x, size_type i )
149156
{
150157
assert( i < size( x ) );
151-
auto const point =
152-
AccessTraits<typename PositionLike::positions_type,
153-
PrimitivesTag>::get( x.data, x.first + i );
154-
return attach( intersects( Sphere{ point, x.radius } ), (int)i );
158+
auto const point = AccessTraits<typename PositionLike::positions_type
159+
#if ARBORX_VERSION < 10799
160+
,
161+
PrimitivesTag
162+
#endif
163+
>::get( x.data, x.first + i );
164+
return attach(
165+
intersects( Sphere{ point, static_cast<float>( x.radius ) } ),
166+
(int)i );
155167
}
156168
};
157169
} // namespace ArborX
@@ -186,6 +198,21 @@ struct CollisionFilter<HalfNeighborTag>
186198
template <typename Tag>
187199
struct NeighborDiscriminatorCallback
188200
{
201+
#if ARBORX_VERSION >= 10799
202+
template <typename Predicate, typename Geometry, typename OutputFunctor>
203+
KOKKOS_FUNCTION void
204+
operator()( Predicate const& predicate,
205+
ArborX::PairValueIndex<Geometry, int> const& value_pair,
206+
OutputFunctor const& out ) const
207+
{
208+
int const primitive_index = value_pair.index;
209+
int const predicate_index = getData( predicate );
210+
if ( CollisionFilter<Tag>::keep( predicate_index, primitive_index ) )
211+
{
212+
out( primitive_index );
213+
}
214+
}
215+
#else
189216
template <typename Predicate, typename OutputFunctor>
190217
KOKKOS_FUNCTION void operator()( Predicate const& predicate,
191218
int primitive_index,
@@ -197,13 +224,28 @@ struct NeighborDiscriminatorCallback
197224
out( primitive_index );
198225
}
199226
}
227+
#endif
200228
};
201229

202230
// Count in the first pass
203231
template <typename Counts, typename Tag>
204232
struct NeighborDiscriminatorCallback2D_FirstPass
205233
{
206234
Counts counts;
235+
#if ARBORX_VERSION >= 10799
236+
template <typename Predicate, typename Geometry>
237+
KOKKOS_FUNCTION void
238+
operator()( Predicate const& predicate,
239+
ArborX::PairValueIndex<Geometry, int> const& value_pair ) const
240+
{
241+
int const primitive_index = value_pair.index;
242+
int const predicate_index = getData( predicate );
243+
if ( CollisionFilter<Tag>::keep( predicate_index, primitive_index ) )
244+
{
245+
++counts( predicate_index ); // WARNING see below**
246+
}
247+
}
248+
#else
207249
template <typename Predicate>
208250
KOKKOS_FUNCTION void operator()( Predicate const& predicate,
209251
int primitive_index ) const
@@ -214,6 +256,7 @@ struct NeighborDiscriminatorCallback2D_FirstPass
214256
++counts( predicate_index ); // WARNING see below**
215257
}
216258
}
259+
#endif
217260
};
218261

219262
// Preallocate and attempt fill in the first pass
@@ -222,6 +265,29 @@ struct NeighborDiscriminatorCallback2D_FirstPass_BufferOptimization
222265
{
223266
Counts counts;
224267
Neighbors neighbors;
268+
#if ARBORX_VERSION >= 10799
269+
template <typename Predicate, typename Geometry>
270+
KOKKOS_FUNCTION void
271+
operator()( Predicate const& predicate,
272+
ArborX::PairValueIndex<Geometry, int> const& value_pair ) const
273+
{
274+
int const primitive_index = value_pair.index;
275+
int const predicate_index = getData( predicate );
276+
auto& count = counts( predicate_index );
277+
if ( CollisionFilter<Tag>::keep( predicate_index, primitive_index ) )
278+
{
279+
if ( count < (int)neighbors.extent( 1 ) )
280+
{
281+
neighbors( predicate_index, count++ ) =
282+
primitive_index; // WARNING see below**
283+
}
284+
else
285+
{
286+
count++;
287+
}
288+
}
289+
}
290+
#else
225291
template <typename Predicate>
226292
KOKKOS_FUNCTION void operator()( Predicate const& predicate,
227293
int primitive_index ) const
@@ -241,6 +307,7 @@ struct NeighborDiscriminatorCallback2D_FirstPass_BufferOptimization
241307
}
242308
}
243309
}
310+
#endif
244311
};
245312

246313
// Fill in the second pass
@@ -249,6 +316,23 @@ struct NeighborDiscriminatorCallback2D_SecondPass
249316
{
250317
Counts counts;
251318
Neighbors neighbors;
319+
#if ARBORX_VERSION >= 10799
320+
template <typename Predicate, typename Geometry>
321+
KOKKOS_FUNCTION void
322+
operator()( Predicate const& predicate,
323+
ArborX::PairValueIndex<Geometry, int> const& value_pair ) const
324+
{
325+
int const primitive_index = value_pair.index;
326+
int const predicate_index = getData( predicate );
327+
auto& count = counts( predicate_index );
328+
if ( CollisionFilter<Tag>::keep( predicate_index, primitive_index ) )
329+
{
330+
assert( count < (int)neighbors.extent( 1 ) );
331+
neighbors( predicate_index, count++ ) =
332+
primitive_index; // WARNING see below**
333+
}
334+
}
335+
#else
252336
template <typename Predicate>
253337
KOKKOS_FUNCTION void operator()( Predicate const& predicate,
254338
int primitive_index ) const
@@ -262,6 +346,7 @@ struct NeighborDiscriminatorCallback2D_SecondPass
262346
primitive_index; // WARNING see below**
263347
}
264348
}
349+
#endif
265350
};
266351

267352
// NOTE** Taking advantage of the knowledge that one predicate is processed by a
@@ -319,7 +404,12 @@ auto makeNeighborList( ExecutionSpace space, Tag, Positions const& positions,
319404

320405
using memory_space = typename Positions::memory_space;
321406

407+
#if ARBORX_VERSION >= 10799
408+
ArborX::BoundingVolumeHierarchy bvh(
409+
space, ArborX::Experimental::attach_indices<int>( positions ) );
410+
#else
322411
ArborX::BVH<memory_space> bvh( space, positions );
412+
#endif
323413

324414
Kokkos::View<int*, memory_space> indices(
325415
Kokkos::view_alloc( "indices", Kokkos::WithoutInitializing ), 0 );
@@ -444,14 +534,23 @@ auto make2DNeighborList( ExecutionSpace space, Tag, Positions const& positions,
444534

445535
using memory_space = typename Positions::memory_space;
446536

537+
#if ARBORX_VERSION >= 10799
538+
ArborX::BoundingVolumeHierarchy bvh(
539+
space, ArborX::Experimental::attach_indices<int>( positions ) );
540+
#else
447541
ArborX::BVH<memory_space> bvh( space, positions );
542+
#endif
448543

449544
auto const predicates =
450545
Impl::makePredicates( positions, first, last, radius );
451546

452547
auto const n_queries =
453-
ArborX::AccessTraits<std::remove_const_t<decltype( predicates )>,
454-
ArborX::PredicatesTag>::size( predicates );
548+
ArborX::AccessTraits<std::remove_const_t<decltype( predicates )>
549+
#if ARBORX_VERSION < 10799
550+
,
551+
ArborX::PredicatesTag
552+
#endif
553+
>::size( predicates );
455554

456555
Kokkos::View<int**, memory_space> neighbors;
457556
Kokkos::View<int*, memory_space> counts( "counts", n_queries );

0 commit comments

Comments
 (0)