Skip to content

Commit

Permalink
[Unity] Fix ccache env for nn.SourceModule (#16257)
Browse files Browse the repository at this point in the history
This PR refactors the compilation of `nn.SourceModule` to enable ccache by using the relative path, instead of using absolute path. Also it adds the ccache env to not hash the directory.
  • Loading branch information
cyx-6 authored and junrushao committed Dec 18, 2023
1 parent 4e66690 commit 45eeb8c
Showing 1 changed file with 12 additions and 5 deletions.
17 changes: 12 additions & 5 deletions python/tvm/relax/frontend/nn/extern.py
Original file line number Diff line number Diff line change
Expand Up @@ -371,17 +371,24 @@ def compile(self, output_path: Path) -> None:
"""Compiles the source code in a provided directory and returns the compiled artifact."""
with tempfile.TemporaryDirectory() as temp_dir_str:
temp_dir = Path(temp_dir_str)
source_path = temp_dir / f"main{self.source_suffix}"
object_path = temp_dir / f"main{self.output_suffix}"
source_filename = f"main{self.source_suffix}"
object_filename = f"main{self.output_suffix}"
source_path = temp_dir / source_filename
object_path = temp_dir / object_filename
with source_path.open("w", encoding="utf-8") as file:
file.write(self.source_code)
_cc.create_shared(
output=str(object_path),
objects=[str(source_path)],
output=object_filename,
objects=[source_filename],
options=self.compile_options,
cc=self.compiler,
cwd=temp_dir,
ccache_env={"CCACHE_COMPILERCHECK": "content"} if shutil.which("ccache") else None,
ccache_env={
"CCACHE_COMPILERCHECK": "content",
"CCACHE_NOHASHDIR": "1",
}
if shutil.which("ccache")
else None,
)
shutil.move(str(object_path), str(output_path))

Expand Down

0 comments on commit 45eeb8c

Please sign in to comment.