Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions dependencies.toml
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ dependencies = [
'py-cpuinfo',
'tqdm',
'typing_extensions',
'xtc-build~=0.1.0',
'pyyaml',
'scikit-learn',
'xdsl~=0.57.1',
Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ dependencies = [
'py-cpuinfo',
'tqdm',
'typing_extensions',
'xtc-build~=0.1.0',
'pyyaml',
'scikit-learn',
'xdsl~=0.57.1',
Expand Down
199 changes: 51 additions & 148 deletions src/xtc/backends/tvm/TVMCompiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,19 +7,26 @@
import tempfile
from pathlib import Path
import shutil
import subprocess
import shlex
import sys
from functools import partial
from packaging.version import Version

from xtc_build import (
Archive,
BuildContext,
ExternalArchive,
ExternalSharedLibrary,
GnuToolchain,
Object,
SharedLibrary,
)

from xtc.targets.host import HostModule

import xtc.backends.tvm as backend
import xtc.itf as itf
from xtc.utils.text import jinja_generate_file
from xtc.utils.tarfile import TarFile
from xtc.utils.files import relative_to
from xtc.utils.ext_tools import cc_opts

from xtc.utils.host_tools import (
Expand Down Expand Up @@ -199,6 +206,9 @@ def compile(self, schedule: itf.schd.Schedule) -> itf.comp.Module:
assert Path(module_file).with_suffix("") == Path(lib_path)
if self.save_temps:
self._save_temp_file(module_file)
for csrc in module_args.get("csrcs", []):
self._save_temp_file(csrc)

if type == "shlib" and self.print_assembly:
disassembly = disassemble(
module_file,
Expand Down Expand Up @@ -440,7 +450,6 @@ def build(
unpacked_lib_dir = Path(lib_fname).parent
unpacked_lib_base = Path(lib_fname).stem
packed_lib_dir = Path(packed_lib_fname).parent
packed_lib_name = Path(packed_lib_fname).stem
packed_ar_name = f"{packed_lib_fname}.a"
assert packed_lib_dir == unpacked_lib_dir, (
f"must generate wrapper at the same location as packed lib"
Expand Down Expand Up @@ -484,69 +493,53 @@ def build(
f"{tvm_libdir}/libtvm_runtime{ext}",
f"{tvm_ffi_libdir}/libtvm_ffi{ext}",
]
elif type == "shlib":
output_dir = unpacked_lib_dir
base_sources = [f"{output_base}.c"] + self._runtime_sources()
xflags = [None] * len(base_sources) + [additional_csrcs_xflags] * len(
else:
assert type in ["shlib", "arlib"]
base_sources = [f"{output_base}.c", *self._runtime_sources()]
sources = [*base_sources, *additional_csrcs]
xflags = [""] * len(base_sources) + [additional_csrcs_xflags] * len(
additional_csrcs
)
object_fnames = [
str(relative_to(fname, output_dir))
for fname in self._build_objects(
base_sources + additional_csrcs,
tdir,
xflags=xflags,
objects = [
Object(
name=f"{index:04d}_{Path(source).stem}",
source=source,
compile_flags=flags,
pic=True,
)
for index, (source, flags) in enumerate(zip(sources, xflags))
]
opts = " ".join(cc_opts)
sh_opts = "--shared -fPIC"
ext = ".so"
if sys.platform == "darwin":
sh_opts += " -undefined dynamic_lookup"
ext = ".dylib"
shlib_fname = f"{unpacked_lib_base}{ext}"
shlib_dest = str(relative_to(shlib_fname, output_dir))
cmd = (
f"{cc_command(self._arch)} {sh_opts} {opts} "
f"{' '.join(object_fnames)} "
f"{relative_to(packed_lib_fname, output_dir)}.a "
f"-o {unpacked_lib_base}{ext}"
)
p = subprocess.run(
shlex.split(cmd),
text=True,
capture_output=True,
cwd=output_dir,
context = BuildContext(
build_dir=tdir,
toolchain=GnuToolchain(
cc=cc_command(self._arch),
ar=binutils_command("ar", self._arch),
),
compile_flags=cc_opts,
)
if p.returncode != 0:
raise RuntimeError(
f"Failed command {cmd} (cwd: {output_dir}:\n"
f"{p.stdout}\n"
f"{p.stderr}\n"
)
module_file = f"{lib_fname}{ext}"
shlibs += [
f"{tvm_libdir}/libtvm_runtime{ext}",
f"{tvm_ffi_libdir}/libtvm_ffi{ext}",
]
else:
assert type == "arlib"
base_sources = [f"{output_base}.c"] + self._runtime_sources()
xflags = [None] * len(base_sources) + [additional_csrcs_xflags] * len(
additional_csrcs
)
archive_fname = self._build_archive(
base_sources + additional_csrcs,
f"{lib_fname}.a",
tdir,
xflags=xflags,
)
module_file = archive_fname
arlibs += [f"{packed_lib_fname}.a"]
shlibs += [
f"{tvm_libdir}/libtvm_runtime{ext}",
f"{tvm_ffi_libdir}/libtvm_ffi{ext}",
]
if type == "shlib":
link_flags = [*cc_opts]
if sys.platform == "darwin":
link_flags.extend(["-undefined", "dynamic_lookup"])
library = SharedLibrary(
unpacked_lib_base,
objects=objects,
archives=[ExternalArchive(packed_ar_name, pic=True)],
libraries=[ExternalSharedLibrary(path) for path in shlibs],
link_flags=link_flags,
)
built_path = library.build(context)
module_file = f"{lib_fname}{ext}"
else:
archive = Archive(unpacked_lib_base, objects=objects)
built_path = archive.build(context)
module_file = f"{lib_fname}.a"
arlibs += [packed_ar_name]
shutil.move(built_path, module_file)
except Exception:
raise
else:
Expand All @@ -564,93 +557,3 @@ def _runtime_sources(self) -> list[str]:
host_runtime_dir = Path(__file__).parents[2] / "csrcs" / "runtimes" / "host"
tvm_runtime_init_c = str(host_runtime_dir / "tvm_runtime_init.c")
return [tvm_runtime_init_c]

def _build_object(
self,
source_fname: str,
object_fname: str,
flags: str | None = None,
xflags: str | None = None,
) -> str:
assert object_fname.endswith(".o")
flags = " ".join(cc_opts) if flags is None else flags
xflags = "" if xflags is None else xflags
pic_flags = "-fPIC"
output_dir = Path(object_fname).parent
object_dest = str(relative_to(object_fname, output_dir))
source_inp = str(relative_to(source_fname, output_dir))
cmd = (
f"{cc_command(self._arch)} -c {pic_flags} {flags} {xflags} "
f"{source_inp} "
f"-o {object_dest}"
)
p = subprocess.run(
shlex.split(cmd),
text=True,
capture_output=True,
cwd=output_dir,
)
if p.returncode != 0:
raise RuntimeError(
f"Failed command {cmd} (cwd: {output_dir} :\n{p.stdout}\n{p.stderr}\n"
)
return object_fname

def _build_objects(
self,
source_fnames: list[str],
output_dir: str,
flags: list[str | None] | str | None = None,
xflags: list[str | None] | str | None = None,
) -> list[str]:
if not isinstance(flags, list):
flags = [flags] * len(source_fnames)
if not isinstance(xflags, list):
xflags = [xflags] * len(source_fnames)
return [
self._build_object(
fname,
str(Path(output_dir) / f"{Path(fname).stem}.o"),
flags,
xflags,
)
for fname, (flags, xflags) in zip(source_fnames, zip(flags, xflags))
]

def _build_archive(
self,
source_fnames: list[str],
archive_fname: str,
flags: list[str | None] | str | None = None,
xflags: list[str | None] | str | None = None,
) -> str:
assert archive_fname.endswith(".a")
output_dir = Path(archive_fname).parent
archive_dest = str(relative_to(archive_fname, output_dir))
tdir = tempfile.mkdtemp(dir=output_dir)
try:
object_fnames = [
str(relative_to(fname, output_dir))
for fname in self._build_objects(source_fnames, tdir, flags, xflags)
]
cmd = (
f"{binutils_command('ar', self._arch)} -crs {archive_dest} "
f"{' '.join(object_fnames)} "
)
p = subprocess.run(
shlex.split(cmd),
text=True,
capture_output=True,
cwd=output_dir,
)
if p.returncode != 0:
raise RuntimeError(
f"Failed command {cmd} (cwd: {output_dir}:\n"
f"{p.stdout}\n"
f"{p.stderr}\n"
)
except Exception:
raise
else:
shutil.rmtree(tdir)
return archive_fname
58 changes: 27 additions & 31 deletions src/xtc/targets/host/HostAREvaluator.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,16 @@
from typing_extensions import override
import tempfile
from pathlib import Path
import subprocess
import shlex
import shutil
import sys

from xtc_build import (
BuildContext,
ExternalArchive,
ExternalSharedLibrary,
SharedLibrary,
)

import xtc.itf as itf
import xtc.targets.host as host
from xtc.utils.ext_tools import cc_opts
Expand Down Expand Up @@ -40,44 +45,35 @@ def evaluate(self) -> tuple[list[float], int, str]:
def module(self) -> itf.comp.Module:
return self._module

def _compile_to_shlib(self, shlib_base: str):
cwd_dir = Path(shlib_base).parent
shlib_name = Path(shlib_base).stem
opts = " ".join(cc_opts)
sh_opts = "--shared -fPIC"
arlibs = [
str(Path(fname).absolute())
for fname in [self._module.file_name] + self._module.arlibs
def _compile_to_shlib(self, build_dir: Path, shlib_name: str) -> Path:
archives = [
ExternalArchive(fname, pic=True)
for fname in [self._module.file_name, *self._module.arlibs]
]
opt_whole, opt_no_whole = "-Wl,--whole-archive", "-Wl,--no-whole-archive"
ext = ".so"
libraries = [ExternalSharedLibrary(fname) for fname in self._module.shlibs]
symbol = self._module.payload_name
if sys.platform == "darwin":
sh_opts += " -undefined dynamic_lookup"
opt_whole, opt_no_whole = "-Wl,-all_load", ""
ext = ".dylib"
cmd = (
f"cc {sh_opts} {opts} "
f"{opt_whole} "
f"{' '.join(arlibs)} "
f"{opt_no_whole} "
f"-o {shlib_name}{ext}"
)
p = subprocess.run(
shlex.split(cmd), text=True, capture_output=True, cwd=cwd_dir
symbol = f"_{symbol}"
link_flags = [*cc_opts, f"-Wl,-u,{symbol}"]
if sys.platform == "darwin":
link_flags.extend(["-undefined", "dynamic_lookup"])
library = SharedLibrary(
shlib_name,
archives=archives,
libraries=libraries,
link_flags=link_flags,
)
if p.returncode != 0:
raise RuntimeError(f"Failed command {cmd}:\n{p.stdout}\n{p.stderr}\n")
return library.build(BuildContext(build_dir=build_dir))

def _build_shlib_module(self) -> None:
self.tmp_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
c_stem = Path(self._module.file_name).stem
shlib_base = str(Path(self.tmp_dir.name) / f"{c_stem}_eval")
self._compile_to_shlib(shlib_base)
ext = ".dylib" if sys.platform == "darwin" else ".so"
shlib_name = f"{c_stem}_eval"
shlib_path = self._compile_to_shlib(Path(self.tmp_dir.name), shlib_name)
self._shlib_module: host.HostModule = host.HostModule(
shlib_base,
shlib_name,
self._module.payload_name,
f"{shlib_base}{ext}",
str(shlib_path),
"shlib",
bare_ptr=self._module._bare_ptr,
graph=self._module._graph,
Expand Down
Loading
Loading