SWPipeline API Reference

SWPipeline 通过两层依赖声明描述任务间的数据关系： SWPipeline uses two layers of dependency declarations to describe data relationships between tasks:

• intra_iter_deps：同一 batch（同一 iter）内的数据依赖。如果任务 B 消费任务 A 的输出，写 intra_iter_deps=[(B, A)]。 intra_iter_deps: Data dependency within the same batch (same iter). If task B consumes A's output, write intra_iter_deps=[(B, A)].
• inter_iter_deps：跨迭代数据依赖。任务 A(iter i) 依赖 B(iter i-1) 的输出，写 inter_iter_deps=[(A, B)]。仅支持 i-1。 inter_iter_deps: Cross-iteration data dependency. Task A(iter i) depends on B(iter i-1), write inter_iter_deps=[(A, B)]. Only i-1 is supported.

两者都纯粹描述数据依赖语义，与 stream、stage 无关。stream 和 stage 带来的约束在推论部分讨论。运行时，每条依赖通过两阶段协议执行： Both purely describe data dependency semantics, independent of stream or stage. Constraints arising from stream and stage are discussed in the Corollaries section. At runtime, each dependency is enforced via a two-phase protocol:

# Phase 1: CPU synchronization — threading.Event
sig = self._cpu_signal(iter_idx, dep_name)
sig.wait(timeout=30)                         # CPU blocks until dep's CPU work is done

# Phase 2: GPU synchronization — torch.cuda.Event
stream.wait_event(self._cuda_event(iter_idx, dep_name))  # GPU waits for dep's kernels

# Signal sender (after task completes):
self._cuda_event(iter_idx, task_def.name).record(stream)
self._cpu_signal(iter_idx, task_def.name).set()

多线程分发与 CPU 同步Multi-Thread Dispatch & CPU Synchronization

SWPipeline 支持多线程提交 kernel。每个任务通过 thread_group 属性指定所属的工作线程。运行时，_enqueue_period 按 _enqueue_order 顺序将任务分发到各线程的 work_queue： SWPipeline supports multi-threaded kernel submission. Each task specifies its worker thread via the thread_group attribute. At runtime, _enqueue_period dispatches tasks to each thread's work_queue following _enqueue_order:

_enqueue_period 分发流程_enqueue_period dispatch flow:
for task_def in self._enqueue_order:           # stage-descending
    self._work_queues[task_def.thread_group].put((iter_idx, task_def))
                      ↓                 ↓
            ┌─────────────────┐  ┌─────────────────┐
            │  Thread "memcpy" │  │  Thread "default" │
            │  work_queue:     │  │  work_queue:      │
            │   H2D            │  │   WaitBatch       │
            │                  │  │   Forward          │
            │                  │  │   Backward         │
            └────────┬─────────┘  └────────┬──────────┘
                     │                     │
              _submit_task()          _submit_task()

CPU 依赖不是独立的依赖层——它是两阶段协议的 Phase 1 自动处理的。当 Thread B 上的 WaitBatch 声明 intra_iter_deps=[H2D] 时： CPU dependency is not a separate layer — it is automatically handled by Phase 1 of the two-phase protocol. When WaitBatch on Thread B declares intra_iter_deps=[H2D]:

1. Thread B 在 _submit_task 开头执行 cpu_signal("H2D").wait() → CPU 阻塞，直到 Thread A 完成 H2D 并调用 .set() Thread B calls cpu_signal("H2D").wait() at the start of _submit_task → CPU blocks until Thread A finishes H2D and calls .set()
2. 解除阻塞后，Thread B 执行 stream.wait_event(cuda_event("H2D")) → GPU 也同步 Once unblocked, Thread B calls stream.wait_event(cuda_event("H2D")) → GPU also synchronized

即使两个任务在同一线程上，Phase 1 的 wait() 也一样执行——如果 dep 已完成则 wait() 立即返回，零开销。因此用户只需声明数据依赖（intra_iter_deps / inter_iter_deps），CPU 线程同步由 Phase 1 自动保证。 Even for tasks on the same thread, Phase 1's wait() still executes — if the dep is already done, wait() returns immediately with zero overhead. So users only need to declare data dependencies (intra_iter_deps / inter_iter_deps); CPU thread synchronization is automatically guaranteed by Phase 1.

