[triton] Triton NVIDIA 백엔드: Boolean reduction을 세 개의 predicate chain으로 접어 74% 빨라진 이유

PR 링크: triton-lang/triton#12650 상태: Merged | 변경: +114 / -5

들어가며

Triton에서 tl.reduce로 any(x == key)나 all(mask) 같은 Boolean reduction을 하면, 각 lane은 먼저 자기 register에 있는 값들을 thread 내부에서 접고(intra-thread reduction), 그 다음 warp shuffle이나 VOTE로 lane 간 결과를 합칩니다. 이 PR은 앞단계, 즉 lane 하나가 수십~수백 개의 i1 값을 접는 부분을 다룹니다.

기존 구현은 register 값을 모두 vals에 넣은 뒤 treeReduce로 이진 트리 reduction을 수행했습니다. 일반적인 정수/실수 reduction에서는 트리가 ILP(instruction-level parallelism)를 높여 좋은 선택이지만, i1에서는 문제가 생깁니다. GPU의 predicate register(P0~P6)는 개수가 매우 적어서, 트리의 중간 결과를 모두 predicate로 들고 있을 수 없습니다. 결국 ptxas는 중간 결과를 일반 register의 bit에 packing했다가 나중에 다시 꺼내는 코드를 만듭니다.

ISETP.NE.U32.AND P5, PT, R5, R0, PT;
ISETP.EQ.U32.OR  P5, PT, R4, R0.reuse, !P5;
@P5 LOP3.LUT R2, R2, 0x1000000, RZ, 0xfc, !PT;       // set bit 24
// ...
LOP3.LUT P4, RZ, R2.reuse, 0x1000000, RZ, 0xc0, !PT; // extract bit 24

비교 한 번에 bit set, bit extract가 따라붙는 셈입니다. 이 PR은 reduction 모양을 "세 개의 선형 chain"으로 바꿔 이 낭비를 없앱니다.

코드 분석

1. TargetInfoBase.h — target hook 추가

  virtual unsigned getReductionTreeArity(Operation *combinerOp) const {
    return 2;
  }
+  // Whether to use three interleaved register chains for a Boolean reduction.
+  virtual bool preferThreePredicateChains() const { return false; }

기본값은 false이므로 AMD 등 다른 backend의 동작은 바뀌지 않습니다. 이 최적화는 ptxas의 instruction selection 특성(ISETP ... OR, PLOP3)에 의존하기 때문에 target별 opt-in으로 두는 것이 맞습니다.

2. NVIDIA TargetInfo.h — opt-in

+  bool preferThreePredicateChains() const override {
+    // ptxas can fuse comparisons into each chain and merge the three
+    // partial predicates with one PLOP3.
+    return true;
+  }

주석이 설계 의도를 그대로 말해줍니다. 숫자 3은 임의가 아닙니다. PLOP3.LUT는 predicate 3개를 입력으로 받아 임의의 논리 함수를 계산하는 명령이므로, 세 chain의 partial 결과는 명령어 하나로 합쳐집니다.

3. ReduceOpToLLVM.cpp — 핵심 변경

먼저 이 변환을 적용해도 안전한지 판별합니다.

+    // Extra operations in the combiner region may have observable effects.
+    bool isPureBooleanCombiner =
+        isa_and_nonnull<arith::AndIOp, arith::OrIOp, arith::XOrIOp,
+                        arith::MinSIOp, arith::MaxSIOp, arith::MinUIOp,
+                        arith::MaxUIOp>(combinerOp) &&
+        combinerOp->getResult(0).getType().isInteger(1) &&
+        llvm::hasSingleElement(combineRegion.front().without_terminator());
+    bool usePredicateChains =
+        isPureBooleanCombiner && targetInfo.preferThreePredicateChains();

세 가지 조건이 있습니다.

• 결합·교환 법칙을 만족하는 연산: i1에서 and/umin/smax는 all, or/umax/smin은 any, xor는 parity입니다. 모두 순서를 바꿔도 결과가 같으므로 트리 대신 interleaved chain으로 재배열해도 됩니다.
• 결과 타입이 i1: predicate register에 올라가는 경우만 대상입니다.
• combiner region에 연산이 하나뿐: tt.print 같은 side effect가 섞여 있으면 combine 호출 순서가 관찰 가능해집니다. 이를 테스트 @reduce_register_any_with_effect가 검증하는데, vprintf가 있으면 기존처럼 X0 or X1 트리가 유지됩니다.

