kernbench2/docs/adr/ADR-0021-pe-pipeline-refactor.md

ADR-0021: PE 파이프라인 리팩토링 — 컴포넌트 분리 + Scheduler 기반 라우팅

Status

Proposed

Context

현재 구조의 문제

pe_accel (SchedulerV2Component)은 5개 하드웨어 블록(DmaIn, DmaWb, Gemm, Math, Tcm)을 단일 컴포넌트 내부에 숨기고 있다.

SchedulerV2Component (단일 topology 노드)
├── DmaInBlock     ← 내부 SimPy Store로 직접 연결
├── DmaWbBlock     ← topology에 안 보임
├── GemmBlock      ← 교체 불가
├── MathBlock      ← 교체 불가
└── TcmBlock       ← 교체 불가

문제점:

• 블록이 다음 블록을 desc.next_block으로 직접 참조 — 하드코딩된 라우팅
• 개별 블록 교체 불가 (ADR-0015 컴포넌트 교체 원칙 위배)
• topology에서 PE 내부 구조가 보이지 않음
• GemmBlock과 MathBlock이 TCM load/store 로직을 각각 중복 구현

실제 하드웨어 구조

HBM ←(DMA)→ TCM ←(Fetch/Store Unit)→ Register File ←→ GEMM/MATH Engine

• DMA: HBM ↔ TCM 전송 (fabric 경유, 수십~수백 ns)
• Fetch/Store Unit: TCM ↔ Register File 전송 (BW 기반, 수 ns)
• GEMM/MATH Engine: Register File 간 연산 (cycle-accurate)
• 완료 신호: PE 내부 1-cycle wire signal (done pin assert)

---

Decision

D1. 각 블록을 독립 컴포넌트로 분리

pe_accel의 내부 블록을 독립 PeEngineBase 컴포넌트로 분리한다. 기존 5개 + Fetch/Store Unit 1개 = 6개 컴포넌트.

컴포넌트	역할	HW 대응
PE_SCHEDULER	plan 생성, tile 상태 관리, stage 라우팅	Scheduler/Sequencer
PE_DMA	HBM ↔ TCM (fabric 경유)	DMA Engine
PE_FETCH_STORE	TCM ↔ Register File	Load/Store Unit
PE_GEMM	MAC compute (register only)	MAC Array
PE_MATH	element-wise/reduction (register only)	SIMD/Vector Unit
PE_TCM	BW-serialized scratchpad	SRAM Bank

각 컴포넌트는 topology 노드로 존재하며, port/wire로 연결된다. impl을 교체하면 개별 블록의 타이밍 모델을 변경할 수 있다.

D2. Token Self-Routing — Scheduler는 dispatch + completion만

컴포넌트가 매 stage마다 scheduler를 경유하지 않는다. Token이 plan을 가지고 있어 컴포넌트가 직접 다음 stage로 체이닝한다.

Scheduler → DMA → Fetch → GEMM → Math → Store → DMA_WB → (done) → Scheduler
              ↑ 체이닝: scheduler 안 거침                    completion만

이는 실제 HW에서 각 블록의 done signal이 다음 블록에 직접 wire로 연결되어 있는 구조와 일치한다. Scheduler는 초기 dispatch + completion aggregation만 담당.

Stage 정의

class StageType(Enum):
    DMA_READ = 0
    FETCH = 1
    GEMM = 2
    MATH = 3
    STORE = 4
    DMA_WRITE = 5

Plan 구조

Scheduler가 CompositeCmd를 받으면 tile 단위 실행 plan을 생성한다. Plan은 각 tile의 stage sequence를 정의한다:

@dataclass
class Stage:
    stage_type: StageType
    component: str       # topology 노드 ID (e.g. "sip0.cube0.pe0.pe_dma")
    params: dict         # stage별 파라미터 (dynamic)

@dataclass(frozen=True)
class TilePlan:
    tile_id: int
    stages: tuple[Stage, ...]  # 순서대로 실행할 stage 목록 (immutable)

Plan에 따라 stage sequence가 달라진다:

# 일반 GEMM: HBM → TCM → Register → Compute → Register → TCM → HBM
stages = (DMA_READ, FETCH, GEMM, STORE, DMA_WRITE)

