[Triton] SWP 루프 로우어링에서 barrier 위치 결정 로직 수정

PR 링크: triton-lang/triton#8732 상태: Merged | 변경: +283 / -22

들어가며

Software Pipelining(SWP)에서 MMA 연산이 "완료"되었음을 표시하는 barrier의 위치는, tmem_load 또는 non-pipelined operand 중 더 늦게 등장하는 것을 기준으로 결정된다. 기존에는 schedule.isOpBefore를 사용했는데, 이는 operand의 스케줄 순서를 비교할 뿐 실제 루프 body에서의 실행 순서를 반영하지 못했다.

핵심 코드 분석

Before

std::optional<Operation *> latestSyncPoint;
for (auto user : alloc->getUsers()) {
  if (auto load = dyn_cast<ttng::TMEMLoadOp>(user)) {
    if (!latestSyncPoint || schedule.isOpBefore(load, *latestSyncPoint)) {
      latestSyncPoint = load;
    }
  }
}

isOpBefore는 스케줄 단계(stage/cluster)만 비교하여, MMA 이전에 스케줄된 operand가 MMA 이후에 실행되는 경우를 잘못 판단할 수 있다.

After

llvm::SmallDenseSet<Operation *> syncCandidates;
for (auto user : alloc->getUsers()) {
  if (auto load = dyn_cast<ttng::TMEMLoadOp>(user)) {
    syncCandidates.insert(load);
  }
}
// ... non-pipelined operand defs도 syncCandidates에 추가

// MMA 이후 linearized schedule에서 첫 번째 sync candidate를 찾음
auto linearizedSchedule = schedule.linearized(forOp, mma);
std::optional<Operation *> latestSyncPoint =
    linearizedSchedule.findNext(
        [&](Operation *op) { return syncCandidates.contains(op); });

LinearizedIterator (Schedule.h)

class LinearizedIterator {
  // MMA에서 시작하여, stage/cluster/IR 순서로 순회
  // stage 경계를 넘을 때 stage limit 증가
  // 원형 순회로 wrap-around 처리
  std::optional<Operation *>
  findNext(std::function<bool(Operation *)> predicate);
};

왜 이게 좋은가

• 정확한 실행 순서: linearized schedule은 (stage, cluster, IR 순서)의 3단계 정렬을 반영하여, 실제 실행 순서에 기반한 판단을 제공한다.
• 재사용 가능한 API: LinearizedIterator는 forOp의 스케줄을 순회하는 범용 도구로, 다른 패스에서도 활용 가능하다.
• MMA 이후 기준: barrier는 반드시 MMA 이후에 등장하는 sync point를 기준으로 배치되어야 한다는 불변 조건을 명시적으로 구현한다.

정리

Software Pipelining에서 동기화 지점의 정확한 배치는 정확성과 성능 모두에 영향을 미친다. 이 PR은 linearized schedule iterator라는 범용 도구를 도입하여, 스케줄 순서와 실행 순서의 불일치 문제를 근본적으로 해결한다.

참고 자료

• PR #8732 전체 diff
• Triton Software Pipelining 설계

---

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