Fix wheel build to remove redundant libraries

.github/workflows/generate-model-report.yml

@@ -14,7 +14,7 @@ jobs:
       fail-fast: false
 
     runs-on:
-      - n150
+      - wormhole_b0
 
     container:
       image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04:latest
@@ -38,6 +38,7 @@ jobs:
       run: |
         echo "work-dir=$(pwd)" >> "$GITHUB_OUTPUT"
         echo "install-dir=$(pwd)/install" >> "$GITHUB_OUTPUT"
+        echo "dist-dir=$(pwd)/dist" >> "$GITHUB_OUTPUT"
         echo "test-output-dir=$(pwd)/results/models/tests/" >> "$GITHUB_OUTPUT"
 
     - name: Git safe dir
@@ -47,7 +48,7 @@ jobs:
       uses: actions/download-artifact@v4
       with:
         name: install-artifacts
-        path: ${{ steps.strings.outputs.install-output-dir }}
+        path: ${{ steps.strings.outputs.install-dir }}
 
     - name: 'Untar install directory'
       shell: bash

.github/workflows/generate-ttnn-md.yml

@@ -25,7 +25,7 @@ jobs:
         shell: bash
         run: |
           echo "work-dir=$(pwd)" >> "$GITHUB_OUTPUT"
-          echo "install-output-dir=$(pwd)/install" >> "$GITHUB_OUTPUT"
+          echo "install-dir=$(pwd)/install" >> "$GITHUB_OUTPUT"
           echo "test-output-dir=$(pwd)/results/models/tests/" >> "$GITHUB_OUTPUT"
 
       - name: Git safe dir

.github/workflows/run-build.yml

@@ -50,7 +50,6 @@ jobs:
       run: |
         echo "work-dir=$(pwd)" >> "$GITHUB_OUTPUT"
         echo "build-output-dir=$(pwd)/build" >> "$GITHUB_OUTPUT"
-        echo "install-output-dir=$(pwd)/install" >> "$GITHUB_OUTPUT"
 
     - name: Git safe dir
       run: git config --global --add safe.directory ${{ steps.strings.outputs.work-dir }}

.github/workflows/run-e2e-tests.yml

@@ -50,6 +50,7 @@ jobs:
       run: |
         echo "work-dir=$(pwd)" >> "$GITHUB_OUTPUT"
         echo "install-dir=$(pwd)/install" >> "$GITHUB_OUTPUT"
+        echo "dist-dir=$(pwd)/dist" >> "$GITHUB_OUTPUT"
         echo "test_report_path_torch=report_torch_$JOB_ID.xml" >> "$GITHUB_OUTPUT"
         echo "test_report_path_models=report_models_$JOB_ID.xml" >> "$GITHUB_OUTPUT"
         echo "test_report_path_onnx=report_onnx_$JOB_ID.xml" >> "$GITHUB_OUTPUT"

env/activate

@@ -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/lib/python3.11/site-packages/tt_mlir:$(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)/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 ARCH_NAME="${ARCH_NAME:-wormhole_b0}"
 export TT_METAL_LOGGER_LEVEL="ERROR"

setup.py

@@ -13,17 +13,6 @@ class install_metal_libs(install_lib):
     def run(self):
         install_lib.run(self)
         install_path = os.path.join(self.install_dir, "tt_mlir")
