LeiWang1999 commented on PR #18124:
URL: https://github.com/apache/tvm/pull/18124#issuecomment-3051120831
@vacu9708 Thanks for your report, I think one possible solution would be to
refactor the Cython build phase into command lines. This would give us a more
flexible method for finding the Cython executable and building Cython modules,
and it would also enable caching
```python
def get_cython_compiler() -> Optional[str]:
"""Return the path to the Cython compiler.
Returns
-------
out: Optional[str]
The path to the Cython compiler, or None if none was found.
"""
cython_names = ["cython", "cython3"]
# Check system PATH
dirs_in_path = list(os.get_exec_path())
# Add user site-packages bin directory
user_base = site.getuserbase()
if user_base:
user_bin = os.path.join(user_base, "bin")
if os.path.exists(user_bin):
dirs_in_path = [user_bin] + dirs_in_path
# If in a virtual environment, add its bin directory
if sys.prefix != sys.base_prefix:
venv_bin = os.path.join(sys.prefix, "bin")
if os.path.exists(venv_bin):
dirs_in_path = [venv_bin] + dirs_in_path
for cython_name in cython_names:
for d in dirs_in_path:
cython_path = os.path.join(d, cython_name)
if os.path.isfile(cython_path) and os.access(cython_path,
os.X_OK):
return cython_path
return None
```
and build .pyx files from command line.
```python
with open(cython_wrapper_path, "r") as f:
cython_wrapper_code = f.read()
cache_dir = get_cache_dir()
source_path = cache_dir / "cython_wrapper.cpp"
library_path = cache_dir / "cython_wrapper.so"
md5_path = cache_dir / "md5.txt"
code_hash = hashlib.sha256(cython_wrapper_code.encode()).hexdigest()
cache_path = cache_dir / f"{code_hash}.so"
lock_file = cache_path.with_suffix('.lock')
# Check if cached version exists and is valid
need_compile = True
if md5_path.exists() and library_path.exists():
with open(md5_path, "r") as f:
cached_hash = f.read().strip()
if cached_hash == code_hash:
logger.debug("Cython jit adapter is up to date, no need to
compile...")
need_compile = False
else:
logger.info("Cython jit adapter is out of date, need to
recompile...")
else:
logger.info("No cached version found for cython jit adapter, need to
compile...")
if need_compile:
logger.info("Waiting for lock to compile cython jit adapter...")
with open(lock_file, 'w') as lock:
fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
try:
# After acquiring the lock, check again if the file has been
compiled by another process
if md5_path.exists() and library_path.exists():
with open(md5_path, "r") as f:
cached_hash = f.read().strip()
if cached_hash == code_hash:
logger.info(
"Another process has already compiled the
file, using it...")
need_compile = False
if need_compile:
logger.info("Compiling cython jit adapter...")
temp_path = cache_dir / f"temp_{code_hash}.so"
with open(md5_path, "w") as f:
f.write(code_hash)
# compile the cython_wrapper.pyx file into .cpp
cython = get_cython_compiler()
if cython is None:
raise Exception("Cython is not installed, please
install it first.")
os.system(f"{cython} {cython_wrapper_path} --cplus -o
{source_path}")
python_include_path = sysconfig.get_path("include")
cc = get_cplus_compiler()
command = f"{cc} -shared -pthread -fPIC -fwrapv -O2
-Wall -fno-strict-aliasing -I{python_include_path} {source_path} -o {temp_path}"
os.system(command)
# rename the temp file to the library file
temp_path.rename(library_path)
except Exception as e:
if 'temp_path' in locals() and temp_path.exists():
temp_path.unlink()
raise Exception(f"Failed to compile cython jit adapter:
{e}") from e
finally:
if lock_file.exists():
lock_file.unlink()
# add the .so file to the sys.path
cache_dir_str = str(cache_dir)
if cache_dir_str not in sys.path:
sys.path.append(cache_dir_str)
```
example codes are from
https://github.com/tile-ai/tilelang/blob/main/tilelang/jit/adapter/cython/adapter.py
cc @Hzfengsy
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
