[triton] Triton 컴파일러의 Reduce 연산 최적화: OptimizeThreadLocality 개선 분석

PR 링크: triton-lang/triton#11370 상태: Merged | 변경: +679 / -212

들어가며

Triton 컴파일러에서 OptimizeThreadLocality 패스는 GPU 커널의 성능을 결정짓는 핵심적인 최적화 단계입니다. 특히 scf.for 루프 내에서 발생하는 reduce 연산을 효율적으로 처리하는 것은 메모리 접근과 스레드 간 통신을 최적화하는 데 매우 중요합니다. 하지만 기존 구현은 패턴 매칭 조건이 지나치게 엄격하고, 일부 버그가 존재하며, 확장성이 부족하다는 문제점이 있었습니다. 이번 PR은 패턴 매칭 로직을 순수 함수 형태로 재작성하고, 기존의 제약 사항들을 완화하여 더 넓은 범위의 커널에서 최적화가 적용될 수 있도록 개선했습니다.

코드 분석

패턴 매칭의 구조적 개선

기존에는 reduce 연산을 먼저 찾고 그에 의존하는 루프를 탐색하는 방식이었으나, 이번 변경에서는 scf.for 루프의 iter_arg를 추적하는 방식으로 변경되었습니다.

Before (기존 방식):

// reduce를 먼저 찾고, 그 사용처를 거슬러 올라가 forOp를 탐색
mod.walk([&](triton::ReduceOp reduce) -> void {
  // ... 복잡한 중첩 조건문들 ...
  auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
});

After (개선된 방식):

// 루프를 먼저 순회하고, 각 iter_arg에 대해 매칭 함수를 호출
struct AccumulatedReduce { triton::ReduceOp reduce; Operation *update; };
std::optional<AccumulatedReduce> matchAccumulatedReduce(scf::ForOp forOp, unsigned argIdx) {
  BlockArgument iterArg = forOp.getRegionIterArgs()[argIdx];
  // ... 추적 로직 ...
}

이러한 구조적 변화를 통해 여러 개의 독립적인 누적기(accumulator)가 하나의 루프 내에 존재하더라도 정확하게 매칭할 수 있게 되었습니다.

수익성 게이트(Profitability Gate) 완화

기존에는 elementPerThread == 1이라는 조건이 있었으나, 이는 사실 최적화의 의도와 반대되는 조건이었습니다. 이번 수정으로 스레드 레이아웃이 리덕션 축을 따라 분산되지 않는 경우를 정확히 판단하도록 변경되었습니다.

// After: 정확한 스레드 로컬리티 판단
bool isThreadLocal = getThreadsPerWarp(srcType)[axis] == 1 &&
                     getWarpsPerCTA(srcType)[axis] == 1 &&
                     getCTASplitNum(srcType.getEncoding())[axis] == 1;

왜 이게 좋은가

1. 범용성 확대: 기존에는 load 연산에서 정의된 reduce만 최적화가 가능했으나, 이제는 reduce가 패스 내에서 변환되지 않는 한 더 넓은 범위의 연산에 대해 최적화가 적용됩니다. 이는 특히 element-wise 커널에서 성능 향상을 기대할 수 있습니다.
2. 지원 연산 확장: arith::getNeutralElement를 활용하여 기존의 6개 부동소수점 연산 외에도 addi, muli, andi 등 모든 가환(commutative) 정수 연산에 대해 최적화가 가능해졌습니다.
3. 코드 유지보수성: 패턴 매칭 로직이 순수 함수로 분리되어 테스트가 용이해졌으며, 12개의 새로운 lit-test가 추가되어 향후 회귀 버그를 방지할 수 있는 견고한 테스트 셋을 확보했습니다.

리뷰 과정에서 논의된 바와 같이, one-use 제약을 완화하고 replaceAllUsesWith를 사용함으로써 컴파일러의 리라이터(rewriter)가 더 유연하게 동작하게 된 점은 매우 훌륭한 설계 개선입니다.

참고 자료

• https://mlir.llvm.org/docs/Dialects/SCF/
• https://github.com/triton-lang/triton

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.