Skip to content

Commit

Permalink
#475: Create a whls release for ttrt (#476)
Browse files Browse the repository at this point in the history
  • Loading branch information
@tapspatel
tapspatel authored Aug 23, 2024
1 parent 239feca commit f8cbb68
Show file tree
Hide file tree
Showing 5 changed files with 165 additions and 29 deletions.
14 changes: 14 additions & 0 deletions docs/src/ttrt.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,18 @@ cmake --build build -- ttrt
ttrt --help
```

## Installing ttrt as python whls (coming soon)
1. Download whls
2. Create a python venv
```bash
python -m venv ttrt_env
source ttrt_env/bin/activate
```
3. Install whls
```bash
pip install *.whl
```

### Building runtime mode
Add the following flags when building the compiler
```bash
Expand All @@ -35,6 +47,8 @@ pip install loguru
pip install pandas
pip install seaborn
pip install graphviz
pip install pyyaml
pip install click
```

## Generate a flatbuffer file from compiler
Expand Down
12 changes: 11 additions & 1 deletion runtime/tools/python/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,19 @@ add_custom_target(ttrt-copy-files
)

add_custom_target(ttrt
COMMAND TTMLIR_ENABLE_RUNTIME=${TTMLIR_ENABLE_RUNTIME} TT_RUNTIME_ENABLE_TTNN=${TT_RUNTIME_ENABLE_TTNN} TT_RUNTIME_ENABLE_TTMETAL=${TT_RUNTIME_ENABLE_TTMETAL} TT_RUNTIME_ENABLE_PERF_TRACE=${TT_RUNTIME_ENABLE_PERF_TRACE} TTMLIR_VERSION_MAJOR=${TTMLIR_VERSION_MAJOR} TTMLIR_VERSION_MINOR=${TTMLIR_VERSION_MINOR} TTMLIR_VERSION_PATCH=${TTMLIR_VERSION_PATCH} SOURCE_ROOT=${TTMLIR_SOURCE_DIR} python -m pip install .
COMMAND TTMLIR_ENABLE_RUNTIME=${TTMLIR_ENABLE_RUNTIME}
TT_RUNTIME_ENABLE_TTNN=${TT_RUNTIME_ENABLE_TTNN}
TT_RUNTIME_ENABLE_TTMETAL=${TT_RUNTIME_ENABLE_TTMETAL}
TT_RUNTIME_ENABLE_PERF_TRACE=${TT_RUNTIME_ENABLE_PERF_TRACE}
TTMLIR_VERSION_MAJOR=${TTMLIR_VERSION_MAJOR}
TTMLIR_VERSION_MINOR=${TTMLIR_VERSION_MINOR}
TTMLIR_VERSION_PATCH=${TTMLIR_VERSION_PATCH}
SOURCE_ROOT=${TTMLIR_SOURCE_DIR}
python -m pip wheel . --wheel-dir build --verbose
COMMAND python -m pip install build/*.whl --force-reinstall
WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}"
COMMENT "python ttrt package"
DEPENDS ttrt-copy-files
)

add_dependencies(ttrt TTRuntime)
103 changes: 95 additions & 8 deletions runtime/tools/python/setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,8 @@
os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", ".."),
)
toolchain = os.environ.get("TTMLIR_TOOLCHAIN_DIR", "/opt/ttmlir-toolchain")
metallibdir = f"{src_dir}/third_party/tt-metal/src/tt-metal-build/lib"
metaldir = f"{src_dir}/third_party/tt-metal/src/tt-metal-build"
ttmetalhome = os.environ.get("TT_METAL_HOME", "")

os.environ["LDFLAGS"] = "-Wl,-rpath,'$ORIGIN'"
enable_runtime = os.environ.get("TTMLIR_ENABLE_RUNTIME", "OFF") == "ON"
Expand Down Expand Up @@ -47,24 +48,30 @@
]

dylibs = []
runlibs = []
perflibs = []
metallibs = []

linklibs = ["TTBinary", "TTRuntimeSysDesc"]
if enable_ttnn:
dylibs += ["_ttnn.so"]
runlibs += ["_ttnn.so"]
linklibs += ["TTRuntimeTTNN", ":_ttnn.so"]

if enable_ttmetal:
dylibs += ["libtt_metal.so"]
runlibs += ["libtt_metal.so"]
linklibs += ["TTRuntimeTTMetal", "tt_metal"]

if enable_perf:
dylibs += ["libtracy.so.0.10.0"]
runlibs += ["libtracy.so.0.10.0"]
perflibs += ["capture-release"]
perflibs += ["csvexport-release"]

if enable_runtime:
assert enable_ttmetal or enable_ttnn, "At least one runtime must be enabled"

for dylib in dylibs:
for dylib in runlibs:
shutil.copy(
f"{metallibdir}/{dylib}",
f"{metaldir}/lib/{dylib}",
f"{src_dir}/build/runtime/tools/python/ttrt/runtime",
)
command = [
Expand All @@ -84,6 +91,72 @@
)
except subprocess.CalledProcessError as e:
print(f"Command failed with return code {e.returncode}")

for dylib in perflibs:
shutil.copy(
f"{metaldir}/tools/profiler/bin/{dylib}",
f"{src_dir}/build/runtime/tools/python/ttrt/runtime",
)
shutil.copy(
f"{metaldir}/tools/profiler/bin/{dylib}",
f"{ttmetalhome}/{dylib}",
)

# copy metal dir folder
shutil.copytree(
f"{ttmetalhome}/tt_metal",
f"{src_dir}/build/runtime/tools/python/ttrt/runtime/tt_metal",
dirs_exist_ok=True,
)

# copy runtime dir folder
shutil.copytree(
f"{ttmetalhome}/runtime",
f"{src_dir}/build/runtime/tools/python/ttrt/runtime/runtime",
dirs_exist_ok=True,
)

# copy kernels
shutil.copytree(
f"{ttmetalhome}/ttnn",
f"{src_dir}/build/runtime/tools/python/ttrt/runtime/ttnn",
dirs_exist_ok=True,
)

# copy test folder
shutil.copytree(
f"{ttmetalhome}/tests",
f"{src_dir}/build/runtime/tools/python/ttrt/runtime/tests",
dirs_exist_ok=True,
)

import os

def package_files(directory):
paths = []
for (path, directories, filenames) in os.walk(directory):
for filename in filenames:
paths.append(os.path.join("..", path, filename))
return paths

extra_files_tt_metal = package_files(
f"{src_dir}/build/runtime/tools/python/ttrt/runtime/tt_metal/"
)
extra_files_runtime = package_files(
f"{src_dir}/build/runtime/tools/python/ttrt/runtime/runtime/"
)
extra_files_ttnn = package_files(
f"{src_dir}/build/runtime/tools/python/ttrt/runtime/ttnn/"
)
extra_files_tests = package_files(
f"{src_dir}/build/runtime/tools/python/ttrt/runtime/tests/"
)

metallibs += extra_files_tt_metal
metallibs += extra_files_runtime
metallibs += extra_files_ttnn
metallibs += extra_files_tests

ext_modules.append(
Pybind11Extension(
"ttrt.runtime._C",
Expand All @@ -102,12 +175,24 @@
f"{src_dir}/build/runtime/lib/ttmetal",
f"{toolchain}/lib",
f"{src_dir}/build/runtime/tools/python/ttrt/runtime",
f"{metallibdir}",
f"{metaldir}/lib",
],
define_macros=[("VERSION_INFO", __version__)],
)
)

dylibs += runlibs
dylibs += perflibs
dylibs += metallibs

packages = ["ttrt", "ttrt.common", "ttrt.binary", "ttrt.runtime"]
package_dir = {}
if enable_perf:
packages += ["tracy"]
packages += ["tt_metal"]
package_dir["tracy"] = f"{ttmetalhome}/ttnn/tracy"
package_dir["tt_metal"] = f"{ttmetalhome}/tt_metal"

setup(
name="ttrt",
version=__version__,
Expand All @@ -118,12 +203,14 @@
long_description="",
ext_modules=ext_modules,
cmdclass={"build_ext": build_ext},
packages=["ttrt", "ttrt.common", "ttrt.binary", "ttrt.runtime"],
packages=packages,
package_dir=package_dir,
install_requires=["pybind11"],
entry_points={
"console_scripts": ["ttrt = ttrt:main"],
},
package_data={"ttrt.runtime": dylibs},
# include_package_data=True,
zip_safe=False,
python_requires=">=3.7",
)
33 changes: 23 additions & 10 deletions runtime/tools/python/ttrt/common/api.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
import time
import socket
from pkg_resources import get_distribution
import sys
import shutil
import atexit

Expand Down Expand Up @@ -1064,8 +1063,12 @@ def __init__(self, args={}, logging=None, artifacts=None):
self.query = API.Query({}, self.logger, self.artifacts)
self.ttnn_binaries = []
self.ttmetal_binaries = []
self.tracy_capture_tool_path = f"{self.globals.get_ttmlir_home_path()}/third_party/tt-metal/src/tt-metal-build/tools/profiler/bin/capture-release"
self.tracy_csvexport_tool_path = f"{self.globals.get_ttmlir_home_path()}/third_party/tt-metal/src/tt-metal-build/tools/profiler/bin/csvexport-release"
self.tracy_capture_tool_path = (
f"{self.globals.get_ttmetal_home_path()}/capture-release"
)
self.tracy_csvexport_tool_path = (
f"{self.globals.get_ttmetal_home_path()}/csvexport-release"
)
self.tracy_capture_tool_process = None

def preprocess(self):
Expand Down Expand Up @@ -1138,6 +1141,9 @@ def check_constraints(self):
def execute(self):
self.logging.debug(f"executing perf API")

sys.path.append(f"{get_ttrt_metal_home_path()}")
sys.path.append(f"{get_ttrt_metal_home_path()}/ttnn")

import tracy
import tracy.tracy_state
from tt_metal.tools.profiler.process_ops_logs import process_ops
Expand Down Expand Up @@ -1237,13 +1243,19 @@ def _execute(binaries):
raise Exception("No available port found")

env_vars = dict(os.environ)
self.globals.add_global_env("TRACY_PORT", port)
self.globals.add_global_env(
"TT_METAL_DEVICE_PROFILER_DISPATCH", "0"
)
env_vars["TRACY_PORT"] = port
env_vars["TT_METAL_DEVICE_PROFILER_DISPATCH"] = "0"

if self["device"]:
self.globals.add_global_env("TT_METAL_DEVICE_PROFILER", "1")
env_vars["TT_METAL_DEVICE_PROFILER"] = "1"

current_pythonpath = env_vars.get("PYTHONPATH", "")
path1 = f"{get_ttrt_metal_home_path()}"
path2 = f"{get_ttrt_metal_home_path()}/ttnn"
new_pythonpath = (
f"{current_pythonpath}{os.pathsep}{path1}{os.pathsep}{path2}"
)
env_vars["PYTHONPATH"] = new_pythonpath

tracy_capture_tool_command = f"{self.tracy_capture_tool_path} -o {PROFILER_LOGS_DIR / TRACY_FILE_NAME} -f -p {port}"
self.tracy_capture_tool_process = subprocess.Popen(
Expand All @@ -1268,11 +1280,12 @@ def _execute(binaries):
command_options += f" {api['name']} {self[name]} "

library_link_path = self.globals.get_ld_path(
f"{Globals.get_ttmlir_venv_path()}/lib/python3.10/site-packages/ttrt/runtime/"
f"{self.globals.get_ttmetal_home_path()}"
)
test_env_flags = f"LD_LIBRARY_PATH={library_link_path}"

test_command = f"{test_env_flags} python -m tracy -p {self.globals.get_ttmlir_venv_path()}/bin/ttrt run {bin.file_path} --save-artifacts {command_options}"
ttrt_executable_path = shutil.which("ttrt")
test_command = f"{test_env_flags} python -m tracy -p {ttrt_executable_path} run {bin.file_path} --save-artifacts {command_options}"
print(f"test command for binary={bin.file_path} is: {test_command}")
testProcess = subprocess.Popen(
[test_command], shell=True, env=env_vars, preexec_fn=os.setsid
Expand Down
32 changes: 22 additions & 10 deletions runtime/tools/python/ttrt/common/util.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import os
import json
import importlib.machinery
import importlib.util
import sys
import signal
import os
Expand All @@ -13,17 +14,36 @@
import time
import socket
from pkg_resources import get_distribution
import sys
import shutil

import ttrt.binary

# environment tweaks
if "LOGGER_LEVEL" not in os.environ:
os.environ["LOGGER_LEVEL"] = "FATAL"
if "TT_METAL_LOGGER_LEVEL" not in os.environ:
os.environ["TT_METAL_LOGGER_LEVEL"] = "FATAL"


def get_ttrt_metal_home_path():
package_name = "ttrt"
spec = importlib.util.find_spec(package_name)
package_path = os.path.dirname(spec.origin)
tt_metal_home = f"{package_path}/runtime"
return tt_metal_home


os.environ["TT_METAL_HOME"] = get_ttrt_metal_home_path()

new_linker_path = f"{get_ttrt_metal_home_path()}/tests"
current_ld_library_path = os.environ.get("LD_LIBRARY_PATH", "")
if current_ld_library_path:
updated_ld_library_path = f"{new_linker_path}:{current_ld_library_path}"
else:
updated_ld_library_path = new_linker_path
os.environ["LD_LIBRARY_PATH"] = updated_ld_library_path


class Logger:
def __init__(self, file_name=""):
import logging
Expand Down Expand Up @@ -80,14 +100,6 @@ def get_ld_path(self, path):
updated_path = f"{path}:{current_path}"
return updated_path

@staticmethod
def get_ttmlir_home_path():
return os.environ.get("TT_MLIR_HOME", f"{os.getcwd()}")

@staticmethod
def get_ttmlir_venv_path():
return os.environ.get("TTMLIR_VENV_DIR", "/opt/ttmlir-toolchain/venv")

@staticmethod
def get_ttmetal_home_path():
return os.environ.get("TT_METAL_HOME", "third_party/tt-metal/src/tt-metal")
Expand Down Expand Up @@ -359,7 +371,7 @@ def __init__(self, logger, file_manager=None, artifacts_folder_path=""):
self.artifacts_folder_path = (
artifacts_folder_path
if artifacts_folder_path != ""
else f"{Globals.get_ttmlir_home_path()}/ttrt-artifacts"
else f"{os.getcwd()}/ttrt-artifacts"
)

self.logging.info(f"setting artifacts folder path={self.artifacts_folder_path}")
Expand Down

0 comments on commit f8cbb68

Please sign in to comment.