# TCM 데이터로 바로 GEMM (DMA read 생략):
stages = (FETCH, GEMM, STORE, DMA_WRITE)

# MATH element-wise:
stages = (DMA_READ, FETCH, MATH, STORE, DMA_WRITE)

# GEMM + accumulation (중간 K-tile, writeback 생략):
stages = (DMA_READ, FETCH, GEMM, STORE)  # store to TCM only

컴포넌트는 다음 컴포넌트를 하드코딩하지 않는다. Token의 plan에서 다음 stage를 읽고, out_port로 직접 전달한다. 네트워크 패킷이 라우팅 헤더를 가지고 있는 것과 같은 패턴이다.

Pipeline Context

@dataclass
class PipelineContext:
    id: str
    total_tiles: int
    completed_tiles: int = 0
    done_event: simpy.Event = None  # 모든 tile 완료 시 succeed

    def complete_tile(self) -> None:
        self.completed_tiles += 1
        if self.completed_tiles == self.total_tiles:
            self.done_event.succeed()

Completion은 exactly-once contract: 각 tile의 마지막 stage는 정확히 한 번만 complete_tile()을 호출해야 한다. 중복 호출은 버그이며, done_event는 단 한 번만 succeed되어야 한다 (SimPy Event 제약).

Scheduler 역할 (축소됨)

Scheduler는 CompositeCmd를 받으면 plan과 PipelineContext를 생성한 뒤, 이를 scheduler 내부의 _pending_feeds FIFO에 enqueue하고 즉시 리턴한다.

실제 tile 투입은 단일 feeder process (_feed_loop)가 담당한다. 이 feeder는 _pending_feeds를 FIFO 순서로 소비하며, composite command 간 tile feed interleaving은 허용하지 않는다. 즉, 한 command의 모든 tile이 첫 stage queue에 투입된 후에만 다음 command의 feed가 시작된다.

Scheduler당 _feed_loop는 정확히 하나만 존재하며, composite command의 tile feed는 이 단일 process를 통해서만 수행된다. Command issue order는 PE_SCHEDULER가 PeInternalTxn을 수신한 순서를 의미한다.

이 구조는 command issue order를 유지하면서도, 첫 stage queue full 시 feeder process만 block되고 scheduler worker의 inbox 처리 자체는 멈추지 않도록 한다.

class PeSchedulerV2(PeEngineBase):
    _pipelines: dict[str, PipelineContext]
    _pending_feeds: simpy.Store   # FIFO of (plan, ctx)

    def start(self, env):
        super().start(env)
        self._pending_feeds = simpy.Store(env)
        env.process(self._feed_loop(env))

    def _dispatch_composite(self, env, pe_txn, cmd):
        plan = generate_plan(cmd)
        ctx = PipelineContext(
            id=next_id(),
            total_tiles=len(plan.tiles),
            done_event=pe_txn.done,
        )
        self._pipelines[ctx.id] = ctx

        # feeder queue에 등록만 하고 즉시 리턴
        yield self._pending_feeds.put((plan, ctx))

    def _feed_loop(self, env):
        """단일 feeder process: composite command를 FIFO 순서로 feed.

        Composite command 간 tile feed interleaving은 허용하지 않는다.
        한 command의 모든 tile이 첫 stage queue에 투입된 후에만
        다음 command의 feed가 시작된다.

        첫 stage queue full 시 이 feeder만 block되며,
        scheduler worker의 inbox 처리는 멈추지 않는다.
        """
        while True:
            plan, ctx = yield self._pending_feeds.get()
            for tile in plan.tiles:
                token = TileToken(
                    tile_id=tile.tile_id,
                    pipeline_ctx=ctx,
                    plan=tile,
                    stage_idx=0,
                    params=tile.stages[0].params,
                )
                yield self.out_ports[tile.stages[0].component].put(token)
                # queue capacity = HW queue depth → full이면 feeder만 block

본 ADR에서 scheduler는 여러 composite command를 수용할 수 있으나, tile submission order는 command 단위 FIFO를 따른다. Command 내부에서는 tile-level pipeline overlap을 허용하지만, command 간 tile feed interleaving은 허용하지 않는다.

