[triton] getTranspositionSelectors 알고리즘 단순화 및 복원

PR 링크: triton-lang/triton#9713 상태: Merged | 변경: +106 / -100

들어가며

GPU에서 레이아웃 변환(convert_layout)은 레지스터와 lane 간의 데이터 교환을 수반합니다. Triton은 이를 permutation의 대수적 분해(algebraic decomposition)로 표현하여, prmt(permute) 명령어와 shuffle로 변환합니다. 이 PR은 다중 mixed transposition(레지스터-lane 간 교환)에서 발생하던 정합성 문제를 수정하고 알고리즘을 단순화합니다.

핵심 코드 분석

핵심은 mixed transposition에서 intra-register bit를 peel off하는 알고리즘입니다.

Before (다중 mixed transposition에서 packing 비활성화):

// TODO: getTranspositionSelectors incorrectly lowers some conversions with
// multiple mixed transpositions. Disable packing for multiple mixed
// transpositions until we are confident.
if (m > 1)
    nPack = 0;

After (올바른 분해로 packing 항상 활성화):

// nPack 제한 제거 - 알고리즘이 올바르게 동작
auto processedTranspos =
    getTranspositionSelectors(mixedTranspositions, regBases, nPack);

알고리즘의 핵심 변경은 비교환적(noncommutative) 경우의 처리입니다:

if (forwardEnd == cycle.end()) {
    // Isolated low bits with single mixed transposition
    if (cycle.size() == 2) postShufLoBits.pop_back();
    head = tail = cycle.back();
} else if (*forwardEnd < nPack) {
    // End at a different paired low bit
    head = rBit;
    tail = next(head);
    preShufLoBits.push_back(head);
} else {
    // End at original or unpaired high bit
    head = partnerBit = *forwardEnd;
}

왜 이게 좋은가

기존의 m > 1일 때 packing을 비활성화하는 workaround는 추가 명령어를 생성하여 레이아웃 변환 성능을 저하시켰습니다. 이 PR은 근본 원인인 비교환적 transposition 순서 문제를 해결하여, stable_partition으로 올바른 적용 순서를 보장하고 모든 경우에 packing을 활성화합니다. 수학적 주석이 크게 개선되어 P = (r3 r0)(r3 l0)(r3 r0) * (r2 l1)(r2 r1)(r2 r0) 같은 분해를 명확히 설명합니다.

정리

다중 mixed transposition에서의 prmt selector 생성 알고리즘을 수정하여 packing 제한을 제거하고, 비교환적 경우의 정합성을 보장하도록 transposition 순서를 올바르게 정렬했습니다.

참고 자료

• triton-lang/triton#9713

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