Add: AllGather and ReduceScatter distributed L3 examples

[georgebisbas]

Summary

Adds two new symmetric collective communication examples following the
allreduce_distributed pattern.

New files

examples/workers/l3/allgather_distributed/

3-phase symmetric AllGather on the HCCL window scratch pattern:

Phase	Description
1 stage-in	input[0..C) → my scratch slot
2 barrier	TNOTIFY/TWAIT pairwise signal matrix
3 gather	for r in 0..N-1: TLOAD(peer scratch) → TSTORE(output[r*C])

• COUNT_PER_RANK = 64, window = 4 KB
• Output is identical on every rank (rank-ordered concatenation)
• Golden: output[r*C + i] = r*100 + i (closed-form)
• CommRemotePtr<T> defined inline (same pattern as allreduce_distributed)

examples/workers/l3/reduce_scatter_distributed/

4-phase symmetric ReduceScatter on the HCCL window scratch pattern:

Phase	Description
1 stage-in	all N chunks → scratch (nranks × C floats)
2 barrier	TNOTIFY/TWAIT pairwise signal matrix
3 reduce	acc = own scratch[my_rank*C]; TADD all peers via CommRemotePtr
4 stage-out	TSTORE acc → output

• COUNT_PER_RANK = 64; scratch window is nranks-dependent (computed at runtime)
• Each rank produces its own shard; golden per dest: nranks*(dest*C+j) + 100*nranks*(nranks-1)/2
• CommRemotePtr<T> defined inline (same pattern as allreduce_distributed)

Testing

Both examples mirror test_allreduce.py:

• Platforms: a2a3sim, a2a3, a5sim
• Runtime: tensormap_and_ringbuffer
• Parametrized: 2 and 4 devices

pytest examples/workers/l3/allgather_distributed/ -v
pytest examples/workers/l3/reduce_scatter_distributed/ -v
# or run standalone:
python examples/workers/l3/allgather_distributed/main.py -p a2a3sim -d 0-1
python examples/workers/l3/reduce_scatter_distributed/main.py -p a2a3sim -d 0-3

[gemini-code-assist]

Code Review

This pull request introduces end-to-end distributed implementations for AllGather and ReduceScatter operations, including AIV kernels, orchestration shims, and Python execution scripts. The feedback identifies critical synchronization risks in the barrier implementation, recommending a sense-reversing or phase-based mechanism to prevent race conditions when memory is reused. Additionally, there are suggestions to improve the clarity of the buffer specifications in the Python scripts by ensuring the element count accurately reflects the total allocated memory, including signal areas.

[gemini-code-assist]

The barrier implementation should use a sense-reversing or phase-based mechanism to prevent race conditions in distributed environments where memory is reused. Each participant must snapshot the current phase before incrementing the arrival counter. This ensures that fast participants do not interfere with the current barrier generation by looping back too quickly, which is a risk if the signal area is not managed correctly.

References

1. When implementing a reusable barrier with atomic counters, use a sense-reversing or phase-based mechanism. Each participant should snapshot the current phase before incrementing the arrival counter to prevent race conditions where fast participants loop back and interfere with the current barrier generation.

[gemini-code-assist]

To ensure robust synchronization when scratch memory is reused, implement a sense-reversing or phase-based barrier. It is critical to snapshot the current phase before incrementing the arrival counter to prevent race conditions between successive barrier generations, ensuring that no rank proceeds based on stale signal values.

References

1. When implementing a reusable barrier with atomic counters, use a sense-reversing or phase-based mechanism. Each participant should snapshot the current phase before incrementing the arrival counter to prevent race conditions where fast participants loop back and interfere with the current barrier generation.

[gemini-code-assist]

The count parameter in CommBufferSpec is set to COUNT_PER_RANK (64), but the actual buffer size SCRATCH_NBYTES (320) includes additional space for signals. While nbytes is correctly provided, having a count that only reflects the data portion and not the full buffer capacity (including the signal tail) can be misleading for maintenance. Consider setting count to the total number of elements or ensuring it aligns with the intended usage of the buffer.

[gemini-code-assist]

The count in CommBufferSpec for the scratch buffer only accounts for the float data (nranks * COUNT_PER_RANK), but the buffer also contains the signal area. It is better to define the buffer specification such that the count and dtype accurately represent the allocated memory or rely solely on nbytes if the buffer is heterogeneous.