[triton] AMD: Warp Pipeline 지원 추가 - Gluon 프론트엔드부터 LLVM lowering까지

PR 링크: triton-lang/triton#8586 상태: Merged | 변경: +1185 / -1

들어가며

GPU 커널에서 compute와 memory 접근의 오버랩은 성능의 핵심입니다. 이번 PR은 AMD GPU에서 서로 다른 warp가 load, compute, store 같은 staggered 스테이지를 실행하는 warp-pipelined 루프를 구현합니다. Gluon 프론트엔드 API부터 LLVM IR lowering까지 전체 2단계 파이프라인을 제공합니다.

핵심 코드 분석

프론트엔드: Gluon API에서 스테이지 경계 표시

# warp_pipeline.py
class warp_pipeline_stage:
    """warp-pipeline 스테이지의 끝을 표시하는 context manager"""
    def __exit__(self, exc_type, exc, tb):
        if self._semantic is not None:
            attr = self.label if self.label else "cluster"
            self._semantic.builder.create_warp_pipeline_border(attr)
        return False

사용 예시:

for k in gl.range(0, K, one):
    with amd.warp_pipeline_stage("load"):
        a = gl.amd.buffer_load(a_ptr, offs_a)
    with amd.warp_pipeline_stage("compute"):
        acc = gl.amd.mfma(a_tile, b_tile, acc)

백엔드: 조건부 배리어로 lowering

IR에서 rocdl.sched.barrier에 triton.warp_pipeline.border 속성이 붙은 마커가 삽입되고, ConvertWarpPipeline 패스가 이를 scf.execute_region 클러스터로 변환한 뒤 amdgpu.cond_barrier로 guarding합니다.

// gluon_ir.cc - 마커 삽입
.def("create_warp_pipeline_border",
     [](GluonOpBuilder &self, const std::string &marker) {
       auto border = self.create<ROCDL::SchedBarrier>(0);
       auto ctx = self.getContext();
       border->setAttr("triton.warp_pipeline.border",
                       StringAttr::get(ctx, marker));
     });

왜 이게 좋은가

1. 사용자 제어 파이프라이닝: 개발자가 스테이지 경계를 명시적으로 지정하여 최적의 compute-memory 오버랩을 달성할 수 있습니다.
2. 2단계 lowering: 프론트엔드(Gluon->TritonGPU)와 백엔드(TritonGPU->LLVM)가 분리되어 각 단계를 독립적으로 최적화할 수 있습니다.
3. AMD 전용 최적화: rocdl.sched.barrier와 rocdl.s.barrier를 활용한 AMD 아키텍처 특화 최적화입니다.

정리

이 PR은 GPU 파이프라이닝의 전체 스택을 구현한 대규모 변경입니다. Gluon의 Python context manager라는 직관적인 API로 시작하여 LLVM IR의 하드웨어 배리어까지 이어지는 깔끔한 추상화 계층 설계가 돋보입니다.

참고 자료

• triton-lang/triton#8586

---

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