[onnxruntime] ONNX Runtime MLAS, AVX-512 최적화로 MobileClip-S0 모델 추론 속도 향상

PR 링크: microsoft/onnxruntime#31958 상태: Merged | 변경: +494 / -0

들어가며

최근 Microsoft의 ONNX Runtime 레포지토리에서는 [MLAS] AVX-512 16-wide Erf kernel and NCHWc reorder transpose for MobileClip-S0 model이라는 제목의 PR이 머지되었습니다. 이 PR은 특히 MobileClip-S0 모델의 FP32 CPU 추론 성능을 향상시키는 데 중점을 두고 있습니다. 핵심은 두 가지 주요 최적화입니다:

1. 16-wide AVX-512 Erf 커널 추가: 기존의 8-wide 구현을 대체하여 AVX-512 지원 하드웨어에서 ONNX Erf 연산의 성능을 개선합니다.
2. 16x16 AVX-512 NCHWc 재정렬 트랜스포즈 구현: 기존의 SSE2 기반 4-wide 서브 트랜스포즈 방식을 대체하여 블록-16 데이터 재정렬을 단일 패스로 처리합니다.

이 글에서는 이 PR의 코드 변경사항을 상세히 분석하고, 각 최적화가 왜 성능 향상에 기여하는지, 그리고 어떤 기술적 교훈을 얻을 수 있는지 살펴보겠습니다.

코드 분석

1. 16-wide AVX-512 Erf 커널 추가

이 최적화는 Erf 연산의 계산 속도를 높이기 위해 AVX-512 명령어셋을 활용합니다. 기존에는 8개의 float 데이터를 한 번에 처리하는 FMA3 기반 커널을 사용했지만, AVX-512를 지원하는 CPU에서는 16개의 float 데이터를 동시에 처리할 수 있는 새로운 커널을 도입했습니다.

변경점

• onnxruntime/core/mlas/lib/intrinsics/avx512/gelu_avx512f.cpp: 새로운 MlasErfKernelAvx512FImpl 함수가 추가되었습니다. 이 함수는 16개의 float 데이터를 __m512 벡터 레지스터로 로드하고, MlasGeluErfAvx512 함수를 사용하여 연산을 수행한 후, 결과를 다시 메모리에 저장합니다. N이 16보다 작을 경우를 대비한 마스크된 로드 및 저장(_mm512_maskz_loadu_ps, _mm512_mask_storeu_ps)도 포함되어 있습니다.

  void
  MlasErfKernelAvx512FImpl(
      const float* Input,
      float* Output,
      size_t N
      )
  {
      const GeluAvx512BroadcastConstants Constants;
      while (N >= 16) {
          const __m512 X = _mm512_loadu_ps(Input);
          const __m512 Result = MlasGeluErfAvx512(X, Constants); // Note: This primitive computes Erf(X)
  
          _mm512_storeu_ps(Output, Result);
  
          Input += 16;
          Output += 16;
          N -= 16;
      }
  
      if (N > 0) {
          const __mmask16 TailMask = __mmask16((1u << static_cast<unsigned>(N)) - 1u);
          const __m512 X = _mm512_maskz_loadu_ps(TailMask, Input);
          const __m512 Result = MlasGeluErfAvx512(X, Constants);
  
          _mm512_mask_storeu_ps(Output, TailMask, Result);
      }
  }
  

• onnxruntime/core/mlas/lib/platform.cpp: MLAS 플랫폼 초기화 시 AVX-512 기능이 감지되면 ErfKernelRoutine 포인터를 새로 구현된 MlasErfKernelAvx512F로 설정합니다.

  if (((Cpuid7[1] & 0x10000) != 0) && ((xcr0 & 0xE0) == 0xE0)) {
      this->GeluErfKernelRoutine = MlasGeluErfKernelAvx512F;
      this->ErfKernelRoutine = MlasErfKernelAvx512F; // Assigns the new 16-wide kernel
      this->SiluKernelRoutine = MlasSiluKernelAvx512F;
      // ... other AVX-512 assignments
  }
  

2. 16x16 NCHWc 재정렬 트랜스포즈 구현

딥러닝 모델에서 데이터 레이아웃 변환은 빈번하게 발생합니다. 특히 NCHW (Batch, Channels, Height, Width) 포맷을 NCHWc (Batch, Channels/c, Height, Width, c) 포맷으로 변환하는 과정은 채널 차원을 블록화하여 연산 효율성을 높입니다. 이 PR에서는 채널 블록 크기가 16일 때, 기존의 SSE2 기반 4x4 서브 트랜스포즈 여러 번을 수행하던 방식을 버리고, AVX-512를 사용하여 16x16 전체를 한 번에 처리하는 새로운 트랜스포즈 커널을 도입했습니다.

변경점

