AMD GPU 프로그래밍 입문에서 전문가까지 (1부) - Composable Kernel(CK)의 TensorDescriptor

섹션별 상세

TensorDescriptor는 논리적 다차원 좌표를 물리적 메모리 주소로 매핑하는 CK의 핵심 추상화 도구이다. CalculateLowerIndex 메서드를 포함한 변환(Transform)들을 트리 구조로 연결하여 상위 좌표를 하위 좌표로 단계적으로 변환하는 메커니즘을 가진다. 이를 통해 개발자는 복잡한 메모리 레이아웃 계산을 직접 수행하지 않고도 선언적으로 정의할 수 있다.

cpp

auto tensor_desc = make_naive_tensor_descriptor(make_tuple(M, K), make_tuple(K, 1));

auto transformed_tensor_desc = transform_tensor_descriptor(
    tensor_desc,
    make_tuple(unmerge, passthrough),
    make_tuple(Sequence<0>{}, Sequence<1>{}), // Lower dimension ids
    make_tuple(Sequence<0, 1>{}, Sequence<2>{}) // Upper dimension ids
);

TensorDescriptor를 생성하고 Unmerge 변환을 적용하여 2D 텐서를 3D 레이아웃으로 변경하는 예시

TensorDescriptor의 계층적 트리 구조 다이어그램 — Diagram논리적 좌표가 여러 단계의 Transform(Unmerge, Passthrough, Embed)을 거쳐 최종 물리적 메모리 주소로 변환되는 과정을 시각화한다. 각 노드는 차원을, 박스는 변환 로직을 나타내어 CK의 추상화 방식을 설명한다.

좌표 변환 프로세스 상세도 — Diagram상위 차원의 좌표(a1, a2, a3)가 각 변환 단계를 통과하며 하위 차원의 인덱스로 계산되는 수식과 흐름을 보여준다. 최종적으로 M*K 형태의 선형 오프셋이 도출되는 메커니즘을 증명한다.

계층적 변환 체이닝을 통해 데이터 구조를 유연하게 재구성할 수 있다. Embed, Unmerge, Merge, PassThrough와 같은 변환을 조합하여 2D 텐서를 특정 차원으로 분할하거나 여러 차원을 하나로 병합하는 조작이 가능하다. 이러한 가변적인 레이아웃 처리는 다양한 GPU 아키텍처와 연산 패턴에 최적화된 커널을 설계하는 데 필수적이다.

고성능 행렬 전치 구현을 위해 각 스레드가 4x4 서브 행렬을 레지스터(VGPR) 내에서 직접 처리하도록 설계한다. 벡터화된 읽기 및 쓰기 명령어를 사용하여 전역 메모리 대역폭 활용도를 높이고, 공유 메모리 사용을 배제하여 스레드 간 동기화 오버헤드를 제거한다. 이 방식은 AMD MI308X GPU에서 PyTorch 구현체(8.4 μs)보다 빠른 5.820 μs의 실행 시간을 달성했다.

cpp

static_for<0, 4, 1>{}([&](auto i){
    a(Number<i>{}) = buf.Get<d4_t>(
        tensor_desc.CalculateOffset(Tuple{x + i, y}), true);
});

// In-register transpose
static_for<1, 4, 1>{}([&](auto i){
    static_for<0, i, 1>{}([&](auto j){
        auto tmp = b(Number<i * 4 + j>{});
        b(Number<i * 4 + j>{}) = b(Number<j * 4 + i>{});
        b(Number<j * 4 + i>{}) = tmp;
    });
});

벡터화된 읽기를 수행한 후 레지스터 내에서 4x4 행렬 전치를 수행하는 핵심 로직

스레드별 4x4 서브 행렬 처리 및 레지스터 전치 과정 — Diagram입력 행렬에서 4x4 블록을 벡터화하여 읽어온 후, 레지스터 내에서 전치(In-register transpose)를 수행하고 다시 벡터화하여 출력 행렬에 쓰는 과정을 도식화한다. 이 방식이 왜 메모리 효율적인지 시각적으로 전달한다.

용어 해설

컴포저블 커널(Composable Kernel (CK)): — AMD GPU에서 고성능 AI 및 HPC 연산을 구현하기 위해 제공되는 C++ 템플릿 라이브러리이다. 복잡한 커널을 재사용 가능한 빌딩 블록으로 분해하고 조합하여 최적의 성능을 낼 수 있도록 설계되었다.
텐서 디스크립터(TensorDescriptor): — 다차원 텐서의 논리적 좌표를 물리적 메모리 주소로 매핑하는 추상화 계층이다. 계층적 변환 트리를 통해 데이터 레이아웃을 정의하며, 복잡한 인덱스 계산을 자동화하여 커널 개발 효율을 높인다.
벡터 범용 레지스터(VGPR): — GPU의 각 스레드가 연산 데이터를 저장하기 위해 사용하는 가장 빠른 메모리 공간이다. 공유 메모리나 전역 메모리보다 접근 속도가 압도적으로 빠르며, 이를 효율적으로 활용하는 것이 커널 성능의 핵심이다.
벡터화된 메모리 접근(Vectorized Memory Access): — 한 번의 메모리 명령으로 여러 개의 데이터(예: float4)를 동시에 읽거나 쓰는 기법이다. 메모리 대역폭 활용도를 극대화하고 명령어 수를 줄여 GPU 추론 및 학습 속도를 향상시킨다.
정적 루프(static_for): — 컴파일 타임에 루프를 완전히 언롤링(Unrolling)하여 실행 시점의 제어 오버헤드를 제거하는 기법이다. GPU 커널에서 반복문을 최적화하여 연산 처리량을 높이는 데 필수적으로 사용된다.

기술

Composable Kernel
ROCm
AMD MI308X
PyTorch
C++

활용 사례

Matrix Transpose Optimization
GEMM (General Matrix Multiply) Preparation
Custom GPU Kernel Development

언급된 리소스

튜토리얼AMD GPU Programming From Beginner to Expert (Part 1)