为何GPU张量执行torch.nonzero会触发CudaMemcpyAsync调用？

以下为代码示例：

import torch
import torch.cuda.nvtx as nvtx_cuda

A = torch.rand(size=(1024*4, 1024*4), device="cuda")
B = torch.rand(size=(1024*4, 1024*4), device="cuda")
C = torch.rand(size=(1024*4, 1024*4), device="cuda")
D = torch.rand(size=(1024*4, 1024*4), device="cuda")
E = torch.rand(size=(1024*4, 1024*4), device="cuda")
F = torch.rand(size=(1024*4, 1024*4), device="cuda")

torch.cuda.cudart().cudaProfilerStart()

idx_mask = torch.rand(1024, device="cuda")
active_idx = torch.rand(1024, device="cuda")
nvtx_cuda.range_push("S0")
active_idx = torch.nonzero(idx_mask, as_tuple=True)
nvtx_cuda.range_pop()

nvtx_cuda.range_push("S1")
C = A.matmul(B)
F = D.matmul(E)
nvtx_cuda.range_pop()

print(active_idx)

torch.cuda.cudart().cudaProfilerStop()

针对上述代码，我十分好奇为何active_idx = torch.nonzero(idx_mask, as_tuple=True)会调用CudaMemcpyAsync。

我的观点如下：

• idx_mask和active_idx初始均位于GPU上。
• 因此所有操作都应在GPU上执行，独立于CPU。
• 我推测Python层的torch.nonzero在底层会调用C++实现的CUDA内核来执行非零值查找操作，但内核总归需要由CPU线程启动，例如GPU需向CPU传递指针等信息，再由CPU启动内核。
• 我还做了一个调用GEMM内核的示例，该内核并未触发任何内存操作（同样是对GPU上的张量执行GEMM运算）。
• 综上，我希望消除GPU到CPU的内存操作，让该操作独立于CPU运行，或者至少将其改为非阻塞操作，即消除cudaSyncThread（对应图中的绿色区块）。

此为非零值操作的示意图。

此为GEMM内核的示意图。

请注意：图2中的GEMM内核并未调用cudaMemasync。

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