Fix venv path, so we no longer require modification to python path after install

CMakeLists.txt

@@ -31,7 +31,7 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 
-if(NOT DEFINED ENV{DONT_OVERRIDE_INSTALL_PATH})
+if(NOT DEFINED ENV{WHEEL_BUILD_CUSTOM_INSTALL_PATH})
     set(CMAKE_INSTALL_PREFIX ${TTTORCH_SOURCE_DIR}/install)
 endif()
 message(STATUS "CMAKE_INSTALL_PREFIX is set to ${CMAKE_INSTALL_PREFIX}")

env/activate

@@ -16,7 +16,7 @@ else
 fi
 
 export TT_TORCH_HOME="$(pwd)"
-export LD_LIBRARY_PATH=$TT_TORCH_HOME/env/venv/lib/python3.11/site-packages/torch/lib:$TTMLIR_TOOLCHAIN_DIR/lib:$TT_TORCH_HOME/install/lib/:$LD_LIBRRARY_PATH
+export LD_LIBRARY_PATH=$TT_TORCH_HOME/env/venv/lib/python3.11/site-packages/torch/lib:$TTMLIR_TOOLCHAIN_DIR/lib:$TT_TORCH_HOME/install/lib/:$TT_TORCH_HOME/env/venv/lib:$LD_LIBRRARY_PATH
 
 export TTMLIR_VENV_DIR="$(pwd)/env/venv"
 if [ -d $TTMLIR_VENV_DIR/bin ]; then
@@ -53,6 +53,6 @@ else
     export TT_METAL_HOME="$(pwd)/third_party/tt-mlir/src/tt-mlir/third_party/tt-metal/src/tt-metal"
 fi
 
-export PYTHONPATH="$(pwd):$(pwd)/env/venv:$(pwd)/env/venv/lib:$(pwd)/install/lib:$(pwd)/.local/toolchain/python_packages/mlir_core:${TT_METAL_HOME}:${TT_METAL_HOME}/tt_eager:${TT_METAL_BUILD_HOME}/tools/profiler/bin"
+export PYTHONPATH="$(pwd):$(pwd)/install/lib:${TT_METAL_HOME}:${TT_METAL_HOME}/tt_eager:${TT_METAL_BUILD_HOME}/tools/profiler/bin"
 export ARCH_NAME="${ARCH_NAME:-wormhole_b0}"
 export TT_METAL_LOGGER_LEVEL="ERROR"

setup.py

@@ -3,10 +3,64 @@
 # SPDX-License-Identifier: Apache-2.0
 from skbuild import setup
 import os
-import glob
 from skbuild.command.install_lib import install_lib
 from setuptools import find_namespace_packages
 import sys
+import shutil
+
+import hashlib
+import zipfile
+
+# TODO, this is quite hacky
+# The install files provided to us from tt-mlir are all in a single folder. Organize them so that
+# everything is where python expects it to be.
+def shuffle_wheel():
+    whl = "dist/tt_torch-0.1-cp311-cp311-linux_x86_64.whl"
+    print(f"Reading in {whl}")
+    with zipfile.ZipFile(whl, "r") as zip_ref:
+        zip_ref.extractall("wheel_contents")
+
+    print("Moving tt-metal")
+    metal_dir = os.path.join(
+        os.getcwd(), "wheel_contents/tt_torch-0.1.data/data/tt-metal"
+    )
+    shutil.copy2(metal_dir, "wheel_contents/")
+    shutil.rmtree(metal_dir)
+
+    record_path = "wheel_contents/tt_torch-0.1.dist-info/RECORD"
+    print("Populating RECORD.")
+    with open(record_path, "w") as record_file:
+        for root, _, files in os.walk("wheel_contents"):
+            for file in files:
+                file_path = os.path.join(root, file)
+                rel_path = os.path.relpath(file_path, "wheel_contents")
+
+                # Skip RECORD itself (it should be left without a hash)
+                if rel_path.endswith("RECORD"):
+                    record_file.write(f"{rel_path},,\n")
+                    continue
+
+                # Compute the hash
+                hasher = hashlib.sha256()
+                with open(file_path, "rb") as f:
+                    hasher.update(f.read())
+                hash_b64 = hasher.digest().hex()
+
+                # Get file size
+                file_size = os.path.getsize(file_path)
+
+                # Write to RECORD
+                record_file.write(f"{rel_path},sha256={hash_b64},{file_size}\n")
+
+    print("Creating wheel.")
+    with zipfile.ZipFile(whl, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk("wheel_contents"):
+            for file in files:
+                file_path = os.path.join(root, file)
+                rel_path = os.path.relpath(file_path, "wheel_contents")
+                zipf.write(file_path, rel_path)
+
+    shutil.rmtree("wheel_contents", ignore_errors=True)
 
 
 class install_metal_libs(install_lib):
@@ -29,7 +83,7 @@ def run(self):
 
 
 # Compile time env vars
-os.environ["DONT_OVERRIDE_INSTALL_PATH"] = "1"
+os.environ["WHEEL_BUILD_CUSTOM_INSTALL_PATH"] = "1"
 
 cmake_args = [
     "-GNinja",
@@ -65,3 +119,5 @@ def run(self):
         "numpy",
     ],
 )
+
+shuffle_wheel()

third_party/CMakeLists.txt

@@ -72,6 +72,12 @@ else()
     set_target_properties(TTMLIRRuntime PROPERTIES EXCLUDE_FROM_ALL TRUE IMPORTED_LOCATION ${TT_MLIR_RUNTIME_LIBRARY_PATH})
     add_dependencies(TTMLIRRuntime tt-mlir)
 
-    install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}" USE_SOURCE_PERMISSIONS)
+    if(DEFINED ENV{WHEEL_BUILD_CUSTOM_INSTALL_PATH})
+        file(MAKE_DIRECTORY ${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages)
+        install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}" USE_SOURCE_PERMISSIONS)
+        # install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}/lib/python3.11/site-packages" USE_SOURCE_PERMISSIONS)
+    else()
+        install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}" USE_SOURCE_PERMISSIONS)
+    endif()
 
 endif()

tt_torch/__init__.py

@@ -11,8 +11,14 @@
     mp.set_start_method("forkserver", force=True)
 
 import os
+import sys
 import importlib.util
 
+# make sure, venv/lib is in the LD_LIBRARY_PATH
+lib_path = os.path.join(os.environ["VIRTUAL_ENV"], "lib")
+sys.path.append(lib_path)
+os.environ["LD_LIBRARY_PATH"] = lib_path + os.pathsep + os.environ["LD_LIBRARY_PATH"]
+
 # find the tt-metal directory, it can either be in the venv if installed from a wheel or in the third_party source tree
 package_name = "tt-metal"
 spec = importlib.util.find_spec(package_name)

tt_torch/csrc/CMakeLists.txt

@@ -80,4 +80,5 @@ set_target_properties(${TARGET_NAME}
     LIBRARY_OUTPUT_NAME ${TARGET_NAME}
     INSTALL_RPATH "$ORIGIN"
 )
+
 install (TARGETS ${TARGET_NAME} LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}/lib)