refactored code for pr and replaced mlir::ModuleOp with mlir::Operation

forge/csrc/passes/lower_to_mlir.cpp

@@ -105,8 +105,8 @@ class MLIRGenerator
         log_info(LogMLIRCompiler, "MLIR module generated successfully.");
         graphModule_.dump();
 
-        // save what's dumped to a file named "{name}.mlir"
-        reportify::dump_mlir("ttir", &graphModule_);
+        // 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

forge/csrc/passes/mlir_compiler.cpp

@@ -77,8 +77,7 @@ runtime::Binary run_mlir_compiler(tt::ForgeGraphModule& module)
     mlir_module->dump();
 
     // save what's dumped to a file named "{name}.mlir"
-    auto mlir_modulee = mlir_module.get();
-    reportify::dump_mlir("ttnn", &mlir_modulee);
+    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());

forge/csrc/reportify/reportify.cpp

@@ -407,17 +407,17 @@ json create_json_for_graph(const graphlib::Graph* graph, std::function<bool(grap
     return this_json;
 }
 
-json create_json_for_mlir(mlir::ModuleOp* module)
+json create_json_for_mlir(const std::string& module_name, mlir::Operation* operation)
 {
     json this_json;
 
-    this_json["module"] = module->getName()->str();
+    this_json["module"] = module_name;
 
     std::string outputString;
     llvm::raw_string_ostream outStream(outputString);
 
     // Put data into string
-    module->print(outStream);
+    operation->print(outStream);
     outStream.flush();
     this_json["content"] = outputString;
 
@@ -437,12 +437,12 @@ JsonNamePairs create_jsons_for_graph(
     return this_json_name_pairs;
 }
 
-JsonNamePairs create_jsons_for_mlir(const std::string& name, mlir::ModuleOp* module)
+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);
-    std::string this_name = name + ".mlir";
+    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);
 
@@ -459,26 +459,33 @@ void dump_graph(
     dump_graph(default_dir, test_name, graph_prefix, graph, report_path);
 }
 
-void dump_mlir(const std::string& name, mlir::ModuleOp* module)
+/**
+ * @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);
 
-    std::string full_report_path = build_report_path(path, module->getName()->str(), report_path);
-
-    JsonNamePairs json_pairs = create_jsons_for_mlir(name, module);
+    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);
 }
 

forge/csrc/reportify/reportify.hpp

@@ -10,7 +10,7 @@
 
 namespace mlir
 {
-class ModuleOp;
+class Operation;
 }  // namespace mlir
 namespace tt
 {
@@ -37,7 +37,7 @@ 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(mlir::ModuleOp* module);
+json create_json_for_mlir(const std::string& module_name, mlir::Operation* operation);
 
 void dump_graph(
     const std::string& test_name,
@@ -59,7 +59,7 @@ void dump_epoch_id_graphs(
     const graphlib::Graph* graph,
     const std::string& directory_path = get_default_reportify_path(""));
 
-void dump_mlir(const std::string& name, mlir::ModuleOp* module);
+void dump_mlir(const std::string& file_name, const std::string& module_name, mlir::Operation* operation);
 
 }  // namespace reportify