Files
lammps/lib/kokkos/core/unit_test/TestConcurrentBitset.hpp
2023-03-03 09:22:33 -07:00

141 lines
4.0 KiB
C++

//@HEADER
// ************************************************************************
//
// Kokkos v. 4.0
// Copyright (2022) National Technology & Engineering
// Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// the U.S. Government retains certain rights in this software.
//
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
// See https://kokkos.org/LICENSE for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//@HEADER
#ifndef TEST_CONCURRENTBITSET_HPP
#define TEST_CONCURRENTBITSET_HPP
#include <gtest/gtest.h>
#include <sstream>
#include <iostream>
#include <impl/Kokkos_ConcurrentBitset.hpp>
namespace Test {
template <class DeviceType>
struct ConcurrentBitset {
using view_unsigned_type = Kokkos::View<uint32_t*, DeviceType>;
using view_int_type = Kokkos::View<int*, DeviceType>;
view_unsigned_type bitset;
view_int_type acquired;
uint32_t bitset_count_lg2;
uint32_t bitset_count_mask;
ConcurrentBitset(const uint32_t arg_bitset_count_lg2,
const view_unsigned_type& arg_bitset,
const view_int_type& arg_acquired)
: bitset(arg_bitset),
acquired(arg_acquired),
bitset_count_lg2(arg_bitset_count_lg2),
bitset_count_mask(uint32_t(1u << arg_bitset_count_lg2) - 1) {}
struct TagAcquire {};
struct TagRelease {};
struct TagReacquire {};
KOKKOS_INLINE_FUNCTION
void operator()(TagAcquire, int i, long& update) const {
unsigned hint = Kokkos::Impl::clock_tic() & bitset_count_mask;
Kokkos::pair<int, int> result =
Kokkos::Impl::concurrent_bitset::acquire_bounded_lg2(
bitset.data(), bitset_count_lg2, hint);
acquired(i) = result.first;
if (0 <= result.first) ++update;
}
KOKKOS_INLINE_FUNCTION
void operator()(TagRelease, int i, long& update) const {
if (0 == (i % 3) && 0 <= acquired(i)) {
Kokkos::Impl::concurrent_bitset::release(bitset.data(), acquired(i));
acquired(i) = -1;
++update;
}
}
KOKKOS_INLINE_FUNCTION
void operator()(TagReacquire, int i, long& update) const {
if (acquired(i) < 0) {
unsigned hint = Kokkos::Impl::clock_tic() & bitset_count_mask;
Kokkos::pair<int, int> result =
Kokkos::Impl::concurrent_bitset::acquire_bounded_lg2(
bitset.data(), bitset_count_lg2, hint);
acquired(i) = result.first;
if (0 <= result.first) ++update;
}
}
};
template <class DeviceType>
void test_concurrent_bitset(int bit_count) {
using Functor = ConcurrentBitset<DeviceType>;
using view_unsigned_type = typename Functor::view_unsigned_type;
using view_int_type = typename Functor::view_int_type;
int bit_count_lg2 = 1;
while ((1 << bit_count_lg2) < bit_count) ++bit_count_lg2;
bit_count = 1 << bit_count_lg2;
const int buffer_length =
Kokkos::Impl::concurrent_bitset::buffer_bound_lg2(bit_count_lg2);
view_unsigned_type bitset("bitset", buffer_length);
// Try to acquire more than available:
const size_t n = (bit_count * 3) / 2;
view_int_type acquired("acquired", n);
typename view_unsigned_type::HostMirror bitset_host =
Kokkos::create_mirror_view(bitset);
Kokkos::deep_copy(bitset, 0u);
long total = 0;
long total_release = 0;
long total_reacquire = 0;
Kokkos::parallel_reduce(
Kokkos::RangePolicy<DeviceType, typename Functor::TagAcquire>(0, n),
Functor(bit_count_lg2, bitset, acquired), total);
ASSERT_EQ(bit_count, total);
Kokkos::parallel_reduce(
Kokkos::RangePolicy<DeviceType, typename Functor::TagRelease>(0, n),
Functor(bit_count_lg2, bitset, acquired), total_release);
Kokkos::parallel_reduce(
Kokkos::RangePolicy<DeviceType, typename Functor::TagReacquire>(0, n),
Functor(bit_count_lg2, bitset, acquired), total_reacquire);
ASSERT_EQ(total_release, total_reacquire);
}
} // namespace Test
#endif /* #ifndef TEST_CONCURRENTBITSET_HPP */