[triton] Triton AMD GPU 백엔드: v_perm 명령어를 활용한 레이아웃 변환 최적화

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

들어가며

Triton은 GPU 프로그래밍을 위한 강력한 프레임워크이지만, 컴파일 과정에서 생성되는 LLVM IR이 항상 하드웨어 수준에서 최적의 기계어로 변환되는 것은 아닙니다. 특히 AMD GPU 환경에서 8비트 데이터의 레이아웃을 변환(convert_layout)할 때, 기존 방식은 여러 개의 shift 및 논리 연산을 조합하여 수행되었습니다. 이는 불필요하게 많은 명령어를 생성하여 성능 저하를 유발했습니다. 본 PR은 AMD GPU의 v_perm 명령어를 직접 활용하는 패턴을 도입하여, 레이아웃 변환 효율을 극대화하는 최적화를 다룹니다.

코드 분석

1. third_party/amd/lib/TritonAMDGPUToLLVM/ConvertLayoutOpToLLVM.cpp

핵심 변경 사항은 ConvertLayoutOpInThreadSwap 클래스의 도입입니다. 이 클래스는 ttg.convert_layout 연산을 가로채어, 하드웨어 수준에서 지원하는 llvm.amdgcn.perm 인트린식을 호출하도록 합니다.

Before (기존 방식): 기존에는 레이아웃 변환 시 비트 연산과 시프트 연산의 조합으로 구현되어, 컴파일러가 이를 최적화하지 못할 경우 명령어 수가 3배 이상 증가하는 문제가 있었습니다.

After (최적화된 방식):

static Value createVPerm(TritonLLVMOpBuilder &b, Value v1, Value v2,
                         ArrayRef<int> shuffleIds) {
  auto loc = b.loc;
  auto &rewriter = *b.builder;
  std::string intrinsic = "llvm.amdgcn.perm";
  int encodedIndices = 0;
  for (int i = 0; i < shuffleIds.size(); ++i) {
    assert(shuffleIds[i] >= 0 && shuffleIds[i] < 8);
    encodedIndices += shuffleIds[i] << (i * 8);
  }
  auto encodedIndicesVal = b.int_val(32, encodedIndices);
  return LLVM::createLLVMIntrinsicCallOp(rewriter, loc, intrinsic, {i32_ty},
                                         {v1, v2, encodedIndicesVal})
      .getResult(0);
}

v_perm은 두 개의 32비트 레지스터에서 바이트 단위로 데이터를 선택하여 재배치할 수 있습니다. 위 코드는 shuffleIds를 기반으로 인코딩된 인덱스를 생성하고, 이를 llvm.amdgcn.perm 인트린식으로 매핑하여 단일 명령어로 처리가 가능하게 합니다.

2. python/test/gluon/test_core.py

최적화가 실제로 적용되었는지 검증하기 위해 테스트 코드가 추가되었습니다. 특히 v_perm 명령어가 생성되는지 확인하는 정규식 검사가 포함되었습니다.

# 테스트 코드 중 일부
assert re.search(r"v_perm", pgm.asm['amdgcn'], re.MULTILINE)

왜 이게 좋은가

1. 명령어 수 감소: 기존의 복잡한 비트 연산 조합을 단일 v_perm 명령어로 대체함으로써, 레지스터 간 데이터 이동 비용을 획기적으로 줄였습니다.
2. 하드웨어 가속: AMD GPU의 하드웨어 Permute 유닛을 직접 활용하므로, 범용적인 비트 연산보다 훨씬 높은 처리량을 보장합니다.
3. 교훈: 컴파일러의 자동 최적화(InstCombine 등)를 무조건 신뢰하기보다, 특정 하드웨어 아키텍처에서 성능 병목이 발생하는 패턴(unpack-reorder-pack)을 식별하고 이를 명시적인 패턴 매칭으로 해결하는 것이 고성능 커널 개발의 핵심임을 보여줍니다.

리뷰 과정에서 논의되었듯이, 장기적으로는 LLVM의 최적화 휴리스틱을 개선하는 것이 근본적인 해결책이지만, 당장의 성능 확보를 위해 Triton 수준에서 명시적인 패턴을 도입하는 것은 매우 실용적인 접근입니다.

참고 자료

• https://llvm.org/docs/AMDGPUUsage.html
• https://github.com/triton-lang/triton

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