Switch stablehlo to embedded build.

tenstorrent · Aug 19, 2024 · cb90754 · cb90754
1 parent d5d51f6
commit cb90754
Show file tree

Hide file tree

Showing 20 changed files with 77 additions and 114 deletions.
diff --git a/.gitmodules b/.gitmodules
@@ -0,0 +1,3 @@
+[submodule "third_party/stablehlo"]
+	path = third_party/stablehlo
+	url = https://github.com/openxla/stablehlo
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -44,17 +44,21 @@ set(TTMLIR_TOOLCHAIN_DIR $ENV{TTMLIR_TOOLCHAIN_DIR})
 set(TTMLIR_SOURCE_DIR ${PROJECT_SOURCE_DIR})
 set(TTMLIR_BINARY_DIR ${PROJECT_BINARY_DIR})
 set(LLVM_LIT_TOOLS_DIR "${TTMLIR_TOOLCHAIN_DIR}/src/llvm-project/llvm/utils/lit")
+include_directories(${PROJECT_SOURCE_DIR}/include)
 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)
 link_directories(${LLVM_BUILD_LIBRARY_DIR})
 add_definitions(${LLVM_DEFINITIONS})
 include(TTMLIRPythonSitePackages)
 
+if (TTMLIR_ENABLE_STABLEHLO)
+  set(STABLEHLO_BUILD_EMBEDDED ON)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third_party/stablehlo ${CMAKE_CURRENT_BINARY_DIR}/stablehlo EXCLUDE_FROM_ALL)
+  include_directories(${CMAKE_CURRENT_SOURCE_DIR}/third_party/stablehlo)
+endif()
+
 add_subdirectory(third_party)
 add_subdirectory(include)
 add_subdirectory(lib)

diff --git a/env/CMakeLists.txt b/env/CMakeLists.txt
@@ -67,24 +67,4 @@ 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)
diff --git a/include/CMakeLists.txt b/include/CMakeLists.txt
@@ -1,3 +1 @@
 add_subdirectory(ttmlir)
-include_directories(${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo)
-include_directories(${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo-build)
diff --git a/include/ttmlir/Conversion/CMakeLists.txt b/include/ttmlir/Conversion/CMakeLists.txt
@@ -1,3 +1,11 @@
+include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo)
+include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo-build)
+include_directories(${TTMLIR_SOURCE_DIR}/include)
+
 set(LLVM_TARGET_DEFINITIONS Passes.td)
+if (TTMLIR_ENABLE_STABLEHLO)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name TTMLIRConversion -DTTMLIR_ENABLE_STABLEHLO)
+else()
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name TTMLIRConversion)
+endif()
 add_public_tablegen_target(TTMLIRConversionPassIncGen)
diff --git a/include/ttmlir/Conversion/Passes.h b/include/ttmlir/Conversion/Passes.h
@@ -8,7 +8,6 @@
 #ifdef TTMLIR_ENABLE_STABLEHLO
 #include "ttmlir/Conversion/StableHLOToTTIR/StableHLOToTTIR.h"
 #endif
-
 #include "ttmlir/Conversion/TTIRToTTNN/TTIRToTTNN.h"
 #include "ttmlir/Conversion/TTNNToEmitC/TTNNToEmitC.h"
 #include "ttmlir/Conversion/TosaToTTIR/TosaToTTIR.h"

diff --git a/include/ttmlir/Conversion/Passes.td b/include/ttmlir/Conversion/Passes.td
@@ -7,12 +7,14 @@
 
 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 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.";

diff --git a/include/ttmlir/Conversion/StableHLOToTTIR/StableHLOToTTIR.h b/include/ttmlir/Conversion/StableHLOToTTIR/StableHLOToTTIR.h
@@ -5,19 +5,14 @@
 #ifndef TTMLIR_CONVERSION_StableHLOToTTIR_StableHLOToTTIR_H
 #define TTMLIR_CONVERSION_StableHLOToTTIR_StableHLOToTTIR_H
 
+
 #include "mlir/IR/BuiltinOps.h"
 #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);
-}
+  std::unique_ptr<OperationPass<ModuleOp>> createConvertStableHLOToTTIRPass();
 #endif
 
 } // namespace mlir::tt

diff --git a/include/ttmlir/Dialect/TTKernel/IR/TTKernelOpsTypes.td b/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();

diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt
@@ -1,40 +1,26 @@
-include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo)
-include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo-build)
-
 add_subdirectory(CAPI)
 add_subdirectory(Conversion)
 add_subdirectory(Dialect)
 add_subdirectory(Target)
 
 # Shared library will include runtime code
 # so we only build it if runtime is enabled
