CUDAで特殊マトリクス乗算を加速する

11848 ワード

原文は突き刺してください.http://galoisplusplus.coding....
マトリクス乗算はGPU加速一般演算を用いた古典的な例であり,NVIDIA公式のCUDA C Programming GuideとCUDA C Best Practices Guideにもマトリクス乗算を加速する方法を説明するモデルコードがある.本スラグでは,マトリクス乗算をどのように加速させるかという特殊なケース,すなわち大きさの異なる2つのマトリクスの乗算について説明する.
ここで少し考えてみると、実はこのテーマは主に本スラグの論文から来ており、本スラグが行っているのは主にGPUでReed-Solomon Codesの符号化を加速させることである.このような符号化の過程は行列乗算であり,例えばcuBLASなどの行列乗算を実現するCUDA libraryを直接用いることができる.しかし、Reed-Solomon Codesで用いられる演算は、いずれも伽羅華域(Galois Field)に定義されており、通常の実数演算とは異なるため、cuBLASなどのlibraryを直接呼び出すことはできない.あいにく、これらのlibraryがオープンしているAPIではoperator overloadingは許可されておらず、マトリクス乗算のsource codeもオープンソースではありません(MAGMAはopen sourceがコードのlibraryを公表していますが、マトリクス乗算の実現はcuBLASの関数を呼び出し、いくつかの特殊な状況を最適化するだけです).これにより、このスラグは、Reed−Solomon Codes符号化のマトリクス乗算をどのように加速させるかを自力更生する必要がある.実現が完了した後、本かすは多くの技術の細部が学術論文の上品な堂に登るのに足りないことを発見し、それらをこの技術文章に記録したいと思っています.我々のReed-Solomon Codesの実際の応用シーンでは、8-bit byteを数値単位とするマトリクスが用いられている.マトリクス乗算は通常、小さなマトリクスに扁長の大きなマトリクスを乗じた形式である.小さなマトリクスの行数と列数は一般的に3桁を超えず、大きなマトリクスの行数は小さなマトリクスの列数であり、大きなマトリクスの列数は一般的に百万の数級より大きい.本スラグの最適化は主にこれらの実情に基づいて行われるが,最適化の技術は一般的な小行列に大行列を乗じた場合に拡張できる.
このスラグの独特な最適化テクニックを紹介する前に、最も簡単な方法から振り返ってみましょう.
Naive Implementation
次の図は、行列乗算の最も一般的な方法を示しています.
{% img/images/mm-CUDA/without-tiling.png %}
CUDA pseudo-codeは以下の通りである.
// input matrix A and B, compute product matrix C
// A: A_height x A_width
// B: A_width x B_width
// C: A_height x B_width
__global__ void matrix_mul(unsigned char *A, unsigned char *B, unsigned char *C, int A_height, int A_width, int B_width)
{
    int bx = blockIdx.x;
    int by = blockIdx.y;
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int row, col;
    //     B       GPU grid size,       while  。
    do {
        unsigned char product = 0;
        row = by * blockDim.y + ty;
        col = bx * blockDim.x + tx;
        if (row < A_height && col < B_width) {
            for (int i = 0; i < A_width; i++) {
                //       + *   operator  overload Galois Field      。
                product += A[row * A_width + i] * B[i * B_width + col];
            }
            C[row * B_width + col] = product;
        }
        bx += gridDim.x;
        col = bx * blockDim.x + tx;
    } while (col < B_width);
}

このような最も一般的なやり方が悪い原因は、一方、マトリクスの要素(すなわち図中の緑色の小格子)がGPUのglobal memory(global memoryはoff-chipのDRAMであり、accessは400~800 cyclesを必要とする)の中に存在するためであり、このようなやり方はaccessのこれらの要素を大量に必要とし、そのため大量のglobal memory transactionsをもたらした.一方,この手法はaccessマトリクスBの要素にcolumn majorを用い,この方式のlocalityが悪く,より高いcache missrateをもたらす.
Square-Tiling Algorithm
最も一般的なアプローチに対する最適化は一般的なアルゴリズムであり、Hennessy and Pattersonの古典教材「Computer Architecture:A Quantitative Approach」などの文献ではblocking、CUDAのoffical guideなどではtilingと呼ばれている.下図に示すように、黄色の正方形はtileと呼ばれているので、ここではこのalgorithmをsquare-tilingアルゴリズムとも呼ぶ.
{% img/images/mm-CUDA/square-tiling.png %}
CUDA pseudo-codeは以下の通りであり、CUDA offical guideにも同様の例が見られる.
// input matrix A and B, compute product matrix C
// A: A_height x A_width
// B: A_width x B_width
// C: A_height x B_width
__global__ void matrix_mul(unsigned char *A, unsigned char *B, unsigned char *C, int A_height, int A_width, int B_width)
{
    // TILE_SIZE   macro
    __shared__ unsigned char rowVector[TILE_SIZE][TILE_SIZE];
    __shared__ unsigned char colVector[TILE_SIZE][TILE_SIZE];
    int bx = blockIdx.x;
    int by = blockIdx.y;
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int row, col;
    //     B       GPU grid size,       while  。
    do {
        unsigned char product = 0;
        row = by * TILE_SIZE + ty;
        col = bx * TILE_SIZE + tx;
        __syncthreads();

        if (row < A_height && col < B_width) {
            // TILE_SIZE   A    ,        bound          。
            for (int i = 0; i < (int) (ceil((float) A_width / TILE_SIZE)); i++) {
                int bound = min(A_width, TILE_SIZE);
                for (int j = tx; j < bound; j += blockDim.x) {
                    rowVector[ty][j] = A[row * A_width + i * bound + j];
                }
                for (int j = ty; j < bound; j += blockDim.y) {
                    colVector[j][tx] = B[col + (i * bound + j) * B_width];
                }
                __syncthreads();
                for (int j = 0; j < bound; j++) {
                    //       + *   operator  overload Galois Field      。
                    product += rowVector[ty][j] * colVector[j][tx];
                }
            }
            C[row * B_width + col] = product;
        }
        bx += gridDim.x;
        col = bx * blockDim.x + tx;
    } while (col < B_width);
}