示例：跨线程依赖Example: Cross-Thread Dependency

Thread "memcpy" (H2D)                    Thread "default" (WaitBatch)
═══════════════════════════════          ═══════════════════════════════════
  ├─ run H2D on memcpy_stream              ├─ cpu_signal("H2D").wait()  ← Phase 1: CPU blocks
  ├─ cuda_event("H2D").record(memcpy)      │   ... Thread B blocked ...
  ├─ cpu_signal("H2D").set()  ─────────────┤   unblocked!
  ▼ (done, process next in queue)          ├─ stream.wait_event(cuda_event("H2D"))  ← Phase 2: GPU waits
                                           ├─ run WaitBatch on default_stream
                                           ▼ (done)

iter_idx = period − stage
pipeline_depth = max(stage) + 1

示例Example: pipeline_depth=3 (SparseDist)
              stage=0(H2D)  stage=1(InputDist)  stage=2(Fwd/Bwd/Opt)
period=2:     iter 2        iter 1              iter 0
period=3:     iter 3        iter 2              iter 1
period=4:     iter 4        iter 3              iter 2

场景 1：单线程 k=1 升序死锁Scenario 1: Single-thread k=1 ascending deadlock

A(s=0, T1), B(s=1, T1), A.inter_iter_deps=[B]  →  同 periodsame period

升序Ascending ❌:  T1: A(s=0) → wait B → BLOCKED → B 永远不提交never submitted 💀

降序Descending ✅:  T1: B(s=1) → set ✓ → A(s=0) → 已 setset ✓

场景 2：多线程交叉死锁Scenario 2: Multi-thread cross deadlock

T1: A(s=0), D(s=1)  A.cross_iter_dep→B(T2)   T2: C(s=0), B(s=1)  C.cross_iter_dep→D(T1)

升序Ascending ❌:  T1: A→wait B→BLOCKED  T2: C→wait D→BLOCKED  → 互锁mutual 💀

降序Descending ✅:  T1: D(s=1)→✓  T2: B(s=1)→✓  T1: A(s=0)→✓  T2: C(s=0)→✓

场景 3：k ≥ 2 结构性死锁Scenario 3: k ≥ 2 structural deadlock

fwd(s=0), bwd(s=1), opt(s=2), fwd.inter_iter_deps=[opt]  →  k=2
           stage=2(opt)   stage=1(bwd)   stage=0(fwd)
period 3:   opt(1)        bwd(2)        fwd(3)
period 4:   opt(2)        bwd(3)        fwd(4)

fwd(3) 需要needs opt(2):  fwd(3) @ period 3,  opt(2) @ period 4  ← 未来 period，任何顺序都死锁future period, deadlock regardless 💀

基本设定：intra_iter_deps = [(bwd, fwd), (opt, bwd)]。分析有/无 fwd.inter_iter_deps = [opt]。 Base: intra_iter_deps = [(bwd, fwd), (opt, bwd)]. With/without fwd.inter_iter_deps = [opt].

无 inter_iter_deps（接受 weight staleness）Without inter_iter_deps (accept staleness)

有 inter_iter_deps → 无 GPU 加速With inter_iter_deps → no GPU speedup

推论 1 + 推论 3 → 仅 3 种合法安排（A/B/C），均因 wait_event 强制串行： Corollary 1 + 3 → only 3 legal arrangements (A/B/C), all forced serial by wait_event:

GPU stream 是 FIFO 队列。考虑以下配置： GPU streams are FIFO queues. Consider this setup:

任务Task   stage   stream   intra_iter_deps
────────────────────────────────────
P      0      X        —              无依赖none
Q      1      Y        [P]            跨 stage(0→1)，不在 period-local 图中cross-stage(0→1), not in period-local graph
R      0      X        [Q]            跨 stage(1→0)，不在 period-local 图中，但 GPU 仍有 wait_eventcross-stage(1→0), not in period-local graph, but GPU still has wait_event

