From fecb19d0764ce650638e1295d5332d7fb3c1512c Mon Sep 17 00:00:00 2001
From: uazizTT <uaziz@tenstorrent.com>
Date: Wed, 7 Aug 2024 10:26:07 -0400
Subject: [PATCH 1/2] Add torch-mlir to tt-mlir build in order to utilize onnx
 to mlir conversion.

---
 CMakeLists.txt               |  1 +
 env/build-requirements.txt   |  3 ++
 third_party/CMakeLists.txt   | 75 ++----------------------------------
 third_party/torch-mlir.cmake | 19 +++++++++
 third_party/tt-metal.cmake   | 73 +++++++++++++++++++++++++++++++++++
 5 files changed, 100 insertions(+), 71 deletions(-)
 create mode 100644 third_party/torch-mlir.cmake
 create mode 100644 third_party/tt-metal.cmake

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 374a31d44..adaa41282 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -6,6 +6,7 @@ if (NOT DEFINED ENV{TTMLIR_ENV_ACTIVATED})
 endif()
 
 option(TT_RUNTIME_ENABLE_PERF_TRACE "Enable performance mode" OFF)
+option(TTMLIR_ENABLE_TORCH_MLIR "Enable TORCH MLIR support" OFF)
 
 set(CMAKE_BUILD_WITH_INSTALL_NAME_DIR ON)
 
diff --git a/env/build-requirements.txt b/env/build-requirements.txt
index 70bd46bca..49ded8f02 100644
--- a/env/build-requirements.txt
+++ b/env/build-requirements.txt
@@ -6,3 +6,6 @@ wheel
 setuptools
 black
 pre-commit
+numpy
+pybind11
+Torch
diff --git a/third_party/CMakeLists.txt b/third_party/CMakeLists.txt
index 7d9493899..6dc263242 100644
--- a/third_party/CMakeLists.txt
+++ b/third_party/CMakeLists.txt
@@ -1,73 +1,6 @@
 include(ExternalProject)
 