ここでは以上のコードを簡単に説明する:マトリクスAとマトリクスBのtileの要素はglobal memoryからshared memory(shared memoryはon-chipのSRAMであり、accessは40 cycles程度を必要とする)にloadされ、その後計算中にshared memoryで直接accessを行う.access shared memoryのlatencyはaccess global memoryのlatencyよりずっと低い.tiling algorithmはshared memoryをmemory accessのcacheとして用い,localityを改善するとともに,use cache中のelementsをreuseすることができ,global memory transactionsの数を低減することができる.
しかし、square-tiling algorithmはマトリクスAとBのサイズが近い場合が一般的であり、Reed-Solomon codesの実際の符号化の場合には適用されない.したがってtileの形状をgeneralizationする必要がある.
Generalized Tiling Algorithm
私のgeneralizationは下図のように{%img/images/mm-CUDA/tiling.png%}
どのようにtileサイズのパラメータ(tileWidthRow,tileWidthCol,tileDepth)を調整するかはpaperの範疇に属し、ここでは省略して、コードを書く方法だけを話します.
最も簡単で乱暴な方法は3つのmacroを使うことです.
// input matrix A and B, compute product matrix C
// A: A_height x A_width
// B: A_width x B_width
// C: A_height x B_width
__global__ void matrix_mul(unsigned char *A, unsigned char *B, unsigned char *C, int A_height, int A_width, int B_width)
{
    // TILE_WIDTH_ROW, TILE_WIDTH_COL, TILE_DEPTH  macro
    __shared__ unsigned char rowVector[TILE_WIDTH_ROW][TILE_DEPTH];
    __shared__ unsigned char colVector[TILE_DEPTH][TILE_WIDTH_COL];
    int bx = blockIdx.x;
    int by = blockIdx.y;
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int row, col;
    //     B       GPU grid size,       while  。
    do {
        unsigned char product = 0;
        row = by * TILE_WIDTH_ROW + ty;
        col = bx * TILE_WIDTH_COL + tx;
        __syncthreads();

        if (row < A_height && col < B_width) {
            for (int i = 0; i < (int) (ceil((float) A_width / TILE_DEPTH)); i++) {
                int bound = min(A_width, TILE_DEPTH);
                for (int j = tx; j < bound; j += blockDim.x) {
                    rowVector[ty][j] = A[row * A_width + i * bound + j];
                }
                for (int j = ty; j < bound; j += blockDim.y) {
                    colVector[j][tx] = B[col + (i * bound + j) * B_width];
                }
                __syncthreads();
                for (int j = 0; j < bound; j++) {
                    //       + *   operator  overload Galois Field      。
                    product += rowVector[ty][j] * colVector[j][tx];
                }
                __syncthreads();
            }
            C[row * B_width + col] = product;
        }
        bx += gridDim.x;
        col = bx * blockDim.x + tx;
    } while (col < B_width);
}

