[onnxruntime] ONNX Runtime WebGPU: Reduce 연산 최적화를 통한 성능 향상

PR 링크: microsoft/onnxruntime#28174 상태: Merged | 변경: +15 / -9

들어가며

최근 대규모 언어 모델이나 비전 모델을 웹 환경에서 실행할 때, WebGPU의 성능은 사용자 경험을 결정짓는 핵심 요소입니다. Microsoft의 ONNX Runtime(ORT)은 다양한 하드웨어 가속을 지원하는데, 이번 PR은 WebGPU 백엔드에서 ReduceMean과 같은 Reduction 연산이 수행될 때 발생하는 비효율적인 메모리 복사 및 Transpose 문제를 해결합니다. 특정 조건에서 'Shared reduction' 방식이 오히려 성능 저하를 유발하는 문제를 발견하고, 이를 'Naive reduction'으로 전환하는 휴리스틱을 도입하여 성능을 크게 개선했습니다.

코드 분석

onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc

기존 코드에서는 Reduction 연산 시 무조건적으로 Shared reduction 방식을 우선시하거나, 복잡한 조건 분기를 통해 연산을 수행했습니다. 특히, 입력 텐서의 차원이 정렬되어 있지 않을 경우, Shared reduction을 위해 Transpose 연산이 선행되어야 했는데, 이 비용이 매우 컸습니다.

Before:

bool use_naive_reduction = name_ == "ArgMin" || name_ == "ArgMax" || (reduce_size < 32 && output_size > 1024) || is_input_empty || input_tensor->Shape().NumDimensions() == 0;

After:

constexpr size_t kReduceNaiveMaxReduceSize = 128;
constexpr size_t kReduceNaiveMinOutputSize = 20000;
// ... (axes 확인 로직) ...
bool use_naive_reduction = name_ == "ArgMin" || name_ == "ArgMax" || (reduce_size < 32 && output_size > 1024) ||
                           (!are_axes_innermost && reduce_size <= kReduceNaiveMaxReduceSize &&
                            output_size > kReduceNaiveMinOutputSize) ||
                           is_input_empty || input_tensor->Shape().NumDimensions() == 0;

핵심 변경 사항은 !are_axes_innermost 조건이 추가된 것입니다. 즉, Reduction 대상 축이 메모리상에서 연속적이지 않을 때(innermost가 아닐 때), Transpose를 강제하는 Shared 방식 대신 Naive 방식을 선택하도록 휴리스틱을 개선했습니다. 또한, kReduceNaiveMaxReduceSize와 kReduceNaiveMinOutputSize라는 임계값을 도입하여, 특정 크기 이상의 연산에서만 이 최적화가 동작하도록 제어했습니다.

왜 이게 좋은가

이번 최적화의 핵심은 '불필요한 데이터 이동 최소화'입니다. 리뷰어 xhcao의 데이터에 따르면, florence-2-base-vision-encoder 모델 실행 시 ReduceMean 노드에서 발생하는 Transpose 연산이 전체 추론 시간의 상당 부분을 차지했습니다.

• 성능 수치: 기존 방식 대비 ReduceMean 연산 자체의 시간은 0.30ms에서 0.11ms로 약 60% 이상 단축되었으며, 전체 모델 추론 시간은 약 340ms에서 270ms로 약 20% 향상되었습니다.
• 교훈: GPGPU 프로그래밍에서 커널의 복잡도(Shared memory 활용 등)보다 더 중요한 것은 메모리 레이아웃입니다. 데이터가 정렬되지 않아 발생하는 Transpose 비용은 커널 자체의 연산 효율보다 훨씬 클 수 있습니다. 따라서, 하드웨어의 특성에 맞는 적절한 알고리즘 선택(Naive vs Shared)이 성능 최적화의 핵심임을 보여줍니다.

리뷰 과정에서 qjia7은 이 로직을 특정 연산에 국한하지 않고 모든 Reduce 연산에 적용할 것을 제안했고, 결과적으로 더 범용적인 최적화 코드가 완성되었습니다. 이는 코드의 일관성을 유지하면서도 성능을 극대화하는 좋은 사례입니다.

참고 자료

• https://github.com/microsoft/onnxruntime
• https://www.w3.org/TR/webgpu/

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