Enable ccache to accelerate contrib compilation

python/tvm/contrib/cc.py

@@ -33,6 +33,10 @@ def _is_linux_like():
     )
 
 
+def _is_windows_like():
+    return sys.platform == "win32"
+
+
 def get_cc():
     """Return the path to the default C/C++ compiler.
 
@@ -58,7 +62,7 @@ def get_cc():
     return None
 
 
-def create_shared(output, objects, options=None, cc=None):
+def create_shared(output, objects, options=None, cc=None, cwd=None, ccache_env=None):
     """Create shared library.
 
     Parameters
@@ -74,13 +78,19 @@ def create_shared(output, objects, options=None, cc=None):
 
     cc : Optional[str]
         The compiler command.
+
+    cwd : Optional[str]
+        The urrent working directory.
+
+    ccache_env : Optional[Dict[str, str]]
+        The environment variable for ccache. Set `None` to disable ccache by default.
     """
     cc = cc or get_cc()
 
     if _is_linux_like():
-        _linux_compile(output, objects, options, cc, compile_shared=True)
-    elif sys.platform == "win32":
-        _windows_compile(output, objects, options)
+        _linux_compile(output, objects, options, cc, cwd, ccache_env, compile_shared=True)
+    elif _is_windows_like():
+        _windows_compile(output, objects, options, cwd, ccache_env)
     else:
         raise ValueError("Unsupported platform")
 
@@ -133,7 +143,7 @@ def create_staticlib(output, inputs, ar=None):
         raise ValueError("Unsupported platform")
 
 
-def create_executable(output, objects, options=None, cc=None):
+def create_executable(output, objects, options=None, cc=None, cwd=None, ccache_env=None):
     """Create executable binary.
 
     Parameters
@@ -149,13 +159,19 @@ def create_executable(output, objects, options=None, cc=None):
 
     cc : Optional[str]
         The compiler command.
+
+    cwd : Optional[str]
+        The urrent working directory.
+
+    ccache_env : Optional[Dict[str, str]]
+        The environment variable for ccache. Set `None` to disable ccache by default.
     """
     cc = cc or get_cc()
 
     if _is_linux_like():
-        _linux_compile(output, objects, options, cc)
+        _linux_compile(output, objects, options, cc, cwd, ccache_env)
     elif sys.platform == "win32":
-        _windows_compile(output, objects, options)
+        _windows_compile(output, objects, options, cwd, ccache_env)
     else:
         raise ValueError("Unsupported platform")
 
@@ -269,7 +285,9 @@ def _fcompile(outputs, objects, options=None):
     return _fcompile
 
 
-def _linux_compile(output, objects, options, compile_cmd, compile_shared=False):
+def _linux_compile(
+    output, objects, options, compile_cmd, cwd=None, ccache_env=None, compile_shared=False
+):
     cmd = [compile_cmd]
     if compile_cmd != "nvcc":
         if compile_shared or output.endswith(".so") or output.endswith(".dylib"):
@@ -288,7 +306,15 @@ def _linux_compile(output, objects, options, compile_cmd, compile_shared=False):
         cmd += objects
     if options:
         cmd += options
-    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    env = None
+    if ccache_env is not None:
+        if shutil.which("ccache"):
+            cmd.insert(0, "ccache")
+            env = os.environ.copy()
+            env.update(ccache_env)
+        else:
+            raise ValueError("ccache not found")
+    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=cwd, env=env)
     (out, _) = proc.communicate()
     if proc.returncode != 0:
         msg = "Compilation error:\n"
@@ -297,7 +323,7 @@ def _linux_compile(output, objects, options, compile_cmd, compile_shared=False):
         raise RuntimeError(msg)
 
 
-def _windows_compile(output, objects, options):
+def _windows_compile(output, objects, options, cwd=None, ccache_env=None):
     cmd = ["clang"]
     cmd += ["-O2"]
 
@@ -312,9 +338,19 @@ def _windows_compile(output, objects, options):
     cmd += objects
     if options:
         cmd += options
+    env = None
+    if ccache_env is not None:
+        if shutil.which("ccache"):
+            cmd.insert(0, "ccache")
+            env = os.environ.copy()
+            env.update(ccache_env)
+        else:
+            raise ValueError("ccache not found")
 
     try:
-        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+        proc = subprocess.Popen(
+            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=cwd, env=env
+        )
         (out, _) = proc.communicate()
     except FileNotFoundError:
         raise RuntimeError(

tests/python/contrib/test_ccache.py

@@ -0,0 +1,79 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Test contrib.cc with ccache"""
+import os
+import pytest
+import shutil
+import tempfile
+import tvm
+from tvm.contrib.cc import create_shared, create_executable, _is_linux_like, _is_windows_like
+
+
+def _src_gen(text):
+    return """
+#include <iostream>
+
+int main() {
+    std::cout << "text";
+    return 0;
+}""".replace(
+        "text", text
+    )
+
+
+def _compile(f_create, text, output):
+    with tempfile.TemporaryDirectory() as temp_dir:
+        src_path = os.path.join(temp_dir, "src.cpp")
+        with open(src_path, "w", encoding="utf-8") as file:
+            file.write(_src_gen(text))
+        log_path = os.path.join(temp_dir, "log.txt")
+        ccache_env = {
+            "CCACHE_COMPILERCHECK": "content",
+            "CCACHE_LOGFILE": log_path,
+        }
+        f_create(output, ["src.cpp"], ["-c"], cwd=temp_dir, ccache_env=ccache_env)
+        with open(log_path, "r", encoding="utf-8") as file:
+            log = file.read()
+        return log
+
+
+@pytest.mark.skipif(shutil.which("ccache") is None, reason="ccache not installed")
+def test_shared():
+    if _is_linux_like():
+        _ = _compile(create_shared, "shared", "main.o")
+        log = _compile(create_shared, "shared", "main.o")
+        assert "Succeeded getting cached result" in log
+    elif _is_windows_like():
+        _ = _compile(create_shared, "shared", "main.obj")
+        log = _compile(create_shared, "shared", "main.obj")
+        assert "Succeeded getting cached result" in log
+
+
+@pytest.mark.skipif(shutil.which("ccache") is None, reason="ccache not installed")
+def test_executable():
+    if _is_linux_like():
+        _ = _compile(create_executable, "executable", "main")
+        log = _compile(create_executable, "executable", "main")
+        assert "Succeeded getting cached result" in log
+    elif _is_windows_like():
+        _ = _compile(create_executable, "executable", "main.exe")
+        log = _compile(create_executable, "executable", "main.exe")
+        assert "Succeeded getting cached result" in log
+
+
+if __name__ == "__main__":
+    tvm.testing.main()