tt-mlir uplift contains LLVM uplift.

Added torch-mlir as a third-party which we now build. This version of torch-mlir will be included in the wheel. Add cxxABI torch as a dependency
tenstorrent · Feb 21, 2025 · 8bff59d · 8bff59d
1 parent 3120d48
commit 8bff59d
Show file tree

Hide file tree

Showing 11 changed files with 52 additions and 11 deletions.
diff --git a/.github/workflows/generate-model-report.yml b/.github/workflows/generate-model-report.yml
@@ -17,7 +17,7 @@ jobs:
       - wormhole_b0
 
     container:
-      image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04:latest
+      image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04 :latest
       options: --user root --device /dev/tenstorrent/0
       volumes:
         - /dev/hugepages:/dev/hugepages

diff --git a/.github/workflows/generate-ttnn-md.yml b/.github/workflows/generate-ttnn-md.yml
@@ -14,7 +14,7 @@ jobs:
     timeout-minutes: 120
     runs-on: ubuntu-latest
     container:
-      image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04:latest
+      image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04 :latest
       options: --user root
 
     steps:

diff --git a/.github/workflows/run-build.yml b/.github/workflows/run-build.yml
@@ -22,7 +22,7 @@ jobs:
       - ${{ matrix.build.runs-on }}
 
     container:
-      image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04:latest
+      image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04 :latest
       options: --user root --device /dev/tenstorrent/0
       volumes:
         - /dev/hugepages:/dev/hugepages

diff --git a/.github/workflows/run-e2e-tests.yml b/.github/workflows/run-e2e-tests.yml
@@ -21,7 +21,7 @@ jobs:
       - ${{ matrix.build.runs-on }}
 
     container:
-      image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04:latest
+      image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04 :latest
       options: --user root --device /dev/tenstorrent/0 --shm-size=2gb
       volumes:
         - /dev/hugepages:/dev/hugepages

diff --git a/.github/workflows/run-op-by-op-model-tests.yml b/.github/workflows/run-op-by-op-model-tests.yml
@@ -218,7 +218,7 @@ jobs:
     name: "tests (${{ matrix.build.runs-on }}, ${{ matrix.build.name }})"
 
     container:
-      image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04:latest
+      image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04 :latest
       options: --user root --device /dev/tenstorrent/0 --shm-size=4gb
       volumes:
         - /dev/hugepages:/dev/hugepages

diff --git a/.github/workflows/run-tests.yml b/.github/workflows/run-tests.yml
@@ -26,7 +26,7 @@ jobs:
       - ${{ matrix.build.runs-on }}
 
     container:
-      image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04:latest
+      image: ghcr.io/tenstorrent/tt-torch/tt-torch-ci-ubuntu-22-04 :latest
       options: --user root --device /dev/tenstorrent/0 --shm-size=2gb
       volumes:
         - /dev/hugepages:/dev/hugepages

