본문으로 건너뛰기

[onnxruntime] ONNX Runtime CPU 성능 향상: YOLOX 모델 최적화 분석

PR 링크: microsoft/onnxruntime#32335 상태: Merged | 변경: +694 / -92

들어가며

최근 Microsoft의 ONNX Runtime 레포지토리에는 YOLOX 모델의 CPU 추론 성능을 크게 향상시키는 PR이 올라왔습니다. 이 PR은 두 가지 주요 최적화 기법을 적용했습니다. 첫째, Slice와 Concat 연산을 SpaceToDepth 연산으로 대체하여 연산 오버헤드를 줄였습니다. 둘째, Convolution 연산의 스레드별 작업 분배를 개선하여 멀티스레딩 성능을 향상시켰습니다. 이 글에서는 해당 PR의 코드 변경 사항을 자세히 살펴보고, 이러한 변경이 왜 성능 향상으로 이어지는지, 그리고 어떤 일반적인 교훈을 얻을 수 있는지 분석해 보겠습니다.

이 PR은 특히 YOLOX-Tiny 모델에서 2스레드 사용 시 약 16.8%, 4스레드 사용 시 약 15.4%의 성능 향상을 보였으며, YOLOX-Nano 모델에서도 2스레드 및 4스레드에서 각각 약 8%의 성능 향상을 기록했습니다. 이러한 주목할 만한 성능 개선의 비결을 코드 레벨에서 파헤쳐 보겠습니다.

코드 분석

이번 PR은 크게 두 부분으로 나누어 볼 수 있습니다. 첫 번째는 Slice와 Concat 연산을 SpaceToDepth 연산으로 융합하는 최적화이고, 두 번째는 MLAS(Microsoft Linear Algebra Subprograms) 라이브러리 내의 Convolution 연산 스케줄링 개선입니다.

1. Slice + Concat을 SpaceToDepth로 융합 (slice_concat_to_space_to_depth_fusion.cc)

이 변경은 ONNX Runtime의 그래프 최적화 단계에서 이루어집니다. 특정 패턴, 즉 여러 개의 Slice 연산과 하나의 Concat 연산이 연속적으로 사용되는 경우, 이를 단일 SpaceToDepth 연산으로 대체하여 연산 그래프를 단순화하고 메모리 접근 오버헤드를 줄이는 것을 목표로 합니다.

기존에는 다음과 같은 연산들이 사용되었을 수 있습니다:

// (가상 코드 예시)
input -> Slice(axis=1, start=0, end=H) -> Concat(axis=1, [slice1, slice2, ...]) -> ...

이 PR은 slice_concat_to_space_to_depth_fusion.cc 파일에서 이러한 패턴을 탐지하고 SpaceToDepth 연산으로 대체하는 로직을 추가하거나 개선했습니다. SpaceToDepth 연산은 입력 텐서의 공간 차원(예: 높이, 너비)을 채널 차원으로 재배열하는 연산으로, 특정 형태의 데이터 재구성에 효율적입니다. 여러 Slice와 Concat을 사용하는 것은 종종 SpaceToDepth가 수행하는 작업과 유사하지만, 별도의 연산으로 처리될 경우 추가적인 메모리 복사 및 연산 오버헤드가 발생합니다.

이 융합을 통해 불필요한 중간 텐서 생성을 피하고, 연산 그래프의 복잡성을 줄여 전반적인 추론 속도를 향상시킬 수 있습니다. 특히 YOLOX와 같이 공간적 특징을 채널 차원으로 재구성하는 연산이 빈번하게 사용되는 모델에서 효과적입니다.

2. MLAS Convolution 스레딩 개선 (mlasi.h, snchwc.cpp)

이 부분은 MLAS 라이브러리 내부의 Convolution 연산 성능을 최적화하는 데 중점을 둡니다. 특히 멀티스레딩 환경에서 스레드 간의 작업 불균형을 해소하여 CPU 활용률을 높이는 것이 목표입니다.

a. MlasNchwcChooseFilterSetSize 함수 추가 (mlasi.h)