-
-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)
+if (TTMLIR_ENABLE_RUNTIME)
+    add_subdirectory(SharedLib)
 endif()
 
 add_mlir_library(TTMLIRStatic STATIC RegisterAll.cpp
-
-  DEPENDS
-  ${ttmlir_dependencies}
-
   LINK_LIBS PUBLIC
-  ${link_libs}
+    MLIR
+    MLIRTTDialect
+    MLIRTTIRDialect
+    MLIRTTIRTransforms
+    TTMLIRConversions
+    MLIRTTIRAnalysis
+    MLIRTTNNDialect
+    MLIRTTNNTransforms
+    MLIRTTKernelDialect
+    MLIRTTMetalDialect
+    MLIRTTMetalTransforms
+    MLIRTTNNPipelines
 )
-
diff --git a/lib/Conversion/CMakeLists.txt b/lib/Conversion/CMakeLists.txt
@@ -1,15 +1,13 @@
 add_library(TTMLIRConversions INTERFACE)
 
+add_subdirectory(TosaToTTIR)
+add_subdirectory(TTNNToEmitC)
+add_subdirectory(TTIRToTTNN)
 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)
+include_directories(${TTMLIR_SOURCE_DIR}/include)
 
 set(link_libs 
 TTMLIRTosaToTTIR;
@@ -22,5 +20,5 @@ list(APPEND link_libs TTMLIRStableHLOToTTIR)
 endif()
 
 target_link_libraries(TTMLIRConversions INTERFACE
-  ${link_libs}
+${link_libs}
 )
diff --git a/lib/Conversion/StableHLOToTTIR/CMakeLists.txt b/lib/Conversion/StableHLOToTTIR/CMakeLists.txt
@@ -1,7 +1,9 @@
 include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo)
 include_directories(${TTMLIR_TOOLCHAIN_DIR}/src/stablehlo-build)
+include_directories(${TTMLIR_SOURCE_DIR}/include)
+include_directories(${PROJECT_SOURCE_DIR}/include)
+include_directories(${PROJECT_SOURCE_DIR}/build/stablehlo)
 
