Structure of Pytorch

1

I’ve made some pull requests to Pytorch’s C++ / CUDA functions, but I was never able to fully comprehend the general structure of Pytorch to the details. I’ve dug into the directory many times, but after a attempts, I thought the best would be to hear from the hoarse’s mouth.

I want to clairfy the general function call stack, from when the end user uses nn.Conv2d all the way down to when cudnnConvolutionForward is called.

nn.Conv2d -> ??? -> cudnnConvolutionForward

Here is what I’ve figured out so far:

  • nn.Conv2d (defined in torch/nn/modules/conv.py) is just a thin wrapper around F.conv2d
  • F.conv2d (defined in torch.nn.functional.py) just calls _add_docstr(torch.conv2d, "really long doc string...") to add docstring.
  • _add_docstr (defined torch/csrc/Module.cpp) function is a C++ function. Looking at the actual definition of the function THPModule_addDocStr (defined in the same file), all it does is check if ml_doc field is occupied, and if not it adds the "really long doc..." to the field. So _add_docstr doesn’t define the function. As the name suggests, it only adds docstring to the pre-defined function.
    • Pytorch seems to use the standard way to extends the CPython’s interpreter described here.
  • This makes me believe that torch.conv2d is another C++ function that is defined somewhere. Here where I’m lost
  • The only fucnction that calls cudnnConvolutionForward is raw_cudnn_convolution_forward_out (defined pytorch/aten/src/ATen/native/cudnn/Conv.cpp). So I assume calling nnConv2d eventually calls raw_cudnn_convolution_forward_out. So I climb up the call stack starting from raw_cudnn_convolution_forward_out.
  • After climbing up a few functions, we reach conv2d (defined in aten/src/ATen/native/Convolution.cpp). conv2d is called in many functions, but I couldn’t find a function that directly adds conv2d to CPython’s interpreter.

To summarize, from top down, I have:

nn.Conv2d -> F.conv2d -> torch.conv2d -> ???

And from bottom up I have:

??? -> conv2d -> convolution -> _convolution-> cudnn_convolution-> cudnn_convolution_forward-> raw_cudnn_convolution_forward_out-> cudnnConvolutionForward

My assumption is torch.conv2d is conv2d defined in aten/src/ATen/native/Convolution.cpp. But this is a specualtion. Please help me out how the two branches connect. Thanks!

1 Like
2

I have an old blog post describing how this works from a “where does my PyTorch function come from?” perspective. The key is the native_functions.yaml. From that the dispatching works to Convolution.cpp via autogenerated wrappers (torch/autograd/generated for the Python bindings, and also to go from at::… to at::native::…).

Best regards

Thomas

2 Likes
3

Your blog post was extremely helpful. Thanks