この方法はcompile timeで3つのパラメータのサイズを決定する必要があり,実際のニーズを満たすことは明らかではない.本スラグのやり方は前の3つのmacroの代わりにparameterを使うので、もちろん、shared memory全体を管理する必要があります.CUDA pseudo-codeは以下の通りである.
// input matrix A and B, compute product matrix C
// A: A_height x A_width
// B: A_width x B_width
// C: A_height x B_width
__global__ void matrix_mul(unsigned char *A, unsigned char *B, unsigned char *C, int A_height, int A_width, int B_width, int tileWidthRow, int tileWidthCol, int tileDepth)
{
    extern __shared__ unsigned char sMem[];
    int rowVectorSize = tileWidthRow * tileDepth;
    int bx = blockIdx.x;
    int by = blockIdx.y;
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int row, col;
    //     B       GPU grid size,       while  。
    do {
        unsigned char product = 0;
        row = by * tileWidthRow + ty;
        col = bx * tileWidthCol + tx;
        __syncthreads();

        if (row < A_height && col < B_width) {
            //        runtime  ,       tileDepth A       。
            for (int j = tx; j < tileDepth; j += blockDim.x) {
                sMem[ty * tileDepth + j] = A[row * A_width + j];
            }
            for (int j = ty; j < tileDepth; j += blockDim.y) {
                sMem[rowVectorSize + j * tileWidthCol + tx] = B[col + j * B_width];
            }
            __syncthreads();
            //       + *   operator  overload Galois Field      。
            for (int j = 0; j < tileDepth; j++) {
                product += sMem[ty * tileDepth + j] * sMem[rowVectorSize + j * tileWidthCol + tx];
            }
            __syncthreads();
            C[row * B_width + col] = product;
        }
        bx += gridDim.x;
        col = bx * blockDim.x + tx;
    } while (col < B_width);
}

Further Improvement
前述したように、Reed-Solomon Codesの行列は、8-bit byteを1つの数値単位とする.行列Bの各行がword−alignedされる場合、このスラグは、ALUのoperations数およびglobal memory transactionsを低減するためにさらに最適化することもできる.
我々の最初の観測はelementsに対して実行される加算演算に重点を置いた.従来は2つの8ビットbyteを加算していたが,GPUではALUのオペランドと結果のレジスタは32ビットであった.我々のやり方は加算をする前に4つのbyteを1つのwordのレジスタに包み、ALUで使用されるレジスタが毎回十分に使用されることを保証することです.
我々の第2の観察は,マトリクスBの黄色矩形のelementsをglobal memory loadからshared memoryに移行する過程でbyteが同時にloadしてきたのに対し,8ビットはglobal memory transactionsの最低単位(global memory transactionsは8ビット,16ビット,32ビット,64ビット,128ビット)であった.loadごとにbitを増やし、transactionsの数を減らす方法を考えなければなりません.
私の実装では、wordの幅が異なる場合に適応するためにtemplateを採用しました.CUDA pseudo-codeは以下の通りである.
// input matrix A and B, compute product matrix C
// A: A_height x A_width
// B: A_width x B_width
// C: A_height x B_width
template 
__global__ void matrix_mul(unsigned char *A, T *B, T *C, int A_height, int A_width, int B_width, int tileWidthRow, int tileWidthCol, int tileDepth)
{
    extern __shared__ unsigned char sMemBytes[];
    extern __shared__ T sMemWords[];
    int rowVectorSize = (int) (ceil((float) tileWidthRow * tileDepth / sizeof(T))) * sizeof(T);
    int bx = blockIdx.x;
    int by = blockIdx.y;
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int row, col;
    //     B       GPU grid size,       while  。
    do {
        T product = 0;
        row = by * tileWidthRow + ty;
        col = bx * tileWidthCol + tx;
        __syncthreads();

        if (row < A_height && col < B_width) {
            //        runtime  ,       tileDepth A       。
            for (int j = tx; j < tileDepth; j += blockDim.x) {
                sMemBytes[ty * tileDepth + j] = A[row * A_width + j];
            }
            for (int j = ty; j < tileDepth; j += blockDim.y) {
                sMemWords[rowVectorSize / sizeof(T) + j * tileWidthCol + tx] = B[col + j * B_width];
            }
            __syncthreads();
            //       + *   operator  overload Galois Field      。
            for (int j = 0; j < tileDepth; j++) {
                T C_word_item = 0;
                unsigned char *C_byte_item = (unsigned char *) &C_word_item;
                for (int k = 0; k < sizeof(T); k++) {
                    C_byte_item[k] += sMemBytes[ty * tileDepth + j] * sMemBytes[rowVectorSize + (j * tileWidthCol + tx) * sizeof(T) + k];
                }
                product += C_word_item;
            }
            __syncthreads();
            C[row * B_width + col] = product;
        }
        bx += gridDim.x;
        col = bx * blockDim.x + tx;
    } while (col < B_width);
}

ここで簡単に説明すると、以下の2つのpointer(sMemBytesとsMemWords)は実際に同じメモリアドレスを指している(この性質は各種CUDA入門資料を参照).ここでdeclareは、同じメモリを異なるタイプで操作しやすいように2回使用します.
    extern __shared__ unsigned char sMemBytes[];
    extern __shared__ T sMemWords[];

後記
本スラグはすでに答弁を通過し、本文に関連する内容は本スラグ答弁時のslidesを参考にして、いくつかのアニメーションを作った.