CubeCL - CUDA, ROCm, WGPU를 위한 Rust 기반 GPU 커널

1 day ago 6

  • CubeCL은 Rust에서 GPU 커널을 작성할 수 있도록 해주는 고성능 멀티플랫폼 언어 확장
  • 함수, 제네릭, 구조체를 완벽하게 지원하며, 특성, 메서드, 타입 추론은 부분적으로 지원
  • WGPU, CUDA, ROCm 기반 런타임을 지원하며, SIMD 명령어를 활용한 최적화된 JIT CPU 런타임도 개발중
  • 자동 벡터화, 컴파일 타임 최적화(Comptime), 오토튠 기능을 제공하여 코드 성능을 극대화함
  • Rust의 프로시저 매크로 시스템을 활용하여 GPU 커널 코드를 파싱하고 확장하는 독특한 방식을 사용함

CubeCL 소개

  • CubeCL은 Rust 언어를 확장하여 GPU 커널을 직접 작성할 수 있도록 지원하는 고성능 컴퓨팅 언어 확장
  • 제로 코스트 추상화를 활용하여 효율적이고 유지보수하기 쉬운 커널 개발을 가능하게 함
  • 현재 함수, 제네릭, 구조체를 완벽 지원하고, 트레이트, 메서드, 타입 추론은 부분 지원함

간단한 사용 예시

  • #[cube] 어트리뷰트를 사용해 GPU에서 실행할 함수를 지정함
  • gelu_array와 gelu_scalar 함수는 SIMD 명령어를 자동으로 사용하여 성능을 향상시킴
  • launch_unchecked 함수를 사용하여 커널을 실행할 수 있음

지원하는 런타임

  • WGPU: Vulkan, Metal, DirectX, WebGPU를 통한 크로스 플랫폼 GPU 지원
  • CUDA: NVIDIA GPU 최적화 지원
  • ROCm/HIP: AMD GPU 지원 (개발 중)
  • JIT CPU 런타임: Cranelift 기반 SIMD 명령어를 활용한 최적화된 JIT CPU 런타임 개발중

CubeCL의 목표

  • 고성능 컴퓨팅 커널을 하드웨어에 구애받지 않고 작성할 수 있도록 지원함
  • 자동 벡터화, 컴타임, 자동 조정 기능을 통해 코드의 구성 가능성, 재사용성, 테스트 가능성, 유지보수성을 향상시킴
  • 메모리 관리 전략을 통해 높은 처리량을 위한 최적화된 버퍼 재사용을 제공함
  • 선형 대수 구성 요소를 개발하여 고성능 및 과학적 컴퓨팅 생태계를 구축하고자 함

CubeCL의 작동 방식

  • Rust의 프로시저 매크로 시스템을 활용하여 GPU 커널 코드를 파싱하고 확장하는 독특한 방식을 사용함
  • 컴타임: 원본 코드를 변환하지 않음으로써 컴파일 타임 최적화를 쉽게 통합할 수 있음
  • 자동 벡터화: CubeCL 함수의 입력을 벡터화하여 각 중간 변수의 벡터화 요소를 결정함
  • Rust 통합: 생성된 코드는 유효한 Rust 코드로 남아 특정 런타임에 대한 의존 없이 번들로 제공될 수 있음

디자인 및 특별 기능

  • CubeCL은 큐브를 중심으로 설계되었으며, 모든 컴퓨팅 API는 하드웨어에 매핑되어 3D 표현을 사용할 수 있음
  • 자동 벡터화: SIMD 명령어를 사용하여 성능을 향상시킬 수 있으며, 벡터화 요소를 커널 내에서 직접 접근 가능함
  • 컴타임: 컴파일러 IR을 런타임에 수정하여 최적화와 유연성을 제공함
  • 자동 조정: 런타임에 작은 벤치마크를 실행하여 최적의 커널과 설정을 선택함

주의사항 및 역사

  • CubeCL은 현재 알파 버전으로, 다듬어지지 않은 부분이 존재함
  • 원래 Burn 프로젝트의 WebGPU 백엔드로 시작되었으나
  • CUDA 런타임을 추가하며 범용 IR과 Rust 프론트엔드로 확장됨
  • CUDA와 WebGPU의 차이를 극복하기 위해 범용 컨셉을 개발함

Read Entire Article