CUTLASS: Fast Linear Algebra in CUDA C++

ijpq發表於2024-03-26
ASTC++

https://developer.nvidia.com/blog/cutlass-linear-algebra-cuda/

Efficient Matrix Multiplication on GPUs

計算密集度 = (時間複雜度/空間複雜度) = O(N^3)/O(N^2) = O(N)

// naive
for (int i = 0; i < M; ++i)
    for (int j = 0; j < N; ++j)
       for (int k = 0; k < K; ++k) 
            C[i][j] += A[i][k] * B[k][j];

image
對於A的每一行，都會讀取一整個B矩陣. 能否達到O(N)的複用次數？B中的元素按照列來訪問，那麼每次重新讀取B[0,0]時是重新進行訪存的。而實際上，讀取一次B[0,0]只進行了1次計算。

// opt 1
for (int k = 0; k < K; ++k)     // K dimension now outer-most loop
    for (int i = 0; i < M; ++i)
        for (int j = 0; j < N; ++j)
            C[i][j] += A[i][k] * B[k][j];

image
對於A[0,0]，只讀取一次，但執行j=0~N-1，共N次計算。
對於B[0,0]，讀取i=0~M-1，j=0, 共M次，執行這M次計算，並且這裡可能會cache住，所以實際讀取次數可能小於M。
對於C[0,0]，讀取k=0~K-1,共K次。執行K次(k=0~K-1, i=0,j=0)計算。
理論上已經比naive更接近於計算密集度O(N)

但這個計算依賴於每一次寫入C中的元素時，這個元素的讀取，寫入速度能和乘法指令一樣快，也就是說這些C中的元素最好能一直在cache中，而不發生trash。否則，這個計算過程就會使依賴於訪存。

// opt 2
for (int m = 0; m < M; m += Mtile)                // iterate over M dimension
    for (int n = 0; n < N; n += Ntile)            // iterate over N dimension
        for (int k = 0; k < K; ++k)
            for (int i = 0; i < Mtile; ++i)       // compute one tile 
                for (int j = 0; j < Ntile; ++j) {
                    int row = m + i;
                    int col = n + j;
                    C[row][col] += A[row][k] * B[k][col];
                }

這樣，在tile夠小的時候，tile中元素就能放到cache中。

相關文章