为何GPU上的Tensor索引会触发GPU->CPU内存拷贝与cudaStreamSync？

问题背景 #

以下是测试代码：

import torch
import torch.cuda.nvtx as nvtx_cuda
torch.cuda.cudart().cudaProfilerStart()

idx = torch.tensor([2,3,4], device="cuda")

tensor_0 = torch.rand(6, device="cuda")
tensor_1 = torch.rand(6, device="cuda")

nvtx_cuda.range_push("S0")
for i in idx:
    tensor_0[i] = tensor_0[i] + tensor_1[i]
nvtx_cuda.range_pop()

nvtx_cuda.range_push("S1")
for i in range(4):
    tensor_0[i] = tensor_0[i] + tensor_1[i]
nvtx_cuda.range_pop()

idx = [2,3,4]
nvtx_cuda.range_push("S2")
for i in idx:
    tensor_0[i] = tensor_0[i] + tensor_1[i]
nvtx_cuda.range_pop()


torch.cuda.cudart().cudaProfilerStop()

通过Nsys profiling得到结果：

观察结果可知：

• S0场景下，尽管所有张量都在GPU上，却触发了GPU→CPU内存拷贝与cudaStreamSync操作
• S1和S2场景下无内存拷贝，仅执行设备间（DtoD）操作

问题解答 #

1. 为何Tensor类型的idx会触发GPU→CPU拷贝与同步？ #

当你在Python的for循环中直接迭代GPU张量idx时，Python解释器需要获取张量中的每个元素值来完成循环逻辑，而Python的循环是在CPU上执行的。每一次迭代i in idx都会隐式触发类似i.item()的操作——这个操作会强制将GPU张量中的单个元素同步拷贝到CPU，转换为Python整数后再参与后续的索引操作。

• 同步操作cudaStreamSync是为了等待GPU端数据准备完成，确保拷贝到CPU的数据是最新的
• 这种逐元素的跨设备拷贝和同步，就是你在Nsys中看到的额外内存开销来源

2. S1和S2场景为何无内存拷贝？它们的idx在GPU上吗？ #

• S1场景：range(4)生成的是CPU端的整数序列，循环中的i是Python原生整数。用整数索引GPU张量时，PyTorch会直接生成对应的GPU内核操作（DtoD）：索引值在CPU上，但框架会将其作为参数传递给GPU内核，整个计算逻辑完全在GPU上执行，不需要跨设备的数据搬运或同步。
• S2场景：idx是CPU端的Python列表，元素均为Python整数。和S1逻辑一致，PyTorch可以直接将整数索引转换为GPU端的内核操作，不需要将列表数据拷贝到GPU，也不需要从GPU拷贝数据到CPU，因此不会触发内存拷贝。

简言之：S1和S2的idx都在CPU上，但因为是Python原生整数/整数列表，PyTorch能直接将索引逻辑映射为纯GPU操作；而S0的idx是GPU张量，Python循环迭代时必须逐元素拷贝到CPU，因此产生了额外开销。

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