[onnxruntime] ONNX Runtime WebGPU: Intel Xe-3LPG를 위한 고성능 GEMM 최적화 분석

PR 링크: microsoft/onnxruntime#29271 상태: Merged | 변경: +168 / -40

들어가며

최근 microsoft/onnxruntime 레포지토리에 Intel의 최신 GPU 아키텍처인 Xe-3LPG를 타겟으로 한 WebGPU GEMM(General Matrix Multiply) 커널 최적화가 병합되었습니다. 기존의 Intel 서브그룹 커널은 A 행렬 로드 시 메모리 대역폭 병목과 K 루프 내의 잦은 workgroupBarrier 호출로 인해 성능 저하가 발생하고 있었습니다. 본 PR은 vec4 로드와 B 타일 더블 버퍼링이라는 두 가지 핵심 기법을 도입하여 이 문제를 해결했습니다.

코드 분석

1. vec4 A 로드 (Cooperative Loading)

K 차원이 4의 배수일 때, 기존의 스칼라 로드 대신 vec4를 사용하여 글로벌 메모리 접근 횟수를 줄였습니다. 서브그룹 내에서 협력적으로 데이터를 로드하여 성능을 극대화합니다.

Before:

// 기존 스칼라 로드 방식 (gemm_subgroup.cc)
mm_readA(batch, globalRowStart + ..., aColV ...);

After:

// vec4 협력 로드 (gemm_subgroup.cc)
// 각 lane이 rows_per_group만큼의 행을 vec4 단위로 로드
for (uint32_t j = 0; j < rows_per_group; j++) {
  a_val_j = mm_readA(batch, globalRowStart + aGroup * rows_per_group + j, aColV);
}

2. B 타일 더블 버퍼링

FP16 데이터 타입에 대해 B 타일을 더블 버퍼링하여 연산과 메모리 로드를 오버랩시켰습니다. 이를 통해 K 루프당 workgroupBarrier를 2회에서 1회로 줄였습니다.

Before:

// 단일 버퍼 사용 시 매 루프마다 배리어 필요
mm_Bsub[inputRow][inputCol] = mm_readB(...);
workgroupBarrier();

After:

// 더블 버퍼링으로 로드 지연 시간 은닉
mm_Bsub[buf][inputRow][inputCol] = mm_readB(...);
// 루프 끝에서만 배리어 호출

왜 이게 좋은가

이번 최적화는 Xe-3LPG 아키텍처에서 평균 12.7%의 성능 향상을 기록했습니다. 특히 jina-clip-v1 모델에서 13.4%, sd-v1.5 텍스트 인코더에서 15.9%의 높은 개선폭을 보였습니다.

교훈:

1. 메모리 접근 최적화: 스칼라 로드를 벡터 로드로 변경하는 것만으로도 메모리 대역폭 효율을 크게 높일 수 있습니다.
2. 지연 시간 은닉(Latency Hiding): 더블 버퍼링을 통해 연산 유닛이 메모리 로드를 기다리는 시간을 줄여 GPU 활용도를 극대화했습니다.
3. 아키텍처별 분기: CacheHint를 통해 아키텍처별로 최적화된 파이프라인을 캐싱함으로써, 특정 아키텍처를 위한 최적화가 다른 아키텍처의 성능을 저하시키지 않도록 설계했습니다.

리뷰어 피드백 분석

리뷰어 hariharans29는 더블 버퍼링 구현이 레이스 컨디션 없이 안전하게 설계되었음을 확인했습니다. 특히 subgroupBroadcast를 활용한 협력 로드 로직과 CacheHint의 정교한 구현을 높게 평가했습니다. 다만, 아키텍처 문자열 하드코딩에 대한 우려가 있었으나, 현재 EP의 표준 패턴임을 감안하여 승인되었습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://www.w3.org/TR/webgpu/
• https://github.com/microsoft/onnxruntime/pull/29271

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