D3. 데이터 전달 vs 완료 신호 — HW 모델링 기준

통신 유형	방식	HW 대응
tile token (작업 지시)	message via out_port	command queue에 enqueue
stage 완료 → 다음 stage	컴포넌트가 직접 out_port.put	done-triggered local enqueue
pipeline 완료 → scheduler	PipelineContext.complete_tile()	completion interrupt

Tile token: out_port.put() 사용. SimPy Store capacity = HW queue depth.

Intra-PE chaining latency: 본 ADR 범위에서는 intra-PE stage trigger에 explicit latency model을 두지 않는다. 컴포넌트 간 체이닝은 PE 내부 wire에 해당하며, scheduler 왕복이 없으므로 artificial hop cost가 발생하지 않는다.

Pipeline 완료: 마지막 stage의 컴포넌트가 pipeline_ctx.complete_tile() 호출. 모든 tile 완료 시 PipelineContext가 done_event.succeed().

D4. 비동기 파이프라인 — 자연스러운 overlap

Scheduler는 CompositeCmd를 비동기로 처리한다. 다만 tile feed는 command마다 독립 process를 만들지 않고, scheduler 내부의 단일 feeder process가 FIFO 순서로 수행한다. 따라서 scheduler는 다음 command를 계속 받을 수 있지만, 첫-stage tile 투입 순서는 command 단위로 보장된다.

SimPy Store capacity = HW queue depth이므로:

• queue가 차면 put()이 자연스럽게 block (backpressure)
• DMA가 tile 0을 처리하는 동안 GEMM은 이미 완료된 tile의 fetch를 시작
• 두 번째 CompositeCmd가 들어오면 DMA queue에 바로 이어서 투입

First-stage feed order (feeder → DMA queue):
  [cmd1:t0][cmd1:t1][cmd1:t2]...[cmd1:tN] | [cmd2:t0][cmd2:t1]...
                                            ↑ cmd1 feed 완료 후 cmd2 시작

Runtime pipeline (downstream overlap):
  PE_DMA:    [cmd1:t0][cmd1:t1][cmd1:t2]...[cmd1:tN][cmd2:t0][cmd2:t1]...
  PE_FETCH:          [cmd1:t0][cmd1:t1]...
  PE_GEMM:                   [cmd1:t0][cmd1:t1]...
                              ↑ 같은 cmd 내부에서 pipeline overlap

이때 overlap은 서로 다른 command의 tile feed interleaving에서 오는 것이 아니라, 먼저 투입된 command의 tile들이 downstream stage로 진행되는 동안 feeder가 다음 tile들을 계속 투입하면서 자연스럽게 발생한다.

예를 들어 cmd1의 모든 tile이 첫 stage queue에 투입되기 전에는 cmd2의 tile feed는 시작되지 않는다. 그러나 cmd1.tile0이 이미 GEMM으로 진행한 상태에서 cmd1.tile1, cmd1.tile2가 DMA/FETCH에 남아 있을 수 있으므로, 같은 command 내부에서는 pipeline overlap이 자연스럽게 발생한다.

컴포넌트 체이닝 패턴

모든 컴포넌트가 동일한 패턴을 따른다:

def _pipeline_worker(self, env):
    while True:
        token = yield self._inbox.get()

        # 자기 stage 처리
        yield from self._process(env, token)

        # 다음 stage로 체이닝 (plan에서 읽음)
        next_idx = token.stage_idx + 1
        if next_idx < len(token.plan.stages):
            next_stage = token.plan.stages[next_idx]
            token.stage_idx = next_idx
            token.params = next_stage.params
            yield self.out_ports[next_stage.component].put(token)
        else:
            # 마지막 stage — pipeline completion
            token.pipeline_ctx.complete_tile()

D5. PE_FETCH_STORE — TCM ↔ Register File 전담

기존에 GemmBlock과 MathBlock이 각각 TCM read/write를 구현했으나, 이를 PE_FETCH_STORE 컴포넌트로 분리한다.

