[triton] Triton NVIDIA 백엔드: full-warp boolean reduction을 warp vote로 lowering하기

PR 링크: triton-lang/triton#12635 상태: Merged | 변경: +30 / -0

들어가며

GPU 커널에서 tensor<...xi1> 같은 boolean 텐서를 축 방향으로 reduce하는 일은 생각보다 자주 일어납니다. "이 행에 조건을 만족하는 원소가 하나라도 있는가(any)", "모든 원소가 조건을 만족하는가(all)" 같은 질문이 대표적입니다. Triton에서는 이를 tt.reduce 안에 arith.ori나 arith.andi(또는 arith.minui/arith.maxui)를 넣어 표현합니다.

문제는 NVIDIA 백엔드가 이런 i1 reduction을 일반 정수 reduction과 똑같이 처리하고 있었다는 점입니다. PR 설명에 따르면 GB300에서 OR reduction 하나가 다음과 같은 SASS 시퀀스로 내려갔습니다.

SEL        // predicate -> 정수 레지스터로 변환
REDUX.OR   // warp 단위 정수 reduction
IMAD.U32   // 결과 정리
LOP3.LUT   // 다시 predicate 형태로 변환

즉, predicate 레지스터에 있던 1비트 값을 굳이 32비트 정수로 꺼내 REDUX로 줄인 뒤 다시 predicate로 되돌리는 왕복이 발생한 것입니다. 하지만 NVIDIA GPU에는 이 작업을 위해 존재하는 전용 명령이 있습니다. 바로 warp vote(vote.sync.any / vote.sync.all)입니다. 이 PR은 full-warp boolean reduction을 vote 명령 하나로 lowering합니다.

코드 분석

TargetInfo.cpp — warpReduce에 vote fast path 추가

핵심 변경은 TargetInfo::warpReduce 안의 단 12줄입니다.

Before

  Value value = acc[0];
  if (value.getType().isInteger(64)) {
    // The i64 tree uses two shuffles per stage for all groups, while REDUX
    // needs one instruction per piece and group. ...

After

  Value value = acc[0];
  if (value.getType().isInteger(1) && reduceLaneIdMask == 31) {
    bool all =
        *kind == NVVM::ReductionKind::AND || *kind == NVVM::ReductionKind::UMIN;
    bool any =
        *kind == NVVM::ReductionKind::OR || *kind == NVVM::ReductionKind::UMAX;
    if (all || any) {
      auto vote = all ? NVVM::VoteSyncKind::all : NVVM::VoteSyncKind::any;
      acc[0] = NVVM::VoteSyncOp::create(rewriter, loc, value.getType(),
                                        b.i32_val(-1), value, vote);
      return true;
    }
  }
  if (value.getType().isInteger(64)) {

하나씩 뜯어보면 다음과 같습니다.

1. 적용 조건: isInteger(1) && reduceLaneIdMask == 31 값이 i1이고, reduction이 warp 내 32개 lane 전체(lane id의 하위 5비트 모두)를 대상으로 할 때만 적용됩니다. vote 명령은 membermask에 포함된 모든 lane의 predicate를 한 번에 합치므로, 일부 lane 그룹끼리만 줄이는 부분 reduction에는 그대로 쓸 수 없습니다. 조건을 보수적으로 잡아 정확성을 지킨 셈입니다.

2. 연산 매핑: AND/UMIN → all, OR/UMAX → any 1비트 unsigned 값에서 min(a, b)는 a & b와 같고, max(a, b)는 a | b와 같습니다. 그래서 프런트엔드가 arith.minui로 표현한 reduction도 vote.all로 내릴 수 있습니다. 테스트에서 일부러 arith.minui를 쓴 이유가 이 동치성을 검증하기 위해서입니다. 반면 signed min/max는 i1에서 1이 -1로 해석되어 의미가 뒤집히므로 매핑에서 제외됐습니다. 이 판단도 정확합니다.

3. membermask -1 b.i32_val(-1)은 0xFFFFFFFF, 즉 32개 lane 전부를 뜻합니다. full-warp 조건과 정확히 맞물립니다.

4. 결과 타입이 i1 그대로 NVVM::VoteSyncOp의 결과 타입을 value.getType()(i1)으로 지정합니다. 입력과 출력이 모두 predicate 레지스터에 남기 때문에 SEL/LOP3 같은 변환 명령이 사라집니다. vote 결과는 warp의 모든 lane에 동일하게 broadcast되므로, 이후 별도의 shuffle 없이 곧바로 return true로 reduction을 끝낼 수 있습니다.

reduce_to_llvm.mlir — lit 테스트

// REDUX: %[[MASK:.*]] = llvm.mlir.constant(-1 : i32)
// REDUX: nvvm.vote.sync any %[[MASK]], {{.*}} -> i1
// REDUX: nvvm.vote.sync all %[[MASK]], {{.*}} -> i1

tensor<4x32xi1>를 axis 1(warp 방향 32개)로 reduce하는 두 함수를 만들고, arith.ori가 vote.sync any로, arith.minui가 vote.sync all로 내려가는지 확인합니다. mask가 -1 상수라는 점까지 검사해 full-warp 의미를 고정합니다.

왜 이게 좋은가

성능 수치

PR 작성자가 GB300에서 측정한 결과는 다음과 같습니다.

항목	Before	After
SASS (OR)	SEL + REDUX.OR + IMAD.U32 + LOP3.LUT	VOTE.ANY
SASS (AND)	유사한 4-명령 시퀀스	VOTE.ALL
OR reduction 64회 체인	4.887 µs	3.593 µs (약 26% 감소)

4개 명령이 1개로 줄었고, 의존성 체인으로 이어진 reduction에서 지연이 크게 줄었습니다. 마스크 계산이나 early-exit 판단처럼 boolean reduction이 루프 안에서 반복되는 커널이라면 체감 효과가 더 큽니다.

일반적 교훈

• 타입 폭에 맞는 명령을 고르자. 범용 정수 경로는 정확하지만 1비트 값에는 과합니다. 데이터가 원래 있는 레지스터 클래스(predicate)에서 벗어나지 않는 것이 가장 싼 경로입니다.
• 대수적 동치를 활용하자. i1에서 umin ≡ and, umax ≡ or라는 성질 덕분에 프런트엔드 표현이 달라도 같은 fast path를 탈 수 있습니다. 반대로 smin/smax처럼 동치가 성립하지 않는 경우는 명확히 배제해야 합니다.
• fast path는 좁게 시작하자. reduceLaneIdMask == 31 조건으로 full-warp만 다루는 것은 의도된 제약입니다. 부분 warp reduction은 vote.ballot과 비트 마스킹으로 확장할 수 있겠지만, 우선 가장 흔하고 확실한 경우부터 처리하는 편이 안전합니다.
• 하드웨어 전용 primitive를 먼저 찾아보자. warp vote는 CUDA 초창기부터 있던 명령입니다. 컴파일러가 이런 primitive를 놓치는 지점을 찾는 것만으로도 큰 이득을 얻을 수 있습니다.

작은 diff지만, 컴파일러 백엔드 최적화가 어디서 이득을 얻는지 잘 보여주는 PR입니다.

참고 자료

• CUDA C++ Programming Guide — Warp Vote Functions — vote.all/any의 의미와 membermask 설명
• PTX ISA — vote.sync — 이 PR이 lowering 대상으로 사용하는 PTX 명령
• MLIR NVVM Dialect — NVVM::VoteSyncOp 등 NVVM dialect 연산 문서
• Triton tl.reduce — tt.reduce에 대응하는 Triton 언어 API

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