From f8cbb68fc5e28ef1ccb078ad1cd84616bdf01890 Mon Sep 17 00:00:00 2001 From: Tapasvi Patel <133996364+tapspatel@users.noreply.github.com> Date: Fri, 23 Aug 2024 14:54:15 -0500 Subject: [PATCH] #475: Create a whls release for ttrt (#476) --- docs/src/ttrt.md | 14 +++ runtime/tools/python/CMakeLists.txt | 12 ++- runtime/tools/python/setup.py | 103 +++++++++++++++++++++-- runtime/tools/python/ttrt/common/api.py | 33 +++++--- runtime/tools/python/ttrt/common/util.py | 32 ++++--- 5 files changed, 165 insertions(+), 29 deletions(-) diff --git a/docs/src/ttrt.md b/docs/src/ttrt.md index e881392ed..26be66bcf 100644 --- a/docs/src/ttrt.md +++ b/docs/src/ttrt.md @@ -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 @@ -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 diff --git a/runtime/tools/python/CMakeLists.txt b/runtime/tools/python/CMakeLists.txt index 8bd54df25..9a5d7c745 100644 --- a/runtime/tools/python/CMakeLists.txt +++ b/runtime/tools/python/CMakeLists.txt @@ -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) diff --git a/runtime/tools/python/setup.py b/runtime/tools/python/setup.py index 12401f15e..5864b12c9 100644 --- a/runtime/tools/python/setup.py +++ b/runtime/tools/python/setup.py @@ -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" @@ -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 = [ @@ -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", @@ -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__, @@ -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", ) diff --git a/runtime/tools/python/ttrt/common/api.py b/runtime/tools/python/ttrt/common/api.py index 005f74334..2f7c20e6f 100644 --- a/runtime/tools/python/ttrt/common/api.py +++ b/runtime/tools/python/ttrt/common/api.py @@ -13,7 +13,6 @@ import time import socket from pkg_resources import get_distribution -import sys import shutil import atexit @@ -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): @@ -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 @@ -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( @@ -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 diff --git a/runtime/tools/python/ttrt/common/util.py b/runtime/tools/python/ttrt/common/util.py index 12d4932a4..d332e3b81 100644 --- a/runtime/tools/python/ttrt/common/util.py +++ b/runtime/tools/python/ttrt/common/util.py @@ -5,6 +5,7 @@ import os import json import importlib.machinery +import importlib.util import sys import signal import os @@ -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 @@ -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") @@ -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}")