allow user passing relative paths in include_dirs within setuptools.setup (#38264)

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/38264

Test Plan: Imported from OSS

Differential Revision: D21509277

Pulled By: glaringlee

fbshipit-source-id: b0bc17d375a89b96b1bdacde5987b4f4baa9468e

test/cpp_extensions/extension.cpp

@@ -1,5 +1,8 @@
 #include <torch/extension.h>
 
+// test include_dirs in setuptools.setup with relative path
+#include <tmp.h>
+
 torch::Tensor sigmoid_add(torch::Tensor x, torch::Tensor y) {
   return x.sigmoid() + y.sigmoid();
 }

test/cpp_extensions/self_compiler_include_dirs_test/tmp.h

@@ -0,0 +1,4 @@
+/* This is an empty header file that is used to test self.compiler.include_dirs with
+ * relative path. To test, just inlucde this file into your source code and
+ * add 'self_compiler_include_dirs_test' into 'include_dirs' in setuptools.setup.
+ */

test/cpp_extensions/setup.py

@@ -59,4 +59,5 @@ setup(
     name='torch_test_cpp_extension',
     packages=['torch_test_cpp_extension'],
     ext_modules=ext_modules,
+    include_dirs='self_compiler_include_dirs_test',
     cmdclass={'build_ext': BuildExtension.with_options(use_ninja=USE_NINJA)})

torch/utils/cpp_extension.py

@@ -359,6 +359,13 @@ class BuildExtension(build_ext, object):
                     ['--compiler-options', "'-fPIC'"] +
                     cflags + _get_cuda_arch_flags(cflags))
 
+        def convert_to_absolute_paths_inplace(paths):
+            # Helper function. See Note [Absolute include_dirs]
+            if paths is not None:
+                for i in range(len(paths)):
+                    if not os.path.isabs(paths[i]):
+                        paths[i] = os.path.abspath(paths[i])
+
         def unix_wrap_single_compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
             # Copy before we make any modifications.
             cflags = copy.deepcopy(extra_postargs)
@@ -407,6 +414,10 @@ class BuildExtension(build_ext, object):
             # Use absolute path for output_dir so that the object file paths
             # (`objects`) get generated with absolute paths.
             output_dir = os.path.abspath(output_dir)
+
+            # See Note [Absolute include_dirs]
+            convert_to_absolute_paths_inplace(self.compiler.include_dirs)
+
             _, objects, extra_postargs, pp_opts, _ = \
                 self.compiler._setup_compile(output_dir, macros,
                                              include_dirs, sources,
@@ -536,6 +547,15 @@ class BuildExtension(build_ext, object):
             if not self.compiler.initialized:
                 self.compiler.initialize()
             output_dir = os.path.abspath(output_dir)
+
+            # Note [Absolute include_dirs]
+            # Convert relative path in self.compiler.include_dirs to absolute path if any,
+            # For ninja build, the build location is not local, the build happens
+            # in a in script created build folder, relative path lost their correctness.
+            # To be consistent with jit extension, we allow user to enter relative include_dirs
+            # in setuptools.setup, and we convert the relative path to absolute path here
+            convert_to_absolute_paths_inplace(self.compiler.include_dirs)
+
             _, objects, extra_postargs, pp_opts, _ = \
                 self.compiler._setup_compile(output_dir, macros,
                                              include_dirs, sources,