-        os.makedirs(install_path, exist_ok=True)
-        lib_dir = os.path.abspath(
-            os.path.join(
-                os.getcwd(), "_skbuild", "linux-x86_64-3.11", "cmake-install", "lib"
-            )
-        )
-        if os.path.exists(lib_dir):
-            for file in glob.glob(os.path.join(lib_dir, "*")):
-                if os.path.isfile(file):
-                    self.copy_file(file, install_path)
-
         # Copy third_party/tt-mlir/src/tt-mlir-build/bin/ttmlir-opt into lib
         ttmlir_opt = os.path.abspath(
             os.path.join(
@@ -38,50 +27,13 @@ def run(self):
         )
         self.copy_file(ttmlir_opt, install_path)
 
-        metal_dir = os.path.abspath(
-            os.path.join(
-                os.getcwd(),
-                "third_party",
-                "tt-mlir",
-                "src",
-                "tt-mlir",
-                "third_party",
-                "tt-metal",
-                "src",
-                "tt-metal",
-            )
-        )
-        os.makedirs(self.install_dir, exist_ok=True)
-
-        tt_metal_folders_to_ignore = [
-            ".cpmcache",
-            ".git",
-        ]
-
-        def ignore(folder, contents):
-            relative_folder = os.path.relpath(folder, start=metal_dir)
-
-            ignored_items = [
-                item
-                for item in contents
-                if any(
-                    os.path.join(relative_folder, item).startswith(ignore)
-                    for ignore in tt_metal_folders_to_ignore
-                )
-            ]
-
-            return ignored_items
-
-        install_path = os.path.join(self.install_dir, "tt_torch_tt_metal_libs")
-        os.makedirs(install_path, exist_ok=True)
-        shutil.copytree(metal_dir, install_path, dirs_exist_ok=True, ignore=ignore)
-
 
 # Compile time env vars
 os.environ["DONT_OVERRIDE_INSTALL_PATH"] = "1"
 
 cmake_args = [
     "-GNinja",
+    "-DBUILD_TTRT=OFF",
 ]
 if "--code_coverage" in sys.argv:
     cmake_args += [

third_party/CMakeLists.txt

@@ -41,11 +41,17 @@ else()
     endif()
 
     include(ExternalProject)
-    set(TT_MLIR_COMPILER_LIBRARY_PATH ${CMAKE_INSTALL_PREFIX}/lib/libTTMLIRCompiler.so)
-    set(TT_MLIR_RUNTIME_LIBRARY_PATH ${CMAKE_INSTALL_PREFIX}/lib/libTTMLIRRuntime.so)
+    set(TTMLIR_PREFIX ${TTTORCH_SOURCE_DIR}/third_party/tt-mlir)
+    set(TTMLIR_INSTALL_PREFIX ${TTMLIR_PREFIX}/install)
+    set(TT_MLIR_COMPILER_LIBRARY_PATH ${TTMLIR_INSTALL_PREFIX}/lib/libTTMLIRCompiler.so)
+    set(TT_MLIR_RUNTIME_LIBRARY_PATH ${TTMLIR_INSTALL_PREFIX}/lib/libTTMLIRRuntime.so)
+
+    set(TTMLIR_BUILD_DIR ${TTMLIR_PREFIX}/src/tt-mlir-build)
+
     ExternalProject_Add(
         tt-mlir
-        PREFIX ${TTTORCH_SOURCE_DIR}/third_party/tt-mlir
+        PREFIX ${TTMLIR_PREFIX}
+        INSTALL_COMMAND ${CMAKE_COMMAND} --install ${TTMLIR_BUILD_DIR} --component SharedLib
         CMAKE_GENERATOR Ninja
         CMAKE_ARGS
           -DCMAKE_BUILD_TYPE=Release
@@ -58,7 +64,7 @@ else()
           -DTTMLIR_ENABLE_RUNTIME=ON
           -DTTMLIR_ENABLE_BINDINGS_PYTHON=OFF
           -DTT_RUNTIME_DEBUG=${TT_RUNTIME_DEBUG}
-          -DCMAKE_INSTALL_PREFIX=${CMAKE_INSTALL_PREFIX}
+          -DCMAKE_INSTALL_PREFIX=${TTMLIR_INSTALL_PREFIX}
           -DTTMLIR_ENABLE_DEBUG_STRINGS=${TT_RUNTIME_DEBUG}
         BUILD_COMMAND ${BUILD_COMMAND}
         GIT_REPOSITORY https://github.com/tenstorrent/tt-mlir.git
@@ -76,7 +82,6 @@ else()
     set_target_properties(TTMLIRRuntime PROPERTIES EXCLUDE_FROM_ALL TRUE IMPORTED_LOCATION ${TT_MLIR_RUNTIME_LIBRARY_PATH})
     add_dependencies(TTMLIRRuntime tt-mlir)
 
-    install(FILES ${TT_MLIR_COMPILER_LIBRARY_PATH} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib)
-    install(FILES ${TT_MLIR_RUNTIME_LIBRARY_PATH} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib)
+    install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}" USE_SOURCE_PERMISSIONS)
 
 endif()

tt_torch/__init__.py

@@ -14,7 +14,7 @@
 import importlib.util
 
 # 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_torch_tt_metal_libs"
+package_name = "tt-metal"
 spec = importlib.util.find_spec(package_name)
 if spec is not None:
     tt_metal_home = os.path.abspath(spec.submodule_search_locations[0])

tt_torch/csrc/CMakeLists.txt

@@ -34,7 +34,7 @@ target_link_libraries(TT_TORCH_MLIR PUBLIC
 )
 target_link_directories(TT_TORCH_MLIR PUBLIC
     ${TTMLIR_TOOLCHAIN_DIR}/lib
-    ${CMAKE_INSTALL_PREFIX}/lib
+    ${TTMLIR_INSTALL_PREFIX}/lib
 )
 
 set(TORCH_INSTALL_PREFIX ${TTTORCH_SOURCE_DIR}/env/venv/lib/python3.11/site-packages/torch)