为何当数组定义在GPU上时，np.einsum()可以在GPU上运行？

• CuPy 实现了 NumPy 官方定义的数组接口协议，同时 NumPy 1.16 及以上版本支持 __array_function__ 分发机制：当非原生 NumPy 数组类型的参数传入 NumPy 公开接口时，NumPy 不会强制将数据转为 CPU 数组，而是会调用参数所属类的对应魔法方法，把运算逻辑转发给第三方实现执行。你把 CuPy 生成的 GPU 数组传入 np.einsum 时，实际执行的是 CuPy 封装好的 GPU 版 einsum 逻辑，全程没有 GPU 到 CPU 的数据拷贝，因此性能和直接调用 cp.einsum 完全一致。
• 两组测试的耗时差达到 5 个数量级，本身就是运算在 GPU 执行的直观佐证，CPU 不可能达到毫秒级完成该规模的张量收缩运算。

• 检查返回值类型
  运算结束后直接打印输出数组的类型，只要返回类型为 CuPy 数组，即可证明运算全程在 GPU 执行，没有回拷到 CPU。
  验证代码：

  print(type(c))
  # 输出 <class 'cupy._core.core.ndarray'> 即为 GPU 侧返回结果
  

• 监控GPU硬件状态
  在 Colab 命令行执行 nvidia-smi pmon 实时监控GPU状态，运行测试代码时观察对应进程的SM利用率，运算执行时GPU利用率会瞬间冲高到接近100%，可确认运算下发到了GPU硬件执行。
• 类型不匹配测试
  得到输出数组 c 后，尝试和原生 NumPy CPU 数组做运算，如果触发类型不匹配报错，即可证明 c 是存储在GPU上的 CuPy 数组，对应之前的运算也必然在GPU执行。
  测试代码：

  # 下述代码会触发类型错误，证明c在GPU上
  d = np.ones((500,500,500))
  e = c + d
  

• 等价接口耗时对比
  将代码中的 np.einsum 替换为 cp.einsum 执行相同测试，两者耗时几乎一致，可进一步证明 NumPy 接口实际转发到了 CuPy 的GPU实现。

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