[onnxruntime] ONNX Runtime CPU Col2im 최적화: 메모리 접근 및 경계 검사 개선

PR 링크: microsoft/onnxruntime#32698 상태: Merged | 변경: +154 / -7

들어가며

딥러닝 모델의 추론 성능을 최적화할 때, 연산 집약적인 ConvTranspose 연산은 항상 주요 타겟이 됩니다. 특히 CPU 환경에서 ConvTranspose는 Col2im (Column to Image) 변환 과정을 거치는데, 이 과정에서의 메모리 접근 효율성은 전체 추론 속도에 큰 영향을 미칩니다. 최근 microsoft/onnxruntime 레포지토리에 병합된 이 PR은 Col2im 함수의 메모리 접근 패턴을 최적화하고, 불필요한 경계 검사(bounds checks)를 제거하여 특정 환경에서 최대 47%의 성능 향상을 달성했습니다.

코드 분석

1. Fast Path 최적화 (onnxruntime/core/util/math_cpu.cc)

가장 큰 변화는 dilation이 없고 padding이 없는 경우에 대한 'Fast Path'가 추가된 것입니다. 특히 2x2 커널과 stride 2를 사용하는 일반적인 경우를 위해 루프를 언롤링(unrolling)하고 메모리 접근을 최적화했습니다.

Before:

Set<float, CPUMathUtil>(narrow<ptrdiff_t>(hwc), 0, data_im, context);
// ... 기존의 범용적인 루프 수행

After:

if (kernel_h == 2 && kernel_w == 2 && stride_h == 2 && stride_w == 2 &&
    height % 2 == 0 && width % 2 == 0) {
  for (int64_t c = 0; c < channels; ++c) {
    for (int64_t h = 0; h < output_h; ++h) {
      const float* src = data_col + c * 4 * output_hw + h * output_w;
      float* dst = data_im + c * hw + h * 2 * width;
      for (int64_t w = 0; w < output_w; ++w) {
        dst[2 * w] = 0.0f + src[w];
        dst[2 * w + 1] = 0.0f + src[output_hw + w];
        dst[width + 2 * w] = 0.0f + src[2 * output_hw + w];
        dst[width + 2 * w + 1] = 0.0f + src[3 * output_hw + w];
      }
    }
  }
  return;
}

이 최적화는 별도의 Set 함수를 통한 0 초기화 과정을 생략하고, 인접한 출력 값을 한 번에 기록함으로써 캐시 효율성을 극대화합니다.

2. 경계 검사 제거 및 루프 효율화

일반적인 경우(Fallback)에서도 루프 내부의 if 문을 제거하여 분기 예측(branch prediction) 실패를 줄였습니다.

Before:

for (int64_t w = w_offset; src < src_we; src++, w += stride_w) {
  if (is_a_ge_zero_and_a_lt_b(w, width)) {
    dst[h + w] += *src;
  }
}

After:

const int64_t count = std::min(output_w - first_col, (width - 1 - first_w) / stride_w + 1);
for (int64_t col = 0; col < count; ++col) {
  dst[h + first_w + col * stride_w] += src[first_col + col];
}

루프 진입 전에 유효한 범위를 미리 계산(first_col, count)함으로써, 루프 내부에서 매번 수행하던 is_a_ge_zero_and_a_lt_b 검사를 완전히 제거했습니다.

왜 이게 좋은가

이번 최적화의 핵심은 '데이터 지역성(Data Locality) 향상'과 '분기 최소화'입니다.

1. 성능 수치: Intel Core Ultra 7 환경에서 4 스레드 기준, ConvTranspose 연산에서 최대 47%의 성능 향상을 보였습니다. 이는 단순히 알고리즘을 바꾼 것이 아니라, CPU가 데이터를 처리하는 방식을 하드웨어 친화적으로 변경했기 때문입니다.
2. 일반적 교훈:
   • Fast Path 설계: 특정 조건(예: stride 2, kernel 2)이 자주 발생한다면, 이를 위한 전용 코드를 작성하는 것이 범용 코드보다 훨씬 빠릅니다.
   • 루프 불변식 제거(Loop Invariant Code Motion): 루프 내부에서 매번 계산되는 경계 조건은 루프 밖으로 빼내는 것만으로도 성능을 크게 개선할 수 있습니다.
   • SIMD 친화적 구조: 루프 내부의 분기를 제거하면 컴파일러가 자동 벡터화(Auto-vectorization)를 적용하기 훨씬 유리해집니다.

이번 PR은 복잡한 라이브러리일수록 작은 루프 최적화가 전체 시스템 성능에 얼마나 큰 영향을 미칠 수 있는지 보여주는 좋은 사례입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html

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