# PE_FETCH_STORE._process()
def _process(self, env, token):
    yield self.out_ports[tcm_id].put(TcmRequest(token.params["direction"], ...))
    yield tcm_done
    # 체이닝은 base class가 처리 (D4 패턴)

장점:

• GEMM/MATH는 순수 compute만 — TCM 접근 로직 없음
• fetch/store BW 경합이 자연스럽게 모델링됨 (PE_TCM의 resource로 serialization)
• prefetch 전략 등 fetch unit 단독 교체로 실험 가능

D6. 각 Compute 컴포넌트의 단순화

GEMM/MATH는 register 데이터가 이미 준비된 상태에서 compute만 수행. 체이닝은 공통 패턴(D4)을 따르므로, _process()만 구현하면 된다:

# PE_GEMM._process()
def _process(self, env, token):
    yield env.timeout(self._mac_latency(token.params))

# PE_MATH._process()
def _process(self, env, token):
    yield env.timeout(self._simd_latency(token.params))

# PE_FETCH_STORE._process()
def _process(self, env, token):
    yield self.out_ports[tcm_id].put(TcmRequest(token.params["direction"], ...))
    yield tcm_done

# PE_DMA._process()
def _process(self, env, token):
    yield from self._do_fabric_dma(token.params)

타이밍 모델만 교체하면 cycle-accurate든 analytical든 자유롭게 변경 가능. 체이닝 로직은 base class에 있으므로 각 컴포넌트는 순수 stage 로직만 구현.

D7. Topology 변경

PE template에 PE_FETCH_STORE 추가:

pe_template:
  components:
    pe_cpu:         { kind: pe_cpu,         impl: pe_cpu_v1, ... }
    pe_scheduler:   { kind: pe_scheduler,   impl: pe_scheduler_v2, ... }
    pe_dma:         { kind: pe_dma,         impl: pe_dma_v1, ... }
    pe_fetch_store: { kind: pe_fetch_store, impl: pe_fetch_store_v1, ... }
    pe_gemm:        { kind: pe_gemm,        impl: pe_gemm_v1, ... }
    pe_math:        { kind: pe_math,        impl: pe_math_v1, ... }
    pe_mmu:         { kind: pe_mmu,         impl: pe_mmu_v1, ... }
    pe_tcm:         { kind: pe_tcm,         impl: pe_tcm_v1, ... }
  links:
    # 기존 links...
    fetch_store_to_tcm_bw_gbs: 512.0
    fetch_store_to_tcm_mm: 0.0

PE 내부 edge 연결:

PE_SCHEDULER → PE_DMA (초기 dispatch)
PE_SCHEDULER → PE_FETCH_STORE (초기 dispatch)
PE_SCHEDULER → PE_GEMM (초기 dispatch)
PE_SCHEDULER → PE_MATH (초기 dispatch)
PE_DMA → PE_FETCH_STORE (체이닝)
PE_FETCH_STORE → PE_GEMM (체이닝)
PE_FETCH_STORE → PE_MATH (체이닝)
PE_GEMM → PE_FETCH_STORE (store 체이닝)
PE_MATH → PE_FETCH_STORE (store 체이닝)
PE_FETCH_STORE → PE_DMA (writeback 체이닝)
PE_FETCH_STORE → PE_TCM (BW 요청)

Topology edge는 control/dispatch visibility + runtime chaining 양쪽을 포함한다. Scheduler → 하위 컴포넌트 edge는 초기 dispatch 경로이며, 컴포넌트 간 edge는 token self-routing에 의한 runtime chaining 경로이다.

D8. 기존 코드 마이그레이션 — builtin 통합

기존 builtin v1 컴포넌트와 pe_accel을 새 builtin으로 교체한다.

마이그레이션 전략

1. 기존 components/builtin/ → components/builtin_legacy/로 백업 (수정 없이 보관)
2. 기존 components/custom/pe_accel/ → 동일하게 백업
3. 새 components/builtin/에 ADR-0021 아키텍처로 재구현
4. topology.yaml은 하나만 유지 (pe_fetch_store 포함)
5. components.yaml은 새 builtin을 가리킴

