diff --git a/forge/csrc/passes/lower_to_mlir.cpp b/forge/csrc/passes/lower_to_mlir.cpp index 311388b0..6d391458 100644 --- a/forge/csrc/passes/lower_to_mlir.cpp +++ b/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 "{name}.mlir" + reportify::dump_mlir("ttir", &graphModule_); + #ifdef DEBUG // Create a string to store the output std::string moduleStr; diff --git a/forge/csrc/passes/mlir_compiler.cpp b/forge/csrc/passes/mlir_compiler.cpp index 1495a062..e120496d 100644 --- a/forge/csrc/passes/mlir_compiler.cpp +++ b/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,10 @@ 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); + // Generate binary from the MLIR module. auto binary = mlir::tt::ttnn::ttnnToFlatbuffer(mlir_module.get()); tt::log_info(LogMLIRCompiler, "Flatbuffer binary generated successfully."); diff --git a/forge/csrc/reportify/paths.cpp b/forge/csrc/reportify/paths.cpp index 64186399..f6786a6d 100644 --- a/forge/csrc/reportify/paths.cpp +++ b/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/"); diff --git a/forge/csrc/reportify/paths.hpp b/forge/csrc/reportify/paths.hpp index ab6353c0..c20a6fa4 100644 --- a/forge/csrc/reportify/paths.hpp +++ b/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(); diff --git a/forge/csrc/reportify/reportify.cpp b/forge/csrc/reportify/reportify.cpp index 741d2b9d..29d8f068 100644 --- a/forge/csrc/reportify/reportify.cpp +++ b/forge/csrc/reportify/reportify.cpp @@ -6,7 +6,6 @@ #include #include #include -#include #include #include @@ -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::functiongetName()->str(); + + std::string outputString; + llvm::raw_string_ostream outStream(outputString); + + // Put data into string + module->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 node_filter) { @@ -416,6 +437,18 @@ JsonNamePairs create_jsons_for_graph( return this_json_name_pairs; } +JsonNamePairs create_jsons_for_mlir(const std::string& name, mlir::ModuleOp* module) +{ + JsonNamePairs this_json_name_pairs; + + json this_json = create_json_for_mlir(module); + std::string this_name = 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 +458,29 @@ void dump_graph( std::string default_dir = get_default_reportify_path(""); dump_graph(default_dir, test_name, graph_prefix, graph, report_path); } + +void dump_mlir(const std::string& name, mlir::ModuleOp* module) +{ + if (env_as("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->getName()->str(), report_path); + + JsonNamePairs json_pairs = create_jsons_for_mlir(name, module); + 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 diff --git a/forge/csrc/reportify/reportify.hpp b/forge/csrc/reportify/reportify.hpp index b53af418..d83791e5 100644 --- a/forge/csrc/reportify/reportify.hpp +++ b/forge/csrc/reportify/reportify.hpp @@ -8,6 +8,10 @@ #include "nlohmann/json_fwd.hpp" #include "reportify/paths.hpp" +namespace mlir +{ +class ModuleOp; +} // namespace mlir namespace tt { @@ -33,6 +37,8 @@ json create_json_for_graph( const graphlib::Graph* graph, std::function node_filter = [](graphlib::Node*) { return true; }); +json create_json_for_mlir(mlir::ModuleOp* module); + 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& name, mlir::ModuleOp* module); + } // namespace reportify } // namespace tt