[onnxruntime] ONNX Runtime MLAS, AVX-512 NCHWc Depthwise Convolution 최적화: Sliding Window 커널 도입

2026년 10월 5일수정: 2026년 10월 5일

PR 링크: microsoft/onnxruntime#32749 상태: Merged | 변경: +840 / -8

들어가며

딥러닝 모델의 추론 속도 향상은 하드웨어의 성능을 최대한 활용하는 최적화된 커널 구현에 달려있습니다. 특히 컨볼루션 연산은 모델의 연산량 대부분을 차지하므로, 이에 대한 최적화는 매우 중요합니다. Microsoft의 ONNX Runtime은 MLAS(Microsoft Linear Algebra Subprograms) 라이브러리를 통해 다양한 하드웨어 가속 기능을 제공합니다.

이번 PR은 ONNX Runtime의 MLAS 라이브러리에서 AVX-512 명령어셋을 활용하는 NCHWc 포맷의 Depthwise Convolution 연산에 대한 새로운 최적화 기법인 'Sliding Window' 커널을 도입합니다. 기존의 AVX-512 커널(MlasConvDepthwiseFloatKernelAvx512F)은 Zen 5 CPU에서 FMA 연산 피크 성능의 31-46% 수준에 머물렀습니다. 이는 FMA 연산의 잠재력을 완전히 활용하지 못하고, 패딩(padding) 영역 처리에 비효율이 있었기 때문입니다. 이 PR은 MlasConvDepthwiseFloatKernelAvx512FSliding이라는 새로운 커널을 도입하여 이러한 문제점을 해결하고, stride 1, dilation 1, 그리고 커널 크기 3, 5, 7에 대해 기존 커널과 비트 단위 동일한 결과를 보장하면서도 성능을 대폭 향상시킵니다.

코드 변경사항 분석

이번 PR의 핵심은 새로운 sconv_depthwise_avx512f.cpp 파일의 추가와 관련 설정 파일들의 수정입니다. 기존 커널의 비효율을 개선하기 위해 Sliding Window 기법을 적용한 AVX-512 커널이 도입되었습니다.

1. MLAS 빌드 시스템 및 설정 변경 (`cmake/onnxruntime_mlas.cmake`, `include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h`, `onnxruntime/core/mlas/inc/mlas.h`)

새로운 AVX-512 커널을 빌드에 포함시키고, 사용자가 이를 제어할 수 있도록 옵션을 추가하는 작업이 이루어졌습니다.

빌드 시스템 통합:

diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake
index 996aa4281f044..7106edbe9aa5a 100644
--- a/cmake/onnxruntime_mlas.cmake
+++ b/cmake/onnxruntime_mlas.cmake
@@ -267,6 +267,7 @@ function(setup_mlas_source_for_windows)
       ${MLAS_SRC_DIR}/intrinsics/avx512/sconv_nchw_depthwise_multiplier_greater_than_1_avx512f.cpp
       ${MLAS_SRC_DIR}/linear_attention_kernel_avx512f.cpp
       ${MLAS_SRC_DIR}/intrinsics/avx512/reorder_avx512f.cpp
+      ${MLAS_SRC_DIR}/intrinsics/avx512/sconv_depthwise_avx512f.cpp
     )

     set_source_files_properties(${mlas_platform_srcs_avx512} PROPERTIES COMPILE_FLAGS "/arch:AVX512")
@@ -1011,6 +1012,7 @@ else()
           ${MLAS_SRC_DIR}/intrinsics/avx512/sconv_nchw_depthwise_multiplier_greater_than_1_avx512f.cpp
           ${MLAS_SRC_DIR}/linear_attention_kernel_avx512f.cpp
           ${MLAS_SRC_DIR}/intrinsics/avx512/reorder_avx512f.cpp
+          ${MLAS_SRC_DIR}/intrinsics/avx512/sconv_depthwise_avx512f.cpp
       )
       set_source_files_properties(${mlas_platform_srcs_avx512f} PROPERTIES COMPILE_FLAGS "-mavx512f")

Windows와 Linux 환경 모두에서 새로운 sconv_depthwise_avx512f.cpp 파일이 MLAS 빌드에 포함되도록 CMake 설정이 업데이트되었습니다. 이를 통해 AVX-512F 명령어셋을 지원하는 플랫폼에서 이 최적화된 커널을 사용할 수 있게 됩니다.

세션 옵션 추가:

diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h
index 977fd1c54571b..a1aa0465313a5 100644
--- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h
+++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h
@@ -602,6 +602,16 @@ static const char* const kOrtSessionOptionsMlasKleidiAiConvIgemmMaxWork = "mlas. 
 // This option exists for perf experimentation; the default may be retuned in future releases.
 static const char* const kOrtSessionOptionsMlasNchwcPointwiseConvMaxInputChannelBatch = "mlas.nchwc_pointwise_conv_max_input_channel_batch";

+// Selects the NCHWc depthwise convolution kernel on AVX-512 platforms. The sliding window kernel keeps
+// each input column of a kernel row in a register across the kernel columns and handles the padding
+// columns with masks. It supports stride 1, dilation 1 and kernel widths 3, 5 and 7 (other shapes use
+// the assembly kernel). Its results are bitwise identical to the assembly kernel, except that a NaN
+// result may carry a different NaN payload or sign.
+// Option values:
+// - "1": Use the sliding window kernel where it applies. [DEFAULT]
+// - "0": Always use the assembly kernel.
+static const char* const kOrtSessionOptionsMlasNchwcDepthwiseSliding = "mlas.nchwc_depthwise_sliding";
+
 // When converting DQ + MatMul -> MatMulNBits, the accuracy level of the MatMulNBits is controlled by this option.
 // Refer to MatMulNBits op schema for more details.
 // If not provided, default is 4.

mlas.nchwc_depthwise_sliding라는 새로운 세션 옵션이 추가되었습니다. 이 옵션은 기본값으로 1 (활성화)이 설정되어 Sliding Window 커널을 사용하도록 합니다. 사용자는 이 옵션을 0으로 설정하여 기존의 어셈블리 커널을 사용하도록 강제할 수 있습니다. 이는 성능 회귀 발생 시 롤백 메커니즘으로 활용될 수 있습니다.

MLAS 내부 구조체 업데이트:

diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h
index f650a4b58de68..e5b2882bce2b3 100644
--- a/onnxruntime/core/mlas/inc/mlas.h
+++ b/onnxruntime/core/mlas/inc/mlas.h
@@ -232,6 +232,7 @@ struct MLAS_BACKEND_KERNEL_SELECTOR_CONFIG {
     bool use_kleidiai = true; /**< Flag to use KleidiAI backend kernels if available */
     size_t kleidiai_conv_igemm_max_work = 0; /**< Optional SME IGEMM route threshold override; 0 uses default */
     size_t nchwc_pointwise_conv_max_input_channel_batch = 0; /**< Optional NCHWc pointwise conv input channel batch override; 0 uses default (128) */
+    bool nchwc_depthwise_sliding_kernel = true; /**< Use the sliding window AVX-512 NCHWc depthwise kernel (bitwise identical apart from NaN payloads) where it applies */

};

// ``` MLAS 내부의 커널 선택 로직에 사용되는 MLAS_BACKEND_KERNEL_SELECTOR_CONFIG 구조체에 nchwc_depthwise_sliding_kernel 필드가 추가되었습니다. 이 필드는 세션 옵션과 연동되어 내부적으로 Sliding Window 커널 사용 여부를 결정합니다.

2. Sliding Window Depthwise Convolution 커널 구현 (`intrinsics/avx512/sconv_depthwise_avx512f.cpp`)

이 파일은 이번 PR의 핵심으로, 새로운 Sliding Window AVX-512 커널을 정의합니다.

기존 커널의 문제점: 기존 MlasConvDepthwiseFloatKernelAvx512F는 다음과 같은 비효율을 가졌습니다:
- 6개의 누산기(accumulator)만 사용하여 FMA 연산의 잠재력을 충분히 활용하지 못함 (8개 필요).
- 커널의 각 탭(tap)마다 입력(input)을 다시 로드함.
- 패딩 영역을 스칼라 연산으로 처리하여 비효율 발생 (커널 시간의 16% 소요).
Sliding Window 기법: 새로운 MlasConvDepthwiseFloatKernelAvx512FSliding 커널은 다음과 같은 방식으로 개선되었습니다:
- 입력 재사용: 커널 행(kernel row)의 입력 열(input column)들을 레지스터에 한 번 로드한 후, 해당 커널 행의 모든 커널 열(kernel column)에서 재사용합니다. 이를 통해 반복적인 메모리 로드를 줄입니다.
- 패딩 처리: 패딩 영역은 마스크(mask)를 사용하여 처리합니다. masked loads와 merge-masked FMAs를 통해 경계 밖의 메모리 접근을 방지하고, 패딩 탭은 0과의 곱셈 대신 건너뛰게 됩니다. 이는 스칼라 처리 대비 효율적입니다.
- 누산기 활용: FMA 연산의 잠재력을 더 잘 활용하기 위해 더 많은 누산기를 사용합니다. (구체적인 개수는 MaximumOutputCount 매크로 참조)

