Как эффективно вычислять сумму всех столбцов и строк матрицы в CUDA?

Question

Я хотел бы рассчитать сумму всех столбцов и сумму всех строк матрицы в CUDA. Один из способов сделать это - использовать подпрограмму SGEMV от BLAS, умножая матрицу на вектор 1s.

Однако это приводит к двум сканированию матрицы, если предположить, что она намного больше, чем кеш L1: одна для строк и другая для столбцов. Кроме того, я планирую дополнительно модифицировать код для других операторов, поэтому я пишу свое собственное ядро.

Мой подход до сих пор заключался в том, чтобы разбить матрицу на подматрицы размером 32 x 32. Каждый блок потоков загружает такую ​​подматрицу в общую память, вычисляет суммы строк и количеств подматрицы и добавляет их атомарно к соответствующему выходу (row и col ниже). Таким образом, данные матрицы должны считываться только с VRAM один раз.

Для простоты, код предполагает, что матрица является n x n, n % 32 == 0 и блок резьбы 32 x 32

__global__ void sum_cols_and_rows(size_t n, const float* matrix, float* col, float* row) 
{ 
    __shared__ float sh[32][32]; 

    size_t x = blockDim.x * blockIdx.x + threadIdx.x; 
    size_t y = blockDim.y * blockIdx.y + threadIdx.y; 

    float sum = matrix[x + n * y]; 
    sh[threadIdx.x][threadIdx.y] = sum; 

    for(unsigned w = 16; w >= 1; w /= 2) 
     sum += __shfl_down(sum, w); 
    const size_t laneID = threadIdx.x & 0x1f; // 32-1 
    if(laneID == 0) 
     atomicAdd(row + y, sum); 
    __syncthreads(); 

    sum = sh[threadIdx.y][threadIdx.x]; // swapped indexes 
    for(unsigned w = 16; w >= 1; w /= 2) 
     sum += __shfl_down(sum, w); 
    if(laneID == 0) 
     atomicAdd(col + blockDim.x * blockIdx.x + threadIdx.y, sum); 
} 

// launch : 
sum_cols_and_rows<<<dim3(n/32, n/32), dim3(32, 32), 32*32*sizeof(float)>>>(n, matrix, col, row); 

Однако производительность довольно неутешительная. Я вижу около 20% теоретической пропускной способности памяти 224 ГБ/с на GTX 980, даже на больших матрицах, , например. 16384x16384.

Есть ли способ сделать этот подход теоретическим пределом пропускной способности?

Answer 1

В вашем решении каждый блок NxN матрицы обрабатывается отдельным блоком потоков NxN. По сути, каждая отдельная нить очень мало работает, поэтому накладные расходы доминируют над фактическими вычислениями. Вы можете улучшить его, если блоки потоков будут обрабатывать более одного блока матрицы.

Но есть более простое решение, использующее только N потоков на матричный блок, где один поток суммирует весь столбец.

Реализация будет похожа на это:

__global__ void sum_cols_and_rows(size_t n, const float* matrix, float* col, float* row) 
{ 
    size_t laneID = threadIdx.x & 31; 

    size_t x = blockDim.x * blockIdx.x + threadIdx.x; 
    size_t y = N_ITERATIONS * blockIdx.y; 

    size_t idx = y * n + x; 

    float vertical = 0; 

    for(int i = 0; i < N_ITERATIONS; i++) { 
     float v = matrix[idx]; 
     vertical += v; 
     for(unsigned w = 16; w >= 1; w /= 2) 
      v += __shfl_down(v, w); 
     if(laneID == 0) 
      atomicAdd(&row[y], v); 
     y++; 
     idx += n; 
    } 

    atomicAdd(&col[x], vertical); 
} 

перестраиваемых параметры здесь количество перекосов в группу потоков и количество строк в каждой матрице блоке (N_ITERATIONS). Большие значения могут уменьшаться накладные расходы за счет параллелизма.

Другая идея поэкспериментировать с является vectorized loading - один из:

float2 v2 = reinterpret_cast<float2*>(matrix)[idx]; 
float v = v2.x + v2.y; 

float4 v4 = reinterpret_cast<float4*>(matrix)[idx]; 
float v = (v4.x + v4.y) + (v4.z + v4.w);