--- a/onnxruntime/core/mlas/lib/mlasi.h
+++ b/onnxruntime/core/mlas/lib/mlasi.h
@@ -2113,6 +2113,54 @@ MlasNchwcCostToWorkIndex(
     return (BatchGroup * FilterSetCount + Set) * OutputHeight + Row;
 }
 
+// 
+// Picks the filter set size for the NCHWc grouped convolution algorithms.
+// Normally MaximumFilterSetSize, but when that would not produce enough work
+// units (BatchCount * GroupCount * FilterSetCount * OutputHeight) to give
+// every thread at least Target units, the size is halved -- repeatedly, down
+// to 1 -- so the operation splits more finely.
+//
+// Pulled out of MLAS_NCHWC_GROUPED_CONV_ALGORITHM::ChooseFilterSetSize as a
+// pure function of its inputs (no environment variable read, no object
+// state) so the halving dispatch can be exercised directly by unit tests
+// across a range of shapes, including ones with ragged block counts.
+//
+
+inline
+size_t
+MlasNchwcChooseFilterSetSize(
+    size_t MaximumFilterSetSize,
+    size_t Target,
+    ptrdiff_t tids,
+    size_t BlockSize,
+    size_t OutputChannels,
+    size_t BatchCount,
+    size_t GroupCount,
+    size_t OutputHeight
+    )
+{
+    if (Target == 0 || tids <= 1) {
+        return MaximumFilterSetSize;
+    }
+
+    const size_t Wanted = static_cast<size_t>(tids) * Target;
+    const size_t Blocks = OutputChannels / BlockSize;
+
+    size_t Size = MaximumFilterSetSize;
+
+    while (Size > 1) {
+        const size_t Sets = (Blocks + Size - 1) / Size;
+
+        if (BatchCount * GroupCount * Sets * OutputHeight >= Wanted) {
+            break;
+        }
+
+        Size /= 2;
+    }
+
+    return Size;
+}
+
 //
 // Define the minimum floating point value (and its bit value equivalent) that
 // has no fractional bits. This number can be used for fast rounding of floating

새로 추가된 MlasNchwcChooseFilterSetSize 함수는 Convolution 연산에서 필터 세트(Filter Set)의 크기를 결정하는 로직을 담당합니다. 기존에는 고정된 FilterSetSize (예: 4)를 사용했지만, 이 함수는 입력 크기, 스레드 수(tids), 목표 작업 단위 수(Target) 등을 고려하여 MaximumFilterSetSize를 동적으로 조절합니다. 만약 현재 설정된 FilterSetSize로는 각 스레드에 충분한 작업량을 할당할 수 없다면, FilterSetSize를 절반으로 줄여 더 작은 단위로 작업을 분할합니다. 이는 특히 출력 높이가 낮거나 채널 수가 적은 경우에 발생할 수 있는 스레드 유휴 시간을 줄이는 데 기여합니다.

b. MLAS_NCHWC_GROUPED_CONV_ALGORITHM 클래스 수정 (snchwc.cpp)

--- a/onnxruntime/core/mlas/lib/snchwc.cpp
+++ b/onnxruntime/core/mlas/lib/snchwc.cpp
@@ -525,10 +525,39 @@ struct MLAS_NCHWC_GROUPED_CONV_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM
     // reused for a given set of input inside the kernel.
     //
 
-    static constexpr size_t FilterSetSize = 4;
+    static constexpr size_t MaximumFilterSetSize = 4;
 
+
+    //
+    // Number of output channel blocks handled by one kernel invocation. Larger
+    // is cheaper per multiply-accumulate because a broadcast input is reused
+    // across every filter block, but it also produces fewer, coarser work units.
+    //
+
+    const size_t FilterSetSize;
     const size_t FilterSetCount;
 
+    //
+    // Distribute the output channel blocks evenly over the filter sets.
+    //
+    // Assigning FilterSetSize blocks to every set leaves the remainder in the
+    // final set, so that set can be up to FilterSetSize times cheaper than its
+    // peers. Work units are partitioned by count rather than by cost, and the
+    // output line index varies fastest, so a whole ragged filter set is handed
+    // to a different thread than the full ones and that thread finishes early.
+    //
+    // Spreading the blocks evenly costs nothing: the number of filter sets is
+    // unchanged, so the kernel still broadcasts each input the same number of
+    // times and issues the same number of multiply-accumulates. Only the
+    // per-thread balance changes.
+    //
+    // Example, 96 channels with a block size of 16: 6 blocks previously split
+    // as 4 + 2 and now split as 3 + 3.
+    //
+
+    const bool BalancedFilterSets;
+    const size_t FilterSetBase;
+    const size_t FilterSetRemainder;
+
     //
     // Stores the current output line, filter cluster, and group that this thread
     // is operating on.
