GPU上快速实现无重叠布尔数组匹配的优化问询

现有大小两个2D NumPy布尔数组，需找出小数组在大数组上所有可放置的位置——要求小数组与大数组对应切片无同时为True的元素（类似无重叠放置物体），结果数组以左上角坐标标记位置且与大数组同形。

已实现的两种GPU方案 #

方案一：PyTorch卷积实现 #

利用卷积操作快速检测重叠，卷积结果为0的位置即为可放置位置：

import torch
import torch.nn.functional as F
import numpy as np

def get_placement_options(large_array, small_array):
    shape = large_array.shape

    # 转换为PyTorch CUDA张量
    large_array = torch.from_numpy(large_array).to('cuda:0').float()
    small_array = torch.from_numpy(small_array).to('cuda:0').float()

    # 卷积计算重叠情况
    possible_locations = (F.conv2d(large_array[None, None], small_array[None, None])[0, 0] < .5).cpu().numpy()
    
    result = np.zeros(shape, dtype='bool')
    result[:possible_locations.shape[0], :possible_locations.shape[1]] = possible_locations
    return result

方案二：CuPy自定义CUDA核+位打包实现 #

通过位打包将布尔数组转为int64，尝试用自定义CUDA核提升效率：

import cupy as cp
import numpy as np
from cupy.lib.stride_tricks import as_strided

class Placer:
    def __init__(self):
        kernel_code = """
        extern "C" {
        __global__ void compute(long long int* large_array, long long int* small_array, int* result, int rn, int rm, int n, int m, int k, int l) {
            int i = blockIdx.x * blockDim.x + threadIdx.x;  // 大数组中待计算的x坐标
            int j = blockIdx.y * blockDim.y + threadIdx.y;  // 大数组中待计算的y坐标
            
            if (i < rn && j < rm) {
            
                int r_i = i % 8;  // x方向的偏移量
                int r_j = j % 8;  // y方向的偏移量
                int sub_array_index = r_i * 8 * n * m + r_j * n * m;

                for (int p = 0; p < k; ++p) {
                    for (int q = 0; q < l; ++q) {
                        if ((small_array[p * l + q] & large_array[sub_array_index + ((i / 8)+p) * m + (j / 8)+q]) != 0) {
                            result[i * rm + j] = 0;
                            return;
                        }
                    }
                }
            }

        }
        }
        """

        # 编译核函数
        self.compiled_kernel = cp.RawKernel(kernel_code, 'compute')

    def __call__(self, large_array: np.ndarray, small_array: np.ndarray):
        
        # 初始化结果数组
        result_np = np.zeros_like(large_array)

        # 补全小数组至8的倍数，同时给大数组补相同的padding（补True，避免边界误判）
        padding = ((0, (8 - small_array.shape[0] % 8) % 8), (0, (8 - small_array.shape[1] % 8) % 8))
        small_array = np.pad(small_array, padding, mode='constant', constant_values=False)
        K, L = small_array.shape
        large_array = np.pad(large_array, padding, mode='constant', constant_values=True) 

        # 补全大数组至8的倍数
        padding = ((0, (8 - large_array.shape[0] % 8) % 8), (0, (8 - large_array.shape[1] % 8) % 8))
        large_array = np.pad(large_array, padding, mode='constant', constant_values=True) 
        N, M = large_array.shape
        
        # 生成大数组的所有8x8滑动窗口，打包为int64并存储所有偏移情况
        large_array_cp = cp.array(large_array)
        large_array_cp = cp.pad(self.sliding_window_view_cp(large_array_cp, (8, 8)), ((0, 7), (0, 7), (0, 0), (0, 0)), 'constant', constant_values=True).transpose((2, 3, 0, 1)).copy()
        large_array_cp = cp.packbits(large_array_cp.transpose((0, 1, 3, 2))).reshape(8, 8, large_array.shape[1], large_array.shape[0] // 8)
        large_array_cp = large_array_cp.transpose((0, 1, 3, 2)).copy().view('int64')

        # 将小数组打包为int64
        small_array = cp.array(np.packbits(small_array.copy(), axis=0).view(np.int64))
        
        # 调用核函数
        block = (32, 32, 1)
        grid = ((N-K+1 + block[0] - 1) // block[0], (M-L+1 + block[1] - 1) // block[1])
        result = cp.ones((N-K+1, M-L+1), dtype=cp.int32)
        self.compiled_kernel(grid=grid, block=block, args=(large_array_cp, small_array, result, N-K+1, M-L+1, N // 8, M // 8, K // 8, L // 8))
        
        # 等待GPU计算完成
        cp.cuda.stream.get_current_stream().synchronize()

        # 转换结果并填充到输出数组
        result = result.astype(cp.bool_).get()
        result_np[:result.shape[0], :result.shape[1]] = result

        return result_np

    @staticmethod
    def sliding_window_view_cp(arr, window_shape):
        output_shape = arr.shape[:-len(window_shape)] + tuple(i - j + 1 for i, j in zip(arr.shape[-len(window_shape):], window_shape))
        strides = arr.strides + arr.strides[-len(window_shape):]
        return as_strided(arr, shape=output_shape + window_shape, strides=strides)

性能问题 #

理论上位打包+自定义CUDA核的方案应具备更高效率，但实际测试中与PyTorch方案速度相近（分别耗时0.06秒、0.05秒），推测自定义CUDA核缺少关键优化，需进一步提速。

测试用例 #

large_array = np.random.choice(a=[False, True], size=(1000, 1000))
small_array = np.zeros((100, 100), dtype=bool)
small_array[-4:, -4:] = True
large_array[-260:, -260:] = False

有效位置示意图：大数组右下角区域存在大片可放置区域，其余位置因随机True值分布零散可放置点。

优化思路与方案 #

• 内存访问优化
  • 改用共享内存预加载大数组的局部块，减少全局内存的重复访问。每个线程块加载对应区域的大数组数据到共享内存，线程间复用数据，降低内存延迟。
  • 调整数组存储布局，确保全局内存访问符合合并访问规则，避免非对齐访问带来的额外开销。
• 位运算并行化
  • 利用64位整数的位运算一次性处理8x8块：将小数组的位打包块与大数组对应块做按位与，只要结果非零就标记不可放置，一次操作替代64次单元素检查，大幅减少循环次数。
• 分支优化
  • 用掩码操作替代原核函数中的提前return分支：用变量记录冲突状态，循环结束后统一赋值结果，减少GPU分支发散带来的性能损耗。
• 预处理优化
  • 对小数组做预处理，只记录其中True值的位置，核函数中仅检查大数组对应位置，避免遍历小数组的大量零位，减少计算量。
• 简化预处理流程
  • 改用CuPy内置的cp.lib.stride_tricks.sliding_window_view替代自定义滑动窗口实现，减少手动内存操作带来的拷贝和性能损耗。

内容的提问来源于stack exchange，提问作者F.Wessels