-get_property(STABLEHLO_LIBS GLOBAL PROPERTY STABLEHLO_LIBS)
 add_mlir_library(TTMLIRStableHLOToTTIR
   StableHLOToTTIR.cpp
 
@@ -12,6 +14,5 @@ add_mlir_library(TTMLIRStableHLOToTTIR
   TTMLIRConversionPassIncGen
 
   LINK_LIBS PUBLIC
-  MLIRTransformUtils
-  ${STABLEHLO_LIBS}
+  MLIR
 )
diff --git a/lib/RegisterAll.cpp b/lib/RegisterAll.cpp
@@ -17,7 +17,10 @@
 #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,
@@ -28,12 +31,14 @@ void mlir::tt::registerAllDialects(mlir::DialectRegistry &registry) {
               mlir::scf::SCFDialect, mlir::cf::ControlFlowDialect,
               mlir::tosa::TosaDialect, mlir::vector::VectorDialect,
               mlir::emitc::EmitCDialect>();
-  mlir::stablehlo::registerAllDialects(registry);
+#if TTMLIR_ENABLE_STABLEHLO
+        mlir::stablehlo::registerAllDialects(registry);
+#endif
 }
 
 void mlir::tt::registerAllPasses() {
   // Register all dialect conversion passes
-
+  //
   mlir::tt::registerTTMLIRConversionPasses();
 
   mlir::tt::ttir::registerPasses();

diff --git a/lib/SharedLib/CMakeLists.txt b/lib/SharedLib/CMakeLists.txt
@@ -20,18 +20,13 @@ set(TTMLIR_LIBS
     TTMLIRTTNNToEmitC
 )
 
-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)
 
 add_dependencies(TTMLIR
     ${TTMLIR_LIBS}
     ${TTNN_RUNTIME_LIBS}
     ${TTNN_LIBS}
-    ${STABLEHLO_LIBS}
 )
 
 target_link_libraries(TTMLIR PRIVATE
@@ -46,7 +41,7 @@ target_link_libraries(TTMLIR PRIVATE
     -Wl,--no-whole-archive
 
     ${TTNN_LIBS}
-    ${STABLEHLO_LIBS}
+
     flatbuffers
 )
 

diff --git a/runtime/lib/binary.cpp b/runtime/lib/binary.cpp
@@ -4,7 +4,6 @@
 
 #include <fstream>
 
-#pragma clang diagnostic ignored "-Wcovered-switch-default"
 #include "flatbuffers/idl.h"
 
 #include "tt/runtime/types.h"

diff --git a/test/ttmlir/ttmlir-opt.mlir b/test/ttmlir/ttmlir-opt.mlir
@@ -1,3 +1,18 @@
 // RUN: ttmlir-opt --show-dialects | FileCheck %s
 // CHECK: Available Dialects:
-// CHECK-SAME: arith,builtin,cf,emitc,func,linalg,ml_program,scf,tensor,tosa,tt,ttir,ttkernel,ttmetal,ttnn,vector
+// CHECK: arith
+// CHECK: builtin
+// CHECK: cf
+// CHECK: emitc
+// CHECK: func
+// CHECK: linalg
+// CHECK: ml_program
+// CHECK: scf
+// CHECK: tensor
+// CHECK: tosa
+// CHECK: tt
+// CHECK: ttir
+// CHECK: ttkernel
+// CHECK: ttmetal
+// CHECK: ttnn
+// CHECK: vector
diff --git a/third_party/CMakeLists.txt b/third_party/CMakeLists.txt
@@ -13,13 +13,6 @@ 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
@@ -53,7 +46,7 @@ ExternalProject_Add(
     -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}
+    -DENABLE_TRACY=${TT_RUNTIME_ENABLE_PERF_TRACE}
   GIT_REPOSITORY https://github.com/tenstorrent/tt-metal.git
   GIT_TAG f6a2e5cb2b857bf4c72401bea68adf98c25bbe47
   GIT_PROGRESS ON

diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt
@@ -1,7 +1,2 @@
-include(ExternalProject)
-
-include_directories(${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo)
-include_directories(${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo-build)
-
 add_subdirectory(ttmlir-opt)
 add_subdirectory(ttmlir-translate)
diff --git a/tools/ttmlir-opt/CMakeLists.txt b/tools/ttmlir-opt/CMakeLists.txt
@@ -1,17 +1,9 @@
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
-set(LIBS ${dialect_libs} ${conversion_libs} StablehloBase TTNNTargetFlatbuffer MLIRTTIRDialect TTMLIRConversions MLIRTransformUtils MLIROptLib MLIRTargetCpp TTMLIRStatic CheckOps ChloCAPI ChloOps InterpreterOps StablehloAssemblyFormat StablehloBase StablehloBroadcastUtils StablehloCAPI StablehloLinalgTransforms StablehloOps StablehloPasses StablehloPortableApi StablehloRegister StablehloSerialization StablehloTOSATransforms StablehloTestUtils StablehloTypeInference Version VhloCAPI VhloOps VhloTypes)
-
-if (TTMLIR_ENABLE_STABLEHLO)
-	list(APPEND LIBS TTMLIRStableHLOToTTIR)
-endif()
-
+set(LIBS ${dialect_libs} ${conversion_libs} MLIROptLib MLIRTargetCpp TTMLIRStatic)
 add_llvm_executable(ttmlir-opt ttmlir-opt.cpp)
 
-target_link_directories(ttmlir-opt PRIVATE ${TTMLIR_PROJECT_DIR}/build/lib)
-target_link_libraries(ttmlir-opt PRIVATE ${LIBS} libStablehloBase.a)
-target_link_directories(ttmlir-opt PRIVATE ${TTMLIR_TOOLCHAIN_DIR}/lib)
-
 llvm_update_compile_flags(ttmlir-opt)
+target_link_libraries(ttmlir-opt PRIVATE ${LIBS})
 
-mlir_check_all_link_libraries(ttmlir-opt ${LIBS})
+mlir_check_all_link_libraries(ttmlir-opt)
diff --git a/tools/ttmlir-opt/ttmlir-opt.cpp b/tools/ttmlir-opt/ttmlir-opt.cpp
@@ -7,9 +7,6 @@
 #include "mlir/InitAllPasses.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
-#ifdef TTMLIR_ENABLE_STABLEHLO
-#include "stablehlo/dialect/Register.h"
-#endif
 
 #include "ttmlir/RegisterAll.h"
 
@@ -19,9 +16,7 @@ int main(int argc, char **argv) {
 
   mlir::DialectRegistry registry;
   mlir::tt::registerAllDialects(registry);
-#ifdef TTMLIR_ENABLE_STABLEHLO  
-  mlir::stablehlo::registerAllDialects(registry);
-#endif
+
   return mlir::asMainReturnCode(
       mlir::MlirOptMain(argc, argv, "ttmlir optimizer driver\n", registry));
 }