P 和 R 都在 stage 0、stream X 上。stage-descending 先提交 Q(stage 1)，再提交 P、R(stage 0)。问题是 R 依赖 Q（跨 stream Y→X），提交时会插入 stream_X.wait_event(Q)。如果 R 被排在 P 前面（如按字母序），P 就被堵在 R 的 wait_event 后面白白空转： P and R are both at stage 0 on stream X. Stage-descending submits Q(stage 1) first, then P and R(stage 0). The problem: R depends on Q (cross-stream Y→X), inserting stream_X.wait_event(Q). If R is queued before P (e.g., alphabetical), P idles behind R's wait_event:

不优顺序（R 先于 P）Suboptimal (R before P):  stream_X: [R.wait_event(Q)] → [R kernels] → [P kernels]
                                                       ↑ P 白等idles

优化顺序（P 先于 R）Optimal (P before R):     stream_X: [P kernels] → [R.wait_event(Q)] → [R kernels]
                                         ↑ P 立即执行runs immediately

GPU stall 定义：stream.wait_event(event) 使 GPU stream 暂停直到 event 被 record。如果 event 来自同 period 内尚在执行的另一个 stream 上的任务，stream 会空转。这段空转时间就是 stall。 GPU stall defined: stream.wait_event(event) pauses the GPU stream until the event is recorded. If the event comes from another stream's task still executing within the same period, the stream idles. This idle time is the stall.

Stage-descending 的局限：它只按 stage 排序，同 stage 内的任务顺序不考虑 stall_cost。ready-first 算法解决这个问题——优先提交无跨 stream 依赖（stall_cost=0）的任务。 Stage-descending limitation: It only sorts by stage; within the same stage, task order ignores stall_cost. Ready-first fixes this — it prioritizes tasks with no cross-stream deps (stall_cost=0).

算法分两步：The algorithm has two steps:

步骤 1：构建 period-local 依赖图Step 1: Build the period-local dependency graph

只纳入同一 period 内实际发生交互的依赖边（参见推论 3 的分类）： Include only dependency edges that interact within the same period (see Corollary 3 classification):

不纳入的边：intra_iter_deps 中 stagedep < stagetask（dep 在更早 period 完成）；inter_iter_deps 中 stagedep ≤ stagetask（dep 在更早 period 完成）。这些依赖的 threading.Event 已 set、cuda_event 已 record，不会造成任何阻塞。 Excluded edges: intra_iter_deps where stagedep < stagetask (dep completed in earlier period); inter_iter_deps where stagedep ≤ stagetask (dep completed in earlier period). These dependencies have threading.Event set and cuda_event recorded — zero blocking.

步骤 2：Stall-Cost 感知的拓扑排序Step 2: Stall-cost-aware topological sort

对 period-local 图做拓扑排序。当多个任务同时 ready（入度 = 0）时，按 stall_cost 升序选取： Topologically sort the period-local graph. When multiple tasks are ready (in-degree = 0), pick by ascending stall_cost:

stall_cost(task) = 该任务在 period-local 图中，有多少个跨 stream 的入边（dep 与 task 不在同一 stream 上）number of cross-stream in-edges for this task in the period-local graph (dep and task on different streams)

排序 keySort key: (stall_cost, task_name)
  → stall_cost = 0 的任务优先：不插入 wait_event，不堵同 stream 后续任务tasks first: no wait_event inserted, won't block subsequent tasks on same stream
  → task_name 做 tie-breaker，保证确定性（跨 rank 一致）as tie-breaker for determinism (consistent across ranks)