• onnxruntime/core/mlas/lib/intrinsics/avx512/reorder_avx512f.cpp: 새로운 파일로, MlasReorderTranspose16x16Avx512F 함수가 구현되었습니다. 이 함수는 Intel의 권장 패턴인 unpacklo/hi_ps, unpacklo/hi_pd (float-double 캐스팅 활용), shuffle_f32x4 등을 조합하여 16x16 float 데이터를 효율적으로 트랜스포즈합니다. 이 함수를 기반으로 NCHW -> NCHWc 변환(MlasReorderInputNchwBlock16Avx512F)과 NCHWc -> NCHW 변환(MlasReorderOutputNchwBlock16Avx512F)을 위한 래퍼 함수가 제공됩니다.

  // Inside MlasReorderTranspose16x16Avx512F
  // ... unpacklo/hi_ps operations ...
  // ... unpacklo/hi_pd operations ...
  // ... shuffle_f32x4 operations ...
  // ... storing results ...
  
  // Example for NCHW -> NCHWc
  void
  MLASCALL
  MlasReorderInputNchwBlock16Avx512F(
      const float* S, // Source (NCHW)
      float* D,       // Destination (NCHWc)
      size_t InputSize // Spatial size (Height * Width)
      )
  {
      // ... loop for spatial tail ...
      for (; p + 16 <= InputSize; p += 16) {
          MlasReorderTranspose16x16Avx512F(S + p, D + p * 16, InputSize, 16);
      }
      // ... scalar tail for remaining spatial elements ...
  }
  

• onnxruntime/core/mlas/lib/reorder.cpp: 기존 MlasReorderInputNchw 및 MlasReorderOutputNchwThreaded 함수 내에 AVX-512 환경에서 채널 블록 크기가 16이고 처리할 채널 수가 정확히 16개일 경우, 새로 구현된 AVX-512 커널로 분기하는 로직이 추가되었습니다.

  // Inside MlasReorderInputNchw
  if (BlockSize == 16 && InputChannelsThisIteration == 16) {
      MlasReorderInputNchwBlock16Avx512F(S, D, InputSize);
      S += BlockSize * InputSize;
      D += BlockSize * InputSize;
      continue;
  }
  

• cmake/onnxruntime_mlas.cmake: 새로 추가된 reorder_avx512f.cpp 파일이 빌드 시스템에 포함되어, AVX-512 타겟 빌드 시 컴파일됩니다.

왜 이게 좋은가?

성능 향상

PR 설명에 포함된 성능 측정 결과는 이 최적화의 효과를 명확히 보여줍니다. STRIX 365 (AMD Ryzen AI 9 365) 환경에서 MobileClip-S0 모델을 FP32로 추론했을 때, 스레드 구성에 따라 최대 1.5배 이상의 성능 향상을 보였습니다. 특히 16개의 float 데이터를 한 번에 처리하는 AVX-512 Erf 커널과 16x16 데이터를 단일 패스로 처리하는 NCHWc 재정렬 트랜스포즈는 SIMD(Single Instruction, Multiple Data) 연산의 이점을 극대화하여 병렬 처리 효율성을 크게 높였습니다.

기술적 교훈

1. 최신 SIMD 명령어셋 적극 활용: AVX-512와 같은 최신 SIMD 명령어셋은 특정 연산에서 기존 명령어셋 대비 훨씬 높은 성능 향상을 제공할 수 있습니다. CPU 아키텍처의 특성을 깊이 이해하고 이를 활용하는 코드를 작성하는 것이 중요합니다.
2. 데이터 재정렬 최적화의 중요성: 딥러닝 모델에서는 데이터 레이아웃 변환이 성능에 큰 영향을 미칩니다. NCHWc와 같은 포맷은 메모리 접근 패턴을 개선하여 SIMD 연산 효율을 높이는 데 기여하며, 이러한 재정렬 연산을 최적화하는 것은 전체 추론 성능 향상에 필수적입니다.
3. 철저한 테스트 커버리지: 새로운 커널을 도입할 때는 기존 커널과의 정확도(ULP - Unit in the Last Place) 비교뿐만 아니라, 표준 라이브러리 함수와의 비교, 특수 값(NaN, Inf 등) 처리, 다양한 데이터 길이 및 경계 조건에 대한 테스트가 필수적입니다. 이 PR의 test_erf.cpp와 test_reorder_input.cpp는 이러한 철저한 테스트의 좋은 예시입니다.
4. 점진적 최적화 및 런타임 디스패치: 모든 환경에서 최신 명령어셋을 강제하는 대신, 런타임에 CPU 기능을 감지하고 최적화된 커널을 선택적으로 사용하는 platform.cpp의 디스패치 방식은 코드의 호환성과 성능을 동시에 확보하는 현명한 방법입니다. 또한, 기존 코드 경로에 새로운 최적화 경로를 추가하고, 특정 조건(예: BlockSize == 16)에서만 새로운 경로를 사용하도록 하는 방식은 안정성을 높입니다.

리뷰 요약 및 추가 고려사항

리뷰어(@hariharans29)는 이 PR이 정확성 측면에서 승인 가능하며, 두 가지 주요 AVX-512 최적화가 잘 구현되었다고 평가했습니다. 특히 Erf 커널의 정확도 테스트와 재정렬 테스트의 철저함을 칭찬했습니다. 다만, 몇 가지 개선 사항이 제안되었습니다:

• CLA (Contributor License Agreement) 확인: 외부 기여자의 경우 CLA 동의가 필요합니다. 이 PR에서는 swetha097님이 `company=

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