[onnxruntime] ONNX Runtime의 ARM SVE i8mm QGEMM 최적화: 휴대용 머신 코드 전략

PR 링크: microsoft/onnxruntime#31146 상태: Merged | 변경: +8031 / -206

들어가며

최신 ARM 프로세서에서 제공하는 SVE(Scalable Vector Extension)와 I8MM(Integer Matrix Multiply) 확장은 행렬 연산 성능을 비약적으로 향상시킬 수 있는 강력한 도구입니다. 하지만 이를 활용한 커널을 작성할 때, 특정 컴파일러 플래그에 의존하거나 플랫폼별로 다른 바이너리를 생성해야 하는 문제는 배포와 유지보수의 큰 걸림돌이 됩니다. 본 PR은 svmmla 명령어를 사용하는 S8S8/U8S8 QGEMM 커널을 도입하면서, KleidiAI 스타일의 '휴대용 머신 코드(Portable Machine Code)' 전략을 통해 이 문제를 해결했습니다.

코드 분석

1. 휴대용 머신 코드 생성 및 빌드 시스템

핵심은 SVE intrinsics로 작성된 참조 구현을 gen_sve_asm.py 스크립트를 통해 기계어 코드(Hex-encoded assembly)로 변환하는 것입니다. 이를 통해 컴파일러가 SVE를 지원하지 않아도 어셈블러만으로 빌드가 가능합니다.

# cmake/CMakeLists.txt
set(ORT_SVE_ABI_FLAGS "-ffixed-x18")

특히 Windows ARM64 환경에서 x18 레지스터는 TEB(Thread Environment Block)를 가리키는 예약된 레지스터입니다. 이를 보호하기 위해 -ffixed-x18 플래그를 강제하여 커널이 해당 레지스터를 오염시키지 않도록 했습니다.

2. SVE i8mm 커널 구현

qgemm_mmla_sve_asm.S 파일은 5,000줄이 넘는 KAI_ASM_INST 매크로로 구성되어 있습니다. 이는 어떤 환경에서도 동일한 바이너리를 보장합니다.

// onnxruntime/core/mlas/lib/aarch64/qgemm_mmla_sve_asm.S
KAI_ASM_GLOBAL(MlasGemmS8S8KernelSmmlaSveImpl)
KAI_ASM_FUNCTION_LABEL(MlasGemmS8S8KernelSmmlaSveImpl)
    KAI_ASM_BTI_C
    KAI_ASM_INST(0x910003eb)  // mov x11, sp
    KAI_ASM_INST(0x0423e3ec)  // cntb x12, all, mul #4

3. M=1 최적화 경로

GEMV 형태의 연산(M=1)에서는 기존 NEON 커널과 동등한 성능을 내기 위해 별도의 최적화 경로를 추가했습니다. 불필요한 메모리 복사와 연산을 제거하여 성능을 개선했습니다.

왜 이게 좋은가

1. 성능 향상: Cortex-X925 환경에서 4스레드 기준 약 1.17배의 성능 향상을 기록했습니다. 특히 M=1 경로에서 기존 대비 parity 수준의 성능을 확보했습니다.
2. 이식성: KAI_ASM_INST를 사용한 기계어 코드는 GNU 어셈블러와 MSVC armasm64 모두에서 컴파일 가능하며, SVE 툴체인 유무와 무관하게 동작합니다.
3. 안정성: x18 레지스터 보호와 같은 ABI 준수는 플랫폼 간 호환성을 보장하며, 33,213개의 테스트 케이스를 통해 NEON 커널과 비트 단위로 동일한 결과를 검증했습니다.

교훈

• ABI 준수의 중요성: 플랫폼 예약 레지스터(x18)를 건드리는 것은 치명적인 런타임 오류를 유발합니다. 컴파일러 플래그와 생성기 단계에서의 검증을 병행하는 'Belt-and-braces' 접근이 필수적입니다.
• 코드 생성 전략: 성능이 중요한 커널은 고수준 언어로 작성하고, 배포 시에는 기계어로 고정(Freeze)하는 방식은 성능과 이식성이라는 두 마리 토끼를 잡는 효과적인 패턴입니다.

참고 자료

• Arm KleidiAI — 휴대용 머신 코드 스타일의 영감이 된 라이브러리
• SVE Intrinsics — SVE 명령어 활용을 위한 공식 문서

참고 자료

• https://developer.arm.com/documentation/101028/0012/SVE-intrinsics

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