# components.yaml — 새 builtin
pe_scheduler_v1: kernbench.components.builtin.pe_scheduler:PeSchedulerComponent
pe_gemm_v1:      kernbench.components.builtin.pe_gemm:PeGemmComponent
pe_math_v1:      kernbench.components.builtin.pe_math:PeMathComponent
pe_dma_v1:       kernbench.components.builtin.pe_dma:PeDmaComponent
pe_fetch_store_v1: kernbench.components.builtin.pe_fetch_store:PeFetchStoreComponent
pe_tcm_v1:       kernbench.components.builtin.pe_tcm:PeTcmComponent

impl 이름(pe_gemm_v1 등)은 유지하되, 구현이 ADR-0021 아키텍처로 교체된다. 기존 벤치마크와 테스트의 topology.yaml 참조는 변경 없이 동작한다.

레이턴시 모델 계승

새 builtin 컴포넌트의 레이턴시 모델링(MAC cycle 계산, SIMD latency, TCM BW serialization, DMA fabric latency 등)은 pe_accel 현재 버전의 구현을 바탕으로 한다. tiling.py의 tile schedule 생성 로직도 그대로 가져온다. 아키텍처(컴포넌트 분리, self-routing)만 변경하고, 타이밍 정확도는 유지한다.

테스트 전략

테스트 계획

1. 기존 테스트 통과 (regression): 마이그레이션 완료 후 기존 테스트(366개)가 전부 통과해야 한다.

2. 레이턴시 regression: pe_accel과 동일한 입력에 대해 새 builtin이 동일 레이턴시를 산출하는지 검증.

3. Phase 1 → Phase 2 end-to-end: SimPy 시뮬레이션(Phase 1)에서 op_log 생성 → DataExecutor(Phase 2)로 실제 numpy 연산 → 결과 정합성 검증까지 통합 테스트.

• GEMM: tl.composite(gemm) → op_log → Phase 2 matmul → allclose 검증
• MATH: tl.exp / tl.add 등 → op_log → Phase 2 numpy op → allclose 검증
• 체이닝: GEMM 출력 → MATH 입력 → 최종 결과 end-to-end 검증

4. TileToken self-routing:

• tile이 plan의 stage sequence를 따라 체이닝되는지 검증
• 마지막 stage에서 PipelineContext.complete_tile() exactly-once 검증
• queue backpressure: DMA queue capacity 초과 시 feeder만 block 검증

5. 비동기 pipeline overlap:

• 동일 command 내 tile 간 stage overlap 발생 검증 (tile0 GEMM 중 tile1 DMA)
• 다중 command: cmd1 feed 완료 후 cmd2 feed 시작 (FIFO 순서) 검증

D9. TileToken 메시지 정의

컴포넌트 간 tile 작업 전달에 사용하는 메시지. Token이 plan과 stage index를 가지고 있어 self-routing이 가능하다.

@dataclass
class TileToken:
    tile_id: int
    pipeline_ctx: PipelineContext    # completion 추적
    plan: TilePlan                   # 이 tile의 전체 stage sequence (immutable)
    stage_idx: int                   # 현재 stage index in plan.stages
    params: dict                     # current stage 파라미터 캐시 (canonical: plan.stages[stage_idx].params)
    data_op: bool = True             # op_log 기록 대상 (ADR-0020)

TileToken은 한 시점에 하나의 컴포넌트에 의해서만 소유되며, 동시에 여러 컴포넌트에 의해 참조되지 않는다 (single-owner).

Token lifecycle:

1. Scheduler가 stage_idx=0으로 생성, 첫 stage 컴포넌트에 put
2. 컴포넌트가 _process() 실행 후 stage_idx 증가, 다음 컴포넌트에 put
3. 마지막 stage 컴포넌트가 pipeline_ctx.complete_tile() 호출
4. 모든 tile 완료 시 PipelineContext가 done_event.succeed()

기존 PeInternalTxn과의 관계:

• PeInternalTxn: PE_CPU → PE_SCHEDULER 간 command 전달 (기존 유지)
• TileToken: PE_SCHEDULER → 하위 컴포넌트 간 tile 단위 작업 전달 (신규, self-routing)

---

