본문으로 건너뛰기

[Triton] AMD StreamPipeliner 리팩터링 - LowerLoops를 Pipeline 패스로 이동

들어가며

Triton의 AMD 백엔드에서 루프 파이프라이닝은 ScheduleLoops와 Pipeline 두 패스로 나뉜다. 기존에는 LDS(Local Data Share) 할당, local/async load 생성 등의 로직이 ScheduleLoops에 섞여 있어 두 패스의 책임이 불분명했다. 이 PR은 이런 lowering 로직을 Pipeline 패스로 이동시키는 대규모 리팩터링이다.

핵심 코드 분석

Before

# ScheduleLoops에 여러 책임이 혼재
amd.passes.ttgpuir.add_schedule_loops(pm, options.num_stages, use_async_copy, use_block_pingpong)
amd.passes.ttgpuir.add_pipeline(pm, use_async_copy)

ScheduleLoops가 num_stages, use_async_copy, use_block_pingpong 세 가지 파라미터를 받았다.

After

# 책임이 명확하게 분리됨
amd.passes.ttgpuir.add_schedule_loops(pm, options.num_stages)
amd.passes.ttgpuir.add_pipeline(pm, use_async_copy, use_block_pingpong)

ScheduleLoops는 num_stages만, Pipeline이 use_async_copy와 use_block_pingpong을 담당한다. 새로 생성된 LowerLoops.cpp(754줄)에 LDS 할당, async copy 생성, schedule 관리 로직이 모여 있다.

// 새 LowerLoops.cpp - async copy chain 생성
AsyncCopyChainOps createAsyncCopy(tt::LoadOp loadOp, Value alloc, Value extractIdx) {
    OpBuilder builder(loadOp);
    auto viewLoad = triton::createSingleBufferView(builder, alloc, extractIdx)
                        .getDefiningOp<ttg::MemDescIndexOp>();
    // ...
}

CoarseSchedule::deSerialize에도 normalizeClusterId 옵션이 추가되어 symbolic cluster와 실제 cluster 간 매핑이 가능해졌다.

왜 이게 좋은가

  • 단일 책임 원칙: ScheduleLoops는 스케줄링만, Pipeline은 코드 변환만 담당한다.
  • symbolic cluster: ScheduleLoops에서 생성한 symbolic cluster ID를 LowerLoops에서 실제 cluster로 매핑하는 명확한 인터페이스가 생겼다.
  • IR attribute 활용: bypassLDS 같은 정보를 IR attribute로 패스 간 전달하여 결합도를 낮췄다.

정리

+936/-718의 대규모 리팩터링이지만, 패스 경계를 재정의하여 각 패스의 테스트와 유지보수가 독립적으로 가능해졌다. StreamPipeliner 리팩터링 시리즈의 핵심 PR이다.

참고 자료

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석 내용은 실제 PR diff를 기반으로 합니다.

댓글