실제 재배열은 아주 작은 diff입니다.

-        vals.push_back(std::move(cur));
+        if (usePredicateChains && i >= 3)
+          accumulate(loc, rewriter, combineRegion, vals[i % 3], cur);
+        else
+          vals.push_back(std::move(cur));

처음 세 값은 vals[0..2]에 들어가 chain의 seed가 되고, 이후 i번째 값은 vals[i % 3]에 순차적으로 누적됩니다. 즉 chain 0은 register 0, 3, 6, 9…, chain 1은 1, 4, 7…을 담당합니다. 마지막엔 원소가 3개뿐인 vals가 그대로 treeReduce에 넘어가므로 기존 lane 간 reduction 로직은 손대지 않습니다.

4. 테스트

// Each lane owns 32 flags. The first chain combines registers 0, 3 and 6.
// PRED-CHAINS: %[[ALL:.*]] = llvm.intr.umin(%[[X0]], %[[X3]])
// PRED-CHAINS: llvm.intr.umin(%[[ALL]], %[[X6]])

7개의 combiner 종류를 모두 확인하고, 16/512개 크기에서도 X0과 X3이 결합되는 것을 검증합니다. 또 compute-capability=80/90/100/103/120 모든 RUN line에 PRED-CHAINS prefix를 붙여 NVIDIA 세대 전반에서 동작함을 보장합니다.

왜 이게 좋은가

변경 후 SASS는 다음과 같습니다.

ISETP.EQ.U32.OR P0, PT, R249, R0.reuse, P0;
ISETP.EQ.U32.OR P1, PT, R251, R0.reuse, P1;
ISETP.EQ.U32.OR P2, PT, R250, R0,       P2;
PLOP3.LUT       P0, PT, P2, P0, P1, 0xfe, 0x0;       // P2 | P0 | P1

ISETP는 비교 결과를 다른 predicate와 AND/OR로 결합하는 기능을 내장하고 있습니다. chain 형태 acc = acc | (x == key)는 이 패턴에 정확히 맞아 비교 + 누적이 명령어 1개가 됩니다. 살아있는 predicate는 chain 3개뿐이라 spill도 없습니다.

성능 수치: GB300에서 lane당 256개 값을 반복해서 fold하는 microbenchmark가 20.27 µs → 5.20 µs (−74.4%). 두 버전 모두 iteration당 VOTE.ANY는 한 번이므로, 차이는 순수하게 thread 내부 reduction에서 나온 것입니다.

그렇다면 chain을 1개가 아니라 3개로 둔 이유는? chain 하나면 모든 ISETP가 직전 결과에 의존하는 직렬 dependency가 되어 latency가 그대로 노출됩니다. 3개를 interleave하면 독립적인 chain이 번갈아 issue되어 pipeline을 채우고, 마지막 merge도 PLOP3 하나로 끝납니다. ILP와 register pressure 사이의 균형점입니다.

일반적 교훈

1. 트리 reduction이 항상 최선은 아니다. 중간값을 담을 저장소가 희소하면(predicate register처럼), 넓은 트리보다 소수의 선형 chain이 낫습니다.
2. 하위 컴파일러의 fusion 패턴에 맞춰 IR 모양을 만들어라. Triton은 SASS를 직접 생성하지 않지만, ptxas가 ISETP.OR와 PLOP3를 잘 쓰도록 LLVM IR의 dataflow를 조정한 것이 핵심입니다.
3. 재결합은 의미 보존 조건을 엄격히. 결합·교환 가능한 단일 연산이면서 side effect가 없을 때만 순서를 바꾸고, 이를 테스트로 고정한 점이 모범적입니다.

참고 자료

• Triton tl.reduce — 이 PR이 lowering을 개선한 reduction 연산
• MLIR arith dialect — combiner로 판별하는 andi/ori/xori/minui/maxui 등 연산 정의
• CUDA Binary Utilities (SASS instruction set) — ISETP, PLOP3, VOTE 등 SASS 명령어 참고
• PTX ISA — predicate register와 setp 명령의 의미

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