Files
lammps/lib/kokkos/core/unit_test/TestAtomicOperations.hpp
2023-11-21 15:02:12 -07:00

454 lines
17 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
#include <Kokkos_Core.hpp>
#include <Kokkos_Pair.hpp>
namespace TestAtomicOperations {
struct AddAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
Kokkos::atomic_add(ptr_op, update);
T old_val = Kokkos::atomic_fetch_add(ptr_fetch_op, update);
T new_val = Kokkos::atomic_add_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return old + update;
}
static const char* name() { return "add"; }
};
struct SubAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
Kokkos::atomic_sub(ptr_op, update);
T old_val = Kokkos::atomic_fetch_sub(ptr_fetch_op, update);
T new_val = Kokkos::atomic_sub_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return old - update;
}
static const char* name() { return "sub"; }
};
struct IncAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T) {
Kokkos::atomic_inc(ptr_op);
T old_val = Kokkos::atomic_fetch_inc(ptr_fetch_op);
T new_val = Kokkos::atomic_inc_fetch(ptr_op_fetch);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T) {
return old + 1;
}
static const char* name() { return "inc"; }
};
struct DecAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T) {
Kokkos::atomic_dec(ptr_op);
T old_val = Kokkos::atomic_fetch_dec(ptr_fetch_op);
T new_val = Kokkos::atomic_dec_fetch(ptr_op_fetch);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T) {
return old - 1;
}
static const char* name() { return "dec"; }
};
struct MaxAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
Kokkos::atomic_max(ptr_op, update);
T old_val = Kokkos::atomic_fetch_max(ptr_fetch_op, update);
T new_val = Kokkos::atomic_max_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return update > old ? update : old;
}
static const char* name() { return "max"; }
};
struct MinAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
Kokkos::atomic_min(ptr_op, update);
T old_val = Kokkos::atomic_fetch_min(ptr_fetch_op, update);
T new_val = Kokkos::atomic_min_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return update < old ? update : old;
}
static const char* name() { return "min"; }
};
struct MulAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
Kokkos::atomic_mul(ptr_op, update);
T old_val = Kokkos::atomic_fetch_mul(ptr_fetch_op, update);
T new_val = Kokkos::atomic_mul_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return old * update;
}
static const char* name() { return "mul"; }
};
struct DivAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
Kokkos::atomic_div(ptr_op, update);
T old_val = Kokkos::atomic_fetch_div(ptr_fetch_op, update);
T new_val = Kokkos::atomic_div_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return old / update;
}
static const char* name() { return "div"; }
};
struct ModAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
// Kokkos::atomic_mod(ptr_op, update);
(void)Kokkos::atomic_fetch_mod(ptr_op, update);
T old_val = Kokkos::atomic_fetch_mod(ptr_fetch_op, update);
T new_val = Kokkos::atomic_mod_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return old % update;
}
static const char* name() { return "mod"; }
};
struct AndAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
Kokkos::atomic_and(ptr_op, update);
T old_val = Kokkos::atomic_fetch_and(ptr_fetch_op, update);
T new_val = Kokkos::atomic_and_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return old & update;
}
static const char* name() { return "and"; }
};
struct OrAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
Kokkos::atomic_or(ptr_op, update);
T old_val = Kokkos::atomic_fetch_or(ptr_fetch_op, update);
T new_val = Kokkos::atomic_or_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return old | update;
}
static const char* name() { return "or"; }
};
struct XorAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
// Kokkos::atomic_xor(ptr_op, update);
(void)Kokkos::atomic_fetch_xor(ptr_op, update);
T old_val = Kokkos::atomic_fetch_xor(ptr_fetch_op, update);
T new_val = Kokkos::atomic_xor_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return old ^ update;
}
static const char* name() { return "xor"; }
};
struct NandAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
// Kokkos::atomic_nand(ptr_op, update);
(void)Kokkos::atomic_fetch_nand(ptr_op, update);
T old_val = Kokkos::atomic_fetch_nand(ptr_fetch_op, update);
T new_val = Kokkos::atomic_nand_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return ~(old & update);
}
static const char* name() { return "nand"; }
};
struct LShiftAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
// Kokkos::atomic_lshift(ptr_op, update);
(void)Kokkos::atomic_fetch_lshift(ptr_op, update);
T old_val = Kokkos::atomic_fetch_lshift(ptr_fetch_op, update);
T new_val = Kokkos::atomic_lshift_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return old << update;
}
static const char* name() { return "lshift"; }
};
struct RShiftAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
// Kokkos::atomic_rshift(ptr_op, update); not implemented
(void)Kokkos::atomic_fetch_rshift(ptr_op, update);
T old_val = Kokkos::atomic_fetch_rshift(ptr_fetch_op, update);
T new_val = Kokkos::atomic_rshift_fetch(ptr_op_fetch, update);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T update) {
return old >> update;
}
static const char* name() { return "rshift"; }
};
struct LoadStoreAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T update) {
T old_val = Kokkos::atomic_load(ptr_op);
Kokkos::atomic_store(ptr_op, update);
Kokkos::atomic_store(ptr_op_fetch, update);
Kokkos::atomic_store(ptr_fetch_op, update);
return Kokkos::pair<T, T>(old_val, update);
}
template <class T>
KOKKOS_FUNCTION static T op(T, T update) {
return update;
}
static const char* name() { return "load/store"; }
};
struct IncModAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T wrap_value) {
// no atomic_inc_mod in desul
(void)desul::atomic_fetch_inc_mod(ptr_op, wrap_value,
desul::MemoryOrderRelaxed(),
desul::MemoryScopeDevice());
T old_val = desul::atomic_fetch_inc_mod(ptr_fetch_op, wrap_value,
desul::MemoryOrderRelaxed(),
desul::MemoryScopeDevice());
// no atomic_inc_mod_fetch in desul
(void)desul::atomic_fetch_inc_mod(ptr_op_fetch, wrap_value,
desul::MemoryOrderRelaxed(),
desul::MemoryScopeDevice());
T new_val = op(old_val, wrap_value);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T wrap_value) {
return old + 1 > wrap_value ? 0 : old + 1;
}
static const char* name() { return "inc_mod"; }
};
struct DecModAtomicTest {
template <class T>
KOKKOS_FUNCTION static auto atomic_op(T* ptr_op, T* ptr_fetch_op,
T* ptr_op_fetch, T wrap_value) {
// no atomic_dec_mod in desul
(void)desul::atomic_fetch_dec_mod(ptr_op, wrap_value,
desul::MemoryOrderRelaxed(),
desul::MemoryScopeDevice());
T old_val = desul::atomic_fetch_dec_mod(ptr_fetch_op, wrap_value,
desul::MemoryOrderRelaxed(),
desul::MemoryScopeDevice());
// no atomic_dec_mod_fetch in desul
(void)desul::atomic_fetch_dec_mod(ptr_op_fetch, wrap_value,
desul::MemoryOrderRelaxed(),
desul::MemoryScopeDevice());
T new_val = op(old_val, wrap_value);
return Kokkos::pair<T, T>(old_val, new_val);
}
template <class T>
KOKKOS_FUNCTION static T op(T old, T wrap_value) {
return ((old == 0) || (old > wrap_value)) ? wrap_value : old - 1;
}
static const char* name() { return "dec_mod"; }
};
template <class Op, class T, class ExecSpace>
bool atomic_op_test(T old_val, T update) {
Kokkos::View<T[3], ExecSpace> op_data("op_data");
Kokkos::deep_copy(op_data, old_val);
int result = 0;
Kokkos::parallel_reduce(
Kokkos::RangePolicy<ExecSpace>(0, 1),
KOKKOS_LAMBDA(int, int& local_result) {
auto fetch_result =
Op::atomic_op(&op_data(0), &op_data(1), &op_data(2), update);
T expected_val = Op::op(old_val, update);
Kokkos::memory_fence();
if (op_data(0) != expected_val) local_result += 1;
if (op_data(1) != expected_val) local_result += 2;
if (op_data(2) != expected_val) local_result += 4;
if (fetch_result.first != old_val) local_result += 8;
if (fetch_result.second != expected_val) local_result += 16;
},
result);
if ((result & 1) != 0)
printf("atomic_%s failed with type %s\n", Op::name(), typeid(T).name());
if ((result & 2) != 0)
printf("atomic_fetch_%s failed with type %s\n", Op::name(),
typeid(T).name());
if ((result & 4) != 0)
printf("atomic_%s_fetch failed with type %s\n", Op::name(),
typeid(T).name());
if ((result & 8) != 0)
printf("atomic_fetch_%s did not return old value with type %s\n",
Op::name(), typeid(T).name());
if ((result & 16) != 0)
printf("atomic_%s_fetch did not return updated value with type %s\n",
Op::name(), typeid(T).name());
return result == 0;
}
//---------------------------------------------------
//--------------atomic_test_control------------------
//---------------------------------------------------
template <class T, class ExecSpace>
bool AtomicOperationsTestIntegralType(int old_val_in, int update_in, int test) {
T old_val = static_cast<T>(old_val_in);
T update = static_cast<T>(update_in);
switch (test) {
case 0: return atomic_op_test<AddAtomicTest, T, ExecSpace>(old_val, update);
case 1: return atomic_op_test<SubAtomicTest, T, ExecSpace>(old_val, update);
case 2: return atomic_op_test<MaxAtomicTest, T, ExecSpace>(old_val, update);
case 3: return atomic_op_test<MinAtomicTest, T, ExecSpace>(old_val, update);
case 4: return atomic_op_test<MulAtomicTest, T, ExecSpace>(old_val, update);
case 5:
return update != 0
? atomic_op_test<DivAtomicTest, T, ExecSpace>(old_val, update)
: true;
case 6:
return update != 0
? atomic_op_test<ModAtomicTest, T, ExecSpace>(old_val, update)
: true;
case 7: return atomic_op_test<AndAtomicTest, T, ExecSpace>(old_val, update);
case 8: return atomic_op_test<OrAtomicTest, T, ExecSpace>(old_val, update);
case 9: return atomic_op_test<XorAtomicTest, T, ExecSpace>(old_val, update);
case 10:
return atomic_op_test<NandAtomicTest, T, ExecSpace>(old_val, update);
case 11:
return update_in >= 0 ? atomic_op_test<LShiftAtomicTest, T, ExecSpace>(
old_val, update)
: true;
case 12:
return update_in >= 0 ? atomic_op_test<RShiftAtomicTest, T, ExecSpace>(
old_val, update)
: true;
case 13:
return atomic_op_test<IncAtomicTest, T, ExecSpace>(old_val, update);
case 14:
return atomic_op_test<DecAtomicTest, T, ExecSpace>(old_val, update);
case 15:
return atomic_op_test<LoadStoreAtomicTest, T, ExecSpace>(old_val, update);
}
return true;
}
template <class T, class ExecSpace>
bool AtomicOperationsTestUnsignedIntegralType(int old_val_in, int update_in,
int test) {
T old_val = static_cast<T>(old_val_in);
T update = static_cast<T>(update_in);
switch (test) {
case 1:
return atomic_op_test<IncModAtomicTest, T, ExecSpace>(old_val, update);
case 2:
return atomic_op_test<DecModAtomicTest, T, ExecSpace>(old_val, update);
}
return true;
}
template <class T, class ExecSpace>
bool AtomicOperationsTestNonIntegralType(int old_val_in, int update_in,
int test) {
T old_val = static_cast<T>(old_val_in);
T update = static_cast<T>(update_in);
switch (test) {
case 0: return atomic_op_test<AddAtomicTest, T, ExecSpace>(old_val, update);
case 1: return atomic_op_test<SubAtomicTest, T, ExecSpace>(old_val, update);
case 2: return atomic_op_test<MaxAtomicTest, T, ExecSpace>(old_val, update);
case 3: return atomic_op_test<MinAtomicTest, T, ExecSpace>(old_val, update);
case 4: return atomic_op_test<MulAtomicTest, T, ExecSpace>(old_val, update);
case 5:
return update != 0
? atomic_op_test<DivAtomicTest, T, ExecSpace>(old_val, update)
: true;
case 6:
return atomic_op_test<LoadStoreAtomicTest, T, ExecSpace>(old_val, update);
}
return true;
}
} // namespace TestAtomicOperations