ask a question

intranode

dispatch receiver :

// Receive channel offset
int total_offset, num_tokens_to_recv;
if (elect_one_sync()) {
while ((total_offset = ld_volatile_global(channel_start_offset.buffer())) == 0)
;
while ((num_tokens_to_recv = ld_volatile_global(channel_end_offset.buffer())) == 0)
;
total_offset = -total_offset - 1, num_tokens_to_recv = -num_tokens_to_recv - 1;
if (recv_warp_id_in_rank == 0) // record in RX side. right shift by one
recv_channel_offset[responsible_rank * num_channels + responsible_channel] = total_offset;
num_tokens_to_recv -= total_offset; // end - start -> num_tokens
}

combine sender:
// Get tasks
// NOTES: `channel_offset` is already shifted
int rank_offset = send_rank_id > 0 ? rank_prefix_matrix[(send_rank_id - 1) * kNumRanks + rank] : 0;
int num_rank_tokens = rank_prefix_matrix[send_rank_id * kNumRanks + rank] - rank_offset;
int channel_offset = channel_prefix_matrix[send_rank_id * num_channels + responsible_channel];
int num_channel_tokens =
(responsible_channel == num_channels - 1 ? num_rank_tokens
: channel_prefix_matrix[send_rank_id * num_channels + responsible_channel + 1]) -
channel_offset;
int token_start_idx = rank_offset + channel_offset, token_end_idx = rank_offset + channel_offset + num_channel_tokens;

question：

what's the benifit of shift total_offset in recv_channel_offset? why not just store the received num_tokens_to_recv, which is the end of the channel