def _build_enqueue_order_ready_first(self) -> List[PipelineTask]:
    # --- Step 1: build period-local graph ---
    adj:       Dict[str, List[str]] = defaultdict(list)   # dep → [tasks]
    in_degree: Dict[str, int]       = {t.name: 0 for t in self._defs.values()}

    for task_name, dep_list in self._intra_iter_deps.items():
        task_stage = self._stage_map[task_name]
        for dep_name in dep_list:
            if self._stage_map[dep_name] == task_stage:   # same stage → same period
                adj[dep_name].append(task_name)
                in_degree[task_name] += 1

    for task_name, dep_list in self._inter_iter_deps.items():
        task_stage = self._stage_map[task_name]
        for dep_name in dep_list:
            if self._stage_map[dep_name] == task_stage + 1:  # k=1 → same period
                adj[dep_name].append(task_name)
                in_degree[task_name] += 1

    # --- Step 2: compute stall_cost ---
    stall_cost: Dict[str, int] = {t.name: 0 for t in self._defs.values()}
    for task_name, dep_list in self._intra_iter_deps.items():
        task_stream = self._stream_map[task_name]
        task_stage  = self._stage_map[task_name]
        for dep_name in dep_list:
            if self._stage_map[dep_name] == task_stage \
               and self._stream_map[dep_name] != task_stream:
                stall_cost[task_name] += 1        # in-period cross-stream dep

    for task_name, dep_list in self._inter_iter_deps.items():
        task_stream = self._stream_map[task_name]
        task_stage  = self._stage_map[task_name]
        for dep_name in dep_list:
            if self._stage_map[dep_name] == task_stage + 1 \
               and self._stream_map[dep_name] != task_stream:
                stall_cost[task_name] += 1        # in-period cross-stream dep

    # --- Step 3: topo sort with stall_cost priority ---
    ready = sorted(
        [n for n, d in in_degree.items() if d == 0],
        key=lambda n: (stall_cost[n], n)
    )
    order: List[str] = []
    while ready:
        name = ready.pop(0)
        order.append(name)
        for succ in adj[name]:
            in_degree[succ] -= 1
            if in_degree[succ] == 0:
                ready.append(succ)
                ready.sort(key=lambda n: (stall_cost[n], n))

    return [self._defs[n] for n in order]

配置回顾（共 8 个任务，3 个 stream，3 个 stage）：Configuration recap (8 tasks, 3 streams, 3 stages):

任务Task            stage   stream        intra_iter_deps            inter_iter_deps
─────────────────────────────────────────────────────────────────────────
H2D              0      memcpy        —                             —
InputDistStart   1      data_dist     H2D (stage 0→1, cross-stage)   —
InputDistWait    1      data_dist     InputDistStart                —
ZeroGrad         2      default       —                             —
WaitBatch        2      default       InputDistWait (stage 1→2),    —
                                      ZeroGrad
Forward          2      default       InputDistWait (stage 1→2),    —
                                      WaitBatch
Backward         2      default       Forward                       —
OptimizerStep    2      default       Backward                      —

步骤 1：Period-Local 图Step 1: Period-Local Graph

只保留 stagedep = stagetask 的 intra_iter_deps 边（无 inter_iter_deps）：Keep only intra_iter_deps edges where stagedep = stagetask (no inter_iter_deps):

保留的边（同 stage = 同 period）Retained edges (same stage = same period):
  InputDistStart → InputDistWait        (stage 1 → 1) ✓
  ZeroGrad       → WaitBatch            (stage 2 → 2) ✓
  WaitBatch      → Forward              (stage 2 → 2) ✓
  Forward        → Backward             (stage 2 → 2) ✓
  Backward       → OptimizerStep        (stage 2 → 2) ✓

排除的边（跨 stage = 跨 period，dep 在更早 period 已完成）Excluded edges (cross-stage = cross-period, dep completed in earlier period):
  H2D → InputDistStart                  (stage 0 → 1) ✗
  InputDistWait → WaitBatch             (stage 1 → 2) ✗
  InputDistWait → Forward               (stage 1 → 2) ✗

步骤 2：计算 stall_costStep 2: Compute stall_cost

stall_cost = period-local 图中跨 stream 的入边数。所有保留的边都是同 stream 内部的：stall_cost = number of cross-stream in-edges in the period-local graph. All retained edges are within the same stream:

任务Task              period-local 入边period-local in-edges          跨 stream?cross-stream?   stall_cost
──────────────────────────────────────────────────────────────────────
H2D              无none                              —            0
InputDistStart   无none                              —            0
InputDistWait    InputDistStart(data_dist)   同 streamsame     0
ZeroGrad         无none                              —            0
WaitBatch        ZeroGrad(default)           同 streamsame     0
Forward          WaitBatch(default)          同 streamsame     0
Backward         Forward(default)            同 streamsame     0
OptimizerStep    Backward(default)           同 streamsame     0

所有 stall_cost = 0，因为跨 stage（跨 stream）的 dep 被排除了——它们来自更早 period，GPU event 早已 record。 All stall_cost = 0 because cross-stage (cross-stream) deps are excluded — they come from earlier periods with GPU events already recorded.

步骤 3：拓扑排序Step 3: Topological Sort