코드 예시 (핵심 로직):

// ... (MLAS_DW_SLIDING_ARGS 구조체 정의) ...

template <size_t KernelWidth, size_t OutputCount, bool Masked>
MLAS_FORCEINLINE
void DepthwiseTile(
    const MLAS_DW_SLIDING_ARGS& Args,
    const float* InputRow,
    ptrdiff_t Column,
    float* Output
    )
{
    // ... (Accumulator 초기화)

    // ... (Mask 계산 - Masked=true일 경우)

    // ... (InputVector 로드 - Masked 로드 또는 일반 로드)
    // InputVector[j] = _mm512_maskz_loadu_ps(Valid[j], Address); // Masked load
    // InputVector[j] = _mm512_loadu_ps(Address); // Normal load

    // ... (FilterVector 로드)

    // ... (FMA 연산 - Masked FMA 또는 일반 FMA)
    // Accumulator[o] = _mm512_mask3_fmadd_ps(FilterVector, InputVector[o + kw], Accumulator[o], Valid[o + kw]); // Masked FMA
    // Accumulator[o] = _mm512_fmadd_ps(FilterVector, InputVector[o + kw], Accumulator[o]); // Normal FMA

    // ... (Accumulate, Bias, ReLU 처리)

    // ... (결과 저장)
    // _mm512_storeu_ps(Output + o * BlockSize, Accumulator[o]);
}

위 코드는 DepthwiseTile 함수의 핵심 부분을 보여줍니다. KernelWidth, OutputCount, Masked 등의 템플릿 파라미터를 사용하여 다양한 경우에 대한 최적화를 수행합니다. _mm512_maskz_loadu_ps와 _mm512_mask3_fmadd_ps와 같은 AVX-512 마스크 관련 내장 함수를 사용하여 패딩 처리를 효율적으로 수행합니다.

비트 단위 동일성 보장: 이 새로운 커널은 기존 어셈블리 커널과 비트 단위로 동일한 결과를 보장합니다 (NaN 페이로드나 부호가 다른 경우 제외). 이는 vfmadd231ps 또는 vfmadd132ps 명령어의 선택에 따라 NaN 전파 방식이 달라질 수 있기 때문입니다. 이는 정확성 측면에서 매우 중요한 부분입니다.

3. 테스트 및 벤치마크 업데이트

새로운 커널의 도입에 따라 관련 테스트 및 벤치마크 코드도 수정되었습니다.

테스트 스킵 조건 수정: 리뷰 댓글에서 지적된 것처럼, AVX/AVX2 x64 및 non-AMD64 NCHWc 플랫폼에서는 새로운 커널이 선택되지 않아 Sliding=0과 Sliding=1이 동일한 커널을 타임하는 문제가 있었습니다. 이를 해결하기 위해 테스트는 AVX-512F 지원 및 특정 블록 크기 조건에서만 실행되도록 수정되었습니다.
```
diff --git a/onnxruntime/test/mlas/bench/bench_nchwc_depthwise.cpp b/onnxruntime/test/mlas/bench/bench_nchwc_depthwise.cpp
# ... (수정된 부분 요약)
```
```
diff --git a/onnxruntime/test/mlas/unittest/test_nchwc_depthwise_sliding.cpp b/onnxruntime/test/mlas/unittest/test_nchwc_depthwise_sliding.cpp
# ... (수정된 부분 요약)
```
포인터 연산 관련 수정: 패딩 영역 처리 시 발생하는 포인터 연산의 정의되지 않은 동작(undefined behavior)을 방지하기 위해, 컬럼 오프셋을 별도로 추적하고 유효한지 확인한 후에 로드 주소를 계산하도록 수정되었습니다.
```
diff --git a/onnxruntime/core/mlas/lib/intrinsics/avx512/sconv_depthwise_avx512f.cpp b/onnxruntime/core/mlas/lib/intrinsics/avx512/sconv_depthwise_avx512f.cpp
# ... (수정된 부분 요약 - ptrdiff_t Column 사용 및 마스크 로직 개선)
```

왜 이게 좋은가?

성능 향상

이 PR의 가장 큰 장점은 명백한 성능 향상입니다. Zen 5 CPU에서 측정된 결과는 다음과 같습니다:

