如何查找 torch.bmm 的 c++ 源代码 pytorch.

How to find c++ source code of torch.bmm of pytorch

本文关键字:bmm c++ 源代码 pytorch torch 何查找 查找      更新时间:2023-10-16

我很难找到torch.bmm()的源代码,这是在 https://pytorch.org/cppdocs/api/function_namespaceat_1aac51f71f807ca70fd210814114520c34.html#exhale-function-namespaceat-1aac51f71f807ca70fd210814114520c34 中定义的。

我相信它位于命名空间中,因为它在其他地方被引用为at::bmm。 我搜索的是:

  1. 目录 Aten https://github.com/pytorch/pytorch/tree/34877448216149024f44cbcab830169fdb2fa7fb/aten/src/ATen
  2. caffe2 的目录 https://github.com/pytorch/pytorch/tree/74b65c32be68b15dc7c9e8bb62459efbfbde33d8/caffe2/core
  3. 在 GitHub 中使用关键字 C++ 文件的 BMM 直接搜索

但一无所获。有没有系统的方法可以在这么大的项目中定位一个功能(在这种情况下,bmm(?

bmm本身没有(单一(来源。来自ATen的自述文件:

ATen "本机"函数

是向 ATen 添加运算符和函数的现代机制(与通过 TH/THC cwrap 元数据绑定的传统函数相比,它们是"本机"函数(。本机函数在 native_functions.yaml 中声明,并在此目录的一个 cpp 文件中定义了实现。

bmm在 aten\src\ATenativeative_functions.yaml 中声明:

- func: bmm(Tensor self, Tensor mat2) -> Tensor
use_c10_dispatcher: full
variants: function, method
dispatch:
CPU: bmm_cpu
CUDA: bmm_cuda
SparseCPU: bmm_sparse_cpu
SparseCUDA: bmm_sparse_cuda
supports_named_tensor: True

实现(如bmm_cpu(可以在atensrcATennativeLinearAlgebra.cpp中找到。