[onnxruntime] ONNX Runtime의 CPU BatchNormalization 성능 최적화: NC 입력 데이터 처리 개선

PR 링크: microsoft/onnxruntime#33113 상태: Merged | 변경: +119 / -13

들어가며

딥러닝 모델을 CPU에서 실행할 때, BatchNormalization은 종종 예상치 못한 병목 지점이 됩니다. 특히 PyTorch에서 내보낸 모델의 [N, C] 형태(NC 입력) 데이터는 메모리상에서 채널이 연속적으로 배치되어 있음에도 불구하고, 기존 ONNX Runtime의 CPU 커널은 이를 공간적 차원이 있는 것처럼 처리하여 불필요한 오버헤드를 발생시켰습니다. 본 PR은 이러한 NC 입력 데이터에 대해 메모리 레이아웃을 최적화하여 연산을 벡터화함으로써 성능을 극적으로 개선했습니다.

코드 분석

onnxruntime/core/providers/cpu/nn/batch_norm.h

핵심 변경 사항은 use_spatial_map 플래그를 도입하여 입력 데이터의 차원 구조에 따라 Eigen 매핑 방식을 분기한 것입니다.

Before:

ConstEigenArrayMap<T> X_arr(X->Data<T>(),
                            is_spatial_ ? sample_size : sample_size_incl_all_channels,
                            is_spatial_ ? N * C : N);

After:

const bool use_spatial_map = is_spatial_ && sample_size != 1;
ConstEigenArrayMap<T> X_arr(X->Data<T>(),
                            use_spatial_map ? sample_size : sample_size_incl_all_channels,
                            use_spatial_map ? N * C : N);

기존에는 is_spatial_ 플래그에만 의존했으나, 이제 sample_size == 1인 경우(즉, NC 입력)를 별도로 식별합니다. 이를 통해 saved_mean과 saved_var를 계산할 때 루프 내부에서 스칼라 연산을 반복하는 대신, Eigen의 벡터화된 연산을 활용할 수 있게 되었습니다.

if (sample_size == 1) {
  for (size_t n = 0; n < N; ++n) {
    saved_mean_arr += X_arr.col(n);
  }
  // ... (분산 계산도 동일하게 벡터화)
}

onnxruntime/test/providers/cpu/nn/batch_norm_op_test.cc

최적화의 안정성을 검증하기 위해 TestBatchNormContiguousChannels 템플릿 함수가 추가되었습니다. 다양한 랭크와 데이터 타입(float, double), 그리고 학습/추론 모드를 포괄하는 72개의 테스트 케이스를 통해 기존 로직과의 결과 일치성을 보장합니다.

왜 이게 좋은가

이 최적화는 '메모리 레이아웃에 맞는 연산 방식 선택'이라는 고전적이지만 강력한 원칙을 따릅니다. 기존 구현은 N * C개의 작은 블록을 처리하며 오버헤드가 컸으나, 변경 후에는 채널 벡터를 한 번에 처리함으로써 CPU의 SIMD(Single Instruction, Multiple Data) 유닛을 최대한 활용할 수 있게 되었습니다.

성능 지표 (AMD EPYC 9654 기준):

• Inference: 0.3334 ms → 0.0297 ms (약 11.23배 향상)
• Training: 0.8225 ms → 0.0625 ms (약 13.15배 향상)

이러한 개선은 별도의 버퍼 할당이나 스레딩 변경 없이 오직 데이터 접근 방식의 최적화만으로 달성되었다는 점에서 매우 효율적입니다. 일반적인 소프트웨어 엔지니어링 관점에서도, 데이터 구조를 알고리즘에 맞추는 것이 아니라 알고리즘을 데이터 구조의 특성(Contiguous Memory)에 맞추는 것이 성능 최적화의 핵심임을 다시 한번 확인시켜 줍니다.

참고 자료

• https://eigen.tuxfamily.org/dox/group__TopicInsideEigenExample.html
• https://onnxruntime.ai/docs/api/c/index.html

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