PyTorch中的C++/CUDA扩展
这是一个编写PyTorch C++/CUDA扩展的示例。相关教程请参见此处。
本仓库演示了如何编写一个名为extension_cpp.ops.mymuladd
的示例自定义操作,该操作同时具有自定义的CPU和CUDA内核。
本仓库中的示例适用于PyTorch 2.4+版本。
构建方法:
pip install .
测试方法:
python test/test_extension.py
对比Python、C++和CUDA性能的方法:
python test/benchmark.py