How to find the c++/cuda implementation of specific operators in pytorch source code

@shuokay the main reason finding the kernel you’re looking for from there is a pain is because op.call(self, other) uses the pytorch dispatcher, and dynamally dispatches to the right kernel (more on the dispatcher here: Let’s talk about the PyTorch dispatcher : ezyang’s blog)

The source of truth (well, 99% of the time) for every ATen operator and the names of its CPU/CUDA kernels in the codebase is in native_functions.yaml: https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml

For div.Tensor: I happen to know that’s it’s implemented as a structured kernel (more on those here Codegen and Structured Kernels · pytorch/pytorch Wiki · GitHub).

I can tell because of the structured_delegate key in native_functions.yaml here: https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml#L2058C3-L2058C22

As a structured kernel, the op aten::div.Tensor has a “meta” function that performs shape error checks and computes the output shape defined here: https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/BinaryOps.cpp#L174

And it has an implementation that uses TensorIterator, defined here: https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/BinaryOps.cpp#L448 (from there you should be able to grep for each function name to find it).

3 Likes