[Triton] AMD PrepareIfCombining 패스 추가 — scf.if 병합 최적화

PR 링크: triton-lang/triton#9253 상태: Merged | 변경: +374 / -11

들어가며

MLIR의 canonicalizer는 동일 조건을 가진 인접 scf.if 연산을 하나로 병합할 수 있다. 하지만 두 scf.if 사이에 다른 연산이 끼어 있으면 병합이 불가능하다. 이 PR은 PrepareIfCombining 패스를 추가하여, 두 scf.if 사이의 연산을 첫 번째 scf.if 앞으로 이동시켜 canonicalizer가 병합할 수 있도록 준비한다.

핵심 코드 분석

Before (IR 구조)

%0 = ttg.local_load %smem
%1 = scf.if %cond -> tensor {
  %mul = arith.mulf %a, %a
  scf.yield %mul
} else { ... }
%2 = tt.trans %0 {order = [1, 0]}    // if 사이에 끼인 연산
%3 = scf.if %cond -> tensor {        // 같은 조건, 병합 불가
  %add = arith.addf %2, %2
  scf.yield %add
} else { ... }

After (PrepareIfCombining 적용 후)

%0 = ttg.local_load %smem
%2 = tt.trans %0 {order = [1, 0]}    // 첫 번째 if 앞으로 이동
%1 = scf.if %cond -> tensor {
  %mul = arith.mulf %a, %a
  scf.yield %mul
} else { ... }
%3 = scf.if %cond -> tensor {        // 이제 인접 → 병합 가능
  %add = arith.addf %2, %2
  scf.yield %add
} else { ... }

canonicalizer가 이후 두 scf.if를 하나로 병합한다.

왜 이게 좋은가

• 중복 분기 제거: 동일 조건의 scf.if가 하나로 병합되면 분기 명령어가 줄어들고, GPU에서 divergent branch 비용이 감소한다.
• 기존 패스 대체: ReorderInstructions 패스의 moveUpTranspose 최적화를 제거하고, 더 일반적인 접근법으로 대체한다.
• 안전성: 이동 가능 여부를 dependency 분석으로 검증하며, 중첩 if, 다른 블록의 if 등 음성 테스트 케이스도 포함한다.

정리

컴파일러 최적화에서 "패스 간 협력"은 중요한 패턴이다. 이 PR은 PrepareIfCombining이 코드를 재배치하고, canonicalizer가 실제 병합을 수행하는 2단계 협력 구조를 보여준다.

참고 자료

• PR #9253 전체 diff
• MLIR SCF Dialect

---

이 글은 AI 도구의 도움을 받아 작성되었습니다.