【NCCL】3. ncclPrepareTasks 到 scheduleCollTasksToPlan 的衔接机制

Commit: 59242d7c

ncclPrepareTasks 到 scheduleCollTasksToPlan 的衔接机制

ncclPrepareTasks 之后如何衔接到 scheduleCollTasksToPlan 的完整流程。关键在于 ncclLaunchPrepare 函数。

完整调用链

ncclGroupEndInternal
groupLaunch
ncclPrepareTasksAndCollPreconnect
ncclPrepareTasks
任务加入 collTaskQueue
ncclTasksRegAndEnqueue
创建 ncclDevWorkColl
doLaunches
ncclLaunchPrepare
scheduleCollTasksToPlan
finishPlan
ncclLaunchKernel

详细流程说明

1. ncclPrepareTasks 的输出

ncclPrepareTasks() 完成后，任务被组织在以下队列中：

• planner->collTaskQueue: 包含所有集合通信任务（已选择算法和协议）
• planner->collWorkQueue : 对应的工作结构体队列（在 ncclTasksRegAndEnqueue 中创建）

ncclResult_t ncclPrepareTasks(...) {
  // 1. 从 collSorter 取出所有任务
  struct ncclTaskColl* task = ncclTaskCollSorterDequeueAll(&planner->collSorter);
  
  // 2. 选择算法和协议
  NCCLCHECK(getAlgoInfo(comm, &agg, ...));
  agg.devFuncId = ncclDevFuncId(agg.func, agg.opDev.op, agg.datatype, agg.algorithm, agg.protocol);
  
  // 3. 将任务加入 collTaskQueue
  ncclIntruQueueTransfer(&planner->collTaskQueue, &collBins[isCollnet][isNvls]);
  
  return ncclSuccess;
}

2. 中间步骤：ncclTasksRegAndEnqueue

在 groupLaunch() 中，调用 ncclTasksRegAndEnqueue 为每个任务创建设备端工作结构体：

// 在 groupLaunch() 中
comm = groupCommHeadMain[ncclGroupTaskTypeCollective];
do {
  NCCLCHECKGOTO(ncclTasksRegAndEnqueue(comm), ret, fail);  // ← 创建 ncclDevWorkColl
  comm = comm->groupNext[ncclGroupTaskTypeCollective];
} while (comm);

ncclTasksRegAndEnqueue() 的作用：

ncclResult_t ncclTasksRegAndEnqueue(struct ncclComm* comm) {
  struct ncclTaskColl *task = ncclIntruQueueHead(&planner->collTaskQueue);
  while (task != nullptr) {
    // 创建 ncclDevWorkColl 结构体
    struct ncclDevWorkColl devWork = {};
    devWork.sendbuff = (void*)task->sendbuff;
    devWork.recvbuff = (void*)task->recvbuff;
    // ...
    
    // 包装成 ncclWorkList 节点
    workNode = ncclMemoryStackAllocInlineArray<ncclWorkList, ncclDevWorkColl>(...);
    workNode->workType = ncclDevWorkTypeColl;
    workNode->size = sizeof(struct ncclDevWorkColl);
    memcpy((void*)(workNode+1), (void*)&devWork, workNode->size);
    
    // 加入 collWorkQueue
    ncclIntruQueueEnqueue(&planner->collWorkQueue, workNode);
    task = task->next;
  }
}

3. 启动内核准备：doLaunches

groupLaunch() 调用 doLaunches：

if (!simInfo && groupCommHeadMain[ncclGroupTaskTypeCollective] != nullptr) {
  NCCLCHECKGOTO(doLaunches(groupCommHeadMain[ncclGroupTaskTypeCollective]), ret, fail);
}

4. 关键衔接点：ncclLaunchPrepare

doLaunches() 中调用 ncclLaunchPrepare：

