GPU에서의 희소행렬 곱셈(spGEMM)(2)

저번 포스팅에서는 간단하게 희소행렬을 저장할 때 사용되는 포맷에 대해 올렸다. 간단히 복기하자면 메모리 효율성을 위해 COO(Coordinate list ;Triplet) 포맷이나 CSR, CSC(Compressed Sparse Row/Column) 포맷을 이용한다는 것이다.

희소행렬 곱셈 알고리즘

보통 '행렬곱셈' 하면 내적을 이용한 곱셈을 떠올리지만 희소행렬 곱셈에서는 내적방식이 잘 쓰이지 않는다. 만약 행렬이 2D 그 이유는 내적연산을 위해선 두 벡터의 "인덱스가 같은 원소"를 곱해줘야 하지만 희소행렬은 각 행값, 또는 열값이 idx array에 저장되어있다. 따라서 idx array를 읽고 두 index가 같은지 체크해줘야하는 비효율이 생긴다.

Row-wise product

희소행렬 곱셈을 수행할 때 row-wise(row-row) product이 널리 쓰인다. Nvidia에서 제공하는 cuSPARSE, CUSP 역시 이 방식으로 곱셈을 진행한다. 구체적인 알고리즘은 위와 같이 A의 행의 0이 아닌 원소와 그 원소의 열값에 해당하는 B의 행 전체를 곱하고 결과로 C의 행 1개가 생성된다. A, B 모두 원소를 row-wise로 접근하기 때문에 두 행렬 모두 CSR 포맷을 이용한다. 간단히 Naive 한 GPU 코드를 작성해보면 다음과 같다.

// row-wise product based spGEMM
__global__ void spGEMM_expansion(
    int* a_ptr, int* a_idx, double* a_val,
    int* b_ptr, int* b_idx, double* b_val,
    int* c_ptr, int* c_idx, double* c_val)
{
    int bid = blockIdx.x;
    int tid = threadIdx.x;
    
    for(int a_iterator = a_ptr[bid]; a_iterator < a_ptr[bid + 1]; ++a_iterator)
    {
        int a_col = a_iter[a_iterator];
        
        for(int b_iterator = b_ptr[a_col]+tid; b_iterator < b_ptr[a_col + 1]; b_iterator+=blockDim.x)
        {
            // Merge?
            c_idx[offset] = b_idx[b_iterator];
            c_val[offset] = a_val[a_iterator]*b_val[b_iterator];   
        }
    }
}

간단하게 1개의 Thread block이 C의 행 1개를 생성하는 코드이다. 알고리즘 자체는 굉장히 간단하지만 몇가지 문제들이 발생한다. 이슈는 이후에 언급하도록 하겠다.

Outer product

row-wise(row-row) product 방식 이외에는 outer-product(column-row) 방식이 쓰인다. 구체적인 알고리즘은 위와 같이 A의 열과 B의 행의 곱셈으로 C의 일부가 생성된다. row-wise product 방식과 차이점은 연산의 결과가 전체 행렬 한판이라는 것이다. 이는 꽤나 중요한 차이지만 역시 이후에 언급. A의 access pattern은 column-wise, B는 row-wise이기 때문에 A는 CSC 포맷, B는 CSR 포맷을 이용한다. 간단히 Naive 한 GPU 코드를 작성해보면 다음과 같다.

// outer product based spGEMM
__global__ void spGEMM_expansion(
    int* a_ptr, int* a_idx, double* a_val,
    int* b_ptr, int* b_idx, double* b_val,
    int* c_ptr, int* c_idx, double* c_val)
{
    int bid = blockIdx.x;
    int tid = threadIdx.x;
    
    for(int a_iterator = a_ptr[bid]; a_iterator < a_ptr[bid + 1]; ++a_iterator)
    {
        int a_col = a_iter[a_iterator];
        
        for(int b_iterator = b_ptr[bid]+tid; b_iterator < b_ptr[bid + 1]; b_iterator+=blockDim.x)
        {
            // Merge?
            c_idx[offset] = b_idx[b_iterator];
            c_val[offset] = a_val[a_iterator]*b_val[b_iterator];   
        }
    }
}

간단하게 1개의 Thread block이 C행렬의 일부가 되는 C 한 판을 생성하게 된다.