Move StableHLO to env dir and resolve build issues.

CMakeLists.txt

@@ -9,6 +9,10 @@ option(TT_RUNTIME_ENABLE_PERF_TRACE "Enable performance mode" OFF)
 option(TTMLIR_ENABLE_RUNTIME "Enable runtime" OFF)
 option(TTMLIR_ENABLE_STABLEHLO "Enable StableHLO support" ON)
 
+if (TTMLIR_ENABLE_STABLEHLO)
+  add_compile_definitions(TTMLIR_ENABLE_STABLEHLO)
+endif()
+
 set(CMAKE_BUILD_WITH_INSTALL_NAME_DIR ON)
 
 set(CMAKE_CXX_STANDARD 17 CACHE STRING "C++ standard to conform to")
@@ -42,6 +46,8 @@ set(TTMLIR_BINARY_DIR ${PROJECT_BINARY_DIR})
 set(LLVM_LIT_TOOLS_DIR "${TTMLIR_TOOLCHAIN_DIR}/src/llvm-project/llvm/utils/lit")
 include_directories(SYSTEM ${LLVM_INCLUDE_DIRS})
 include_directories(SYSTEM ${MLIR_INCLUDE_DIRS})
+include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo)
+include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo-build)
 include_directories(${TTMLIR_SOURCE_DIR}/include)
 include_directories(${TTMLIR_BINARY_DIR}/include)
 include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/flatbuffers/include)

env/CMakeLists.txt

@@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.20.0)
 project(ttmlir-toolchain LANGUAGES CXX C)
 
 set(FLATBUFFERS_VERSION "fb9afbafc7dfe226b9db54d4923bfb8839635274")
-set(LLVM_PROJECT_VERSION "a6d7828f4c50c1ec7b0b5f61fe59d7a768175dcc")
+set(LLVM_PROJECT_VERSION "9ddfe62f5c11e3f65f444209f514029ded2d58b9")
 
 include(ExternalProject)
 
@@ -67,4 +67,24 @@ ExternalProject_Add(
     DEPENDS python-venv
 )
 
+set(STABLE_HLO_VERSION "v1.5.0")
+
+ExternalProject_Add(
+    stablehlo
+    PREFIX ${TTMLIR_TOOLCHAIN_DIR}
+    CMAKE_GENERATOR Ninja
+    CMAKE_ARGS
+      -DCMAKE_BUILD_TYPE=Release
+      -DCMAKE_C_COMPILER=${CMAKE_C_COMPILER}
+      -DCMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER}
+      -DCMAKE_INSTALL_PREFIX=${TTMLIR_TOOLCHAIN_DIR}
+      -DCMAKE_CXX_COMPILER_LAUNCHER=${CMAKE_CXX_COMPILER_LAUNCHER}
+      # TODO: Figure out if we need this (or benefit from it) when it comes to linking stablehlo.
+      #-DSTABLEHLO_BUILD_EMBEDDED=ON
+      -DMLIR_DIR=${TTMLIR_TOOLCHAIN_DIR}/lib/cmake/mlir
+    GIT_REPOSITORY https://github.com/openxla/stablehlo
+    GIT_TAG ${STABLE_HLO_VERSION}
+    GIT_PROGRESS ON
+)
+
 add_custom_target(llvm-lit ALL COMMAND cp ${TTMLIR_TOOLCHAIN_DIR}/src/llvm-project-build/bin/llvm-lit ${TTMLIR_TOOLCHAIN_DIR}/bin/llvm-lit DEPENDS llvm-project)

include/ttmlir/Conversion/Passes.td

@@ -7,15 +7,12 @@
 
 include "mlir/Pass/PassBase.td"
 
-
-#ifdef TTMLIR_ENABLE_STABLEHLO
 def ConvertStableHLOToTTIR : Pass<"convert-stablehlo-to-ttir", "::mlir::ModuleOp"> {
   let summary = "Convert StableHLO dialect to TTIR dialect.";
   let constructor = "createConvertStableHLOToTTIRPass()";
   // TODO(mrakita): Probably will need to add some include here for StableHLO dialect.
   let dependentDialects = ["mlir::stablehlo::StablehloDialect", "mlir::tt::ttir::TTIRDialect"];
 }
