Enclose StableHLO features to be only enabled when TTMLIR_ENABLE_STAB…

…LEHLO flag is set.
tenstorrent · Aug 19, 2024 · bc02c9b · bc02c9b
1 parent d9dd557
commit bc02c9b
Show file tree

Hide file tree

Showing 9 changed files with 61 additions and 14 deletions.
diff --git a/include/ttmlir/Conversion/Passes.h b/include/ttmlir/Conversion/Passes.h
@@ -5,7 +5,9 @@
 #ifndef TTMLIR_CONVERSION_PASSES_H
 #define TTMLIR_CONVERSION_PASSES_H
 
+#ifdef TTMLIR_ENABLE_STABLEHLO
 #include "ttmlir/Conversion/StableHLOToTTIR/StableHLOToTTIR.h"
+#endif
 
 #include "ttmlir/Conversion/TTIRToTTNN/TTIRToTTNN.h"
 #include "ttmlir/Conversion/TTNNToEmitC/TTNNToEmitC.h"

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

diff --git a/include/ttmlir/Conversion/StableHLOToTTIR/StableHLOToTTIR.h b/include/ttmlir/Conversion/StableHLOToTTIR/StableHLOToTTIR.h
@@ -9,9 +9,9 @@
 #include "mlir/Pass/Pass.h"
 
 namespace mlir::tt {
-
+#ifdef TTMLIR_ENABLE_STABLEHLO
 std::unique_ptr<OperationPass<ModuleOp>> createConvertStableHLOToTTIRPass();
-
+#endif
 } // namespace mlir::tt
 
 #endif
diff --git a/lib/Conversion/StableHLOToTTIR/CMakeLists.txt b/lib/Conversion/StableHLOToTTIR/CMakeLists.txt
@@ -2,7 +2,7 @@ include_directories(${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo)
 include_directories(${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo-build)
 
 get_property(STABLEHLO_LIBS GLOBAL PROPERTY STABLEHLO_LIBS)
-
+if (TTMLIR_ENABLE_STABLEHLO)
 add_mlir_library(TTMLIRStableHLOToTTIR
   StableHLOToTTIR.cpp
 
@@ -20,3 +20,4 @@ add_mlir_library(TTMLIRStableHLOToTTIR
   MLIRTransformUtils
   ${STABLEHLO_LIBS}
 )
+endif ()
diff --git a/lib/Conversion/StableHLOToTTIR/StableHLOToTTIR.cpp b/lib/Conversion/StableHLOToTTIR/StableHLOToTTIR.cpp
@@ -19,9 +19,8 @@
 #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;
 
@@ -102,3 +101,4 @@ std::unique_ptr<OperationPass<ModuleOp>> createConvertStableHLOToTTIRPass() {
 }
 
 } // namespace mlir::tt
+#endif
diff --git a/lib/RegisterAll.cpp b/lib/RegisterAll.cpp
@@ -17,9 +17,9 @@
 #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,7 +30,9 @@ 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() {

diff --git a/lib/SharedLib/CMakeLists.txt b/lib/SharedLib/CMakeLists.txt
@@ -18,7 +18,9 @@ set(TTMLIR_LIBS
     MLIRTTIRAnalysis
     MLIRTTNNPipelines
     TTMLIRTTNNToEmitC
+if (TTMLIR_ENABLE_STABLEHLO)
     TTMLIRStableHLOToTTIR
+endif ()
 )
 
 # We supply empty.cpp because CMake does not allow creating a library without sources.

diff --git a/tools/ttmlir-opt/CMakeLists.txt b/tools/ttmlir-opt/CMakeLists.txt
@@ -1,6 +1,3 @@
-
-include(CMakePrintHelpers)
-
 set(stablehlo_libs 
 CheckOps
 ChloCAPI
@@ -14,6 +11,21 @@ StablehloLinalgTransforms
 StablehloOps
 StablehloPasses
 StablehloPortableApi
+StablehloReferenceApi
+StablehloReferenceAxes
+StablehloReferenceConfiguration
+StablehloReferenceElement
+StablehloReferenceErrors
+StablehloReferenceIndex
+StablehloReferenceNumPy
+StablehloReferenceOps
+StablehloReferenceProcess
+StablehloReferenceProcessGrid
+StablehloReferenceScope
+StablehloReferenceTensor
+StablehloReferenceToken
+StablehloReferenceTypes
+StablehloReferenceValue
 StablehloRegister
 StablehloSerialization
 StablehloTOSATransforms
@@ -26,13 +38,36 @@ VhloTypes)
 
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
-set(LIBS ${stablehlo_libs} ${dialect_libs} ${conversion_libs} MLIROptLib MLIRTargetCpp TTMLIRStatic)
+set(LIBS TTMLIRStableHLOToTTIR TTMLIRCAPI ${stablehlo_libs} ${dialect_libs} ${conversion_libs} MLIROptLib MLIRTargetCpp 
+MLIRTTIRAnalysis
+TTMLIRTTIRToTTNN
+MLIRTTNNTransforms
+TTMLIRTosaToTTIR
+TTMLIRTTNNToEmitC
+MLIRTTMetalTransforms
+MLIRTTIRTransforms
+MLIRTTMetalDialect
+MLIRTTKernelDialect
+MLIRTTIRDialect
+MLIRTTNNDialect
+MLIRTTNNPipelines
+MLIRTTDialect
+TTNNTargetFlatbuffer
+TTMLIRCAPI
+TTMLIRStableHLOToTTIR
+TTMLIRStatic)
 
 add_llvm_executable(ttmlir-opt ttmlir-opt.cpp)
 add_dependencies(ttmlir-opt stablehlo)
 
 llvm_update_compile_flags(ttmlir-opt)
-target_link_directories(ttmlir-opt PRIVATE ${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo-build/lib)
+target_link_directories(ttmlir-opt PUBLIC ${TTMLIR_TOOLCHAIN_DIR}/lib)
+target_link_directories(ttmlir-opt PUBLIC ${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo-build/lib)
+
+ADD_CUSTOM_COMMAND(TARGET ttmlir-opt
+	PRE_BUILD
+        COMMAND cp ${PROJECT_SOURCE_DIR}/third_party/stablehlo/src/stablehlo-build/lib/*.a ${TTMLIR_TOOLCHAIN_DIR}/lib
+	DEPENDS stablehlo)
 
 target_link_libraries(ttmlir-opt PRIVATE ${LIBS})
 

diff --git a/tools/ttmlir-opt/ttmlir-opt.cpp b/tools/ttmlir-opt/ttmlir-opt.cpp
@@ -7,8 +7,9 @@
 #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"
 
@@ -18,8 +19,9 @@ 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));
 }