-if ("$ENV{ARCH_NAME}" STREQUAL "grayskull")
-  set(ARCH_NAME "grayskull")
-  set(ARCH_EXTRA_DIR "grayskull")
-elseif ("$ENV{ARCH_NAME}" STREQUAL "wormhole_b0")
-  set(ARCH_NAME "wormhole")
-  set(ARCH_EXTRA_DIR "wormhole/wormhole_b0_defines")
-elseif ("$ENV{ARCH_NAME}" STREQUAL "blackhole")
-  set(ARCH_NAME "blackhole")
-  set(ARCH_EXTRA_DIR "blackhole")
-else()
-  message(FATAL_ERROR "Unsupported ARCH_NAME: $ENV{ARCH_NAME}")
-endif()
-
-if (TT_RUNTIME_ENABLE_PERF_TRACE)
-  add_compile_definitions(TRACY_ENABLE)
-  set(ENV{ENABLE_TRACY} "1")
-else()
-  set(ENV{ENABLE_TRACY} "0")
-endif()
-
-set(TTMETAL_INCLUDE_DIRS
-  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/ttnn/cpp
-  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal
-  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal
-  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal/third_party/umd
-  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal/third_party/fmt
-  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal/hw/inc
-  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal/hw/inc/${ARCH_NAME}
-  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal/hw/inc/${ARCH_EXTRA_DIR}
-  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal/third_party/umd/src/firmware/riscv/${ARCH_NAME}
-  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_eager
-  PARENT_SCOPE
-)
-
-set(TTMETAL_LIBRARY_DIR ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal-build/lib)
-set(TTNN_LIBRARY_PATH ${TTMETAL_LIBRARY_DIR}/_ttnn.so)
-set(TTMETAL_LIBRARY_PATH ${TTMETAL_LIBRARY_DIR}/libtt_metal.so)
-set(TTEAGER_LIBRARY_PATH ${TTMETAL_LIBRARY_DIR}/libtt_eager.so)
-
-set(TTMETAL_LIBRARY_DIR ${TTMETAL_LIBRARY_DIR} PARENT_SCOPE)
-set(TTNN_LIBRARY_PATH ${TTNN_LIBRARY_PATH} PARENT_SCOPE)
-set(TTMETAL_LIBRARY_PATH ${TTMETAL_LIBRARY_PATH} PARENT_SCOPE)
-set(TTEAGER_LIBRARY_PATH ${TTEAGER_LIBRARY_PATH} PARENT_SCOPE)
-
-
-ExternalProject_Add(
-  tt-metal
-  PREFIX ${TTMLIR_SOURCE_DIR}/third_party/tt-metal
-  CMAKE_GENERATOR Ninja
-  CMAKE_ARGS
-    -DCMAKE_BUILD_TYPE=Release
-    -DCMAKE_C_COMPILER=${CMAKE_C_COMPILER}
-    -DCMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER}
-    -DCMAKE_CXX_COMPILER_LAUNCHER=${CMAKE_CXX_COMPILER_LAUNCHER}
-    -DTRACY_ENABLE=${TT_RUNTIME_ENABLE_PERF_TRACE}
-  GIT_REPOSITORY https://github.com/tenstorrent/tt-metal.git
-  GIT_TAG v0.49.0
-  GIT_PROGRESS ON
-  BUILD_BYPRODUCTS ${TTNN_LIBRARY_PATH} ${TTMETAL_LIBRARY_PATH} ${TTEAGER_LIBRARY_PATH}
-)
-
-set_target_properties(tt-metal PROPERTIES EXCLUDE_FROM_ALL TRUE)
-
-list(APPEND library_names TTNN_LIBRARY TTEAGER_LIBRARY TTMETAL_LIBRARY)
-list(APPEND library_paths ${TTNN_LIBRARY_PATH} ${TTMETAL_LIBRARY_PATH} ${TTEAGER_LIBRARY_PATH})
-
-foreach(lib_name lib_path IN ZIP_LISTS library_names library_paths)
-  add_library(${lib_name} SHARED IMPORTED GLOBAL)
-  set_target_properties(${lib_name} PROPERTIES EXCLUDE_FROM_ALL TRUE IMPORTED_LOCATION ${lib_path})
-  add_dependencies(${lib_name} tt-metal)
-endforeach()
+include(tt-metal.cmake)
+if (TTMLIR_ENABLE_TORCH_MLIR)
+  include(torch-mlir.cmake)
+endif ()
diff --git a/third_party/torch-mlir.cmake b/third_party/torch-mlir.cmake
new file mode 100644
index 000000000..a15dafd22
--- /dev/null
+++ b/third_party/torch-mlir.cmake
@@ -0,0 +1,19 @@
+include(ExternalProject)
+
+set(TORCH_MLIR_VERSION "b1a232222f12d4e5640fd62320d02f6c832bdc4e")
+
+ExternalProject_Add(
+    torch-mlir
+    PREFIX ${TTMLIR_TOOLCHAIN_DIR}
+    CMAKE_GENERATOR Ninja
+    CMAKE_ARGS
+    -DCMAKE_BUILD_TYPE=Release
+    -DPython3_FIND_VIRTUALENV=ONLY
+    -DCMAKE_INSTALL_PREFIX=${TTMLIR_TOOLCHAIN_DIR}
+    -DMLIR_DIR=${TTMLIR_TOOLCHAIN_DIR}/llvm-project/lib/cmake/mlir/
+    -DLLVM_DIR=${TTMLIR_TOOLCHAIN_DIR}/llvm-project/lib/cmake/llvm/
+    -DMLIR_ENABLE_BINDINGS_PYTHON=ON
+    -DLLVM_TARGETS_TO_BUILD=host
+    GIT_REPOSITORY https://github.com/llvm/torch-mlir
+    GIT_TAG ${TORCH_MLIR_VERSION}
+)
diff --git a/third_party/tt-metal.cmake b/third_party/tt-metal.cmake
new file mode 100644
index 000000000..7d9493899
--- /dev/null
+++ b/third_party/tt-metal.cmake
@@ -0,0 +1,73 @@
+include(ExternalProject)
+
+if ("$ENV{ARCH_NAME}" STREQUAL "grayskull")
+  set(ARCH_NAME "grayskull")
+  set(ARCH_EXTRA_DIR "grayskull")
+elseif ("$ENV{ARCH_NAME}" STREQUAL "wormhole_b0")
+  set(ARCH_NAME "wormhole")
+  set(ARCH_EXTRA_DIR "wormhole/wormhole_b0_defines")
+elseif ("$ENV{ARCH_NAME}" STREQUAL "blackhole")
+  set(ARCH_NAME "blackhole")
+  set(ARCH_EXTRA_DIR "blackhole")
+else()
+  message(FATAL_ERROR "Unsupported ARCH_NAME: $ENV{ARCH_NAME}")
+endif()
+
+if (TT_RUNTIME_ENABLE_PERF_TRACE)
+  add_compile_definitions(TRACY_ENABLE)
+  set(ENV{ENABLE_TRACY} "1")
+else()
+  set(ENV{ENABLE_TRACY} "0")
+endif()
+
+set(TTMETAL_INCLUDE_DIRS
+  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/ttnn/cpp
+  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal
+  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal
+  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal/third_party/umd
+  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal/third_party/fmt
+  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal/hw/inc
+  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal/hw/inc/${ARCH_NAME}
+  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal/hw/inc/${ARCH_EXTRA_DIR}
+  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_metal/third_party/umd/src/firmware/riscv/${ARCH_NAME}
+  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/tt_eager
+  PARENT_SCOPE
+)
+
+set(TTMETAL_LIBRARY_DIR ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal-build/lib)
+set(TTNN_LIBRARY_PATH ${TTMETAL_LIBRARY_DIR}/_ttnn.so)
+set(TTMETAL_LIBRARY_PATH ${TTMETAL_LIBRARY_DIR}/libtt_metal.so)
+set(TTEAGER_LIBRARY_PATH ${TTMETAL_LIBRARY_DIR}/libtt_eager.so)
+
+set(TTMETAL_LIBRARY_DIR ${TTMETAL_LIBRARY_DIR} PARENT_SCOPE)
+set(TTNN_LIBRARY_PATH ${TTNN_LIBRARY_PATH} PARENT_SCOPE)
+set(TTMETAL_LIBRARY_PATH ${TTMETAL_LIBRARY_PATH} PARENT_SCOPE)
+set(TTEAGER_LIBRARY_PATH ${TTEAGER_LIBRARY_PATH} PARENT_SCOPE)
+
+
+ExternalProject_Add(
+  tt-metal
+  PREFIX ${TTMLIR_SOURCE_DIR}/third_party/tt-metal
+  CMAKE_GENERATOR Ninja
+  CMAKE_ARGS
+    -DCMAKE_BUILD_TYPE=Release
+    -DCMAKE_C_COMPILER=${CMAKE_C_COMPILER}
+    -DCMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER}
+    -DCMAKE_CXX_COMPILER_LAUNCHER=${CMAKE_CXX_COMPILER_LAUNCHER}
+    -DTRACY_ENABLE=${TT_RUNTIME_ENABLE_PERF_TRACE}
+  GIT_REPOSITORY https://github.com/tenstorrent/tt-metal.git
+  GIT_TAG v0.49.0
+  GIT_PROGRESS ON
+  BUILD_BYPRODUCTS ${TTNN_LIBRARY_PATH} ${TTMETAL_LIBRARY_PATH} ${TTEAGER_LIBRARY_PATH}
+)
+
+set_target_properties(tt-metal PROPERTIES EXCLUDE_FROM_ALL TRUE)
+
+list(APPEND library_names TTNN_LIBRARY TTEAGER_LIBRARY TTMETAL_LIBRARY)
+list(APPEND library_paths ${TTNN_LIBRARY_PATH} ${TTMETAL_LIBRARY_PATH} ${TTEAGER_LIBRARY_PATH})
+
+foreach(lib_name lib_path IN ZIP_LISTS library_names library_paths)
+  add_library(${lib_name} SHARED IMPORTED GLOBAL)
+  set_target_properties(${lib_name} PROPERTIES EXCLUDE_FROM_ALL TRUE IMPORTED_LOCATION ${lib_path})
+  add_dependencies(${lib_name} tt-metal)
+endforeach()

