From 0be826ae5110d3b83a7b5ec855c269fc2ba72f3b Mon Sep 17 00:00:00 2001 From: Yaxing Cai Date: Tue, 28 Nov 2023 23:27:25 +0000 Subject: [PATCH 1/3] Enable ccache to accelerate contrib compilation This PR adds the interface for ccache in `contrib.cc` to enable ccache when creating libs or exectuables. --- python/tvm/contrib/cc.py | 40 ++++++++++++++++++++++++++++++---------- 1 file changed, 30 insertions(+), 10 deletions(-) diff --git a/python/tvm/contrib/cc.py b/python/tvm/contrib/cc.py index ad6a82c49c14..50a65ba208d8 100644 --- a/python/tvm/contrib/cc.py +++ b/python/tvm/contrib/cc.py @@ -58,7 +58,7 @@ def get_cc(): return None -def create_shared(output, objects, options=None, cc=None): +def create_shared(output, objects, options=None, cc=None, ccache_env=None): """Create shared library. Parameters @@ -74,13 +74,16 @@ def create_shared(output, objects, options=None, cc=None): cc : Optional[str] The compiler command. + + 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) + _linux_compile(output, objects, options, cc, ccache_env, compile_shared=True) elif sys.platform == "win32": - _windows_compile(output, objects, options) + _windows_compile(output, objects, options, ccache_env) else: raise ValueError("Unsupported platform") @@ -133,7 +136,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, ccache_env=None): """Create executable binary. Parameters @@ -149,13 +152,16 @@ def create_executable(output, objects, options=None, cc=None): cc : Optional[str] The compiler command. + + 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, ccache_env) elif sys.platform == "win32": - _windows_compile(output, objects, options) + _windows_compile(output, objects, options, ccache_env) else: raise ValueError("Unsupported platform") @@ -269,7 +275,7 @@ 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, 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 +294,13 @@ 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) + if ccache_env is not None: + if shutil.which("ccache"): + cmd.insert(0, "ccache") + ccache_env = os.environ.copy().update(ccache_env) + else: + raise ValueError("ccache not found") + proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=ccache_env) (out, _) = proc.communicate() if proc.returncode != 0: msg = "Compilation error:\n" @@ -297,7 +309,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, ccache_env=None): cmd = ["clang"] cmd += ["-O2"] @@ -312,9 +324,17 @@ def _windows_compile(output, objects, options): cmd += objects if options: cmd += options + if ccache_env is not None: + if shutil.which("ccache"): + cmd.insert(0, "ccache") + ccache_env = os.environ.copy().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, env=ccache_env + ) (out, _) = proc.communicate() except FileNotFoundError: raise RuntimeError( From 4e3d779746b982e0e9fec64f91fab0dd0911861d Mon Sep 17 00:00:00 2001 From: Yaxing Cai Date: Wed, 29 Nov 2023 00:13:55 +0000 Subject: [PATCH 2/3] add unit test --- python/tvm/contrib/cc.py | 42 +++++++++++----- tests/python/contrib/test_ccache.py | 75 +++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 13 deletions(-) create mode 100644 tests/python/contrib/test_ccache.py diff --git a/python/tvm/contrib/cc.py b/python/tvm/contrib/cc.py index 50a65ba208d8..918e3c8f7261 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, ccache_env=None): +def create_shared(output, objects, options=None, cc=None, cwd=None, ccache_env=None): """Create shared library. Parameters @@ -75,15 +79,18 @@ def create_shared(output, objects, options=None, cc=None, ccache_env=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, ccache_env, compile_shared=True) - elif sys.platform == "win32": - _windows_compile(output, objects, options, ccache_env) + _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") @@ -136,7 +143,7 @@ def create_staticlib(output, inputs, ar=None): raise ValueError("Unsupported platform") -def create_executable(output, objects, options=None, cc=None, ccache_env=None): +def create_executable(output, objects, options=None, cc=None, cwd=None, ccache_env=None): """Create executable binary. Parameters @@ -153,15 +160,18 @@ def create_executable(output, objects, options=None, cc=None, ccache_env=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, ccache_env) + _linux_compile(output, objects, options, cc, cwd, ccache_env) elif sys.platform == "win32": - _windows_compile(output, objects, options, ccache_env) + _windows_compile(output, objects, options, cwd, ccache_env) else: raise ValueError("Unsupported platform") @@ -275,7 +285,9 @@ def _fcompile(outputs, objects, options=None): return _fcompile -def _linux_compile(output, objects, options, compile_cmd, ccache_env=None, 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"): @@ -294,13 +306,15 @@ def _linux_compile(output, objects, options, compile_cmd, ccache_env=None, compi cmd += objects if options: cmd += options + env = None if ccache_env is not None: if shutil.which("ccache"): cmd.insert(0, "ccache") - ccache_env = os.environ.copy().update(ccache_env) + env = os.environ.copy() + env.update(ccache_env) else: raise ValueError("ccache not found") - proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=ccache_env) + 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" @@ -309,7 +323,7 @@ def _linux_compile(output, objects, options, compile_cmd, ccache_env=None, compi raise RuntimeError(msg) -def _windows_compile(output, objects, options, ccache_env=None): +def _windows_compile(output, objects, options, cwd=None, ccache_env=None): cmd = ["clang"] cmd += ["-O2"] @@ -324,16 +338,18 @@ def _windows_compile(output, objects, options, ccache_env=None): cmd += objects if options: cmd += options + env = None if ccache_env is not None: if shutil.which("ccache"): cmd.insert(0, "ccache") - ccache_env = os.environ.copy().update(ccache_env) + 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, env=ccache_env + cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=cwd, env=env ) (out, _) = proc.communicate() except FileNotFoundError: diff --git a/tests/python/contrib/test_ccache.py b/tests/python/contrib/test_ccache.py new file mode 100644 index 000000000000..c31e56ed5efe --- /dev/null +++ b/tests/python/contrib/test_ccache.py @@ -0,0 +1,75 @@ +# 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 tempfile +from tvm.contrib.cc import create_shared, create_executable, _is_linux_like, _is_windows_like + + +def _src_gen(text): + return """ +#include + +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 + + +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 + + +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__": + test_shared() + test_executable() From 0eb2044f43ddc84b046de94902e5ac44eb0d8344 Mon Sep 17 00:00:00 2001 From: Yaxing Cai Date: Wed, 29 Nov 2023 19:19:04 +0000 Subject: [PATCH 3/3] fix test --- tests/python/contrib/test_ccache.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/python/contrib/test_ccache.py b/tests/python/contrib/test_ccache.py index c31e56ed5efe..c46e72c648c1 100644 --- a/tests/python/contrib/test_ccache.py +++ b/tests/python/contrib/test_ccache.py @@ -16,7 +16,10 @@ # 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 @@ -48,6 +51,7 @@ def _compile(f_create, text, output): 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") @@ -59,6 +63,7 @@ def test_shared(): 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") @@ -71,5 +76,4 @@ def test_executable(): if __name__ == "__main__": - test_shared() - test_executable() + tvm.testing.main()