[triton] Boolean 서브그룹 reduction을 warp vote와 ballot 한 번으로 끝내기: Triton NVIDIA 백엔드 최적화

PR 링크: triton-lang/triton#12647 상태: Merged | 변경: +132 / -24

들어가며

Triton에서 tt.reduce로 Boolean(i1) 텐서를 축소하는 일은 생각보다 흔합니다. 마스크의 any/all, early-exit 조건 검사, 정렬·탐색 커널의 수렴 판정 등이 모두 i1 reduction입니다.

기존 NVIDIA 백엔드는 warp 전체(32 lane)가 하나의 reduction 그룹일 때만 vote.sync.any / vote.sync.all을 사용했습니다. 조건은 reduceLaneIdMask == 31이었죠. 그런데 실제 레이아웃에서는 2, 4, 8, 16개 lane이 하나의 행(row)을 나눠 갖는 부분(subgroup) reduction이 자주 등장합니다. 이 경우 컴파일러는 shfl.sync.bfly 기반의 butterfly shuffle 트리를 log2(lanes) 단계만큼 펼쳤습니다. 1비트 값을 줄이는데 여러 번의 shuffle + 연산이 직렬로 이어지는 셈입니다.

이 PR은 두 가지 관찰로 이 문제를 해결합니다.

1. 복제(replica)는 AND/OR 결과를 바꾸지 않는다. warp 안에 한 그룹의 복사본만 있다면 warp 전체에 vote.sync.any/all을 그대로 쓰면 됩니다.
2. 독립 그룹이 여러 개라면 vote.sync.ballot 한 번으로 32개 predicate를 32비트 정수로 모은 뒤, 각 lane이 자기 그룹의 lane mask로 AND 연산하여 판정하면 됩니다.

대상은 i1의 AND/OR, 그리고 i1에서 의미상 동일한 minui(=all)/maxui(=any)입니다.

코드 분석

1. TargetInfo.cpp — vote 경로의 일반화