@@ -540,15 +669,63 @@ struct MLAS_NCHWC_GROUPED_CONV_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM
     size_t WorkRemaining;
     size_t FilterCount;
 
+    //
+    // Picks the filter set size. Normally the maximum, but when the resulting
+    // work unit count would not give every thread a reasonable number of units
+    // the size is halved so the operation splits more finely.
+    //
+
+    size_t ChooseFilterSetSize(void) const
+    {
+        //
+        // MlasNchwcGetBlockSize() returns one rather than zero on platforms
+        // without NCHWc support, precisely so this division is always safe.
+        //
+
+        assert(BlockSize > 0);
+
+        return MlasNchwcChooseFilterSetSize(MaximumFilterSetSize, MlasNchwcFilterSetTarget(),
+                                            WorkBlock->tids, BlockSize, OutputChannels,
+                                            BatchCount, GroupCount, OutputHeight);
+    }
+
     MLAS_NCHWC_GROUPED_CONV_ALGORITHM(const MLAS_NCHWC_CONV_WORK_BLOCK* WorkBlock) :
         MLAS_NCHWC_CONV_ALGORITHM(WorkBlock),
-        FilterSetCount((OutputChannels + (BlockSize * FilterSetSize) - 1) / (BlockSize * FilterSetSize))
+        FilterSetSize(ChooseFilterSetSize()),
+        FilterSetCount((OutputChannels + (BlockSize * FilterSetSize) - 1) / (BlockSize * FilterSetSize)),
+        BalancedFilterSets(MlasNchwcUseBalancedFilterSets()),
+        FilterSetBase((OutputChannels / BlockSize) / FilterSetCount),
+        FilterSetRemainder((OutputChannels / BlockSize) % FilterSetCount)
     {
     }
 
+    //
+    // Returns the index of the first output channel block owned by a filter set.
+    //
+
+    size_t FilterSetBlockOffset(size_t Set) const
+    {
+        if (!BalancedFilterSets) {
+            return Set * FilterSetSize;
+        }
+
+        return Set * FilterSetBase + std::min(Set, FilterSetRemainder);
+    }
+
     void ComputeFilterCount(void)
     {
-        FilterCount = std::min(FilterSetSize, (OutputChannels / BlockSize) - FilterSet * FilterSetSize);
+        if (!BalancedFilterSets) {
+            FilterCount = std::min(FilterSetSize, (OutputChannels / BlockSize) - FilterSet * FilterSetSize);
+            return;
+        }
+
+        //
+        // FilterSetCount is the ceiling of the block count over FilterSetSize,
+        // so FilterSetBase never exceeds FilterSetSize and the kernel dispatch
+        // table is always in range.
+        //
+
+        FilterCount = FilterSetBase + ((FilterSet < FilterSetRemainder) ? 1 : 0);
     }
 
     void SeekToWork(size_t WorkIndex)
@@ -574,14 +738,16 @@ struct MLAS_NCHWC_GROUPED_CONV_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM
         Input += BatchGroup * InputChannels * InputSize;
 
         Output += BatchGroup * OutputChannels * OutputSize;
-        Output += BlockSize * FilterSet * FilterSetSize * OutputSize;
+        const size_t FilterBlockOffset = FilterSetBlockOffset(FilterSet);
+
+        Output += BlockSize * FilterBlockOffset * OutputSize;
 
         Filter += Group * OutputChannels * InputChannels * KernelSize;
-        Filter += BlockSize * FilterSet * FilterSetSize * InputChannels * KernelSize;
+        Filter += BlockSize * FilterBlockOffset * InputChannels * KernelSize;
 
         if (Bias != nullptr) {
             Bias += Group * OutputChannels;
-            Bias += BlockSize * FilterSet * FilterSetSize;
+            Bias += BlockSize * FilterBlockOffset;
         }
 
         //
