[onnxruntime] ONNX Runtime WebGPU: Subgroup Matrix 연산을 활용한 Conv 최적화

PR 링크: microsoft/onnxruntime#32252 상태: Merged | 변경: +5191 / -825

들어가며

최신 비전 모델에서 Convolution(Conv) 연산은 전체 추론 시간의 상당 부분을 차지합니다. 기존의 ONNX Runtime WebGPU EP는 MatMul 연산에 대해서는 Intel Xe2/Xe3 아키텍처의 Subgroup Matrix 가속 기능을 활용하고 있었으나, Conv 연산은 이를 충분히 활용하지 못하고 있었습니다. 본 PR은 Conv 연산을 암시적 GEMM(Implicit GEMM)으로 변환하고, Subgroup Matrix 연산을 적용하여 성능을 최적화하는 내용을 담고 있습니다.

코드 분석

1. Subgroup Matrix 가속을 위한 Padding 로직 분리

기존에는 subgroup_matrix_matmul.cc 내부에 하드코딩되어 있던 Padding 로직을 별도의 유틸리티 함수로 추출하여 재사용성을 높였습니다.

Before:

// subgroup_matrix_matmul.cc 내부의 로직
class SubgroupMatrixMatMulPadBProgram final : public Program<SubgroupMatrixMatMulPadBProgram> { ... };

After:

// core/providers/webgpu/math/subgroup_matrix_pad_b.cc
Status PadSubgroupMatrixWeightColumns(ComputeContext& context, const Tensor& b, const TensorShape& b_shape, uint32_t N, uint32_t n_b, std::unique_ptr<Tensor>& padded) {
  // ... PadProgram을 활용한 범용적인 패딩 로직 구현
}

이 변경을 통해 Conv 연산에서도 동일한 패딩 로직을 사용하여 가중치 정렬을 맞출 수 있게 되었습니다.

2. Conv 연산의 Implicit GEMM 경로 추가

conv.cc에 SubgroupMatrixConv 경로를 추가하여, 특정 조건(F16, channels-last, group=1 등)을 만족하는 경우 가속 경로를 타도록 설계했습니다.

// onnxruntime/core/providers/webgpu/nn/conv.cc
#include "core/providers/webgpu/nn/subgroup_matrix_conv.h"
// ...
if (SubgroupMatrixConv::IsSupported(...)) {
  return SubgroupMatrixConv::Compute(...);
}

왜 이게 좋은가

이번 최적화는 특히 비전 모델(ResNet, YOLO, Stable Diffusion 등)에서 괄목할 만한 성능 향상을 보여줍니다. 리뷰 과정에서 공유된 벤치마크 데이터에 따르면, resnet50-v1-f16 모델의 경우 기존 대비 최대 50% 이상의 성능 향상을 기록했습니다.

주요 교훈

1. 하드웨어 가속 활용: GPU의 특정 기능(Subgroup Matrix)을 활용하기 위해 데이터를 정렬(Padding)하는 비용보다, 연산 속도 향상으로 얻는 이득이 훨씬 큽니다.
2. 코드 재사용성: PadSubgroupMatrixWeightColumns와 같이 유틸리티를 분리함으로써, MatMul과 Conv 간의 로직 중복을 제거하고 유지보수성을 높였습니다.
3. 정밀도 이슈 대응: 리뷰 과정에서 F16 연산 시 발생할 수 있는 오버플로우 문제를 발견하고, F32 누적기를 사용하는 방식으로 대응한 점은 실제 프로덕션 환경에서 매우 중요한 기술적 통찰입니다.

결론

이번 PR은 WebGPU 환경에서 고성능 딥러닝 추론을 가능하게 하는 중요한 이정표입니다. 특히 하드웨어의 특성을 고려한 타일링 전략과 정밀도 관리 전략은 향후 다른 연산 최적화에도 훌륭한 레퍼런스가 될 것입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://www.w3.org/TR/webgpu/

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