From bb33f43b5cd2d6afe615f62a5d1f3e602795fb74 Mon Sep 17 00:00:00 2001
From: richagadgil <richa.gadgil@amd.com>
Date: Wed, 23 Oct 2024 21:15:03 +0000
Subject: [PATCH] first pass

---
 src/targets/gpu/include/migraphx/gpu/mlir.hpp |  1 +
 src/targets/gpu/jit/mlir.cpp                  |  7 ++++
 src/targets/gpu/mlir.cpp                      | 32 +++++++++++++++++++
 3 files changed, 40 insertions(+)

diff --git a/src/targets/gpu/include/migraphx/gpu/mlir.hpp b/src/targets/gpu/include/migraphx/gpu/mlir.hpp
index 13d8fdfa916..5541ada6560 100644
--- a/src/targets/gpu/include/migraphx/gpu/mlir.hpp
+++ b/src/targets/gpu/include/migraphx/gpu/mlir.hpp
@@ -40,6 +40,7 @@ namespace gpu {
 
 MIGRAPHX_GPU_EXPORT std::string dump_mlir(module m);
 MIGRAPHX_GPU_EXPORT std::string dump_mlir(module m, const std::vector<shape>& inputs);
+MIGRAPHX_GPU_EXPORT std::string dump_mlir(module m, const std::vector<shape>& inputs, const fs::path& location);
 
 MIGRAPHX_GPU_EXPORT bool
 is_module_fusible(const module& m, const context& migraphx_ctx, const value& solution);
diff --git a/src/targets/gpu/jit/mlir.cpp b/src/targets/gpu/jit/mlir.cpp
index 06bb69fbd11..1c4c131c8b3 100644
--- a/src/targets/gpu/jit/mlir.cpp
+++ b/src/targets/gpu/jit/mlir.cpp
@@ -39,6 +39,7 @@ inline namespace MIGRAPHX_INLINE_NS {
 namespace gpu {
 
 MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_MLIR_DUMP_TO_MXR);
+MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_MLIR_DUMP_TO_FILE);
 
 static module create_pointwise_module(module_ref in_mod)
 {
@@ -212,6 +213,7 @@ struct mlir_compiler : compiler<mlir_compiler>
                                               bool exhaustive) const
     {
         static const auto mxr_loc = string_value_of(MIGRAPHX_MLIR_DUMP_TO_MXR{});
+        static const auto mlir_loc = string_value_of(MIGRAPHX_MLIR_DUMP_TO_FILE{});
 
         auto shapes = to_shapes(ins->inputs());
         auto* smod  = ins->module_inputs().front();
@@ -219,11 +221,16 @@ struct mlir_compiler : compiler<mlir_compiler>
         {
             dump_mlir_to_mxr(*smod, ins->inputs(), mxr_loc);
         }
+        if(not mlir_loc.empty())
+        {
+            dump_mlir(*smod, shapes, mlir_loc);
+        }
         return get_tuning_config_mlir(ctx, *smod, shapes, exhaustive);
     }
 
     static void trace(std::ostream& os, instruction_ref ins)
     {
+        std::cout << "tracing" << std::endl;
         auto shapes = to_shapes(ins->inputs());
         auto* smod  = ins->module_inputs().front();
         os << dump_mlir(*smod, shapes);
diff --git a/src/targets/gpu/mlir.cpp b/src/targets/gpu/mlir.cpp
index e9d06ce0a9b..09d40ee4f11 100644
--- a/src/targets/gpu/mlir.cpp
+++ b/src/targets/gpu/mlir.cpp
@@ -1096,6 +1096,37 @@ std::string dump_mlir(module m, const std::vector<shape>& inputs)
     return mlir_print(&mlirOperationPrint, mod_op);
 }
 
+std::string dump_mlir(module m, const std::vector<shape>& inputs, const fs::path& location)
+{
+    auto name =
+        mlir_program::get_symbol_name(m) + ".mlir";
+    replace_string_inplace(name, ", ", "_");
+    replace_string_inplace(name, ":", "s");
+    auto f = location / name;
+
+    const_module_ref mr = &m;
+    if(not inputs.empty())
+    {
+        adjust_param_shapes(m, inputs);
+    }
+    rewrite_reduce(m);
+    mlir_program mp;
+    mp.parse(*mr);
+    auto mod_op = mlirModuleGetOperation(mp.mmodule.get());
+
+    std::string mlir_str = mlir_print(&mlirOperationPrint, mod_op);
+    
+    std::ofstream outfile(f);
+    if (outfile.is_open()) {
+        outfile << mlir_str;
+        outfile.close();
+    } else {
+        throw std::runtime_error("Unable to open file");
+    }
+
+    return mlir_str;
+}
+
 std::string dump_mlir(module m) { return dump_mlir(std::move(m), {}); }
 
 mlir_code_object compile_mlir(const context& migraphx_ctx,
@@ -1122,6 +1153,7 @@ mlir_code_object compile_mlir(const context& migraphx_ctx,
     if(trace)
     {
         const std::lock_guard<std::mutex> lock(mutex);
+        std::cout << "printing mlir" << std::endl;
         std::cout << mlir_print(&mlirOperationPrint, mod_op) << std::endl;
     }
     auto co            = mp.compile(solution);