@@ -674,6 +840,12 @@ struct MLAS_NCHWC_GROUPED_CONV_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM
 
     void PrepareWorkWeighted(ptrdiff_t Index)
     {
+        // Balanced sets are all equal size, so the uniform partitioner is exact.
+        if (BalancedFilterSets) {
+            PrepareWork(Index);
+            return;
+        }
+
         const size_t TotalBlockedFilters = OutputChannels / BlockSize;
         const size_t LastSetFilterCount =
             TotalBlockedFilters - (FilterSetCount - 1) * FilterSetSize;
@@ -705,7 +877,7 @@ struct MLAS_NCHWC_GROUPED_CONV_ALGORITHM : MLAS_NCHWC_CONV_ALGORITHM
     }
 };
 
-constexpr size_t MLAS_NCHWC_GROUPED_CONV_ALGORITHM::FilterSetSize;
+constexpr size_t MLAS_NCHWC_GROUPED_CONV_ALGORITHM::MaximumFilterSetSize;
 
 //
 // Implementation of the direct convolution algorithm where the input buffer is

MLAS_NCHWC_GROUPED_CONV_ALGORITHM 클래스는 MLAS 라이브러리에서 NCHWc 포맷의 그룹화된 Convolution 연산을 처리하는 핵심 로직을 담고 있습니다. 이 클래스의 생성자(MLAS_NCHWC_GROUPED_CONV_ALGORITHM)는 이제 MaximumFilterSetSize와 새로 도입된 MlasNchwcChooseFilterSetSize 함수를 사용하여 동적으로 FilterSetSize를 결정합니다. 또한, BalancedFilterSets, FilterSetBase, FilterSetRemainder와 같은 멤버 변수가 추가되었습니다. 이들은 출력 채널 블록을 각 필터 세트에 균등하게 분배하기 위한 로직을 구현합니다. 이전에는 마지막 필터 세트에 남은 블록들이 몰려 있어 작업 불균형을 야기할 수 있었지만, 이제는 블록들을 최대한 고르게 분산시켜 각 스레드가 비슷한 양의 작업을 수행하도록 합니다.

ComputeFilterCount 함수는 BalancedFilterSets 플래그에 따라 필터 개수를 계산하는 방식을 분기합니다. SeekToWork 함수에서는 FilterSetBlockOffset 함수를 호출하여 현재 처리할 필터 블록의 시작 오프셋을 계산합니다. 이 오프셋 계산 역시 BalancedFilterSets 여부에 따라 달라지며, 균등 분배 로직을 반영합니다.

PrepareWorkWeighted 함수는 스레드별 작업 준비 로직을 담당하는데, BalancedFilterSets가 활성화된 경우 PrepareWork 함수를 직접 호출하여 균등 분배된 작업량을 사용하도록 수정되었습니다. 이는 기존의 FilterSetSize 기반의 비용 모델 대신, 실제 균등 분배된 작업량을 기반으로 스레드를 할당하기 위함입니다.

이러한 변경들은 Convolution 연산 내부에서 스레드들이 수행하는 작업량을 더 균등하게 만들어, 특정 스레드가 병목 현상을 일으키는 것을 방지하고 전체적인 연산 속도를 향상시킵니다.

왜 이게 좋은가?