**MLAS 마이크로벤치마크 (단일 스레드, 동일 패딩):
- 7x7 커널: -55% ~ -64% 성능 향상 (FMA 피크의 81-87% 달성, 기존 31-46% 대비)
- 5x5 커널: -43% ~ -56% 성능 향상
- 3x3 커널: -10% ~ -47% 성능 향상 (L2 스트리밍 바운드 발생 시 약 45% 피크 성능)
**실제 모델 성능 (Sliding On vs Off):
- mobileclip: 1T -2.89%, 4T -2.71%, 8T -2.73%
- mv3_small: 1T -5.44%, 4T -3.30%
- mv3_large: 1T -2.09%, 4T -1.27%

이러한 성능 향상은 AVX-512 명령어셋의 SIMD(Single Instruction, Multiple Data) 능력을 최대한 활용하고, 메모리 로드 및 패딩 처리의 비효율을 제거함으로써 달성되었습니다.

일반적인 교훈

하드웨어 특화 최적화의 중요성: AVX-512와 같은 최신 명령어셋을 적극적으로 활용하는 것은 고성능 컴퓨팅에서 필수적입니다. 특히 컨볼루션과 같이 반복적인 연산에서는 그 효과가 극대화됩니다.
메모리 접근 패턴 최적화: 데이터를 한 번 로드하여 여러 번 재사용하는 'Sliding Window' 기법은 메모리 대역폭 병목 현상을 줄이는 데 효과적입니다. 이는 특히 GPU나 SIMD 유닛을 활용하는 CPU에서 중요합니다.
패딩 처리 효율화: 컨볼루션 연산에서 패딩 처리는 불가피하지만, 이를 효율적으로 (예: 마스크 연산 활용) 처리하는 것이 전체 성능에 큰 영향을 미칩니다.
성능 회귀 방지를 위한 옵션 제공: 새로운 최적화 기법 도입 시, 기존 커널로 롤백할 수 있는 옵션을 제공하는 것은 안정성을 높이는 좋은 전략입니다. 이는 다양한 하드웨어 환경에서의 예기치 못한 성능 저하를 방지하고, 문제 발생 시 원인 분석을 용이하게 합니다.
정확성 보장: 성능 향상과 더불어 비트 단위 동일성(NaN 예외)을 보장하는 것은 라이브러리의 신뢰성을 높이는 데 매우 중요합니다.

리뷰 피드백 반영

리뷰어들은 주로 테스트 커버리지와 코드의 정의되지 않은 동작(undefined behavior)에 대한 우려를 제기했습니다. PR 작성자는 이러한 피드백을 반영하여:

테스트가 실제 새로운 커널을 실행하는지 확인하고, 그렇지 않은 경우 테스트를 스킵하도록 수정했습니다.
포인터 연산의 정의되지 않은 동작을 방지하기 위해 컬럼 오프셋을 별도로 관리하고 마스크 로드를 안전하게 사용하도록 코드를 개선했습니다.
새로운 세션 옵션(mlas.nchwc_depthwise_sliding)에 대한 테스트 커버리지를 강화했습니다.

이러한 리뷰 과정은 코드의 품질과 안정성을 크게 향상시키는 데 기여했습니다.

결론

이번 PR은 ONNX Runtime의 MLAS 라이브러리에서 AVX-512 NCHWc Depthwise Convolution 연산의 성능을 크게 향상시키는 중요한 개선 사항입니다. Sliding Window 기법과 마스크 연산을 효과적으로 활용하여 기존 커널의 병목 현상을 해결하고, Zen 5 CPU에서 상당한 성능 향상을 달성했습니다. 또한, 롤백 옵션 제공 및 철저한 테스트를 통해 안정성과 신뢰성을 확보했습니다. 이러한 최적화는 딥러닝 모델의 추론 속도를 개선하는 데 크게 기여할 것입니다.

참고 자료

https://pytorch.org/docs/stable/generated/torch.compile.html

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer, Kimi-K3 어텐션의 성능 최적화: SM100/SM103 아키텍처를 위한 중간 크기 라우팅 개선
현재글 : [onnxruntime] ONNX Runtime MLAS, AVX-512 NCHWc Depthwise Convolution 최적화: Sliding Window 커널 도입
다음글 [flashinfer] [FlashInfer] Hopper SM90을 위한 MegaMoE 최적화: Combine Wire 및 GEMM 에필로그 개선

댓글

관련 포스트

PR Analysis 의 다른글