Skip to content

Latest commit

 

History

History
56 lines (50 loc) · 1.86 KB

File metadata and controls

56 lines (50 loc) · 1.86 KB

C++/CUDA Extensions in PyTorch

Table of Contents

  1. C++/CUDA Extensions in PyTorch
  2. Installation
  3. To build
  4. To test
  5. To benchmark Python vs. C++ vs. CUDA
  6. How to Contribute
  7. Authors An example of writing a C++/CUDA extension for PyTorch. See here for the accompanying tutorial. This repo demonstrates how to write an example extension_cpp.ops.mymuladd custom op that has both custom CPU and CUDA kernels.

The examples in this repo work with PyTorch 2.4+.

Installation

To use the following extension you must have the following:

  • PyTorch 2.4+: Must have a compatible version installed.
  • CUDA Toolkit: Required for building and running the CUDA kernels if using GPU acceleration.
  • GCC or Clang: Necessary for compiling the C++ extension.
  • Python 3.7+: Check you are using a compatible version of Python.

You can install the required packages using:

pip install torch
pip install -r requirements.txt  # If any additional requirements are specified

To build:

pip install .

To test:

python test/test_extension.py

To benchmark Python vs. C++ vs. CUDA:

python test/benchmark.py

How to Contribute

Contributions are always welcome, to contribute remember to do the following:

1. Fork the Repository

  • Click the "Fork" button at the top of this repository to create your own copy.

2. Clone Your Fork

git clone https://github.com/<your-username>/extension-cpp.git
cd extension-cpp

Authors

Peter Goldsborough, Richard Zou