【NCCL】1. nccl通信API 到设备端 RunWorkColl 的完整调用链

原创已于 2025-12-25 10:57:55 修改 · 837 阅读

15 ·

本内容遵循CC 4.0 BY-SA版权协议

GEO检测

收录于

NCCL

于 2025-12-19 17:48:51 首次发布

Commit: 59242d7c

ncclInfo

chunkSteps（时间维度 - 流水线）

作用：控制多少个chunk可以同时在流水线中处理

// 示例：chunkSteps = 4
时序图：
时间 │ Chunk0 │ Chunk1 │ Chunk2 │ Chunk3 │
t=0  │ RECV   │        │        │        │
t=1  │ COMPUTE│ RECV   │        │        │
t=2  │ SEND   │ COMPUTE│ RECV   │        │
t=3  │ DONE   │ SEND   │ COMPUTE│ RECV   │
t=4  │        │ DONE   │ SEND   │ COMPUTE│

特点：

时间上的重叠：通信、计算、数据传输可以同时进行
隐藏延迟：当Chunk0在发送时，Chunk1已经在计算，Chunk2在接收
内存开销：需要为每个在流水线中的chunk分配buffer

sliceSteps（空间维度 - 并行）

作用：控制每个chunk被分成多少份由不同线程并行处理

// 示例：sliceSteps = 4
Chunk i（128个元素）：
Slice 0: 元素[0-31]   ← 线程组0处理
Slice 1: 元素[32-63]  ← 线程组1处理  
Slice 2: 元素[64-95]  ← 线程组2处理
Slice 3: 元素[96-127] ← 线程组3处理

特点：

空间上的并行：多个线程同时处理同一个chunk的不同部分
负载均衡：均匀分配计算工作
线程协作：线程间可能需要同步

ncclAllGather 到设备端 RunWorkColl 的完整调用链

追踪出了从主机端 API ncclAllGather 到设备端模板函数 RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE>::run() 的完整调用路径。

调用链概览

详细调用流程

第一阶段：主机端任务创建（Host-side Task Creation）

1. `ncclAllGather()` - 用户API

位置: src/collectives.cc:82-92
功能: NCCL AllGather 集合通信的公共API入口
操作: 创建 ncclInfo 结构体并调用 ncclEnqueueCheck

2. `ncclEnqueueCheck()` - 参数验证

位置: src/enqueue.cc:2620
功能: 验证参数并启动组操作
关键调用: taskAppend(info->comm, info)

3. `taskAppend()` - 任务分发

位置: src/enqueue.cc:2548
功能: 根据操作类型分发到不同处理函数
关键调用: collTaskAppend(comm, info, opDev)

4. `collTaskAppend()` - 创建集合通信任务

位置: src/enqueue.cc:2453
功能: 创建 ncclTaskColl 任务并加入排序器

关键操作:

struct ncclTaskColl* t = ncclMemoryPoolAlloc<struct ncclTaskColl>(...);
// 填充任务信息
t->func = ncclFuncAllGather;
t->sendbuff = info->sendbuff;
t->recvbuff = info->recvbuff;
// ...
ncclTaskCollSorterInsert(&planner->collSorter, t, t->trafficBytes);

第二阶段：任务调度与计划（Task Scheduling & Planning）

5. `ncclPrepareTasks()` - 准备任务

位置: src/enqueue.cc:348
功能: 从排序器取出任务，选择算法和协议
关键操作:
- 调用 getAlgoInfo() 选择最优算法（如 NCCL_ALGO_PAT）和协议（如 NCCL_PROTO_SIMPLE）
- 计算 devFuncId = ncclDevFuncId(func, op, datatype, algorithm, protocol)
- 将任务加入 planner->collTaskQueue

6. `scheduleCollTasksToPlan()` - 调度任务到执行计划

位置: src/enqueue.cc:519
功能: 将任务分配到通道并创建内核计划
关键操作:
- 调用 addWorkBatchToPlan() 添加工作批次
- 设置 plan->kernelFn = ncclDevKernelForFunc[task->devFuncId]
- 将任务加入 plan->workQueue

7. `finishPlan()` - 完成计划

位置: src/enqueue.cc:182
功能: 准备内核参数和工作批次
关键操作: 设置 plan->kernelArgs 结构体

8. `uploadWork()` - 上传工作数据

位置: src/enqueue.cc:1151
功能: 将工作结构体上传到设备内存

第三阶段：内核启动（Kernel Launch）

9. `ncclLaunchKernel()` - 启动CUDA内核

位置: src/enqueue.cc:1565
功能: 启动CUDA内核

关键代码:

void* sym = plan->kernelFn;  // 指向 ncclDevKernel_AllGather_PAT_SIMPLE
CUfunction fn;
cudaGetFuncBySymbol(&fn, sym);
cuLaunchKernelEx(&launchConfig, fn, nullptr, extra);

第四阶段：设备端执行（Device-side Execution）

10. 生成的内核函数 - `ncclDevKernel_AllGather_PAT_SIMPLE`

生成位置: src/device/generate.py 脚本生成
实际代码: 在编译时生成到 gensrc/all_gather.cu

