[triton] AMD TargetInfo에 16/32비트 Elementwise 벡터화 지원 추가

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

들어가며

Triton의 코드 생성기는 타겟 GPU의 하드웨어 특성을 TargetInfo 인터페이스를 통해 질의합니다. Elementwise 연산의 벡터화 가능 여부도 이 인터페이스를 통해 결정되는데, AMD GPU에서 이 플래그가 비활성화되어 있어 reduction 등에서 벡터 명령어를 활용하지 못했습니다.

핵심 코드 분석

Before: (해당 메서드 없음 - 기본값 false 반환)

After:

bool TargetInfo::supportBitwidth16Elementwise() const { return true; }

bool TargetInfo::supportBitwidth32Elementwise() const {
  return getISAFamily() == ISAFamily::GFX1250;
}

모든 AMD GPU에서 16비트(f16) elementwise 벡터화를 활성화하고, GFX1250에서만 32비트(f32) 벡터화를 추가로 활성화합니다. 테스트는 reduction tree에서 벡터 연산이 생성되는지 확인합니다:

// GFX942: llvm.fadd {{.*}} : vector<2xf16>   -- 16비트 벡터화 확인
// GFX942-NOT: llvm.fadd {{.*}} : vector<2xf32> -- 32비트는 미지원

// GFX1250: llvm.fadd {{.*}} : vector<2xf16>   -- 16비트 벡터화
// GFX1250: llvm.fadd {{.*}} : vector<2xf32>   -- 32비트도 벡터화

왜 이게 좋은가

이 변경은 작지만 누적 효과가 큰 최적화입니다. Reduction은 트랜스포머 모델의 LayerNorm, Softmax 등에서 빈번하게 사용됩니다. 16비트 벡터화로 reduction tree의 각 단계에서 2개의 요소를 동시에 처리하므로, 레지스터 수준의 연산량이 절반으로 줄어듭니다. GFX1250에서는 f32 벡터 연산도 하드웨어가 지원하므로 추가 활성화하되, 이전 아키텍처에서는 보수적으로 비활성 상태를 유지합니다.

정리

• 모든 AMD GPU에서 supportBitwidth16Elementwise 활성화
• GFX1250에서 supportBitwidth32Elementwise 추가 활성화
• Reduction tree에서 vector<2xf16> / vector<2xf32> 벡터 연산 생성
• GFX942, GFX950, GFX1250 각각에 대한 테스트 포함

참고 자료

• triton-lang/triton#9509

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