static ncclResult_t doLaunches(struct ncclComm* head) {
  do {
    struct ncclComm* comm = cliqueHead;
    do {
      CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), result, failure);
      NCCLCHECKGOTO(ncclLaunchPrepare(comm), result, failure);  // ← 关键调用
      // ...
      comm = comm->groupNext[ncclGroupTaskTypeCollective];
    } while (comm != nullptr && comm->intraComm0 == cliqueHead->intraComm0);
    
    // 然后启动内核
    while (true) {
      comm = cliqueHead;
      do {
        struct ncclKernelPlan* plan = comm->planner.unlaunchedPlansHead;
        if (plan != nullptr) {
          NCCLCHECKGOTO(ncclLaunchKernel(comm, plan), result, failure);  // ← 启动内核
        }
        comm = next;
      } while (comm != cliqueNextHead);
    }
  } while (cliqueHead != nullptr);
}

5. 调度任务到计划：ncclLaunchPrepare

ncclLaunchPrepare() - 这里是关键衔接点！

ncclResult_t ncclLaunchPrepare(struct ncclComm* comm) {
  struct ncclKernelPlanner* planner = &comm->planner;
  
  // 只要还有任务，就继续创建计划
  if (planner->nTasksColl + planner->nTasksP2p != 0 || ...) {
    do {
      // 创建新的内核计划
      struct ncclKernelPlan* plan = ncclMemoryPoolAlloc<struct ncclKernelPlan>(...);
      plan->comm = comm;
      plan->persistent = persistent;
      
      // 设置预算
      struct ncclKernelPlanBudget budget;
      budget.inArgsBytes = comm->workArgsBytes - sizeof(struct ncclDevKernelArgs);
      budget.outArgsBytes = plan->persistent ? (1<<30) : comm->workFifoBytes/2;
      
      // ← 关键调用：调度集合通信任务到计划
      if (planner->nTasksColl != 0) {
        NCCLCHECKGOTO(scheduleCollTasksToPlan(comm, plan, &budget), result, failure);
      }
      
      // 调度 P2P 任务
      if (planner->nTasksColl == 0 && planner->nTasksP2p != 0) {
        NCCLCHECKGOTO(scheduleP2pTasksToPlan(comm, plan, &budget), result, failure);
      }
      
      // 完成计划
      finishPlan(comm, plan);
      
      // 将计划加入队列
      if (plan->workBytes != 0) {
        ncclIntruQueueEnqueue(&planner->planQueue, plan);
        nPlans += 1;
      }
      
    } while (planner->nTasksColl + planner->nTasksP2p != 0 || ...);
    
    // 保存未启动的计划头
    planner->unlaunchedPlansHead = ncclIntruQueueHead(&planner->planQueue);
  }
  
  return ncclSuccess;
}

6. 调度任务到计划：scheduleCollTasksToPlan

scheduleCollTasksToPlan() 的核心功能：

static ncclResult_t scheduleCollTasksToPlan(
    struct ncclComm* comm, struct ncclKernelPlan* plan, struct ncclKernelPlanBudget* budget) {
  
  struct ncclKernelPlanner* planner = &comm->planner;
  
  // 从 collTaskQueue 和 collWorkQueue 取出任务
  while (nPlanColls!=0 && !ncclIntruQueueEmpty(&planner->collTaskQueue)) {
    struct ncclTaskColl* task = ncclIntruQueueHead(&planner->collTaskQueue);
    struct ncclWorkList* workNode = ncclIntruQueueHead(&planner->collWorkQueue);
    struct ncclDevWorkColl* devWork = (struct ncclDevWorkColl*)(workNode+1);
    
    // 分配通道
    devWork->channelLo = channelId;
    devWork->channelHi = channelId + nChannels-1;
    devWork->cbd.countLo = countLo;
    devWork->cbd.countMid = countMid;
    devWork->cbd.countHi = countHi;
    
    // 添加工作批次到计划
    addWorkBatchToPlan(comm, plan, c, workNode->workType, task->devFuncId, plan->workBytes);
    
    // 设置内核函数指针
    plan->kernelFn = ncclDevKernelForFunc[task->devFuncId];  // ← 设置要启动的内核
    
    // 从队列中移除任务，加入计划
    ncclIntruQueueDequeue(&planner->collTaskQueue);
    ncclIntruQueueDequeue(&planner->collWorkQueue);
    ncclIntruQueueEnqueue(&plan->collTaskQueue, task);
    ncclIntruQueueEnqueue(&plan->workQueue, workNode);
    
    plan->workBytes += workNode->size;
  }
  
  return ncclSuccess;
}

