[onnxruntime] ONNX Runtime WebGPU: NHWC Depthwise Convolution 최적화 분석

PR 링크: microsoft/onnxruntime#32498 상태: Merged | 변경: +182 / -11

들어가며

ONNX Runtime의 WebGPU 실행 공급자(Execution Provider)는 다양한 모델의 추론을 가속화하기 위해 최적화된 셰이더를 사용합니다. 하지만 기존의 GroupedConvProgram은 일반적인 그룹 컨볼루션(Grouped Convolution)을 처리하도록 설계되어 있어, 특정 상황에서 성능 병목이 발생했습니다. 특히 NHWC(Channels-Last) 레이아웃에서의 Depthwise Convolution은 입력과 출력 채널이 1:1로 대응됨에도 불구하고, 기존 로직이 이를 일반적인 그룹 컨볼루션으로 간주하여 스칼라(Scalar) 방식으로 처리하는 비효율이 있었습니다. 본 PR은 이 문제를 해결하기 위해 Depthwise 전용 벡터화 경로를 추가하여 성능을 개선했습니다.

코드 분석

1. onnxruntime/core/providers/webgpu/nn/conv.cc: 최적화 경로 활성화

기존에는 output_channels_per_group >= 4인 경우에만 벡터화를 수행했습니다. Depthwise 컨볼루션은 그룹당 채널이 1개이므로 이 조건에 걸리지 않아 벡터화 혜택을 받지 못했습니다.

Before

auto components = static_cast<int>(is_channels_last && output_channels_per_group >= 4 ? GetMaxComponents(output_channels) : 1);

After

const bool is_depthwise_vec = is_channels_last && output_channels_per_group == 1 &&
                              conv_attrs_.group == input_channels && GetMaxComponents(output_channels) > 1;
auto components = static_cast<int>(is_channels_last && (output_channels_per_group >= 4 || is_depthwise_vec)
                                       ? GetMaxComponents(output_channels)
                                       : 1);

이 수정으로 Depthwise 컨볼루션이 벡터화 가능한 조건(is_depthwise_vec)을 만족할 때, 4개 채널을 한 번에 처리할 수 있게 되었습니다.

2. onnxruntime/core/providers/webgpu/nn/grouped_conv.cc: 전용 셰이더 로직 구현

일반적인 경로와 달리, Depthwise 전용 경로는 오프셋을 직접 계산하여 인덱스 재구성을 최소화합니다.

After (신규 추가된 CalculateResultDepthwiseVec)

// 오프셋을 단계별로 계산하여 4차원 인덱스 재구성을 방지
let x_base = ((x_row_base + u32(xHeight)) * uniforms.x_shape[2]) * uniforms.x_shape[3] + output_channel;
let w_base = wHeight * uniforms.w_shape[1] * w_plane + output_channel;
// ... 내부 루프에서 벡터 연산 수행
value += x.GetByOffset("x_base + u32(xWidth) * uniforms.x_shape[3]") * w.GetByOffset("w_base + wWidth * w_plane");

왜 이게 좋은가

1. 벡터화(Vectorization): NHWC 레이아웃에서 채널이 가장 안쪽에 위치하므로, 4개의 채널을 하나의 벡터로 묶어 처리함으로써 메모리 대역폭 효율을 극대화했습니다.
2. 연산 오버헤드 감소: 기존의 일반적인 그룹 컨볼루션은 매 탭(tap)마다 복잡한 4차원 인덱스 계산을 수행해야 했습니다. 반면, 이 최적화는 오프셋을 단계적으로 증가시키는 방식을 사용하여 루프 내부의 산술 연산을 획기적으로 줄였습니다.
3. 성능 수치: RTX 3060 환경에서 문서 레이아웃 모델 추론 시, 전체 16ms 중 0.16ms를 차지하던 컨볼루션 연산에서 유의미한 개선을 확인했습니다.

교훈: 범용적인 알고리즘이 항상 최적은 아닙니다. 특정 레이아웃(NHWC)과 연산 특성(Depthwise)이 결합된 경우, 전용 경로를 분리하여 구현하는 것이 성능 최적화의 핵심입니다.

리뷰 피드백 반영

리뷰 과정에서 xRCCorner가 vec2<i32>로 변경됨에 따라, 기존의 unsigned wraparound를 이용한 경계 검사가 무효화되는 문제가 발견되었습니다. 이를 명시적인 if (xHeight < 0 || ...) 체크로 수정하여 패딩(Padding)이 포함된 케이스에서도 정확한 동작을 보장하도록 개선했습니다.

참고 자료

• https://www.w3.org/TR/WGSL/

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