[onnxruntime] ONNX Runtime WebGPU: Windows D3D12를 활용한 모델 로딩 가속화

PR 링크: microsoft/onnxruntime#33055 상태: Merged | 변경: +4579 / -58

들어가며

대규모 ONNX 모델을 로딩할 때, 외부 파일에 저장된 가중치를 처리하는 방식은 성능의 병목이 되곤 합니다. 기존의 WebGPU 로딩 경로는 가중치를 CPU 메모리에 먼저 로드한 뒤 GPU로 복사하는 방식을 취하는데, 이는 모델 크기가 커질수록 CPU 메모리 압박을 가중시키고 세션 생성 시간을 지연시킵니다. 이번 Microsoft/onnxruntime의 업데이트는 Windows 환경에서 D3D12를 활용하여 가중치를 CPU 메모리 거침없이 GPU로 직접 스트리밍하는 가속 경로를 도입했습니다.

코드 분석

1. D3D12 파일 로더 도입 (onnxruntime_d3d12_file_loader.cmake)

새로운 로더인 D3D12FileBufferLoader는 플랫폼 독립적인 로직과 분리되어 Windows 전용으로 구현되었습니다. 이 로더는 12개의 16MiB D3D12 업로드 버퍼 풀을 사용하여 파일 읽기와 GPU 복사를 오버랩시킵니다.

// onnxruntime/core/platform/windows/d3d12_file_loader/d3d12_file_buffer_loader.h
class D3D12FileBufferLoader {
  // 16 MiB 버퍼 풀 관리 및 비동기 파일 읽기/복사 로직
  // ...
};

2. 세션 옵션 확장 (js/common/lib/inference-session.ts)

사용자는 weightLoadAcceleration 옵션을 통해 이 기능을 제어할 수 있습니다. preferred 모드를 사용하면 가속이 가능한 환경에서 자동으로 최적화된 경로를 선택합니다.

// Before: 옵션 없음
// After:
weightLoadAcceleration?: 'off' | 'preferred' | 'required';

3. WebGPU 컨텍스트 통합 (onnxruntime/core/providers/webgpu/d3d12_external_data_loader.cc)

Dawn(WebGPU 구현체)과 D3D12 리소스를 연결하기 위해 SharedBufferMemoryD3D12ResourceDescriptor를 사용하여 GPU 리소스를 직접 임포트합니다.

// D3D12 리소스를 Dawn으로 임포트하는 핵심 로직
auto descriptor = wgpu::SharedBufferMemoryD3D12ResourceDescriptor{};
descriptor.resource = tensor.resource.Get();

왜 이게 좋은가

이 최적화의 핵심은 'Zero-Copy'에 가까운 스트리밍입니다. 기존 방식은 파일 -> CPU 메모리 -> GPU 메모리의 2단계 복사를 거쳤으나, 이제는 파일 -> D3D12 업로드 버퍼 -> GPU 로컬 리소스로 직접 이동합니다.

성능 개선 수치

• NVIDIA GeForce RTX 4070: 기존 1.562s 대비 0.715s로 54.2% 단축
• Intel Graphics: 기존 1.420s 대비 0.669s로 52.9% 단축

일반적 교훈

1. 메모리 경계 설정: 대규모 데이터를 처리할 때 전체를 메모리에 올리지 않고 고정된 크기의 버퍼 풀(16MiB x 12)을 사용하여 메모리 사용량을 예측 가능하게 유지해야 합니다.
2. I/O와 GPU 작업의 오버랩: Windows 이벤트와 D3D12 펜스를 활용하여 파일 I/O가 진행되는 동안 GPU 복사 작업을 병렬로 수행함으로써 유휴 시간을 최소화할 수 있습니다.
3. 플랫폼 특화 최적화: 범용적인 추상화도 중요하지만, 특정 OS(Windows)와 그래픽 API(D3D12)의 기능을 활용하면 성능을 극적으로 끌어올릴 수 있습니다.

리뷰어 피드백 반영

리뷰 과정에서 std::async 사용에 따른 스레드 생성 오버헤드나, FILE_FLAG_NO_BUFFERING 사용 시의 정렬 문제 등 세밀한 기술적 논의가 있었습니다. 특히 D3D12FileBufferLoader의 파일 경로 처리 시 대소문자 구분 문제나, 에러 발생 시의 10ms 폴링 대기 시간 등이 지적되었으며, 이는 향후 더 정교한 이벤트 기반 처리로 개선될 여지를 남겨두었습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://learn.microsoft.com/en-us/windows/win32/direct3d12/direct3d-12-programming-guide

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