关键数据流转

任务在不同队列中的流转：

1. collSorter (在 collTaskAppend 中)

   • 任务按流量大小排序存储

2. collTaskQueue (在 ncclPrepareTasks 中)

   • 从 collSorter 取出，选择算法/协议后存入
   • 按 (collnet, nvls) 分类排序

3. collWorkQueue (在 ncclTasksRegAndEnqueue 中)

   • 为每个任务创建 ncclDevWorkColl 结构体
   • 与 collTaskQueue 一一对应

4. plan->workQueue (在 scheduleCollTasksToPlan 中)

   • 从 collTaskQueue 和 collWorkQueue 取出
   • 分配通道，加入具体的执行计划

完整调用链总结

ncclGroupEndInternal()
    ↓
groupLaunch()
    ↓
ncclPrepareTasksAndCollPreconnect()
    ↓
ncclPrepareTasks()  // 从 collSorter 取出任务，选择算法/协议，加入 collTaskQueue
    ↓
[返回到 groupLaunch]
    ↓
ncclTasksRegAndEnqueue()  // 为任务创建 ncclDevWorkColl，加入 collWorkQueue
    ↓
[返回到 groupLaunch]
    ↓
doLaunches()
    ↓
ncclLaunchPrepare()  // ← 这里是衔接点！
    ↓
scheduleCollTasksToPlan()  // 从 collTaskQueue/collWorkQueue 取出，调度到 plan
    ↓
finishPlan()  // 完成计划，准备内核参数
    ↓
[返回到 doLaunches]
    ↓
ncclLaunchKernel()  // 启动 CUDA 内核

关键设计要点

1. 多阶段处理

• 准备阶段 (ncclPrepareTasks): 选择算法和协议
• 注册阶段 (ncclTasksRegAndEnqueue): 创建设备端工作结构体
• 调度阶段 (scheduleCollTasksToPlan): 分配通道，创建执行计划
• 启动阶段 (ncclLaunchKernel): 启动CUDA内核

2. 循环创建多个计划

ncclLaunchPrepare 中的 do-while 循环会持续创建计划，直到所有任务都被调度：

do {
  struct ncclKernelPlan* plan = ...;
  scheduleCollTasksToPlan(comm, plan, &budget);  // 尽可能多地调度任务
  finishPlan(comm, plan);
  ncclIntruQueueEnqueue(&planner->planQueue, plan);
} while (planner->nTasksColl + planner->nTasksP2p != 0);

3. 预算控制

每个计划都有预算限制（ncclKernelPlanBudget）：

• inArgsBytes: 内核参数空间
• outArgsBytes: FIFO或持久化缓冲区空间

scheduleCollTasksToPlan 会根据预算决定每个计划包含多少任务。

代码位置总结

步骤	函数	位置	功能
1	ncclPrepareTasks	src/enqueue.cc:348	选择算法/协议，任务→collTaskQueue
2	ncclTasksRegAndEnqueue	src/enqueue.cc:285	创建ncclDevWorkColl→collWorkQueue
3	doLaunches	src/group.cc:259	启动内核的主循环
4	ncclLaunchPrepare	src/enqueue.cc:1417	衔接点：创建计划
5	scheduleCollTasksToPlan	src/enqueue.cc:519	调度任务到计划
6	finishPlan	src/enqueue.cc:182	完成计划，准备内核参数
7	ncclLaunchKernel	src/enqueue.cc:1565	启动CUDA内核

关键衔接点 是 ncclLaunchPrepare() 函数，它在第 1470 行调用 scheduleCollTasksToPlan，将 collTaskQueue 中的任务调度到具体的执行计划中。