tracel-ai/cubecl

해결하는 문제

CubeCL은 다양한 하드웨어에서 고성능 컴퓨트 커널을 작성하는 어려움을 해결합니다. 기존에는 개발자들이 NVIDIA용 CUDA, AMD용 HIP, Apple용 Metal 등의 GPU와 CPU용 별도의 코드를 작성해야 했거나, 최고 성능을 희생할 수 있는 고수준 추상화를 사용해야 했습니다. CubeCL은 Rust로 단일 커널을 작성하면 여러 백엔드에서 최적화된 방식으로 실행되도록 변환할 수 있게 해줍니다.

작동 방식

CubeCL은 Rust 언어 확장 및 즉시 컴파일러(JIT)로 작동합니다. #[cube] 매크로를 사용해 Rust 함수를 중간 표현(IR)인 pliron으로 재작성한 후, PTX, MSL, SPIR-V, WGSL과 같은 플랫폼별 언어로 컴파일합니다.

하드웨어 병렬성을 네 가지 수준으로 정리합니다:

• Unit: 개별 레지스터(스레드/호출).
• Plane: 동기 실행 명령(워프/서브그룹).
• Cube: 공유 메모리(블록/워크그룹).
• Hypercube: 여러 큐브에 걸친 글로벌 메모리.

대상 사용자

저수준 컴퓨트 커널, 과학 계산 라이브러리, 딥러닝 프레임워크(Burn 등)를 개발하는 개발자들을 위한 것입니다. 다양한 GPU/CPU 벤더용 별도의 코드베이스를 유지하지 않고도 최대한의 하드웨어 성능을 얻고 싶은 분들에게 적합합니다.

주요 특징

• 단일 코드베이스: Rust로 커널을 작성하고 CUDA, HIP, Metal, Vulkan, WebGPU, CPU에 배포 가능.
• 자동 벡터화: Vector<F, N> 타입을 사용해 폭 N을 실행 시점에 결정하고, 네이티브 SIMD 명령어로 하강.
• Comptime: 첫 번째 컴파일 시 컴파일러 IR을 수정해 특정 하드웨어 명령어, 형태, 루프 전개에 맞게 커널을 특화 가능.
• 자동 튜닝: 런타임에 다양한 커널 설정을 벤치마크해 현재 디바이스에 가장 성능이 좋은 버전을 자동 선택.
• Rust 통합: 커널은 Rust 컴파일러에 의해 타입 체크 및 빌로우 체크가 되며, cargo test로 CPU에서 테스트 가능.