Dumping ttir and ttnn into json files (#647)

adding functionality to reportify for formatting and storing MLIR input (ttir) and MLIR output (ttnn) in a JSON file. With this update, running any module will save MLIR information within the path $REPORTIFY_PATH$/$OPERATION_NAME$/mlir_reports/ttir.mlir as ttir.mlir and ttnn.mlir.

forge/csrc/passes/lower_to_mlir.cpp

@@ -40,6 +40,9 @@
 #include "ttmlir/Dialect/TT/IR/TTOpsTypes.h"
 #include "ttmlir/Dialect/TTIR/IR/TTIROps.h"
 
+// Reportify headers
+#include "reportify/reportify.hpp"
+
 namespace
 {
 using namespace tt;
@@ -102,6 +105,9 @@ class MLIRGenerator
         log_info(LogMLIRCompiler, "MLIR module generated successfully.");
         graphModule_.dump();
 
+        // save what's dumped to a file named "{file_name}.mlir"
+        reportify::dump_mlir("ttir", graphModule_.getNameAttr().getValue().str(), graphModule_.getOperation());
+
 #ifdef DEBUG
         // Create a string to store the output
         std::string moduleStr;

forge/csrc/passes/mlir_compiler.cpp

@@ -34,6 +34,9 @@
 #include "ttmlir/Dialect/TTNN/IR/TTNN.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.
@@ -73,6 +76,9 @@ runtime::Binary run_mlir_compiler(tt::ForgeGraphModule& module)
 
     mlir_module->dump();
 
+    // save what's dumped 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.");

forge/csrc/reportify/paths.cpp

@@ -53,6 +53,13 @@ std::string get_pass_reports_relative_directory()
     return retstring;
 }
 
+std::string get_mlir_reports_relative_directory()
+{
+    std::string retstring("mlir_reports");
+
+    return retstring;
+}
+
 std::string get_router_report_relative_directory()
 {
     std::string retstring("/router_reports/EpochType/");

forge/csrc/reportify/paths.hpp

@@ -13,6 +13,7 @@ std::string build_report_path(std::string const& base_path, std::string const& t
 bool initalize_reportify_directory(const std::string& reportify_dir, const std::string& test_name);
 std::string get_default_reportify_path(const std::string& test_name);
 std::string get_pass_reports_relative_directory();
+std::string get_mlir_reports_relative_directory();
 std::string get_router_report_relative_directory();
 std::string get_memory_report_relative_directory();
 std::string get_epoch_type_report_relative_directory();

forge/csrc/reportify/reportify.cpp

@@ -6,7 +6,6 @@
 #include <filesystem>
 #include <fstream>
 #include <iomanip>
-#include <ostream>
 #include <sstream>
 #include <string>
 
@@ -17,8 +16,13 @@
 #include "reportify/to_json.hpp"
 #include "utils/logger.hpp"
 
+// MLIR headers
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wunused-local-typedef"
+#include "mlir/IR/BuiltinOps.h"
+#pragma clang diagnostic pop
+
 using json = nlohmann::json;
-using tt::LogReportify;
 
 namespace tt
 {
@@ -403,6 +407,23 @@ json create_json_for_graph(const graphlib::Graph* graph, std::function<bool(grap
     return this_json;
 }
 
+json create_json_for_mlir(const std::string& module_name, mlir::Operation* operation)
+{
+    json this_json;
+
+    this_json["module"] = module_name;
+
+    std::string outputString;
+    llvm::raw_string_ostream outStream(outputString);
+
+    // Put data into string
+    operation->print(outStream);
+    outStream.flush();
+    this_json["content"] = outputString;
+
+    return this_json;
+}
+
 JsonNamePairs create_jsons_for_graph(
     const std::string& graph_prefix, const graphlib::Graph* graph, std::function<bool(graphlib::Node*)> node_filter)
 {
@@ -416,6 +437,19 @@ JsonNamePairs create_jsons_for_graph(
     return this_json_name_pairs;
 }
 
+JsonNamePairs create_jsons_for_mlir(
+    const std::string& file_name, const std::string& module_name, mlir::Operation* operation)
+{
+    JsonNamePairs this_json_name_pairs;
+
+    json this_json = create_json_for_mlir(module_name, operation);
+    std::string this_name = file_name + ".mlir";
+    JsonNamePair this_json_name_pair = std::make_pair(this_json, this_name);
+    this_json_name_pairs.push_back(this_json_name_pair);
+
+    return this_json_name_pairs;
+}
+
 void dump_graph(
     const std::string& test_name,
     const std::string& graph_prefix,
@@ -425,5 +459,37 @@ void dump_graph(
     std::string default_dir = get_default_reportify_path("");
     dump_graph(default_dir, test_name, graph_prefix, graph, report_path);
 }
+
+/**
+ * @brief Dumps the MLIR's Operation to a JSON file.
+ *
+ * This function generates a JSON representation of the given MLIR operation and writes it to a file.
+ * The file path is following: `$REPORTIFY_PATH$/$OPERATION_NAME$/mlir_reports/$FILE_NAME$.mlir`.
+ * If the environment variable "FORGE_DISABLE_REPORTIFY_DUMP" is set to true, the function returns without performing
+ * any action.
+ *
+ * @param name The name of the file to be saved (currently in use are 'ttir' and 'ttnn').
+ * @param operation A pointer to the MLIR operation to be dumped.
+ */
+void dump_mlir(const std::string& file_name, const std::string& module_name, mlir::Operation* operation)
+{
+    if (env_as<bool>("FORGE_DISABLE_REPORTIFY_DUMP"))
+        return;
+
+    std::string path = get_default_reportify_path("");
+    std::string report_path = get_mlir_reports_relative_directory();
+    std::string full_report_path = build_report_path(path, module_name, report_path);
+
+    JsonNamePairs json_pairs = create_jsons_for_mlir(file_name, module_name, operation);
+    json root_json = json_pairs.back().first;
+
+    std::string root_json_name = json_pairs.back().second;
+    std::transform(root_json_name.begin(), root_json_name.end(), root_json_name.begin(), ::tolower);
+    std::string root_json_path = full_report_path + root_json_name;
+
+    std::filesystem::create_directories(std::filesystem::path(full_report_path));
+    write_json_to_file(root_json_path, root_json);
+}
+
 }  // namespace reportify
 }  // namespace tt

forge/csrc/reportify/reportify.hpp

@@ -8,6 +8,10 @@
 #include "nlohmann/json_fwd.hpp"
 #include "reportify/paths.hpp"
 
+namespace mlir
+{
+class Operation;
+}  // namespace mlir
 namespace tt
 {
 
@@ -33,6 +37,8 @@ json create_json_for_graph(
     const graphlib::Graph* graph,
     std::function<bool(graphlib::Node*)> node_filter = [](graphlib::Node*) { return true; });
 
+json create_json_for_mlir(const std::string& module_name, mlir::Operation* operation);
+
 void dump_graph(
     const std::string& test_name,
     const std::string& graph_prefix,
@@ -52,6 +58,9 @@ void dump_epoch_id_graphs(
     const std::string& graph_prefix,
     const graphlib::Graph* graph,
     const std::string& directory_path = get_default_reportify_path(""));
+
+void dump_mlir(const std::string& file_name, const std::string& module_name, mlir::Operation* operation);
+
 }  // namespace reportify
 
 }  // namespace tt