add _GLIBCXX_USE_CXX11_ABI=0 to cpp_extensions when binary built pytorch is detected

soumith · soumith · commit f08f222db3b2 · 2018-07-25T17:22:33.000-07:00
diff --git a/torch/utils/cpp_extension.py b/torch/utils/cpp_extension.py
@@ -64,6 +64,9 @@ def _find_cuda_home():
 # it the below pattern.
 BUILT_FROM_SOURCE_VERSION_PATTERN = re.compile(r'\d+\.\d+\.\d+\w+\+\w+')
 
+def is_binary_build():
+    return not BUILT_FROM_SOURCE_VERSION_PATTERN.match(torch.version.__version__)
+
 
 def check_compiler_abi_compatibility(compiler):
     '''
@@ -77,7 +80,7 @@ def check_compiler_abi_compatibility(compiler):
         False if the compiler is (likely) ABI-incompatible with PyTorch,
         else True.
     '''
-    if BUILT_FROM_SOURCE_VERSION_PATTERN.match(torch.version.__version__):
+    if not is_binary_build():
         return True
     try:
         check_cmd = '{}' if sys.platform == 'win32' else '{} --version'
@@ -134,6 +137,7 @@ def build_extensions(self):
         self._check_abi()
         for extension in self.extensions:
             self._define_torch_extension_name(extension)
+            self._add_gnu_abi_flag_if_binary(extension)
 
         # Register .cu and .cuh as valid source extensions.
         self.compiler.src_extensions += ['.cu', '.cuh']
@@ -266,6 +270,21 @@ def _define_torch_extension_name(self, extension):
         else:
             extension.extra_compile_args.append(define)
 
+    def _add_gnu_abi_flag_if_binary(self, extension):
+        # If the version string looks like a binary build,
+        # we know that PyTorch was compiled with gcc 4.9.2.
+        # if the extension is compiled with gcc >= 5.1,
+        # then we have to define _GLIBCXX_USE_CXX11_ABI=0
+        # so that the std::string in the API is resolved to
+        # non-C++11 symbols
+        define = '-D_GLIBCXX_USE_CXX11_ABI=0'
+        if is_binary_build():
+            if isinstance(extension.extra_compile_args, dict):
+                for args in extension.extra_compile_args.values():
+                    args.append(define)
+            else:
+                extension.extra_compile_args.append(define)
+
 
 def CppExtension(name, sources, *args, **kwargs):
     '''
@@ -785,6 +804,9 @@ def _write_ninja_file(path,
     common_cflags = ['-DTORCH_EXTENSION_NAME={}'.format(name)]
     common_cflags += ['-I{}'.format(include) for include in includes]
 
+    if is_binary_build():
+        common_cflags += ['-D_GLIBCXX_USE_CXX11_ABI=0']
+
     cflags = common_cflags + ['-fPIC', '-std=c++11'] + extra_cflags
     if sys.platform == 'win32':
         from distutils.spawn import _nt_quote_args