Before: full warp일 때만 vote를 사용합니다.

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) {

After: lane mask 조건이 사라지고, reduction lane과 broadcast(복제) lane을 합친 "varying bits"로 그룹을 계산합니다.

bool all =
    *kind == NVVM::ReductionKind::AND || *kind == NVVM::ReductionKind::UMIN;
bool any =
    *kind == NVVM::ReductionKind::OR || *kind == NVVM::ReductionKind::UMAX;
if (value.getType().isInteger(1) && (all || any)) {
  // Repeating predicates preserves AND/OR, so replicas share a group.
  unsigned varyingBits = reduceLaneIdMask | bro...

주석이 핵심입니다. a AND a = a, a OR a = a(멱등성)이므로 같은 값이 여러 lane에 복제되어 있어도 결과가 그대로입니다. 그래서 broadcast lane도 "같은 그룹"으로 편입할 수 있고, varying bits가 31이 되면 결국 full-warp vote 하나로 끝납니다.

그룹이 여러 개라면 ballot을 씁니다. 테스트가 생성 코드를 잘 보여줍니다.

// REDUX-DAG: %[[LEADER:.*]] = llvm.and %{{.*}}, %[[BIT]]
// REDUX-DAG: %[[MEMBERS:.*]] = llvm.shl %[[LOW]], %[[LEADER]]
// REDUX-DAG: %[[NEG:.*]] = llvm.xor
// REDUX-DAG: %[[BALLOT:.*]] = nvvm.vote.sync ballot %[[FULL]], %[[NEG]] -> i32
// REDUX-DAG: %[[HIT:.*]] = llvm.and %[[BALLOT]], %[[MEMBERS]]
// REDUX: llvm.icmp "eq" %[[HIT]], %[[ZERO]]

• LEADER = laneId & 16: 그룹을 구분하는 고정 비트만 남겨 그룹 리더 위치를 얻습니다.
• MEMBERS = 0xFFFF << LEADER: "그룹 0의 lane mask"를 리더 위치로 shift해 내 그룹의 mask를 만듭니다.
• all은 !p를 ballot하고 "내 그룹에 false가 하나도 없음(== 0)"으로 판정합니다. any는 p를 ballot하고 != 0으로 판정합니다.

even/odd 레이아웃에서는 그룹 0 mask가 1431655765(=0x55555555)가 되어 짝수 lane을, << 1로 홀수 lane을 고릅니다. shuffle 트리가 vote.ballot 1회 + 정수 연산 몇 개로 바뀝니다.

2. LayoutUtils.cpp — getGroupZeroMask 추출

"varying bit 집합의 모든 부분합으로 이뤄진 ID 집합"을 32비트 mask로 만드는 헬퍼입니다.

uint32_t getGroupZeroMask(unsigned varyingIdBits) {
  assert(varyingIdBits < 32);
  uint32_t pattern = 1;
  for (unsigned stride = 1; stride <= varyingIdBits; stride <<= 1)
    if (varyingIdBits & stride)
      pattern |= pattern << stride;
  return pattern;
}

주석 예시대로 bits 0과 2가 varying이면 ID {0,1,4,5}, 즉 0b110011이 됩니다. 각 varying 비트마다 현재 패턴을 해당 stride만큼 복사하는 방식이라 루프는 최대 5회입니다.

3. Ops.cpp — 기존 TMA multicast 코드 재사용

흥미롭게도 동일한 알고리즘이 이미 TMA multicast mask 계산에 있었습니다.

Before:

int blockBits = llvm::Log2_32(numCTAs);
uint32_t fixedBits = (~broadcastBits) & (numCTAs - 1);
uint32_t pattern = 1;
for (int i = 0; i < blockBits; ++i) {
  if ((fixedBits & (1u << i)) == 0)
    pattern |= (pattern << (1u << i));
}
return {fixedBits, pattern};

After:

uint32_t varyingBits = broadcastBits & (numCTAs - 1);
return {(numCTAs - 1u) ^ varyingBits, getGroupZeroMask(varyingBits)};

CTA multicast 그룹과 warp lane 그룹은 "일부 ID 비트는 고정, 나머지는 변하는 집합"이라는 같은 구조입니다. 리뷰어 Jokeren도 "이 함수는 Triton dialect와 무관하게 일반적"이라고 지적했고, 결과적으로 triton/Tools/LayoutUtils로 옮겨져 두 곳에서 공유하게 됐습니다.

4. matchReduxKind — 단일 combiner 조건 강화

// A warp instruction replaces the whole combine region.
if (!reduceOp ||
    !llvm::hasSingleElement(op.getCombineOp().front().without_terminator()))
  return std::nullopt;

getSingleCombiner()의 의미가 "반환되는 binary op"로 바뀌면서 region에 다른 op(예: tt.print)가 섞여 있어도 combiner가 반환될 수 있게 됐습니다. 하지만 warp 명령 하나가 region 전체를 대체하면 부수효과가 사라지므로, 여기서 op가 정확히 하나인지 별도로 확인합니다. @reduce_predicate_with_print 테스트가 이 경우 shfl.sync bfly와 vprintf가 유지됨을 검증합니다. 이와 관련해 Jokeren은 getSingleCombiner 내부의 일부 검사가 "이미 아래에서 operand가 block argument인지 확인하므로 중복일 수 있다"고 리뷰했습니다 — 의미 변경 시 검증 로직의 책임을 어디에 둘지에 대한 좋은 논의입니다.

5. XOR 테스트의 디테일

// Each warp holds eight copies of its four predicates. XOR counts each once.
// REDUX: %[[UNIQUE:.*]] = llvm.select
// REDUX: nvvm.redux.sync xor %[[UNIQUE]], %[[MASK]]

XOR은 멱등이 아니므로(a ^ a = 0) 복제 lane을 그대로 섞으면 안 됩니다. 대표 lane만 값을 남기는 select 후 redux.sync xor를 쓰는 것을 확인합니다. 최적화가 AND/OR에만 적용되는 이유를 정확히 보여주는 대조군입니다.

왜 이게 좋은가

GB300/GB200/H100에서 4 warps/CTA, 연속 lane 그룹 Boolean row reduction의 커널 시간 변화(음수가 빠름):

Lanes	GB300 any	GB300 all	H100 any	H100 all
4	-1.4%	-1.8%	-1.5%	-1.8%
8	-2.4%	-3.6%	-2.5%	-3.3%
16	-0.6%	-0.6%	-1.3%	-1.7%
32	0.0%	+0.3%	0.0%	+0.1%

8-lane 그룹에서 최대 3.6% 개선되었고, full-warp 대조군은 기계어가 동일하므로 변화는 측정 노이즈입니다. 커널 전체 시간 기준이므로 reduction 부분만 보면 개선폭은 더 큽니다.

일반적 교훈:

• 연산의 대수적 성질(멱등성)을 활용하라. AND/OR은 복제에 불변이므로 레이아웃의 broadcast 차원을 그룹에 흡수할 수 있습니다. XOR은 그렇지 않다는 대조 테스트까지 두는 것이 좋은 습관입니다.
• 1비트 데이터는 bitmask로 다뤄라. ballot은 warp 전체 상태를 32비트 정수 하나로 모아 주므로, 여러 독립 그룹의 판정이 정수 AND 한 번으로 끝납니다. log2(n)단계 shuffle의 직렬 의존성이 사라집니다.
• 같은 구조의 문제는 같은 헬퍼로. CTA multicast와 lane 그룹 mask가 동일한 "고정/가변 비트" 구조임을 알아보고 공통 유틸리티로 추출했습니다.
• 명령어 치환 시 부수효과를 보존하라. region 전체를 하드웨어 명령 하나로 바꿀 때는 region 내 다른 op가 없음을 반드시 확인해야 합니다.

참고 자료

• PTX vote instruction — vote.sync.any/all/ballot 명령 공식 문서
• CUDA Warp Vote Functions — __any_sync, __all_sync, __ballot_sync 설명
• triton.language.reduce — Triton reduce API 문서

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