[onnxruntime] ONNX Runtime의 RISC-V RVV 커널 최적화: 추론 성능 극대화

PR 링크: microsoft/onnxruntime#32540 상태: Merged | 변경: +1905 / -404

들어가며

ONNX Runtime의 MLAS(Microsoft Linear Algebra Subprograms) 라이브러리는 다양한 하드웨어 아키텍처에서 고성능 연산을 제공합니다. 이번 PR은 RISC-V 64비트 환경에서 RVV(RISC-V Vector) 확장을 활용하여 FP32 추론 성능을 대폭 개선했습니다. 특히 Softmax, NCHWc Convolution, 그리고 Activation 루틴을 최적화하여 기존 대비 최대 3배 이상의 성능 향상을 달성했습니다.

코드 분석

1. Activation 루틴의 벡터화 (conv_activation_kernel_rvv.cpp)

기존에는 Activation 연산이 스칼라 방식으로 처리되었으나, 이제 RVV를 사용하여 벡터 단위로 연산을 수행합니다. __riscv_vsetvl_e32m4를 통해 벡터 길이를 설정하고, vfloat32m4_t 타입을 사용하여 4배의 벡터 레지스터 그룹(LMUL 4)을 활용합니다.

// Before: 스칼라 기반 루프 (생략)

// After: RVV 기반 루프
while (n > 0) {
    const size_t vl = __riscv_vsetvl_e32m4(n);
    vfloat32m4_t Vector = __riscv_vle32_v_f32m4(buffer, vl);
    if (AddBias) {
        Vector = __riscv_vfadd_vf_f32m4(Vector, bias, vl);
    }
    __riscv_vse32_v_f32m4(buffer, ActivationFunction.Activate(Vector, vl), vl);
    buffer += vl;
    n -= vl;
}

2. Depthwise Convolution 최적화 (convolve.cpp)

기존의 im2col 방식은 오버헤드가 컸습니다. 이번 변경으로 특정 조건(커널 크기 16 이하 등)에서 RVV 전용 커널을 직접 호출하도록 라우팅을 추가했습니다.

+        if (GetMlasPlatform().ConvNchwFloatKernel != nullptr
+                && Dimensions == 2
+                && Parameters->FilterCount == 1 && Parameters->InputChannels == 1
+                && Parameters->KernelShape[1] <= kDepthwiseGeneralMaxKernelWidth
+                && ... ) {
+            *WorkingBufferSize = Parameters->InputShape[1] + 2;
+            Parameters->Algorithm = MlasConvAlgorithmDepthwise;
+            return;
+        }

왜 이게 좋은가

이번 최적화는 단순히 연산 속도만 높인 것이 아니라, 메모리 접근 패턴과 레지스터 활용도를 극대화했습니다.

1. LMUL 활용: LMUL 4를 사용하여 벡터 레지스터 그룹을 효율적으로 구성함으로써, 더 많은 데이터를 한 번에 처리하고 명령어 파이프라인의 효율을 높였습니다.
2. 독립적 누적기(Independent Accumulators): NCHWc Convolution에서 4개의 인접 출력을 독립적인 누적기로 계산하여 데이터 의존성을 줄이고 처리량을 개선했습니다.
3. 성능 수치: SpacemiT X100 보드 기준, ResNet50 모델에서 약 3.01배, MobileNetV2에서 3.09배의 속도 향상을 기록했습니다.

교훈: 하드웨어 가속기(RVV)를 사용할 때는 단순히 명령어를 벡터화하는 것을 넘어, 데이터 레이아웃(NCHWc)과 레지스터 그룹화(LMUL)를 하드웨어 특성에 맞게 조정하는 것이 성능의 핵심입니다.

리뷰어 피드백 반영

리뷰 과정에서 ActivationRoutine이 MLAS_USE_RVV가 꺼진 환경에서도 컴파일될 수 있도록 하는 빌드 가드 문제가 지적되었습니다. 또한, 벡터 연산 결과의 정확성을 검증하기 위해 NaN/±0 처리를 고려한 단위 테스트 강화가 제안되었습니다. 이는 성능뿐만 아니라 안정성까지 고려한 엔지니어링의 중요성을 보여줍니다.

참고 자료

• https://github.com/riscv/riscv-v-spec
• https://onnxruntime.ai/docs/performance/model-optimizations.html

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