初始 ready（入度=0，按 (stall_cost, name) 排序）Initial ready (in-degree=0, sorted by (stall_cost, name)):
  [(0,"H2D"), (0,"InputDistStart"), (0,"ZeroGrad")]

出队过程Dequeue process:
  1. H2D           → 无后继在 period-local 图中no successors in period-local graph
  2. InputDistStart → 释放releases InputDistWait
  3. ZeroGrad       → 释放releases WaitBatch
  4. InputDistWait  → 无 period-local 后继no period-local successors
  5. WaitBatch      → 释放releases Forward
  6. Forward        → 释放releases Backward
  7. Backward       → 释放releases OptimizerStep
  8. OptimizerStep

最终提交顺序Final submission order:

 #  任务Task             stream       stage   stall_cost
 1. H2D              memcpy        0       0
 2. InputDistStart   data_dist     1       0
 3. ZeroGrad         default       2       0
 4. InputDistWait    data_dist     1       0
 5. WaitBatch        default       2       0
 6. Forward          default       2       0
 7. Backward         default       2       0
 8. OptimizerStep    default       2       0

考虑一个假设的 pipeline，同一 stream 上有任务分布在不同 stage，且存在 inter_iter_deps（k=1）：Consider a hypothetical pipeline where the same stream has tasks at different stages, with inter_iter_deps (k=1):

配置Setup:
  P: stage=0, stream=X               # stall_cost = 0
  Q: stage=1, stream=Y, intra_iter_deps=[P]     # 跨 stage：排除出 period-local 图cross-stage: excluded from period-local graph
  R: stage=0, stream=X, intra_iter_deps=[Q]     # 跨 stage：排除出 period-local 图cross-stage: excluded from period-local graph
  A: stage=0, stream=Z, inter_iter_deps=[Q]   # k=1: Q(stage 1) → 进入 period-local 图enters period-local graph

Period-Local 图Period-Local Graph（只有 k=1 的 inter_iter_deps 边） (only the k=1 inter_iter_deps edge):

  Q(stage 1) ──→ A(stage 0)       # cross_iter_dep, k=1, 同 periodsame period
  P, R: 无 period-local 边no period-local edges

stall_cost:

  P: 0    （无 period-local 入边）(no period-local in-edge)
  Q: 0    （无 period-local 入边）(no period-local in-edge)
  R: 0    （无 period-local 入边）(no period-local in-edge)
  A: 1    （Q→A 跨 stream: Y→Z）(Q→A cross-stream: Y→Z)

拓扑排序Topological sort:

初始 readyInitial ready: [(0,"P"), (0,"Q"), (0,"R")]   A 入度=1，不在 ready 中in-degree=1, not ready

  1. P  (stall=0, stream=X) → 立即执行，不等任何人runs immediately, no waits
  2. Q  (stall=0, stream=Y) → 立即执行runs immediately  → 释放 Areleases A
  3. R  (stall=0, stream=X) → 立即执行（P 已在 X 上，FIFO 保证顺序）runs immediately (P already on X, FIFO guarantees order)
  4. A  (stall=1, stream=Z) → wait_event(Q)，但 Q 已提交，延迟最小wait_event(Q), but Q already submitted, minimal delay

最终顺序Final order:  P(X,s0) → Q(Y,s1) → R(X,s0) → A(Z,s0)

GPU 时间线timeline:
  stream_X: ▓▓P▓▓▓▓R▓▓▓▓▓▓           ← P,R 连续执行，零空转run back-to-back, zero stall
  stream_Y: ▓▓Q▓▓▓▓▓▓▓▓▓▓▓
  stream_Z: ░░░░░░▓▓A▓▓▓▓▓           ← A 等 Q 完成后执行waits for Q then runs

Ready-first 不仅优化性能，还从结构上保证了无死锁。死锁的根本原因是循环等待——任务 X 等待任务 Y，而 Y 又直接或间接等待 X。Ready-first 通过以下两层保证消除循环等待： Ready-first is not merely a performance optimization — it structurally guarantees deadlock freedom. Deadlock arises from circular waits — task X waits for Y while Y (directly or indirectly) waits for X. Ready-first eliminates circular waits through two layers of guarantees:

