如何利用GPU实现稀疏与密集数组的快速张量点积？

问题概述 #

需要计算密集数组arr与块稀疏数组cost的张量点积：

• arr尺寸：(width, height, m, n)，其中width/height约5000，m<100，n<10
• cost尺寸：(width, height, width, height)，为严格块稀疏结构——每个cost[i,j,:,:]仅包含一个k*k块（k<50）
• 目标计算：result = np.tensordot(cost, arr, axes=[[2,3],[0,1]])，结果尺寸与arr一致
• 现有瓶颈：CPU多线程版本速度不足，GPU密集版本无法适配稀疏结构

实现方案 #

方案1：CUDA C++ 手动块稀疏计算 #

利用cost的块特性，跳过完整张量存储，直接基于块信息实现GPU加速：

1. 稀疏存储格式定义
   只存储非零块的核心信息，避免冗余：

   struct Block {
       int i, j;          // 块对应的(i,j)索引
       int block_x, block_y;  // 块在cost[i,j,:,:]中的起始坐标
       float data[50][50];    // 块内数据（按最大k=50分配固定空间）
   };
   

   同时预计算二维映射数组block_idx[i][j]，快速查找(i,j)对应的块索引（无块则为-1）。

2. CUDA核函数实现
   按(i,j,m,n)维度分配线程，每个线程负责一个输出元素的计算：

   __global__ void sparse_tensordot_kernel(
       const Block* blocks,
       const int* block_idx,
       const float* arr,
       float* result,
       int width, int height, int m, int n, int max_k
   ) {
       int i = blockIdx.x * blockDim.x + threadIdx.x;
       int j = blockIdx.y * blockDim.y + threadIdx.y;
       int m_idx = blockIdx.z * blockDim.z + threadIdx.z;
       int n_idx = blockIdx.w * blockDim.w + threadIdx.w;
   
       if (i >= width || j >= height || m_idx >= m || n_idx >= n) return;
   
       int idx = block_idx[i * height + j];
       if (idx == -1) {
           result[i * height * m * n + j * m * n + m_idx * n + n_idx] = 0.0f;
           return;
       }
   
       const Block& b = blocks[idx];
       float sum = 0.0f;
       // 遍历块内有效元素
       for (int dx = 0; dx < max_k; dx++) {
           int x = b.block_x + dx;
           if (x >= width) break;
           for (int dy = 0; dy < max_k; dy++) {
               int y = b.block_y + dy;
               if (y >= height) break;
               float c_val = b.data[dx][dy];
               float a_val = arr[x * height * m * n + y * m * n + m_idx * n + n_idx];
               sum += c_val * a_val;
           }
       }
       result[i * height * m * n + j * m * n + m_idx * n + n_idx] = sum;
   }
   

3. 调用流程

   • 将blocks、block_idx、arr拷贝至GPU显存
   • 配置线程块（如(16,16,4,2)，根据GPU算力调整）
   • 启动核函数后将result拷贝回CPU

方案2：PyTorch 稀疏张量+广播优化 #

利用PyTorch内置的稀疏张量支持，结合einsum简化计算：

1. 构造稀疏cost张量
   提取所有非零块的坐标与值，无需存储完整四维张量：
   • 遍历每个(i,j)，生成对应k*k块内所有(i,j,x,y)的坐标
   • 将坐标整理为(N,4)的LongTensor，块内值整理为(N,)的FloatTensor，构造稀疏张量
2. 执行张量点积

   import torch
   
   # 假设已构造好稀疏cost张量：sparse_cost (shape: [width, height, width, height], sparse)
   arr = torch.tensor(arr, device='cuda')
   
   # 等价于tensordot(axes=[[2,3],[0,1]])的计算
   result = torch.einsum('ijxy,xymn->ijmn', sparse_cost, arr)
   

   PyTorch会自动针对稀疏张量的块结构做计算优化，减少无效运算。

方案3：CuPy 自定义RawKernel实现 #

基于Python生态，用CuPy的RawKernel实现类似CUDA C++的逻辑：

import cupy as cp

kernel_code = r'''
extern "C" __global__ void sparse_tensordot(
    const int* block_info,  // 每个块存储[i,j,block_x,block_y]，按顺序排列
    const float* block_data,  // 所有块的数据，每个块占max_k*max_k个元素
    const float* arr,
    float* result,
    int width, int height, int m, int n, int max_k, int num_blocks
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;
    int m_idx = blockIdx.z * blockDim.z + threadIdx.z;
    int n_idx = blockIdx.w * blockDim.w + threadIdx.w;

    if (i >= width || j >= height || m_idx >= m || n_idx >= n) return;

    // 查找当前(i,j)对应的块索引（可提前预计算映射数组优化此步骤）
    int block_idx = -1;
    for (int b=0; b<num_blocks; b++) {
        if (block_info[b*4] == i && block_info[b*4+1] == j) {
            block_idx = b;
            break;
        }
    }

    if (block_idx == -1) {
        result[i*height*m*n + j*m*n + m_idx*n + n_idx] = 0.0f;
        return;
    }

    int block_x = block_info[block_idx*4 + 2];
    int block_y = block_info[block_idx*4 + 3];
    const float* data = block_data + block_idx*max_k*max_k;

    float sum = 0.0f;
    for (int dx=0; dx<max_k; dx++) {
        int x = block_x + dx;
        if (x >= width) break;
        for (int dy=0; dy<max_k; dy++) {
            int y = block_y + dy;
            if (y >= height) break;
            float c_val = data[dx*max_k + dy];
            float a_val = arr[x*height*m*n + y*m*n + m_idx*n + n_idx];
            sum += c_val * a_val;
        }
    }
    result[i*height*m*n + j*m*n + m_idx*n + n_idx] = sum;
}
'''

# 编译核函数
kernel = cp.RawKernel(kernel_code, 'sparse_tensordot')

# 准备GPU数据
block_info = cp.array(block_info_list, dtype=cp.int32)
block_data = cp.array(block_data_list, dtype=cp.float32)
arr_gpu = cp.array(arr, dtype=cp.float32)
result_gpu = cp.zeros_like(arr_gpu)

# 配置线程块与网格
block_dim = (16, 16, 4, 2)
grid_dim = (
    (width + block_dim[0] - 1) // block_dim[0],
    (height + block_dim[1] - 1) // block_dim[1],
    (m + block_dim[2] - 1) // block_dim[2],
    (n + block_dim[3] - 1) // block_dim[3]
)

# 启动核函数
kernel(grid_dim, block_dim, (block_info, block_data, arr_gpu, result_gpu, width, height, m, n, 50, len(block_info_list)))

# 结果回传CPU
result = result_gpu.get()

关键优化建议 #

• 预计算块索引映射：提前在CPU生成(i,j)到块的索引数组，拷贝到GPU后直接查询，避免核函数内循环查找
• 线程维度匹配：按(i,j,m,n)分配线程，让每个线程负责一个输出元素，最大化GPU核心利用率
• 优先框架方案：PyTorch/CuPy的内置稀疏操作已做底层优化，无需手动写核函数时优先选择

内容的提问来源于stack exchange，提问作者C.K.