Brandon Potter
@@ -0,0 +1,174 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved.
|
||||
*
|
||||
* Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
* of this software and associated documentation files (the "Software"), to
|
||||
* deal in the Software without restriction, including without limitation the
|
||||
* rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
|
||||
* sell copies of the Software, and to permit persons to whom the Software is
|
||||
* furnished to do so, subject to the following conditions:
|
||||
*
|
||||
* The above copyright notice and this permission notice shall be included in
|
||||
* all copies or substantial portions of the Software.
|
||||
*
|
||||
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
|
||||
* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
|
||||
* IN THE SOFTWARE.
|
||||
*****************************************************************************/
|
||||
|
||||
#include "wf_coal_policy.hpp"
|
||||
|
||||
#include "util.hpp"
|
||||
|
||||
namespace rocshmem {
|
||||
|
||||
/**
|
||||
* TODO: Determine tradeoffs between number of instructions, LDS utilization,
|
||||
* and quality of coalescing.
|
||||
*/
|
||||
|
||||
/**
|
||||
* About the current algorithm:
|
||||
* 1) The algorithm is low overhead. It uses no LDS (shared) space or
|
||||
* rounds of a tree reduction (which may cause it to miss some
|
||||
* scenarios).
|
||||
* 2) The algorithm can identify multiple groups of coalescable packets
|
||||
* within a wavefront.
|
||||
* As an example:
|
||||
* Threads 0-8 are coalescable.
|
||||
* Threads 10-63 are coalescable.
|
||||
* The logic will generate 3 messages.
|
||||
* +-----------------------+----------+-------------------------+
|
||||
* | Thread_0 ... Thread_8 | Thread_9 | Thread_10 ... Thread_63 |
|
||||
* +-----------------------+----------+-------------------------+
|
||||
* | msg_1 | msg 2 | msg 3 |
|
||||
* +-----------------------+----------+-------------------------+
|
||||
* However, threads must be contiguous by thread id to be coalesced.
|
||||
* For example, Thread_0 and Thread_2 are coalescable, but Thread_1 is
|
||||
* not coalescable (see coalescability conditions) then none of the
|
||||
* three threads can be coalesced.
|
||||
* 3) The algorithm misses opportunities when threads have coalescable
|
||||
* messages but the message sizes are different.
|
||||
*/
|
||||
__device__ bool WfCoalOn::coalesce(int pe, const void *source, const void *dest,
|
||||
size_t *size) {
|
||||
const uint64_t src = (const uint64_t)source;
|
||||
const uint64_t dst = (const uint64_t)dest;
|
||||
|
||||
/**
|
||||
* Split 64-bit values into high and low for 32-bit shuffles.
|
||||
* Unfortunately, the shuffle operations only support 32-bit widths.
|
||||
*/
|
||||
uint32_t src_low = uint32_t(src & 0xFFFFFFFF);
|
||||
uint32_t src_high = uint32_t((src >> 32) & 0xFFFFFFFF);
|
||||
uint32_t dst_low = uint32_t(dst & 0xFFFFFFFF);
|
||||
uint32_t dst_high = uint32_t((dst >> 32) & 0xFFFFFFFF);
|
||||
|
||||
/**
|
||||
* Shuffle message info to upwards neighboring threads.
|
||||
* +----------------------------------------------------------------+
|
||||
* | Thread_0 Thread_1 ... ... ... ... Thread_62 Thread_63 |
|
||||
* +----------------------------------------------------------------+
|
||||
* | Upwards Downwards |
|
||||
* +----------------------------------------------------------------+
|
||||
*
|
||||
* The implementation of __shfl_up comes from the hip header files.
|
||||
* In rocm 2.10, the filename is device_functions.h.
|
||||
*/
|
||||
uint64_t lower_src_low = __shfl_up(src_low, 1);
|
||||
uint64_t lower_src_high = __shfl_up(src_high, 1);
|
||||
uint64_t lower_dst_low = __shfl_up(dst_low, 1);
|
||||
uint64_t lower_dst_high = __shfl_up(dst_high, 1);
|
||||
int lower_pe = __shfl_up(pe, 1);
|
||||
size_t lower_size = __shfl_up((unsigned int)*size, 1);
|
||||
|
||||
/**
|
||||
* Recombine the incoming 64-bit values from neighbor.
|
||||
*/
|
||||
uint64_t lower_src = (lower_src_high << 32) | lower_src_low;
|
||||
uint64_t lower_dst = (lower_dst_high << 32) | lower_dst_low;
|
||||
|
||||
/**
|
||||
* The mask variable tells us which threads are active in the wavefront.
|
||||
* An active thread will call into this function with a value '1' which
|
||||
* notifies the other threads that the lane is active.
|
||||
*/
|
||||
uint64_t mask = __ballot(1);
|
||||
|
||||
/**
|
||||
* The wv_id variable holds the wavefront id number. To set it, we
|
||||
* flatten the thread block out (to make it one-dimensional) and then
|
||||
* modulo based of the wavefront size (which is a characteristic of
|
||||
* the hardware).
|
||||
*/
|
||||
int wv_id = get_flat_block_id() % WF_SIZE;
|
||||
|
||||
/**
|
||||
* If coalescable evaluates to true, this thread is __NOT__ responsible
|
||||
* for sending a message (another thread will send the message on its
|
||||
* behalf). In other words, the thread's message is coalesced, yay.
|
||||
*
|
||||
* If coalescable evaluates to false, this thread is responsible for
|
||||
* sending a message (which means that it is not coalescable with its
|
||||
* lower neighbor).
|
||||
*/
|
||||
bool coalescable =
|
||||
(mask & (1LL << (wv_id - 1))) && // Ensure lower lane is active
|
||||
(lower_size == *size) && // Ensure lower lane size is equal
|
||||
((lower_src + *size) == src) && // Ensure I cover lower src
|
||||
((lower_dst + *size) == dst) && // Ensure I cover lower dst
|
||||
(pe == lower_pe) && // Must be sending to the same pe
|
||||
(wv_id != 0); // Thread_0 is never coalescable
|
||||
|
||||
/**
|
||||
* Share the lower neighbor coalescability status with all the active
|
||||
* threads in the wavefront.
|
||||
*
|
||||
* Inactive threads will not participate in the ballot which returns '0'
|
||||
* in their position.
|
||||
*/
|
||||
uint64_t lowerNeighborCoal = __ballot(coalescable);
|
||||
|
||||
/**
|
||||
* If the thread is not coalescable, it must send a message.
|
||||
* It needs to check how many threads are considered coalesced above it
|
||||
* to adjust its message size.
|
||||
*
|
||||
* Do this by counting the number of contiguous '1's greater than its
|
||||
* thread ID from the ballot function. It will coalesce the messages for
|
||||
* all contiguous higher threads which report that they are coalescable
|
||||
* with their immediate lower neighbor.
|
||||
*/
|
||||
if (!coalescable) {
|
||||
int coal_size = *size;
|
||||
|
||||
/**
|
||||
* Ignore the lower threads and the thread's own position.
|
||||
*/
|
||||
lowerNeighborCoal >>= (wv_id + 1);
|
||||
|
||||
/**
|
||||
* Invert and find the first bit index set to zero. This bit
|
||||
* indicates the first higher thread which is not coalescable with
|
||||
* its lower neighbor.
|
||||
*
|
||||
* This thread is now responsible for coalescing everything
|
||||
* between its own index and that one.
|
||||
*/
|
||||
uint32_t coalMsgs = __ffsll((unsigned long long)~lowerNeighborCoal);
|
||||
|
||||
if (coalMsgs) {
|
||||
coal_size += *size * (coalMsgs - 1);
|
||||
}
|
||||
|
||||
*size = coal_size;
|
||||
}
|
||||
|
||||
return !coalescable;
|
||||
}
|
||||
|
||||
} // namespace rocshmem
|
||||
Reference in New Issue
Block a user