在PyTorch中如何在CUDA GPU上并行化重复执行的布尔函数？

针对你的需求，要让代码中标记为IN PARALLEL的布尔运算在CUDA GPU上强制并行，最直接高效的方式是将独立运算合并为张量级操作，利用PyTorch对GPU并行的原生优化，同时避免多次kernel调用的开销。

核心思路 #

原代码中逐元素赋值的方式会触发多次独立的GPU kernel调用，GPU调度这些小kernel时会有额外开销，且无法保证并行执行。通过把三个独立的布尔运算打包成一次张量操作，PyTorch会生成一个单一的GPU kernel，让三个运算在不同的CUDA线程上并行执行，充分利用GPU的SIMT架构优势。

修改后的代码 #

import torch

A = torch.tensor([True, False, True]).to('cuda')
B = torch.tensor([False, True, True]).to('cuda')

n_steps = 100

def compute_next(input_tensor):
    # 一次性计算三个输出值，GPU会并行处理这三个独立运算
    y0 = torch.logical_and(torch.logical_or(input_tensor[0], input_tensor[1]), input_tensor[2])
    y1 = torch.logical_or(torch.logical_or(input_tensor[0], input_tensor[1]), input_tensor[2])
    y2 = torch.logical_and(torch.logical_and(input_tensor[0], input_tensor[1]), input_tensor[2])
    # 将三个结果合并为一个张量返回，完成一次并行计算
    return torch.stack([y0, y1, y2])

for step in range(n_steps):
    # 并行计算B的所有元素，替代原代码中三次单独赋值
    B = compute_next(A)
    
    # 可选：显式同步（PyTorch会自动处理运算依赖，通常无需手动添加）
    # torch.cuda.synchronize()
    
    # 并行计算A的所有元素，逻辑与上面完全一致
    A = compute_next(B)

为什么这样能实现并行？ #

1. 张量操作的并行性：torch.stack会将三个独立的布尔运算结果打包成一个张量，PyTorch会将这三个运算映射到不同的CUDA线程上同步执行，强制实现并行。
2. 减少kernel开销：原代码中每个元素赋值对应一次GPU kernel启动，现在整个compute_next函数只触发一次kernel调用，大幅降低调度开销。
3. 自动依赖管理：PyTorch的异步执行机制会自动保证B = compute_next(A)完成后，才会执行A = compute_next(B)，完全符合你要求的"前一步全部完成后再执行下一步"的逻辑。

进阶优化（可选） #

如果需要极致的性能控制，可以编写自定义CUDA核，直接控制CUDA线程的分配：

1. 编写CUDA核代码（compute_kernel.cu）：

#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void compute_next_kernel(const bool* input, bool* output) {
    // 每个线程负责计算一个输出元素
    int idx = threadIdx.x;
    switch(idx) {
        case 0:
            output[0] = (input[0] || input[1]) && input[2];
            break;
        case 1:
            output[1] = input[0] || input[1] || input[2];
            break;
        case 2:
            output[2] = input[0] && input[1] && input[2];
            break;
    }
}

torch::Tensor compute_next_cuda(torch::Tensor input) {
    torch::Tensor output = torch::empty_like(input);
    // 启动1个block，3个线程（对应三个输出元素）
    compute_next_kernel<<<1, 3>>>(input.data_ptr<bool>(), output.data_ptr<bool>());
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("compute_next", &compute_next_cuda, "Parallel compute next step (CUDA)");
}

2. 编译为PyTorch扩展后，在Python中调用即可。不过对于当前场景，原生张量操作的性能已经足够，自定义核仅适合极端优化需求。

内容的提问来源于stack exchange，提问作者mattroos