-#endif
 
 def ConvertTosaToTTIR : Pass<"convert-tosa-to-ttir", "::mlir::ModuleOp"> {
   let summary = "Convert TOSA dialect to TTIR dialect.";

include/ttmlir/Conversion/StableHLOToTTIR/StableHLOToTTIR.h

@@ -9,9 +9,17 @@
 #include "mlir/Pass/Pass.h"
 
 namespace mlir::tt {
+
 #ifdef TTMLIR_ENABLE_STABLEHLO
 std::unique_ptr<OperationPass<ModuleOp>> createConvertStableHLOToTTIRPass();
+#else
+// define an empty function in case conversion definition does not exist. 
+std::unique_ptr<OperationPass<ModuleOp>> createConvertStableHLOToTTIRPass()
+{
+return std::unique_ptr<OperationPass<ModuleOp>>(nullptr);
+}
 #endif
+
 } // namespace mlir::tt
 
 #endif

include/ttmlir/Dialect/TTKernel/IR/TTKernelOpsTypes.td

@@ -31,8 +31,8 @@ def TTKernel_CB : TTKernel_Type<"CB", "cb"> {
         static CBType get(::mlir::MLIRContext *context,
                               uint64_t address,
                               uint64_t port,
-                              MemRefType memref) {
-          uint64_t numBuffers = 1;
+                              MemRefType memref,
+                              uint64_t numBuffers = 1) {
           uint64_t pageSize = 0;
           if (::mlir::isa<::mlir::tt::TileType>(memref.getElementType())) {
             pageSize = ::mlir::cast<::mlir::tt::TileType>(memref.getElementType()).getSizeBytes();

lib/CMakeLists.txt

@@ -1,5 +1,5 @@
-include_directories(${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo)
-include_directories(${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo-build)
+include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo)
+include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo-build)
 
 add_subdirectory(CAPI)
 add_subdirectory(Conversion)
@@ -8,29 +8,33 @@ add_subdirectory(Target)
 
 # Shared library will include runtime code
 # so we only build it if runtime is enabled
-if (TTMLIR_ENABLE_RUNTIME)
-    add_subdirectory(SharedLib)
+
+set(ttmlir_dependencies "")
+set(link_libs 
+MLIR;
+MLIRTTDialect;
+MLIRTTIRDialect;
+MLIRTTIRTransforms;
+TTMLIRConversions;
+MLIRTTIRAnalysis;
+MLIRTTNNDialect;
+MLIRTTNNTransforms;
+MLIRTTKernelDialect;
+MLIRTTMetalDialect;
+MLIRTTMetalTransforms;
+MLIRTTNNPipelines
+)
+
+if (TTMLIR_ENABLE_STABLEHLO)
+    list(APPEND link_libs TTMLIRStableHLOToTTIR)
 endif()
 
 add_mlir_library(TTMLIRStatic STATIC RegisterAll.cpp
 
   DEPENDS
-  stablehlo 
+  ${ttmlir_dependencies}
 
   LINK_LIBS PUBLIC
-    MLIR
-    MLIRTTDialect
-    MLIRTTIRDialect
-    MLIRTTIRTransforms
-    TTMLIRConversions
-    MLIRTTIRAnalysis
-    MLIRTTNNDialect
-    MLIRTTNNTransforms
-    MLIRTTKernelDialect
-    MLIRTTMetalDialect
-    MLIRTTMetalTransforms
-    MLIRTTNNPipelines
-    TTMLIRStableHLOToTTIR
+  ${link_libs}
 )
 
-add_dependencies(TTMLIRStatic stablehlo)

lib/Conversion/CMakeLists.txt

@@ -1,18 +1,26 @@
-include_directories(${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo)
-include_directories(${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo-build)
+add_library(TTMLIRConversions INTERFACE)
 
+if (TTMLIR_ENABLE_STABLEHLO)
+include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo)
+include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo-build)
 add_subdirectory(StableHLOToTTIR)
+add_dependencies(TTMLIRConversions stablehlo)
+endif()
+
 add_subdirectory(TosaToTTIR)
 add_subdirectory(TTNNToEmitC)
 add_subdirectory(TTIRToTTNN)
 
-add_library(TTMLIRConversions INTERFACE)
+set(link_libs 
+TTMLIRTosaToTTIR;
+TTMLIRTTNNToEmitC;
+TTMLIRTTIRToTTNN
+)
 
-add_dependencies(TTMLIRConversions stablehlo)
+if (TTMLIR_ENABLE_STABLEHLO)
+list(APPEND link_libs TTMLIRStableHLOToTTIR)
+endif()
 
 target_link_libraries(TTMLIRConversions INTERFACE
-  TTMLIRStableHLOToTTIR
-  TTMLIRTosaToTTIR
-  TTMLIRTTNNToEmitC
-  TTMLIRTTIRToTTNN
+  ${link_libs}
 )

lib/Conversion/StableHLOToTTIR/CMakeLists.txt

@@ -1,23 +1,17 @@
-include_directories(${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo)
-include_directories(${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo-build)
+include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo)
+include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo-build)
 
 get_property(STABLEHLO_LIBS GLOBAL PROPERTY STABLEHLO_LIBS)
-if (TTMLIR_ENABLE_STABLEHLO)
 add_mlir_library(TTMLIRStableHLOToTTIR
   StableHLOToTTIR.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/include/ttmlir/Conversion/StableHLOToTTIR
-  # StableHLO includes
-  ${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo
-  ${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo-build
 
   DEPENDS
   TTMLIRConversionPassIncGen
-  stablehlo
 
   LINK_LIBS PUBLIC
   MLIRTransformUtils
   ${STABLEHLO_LIBS}
 )
-endif ()

lib/Conversion/StableHLOToTTIR/StableHLOToTTIR.cpp

@@ -19,7 +19,7 @@
 #include <mlir/Pass/Pass.h>
 #include <mlir/Support/LogicalResult.h>
 #include <mlir/Transforms/DialectConversion.h>
-#ifdef TTMLIR_ENABLE_STABLEHLO
+
 #include "stablehlo/dialect/StablehloOps.h"
 using namespace mlir;
 using namespace tt;
@@ -101,4 +101,3 @@ std::unique_ptr<OperationPass<ModuleOp>> createConvertStableHLOToTTIRPass() {
 }
 
 } // namespace mlir::tt
-#endif

lib/RegisterAll.cpp

@@ -17,9 +17,7 @@
 #include "ttmlir/Dialect/TTNN/IR/TTNN.h"
 #include "ttmlir/Dialect/TTNN/Pipelines/Passes.h"
 #include "ttmlir/Dialect/TTNN/Transforms/Passes.h"
-#ifdef TTMLIR_ENABLE_STABLEHLO
 #include "stablehlo/dialect/Register.h"
-#endif
 void mlir::tt::registerAllDialects(mlir::DialectRegistry &registry) {
   registry
       .insert<mlir::tt::TTDialect, mlir::tt::ttir::TTIRDialect,
@@ -30,9 +28,7 @@ void mlir::tt::registerAllDialects(mlir::DialectRegistry &registry) {
               mlir::scf::SCFDialect, mlir::cf::ControlFlowDialect,
               mlir::tosa::TosaDialect, mlir::vector::VectorDialect,
               mlir::emitc::EmitCDialect>();
-#ifdef TTMLIR_ENABLE_STABLEHLO 
   mlir::stablehlo::registerAllDialects(registry);
-#endif
 }
 
 void mlir::tt::registerAllPasses() {

lib/SharedLib/CMakeLists.txt

@@ -18,11 +18,12 @@ set(TTMLIR_LIBS
     MLIRTTIRAnalysis
     MLIRTTNNPipelines
     TTMLIRTTNNToEmitC
-if (TTMLIR_ENABLE_STABLEHLO)
-    TTMLIRStableHLOToTTIR
-endif ()
 )
 
+if (TTMLIR_ENABLE_STABLEHLO)
+list(APPEND TTMLIR_LIBS TTMLIRStableHLOToTTIR)
+endif()
+
 # We supply empty.cpp because CMake does not allow creating a library without sources.
 add_library(TTMLIR SHARED empty.cpp)
 

third_party/CMakeLists.txt

@@ -1,6 +1,72 @@
 include(ExternalProject)
 
-include(tt-metal.cmake)
-if (TTMLIR_ENABLE_STABLEHLO)
-  include(stablehlo.cmake)
+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/ttnn/cpp/ttnn/deprecated
+  ${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
+  ${PROJECT_SOURCE_DIR}/third_party/tt-metal/src/tt-metal/.cpmcache/reflect/e75434c4c5f669e4a74e4d84e0a30d7249c1e66f
+  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(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)
+
+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 f6a2e5cb2b857bf4c72401bea68adf98c25bbe47
+  GIT_PROGRESS ON
+  BUILD_BYPRODUCTS ${TTNN_LIBRARY_PATH} ${TTMETAL_LIBRARY_PATH}
+)
+
+set_target_properties(tt-metal PROPERTIES EXCLUDE_FROM_ALL TRUE)
+
+list(APPEND library_names TTNN_LIBRARY TTMETAL_LIBRARY)
+list(APPEND library_paths ${TTNN_LIBRARY_PATH} ${TTMETAL_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()