Support for custom operator incremental compilation. (#67659)

risemeup1 · web-flow · commit 68cd8a021974 · 2024-08-25T15:24:04.000+08:00
* fix

* fix
diff --git a/python/paddle/utils/cpp_extension/cpp_extension.py b/python/paddle/utils/cpp_extension/cpp_extension.py
@@ -29,6 +29,7 @@
 from .extension_utils import (
     add_compile_flag,
     find_cuda_home,
+    find_ccache_home,
     find_rocm_home,
     normalize_extension_kwargs,
 )
@@ -85,6 +86,8 @@
     ROCM_HOME = find_rocm_home()
     CUDA_HOME = ROCM_HOME
 
+CCACHE_HOME = find_ccache_home()
+
 
 def setup(**attr: Any) -> None:
     """
@@ -455,8 +458,11 @@ def unix_custom_compile_single_file(
                             ROCM_HOME is not None
                         ), "Not found ROCM runtime, \
                             please use `export ROCM_PATH= XXX` to specify it."
-
-                        hipcc_cmd = os.path.join(ROCM_HOME, 'bin', 'hipcc')
+                        if CCACHE_HOME is not None:
+                            hipcc_cmd = os.path.join(ROCM_HOME, 'bin', 'hipcc')
+                            hipcc_cmd = f'{CCACHE_HOME} {hipcc_cmd}'
+                        else:
+                            hipcc_cmd = os.path.join(ROCM_HOME, 'bin', 'hipcc')
                         self.set_executable('compiler_so', hipcc_cmd)
                         # {'nvcc': {}, 'cxx: {}}
                         if isinstance(cflags, dict):
@@ -466,17 +472,27 @@ def unix_custom_compile_single_file(
                             CUDA_HOME is not None
                         ), "Not found CUDA runtime, \
                             please use `export CUDA_HOME= XXX` to specify it."
-
-                        nvcc_cmd = os.path.join(CUDA_HOME, 'bin', 'nvcc')
+                        if CCACHE_HOME is not None:
+                            nvcc_cmd = os.path.join(CUDA_HOME, 'bin', 'nvcc')
+                            nvcc_cmd = f'{CCACHE_HOME} {nvcc_cmd}'
+                        else:
+                            nvcc_cmd = os.path.join(CUDA_HOME, 'bin', 'nvcc')
                         self.set_executable('compiler_so', nvcc_cmd)
                         # {'nvcc': {}, 'cxx: {}}
                         if isinstance(cflags, dict):
                             cflags = cflags['nvcc']
 
                     cflags = prepare_unix_cudaflags(cflags)
                 # cxx compile Cpp source
-                elif isinstance(cflags, dict):
-                    cflags = cflags['cxx']
+                else:
+                    if CCACHE_HOME is not None:
+                        # self.set_executable('compiler_so', [CCACHE_HOME, *self.executables['compiler_so']])
+                        self.set_executable(
+                            'compiler_so', [CCACHE_HOME, *self.compiler_so]
+                        )
+
+                    if isinstance(cflags, dict):
+                        cflags = cflags['cxx']
 
                 # Note(qili93): HIP require some additional flags for CMAKE_C_FLAGS
                 if core.is_compiled_with_rocm():
diff --git a/python/paddle/utils/cpp_extension/extension_utils.py b/python/paddle/utils/cpp_extension/extension_utils.py
@@ -657,6 +657,37 @@ def create_sym_link_if_not_exist():
         return raw_core_name[:-3]
 
 
+def find_ccache_home():
+    """
+    Use heuristic method to find ccache path
+    """
+    ccache_path = None
+    # step 1. find in $PATH environment variable
+    paths = os.environ.get('PATH', '').split(os.pathsep)
+    for path in paths:
+        ccache_candidate = os.path.join(path, 'ccache')
+        if os.path.exists(ccache_candidate):
+            ccache_path = ccache_candidate
+            break
+
+    # step 2. find ccache path by `which ccache` command
+    if ccache_path is None:
+        which_cmd = 'where' if os.name == 'nt' else 'which'
+        try:
+            ccache_path = (
+                subprocess.check_output([which_cmd, 'ccache']).decode().strip()
+            )
+        except:
+            ccache_path = None
+
+    if ccache_path is None:
+        warning_message = "No ccache found. Please be aware that recompiling all source files may be required. "
+        warning_message += "You can download and install ccache from: https://github.com/ccache/ccache/blob/master/doc/INSTALL.md"
+        warnings.warn(warning_message)
+
+    return ccache_path
+
+
 def find_cuda_home():
     """
     Use heuristic method to find cuda path