宏定义: src/device/common.h:406-409

#define DEFINE_ncclDevKernel(suffix, coll, redop, ty, algo, proto, specializedFnId)
  __global__ void ncclDevKernel_##suffix(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
    ncclKernelMain<specializedFnId, RunWorkBatch<coll, ty, redop<ty>, algo, proto>>(&args4K.args);
  }

展开后:

__global__ void ncclDevKernel_AllGather_PAT_SIMPLE(ncclDevKernelArgs4K const args4K) {
  ncclKernelMain<fnId, RunWorkBatch<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE>>(&args4K.args);
}

11. `ncclKernelMain()` - 内核主函数

位置: src/device/common.h:332
功能: 设备端内核的主入口
关键操作:
- 加载内核参数到共享内存
- 加载通信器和通道信息
- 调用 loadWorkBatchToShmem() 加载工作批次
- 调用 SpecializedRunWorkBatch().run() 或 ncclDevFuncTable[funcId]()

12. `RunWorkBatch<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE>::run()`

位置: src/device/common.h:265
功能: 执行工作批次

关键代码:

template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto>
struct RunWorkBatch {
  __device__ __forceinline__ void run() {
    for (int w=0; w < ncclShmem.nWorks; w++) {
      struct ncclDevWorkColl* work = (struct ncclDevWorkColl*)(ncclShmem.workStorage + w*ncclShmem.workSize);
      int subtn = work->nWarps*WARP_SIZE;
      if (tid < subtn) 
        RunWorkColl<Fn, T, RedOp, Algo, Proto>().run(tid, subtn, work);  // ← 调用这里
    }
  }
};

13. `RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE>::run()` - 目标函数

位置: src/device/all_gather.h:112-172
功能: 执行 AllGather 的 PAT 算法

关键代码:

template<typename T, typename RedOp>
struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE> {
  __device__ __forceinline__ void run(int tid, int nthreads, struct ncclDevWorkColl* work) {
    // PAT (Parallel Aggregated Trees) 算法实现
    // 使用 PatAGAlgorithm 计算步骤
    // 使用 Primitives 执行数据传输
    PatAGAlgorithm<T> patAlgo(...);
    Primitives<T, RedOp, ...> prims(...);
    // 执行 AllGather 操作
    prims.patCopy(ps, shmem);
  }
};

代码生成机制

NCCL 使用 Python 脚本 generate.py 在编译时生成大量的内核和函数变体：

生成的文件

gensrc/all_gather.cu - 包含 AllGather 的所有算法/协议组合
gensrc/device_table.cu - 设备端函数表 ncclDevFuncTable[]
gensrc/host_table.cc - 主机端内核列表 ncclDevKernelList[]

生成的内核示例

// 为 AllGather + PAT + SIMPLE 生成的内核
DEFINE_ncclDevKernel(AllGather_PAT_SIMPLE, 
                     ncclFuncAllGather, 
                     FuncCopy, 
                     int8_t, 
                     NCCL_ALGO_PAT, 
                     NCCL_PROTO_SIMPLE, 
                     fnId)

函数ID映射

ncclDevFuncId(): 根据 (func, op, datatype, algo, proto) 计算函数ID
ncclDevFuncRowToId[]: 函数行到主函数ID的映射
ncclDevKernelForFunc[]: 主函数ID到内核指针的映射
ncclDevFuncTable[]: 设备端函数指针表

关键数据结构

主机端

ncclInfo: 用户API参数
ncclTaskColl: 集合通信任务
ncclKernelPlan: 内核执行计划
ncclDevKernelArgs: 内核参数

设备端

ncclShmemData: 共享内存数据（包含通信器、通道、工作批次）
ncclDevWorkColl: 设备端工作结构体
ncclDevWorkBatch: 工作批次描述符

总结

从 ncclAllGather 到 RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_PAT, NCCL_PROTO_SIMPLE>::run() 的完整调用链包括：

API层 (主机): ncclAllGather → ncclEnqueueCheck → taskAppend → collTaskAppend
调度层 (主机): ncclPrepareTasks → scheduleCollTasksToPlan → finishPlan
启动层 (主机): uploadWork → ncclLaunchKernel
内核层 (设备): ncclDevKernel_AllGather_PAT_SIMPLE → ncclKernelMain
执行层 (设备): RunWorkBatch::run → RunWorkColl::run ← 目标函数

整个过程通过模板特化和代码生成实现了高度优化的设备端执行，支持多种算法（RING, TREE, PAT, NVLS等）和协议（SIMPLE, LL, LL128）的组合。

算法成本

ncclCommInitRankFunc => initTransportsRank => ncclTopoTuneModel 中设置基础的 latencies。
然后再在每一次通信api调用的时候，在 ncclTopoGetAlgoTime 中计算成本
*time = lat * latCount + nBytes / (1000 * bw);，这里面 lat 是从 latencies 多维数组中读取的值，latCount 是根据拓扑和算法计算的，再加上数据量除以带宽计算出一个理论时间成本。
当前通信op，各个算法的理论成本中最小的，就是要选的算法。

标签

#GPU #NCCL #通信库