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

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

GPU 커널 프로그래밍에서 다차원 데이터를 물리적 메모리 주소로 매핑하는 효율적인 구조 설계는 성능 최적화의 핵심이다. AMD의 Composable Kernel(CK)은 TensorDescriptor라는 추상화 계층을 도입하여 논리적 좌표를 물리적 오프셋으로 변환하는 계층적 변환 트리를 제공한다. 본문은 Embed, Unmerge, Merge 등의 변환을 조합하여 복잡한 레이아웃을 구성하는 방법과 이를 실제 행렬 전치 커널에 적용하는 과정을 상세히 설명한다. AMD MI308X 환경에서 테스트한 결과, 이 방식을 적용한 커널은 PyTorch 대비 약 44.3%의 성능 향상을 기록했다.

배경

C++ Template Programming, GPU Architecture (CU, VGPR, Global Memory), Matrix Operations Basics

대상 독자

AMD ROCm 플랫폼에서 고성능 GPU 커널을 직접 설계하고 최적화하려는 AI/HPC 개발자

의미 / 영향

이 기술은 AMD GPU의 하드웨어 잠재력을 최대한 끌어내어 PyTorch와 같은 범용 프레임워크보다 높은 성능의 커널을 구축할 수 있게 한다. 특히 대규모 언어 모델(LLM)이나 복잡한 행렬 연산이 필요한 AI 워크로드에서 비용 대비 성능을 최적화하는 데 중요한 역할을 할 것으로 기대된다.

섹션별 상세

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)를 수행하고 다시 벡터화하여 출력 행렬에 쓰는 과정을 도식화한다. 이 방식이 왜 메모리 효율적인지 시각적으로 전달한다.

실무 Takeaway

TensorDescriptor를 활용하면 복잡한 다차원 인덱스 계산 로직을 추상화하여 커널 코드의 가독성과 유지보수성을 획기적으로 개선할 수 있다.
GPU 커널 최적화 시 static_for와 같은 컴파일 타임 루프 언롤링을 적용하여 런타임 제어 오버헤드를 최소화하고 연산 처리량을 극대화해야 한다.
성능 병목을 해결하기 위해 공유 메모리 대신 레지스터 수준에서 데이터를 처리하고 벡터화된 메모리 접근 패턴을 우선적으로 고려하는 설계가 유효하다.

언급된 리소스

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

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 );

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; }); });

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

핵심 요약

배경

대상 독자

의미 / 영향

섹션별 상세

실무 Takeaway

언급된 리소스

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

핵심 요약

배경

대상 독자

의미 / 영향

섹션별 상세

실무 Takeaway

언급된 리소스

관련 토론

댓글

관련 피드

관련 토론

댓글

관련 피드