diff --git a/.gitignore b/.gitignore
@@ -10,3 +10,4 @@ results/*.xlsx
 third_party/toolchain
 third_party/tt-mlir
 third_party/vision
+third_party/torch-mlir
diff --git a/env/activate b/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:$(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)/env/venv:$(pwd)/env/venv/lib:$(pwd)/install/lib:$(pwd)/install:$(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"
diff --git a/requirements.txt b/requirements.txt
@@ -1,6 +1,4 @@
-torch-mlir
 torch@https://download.pytorch.org/whl/cpu-cxx11-abi/torch-2.5.0%2Bcpu.cxx11.abi-cp311-cp311-linux_x86_64.whl
--f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels
 black
 mdutils
 ninja
@@ -36,3 +34,4 @@ opencv-contrib-python
 setuptools
 wheel
 scikit-build
+nanobind
diff --git a/setup.py b/setup.py
@@ -3,7 +3,6 @@
 # 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
@@ -55,6 +54,10 @@ def run(self):
     description="TT PyTorch FrontEnd",
     long_description=long_description,
     long_description_content_type="text/markdown",
+    install_requires=[
+        "torch@https://download.pytorch.org/whl/cpu-cxx11-abi/torch-2.5.0%2Bcpu.cxx11.abi-cp311-cp311-linux_x86_64.whl",
+        "numpy",
+    ],
     cmake_args=cmake_args,
     cmdclass={
         "install_lib": install_metal_libs,

diff --git a/third_party/CMakeLists.txt b/third_party/CMakeLists.txt
@@ -3,7 +3,8 @@
 # SPDX-License-Identifier: Apache-2.0
 #
 
-set(TT_MLIR_VERSION "4b6791a3421cd9f227dc8f90acc7fc4d6eb19268")
+set(TT_MLIR_VERSION "625e478da24235d1aa236df10eccaa85d605470f")
+set(TORCH_MLIR_VERSION "98b6aee9f5c797c519b353de7bd580241d812cf5")
 
 if (BUILD_TOOLCHAIN)
     cmake_minimum_required(VERSION 3.20)
@@ -74,4 +75,41 @@ else()
 
     install(DIRECTORY ${TTMLIR_INSTALL_PREFIX}/ DESTINATION "${CMAKE_INSTALL_PREFIX}" USE_SOURCE_PERMISSIONS)
 
+    # torch-mlir build
+
+    set(TORCH_MLIR_PREFIX ${TTTORCH_SOURCE_DIR}/third_party/torch-mlir)
+    set(TORCH_MLIR_INSTALL_PREFIX ${TORCH_MLIR_PREFIX}/install)
+
+    set(TORCH_MLIR_BUILD_DIR ${TORCH_MLIR_PREFIX}/src/torch-mlir-build)
+    ExternalProject_Add(
+        torch-mlir
+        PREFIX ${TORCH_MLIR_PREFIX}
+        INSTALL_COMMAND ${CMAKE_COMMAND} --install ${TORCH_MLIR_BUILD_DIR}
+        CMAKE_GENERATOR Ninja
+        CMAKE_ARGS
+            -DCMAKE_BUILD_TYPE=Release
+            -DPython3_FIND_VIRTUALENV=ONLY
+            -DMLIR_DIR=${TTMLIR_TOOLCHAIN_DIR}/lib/cmake/mlir/
+            -DLLVM_DIR=${TTMLIR_TOOLCHAIN_DIR}/lib/cmake/llvm/
+            -DCMAKE_C_COMPILER=clang-17
+            -DCMAKE_CXX_COMPILER=clang++-17
+            -DCMAKE_CXX_COMPILER_LAUNCHER=ccache
+            -DTORCH_MLIR_ENABLE_STABLEHLO=ON
+            -DMLIR_ENABLE_BINDINGS_PYTHON=ON
+            -DLLVM_TARGETS_TO_BUILD=host
+            -DLLVM_ENABLE_ZSTD=OFF
+            -DTORCH_MLIR_ENABLE_LTC=ON
+            -DCMAKE_VISIBILITY_INLINES_HIDDEN=ON
+            -DCMAKE_C_VISIBILITY_PRESET=hidden
+            -DCMAKE_CXX_VISIBILITY_PRESET=hidden
+            -DCMAKE_INSTALL_PREFIX=${TORCH_MLIR_INSTALL_PREFIX}
+
+        GIT_SUBMODULES "externals/stablehlo"
+        GIT_PROGRESS ON
+        GIT_REPOSITORY https://github.com/llvm/torch-mlir.git
+        GIT_TAG ${TORCH_MLIR_VERSION}
+    )
+
+    install(DIRECTORY ${TORCH_MLIR_INSTALL_PREFIX}/python_packages/torch_mlir/torch_mlir DESTINATION "${CMAKE_INSTALL_PREFIX}" USE_SOURCE_PERMISSIONS)
+
 endif()