Skip to main content
Back to path

Welcome to a Grasp lesson

Create a free account to keep studying and create your own study path.

Lesson 78

Building PyTorch Extensions with `setup.py`

Recommended resources

Introduction

In our last lesson, you learned how to bridge the gap between Python and CUDA by wrapping a kernel in a C++ extension using PyTorch's "Just-In-Time" (JIT) compilation features. We saw how load_inline could take C++ and CUDA code as simple strings and dynamically compile them into a callable Python module. This JIT approach is excellent for rapid prototyping and learning.

Now, we will transition to the more robust and standard method used for real-world projects: "Ahead-of-Time" (AOT) compilation. This lesson focuses on using a setup.py script to build, link, and install your custom extension as a formal Python package. Mastering this workflow is essential for creating shareable, high-performance components and is the industry-standard practice for integrating custom operators into larger deep learning codebases. By the end of this lesson, you will be able to package your CUDA kernels into a persistent, importable Python module.

From JIT to AOT: The Role of setup.py

While JIT compilation is convenient, it has its limitations. The build happens every time you run the script (if sources have changed), which can add overhead. More importantly, it doesn't create a distributable package and can become cumbersome as the number of source files grows.

The AOT approach solves this by using a dedicated build script, typically named setup.py. This script leverages Python's standard setuptools library, which is the cornerstone of Python packaging. PyTorch provides specialized helper classes that integrate with setuptools to handle the complexities of compiling mixed C++ and CUDA code, such as invoking nvcc with the correct flags and linking against the necessary PyTorch libraries. This creates a permanent, installed module that your Python code can import without needing to re-compile on every run.

Anatomy of an Extension Build Script

Let's dissect the structure of a typical setup.py file for a mixed C++/CUDA extension. The most effective way to understand this is to walk through a complete, practical example. We'll use a simple scaled_add function as our target.

A well-organized project will have a clear file structure. For our scaled_add extension, it would look like this:

simple_cuda_extension/
├── setup.py                 # The build script
└── src/
    ├── scaled_add.cpp       # The C++ binding and wrapper code
    └── scaled_add_kernel.cu # The CUDA kernel implementation

The resource below provides a complete, step-by-step guide to building this exact extension. We will focus on the setup.py file first.

Practice: Building a Simple CUDA Extension - ApX Machine Learning

This tutorial from ApX Machine Learning provides a clear, self-contained example of building a CUDA extension using setup.py.

First, briefly familiarize yourself with the file structure and objective of the exercise. Then, focus your attention on the setup.py script. The explanation breaks down each critical component of the script.

Let's review the key components from the setup.py file you just read about:

  1. from torch.utils.cpp_extension import BuildExtension, CUDAExtension: This line imports the two most important helpers from PyTorch.

    • CUDAExtension: A specialized version of setuptools.Extension that is pre-configured for CUDA projects. It tells the build system that both C++ (g++/cl.exe) and CUDA (nvcc) compilers will be needed.
    • BuildExtension: A custom command class that replaces the default build_ext command in setuptools. This class contains the logic to correctly invoke the compilers, find PyTorch headers, and link everything together.

  2. setup(...): This is the main function from setuptools.

  3. ext_modules=[CUDAExtension(...)]: This is where you define the extension(s) to be built.

    • name='simple_cuda_extension_cpp': This is arguably the most crucial parameter. It defines the name of the final Python module that you will import.
    • sources=['src/scaled_add.cpp', 'src/scaled_add_kernel.cu']: This is a list of all the source files required to build the extension. The build system is smart enough to compile .cpp files with a C++ compiler and .cu files with nvcc.

  4. cmdclass={'build_ext': BuildExtension}: This directive tells setuptools to use PyTorch's BuildExtension to handle the build process for any extensions defined in ext_modules.

Linking the Build Script and C++ Code

In the previous lesson, you saw the PYBIND11_MODULE macro used to create the Python binding. How does the name you provide in setup.py connect to that C++ code?

The answer lies in another macro: TORCH_EXTENSION_NAME. When you use BuildExtension, it automatically defines TORCH_EXTENSION_NAME as a preprocessor macro whose value is the string you passed to the name argument of CUDAExtension.

Take a look at the C++ binding file from the same tutorial.

Practice: Building a Simple CUDA Extension - ApX Machine Learning

This file provides the C++ interface to our CUDA kernel.

Read through the code in the section Step 2: Write the C++ Binding. Pay special attention to the last few lines, which define the pybind11 module.

You'll see this line:
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)

During compilation of our example, BuildExtension effectively replaces TORCH_EXTENSION_NAME with "simple_cuda_extension_cpp". This elegant mechanism ensures your C++ code is automatically synchronized with your build script, preventing hard-to-debug errors from name mismatches.

Building and Using the Extension

With the source files and setup.py in place, compiling the extension is a single command-line operation. The video below provides a live demonstration of building and testing a C++ extension (without CUDA, but the process is identical).

Pytorch+cpp/cuda extension 教學 tutorial 1 - English CC -

This video from the "AI葵" YouTube channel walks through the AOT build and test process.

Watch the following two segments: First, see how the setup.py file is configured from this timestamp. The speaker then runs pip install . in the terminal, which triggers the build. Next, watch how the newly installed module is imported and tested in a separate Python script from this timestamp. Notice the crucial import torch before importing the custom module.

As demonstrated, you compile and install your extension from the directory containing setup.py using one of two common commands:

  • python setup.py install or pip install .

    • This command compiles the extension and installs it into your Python environment's site-packages directory. Once installed, the module can be imported from any Python script in that environment, just like numpy or torch. This is the standard method for a "final" installation.

  • python setup.py build_ext --inplace

    • This command compiles the extension but places the resulting shared library file (e.g., simple_cuda_extension_cpp.cpython-39-x86_64-linux-gnu.so on Linux) directly in your current directory. This is extremely useful during development because you can re-compile and immediately test without having to reinstall the package each time.

Once built, you can call your high-performance CUDA kernel from Python as if it were a native Python function. The final step is to write a test script to verify its functionality.

Practice: Building a Simple CUDA Extension - ApX Machine Learning

This final section shows how to use and verify the compiled extension.

Please read the code in the section Step 5: Use the Extension in Python. This demonstrates the complete workflow: importing the module, moving data to the GPU, calling the custom function, and, most importantly, verifying its output against a trusted baseline (standard PyTorch operations).

Conclusion

You have now mastered the standard, "Ahead-of-Time" workflow for integrating custom CUDA kernels into PyTorch. This process—organizing source files, writing a setup.py build script, compiling, and importing in Python—is the foundation for building and distributing high-performance custom operators for any serious AI project. While more involved than the JIT methods, it is far more powerful, scalable, and is the expected skill for any GPU kernel engineer.

Key Takeaways:

  • AOT Compilation: Using a setup.py script is the standard method for building permanent, installable Python extensions from C++/CUDA source code.
  • CUDAExtension and BuildExtension: These PyTorch utilities are the core of the build process, automating the complex steps of invoking nvcc and linking against PyTorch libraries.
  • TORCH_EXTENSION_NAME: This macro provides a robust link between your build script's module name and your C++ pybind11 module definition.
  • Build Commands: python setup.py install creates a permanent installation, while python setup.py build_ext --inplace is ideal for rapid development cycles.
  • Verification is Key: Always test your custom kernel's output against a naive, trusted implementation to ensure correctness.

In the next lesson, we will explore Triton, a Python-based language from OpenAI for writing GPU kernels. This will provide a fascinating contrast to the low-level C++/CUDA approach you have been mastering, highlighting the trade-offs between programmer productivity and fine-grained hardware control.