이 PR의 최적화는 다음과 같은 이유로 훌륭합니다.

  1. 연산 오버헤드 감소: Slice + Concat 패턴을 SpaceToDepth로 융합함으로써, 불필요한 중간 데이터 생성 및 복사를 제거했습니다. 이는 메모리 대역폭 사용량을 줄이고, 연산 그래프를 단순화하여 전반적인 추론 시간을 단축시킵니다. 특히 모델의 복잡성이 높아질수록 이러한 융합의 효과는 더욱 커집니다.
  2. 멀티스레딩 효율성 극대화: MLAS의 Convolution 스레딩 개선은 CPU 코어를 최대한 활용하도록 설계되었습니다. MlasNchwcChooseFilterSetSize 함수를 통한 동적 필터 세트 크기 조절과 BalancedFilterSets 옵션을 통한 작업 균등 분배는 다음과 같은 이점을 제공합니다:
    • 스레드 유휴 시간 감소: 작업량이 적은 경우에도 스레드들이 효율적으로 작업을 분배받아 유휴 상태를 최소화합니다.
    • 병목 현상 완화: 특정 스레드에 과도한 작업이 몰리는 것을 방지하여 전체 처리량을 높입니다.
    • 일관된 성능: 다양한 입력 크기와 스레드 수 조합에서도 안정적인 성능을 기대할 수 있습니다.
  3. 실질적인 성능 향상: PR 설명에 명시된 YOLOX 모델에서의 구체적인 성능 향상 수치(YOLOX-Tiny 16.8% @ 2T, 15.4% @ 4T)는 이러한 최적화가 실제 환경에서 매우 효과적임을 입증합니다. 이는 단순히 이론적인 개선이 아니라, 실제 애플리케이션에서 체감할 수 있는 성능 향상입니다.

일반적인 교훈:

  • 패턴 기반 융합의 중요성: 반복적으로 나타나는 연산 패턴을 찾아 더 효율적인 단일 연산으로 대체하는 것은 그래프 최적화의 핵심입니다. 이는 연산 오버헤드를 줄이고 메모리 접근을 최적화하는 데 매우 효과적입니다.
  • 동적 스케줄링 및 작업 분배: 멀티스레딩 환경에서는 고정된 방식보다는 입력 데이터의 특성이나 시스템 환경에 맞춰 동적으로 스레드별 작업량을 조절하고 균등하게 분배하는 것이 성능 향상의 열쇠입니다. 특히 CPU 집약적인 연산에서 이러한 최적화는 필수적입니다.
  • MLAS와 같은 저수준 라이브러리 최적화의 영향: ONNX Runtime과 같은 고수준 프레임워크는 MLAS와 같은 저수준 최적화 라이브러리의 성능 개선으로부터 큰 혜택을 받습니다. CPU 성능 최적화는 이러한 라이브러리 레벨에서의 세심한 튜닝에서 시작됩니다.

리뷰 피드백 반영

리뷰 과정에서 몇 가지 중요한 피드백이 있었습니다:

  • CI 빌드 및 Copilot 주석: 초기 버전에서는 CI 빌드 오류와 Copilot이 제안한 주석 수정이 필요했습니다. 이는 코드 품질과 안정성을 높이는 데 필수적인 과정입니다.
  • PR 분리 제안: SpaceToDepth 융합과 Convolution 성능 개선을 별도의 PR로 분리하자는 제안이 있었으나, 두 변경 사항이 동일한 YOLOX CPU 성능 문제를 해결하며 함께 적용될 때 전체 성능 향상을 명확히 보여줄 수 있다는 이유로 병합되었습니다. 이는 특정 성능 문제에 대한 여러 최적화가 시너지를 낼 때, 이를 통합하여 보여주는 것이 더 효과적일 수 있음을 시사합니다.
  • MLAS 독립 빌드 호환성: ORT_ENFORCE와 같은 ONNX Runtime 전용 매크로가 MLAS 독립 빌드 환경에서 문제를 일으킬 수 있다는 지적이 있었습니다. 이에 대한 수정으로 assert나 MLAS 자체에서 안전한 검증 방식을 사용하도록 변경되었습니다. 이는 라이브러리의 재사용성과 호환성을 높이는 중요한 개선입니다.
  • 테스트 커버리지: 새로운 최적화 로직, 특히 동적 필터 세트 크기 조절 및 균등 분배 로직에 대한 테스트 커버리지를 강화하라는 피드백이 있었습니다. MLAS_NCHWC_FSS_TARGET 환경 변수를 활용하여 다양한 시나리오를 테스트하고, 불완전했던 테스트 코드를 수정하여 새로운 최적화 경로가 제대로 동작하는지 검증하도록 개선되었습니다.

이러한 리뷰 피드백들은 코드의 견고성, 테스트 커버리지, 그리고 다양한 빌드 환경에서의 호환성을 높이는 데 크게 기여했습니다.

References

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글