NVlabs/CuTe

해결하는 문제

PyCuTe는 NVIDIA의 CUTLASS 3.x에서 사용되는 계층적 레이아웃 및 텐서 대수의 순수한 Python 구현체를 제공합니다. 개발자들은 GPU나 원래 CuTe 라이브러리에 필요한 C++ 환경 없이도 복잡한 대수를 학습하고, 새로운 데이터 변환을 프로토타이핑하며, 테스트 벡터를 생성할 수 있습니다.

작동 방식

핵심적으로 PyCuTe는 Shape(정의역)과 Stride(매핑 로직)를 사용하여 좌표를 오프셋으로 매핑하는 함수로 Layout을 정의합니다. coalesce, composition, logical_divide(타일링), logical_product(반복)와 같은 포괄적인 대수 연산을 구현하여 이러한 레이아웃을 조작할 수 있습니다. 또한 공유 메모리의 백엔드 충돌을 처리하기 위해 F2(XOR-swizzle) 스트라이드도 지원합니다. Layout과 Accessor(예: 포인터)를 결합하여 Tensor를 생성함으로써 데이터를 읽거나 쓸 수 있습니다.

대상 사용자

GPU 텐서 계산을 다루는 개발자 및 연구자, Python에서 레이아웃 변환을 프로토타이핑하고 싶은 CUTLASS 사용자, 계층적 텐서 레이아웃의 수학적 기초를 공부하고자 하는 사람들에게 적합합니다.

주요 특징

• 순수한 Python: 핵심 대수에는 GPU나 CUDA 종속성이 필요하지 않습니다.
• 레이아웃 대수: 복잡한 타일링 및 분할을 위한 컴포지션, 컴플리먼트, 널스페이스 등의 연산을 완전히 구현했습니다.
• 시각화 도구: 레이아웃을 ASCII 테이블, 색상 기반 SVG, LaTeX/PDF로 렌더링할 수 있는 내장 지원을 통해 데이터 이동과 백엔드 충돌을 시각화할 수 있습니다.
• 참조 구현체: C++ CuTe 라이브러리와 CuTe DSL에 대한 충실한 Python 대응 구현체입니다.