본문 바로가기

dev_AI_framework

현재 GEMM 구현 내용

#pragma once
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include "../ge/cuda_check.cuh"

// Row-major 매핑 래퍼들

// 단일 배치 GEMM (TF32)
inline void gemm_rm_tf32(
    cublasHandle_t h,
    bool transA, bool transB,
    int M, int N, int K,
    const float* A, int lda,
    const float* B, int ldb,
    float* C, int ldc,
    float alpha=1.f, float beta=0.f)
{
    cublasOperation_t opA = transA ? CUBLAS_OP_T : CUBLAS_OP_N;
    cublasOperation_t opB = transB ? CUBLAS_OP_T : CUBLAS_OP_N;

    CUBLAS_CHECK(
        cublasGemmEx(
            h,
            /*opB,opA*/ opB, opA,
            /*m,n,k*/   N,   M,   K,
            &alpha,
            /*B*/ B, CUDA_R_32F, ldb,
            /*A*/ A, CUDA_R_32F, lda,
            &beta,
            /*C*/ C, CUDA_R_32F, ldc,
            CUBLAS_COMPUTE_32F_FAST_TF32,
            CUBLAS_GEMM_DEFAULT_TENSOR_OP
        )
    );
}

// Strided-batched GEMM (TF32)
inline void gemm_rm_strided_batched_tf32(
    cublasHandle_t h,
    bool transA, bool transB,
    int M, int N, int K,
    const float* A, int lda, long long strideA,
    const float* B, int ldb, long long strideB,
    float* C, int ldc, long long strideC,
    int batch,
    float alpha=1.f, float beta=0.f)
{
    cublasOperation_t opA = transA ? CUBLAS_OP_T : CUBLAS_OP_N;
    cublasOperation_t opB = transB ? CUBLAS_OP_T : CUBLAS_OP_N;

    CUBLAS_CHECK(
        cublasGemmStridedBatchedEx(
            h,
            /*opB,opA*/ opB, opA,
            /*m,n,k*/   N,   M,   K,
            &alpha,
            /*B*/ B, CUDA_R_32F, ldb, strideB,
            /*A*/ A, CUDA_R_32F, lda, strideA,
            &beta,
            /*C*/ C, CUDA_R_32F, ldc, strideC,
            /*batch*/ batch,
            CUBLAS_COMPUTE_32F_FAST_TF32,
            CUBLAS_GEMM_DEFAULT_TENSOR_OP
        )
    );
}

 

1. 행렬 곱의 메모리 접근

  • GPU 메모리에서 coalesced access 가 중요, 
    • Row-major 와 Col-major 가 어떻게 정의되느냐,
    • leading dimension 와 stride 가 얼마냐에 따라 메모리 접근 패턴이 바뀐다.

단순히 열을 따라 원소 하나하나 접근하면 stride 때문에 non-contiguous access 이 생겨 성능이 떨어질 수 있다.

 

2. cuBLAS 내부 최적화

  • cuBLAS는 입력 버퍼를 자도응로 전치하지는 않는다.
    • CUBLAS_OP_T, CUBLAS_OP_N 플래그는 단순히 인덱스 계산 방식을 바꿔서 논리적 전치를 수행한다.
  • 대신 내부적으로 고성능 커널을 사용한다.
    • 데이터 블록을 공유 메모리에 불러와 타일링 방식으로 곱셈
    • 읽기 stride 가 길더라도, warp 단위로 메모리 coalescing 을 최대화하도록 최적화된 접근 패턴

즉, stride 자체를 없애는 건 아니지만, stride 접근 비용을 최소화한다.

 

결론

  • cuBLAS 기본 GEMM은 데이터 복사 없이 논리 전치 + 고성능 타일링으로 “대부분의 케이스”에 최적입니다.
  • B를 물리 전치/패킹하는 전략은 B를 여러 번 재사용할 때 큰 이득을 줍니다. 그럴 땐 **cuBLASLt(알고리즘+workspace)**나 CUTLASS로 가세요.
  • 즉, 단발성은 기본 cuBLAS, 재사용 많은 워크로드는 전치/패킹+cuBLASLt가 베스트 프랙티스입니다.