Skip to content

[PAL] Support Device API Transport #445

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MC952-arch wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[PAL] Support Device API Transport #445

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
84a262a
Unify barrier/team tags and rename Net→Transport in device API
MC952-arch Apr 3, 2026
22497ed
Fix several issues
MC952-arch Apr 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion flagcx/adaptor/flagcx_device.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -696,6 +696,17 @@ flagcxResult_t flagcxDevCommCreate(flagcxComm_t comm,
if (ret == flagcxSuccess)
handle->devComm = innerDevComm;
}
if (handle->devComm != nullptr) {
int nNodes = 0;
if (comm->heteroComm != nullptr && comm->heteroComm->nNodes > 0) {
nNodes = comm->heteroComm->nNodes;
} else if (handle->intraSize > 0 &&
handle->nRanks % handle->intraSize == 0) {
nNodes = handle->nRanks / handle->intraSize;
}
if (nNodes > 0)
handle->nInterPeers = nNodes - 1;
}
if (handle->devComm == nullptr) {
// ---- Fallback path: IPC barriers + inter-node signal relay + one-sided
// ----
Expand Down Expand Up @@ -1189,4 +1200,4 @@ flagcxResult_t flagcxInterBarrierCreateRequirement(
(void)outHandle;
(void)outReq;
return flagcxNotSupported;
}
}
48 changes: 23 additions & 25 deletions flagcx/adaptor/include/device_api/comm_traits.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
*
* Architecture:
* PlatformTraits<P> — platform-level: Intrin, Atomic
* CommTraits<D> — backend-level: Window, DevComm, Team, ...
* CommTraits<D> — backend-level: Window, Comm, Team, ...
* Fallback<PlatformTag> — common IPC fallback (partial specialization)
*
* CommTraits pulls in platform capabilities via using-aliases (not
Expand All @@ -14,7 +14,7 @@
* types that work with any platform.
*
* Selection:
* NVIDIA + NCCL > 2.28: DeviceAPI = CommTraits<NvidiaVendor>
* NVIDIA + NCCL > 2.28: DeviceAPI = CommTraits<NvidiaVendor>
* NVIDIA + fallback: DeviceAPI = CommTraits<Fallback<NvidiaPlatform>>
*
* Kernel code uses DeviceAPI::* exclusively, no #ifdef branches.
Expand All @@ -36,51 +36,49 @@ template <typename PlatformTag>
struct Fallback {};

// ============================================================
// Action types for one-sided operations (needed by traits Net types).
// Action types for one-sided operations (needed by traits Transport types).
// Pure POD structs with no device builtins.
// ============================================================
typedef uint32_t flagcxDevNetSignal_t;
typedef uint32_t flagcxDevNetCounter_t;
typedef uint32_t flagcxDevTransportSignal_t;
typedef uint32_t flagcxDevTransportCounter_t;

struct flagcxDevNet_None {};
struct flagcxDevNet_SignalInc {
flagcxDevNetSignal_t signal;
struct flagcxDevTransport_None {};
struct flagcxDevTransport_SignalInc {
flagcxDevTransportSignal_t signal;
};
struct flagcxDevNet_SignalAdd {
flagcxDevNetSignal_t signal;
struct flagcxDevTransport_SignalAdd {
flagcxDevTransportSignal_t signal;
uint64_t value;
};
struct flagcxDevNet_CounterInc {
flagcxDevNetCounter_t counter;
struct flagcxDevTransport_CounterInc {
flagcxDevTransportCounter_t counter;
};

// Shared memory descriptor for NIC descriptor optimization.
// Uses void* on all paths; vendor Net casts to native type in toNccl().
// Uses void* on all paths; vendor Transport casts to native type in toNccl().
struct flagcxDescriptorSmem {
void *_impl = nullptr;
};

struct flagcxDevNet_DescriptorSmem {
struct flagcxDevTransport_DescriptorSmem {
flagcxDescriptorSmem smem;
};

// Fence level enum — available on all tiers for unified barrier API
enum class flagcxGinFenceLevel { Relaxed };
enum class flagcxTransportFenceLevel { Relaxed };

// ============================================================
// Barrier tag types for DevBarrier<Backend, Tag> dispatch.
// Unified team/barrier tag types.
// Used as both Barrier<Backend, Tag> template parameter
// and as ctor dispatch tags — eliminating the old two-tag redundancy.
// ============================================================
struct flagcxBarrierIntra {};
struct flagcxBarrierInter {};
struct flagcxBarrierWorld {
struct World {}; // tag for world-barrier ctor
struct Intra {}; // tag for intra-only ctor
struct Inter {}; // tag for inter-only ctor
};
struct flagcxTeamTagIntra {};
struct flagcxTeamTagInter {};
struct flagcxTeamTagWorld {};

// Primary template — each backend provides specializations
template <typename Backend, typename BarrierTag, typename Coop>
struct DevBarrier;
template <typename Backend, typename Tag, typename Coop>
struct Barrier;

// Vendor specializations + DeviceAPI selection
#if defined(USE_NVIDIA_ADAPTOR)
Expand Down
Loading