[triton] Warp Specialization: OptimizePartitionWarps와 SWP 순서 교환으로 어노테이션 보존

PR 링크: triton-lang/triton#8415 상태: Merged | 변경: +29 / -14

들어가며

Warp Specialization에서 OptimizePartitionWarps 패스는 각 파티션의 warp 수를 최적화합니다. 그런데 이 패스 내부에서 실행되는 RemoveLayoutConversions가 local_load 연산의 loop.cluster와 loop.stage 어노테이션을 삭제하는 경우가 있었습니다. 이 어노테이션은 Software Warp Pipelining(SWP)에 필수적이어서, 삭제되면 segfault가 발생할 수 있습니다.

핵심 코드 분석

Before: OptimizePartitionWarps가 SWP 이전에 실행

// AutomaticWarpSpecialization.cpp
pm.addPass(createNVWSLowerWarpGroup());
pm.addPass(createTritonGPUOptimizePartitionWarps()); // SWP 이전
pm.addPass(createTritonGPUScheduleLoops());           // SWP

After: SWP 이후로 이동하고, 빈 모듈 조기 반환 추가

// AutomaticWarpSpecialization.cpp
pm.addPass(createNVWSLowerWarpGroup());
pm.addPass(createTritonGPUScheduleLoops());           // SWP 먼저
// OptimizePartitionWarps는 나중에 별도로 실행

// OptimizePartitionWarps.cpp - 빈 모듈 조기 반환
void OptimizePartitionWarps::runOnOperation() {
  SmallVector<WarpSpecializeOp> wsOps;
  getOperation().walk([&](WarpSpecializeOp wsOp) { wsOps.push_back(wsOp); });
  if (wsOps.empty()) {
    return;  // WS 없으면 조기 반환
  }
  // ...
  for (auto wsOp : wsOps) {
    if (failed(optimizePartitionNumWarps(axisInfo, wsOp, runPipelineFn))) {
      return signalPassFailure();
    }
  }
}

왜 이게 좋은가

1. 어노테이션 보존: SWP가 먼저 실행되어 필요한 어노테이션이 이미 소비된 후 OptimizePartitionWarps가 실행되므로 삭제 문제가 없습니다.
2. 조기 반환 최적화: WarpSpecializeOp이 없는 모듈에서는 불필요한 분석을 건너뜁니다.
3. nested-loop 지원 기반: 단순한 경우에는 두 번째 ScheduleLoops가 어노테이션을 복원하지만, nested-loop에서는 복원이 실패하여 이 순서 변경이 필수적입니다.

정리

컴파일러 패스 순서는 미묘하지만 중요합니다. 이 PR은 패스 간 의존성을 명확히 하고, 한 패스가 다른 패스에 필요한 정보를 파괴하지 않도록 순서를 조정한 좋은 사례입니다.

참고 자료

• triton-lang/triton#8415

---

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