This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 1994f402e6 Enable ccache to accelerate contrib compilation (#16176)
1994f402e6 is described below
commit 1994f402e69281a53df37a77acf935798eb856bb
Author: Yaxing Cai <[email protected]>
AuthorDate: Thu Nov 30 14:25:22 2023 -0800
Enable ccache to accelerate contrib compilation (#16176)
This PR adds the interface for ccache in `contrib.cc` to enable ccache when
creating libs or exectuables.
---
python/tvm/contrib/cc.py | 58 +++++++++++++++++++++------
tests/python/contrib/test_ccache.py | 79 +++++++++++++++++++++++++++++++++++++
2 files changed, 126 insertions(+), 11 deletions(-)
diff --git a/python/tvm/contrib/cc.py b/python/tvm/contrib/cc.py
index ad6a82c49c..918e3c8f72 100644
--- a/python/tvm/contrib/cc.py
+++ b/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 cross_compiler(
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(
diff --git a/tests/python/contrib/test_ccache.py
b/tests/python/contrib/test_ccache.py
new file mode 100644
index 0000000000..c46e72c648
--- /dev/null
+++ b/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
+
+
[email protected](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
+
+
[email protected](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()
