[api] add API for exporting model as c++ source

forge/csrc/forge_bindings.cpp

@@ -186,6 +186,7 @@ PYBIND11_MODULE(_C, m)
         py::arg("graph"),
         py::arg("default_df_override") = std::optional<DataFormat>{});
     m.def("run_mlir_compiler", &passes::run_mlir_compiler);
+    m.def("run_mlir_compiler_to_cpp", &passes::run_mlir_compiler_to_cpp);
     m.def("split_graph", &passes::split_graph);
 
     m.def(

forge/csrc/passes/mlir_compiler.cpp

@@ -27,20 +27,24 @@
 #pragma clang diagnostic pop
 
 // TTMLIR headers
+#include "mlir/Target/Cpp/CppEmitter.h"
 #include "tt/runtime/types.h"
 #include "tt_torch_device/tt_device.hpp"
 #include "ttmlir/Dialect/TT/IR/TT.h"
 #include "ttmlir/Dialect/TTIR/IR/TTIR.h"
 #include "ttmlir/Dialect/TTNN/IR/TTNN.h"
+#include "ttmlir/Dialect/TTNN/Transforms/TTNNToCpp.h"
 #include "ttmlir/Target/TTNN/TTNNToFlatbuffer.h"
 
 // Reportify headers
 #include "reportify/reportify.hpp"
 
 namespace tt::passes
 {
-/// Public API for lowering to MLIR, running MLIR passes and generate runtime binary.
-runtime::Binary run_mlir_compiler(tt::ForgeGraphModule& module)
+
+// Template function to run the MLIR compiler pipeline, depending on the desired output.
+template <MLIROutputKind output>
+auto run_mlir_compiler_generic(tt::ForgeGraphModule& module)
 {
     // Register all the required dialects.
     mlir::DialectRegistry registry;
@@ -70,24 +74,55 @@ runtime::Binary run_mlir_compiler(tt::ForgeGraphModule& module)
     // Generate MLIR from the Forge graph.
     mlir::OwningOpRef<mlir::ModuleOp> mlir_module = lower_to_mlir(module, context);
 
-    // Run MLIR registered passes.
-    run_mlir_passes(mlir_module);
+    // Run MLIR pipeline.
+    run_mlir_passes<output>(mlir_module);
+
     tt::log_info(LogMLIRCompiler, "MLIR passes run successfully.");
 
     mlir_module->dump();
 
-    // save what's dumped to a file named "{name}.mlir"
-    reportify::dump_mlir("ttnn", mlir_module->getName()->str(), mlir_module.get());
+    if constexpr (output == MLIROutputKind::Flatbuffer)
+    {
+        // Save generated ttnn module to a file named "{name}.mlir".
+        reportify::dump_mlir("ttnn", mlir_module->getName()->str(), mlir_module.get());
+
+        // Generate binary from the MLIR module.
+        auto binary = mlir::tt::ttnn::ttnnToFlatbuffer(mlir_module.get());
+        tt::log_info(LogMLIRCompiler, "Flatbuffer binary generated successfully.");
 
-    // Generate binary from the MLIR module.
-    auto binary = mlir::tt::ttnn::ttnnToFlatbuffer(mlir_module.get());
-    tt::log_info(LogMLIRCompiler, "Flatbuffer binary generated successfully.");
+        if (binary == nullptr)
+        {
+            throw std::runtime_error("Failed to generate flatbuffer binary.");
+        }
 
-    if (binary == nullptr)
+        return binary;
+    }
+    else if constexpr (output == MLIROutputKind::Cpp)
     {
-        throw std::runtime_error("Failed to generate flatbuffer binary.");
+        std::string cpp_source;
+        llvm::raw_string_ostream rso(cpp_source);
+
+        log_info(LogMLIRCompiler, "Generating C++ code from MLIR module.");
+        auto res = mlir::emitc::translateToCpp(mlir_module.get(), rso);
+        if (mlir::failed(res))
+        {
+            throw std::runtime_error("Failed to generate C++ code.");
+        }
+
+        rso.flush();
+
+        tt::log_info(LogMLIRCompiler, "C++ code generated successfully.");
+        return cpp_source;
     }
+}
 
-    return binary;
+runtime::Binary run_mlir_compiler(tt::ForgeGraphModule& module)
+{
+    return run_mlir_compiler_generic<MLIROutputKind::Flatbuffer>(module);
+}
+
+std::string run_mlir_compiler_to_cpp(tt::ForgeGraphModule& module)
+{
+    return run_mlir_compiler_generic<MLIROutputKind::Cpp>(module);
 }
 }  // namespace tt::passes

forge/csrc/passes/mlir_compiler.hpp

@@ -14,4 +14,7 @@ namespace tt::passes
 {
 /// Public API for running MLIR passes and generating binary.
 runtime::Binary run_mlir_compiler(tt::ForgeGraphModule& module);
+
+/// Public API for lowering to MLIR, running MLIR passes and generate C++ code.
+std::string run_mlir_compiler_to_cpp(tt::ForgeGraphModule& module);
 }  // namespace tt::passes

forge/csrc/passes/mlir_passes.cpp

@@ -17,9 +17,10 @@
 
 namespace tt::passes
 {
-/// Public API for running MLIR passes and generating binary.
-void run_mlir_passes(mlir::OwningOpRef<mlir::ModuleOp> &mlir_module)
+
+void register_mlir_passes()
 {
+    // Static (only once) initialization of the MLIR passes.
     static bool _ = []()
     {
         // Register required passes
@@ -34,13 +35,26 @@ void run_mlir_passes(mlir::OwningOpRef<mlir::ModuleOp> &mlir_module)
         return true;
     }();
     (void)_;
+}
+
+template <MLIROutputKind output>
+void run_mlir_passes(mlir::OwningOpRef<mlir::ModuleOp> &mlir_module)
+{
+    // Register the MLIR passes.
+    register_mlir_passes();
 
     // Create a pass manager.
     mlir::PassManager pm(mlir_module.get()->getName());
 
     // Get the pipeline info for the wanted pipeline.
-    const auto pipelineInfo = mlir::PassPipelineInfo::lookup("ttir-to-ttnn-backend-pipeline");
-
+    static_assert(
+        output == MLIROutputKind::Flatbuffer || output == MLIROutputKind::Cpp,
+        "Handling only Flatbuffer and Cpp output correctly.");
+    constexpr auto pipeline_name =
+        (output == MLIROutputKind::Flatbuffer) ? "ttir-to-ttnn-backend-pipeline" : "ttir-to-emitc-pipeline";
+    const auto pipelineInfo = mlir::PassPipelineInfo::lookup(pipeline_name);
+
+    // Error handler for the pipeline. Will be called if there is an error during parsing of the pipeline options.
     auto err_handler = [](const mlir::Twine &location)
     {
         log_error(LogMLIRCompiler, "Error during parsing pipeline options: {}", location.str());
@@ -77,4 +91,9 @@ void run_mlir_passes(mlir::OwningOpRef<mlir::ModuleOp> &mlir_module)
     log_trace(LogMLIRCompiler, "MLIR module after running passes:\n{}", moduleStr);
 #endif
 }
+
+// Explicit templates instantiation.
+template void run_mlir_passes<MLIROutputKind::Flatbuffer>(mlir::OwningOpRef<mlir::ModuleOp> &mlir_module);
+template void run_mlir_passes<MLIROutputKind::Cpp>(mlir::OwningOpRef<mlir::ModuleOp> &mlir_module);
+
 }  // namespace tt::passes

forge/csrc/passes/mlir_passes.hpp

@@ -10,6 +10,15 @@ class OwningOpRef;
 
 namespace tt::passes
 {
-/// Public API for running MLIR passes and generating binary.
+
+enum class MLIROutputKind
+{
+    Flatbuffer,
+    Cpp
+};
+
+/// Public API for running MLIR passes (pipeline) depending on the desired output.
+template <MLIROutputKind output>
 void run_mlir_passes(mlir::OwningOpRef<mlir::ModuleOp> &mlir_module);
+
 }  // namespace tt::passes

forge/forge/compile.py

@@ -393,6 +393,7 @@ def forge_compile_from_context(context: CompileContext) -> CompiledModel:
     assert context.compiled_binary is not None
 
     compiled_module = CompiledModel(
+        context.forge_module,
         fwd_compiled_graph_state,
         bwd_compiled_graph_state,
         opt_compiled_graph_state,

forge/forge/compiled_graph_state.py

@@ -10,9 +10,11 @@
 from typing import Dict, List, Any, Optional
 
 
+from forge._C import ForgeGraphModule
 from forge._C.graph import Graph
 import forge._C.graph as pygraph
 from forge._C.runtime import run_binary, Binary
+from forge._C import run_mlir_compiler_to_cpp
 from forge.tensor import Tensor, get_post_const_eval_tensors, to_pt_tensors, cast_unsupported_torch_dtype, AnyTensor
 from forge.module import Module, PyTorchModule, AnyModule
 from forge.execution_tracker import ExecutionPhase, record_execution_phase_and_stage
@@ -159,6 +161,7 @@ class CompiledModel:
 
     fwd_compiled_graph_state: CompiledGraphState
     bwd_compiled_graph_state: Optional[CompiledGraphState]
+    opt_compiled_graph_state: Optional[CompiledGraphState]
 
     # Compiled flatbuffer binary composed of programs which execute compiled graphs (e.g., forward, backward, etc.)
     compiled_binary: Binary
@@ -170,6 +173,11 @@ class CompiledModel:
     # Original user-defined module.
     framework_module: AnyModule
 
+    # Forge graph module, currently used for exporting the model to a cpp file.
+    # Needed by the lower to MLIR logic.
+    # Issue(#1350): current state of `CompiledModel` is a bit messy, we should clean it up.
+    forge_graph_module: ForgeGraphModule
+
     # Gradients to be passed into the backward pass.
     # Used when CompiledModel.backward() is part of a chain of backward passes.
     gradient_inputs: List[Optional[torch.Tensor]]
@@ -179,13 +187,15 @@ class CompiledModel:
 
     def __init__(
         self,
+        forge_graph_module: ForgeGraphModule,
         fwd_compiled_graph_state: CompiledGraphState,
         bwd_compiled_graph_state: Optional[CompiledGraphState],
         opt_compiled_graph_state: Optional[CompiledGraphState],
         compiled_binary: Binary,
         framework_module: AnyModule,
         attached_module: Optional["CompiledModel"] = None,
     ):
+        self.forge_graph_module = forge_graph_module
         self.fwd_compiled_graph_state = fwd_compiled_graph_state
         self.bwd_compiled_graph_state = bwd_compiled_graph_state
         self.opt_compiled_graph_state = opt_compiled_graph_state
@@ -392,3 +402,28 @@ def step(self) -> None:
                     ) is self.fwd_compiled_graph_state.get_parameter_tensor(weight_name)
                     assert self.fwd_compiled_graph_state.get_parameter_tensor(weight_name) is val
         self.gradient_outputs = []
+
+    def export_to_cpp(self, export_path: str) -> None:
+        """
+        Export the model to a cpp file.
+
+        Parameters
+        ----------
+        export_path: str
+            Path to the file where the model c++ code will be exported.
+        """
+
+        logger.info(f"Exporting model {self.framework_module.get_name()} to cpp file...")
+        cpp_code = run_mlir_compiler_to_cpp(self.forge_graph_module)
+
+        with open(export_path, "w") as f:
+            f.write(cpp_code)
+
+        logger.info(f'Exported model as cpp file to "{export_path}"')
+        logger.info(
+            f"To compile and run this code, one can utilize the ttnn-standalone tool within the tt-mlir project. \
+            It provides all the necessary build and run scripts. Copy the contents of the .cpp to ttnn-standalone.cpp \
+            and use `./run` to compile&run the code."
+        )
+        logger.info(f"    Tool: https://github.com/tenstorrent/tt-mlir/tree/main/tools/ttnn-standalone")
+        logger.info(f"    Docs: https://docs.tenstorrent.com/tt-mlir/ttnn-standalone.html")

forge/test/test_api.py

@@ -88,3 +88,27 @@ def forward(self, x, y):
 
     if not torch.allclose(output[0], golden.to_pytorch(), rtol=1e-1):
         raise ValueError("Output does not match the golden output")
+
+
+def test_export_to_cpp():
+    class Add(nn.Module):
+        def __init__(self):
+            super().__init__()
+
+        def forward(self, x1, x2):
+            return torch.add(x1, x2)
+
+    model = Add()
+    shape = (1, 1024, 32)
+    inputs = [torch.rand(shape), torch.rand(shape)]
+
+    compiled_model = forge.compile(model, sample_inputs=[torch.rand(shape), torch.rand(shape)])
+
+    file_path = "generated_export_add.cpp"
+    compiled_model.export_to_cpp(file_path)
+
+    assert os.path.exists(file_path)
+    with open(file_path, "r") as f:
+        print(f.read())
+
+    os.remove(file_path)