Update atomic functional tests (#262)

* feat: implement function to return number of blocks in grid.

* test: update atomics functional tests
  - Standard atomic tests: `atomic_add`, `atomic_inc`, `fetch_atomic_add`, `fetch_atomic_inc`, and `fetch_compare_and_swap`
  - Bitwise atomic tests:    `atomic_and`, `atomic_or`, `atomic_xor`, fetch_atomic_and`, `fetch_atomic_or`, and `fetch_atomic_xor`
  - Extended atomic tests: `atomic_fetch`, `atomic_set`, and `atomic_swap`

* Added two different address modes for atomics.
* Added all supported data types for atomics tests.


[ROCm/rocshmem commit: 0a4f8a83b9]
This commit is contained in:
Avinash Kethineedi authored and GitHub committed 2025-10-06 10:50:50 -05:00
1 parent ac13b22edc
commit e31b4d42e5
9 files changed
+676 -265

No files matched your search

+7
View File
@@ -191,6 +191,13 @@ __device__ __forceinline__ int get_flat_block_id() {
hipThreadIdx_z * hipBlockDim_x * hipBlockDim_y; hipThreadIdx_z * hipBlockDim_x * hipBlockDim_y;
} }
/*
* Returns the number of blocks in the caller's flattened grid.
*/
__device__ __forceinline__ int get_grid_num_blocks() {
return hipGridDim_x * hipGridDim_y * hipGridDim_z;
}
/* /*
* Returns the flattened block index that the calling thread is a member of in * Returns the flattened block index that the calling thread is a member of in
* in the grid. Callers from the same block will have the same index. * in the grid. Callers from the same block will have the same index.
@@ -32,33 +32,54 @@ using namespace rocshmem;
/* Declare the global kernel template with a generic implementation */ /* Declare the global kernel template with a generic implementation */
template <typename T> template <typename T>
__global__ void AMOBitwiseTest(int loop, int skip, long long int *start_time, __global__ void AMOBitwiseTest(int loop, int skip, long long int *start_time,
long long int *end_time, char *r_buf, long long int *end_time, T *dest, T *ret_val,
T *s_buf, T *ret_val, TestType type, AddrMode addr_mode, TestType type,
ShmemContextType ctx_type) { ShmemContextType ctx_type) {
return; return;
} }
template <class T>
__device__ inline T* compute_target_ptr(T* base_ptr, AddrMode addr_mode,
int wg_idx, int itr, int n_wgs) {
// PerBlock: element = wg_idx, with n_wgs elements per loop
// PerGrid : single element shared by the whole grid per loop
if (addr_mode == AddrMode::PerBlock) {
size_t offset = wg_idx + itr * n_wgs;
return base_ptr + offset;
} else { // PerGrid
return base_ptr + itr;
}
}
/****************************************************************************** /******************************************************************************
* HOST TESTER CLASS METHODS * HOST TESTER CLASS METHODS
*****************************************************************************/ *****************************************************************************/
template <typename T> template <typename T>
AMOBitwiseTester<T>::AMOBitwiseTester(TesterArguments args) : Tester(args) { AMOBitwiseTester<T>::AMOBitwiseTester(TesterArguments args) : Tester(args) {
CHECK_HIP(hipMalloc((void **)&_ret_val, args.max_msg_size * args.num_wgs)); n_out = (args.addr_mode == AddrMode::PerBlock) ? args.num_wgs : 1;
_r_buf = (char *)rocshmem_malloc(args.max_msg_size); n_in = args.num_wgs * args.wg_size;
_s_buf = (T *)rocshmem_malloc(args.max_msg_size * args.num_wgs); n_loops = args.loop + args.skip;
// One return per *thread* per loop
CHECK_HIP(hipMalloc((void **)&ret_val, args.max_msg_size * n_in * n_loops));
dest = (T *)rocshmem_malloc(args.max_msg_size * n_out * n_loops);
if (dest == nullptr) {
std::cerr << "Error allocating memory from symmetric heap" << std::endl;
std::cerr << "dest: " << (void*)dest << std::endl;
}
} }
template <typename T> template <typename T>
AMOBitwiseTester<T>::~AMOBitwiseTester() { AMOBitwiseTester<T>::~AMOBitwiseTester() {
rocshmem_free(_r_buf); CHECK_HIP(hipFree(ret_val));
CHECK_HIP(hipFree(_ret_val)); rocshmem_free(dest);
} }
template <typename T> template <typename T>
void AMOBitwiseTester<T>::resetBuffers(size_t size) { void AMOBitwiseTester<T>::resetBuffers(size_t size) {
memset(_r_buf, 0, args.max_msg_size); memset(ret_val, 0, args.max_msg_size * n_in * n_loops);
memset(_ret_val, 0, args.max_msg_size * args.num_wgs); memset(dest, 0, args.max_msg_size * n_out * n_loops);
memset(_s_buf, 0, args.max_msg_size * args.num_wgs);
} }
template <typename T> template <typename T>
@@ -67,59 +88,158 @@ void AMOBitwiseTester<T>::launchKernel(dim3 gridsize, dim3 blocksize, int loop,
size_t shared_bytes = 0; size_t shared_bytes = 0;
hipLaunchKernelGGL(AMOBitwiseTest, gridsize, blocksize, shared_bytes, stream, hipLaunchKernelGGL(AMOBitwiseTest, gridsize, blocksize, shared_bytes, stream,
loop, args.skip, start_time, end_time, _r_buf, _s_buf, args.loop, args.skip, start_time, end_time, dest,
_ret_val, _type, _shmem_context); ret_val, args.addr_mode, _type, _shmem_context);
_gridSize = gridsize; num_msgs = n_loops * gridsize.x * blocksize.x;
num_msgs = (loop + args.skip) * gridsize.x; num_timed_msgs = args.loop * gridsize.x * blocksize.x;
num_timed_msgs = loop; }
template <typename G>
void fail_eq(const G& got, const G& exp) {
std::cerr << "data validation error\n"
<< "got " << got << ", expected " << exp << std::endl;
std::exit(-1);
}
// Map (loop, elem_idx) -> dest[] index for current address mode.
template <typename T>
int AMOBitwiseTester<T>::destIndex(int l, int elem_idx) const {
return (args.addr_mode == AddrMode::PerBlock)
? l * static_cast<int>(args.num_wgs) + elem_idx
: l; // PerGrid has a single element per loop
}
// Number of output elements to check per loop for current address mode.
template <typename T>
int AMOBitwiseTester<T>::numElems() const {
return (args.addr_mode == AddrMode::PerBlock)
? static_cast<int>(args.num_wgs)
: 1; // PerGrid
}
// Return pointer to the start of the ret_val “chunk” for (loop, elem_idx)
// plus the chunk length for this address mode.
template <typename T>
std::pair<T*, int> AMOBitwiseTester<T>::retChunk(int l, int elem_idx) const {
if (args.addr_mode == AddrMode::PerBlock) {
// One chunk per element (workgroup): wg_size returns
T* p = ret_val + l * n_in + elem_idx * args.wg_size;
int sz = static_cast<int>(args.wg_size);
return {p, sz};
}
// PerGrid: one big chunk per loop (all threads)
T* p = ret_val + l * n_in;
int sz = static_cast<int>(n_in);
return {p, sz};
}
template <typename T>
void AMOBitwiseTester<T>::verifyDestValues() {
const int loops = static_cast<int>(n_loops);
const int n_elems = numElems();
auto check_equal_all = [&](T expected) {
for (int l = 0; l < loops; ++l) {
for (int elem = 0; elem < n_elems; ++elem) {
const int idx = destIndex(l, elem);
if (dest[idx] != expected) fail_eq(dest[idx], expected);
}
}
};
// Use all-ones mask for type T
const T MASK = static_cast<T>(~T{0});
switch (_type) {
case AMO_AndTestType:
case AMO_FetchAndTestType: {
// Start at 0; 0 & MASK == 0 regardless of writer count.
check_equal_all(static_cast<T>(0));
break;
}
case AMO_OrTestType:
case AMO_FetchOrTestType: {
// final value is MASK.
check_equal_all(MASK);
break;
}
case AMO_XorTestType:
case AMO_FetchXorTestType: {
// PerBlock: K = wg_size; PerGrid: K = num_wgs * wg_size
const int K = (args.addr_mode == AddrMode::PerBlock)
? static_cast<int>(args.wg_size)
: static_cast<int>(args.num_wgs * args.wg_size);
const T expected = (K & 1) ? MASK : static_cast<T>(0);
check_equal_all(expected);
break;
}
default:
break;
}
}
template <typename T>
void AMOBitwiseTester<T>::verifyReturnValues() {
// Only “fetch-*” types produce return values to validate.
if (_type == AMO_AndTestType || _type == AMO_OrTestType ||
_type == AMO_XorTestType) return;
const int loops = static_cast<int>(n_loops);
const int n_elems = numElems();
const T MASK = static_cast<T>(~T{0});
for (int l = 0; l < loops; ++l) {
for (int elem = 0; elem < n_elems; ++elem) {
auto [p, cnt] = retChunk(l, elem);
// Count distribution of observed old values in this chunk
int zeros = 0, masks = 0;
for (int i = 0; i < cnt; ++i) {
zeros += (p[i] == static_cast<T>(0));
masks += (p[i] == MASK);
}
if (zeros + masks != cnt) {
fail_eq(zeros + masks, cnt); // unexpected values present
}
switch (_type) {
case AMO_FetchAndTestType:
// Old value is 0 (dest stays 0)
if (!(zeros == cnt && masks == 0)) fail_eq(zeros, cnt);
break;
case AMO_FetchOrTestType:
// Exactly one 0 (the first OR), rest MASK
if (!(zeros == 1 && masks == cnt - 1)) fail_eq(zeros, 1);
break;
case AMO_FetchXorTestType: {
// returns multiset = { ceil(K/2) zeros, floor(K/2) MASKs }
const int exp_zeros = (cnt + 1) / 2; // ceil(cnt/2)
const int exp_masks = cnt / 2; // floor(cnt/2)
if (!(zeros == exp_zeros && masks == exp_masks)) {
fail_eq(zeros, exp_zeros);
}
// cross-check
if ((cnt & 1) && zeros != masks + 1) fail_eq(zeros, masks + 1);
if (!(cnt & 1) && zeros != masks) fail_eq(zeros, masks);
break;
}
default:
break;
}
}
}
} }
template <typename T> template <typename T>
void AMOBitwiseTester<T>::verifyResults(size_t size) { void AMOBitwiseTester<T>::verifyResults(size_t size) {
T ret; // PE 0 checks returns; target PE checks dest.
if (args.myid == 0) { if (args.myid) {
T expected_val = 0; verifyDestValues();
} else {
switch (_type) { verifyReturnValues();
case AMO_FetchAndTestType:
expected_val = 0;
break;
case AMO_AndTestType:
expected_val = 0;
break;
case AMO_FetchOrTestType:
expected_val = 0xFFFF;
break;
case AMO_OrTestType:
expected_val = 0xFFFF;
break;
case AMO_FetchXorTestType:
expected_val = 0xFFFF;
break;
case AMO_XorTestType:
expected_val = (num_msgs % 2) ? 0xFFFF : 0;
break;
default:
break;
}
int fetch_op =
(_type == AMO_FetchAndTestType || _type == AMO_FetchOrTestType ||
_type == AMO_FetchXorTestType)
? 1
: 0;
if (fetch_op == 1) {
ret = *std::max_element(_ret_val, _ret_val + args.num_wgs);
} else {
ret = *std::max_element(_s_buf, _s_buf + args.num_wgs);
}
if (ret != expected_val) {
std::cerr << "data validation error\n";
std::cerr << "got " << ret << ", expected " << expected_val << std::endl;
exit(-1);
}
} }
} }
@@ -127,50 +247,51 @@ void AMOBitwiseTester<T>::verifyResults(size_t size) {
template <> \ template <> \
__global__ void AMOBitwiseTest<T>( \ __global__ void AMOBitwiseTest<T>( \
int loop, int skip, long long int *start_time, \ int loop, int skip, long long int *start_time, \
long long int *end_time, char *r_buf, T *s_buf, T *ret_val, \ long long int *end_time, T *dest, T *ret_val, \
TestType type, ShmemContextType ctx_type) { \ AddrMode addr_mode, TestType type, ShmemContextType ctx_type) { \
__shared__ rocshmem_ctx_t ctx; \ __shared__ rocshmem_ctx_t ctx; \
int wg_id = get_flat_grid_id(); \ int wg_id = get_flat_grid_id(); \
int global_id = get_flat_id(); \
int n_threads = get_flat_grid_size(); \
int n_wgs = get_grid_num_blocks(); \
rocshmem_wg_init(); \ rocshmem_wg_init(); \
rocshmem_wg_ctx_create(ctx_type, &ctx); \ rocshmem_wg_ctx_create(ctx_type, &ctx); \
if (hipThreadIdx_x == 0) { \ for (int i = 0; i < loop + skip; i++) { \
T *ptr = compute_target_ptr<T>(dest, addr_mode, wg_id, i, n_wgs); \
T ret = 0; \ T ret = 0; \
T cond = 0; \ if (i == skip) { \
for (int i = 0; i < loop + skip; i++) { \ start_time[wg_id] = wall_clock64(); \
if (i == skip) { \
start_time[wg_id] = wall_clock64(); \
} \
switch (type) { \
case AMO_FetchAndTestType: \
ret = rocshmem_ctx_##TNAME##_atomic_fetch_and(ctx, (T *)r_buf, \
0xFFFF, 1); \
break; \
case AMO_AndTestType: \
rocshmem_ctx_##TNAME##_atomic_and(ctx, (T *)r_buf, 0xFFFF, 1); \
break; \
case AMO_FetchOrTestType: \
ret = rocshmem_ctx_##TNAME##_atomic_fetch_or(ctx, (T *)r_buf, \
0xFFFF, 1); \
break; \
case AMO_OrTestType: \
rocshmem_ctx_##TNAME##_atomic_or(ctx, (T *)r_buf, 0xFFFF, 1); \
break; \
case AMO_FetchXorTestType: \
ret = rocshmem_ctx_##TNAME##_atomic_fetch_xor(ctx, (T *)r_buf, \
0xFFFF, 1); \
break; \
case AMO_XorTestType: \
rocshmem_ctx_##TNAME##_atomic_xor(ctx, (T *)r_buf, 0xFFFF, 1); \
break; \
default: \
break; \
} \
} \ } \
rocshmem_ctx_quiet(ctx); \ switch (type) { \
end_time[wg_id] = wall_clock64(); \ case AMO_FetchAndTestType: \
ret_val[wg_id] = ret; \ ret = rocshmem_ctx_##TNAME##_atomic_fetch_and(ctx, ptr, \
rocshmem_ctx_getmem(ctx, &s_buf[wg_id], r_buf, sizeof(T), 1); \ (T)~(T)0, 1); \
break; \
case AMO_AndTestType: \
rocshmem_ctx_##TNAME##_atomic_and(ctx, ptr, (T)~(T)0, 1); \
break; \
case AMO_FetchOrTestType: \
ret = rocshmem_ctx_##TNAME##_atomic_fetch_or(ctx, ptr, \
(T)~(T)0, 1); \
break; \
case AMO_OrTestType: \
rocshmem_ctx_##TNAME##_atomic_or(ctx, ptr, (T)~(T)0, 1); \
break; \
case AMO_FetchXorTestType: \
ret = rocshmem_ctx_##TNAME##_atomic_fetch_xor(ctx, ptr, \
(T)~(T)0, 1); \
break; \
case AMO_XorTestType: \
rocshmem_ctx_##TNAME##_atomic_xor(ctx, ptr, (T)~(T)0, 1); \
break; \
default: \
break; \
} \
ret_val[global_id + i * n_threads] = ret; \
} \ } \
rocshmem_ctx_quiet(ctx); \
end_time[wg_id] = wall_clock64(); \
__syncthreads(); \
rocshmem_wg_ctx_destroy(&ctx); \ rocshmem_wg_ctx_destroy(&ctx); \
rocshmem_wg_finalize(); \ rocshmem_wg_finalize(); \
} \ } \
@@ -179,3 +300,5 @@ void AMOBitwiseTester<T>::verifyResults(size_t size) {
AMO_BITWISE_DEF_GEN(unsigned int, uint) AMO_BITWISE_DEF_GEN(unsigned int, uint)
AMO_BITWISE_DEF_GEN(unsigned long, ulong) AMO_BITWISE_DEF_GEN(unsigned long, ulong)
AMO_BITWISE_DEF_GEN(unsigned long long, ulonglong) AMO_BITWISE_DEF_GEN(unsigned long long, ulonglong)
AMO_BITWISE_DEF_GEN(int32_t, int32)
AMO_BITWISE_DEF_GEN(int64_t, int64)
@@ -44,10 +44,19 @@ class AMOBitwiseTester : public Tester {
virtual void verifyResults(size_t size) override; virtual void verifyResults(size_t size) override;
dim3 _gridSize{}; void verifyDestValues();
char *_r_buf; void verifyReturnValues();
T *_ret_val;
T *_s_buf; int destIndex(int l, int elem_idx) const;
int numElems() const;
std::pair<T*, int> retChunk(int l, int elem_idx) const;
T* dest{nullptr}; // symmetric target buffer [loop][elem]
T* ret_val{nullptr}; // device returns [loop][thread]
size_t n_in{0}; // num_wgs * wg_size
size_t n_out{0}; // elements per loop: PerBlock->num_wgs, PerGrid->1
size_t n_loops{0}; // loop + skip
}; };
#endif #endif
@@ -32,33 +32,54 @@ using namespace rocshmem;
/* Declare the global kernel template with a generic implementation */ /* Declare the global kernel template with a generic implementation */
template <typename T> template <typename T>
__global__ void AMOExtendedTest(int loop, int skip, long long int *start_time, __global__ void AMOExtendedTest(int loop, int skip, long long int *start_time,
long long int *end_time, char *r_buf, long long int *end_time, T *dest, T *ret_val,
T *s_buf, T *ret_val, TestType type, AddrMode addr_mode, TestType type,
ShmemContextType ctx_type) { ShmemContextType ctx_type) {
return; return;
} }
template <class T>
__device__ inline T* compute_target_ptr(T* base_ptr, AddrMode addr_mode,
int wg_idx, int itr, int n_wgs) {
// PerBlock: element = wg_idx, with n_wgs elements per loop
// PerGrid : single element shared by the whole grid per loop
if (addr_mode == AddrMode::PerBlock) {
size_t offset = wg_idx + itr * n_wgs;
return base_ptr + offset;
} else { // PerGrid
return base_ptr + itr;
}
}
/****************************************************************************** /******************************************************************************
* HOST TESTER CLASS METHODS * HOST TESTER CLASS METHODS
*****************************************************************************/ *****************************************************************************/
template <typename T> template <typename T>
AMOExtendedTester<T>::AMOExtendedTester(TesterArguments args) : Tester(args) { AMOExtendedTester<T>::AMOExtendedTester(TesterArguments args) : Tester(args) {
CHECK_HIP(hipMalloc((void **)&_ret_val, args.max_msg_size * args.num_wgs)); n_out = (args.addr_mode == AddrMode::PerBlock) ? args.num_wgs : 1;
_r_buf = (char *)rocshmem_malloc(args.max_msg_size); n_in = args.num_wgs * args.wg_size;
_s_buf = (T *)rocshmem_malloc(args.max_msg_size * args.num_wgs); n_loops = args.loop + args.skip;
// One return per *thread* per loop
CHECK_HIP(hipMalloc((void **)&ret_val, args.max_msg_size * n_in * n_loops));
dest = (T *)rocshmem_malloc(args.max_msg_size * n_out * n_loops);
if (dest == nullptr) {
std::cerr << "Error allocating memory from symmetric heap" << std::endl;
std::cerr << "dest: " << (void*)dest << std::endl;
}
} }
template <typename T> template <typename T>
AMOExtendedTester<T>::~AMOExtendedTester() { AMOExtendedTester<T>::~AMOExtendedTester() {
rocshmem_free(_r_buf); CHECK_HIP(hipFree(ret_val));
CHECK_HIP(hipFree(_ret_val)); rocshmem_free(dest);
} }
template <typename T> template <typename T>
void AMOExtendedTester<T>::resetBuffers(size_t size) { void AMOExtendedTester<T>::resetBuffers(size_t size) {
memset(_r_buf, 0, args.max_msg_size); memset(ret_val, 0, args.max_msg_size * n_in * n_loops);
memset(_ret_val, 0, args.max_msg_size * args.num_wgs); memset(dest, 0, args.max_msg_size * n_out * n_loops);
memset(_s_buf, 0, args.max_msg_size * args.num_wgs);
} }
template <typename T> template <typename T>
@@ -67,87 +88,173 @@ void AMOExtendedTester<T>::launchKernel(dim3 gridsize, dim3 blocksize, int loop,
size_t shared_bytes = 0; size_t shared_bytes = 0;
hipLaunchKernelGGL(AMOExtendedTest, gridsize, blocksize, shared_bytes, stream, hipLaunchKernelGGL(AMOExtendedTest, gridsize, blocksize, shared_bytes, stream,
loop, args.skip, start_time, end_time, _r_buf, _s_buf, args.loop, args.skip, start_time, end_time, dest,
_ret_val, _type, _shmem_context); ret_val, args.addr_mode, _type, _shmem_context);
_gridSize = gridsize; num_msgs = n_loops * gridsize.x * blocksize.x;
num_msgs = (loop + args.skip) * gridsize.x; num_timed_msgs = args.loop * gridsize.x * blocksize.x;
num_timed_msgs = loop; }
template <typename G>
void fail_eq(const G& got, const G& exp) {
std::cerr << "data validation error\n"
<< "got " << got << ", expected " << exp << std::endl;
std::exit(-1);
}
// Map (loop, elem_idx) -> dest[] index for current address mode.
template <typename T>
int AMOExtendedTester<T>::destIndex(int l, int elem_idx) const {
return (args.addr_mode == AddrMode::PerBlock)
? l * static_cast<int>(args.num_wgs) + elem_idx
: l; // PerGrid has a single element per loop
}
// Number of output elements to check per loop for current address mode.
template <typename T>
int AMOExtendedTester<T>::numElems() const {
return (args.addr_mode == AddrMode::PerBlock)
? static_cast<int>(args.num_wgs)
: 1; // PerGrid
}
// Return pointer to the start of the ret_val “chunk” for (loop, elem_idx)
// plus the chunk length for this address mode.
template <typename T>
std::pair<T*, int> AMOExtendedTester<T>::retChunk(int l, int elem_idx) const {
if (args.addr_mode == AddrMode::PerBlock) {
// One chunk per element (workgroup): wg_size returns
T* p = ret_val + l * n_in + elem_idx * args.wg_size;
int sz = static_cast<int>(args.wg_size);
return {p, sz};
}
// PerGrid: one big chunk per loop (all threads)
T* p = ret_val + l * n_in;
int sz = static_cast<int>(n_in);
return {p, sz};
} }
template <typename T> template <typename T>
void AMOExtendedTester<T>::verifyResults(size_t size) { void AMOExtendedTester<T>::verifyDestValues() {
T ret; const int loops = static_cast<int>(n_loops);
if (args.myid == 0) { const int n_elems = numElems();
T expected_val = 0;
switch (_type) { auto check_equal_all = [&](T expected) {
case AMO_FetchTestType: for (int l = 0; l < loops; ++l) {
expected_val = 0; for (int elem = 0; elem < n_elems; ++elem) {
break; const int idx = destIndex(l, elem);
case AMO_SetTestType: if (dest[idx] != expected) fail_eq(dest[idx], expected);
expected_val = 44; }
break;
case AMO_SwapTestType:
expected_val = num_msgs / 2;
break;
default:
break;
} }
};
int fetch_op = auto check_nonzero_all = [&]() {
(_type == AMO_FetchTestType || _type == AMO_SwapTestType) ? 1 : 0; for (int l = 0; l < loops; ++l) {
for (int elem = 0; elem < n_elems; ++elem) {
const int idx = destIndex(l, elem);
if (dest[idx] == T{0}) fail_eq(dest[idx], T{1});
}
}
};
if (fetch_op == 1) { switch (_type) {
ret = *std::max_element(_ret_val, _ret_val + args.num_wgs); case AMO_FetchTestType:
} else { // fetch does not modify dest -> stays 0
ret = *std::max_element(_s_buf, _s_buf + args.num_wgs); check_equal_all(T{0});
} break;
if (ret != expected_val) {
std::cerr << "data validation error\n"; case AMO_SetTestType:
std::cerr << "got " << ret << ", expected " << expected_val << std::endl; // set writes a constant (17)
exit(-1); check_equal_all(static_cast<T>(17));
break;
case AMO_SwapTestType:
// swap writes non-zero values -> final must be non-zero
check_nonzero_all();
break;
default:
break;
}
}
template <typename T>
void AMOExtendedTester<T>::verifyReturnValues() {
// Only fetch/swap produce return values to validate
if (_type == AMO_SetTestType) return;
const int loops = static_cast<int>(n_loops);
const int n_elems = numElems();
for (int l = 0; l < loops; ++l) {
for (int elem = 0; elem < n_elems; ++elem) {
auto [p, cnt] = retChunk(l, elem);
int zeros = 0;
for (int i = 0; i < cnt; ++i) {
zeros += (p[i] == T{0});
}
if (_type == AMO_FetchTestType) {
// fetch returns the current value (initially 0), dest unchanged
if (zeros != cnt) fail_eq(zeros, cnt);
} else { // AMO_SwapTestType
// For a single element per (loop,elem), exactly one atomic_swap
// observes old==0 (the first arriving swap). The rest see non-zero.
if (zeros != 1) fail_eq(zeros, 1);
}
} }
} }
} }
template <typename T>
void AMOExtendedTester<T>::verifyResults(size_t /*size*/) {
// PE 0 checks returns; target PE checks dest.
if (args.myid) {
verifyDestValues();
} else {
verifyReturnValues();
}
}
#define AMO_EXTENDED_DEF_GEN(T, TNAME) \ #define AMO_EXTENDED_DEF_GEN(T, TNAME) \
template <> \ template <> \
__global__ void AMOExtendedTest<T>( \ __global__ void AMOExtendedTest<T>( \
int loop, int skip, long long int *start_time, \ int loop, int skip, long long int *start_time, \
long long int *end_time, char *r_buf, T *s_buf, T *ret_val, \ long long int *end_time, T *dest, T *ret_val, \
TestType type, ShmemContextType ctx_type) { \ AddrMode addr_mode, TestType type, ShmemContextType ctx_type) { \
__shared__ rocshmem_ctx_t ctx; \ __shared__ rocshmem_ctx_t ctx; \
int wg_id = get_flat_grid_id(); \ int wg_id = get_flat_grid_id(); \
int global_id = get_flat_id(); \
int t_id = get_flat_block_id(); \
int n_threads = get_flat_grid_size(); \
int n_wgs = get_grid_num_blocks(); \
rocshmem_wg_init(); \ rocshmem_wg_init(); \
rocshmem_wg_ctx_create(ctx_type, &ctx); \ rocshmem_wg_ctx_create(ctx_type, &ctx); \
if (hipThreadIdx_x == 0) { \ for (int i = 0; i < loop + skip; i++) { \
T ret = 0; \ T *ptr = compute_target_ptr<T>(dest, addr_mode, wg_id, i, n_wgs); \
T cond = 0; \ T ret = 0; \
for (int i = 0; i < loop + skip; i++) { \ if (i == skip) { \
if (i == skip) { \ start_time[wg_id] = wall_clock64(); \
start_time[wg_id] = wall_clock64(); \
} \
switch (type) { \
case AMO_FetchTestType: \
ret = rocshmem_ctx_##TNAME##_atomic_fetch(ctx, (T *)r_buf, 1); \
break; \
case AMO_SetTestType: \
rocshmem_ctx_##TNAME##_atomic_set(ctx, (T *)r_buf, 44, 1); \
break; \
case AMO_SwapTestType: \
ret = rocshmem_ctx_##TNAME##_atomic_swap(ctx, (T *)r_buf, \
ret + 1, 1); \
break; \
default: \
break; \
} \
} \ } \
rocshmem_ctx_quiet(ctx); \ switch (type) { \
end_time[wg_id] = wall_clock64(); \ case AMO_FetchTestType: \
ret_val[wg_id] = ret; \ ret = rocshmem_ctx_##TNAME##_atomic_fetch(ctx, ptr, 1); \
rocshmem_ctx_getmem(ctx, &s_buf[wg_id], r_buf, sizeof(T), 1); \ break; \
case AMO_SetTestType: \
rocshmem_ctx_##TNAME##_atomic_set(ctx, ptr, (T)17, 1); \
break; \
case AMO_SwapTestType: \
ret = rocshmem_ctx_##TNAME##_atomic_swap(ctx, ptr, (T)(t_id + 1),1);\
break; \
default: \
break; \
} \
ret_val[global_id + i * n_threads] = ret; \
} \ } \
rocshmem_ctx_quiet(ctx); \
end_time[wg_id] = wall_clock64(); \
__syncthreads(); \
rocshmem_wg_ctx_destroy(&ctx); \ rocshmem_wg_ctx_destroy(&ctx); \
rocshmem_wg_finalize(); \ rocshmem_wg_finalize(); \
} \ } \
@@ -44,10 +44,19 @@ class AMOExtendedTester : public Tester {
virtual void verifyResults(size_t size) override; virtual void verifyResults(size_t size) override;
dim3 _gridSize{}; void verifyDestValues();
char *_r_buf; void verifyReturnValues();
T *_ret_val;
T *_s_buf; int destIndex(int l, int elem_idx) const;
int numElems() const;
std::pair<T*, int> retChunk(int l, int elem_idx) const;
T* dest{nullptr}; // symmetric target buffer [loop][elem]
T* ret_val{nullptr}; // device returns [loop][thread]
size_t n_in{0}; // num_wgs * wg_size
size_t n_out{0}; // elements per loop: PerBlock->num_wgs, PerGrid->1
size_t n_loops{0}; // loop + skip
}; };
#endif #endif
@@ -33,33 +33,54 @@ using namespace rocshmem;
/* Declare the global kernel template with a generic implementation */ /* Declare the global kernel template with a generic implementation */
template <typename T> template <typename T>
__global__ void AMOStandardTest(int loop, int skip, long long int *start_time, __global__ void AMOStandardTest(int loop, int skip, long long int *start_time,
long long int *end_time, char *r_buf, long long int *end_time, T *dest,
T *s_buf, T *ret_val, TestType type, T *ret_val, AddrMode addr_mode,
ShmemContextType ctx_type) { TestType type, ShmemContextType ctx_type) {
return; return;
} }
template <class T>
__device__ inline T* compute_target_ptr(T* base_ptr, AddrMode addr_mode,
int wg_idx, int itr, int n_wgs) {
// PerBlock: element = wg_idx, with n_wgs elements per loop
// PerGrid : single element shared by the whole grid per loop
if (addr_mode == AddrMode::PerBlock) {
size_t offset = wg_idx + itr * n_wgs;
return base_ptr + offset;
} else { // PerGrid
return base_ptr + itr;
}
}
/****************************************************************************** /******************************************************************************
* HOST TESTER CLASS METHODS * HOST TESTER CLASS METHODS
*****************************************************************************/ *****************************************************************************/
template <typename T> template <typename T>
AMOStandardTester<T>::AMOStandardTester(TesterArguments args) : Tester(args) { AMOStandardTester<T>::AMOStandardTester(TesterArguments args) : Tester(args) {
CHECK_HIP(hipMalloc((void **)&_ret_val, args.max_msg_size * args.num_wgs)); n_out = (args.addr_mode == AddrMode::PerBlock) ? args.num_wgs : 1;
_r_buf = (char *)rocshmem_malloc(args.max_msg_size); n_in = args.num_wgs * args.wg_size;
_s_buf = (T *)rocshmem_malloc(args.max_msg_size * args.num_wgs); n_loops = args.loop + args.skip;
// One return per *thread* per loop
CHECK_HIP(hipMalloc((void **)&ret_val, args.max_msg_size * n_in * n_loops));
dest = (T *)rocshmem_malloc(args.max_msg_size * n_out * n_loops);
if (dest == nullptr) {
std::cerr << "Error allocating memory from symmetric heap" << std::endl;
std::cerr << "dest: " << dest << std::endl;
}
} }
template <typename T> template <typename T>
AMOStandardTester<T>::~AMOStandardTester() { AMOStandardTester<T>::~AMOStandardTester() {
rocshmem_free(_r_buf); CHECK_HIP(hipFree(ret_val));
CHECK_HIP(hipFree(_ret_val)); rocshmem_free(dest);
} }
template <typename T> template <typename T>
void AMOStandardTester<T>::resetBuffers(size_t size) { void AMOStandardTester<T>::resetBuffers(size_t size) {
memset(_r_buf, 0, args.max_msg_size); memset(ret_val, 0, args.max_msg_size * n_in * n_loops);
memset(_ret_val, 0, args.max_msg_size * args.num_wgs); memset(dest, 0, args.max_msg_size * n_out * n_loops);
memset(_s_buf, 0, args.max_msg_size * args.num_wgs);
} }
template <typename T> template <typename T>
@@ -68,52 +89,162 @@ void AMOStandardTester<T>::launchKernel(dim3 gridsize, dim3 blocksize, int loop,
size_t shared_bytes = 0; size_t shared_bytes = 0;
hipLaunchKernelGGL(AMOStandardTest, gridsize, blocksize, shared_bytes, stream, hipLaunchKernelGGL(AMOStandardTest, gridsize, blocksize, shared_bytes, stream,
loop, args.skip, start_time, end_time, _r_buf, _s_buf, args.loop, args.skip, start_time, end_time, dest, ret_val,
_ret_val, _type, _shmem_context); args.addr_mode, _type, _shmem_context);
_gridSize = gridsize; num_msgs = n_loops * gridsize.x * blocksize.x;
num_msgs = (loop + args.skip) * gridsize.x; num_timed_msgs = args.loop * gridsize.x * blocksize.x;
num_timed_msgs = loop; }
template <typename G>
void fail_eq(const G& got, const G& exp) {
std::cerr << "data validation error\n"
<< "got " << got << ", expected " << exp << std::endl;
std::exit(-1);
}
template <typename G>
void fail_nonzero(const G& got) {
std::cerr << "data validation error\n"
<< "got " << got << ", expected non-zero" << std::endl;
std::exit(-1);
}
// Map (loop, elem_idx) -> dest[] index for current address mode.
template <typename T>
int AMOStandardTester<T>::destIndex(int l, int elem_idx) const {
return (args.addr_mode == AddrMode::PerBlock)
? l * args.num_wgs + elem_idx
: l; // PerGrid has a single element per loop
}
// Number of output elements to check per loop for current address mode.
template <typename T>
int AMOStandardTester<T>::numElems() const {
return (args.addr_mode == AddrMode::PerBlock)
? static_cast<int>(args.num_wgs)
: 1; // PerGrid
}
// Return pointer to the start of the ret_val “chunk” for (loop, elem_idx)
// plus the chunk length for this address mode.
template <typename T>
std::pair<T*, int> AMOStandardTester<T>::retChunk(int l, int elem_idx) const {
if (args.addr_mode == AddrMode::PerBlock) {
// One chunk per element (workgroup): wg_size returns
T* p = ret_val + l * n_in + elem_idx * args.wg_size;
int sz = static_cast<int>(args.wg_size);
return {p, sz};
}
// PerGrid: one big chunk per loop
T* p = ret_val + l * n_in;
int sz = static_cast<int>(n_in);
return {p, sz};
}
template <typename T>
void AMOStandardTester<T>::verifyDestValues() {
const int loops = static_cast<int>(n_loops);
const int n_elems = numElems();
auto check_equal_all = [&](T expected) {
for (int l = 0; l < loops; ++l) {
for (int elem = 0; elem < n_elems; ++elem) {
const int idx = destIndex(l, elem);
if (dest[idx] != expected) fail_eq(dest[idx], expected);
}
}
};
auto check_nonzero_all = [&]() {
for (int l = 0; l < loops; ++l) {
for (int elem = 0; elem < n_elems; ++elem) {
const int idx = destIndex(l, elem);
if (dest[idx] == T{0}) fail_nonzero(dest[idx]);
}
}
};
switch (_type) {
case AMO_AddTestType:
case AMO_FAddTestType: {
const T expected = (args.addr_mode == AddrMode::PerBlock)
? static_cast<T>(args.wg_size * 2)
: static_cast<T>(args.wg_size * args.num_wgs * 2);
check_equal_all(expected);
break;
}
case AMO_IncTestType:
case AMO_FIncTestType: {
const T expected = (args.addr_mode == AddrMode::PerBlock)
? static_cast<T>(args.wg_size)
: static_cast<T>(args.wg_size * args.num_wgs);
check_equal_all(expected);
break;
}
case AMO_FCswapTestType:
check_nonzero_all();
break;
default:
break;
}
}
template <typename T>
void AMOStandardTester<T>::verifyReturnValues() {
// Only “fetch-*” types produce return values to validate.
if (_type == AMO_AddTestType || _type == AMO_IncTestType) return;
const int loops = static_cast<int>(n_loops);
const int n_elems = numElems();
auto check_sorted_sequence = [&](auto value_of_i) {
for (int l = 0; l < loops; ++l) {
for (int elem = 0; elem < n_elems; ++elem) {
auto [p, cnt] = retChunk(l, elem);
std::sort(p, p + cnt);
for (int i = 0; i < cnt; ++i) {
const T expected = static_cast<T>(value_of_i(i));
if (p[i] != expected) fail_eq(p[i], expected);
}
}
}
};
auto check_single_success_zero = [&]() {
for (int l = 0; l < loops; ++l) {
for (int elem = 0; elem < n_elems; ++elem) {
auto [p, cnt] = retChunk(l, elem);
unsigned success = 0;
for (int i = 0; i < cnt; ++i) if (!p[i]) ++success;
if (success != 1u) fail_eq(success, 1u);
}
}
};
switch (_type) {
case AMO_FAddTestType:
check_sorted_sequence([](int i) { return i * 2; });
break;
case AMO_FIncTestType:
check_sorted_sequence([](int i) { return i; });
break;
case AMO_FCswapTestType:
check_single_success_zero();
break;
default:
break;
}
} }
template <typename T> template <typename T>
void AMOStandardTester<T>::verifyResults(size_t size) { void AMOStandardTester<T>::verifyResults(size_t size) {
T ret; // PE 0 checks returns; target PE checks dest.
if (args.myid == 0) { if (args.myid) {
T expected_val = 0; verifyDestValues();
} else {
switch (_type) { verifyReturnValues();
case AMO_FAddTestType:
expected_val = 2 * (num_msgs - 1);
break;
case AMO_FIncTestType:
expected_val = num_msgs - 1;
break;
case AMO_AddTestType:
expected_val = 2 * num_msgs;
break;
case AMO_IncTestType:
expected_val = num_msgs;
break;
case AMO_FCswapTestType:
expected_val = (num_msgs - 2) / _gridSize.x;
break;
default:
break;
}
int fetch_op = (_type == AMO_FAddTestType || _type == AMO_FIncTestType || _type == AMO_FCswapTestType) ? 1: 0;
if (fetch_op == 1) {
ret = *std::max_element(_ret_val, _ret_val + args.num_wgs);
} else {
ret = *std::max_element(_s_buf, _s_buf + args.num_wgs);
}
if (ret != expected_val) {
std::cerr << "data validation error\n";
std::cerr << "got " << ret << ", expected " << expected_val << std::endl;
exit(-1);
}
} }
} }
@@ -121,48 +252,47 @@ void AMOStandardTester<T>::verifyResults(size_t size) {
template <> \ template <> \
__global__ void AMOStandardTest<T>( \ __global__ void AMOStandardTest<T>( \
int loop, int skip, long long int *start_time, \ int loop, int skip, long long int *start_time, \
long long int *end_time, char *r_buf, T *s_buf, T *ret_val, \ long long int *end_time, T *dest, T *ret_val, \
TestType type, ShmemContextType ctx_type) { \ AddrMode addr_mode, TestType type, ShmemContextType ctx_type) { \
__shared__ rocshmem_ctx_t ctx; \ __shared__ rocshmem_ctx_t ctx; \
int wg_id = get_flat_grid_id(); \ int wg_id = get_flat_grid_id(); \
int global_id = get_flat_id(); \
int t_id = get_flat_block_id(); \
int n_threads = get_flat_grid_size(); \
int n_wgs = get_grid_num_blocks(); \
rocshmem_wg_init(); \ rocshmem_wg_init(); \
rocshmem_wg_ctx_create(ctx_type, &ctx); \ rocshmem_wg_ctx_create(ctx_type, &ctx); \
if (hipThreadIdx_x == 0) { \ for (int i = 0; i < loop + skip; i++) { \
T *ptr = compute_target_ptr<T>(dest, addr_mode, wg_id, i, n_wgs); \
T ret = 0; \ T ret = 0; \
T cond = 0; \
for (int i = 0; i < loop + skip; i++) { \
if (i == skip) { \ if (i == skip) { \
start_time[wg_id] = wall_clock64(); \ start_time[wg_id] = wall_clock64(); \
} \
switch (type) { \
case AMO_FAddTestType: \
ret = rocshmem_ctx_##TNAME##_atomic_fetch_add(ctx, (T *)r_buf, 2, \
1); \
break; \
case AMO_FIncTestType: \
ret = \
rocshmem_ctx_##TNAME##_atomic_fetch_inc(ctx, (T *)r_buf, 1); \
break; \
case AMO_FCswapTestType: \
ret = rocshmem_ctx_##TNAME##_atomic_compare_swap(ctx, (T *)r_buf, \
cond, (T)i, 1); \
cond = i; \
break; \
case AMO_AddTestType: \
rocshmem_ctx_##TNAME##_atomic_add(ctx, (T *)r_buf, 2, 1); \
break; \
case AMO_IncTestType: \
rocshmem_ctx_##TNAME##_atomic_inc(ctx, (T *)r_buf, 1); \
break; \
default: \
break; \
} \
} \ } \
rocshmem_ctx_quiet(ctx); \ switch (type) { \
end_time[wg_id] = wall_clock64(); \ case AMO_FAddTestType: \
ret_val[wg_id] = ret; \ ret = rocshmem_ctx_##TNAME##_atomic_fetch_add(ctx, (T *)ptr, 2, 1); \
rocshmem_ctx_getmem(ctx, &s_buf[wg_id], r_buf, sizeof(T), 1); \ break; \
case AMO_FIncTestType: \
ret = rocshmem_ctx_##TNAME##_atomic_fetch_inc(ctx, (T *)ptr, 1); \
break; \
case AMO_FCswapTestType: \
ret = rocshmem_ctx_##TNAME##_atomic_compare_swap(ctx, (T *)ptr, 0, \
(T)(t_id + 1), 1); \
break; \
case AMO_AddTestType: \
rocshmem_ctx_##TNAME##_atomic_add(ctx, (T *)ptr, 2, 1); \
break; \
case AMO_IncTestType: \
rocshmem_ctx_##TNAME##_atomic_inc(ctx, (T *)ptr, 1); \
break; \
default: \
break; \
} \
ret_val[global_id + i * n_threads] = ret; \
} \ } \
rocshmem_ctx_quiet(ctx); \
end_time[wg_id] = wall_clock64(); \
__syncthreads(); \
rocshmem_wg_ctx_destroy(&ctx); \ rocshmem_wg_ctx_destroy(&ctx); \
rocshmem_wg_finalize(); \ rocshmem_wg_finalize(); \
} \ } \
@@ -44,10 +44,19 @@ class AMOStandardTester : public Tester {
virtual void verifyResults(size_t size) override; virtual void verifyResults(size_t size) override;
dim3 _gridSize{}; void verifyDestValues();
char *_r_buf; void verifyReturnValues();
T *_ret_val;
T *_s_buf; int destIndex(int l, int elem_idx) const;
int numElems() const;
std::pair<T*, int> retChunk(int l, int elem_idx) const;
T* dest{nullptr}; // symmetric target buffer [loop][elem]
T* ret_val{nullptr}; // device returns [loop][thread]
size_t n_in{0}; // num_wgs * wg_size
size_t n_out{0}; // elements per loop: PerBlock->num_wgs, PerGrid->1
size_t n_loops{0}; // loop + skip
}; };
#endif #endif
@@ -62,6 +62,13 @@ TesterArguments::TesterArguments(int argc, char *argv[]) {
} else if (arg == "-x") { } else if (arg == "-x") {
i++; i++;
shmem_context = atoi(argv[i]); shmem_context = atoi(argv[i]);
} else if (arg == "-m") {
int atomics_addr_mode = atoi(argv[i]);
if(atomics_addr_mode >= static_cast<int>(AddrMode::PerGrid) &&
atomics_addr_mode <= static_cast<int>(AddrMode::PerBlock)) {
addr_mode = static_cast<AddrMode>(atomics_addr_mode);
}
i++;
} else { } else {
show_usage(argv[0]); show_usage(argv[0]);
exit(-1); exit(-1);
@@ -141,6 +148,7 @@ void TesterArguments::show_usage(std::string executable_name) {
std::cout << "\t-o <Operation type for the random_access test>\n"; std::cout << "\t-o <Operation type for the random_access test>\n";
std::cout << "\t-ta <Number of Thread Accessing the communication>\n"; std::cout << "\t-ta <Number of Thread Accessing the communication>\n";
std::cout << "\t-x <shmem context>\n"; std::cout << "\t-x <shmem context>\n";
std::cout << "\t-m Atomics Address mode\n";
} }
void TesterArguments::get_rocshmem_arguments() { void TesterArguments::get_rocshmem_arguments() {
@@ -39,6 +39,14 @@ enum TeamSplitType {
ROCSHMEM_TEST_TEAM_ODDEVEN, // odd-even splitting ROCSHMEM_TEST_TEAM_ODDEVEN, // odd-even splitting
}; };
/*-----------------------------------------
* Atomics Addressing modes (contention model)
*-----------------------------------------*/
enum class AddrMode : int {
PerGrid, // all WGs -> same address
PerBlock, // each WG -> its own address (default)
};
class TesterArguments { class TesterArguments {
public: public:
TesterArguments(int argc, char *argv[]); TesterArguments(int argc, char *argv[]);
@@ -69,6 +77,7 @@ public:
unsigned coal_coef = 64; unsigned coal_coef = 64;
unsigned op_type = 0; unsigned op_type = 0;
unsigned shmem_context = rocshmem::ROCSHMEM_CTX_WG_PRIVATE; unsigned shmem_context = rocshmem::ROCSHMEM_CTX_WG_PRIVATE;
AddrMode addr_mode = AddrMode::PerBlock;
/** /**
* Arguments obtained from rocshmem * Arguments obtained from rocshmem