diff --git a/.gitignore b/.gitignore index 9ea9297330e..324477cfac2 100644 --- a/.gitignore +++ b/.gitignore @@ -57,3 +57,4 @@ test/data/linear.pt .ninja_deps .ninja_log compile_commands.json +*.egg-info/ diff --git a/setup.py b/setup.py index eb14af168f7..c99f64f8890 100644 --- a/setup.py +++ b/setup.py @@ -348,6 +348,19 @@ class install(setuptools.command.install.install): def run(self): if not self.skip_build: self.run_command('build_deps') + + # Copy include directories necessary to compile C++ extensions. + def copy_and_overwrite(src, dst): + print('copying {} -> {}'.format(src, dst)) + if os.path.exists(dst): + shutil.rmtree(dst) + shutil.copytree(src, dst) + + copy_and_overwrite('torch/csrc', 'torch/lib/include/torch/csrc/') + copy_and_overwrite('torch/lib/pybind11/include/pybind11/', + 'torch/lib/include/pybind11') + shutil.copy2('torch/torch.h', 'torch/lib/include/torch/torch.h') + setuptools.command.install.install.run(self) @@ -736,20 +749,35 @@ cmdclass = { } cmdclass.update(build_dep_cmds) - if __name__ == '__main__': - setup(name="torch", version=version, - description="Tensors and Dynamic neural networks in Python with strong GPU acceleration", - ext_modules=extensions, - cmdclass=cmdclass, - packages=packages, - package_data={'torch': [ - 'lib/*.so*', 'lib/*.dylib*', 'lib/*.dll', 'lib/*.lib', - 'lib/torch_shm_manager', - 'lib/*.h', - 'lib/include/TH/*.h', 'lib/include/TH/generic/*.h', - 'lib/include/THC/*.h', 'lib/include/THC/generic/*.h', - 'lib/include/ATen/*.h', - ]}, - install_requires=['pyyaml', 'numpy'], - ) + setup( + name="torch", + version=version, + description=("Tensors and Dynamic neural networks in " + "Python with strong GPU acceleration"), + ext_modules=extensions, + cmdclass=cmdclass, + packages=packages, + package_data={ + 'torch': [ + 'lib/*.so*', + 'lib/*.dylib*', + 'lib/*.dll', + 'lib/*.lib', + 'lib/torch_shm_manager', + 'lib/*.h', + 'lib/include/ATen/*.h', + 'lib/include/pybind11/*.h', + 'lib/include/pybind11/detail/*.h', + 'lib/include/TH/*.h', + 'lib/include/TH/generic/*.h', + 'lib/include/THC/*.h', + 'lib/include/THC/generic/*.h', + 'lib/include/torch/csrc/*.h', + 'lib/include/torch/csrc/autograd/*.h', + 'lib/include/torch/csrc/jit/*.h', + 'lib/include/torch/csrc/utils/*.h', + 'lib/include/torch/torch.h', + ] + }, + install_requires=['pyyaml', 'numpy'], ) diff --git a/test/cpp_extensions/extension.cpp b/test/cpp_extensions/extension.cpp new file mode 100644 index 00000000000..02ae7f11910 --- /dev/null +++ b/test/cpp_extensions/extension.cpp @@ -0,0 +1,30 @@ +#include + +using namespace at; + +Tensor sigmoid_add(Tensor x, Tensor y) { + return x.sigmoid() + y.sigmoid(); +} + +struct MatrixMultiplier { + MatrixMultiplier(int A, int B) { + tensor_ = CPU(kDouble).ones({A, B}); + } + Tensor forward(Tensor weights) { + return tensor_.mm(weights); + } + Tensor get() const { + return tensor_; + } + +private: + Tensor tensor_; +}; + +PYBIND11_MODULE(torch_test_cpp_extensions, m) { + m.def("sigmoid_add", &sigmoid_add, "sigmoid(x) + sigmoid(y)"); + py::class_(m, "MatrixMultiplier") + .def(py::init()) + .def("forward", &MatrixMultiplier::forward) + .def("get", &MatrixMultiplier::get); +} diff --git a/test/cpp_extensions/setup.py b/test/cpp_extensions/setup.py new file mode 100644 index 00000000000..121dc272d0f --- /dev/null +++ b/test/cpp_extensions/setup.py @@ -0,0 +1,14 @@ +from setuptools import setup, Extension +import torch.utils.cpp_extension + +ext_modules = [ + Extension( + 'torch_test_cpp_extensions', ['extension.cpp'], + include_dirs=torch.utils.cpp_extension.include_paths(), + language='c++'), +] + +setup( + name='torch_test_cpp_extensions', + ext_modules=ext_modules, + cmdclass={'build_ext': torch.utils.cpp_extension.BuildExtension}) diff --git a/test/run_test.sh b/test/run_test.sh index 96b7b6454ad..8dae5eed55d 100755 --- a/test/run_test.sh +++ b/test/run_test.sh @@ -64,6 +64,16 @@ $PYCMD test_cuda.py $@ echo "Running NCCL tests" $PYCMD test_nccl.py $@ +echo "Running C++ Extensions tests" +cd cpp_extensions +$PYCMD setup.py install --root ./install +previous_pythonpath="$PYTHONPATH" +export PYTHONPATH="$PWD/$(find ./install -name site-packages):$PYTHONPATH" +cd .. +$PYCMD test_cpp_extensions.py $@ +export PYTHONPATH="$previous_pythonpath" +rm -rf cpp_extensions/install + # Skipping test_distributed for Windows because it doesn't have fcntl if [[ "$OSTYPE" != "msys" ]]; then distributed_set_up() { diff --git a/test/test_cpp_extensions.py b/test/test_cpp_extensions.py new file mode 100644 index 00000000000..345eb7e79a2 --- /dev/null +++ b/test/test_cpp_extensions.py @@ -0,0 +1,23 @@ +import torch +import torch_test_cpp_extensions as cpp_extension + +import common + + +class TestCppExtension(common.TestCase): + def test_extension_function(self): + x = torch.randn(4, 4) + y = torch.randn(4, 4) + z = cpp_extension.sigmoid_add(x, y) + self.assertEqual(z, x.sigmoid() + y.sigmoid()) + + def test_extension_module(self): + mm = cpp_extension.MatrixMultiplier(4, 8) + weights = torch.rand(8, 4) + expected = mm.get().mm(weights) + result = mm.forward(weights) + self.assertEqual(expected, result) + + +if __name__ == '__main__': + common.run_tests() diff --git a/torch/torch.h b/torch/torch.h new file mode 100644 index 00000000000..0ed6043ae66 --- /dev/null +++ b/torch/torch.h @@ -0,0 +1,6 @@ +#pragma once + +#include +#include +#include +#include diff --git a/torch/utils/cpp_extension.py b/torch/utils/cpp_extension.py new file mode 100644 index 00000000000..24fa7a1b809 --- /dev/null +++ b/torch/utils/cpp_extension.py @@ -0,0 +1,18 @@ +import os.path + +from setuptools.command.build_ext import build_ext + + +class BuildExtension(build_ext): + """A custom build extension for adding compiler-specific options.""" + + def build_extensions(self): + for extension in self.extensions: + extension.extra_compile_args = ['-std=c++11'] + build_ext.build_extensions(self) + + +def include_paths(): + here = os.path.abspath(__file__) + torch_path = os.path.dirname(os.path.dirname(here)) + return [os.path.join(torch_path, 'lib', 'include')]