如何将cudaArray转换为LibTorch张量？

cudaArray属于纹理/表面内存，并非线性设备内存，而torch::from_blob()仅支持绑定线性内存（如cudaMalloc分配的内存），因此无法直接通过强转指针的方式将cudaArray转为Torch张量。以下是两种可行的实现方案：

方案一：通过内存拷贝转换（简单直接） #

将cudaArray的数据拷贝到线性CUDA内存后，再绑定为Torch张量。适合对实现复杂度要求低的场景，仅需额外一次设备内内存拷贝。

代码示例（适配你的互操作流程） #

// 注册并映射D3D11-CUDA互操作资源
cudaGraphicsResource* cu_arr_interop;
// ... 此处省略cu_arr_interop的初始化代码 ...
cudaGraphicsMapResources(1, &cu_arr_interop, nullptr);

// 获取映射后的cudaArray
cudaArray* cu_arr;
cudaGraphicsSubResourceGetMappedArray(&cu_arr, cu_arr_interop, 0, 0);

// 查询cudaArray的尺寸与格式属性
cudaChannelFormatDesc desc;
cudaArrayGetChannelDesc(&desc, cu_arr);
int width, height;
cudaArrayGetDimensions(&width, &height, nullptr, cu_arr);

// 1. 创建匹配尺寸与类型的空CUDA张量（自动分配线性内存）
auto tensor_options = torch::TensorOptions()
    .dtype(torch::kFloat32) // 根据cudaArray的实际格式调整（如desc.f对应的数据类型）
    .device(torch::kCUDA, 0);
// 假设纹理为HWC格式的RGB图像，根据实际维度调整
auto tensor_in = torch::empty({height, width, 3}, tensor_options);

// 2. 将cudaArray数据拷贝到张量的线性内存
cudaMemcpy2DFromArray(
    tensor_in.data_ptr(),                // 目标线性内存指针
    tensor_in.stride(0) * sizeof(float), // 目标行步长（字节数）
    cu_arr,                              // 源cudaArray
    0, 0,                                // 源数据起始坐标(x,y)
    width * 3 * sizeof(float),           // 每行拷贝的字节数
    height,                              // 拷贝的行数
    cudaMemcpyDeviceToDevice             // 设备到设备拷贝
);

// 执行模型推理
auto tensor_out = module.forward({tensor_in}).toTensor();

// 后续结果处理...
cudaGraphicsUnmapResources(1, &cu_arr_interop, nullptr);

方案二：通过CUDA内核直接写入（性能优化） #

如果已经在使用cudaSurfaceObject_t处理cudaArray，可以自定义CUDA内核，直接将纹理数据写入Torch张量的线性内存，避免显式的cudaMemcpy调用，同时可利用纹理读取的缓存优化。

代码示例 #

// 注册并映射资源（同方案一）
cudaGraphicsResource* cu_arr_interop;
cudaGraphicsMapResources(1, &cu_arr_interop, nullptr);
cudaArray* cu_arr;
cudaGraphicsSubResourceGetMappedArray(&cu_arr, cu_arr_interop, 0, 0);

// 获取cudaArray尺寸
int width, height;
cudaArrayGetDimensions(&width, &height, nullptr, cu_arr);

// 创建surface对象用于读取cudaArray
cudaSurfaceObject_t surf_obj;
cudaResourceDesc res_desc = {};
res_desc.resType = cudaResourceTypeArray;
res_desc.res.array.array = cu_arr;
cudaCreateSurfaceObject(&surf_obj, &res_desc);

// 创建空CUDA张量
auto tensor_options = torch::TensorOptions()
    .dtype(torch::kFloat32)
    .device(torch::kCUDA, 0);
auto tensor_in = torch::empty({height, width, 3}, tensor_options);

// 启动CUDA内核，将surface数据写入张量
dim3 block(16, 16);
dim3 grid((width + block.x - 1)/block.x, (height + block.y - 1)/block.y);
copy_surface_to_tensor<<<grid, block>>>(surf_obj, tensor_in.data_ptr<float>(), width, height);

// 同步内核（若使用CUDA流，需同步对应流）
cudaDeviceSynchronize();

// 销毁surface对象
cudaDestroySurfaceObject(surf_obj);

// 模型推理与后续处理...
cudaGraphicsUnmapResources(1, &cu_arr_interop, nullptr);

对应的CUDA内核实现 #

__global__ void copy_surface_to_tensor(cudaSurfaceObject_t surf_obj, float* output, int width, int height) {
    const int x = blockIdx.x * blockDim.x + threadIdx.x;
    const int y = blockIdx.y * blockDim.y + threadIdx.y;
    
    if (x >= width || y >= height) return;

    // 读取surface像素（假设为float4格式，对应RGBA，根据实际格式调整）
    float4 pixel = surf2Dread<float4>(surf_obj, x * sizeof(float4), y);
    
    // 将RGB数据写入HWC格式的张量
    const int idx = y * width * 3 + x * 3;
    output[idx] = pixel.x;
    output[idx + 1] = pixel.y;
    output[idx + 2] = pixel.z;
}

关键注意事项 #

• 格式匹配：必须确保cudaArray的数据格式（如通道数、数据类型）与Torch张量的dtype、维度顺序（HWC/CHW）完全一致，否则会出现数据错乱。
• 资源生命周期：在拷贝或内核执行期间，必须保证cudaGraphicsResource处于映射状态，操作完成后及时解映射。
• 零拷贝限制：cudaArray为纹理内存布局，不满足LibTorch张量对线性内存的要求，因此无法实现真正的零拷贝，必须转换为线性内存后才能被LibTorch接管。

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