多线程场景：每个 thread group 有独立的 FIFO 队列。如果任务 X 排在队头但阻塞等待同 queue 中更后面的 Y → head-of-line 死锁。Ready-first 的拓扑排序确保 Y 永远排在 X 前面，消除了这种死锁模式。 Multi-threaded scenario: Each thread group has its own FIFO queue. If task X is at the head but blocks waiting for Y further back in the same queue → head-of-line deadlock. Ready-first's topological sort ensures Y is always ahead of X, eliminating this deadlock pattern.

❌ 可能死锁的入队顺序（未按拓扑排序）❌ Potentially deadlocking enqueue order (not topologically sorted):
  thread_queue: [X, Y]      X 先出队，等 Y 的 CPU signal → Y 永远出不了队 → 死锁dequeued first, waits for Y's CPU signal → Y never dequeues → deadlock

✓ Ready-first 保证的入队顺序（拓扑排序）✓ Ready-first guaranteed enqueue order (topologically sorted):
  thread_queue: [Y, X]      Y 先出队执行 → set CPU signal → X 出队后立即获得 signaldequeued first → set CPU signal → X dequeues and immediately gets signal

形式化：period-local 依赖图是 DAG（无环有向图）——_validate 确保 intra_iter_deps 满足 dep.stage ≤ task.stage，inter_iter_deps 满足 dep.stage - task.stage ≤ 1（推论 1、3）。DAG 的拓扑排序一定存在且无环，因此 ready-first 产生的入队顺序必然无死锁。 Formal argument: The period-local dependency graph is a DAG — _validate ensures intra_iter_deps satisfy dep.stage ≤ task.stage and inter_iter_deps satisfy dep.stage - task.stage ≤ 1 (Corollaries 1 & 3). A DAG always admits a topological sort with no cycles, so the ready-first enqueue order is guaranteed deadlock-free.

Period-local 图排除了 intra_iter_deps 中 stagedep < stagetask 的边。这是否意味着这些同 batch 的依赖得不到保证？以下定理证明不会。 The period-local graph excludes intra_iter_deps edges where stagedep < stagetask. Does this mean these same-batch dependencies are unsatisfied? The following theorem proves they are not.

定理：对于任意 intra_iter_deps 边 (T, D)（T 依赖 D，同一迭代 i），ready-first 入队顺序保证 D 的 CPU signal 和 CUDA event 在 T 执行时已可用。 Theorem: For any intra_iter_deps edge (T, D) where T depends on D within iteration i, the ready-first enqueue order guarantees D's CPU signal and CUDA event are available when T executes.

证明：设 D 的 stage 为 s_D，T 的 stage 为 s_T。推论 1（_validate 强制）要求 s_D ≤ s_T。分三种情况： Proof: Let s_D and s_T be the stages of D and T. Corollary 1 (_validate enforced) requires s_D ≤ s_T. Three cases:

具体示例（FusedSparseDist）：InputDistStart(stage=1) 依赖 H2D(stage=0)，属于情况 B。 Concrete example (FusedSparseDist): InputDistStart(stage=1) depends on H2D(stage=0) — this is Case B.

迭代Iteration i = 3:
  H2D(i=3):            stage=0, period = 3+0 = 3
  InputDistStart(i=3): stage=1, period = 3+1 = 4

时间线Timeline:
  period 3: 提交enqueue H2D(i=3)  →  cpu_signal(3,"H2D").set()  →  cuda_event(3,"H2D").record()
  period 4: 提交enqueue InputDistStart(i=3)
            → cpu_signal(3,"H2D").wait()   →  已 set，立即返回already set, returns immediately
            → stream.wait_event(cuda_event) →  已 record，GPU 保证顺序already recorded, GPU guarantees order

QED：推论 1 排除情况 C；情况 A 由 period-local 拓扑排序保证；情况 B 由 period 顺序执行保证。三者合一，ready-first 对所有 intra_iter_deps 必然正确。对 inter_iter_deps 的证明类似（用推论 3 替换推论 1，证明结构相同）。 QED: Corollary 1 rules out Case C; Case A is guaranteed by period-local topological sort; Case B is guaranteed by sequential period processing. Combined, ready-first is correct for all intra_iter_deps. The proof for inter_iter_deps is analogous (substitute Corollary 3 for Corollary 1; proof structure is identical).