From 0689cec98b02fa1771d43a619bfdf4a7ea203067 Mon Sep 17 00:00:00 2001
From: uazizTT <uaziz@tenstorrent.com>
Date: Wed, 7 Aug 2024 12:31:43 -0400
Subject: [PATCH 2/2] Remove extra dependencies and enable torch-mlir build in
 github workflow.

---
 .github/workflows/build.yml | 1 +
 env/build-requirements.txt  | 2 --
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index fc6fb1c3f..eb46daa3d 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -62,6 +62,7 @@ jobs:
         -DCMAKE_BUILD_TYPE=${{ matrix.build.build_type }} \
         -DTTMLIR_ENABLE_RUNTIME=${{ matrix.build.enable_runtime }} \
         -DTTMLIR_ENABLE_RUNTIME_TESTS=${{ matrix.build.enable_runtime }} \
+        -DTTMLIR_ENABLE_TORCH_MLIR=ON \
         -S ${{ github.workspace }}
 
     - name: Build
diff --git a/env/build-requirements.txt b/env/build-requirements.txt
index 49ded8f02..8103c3b85 100644
--- a/env/build-requirements.txt
+++ b/env/build-requirements.txt
@@ -6,6 +6,4 @@ wheel
 setuptools
 black
 pre-commit
-numpy
-pybind11
 Torch