GEMM Transpose 처리 설계 정리 ( AICF CUDA Backend )

요약

• transpose 는 커널 종류를 늘려서 처리하지 않는다 ( NN / TN / NT / TT )
• 런치 단계에서 stride 기반 논리 뷰를 생성한다
• 커널은 항상 논리적으로만 계산
• global - shared memory 로딩 단계에서 이미 transpose 의미가 반영된 인덱싱으로 데이터를 읽어온다
• WMMA 에서는 shared memory 에 목표 레이아웃으로 packing 한다

transpose 는 연산 로직 변경이 아니라 메모리 해석 방식의 문제

 

현재 구조 핵심 아이디어

1. Transpose 는 데이터를 바꾸는 것이 아님

• 실제 메모리를 바꾸는 연산이 아니라
• row, col - 물리 주소 계산 방식이 바뀌는 것

핵심은

logical index (r, c)
→ physical address = r * stride_row + c * stride_col

 

2. 런처에서 논리 뷰를 만든다

struct MatView2D {
  const void* data;
  int64_t rows;
  int64_t cols;
  int64_t rs;   // row stride
  int64_t cs;   // col stride
};

transpose = false

rows = shape[0]
cols = shape[1]
rs   = stride[0]
cs   = stride[1]

transpose = true

rows = shape[1]
cols = shape[0]
rs   = stride[1]
cs   = stride[0]

stride swap 으로 표현된다.

 

3. 커널 입장에서는 항상 동일한 수식

• A 와 B 의 rs / cs 만 다르지, NN / TN / NT / TT 의 구분은 없음

 

Global - Shared Memory 단계에서 transpose 가 이미 반영됨

transpose 처리는 shared memory 에서가 아니라, global memory load 시점에서 이미 끝난다

A_val = A_base[r * Ars + c * Acs];

• rs / cs 가 어떻게 설정되었는지에 따라
• 같은 r, c 라도
  • 정상 row - major
  • transpose view
  • slice 된 non-contig view

여기서 transpose 의미가 적용됨

 

WMMA 경로에서의 shared memory packing

WMMA 는 레이아웃 요구사항이 있음

• matrix_a : row-major
• matrix_b : co-major

그래서 shared memory 에는 WMMA 가 원하는 형태로만 packing 한다.

 

• 이 transpose packing 은 WMMA 레이아웃 요구사항 때문
• 수학적 transpose 를 처리하기 위한 것이 아님

 

이 경우 커널 수가 폭증하지 않음

non-contiguous view 도 자연스럽게 지원

python / autograd 와의 역할 분리가 명확

 

NCU 를 통한 실제 검증 확인