Non-goals

• PE_CPU 변경: PE_CPU → PE_SCHEDULER 인터페이스는 변경하지 않음 (PeInternalTxn 기반, ADR-0014 유지)
• 다중 pipeline 간 자원 경합 모델: 현재 범위에서는 단일 pipeline의 정확한 모델링에 집중. 다중 pipeline 간 TCM bank conflict 등은 future work.
• builtin_legacy 유지보수: 백업 목적이며, 버그 수정이나 기능 추가 대상이 아님.

Open Questions

• Register File 용량 모델: fetch unit이 register에 로드할 때 용량 제한을 모델링할지. 용량은 바이트 단위(register_file_bytes)로 표현하며, 동시에 보유 가능한 tile 수는 tile 크기에 따라 결정된다. 용량 초과 시 fetch가 stall되어 자연스러운 backpressure가 발생한다.
• Prefetch 전략: 본 ADR에서는 composite command 간 tile feed interleaving을 허용하지 않는다. 따라서 overlap은 command 간 선행 투입이 아니라, 같은 command 내부 tile들의 pipeline progression에서 자연스럽게 발생한다. 추가적인 prefetch가 필요하면 command 간 투입이 아니라, 같은 command 내부에서의 tile ordering 또는 fetch/store unit policy 차원에서 검토한다.
• PE_DMA coalescing: tile 단위 DMA는 fragmentation 발생 가능. DMA 내부에서 merge/coalesce하되 scheduler는 관여하지 않는 방향.
• 동기 실행 모드: 본 ADR에서는 비동기 pipeline을 기본/유일 execution model로 채택한다. 디버그 또는 validation 목적의 sync mode가 필요하면 future ADR에서 검토.
• 다중 pipeline 간 TCM bank conflict: 현재 단일 pipeline 기준. 다중 pipeline이 동시에 TCM에 접근할 때의 bank conflict 모델은 future work.

---

Consequences

긍정적

• 각 블록이 독립 컴포넌트 — 개별 교체 가능 (ADR-0015 준수)
• topology에서 PE 내부 구조 가시화
• 컴포넌트가 다음 컴포넌트를 모름 — plan 기반 라우팅으로 유연성 확보
• DMA와 compute의 자연스러운 파이프라인 overlap (SimPy Store backpressure)
• HW 모델링 정확도 향상 (done signal = Event, data transfer = message)
• fetch/store 분리로 TCM BW 경합 정확히 모델링

부정적

• PE 내부 컴포넌트 수 증가 (5 → 6) — topology 노드/edge 증가
• 컴포넌트 분리로 인해 intra-PE token forwarding이 이전 대비 더 명시적으로 드러남
• 기존 builtin/pe_accel과의 breaking change — 마이그레이션 필요

---

영향받는 파일

파일	변경
topology.yaml	pe_fetch_store 컴포넌트 추가, 체이닝 edge 추가
components.yaml	새 builtin 컴포넌트 등록
src/kernbench/topology/builder.py	PE 내부 edge에 fetch_store + 체이닝 edge 추가
src/kernbench/common/pe_commands.py	TileToken 정의 추가
src/kernbench/components/builtin/pe_scheduler.py	재구현 (feeder + plan 기반 dispatch)
src/kernbench/components/builtin/pe_gemm.py	재구현 (TileToken, _process 패턴)
src/kernbench/components/builtin/pe_math.py	재구현 (TileToken, _process 패턴)
src/kernbench/components/builtin/pe_dma.py	재구현 (TileToken, _process 패턴)
src/kernbench/components/builtin/pe_fetch_store.py	신규
src/kernbench/components/builtin/pe_tcm.py	재구현 (TcmRequest 서비스)
src/kernbench/components/builtin/types.py	신규: TilePlan, Stage, StageType, PipelineContext, TileToken
src/kernbench/components/builtin/tiling.py	pe_accel에서 이식: plan 생성 로직

백업: | src/kernbench/components/builtin_legacy/ | 기존 builtin 전체 백업 (수정 없이 보관) | | src/kernbench/components/custom/pe_accel/ | 기존 pe_accel 백업 (수정 없이 보관) |