Workarounding flatbuffer generation API (#121)

tenstorrent · Jul 12, 2024 · dbe0aec · dbe0aec
1 parent d8b6ced
commit dbe0aec
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 1 deletion.
diff --git a/include/ttmlir/Dialect/TTNN/Transforms/TTNNToSerializedBinary.h b/include/ttmlir/Dialect/TTNN/Transforms/TTNNToSerializedBinary.h
@@ -0,0 +1,15 @@
+// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef TTMLIR_DIALECT_TTNN_TRANSFORMS_TTNNTOSERIALIZEDBINARY_H
+#define TTMLIR_DIALECT_TTNN_TRANSFORMS_TTNNTOSERIALIZEDBINARY_H
+
+#include <memory>
+
+#include "mlir/IR/BuiltinOps.h"
+
+namespace mlir::tt::ttnn {
+std::shared_ptr<void> emitTTNNAsFlatbuffer(OwningOpRef<ModuleOp> &moduleOp);
+} // namespace mlir::tt::ttnn
+#endif
diff --git a/lib/Dialect/TTNN/Transforms/CMakeLists.txt b/lib/Dialect/TTNN/Transforms/CMakeLists.txt
@@ -1,7 +1,7 @@
 add_mlir_dialect_library(MLIRTTNNTransforms
         Passes.cpp
         TTNNToCpp.cpp
-        SerializeToBinary.cpp
+        TTNNToSerializedBinary.cpp
 
         ADDITIONAL_HEADER_DIRS
         ${PROJECT_SOURCE_DIR}/include/ttmlir

diff --git a/...ect/TTNN/Transforms/SerializeToBinary.cpp → ...TNN/Transforms/TTNNToSerializedBinary.cpp b/...ect/TTNN/Transforms/SerializeToBinary.cpp → ...TNN/Transforms/TTNNToSerializedBinary.cpp
@@ -17,6 +17,7 @@
 #include "ttmlir/Dialect/TTNN/IR/TTNNOpsTypes.h"
 #include "ttmlir/Dialect/TTNN/Passes.h"
 #include "ttmlir/Dialect/TTNN/Transforms/TTNNToCpp.h"
+#include "ttmlir/Dialect/TTNN/Transforms/TTNNToSerializedBinary.h"
 #include "ttmlir/Target/TTNN/Target.h"
 #include "ttmlir/Target/Utils/FlatbufferObjectCache.h"
 #include "ttmlir/Target/Utils/FuncOpToProgram.h"
@@ -263,4 +264,23 @@ class TTNNSerializeToBinary
   std::shared_ptr<void> serializedBinary;
 };
 
+std::shared_ptr<void> emitTTNNAsFlatbuffer(OwningOpRef<ModuleOp> &moduleOp) {
+  auto pm = PassManager::on<ModuleOp>(moduleOp.get().getContext());
+  auto pass = createTTNNSerializeToBinary();
+  Pass *basePass = pass.get();
+  pm.addPass(std::move(pass));
+
+  // Run the pass manager.
+  if (failed(pm.run(moduleOp.get()))) {
+    return nullptr;
+  }
+
+  auto *derivedPass = llvm::dyn_cast<TTNNSerializeToBinary>(basePass);
+  if (!derivedPass) {
+    return nullptr;
+  }
+
+  return derivedPass->getBinary();
+}
+
 } // namespace mlir::tt::ttnn