Skip to content

[TLE] Add device intra node barrier#739

Open
Galaxy1458 wants to merge 23 commits into
mainfrom
add_intra_node_barrier
Open

[TLE] Add device intra node barrier#739
Galaxy1458 wants to merge 23 commits into
mainfrom
add_intra_node_barrier

Conversation

@Galaxy1458

@Galaxy1458 Galaxy1458 commented Jul 2, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR adds support for GPU-to-GPU (device-to-device) barriers within a node in the TLE backend, enabling synchronization across multiple GPUs through the FlagCX runtime.

Changes

Python frontend

  • Extend check_and_handle_device_intra_barrier() to recognize device barriers.
  • Lower distributed_barrier(..., space="device", comm_ptr=...) into the tle.device_intra_barrier operation.

MLIR

Introduce a new operation:

  • tle.device_intra_barrier for FlagCxOp

The operation represents an intra-node GPU-to-GPU synchronization barrier with configurable:

  • barrier type (arrive, wait, sync)
  • cooperative group scope
  • barrier index
  • memory ordering
  • multimem mode

LLVM Lowering

Lower tle.device_intra_barrier to the corresponding FlagCX runtime APIs:

  • flagcxIntraBarrierArriveS
  • flagcxIntraBarrierWaitS
  • flagcxIntraBarrierSyncS

The runtime function is selected based on the requested barrier type.

Motivation

This PR introduces a dedicated IR representation for intra-node GPU-to-GPU synchronization, instead of representing it as a generic distributed barrier. It also establishes the infrastructure for future multi-GPU synchronization primitives built on the FlagCX runtime.

Example:

 tle.distributed_barrier(
              comm_ptr=dev_comm_dptr, space="device", 
              group_kind="block", order="acqrel",barrier_kind="sync"
)

@Galaxy1458 Galaxy1458 requested a review from sunnycase as a code owner July 2, 2026 09:50
@Galaxy1458 Galaxy1458 marked this pull request as draft July 2, 2026 09:50
@github-actions github-actions Bot added the nvidia label Jul 3, 2026
@Galaxy1458 Galaxy1458 marked this pull request as ready for review July 9, 2026 03:08
@github-actions github-actions Bot added the main label Jul 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant