Show mlir program when tracing benchmarking

src/targets/gpu/compile_ops.cpp

@@ -80,6 +80,12 @@ struct compiled_result
 {
     compiler_replace replace;
     instruction_ref ins;
+
+    friend std::ostream& operator<<(std::ostream& os, const compiled_result& cr)
+    {
+        cr.replace.trace(os, cr.ins);
+        return os;
+    }
 };
 
 struct compile_plan
@@ -179,6 +185,8 @@ struct compile_plan
                                    std::cout << "No binary" << std::endl;
                                return std::numeric_limits<double>::max();
                            }
+                           if(trace_level > 2)
+                               std::cout << *cr << std::endl;
                            auto t = time_op(
                                *ctx, cr->replace.code_object, to_shapes(cr->ins->inputs()), 20);
                            if(trace_level > 1)

src/targets/gpu/include/migraphx/gpu/compiler.hpp

@@ -48,17 +48,28 @@ struct compiler_replace
     compiler_replace(const operation& op) : code_object{op} {}
 
     template <class F>
-    compiler_replace(const operation& op, F f)
-        : code_object{op},
-          replace_fn([=](const compiler_replace& cr, module& m, instruction_ref ins) {
-              f(m, ins, cr.code_object);
-          })
+    compiler_replace(const operation& op, F f) : code_object{op}, replace_fn(make_replace(f))
+    {
+    }
+
+    template <class F, class Trace>
+    compiler_replace(const operation& op, F f, Trace t)
+        : code_object{op}, replace_fn(make_replace(f)), trace_fn(t)
     {
     }
 
     operation code_object = {};
     std::function<void(const compiler_replace& cr, module& m, instruction_ref ins)> replace_fn =
         nullptr;
+    std::function<void(std::ostream& os, instruction_ref ins)> trace_fn = nullptr;
+
+    template <class F>
+    static auto make_replace(F f)
+    {
+        return [=](const compiler_replace& cr, module& m, instruction_ref ins) {
+            f(m, ins, cr.code_object);
+        };
+    }
 
     void replace(module& m, instruction_ref ins) const
     {
@@ -67,6 +78,12 @@ struct compiler_replace
         else
             m.replace_instruction(ins, code_object, ins->inputs());
     }
+
+    void trace(std::ostream& os, instruction_ref ins) const
+    {
+        if(trace_fn)
+            trace_fn(os, ins);
+    }
 };
 
 using compiler_compile =

src/targets/gpu/include/migraphx/gpu/mlir.hpp

@@ -37,6 +37,7 @@ struct module;
 namespace gpu {
 
 MIGRAPHX_GPU_EXPORT std::string dump_mlir(const module& m);
+MIGRAPHX_GPU_EXPORT std::string dump_mlir(const module& m, const std::vector<shape>& inputs);
 MIGRAPHX_GPU_EXPORT code_object_op compile_mlir(const context& migraphx_ctx,
                                                 module m,
                                                 const std::vector<instruction_ref>& inputs,

src/targets/gpu/jit/mlir.cpp

@@ -46,10 +46,12 @@ struct mlir_compiler : compiler<mlir_compiler>
 
     compiler_replace insert(code_object_op co) const
     {
-        return {std::move(co), [](module& m, instruction_ref ins, const operation& op) {
+        return {std::move(co),
+                [](module& m, instruction_ref ins, const operation& op) {
                     auto mlir = insert_mlir(m, ins, any_cast<code_object_op>(op), ins->inputs());
                     m.replace_instruction(ins, mlir);
-                }};
+                },
+                &trace};
     }
 
     optional<tuning_config> get_tuning_config(const context& ctx,
@@ -61,6 +63,13 @@ struct mlir_compiler : compiler<mlir_compiler>
         auto* smod  = ins->module_inputs().front();
         return get_tuning_config_mlir(ctx, *smod, shapes, exhaustive);
     }
+
+    static void trace(std::ostream& os, instruction_ref ins)
+    {
+        auto shapes = to_shapes(ins->inputs());
+        auto* smod  = ins->module_inputs().front();
+        os << dump_mlir(*smod, shapes);
+    }
 };
 
 } // namespace gpu

src/targets/gpu/mlir.cpp

@@ -934,14 +934,6 @@ struct mlir_program
     std::string sym_name;
 };
 
-std::string dump_mlir(const module& m)
-{
-    mlir_program mp;
-    mp.parse(m);
-    auto mod_op = mlirModuleGetOperation(mp.mmodule.get());
-    return mlir_print(&mlirOperationPrint, mod_op);
-}
-
 void adjust_param_shapes(module& m, const std::vector<shape>& inputs)
 {
     auto names = m.get_parameter_names();
@@ -960,6 +952,24 @@ void adjust_param_shapes(module& m, const std::vector<shape>& inputs)
     }
 }
 
+std::string dump_mlir(const module& m, const std::vector<shape>& inputs)
+{
+    module mm;
+    const_module_ref mr = &m;
+    if(not inputs.empty())
+    {
+        mm = m;
+        mr = &mm;
+        adjust_param_shapes(mm, inputs);
+    }
+    mlir_program mp;
+    mp.parse(*mr);
+    auto mod_op = mlirModuleGetOperation(mp.mmodule.get());
+    return mlir_print(&mlirOperationPrint, mod_op);
+}
+
+std::string dump_mlir(const module& m) { return dump_mlir(m, {}); }
+
 code_object_op compile_mlir(const context& migraphx_ctx,
                             module m,
                             const std::vector<instruction_ref>& inputs,
@@ -1030,13 +1040,20 @@ tuning_config get_tuning_config_mlir(const context& migraphx_ctx,
 
 #else
 
-std::string dump_mlir(const module&) { return {}; }
-
 template <class T>
 void use(T&)
 {
 }
 
+std::string dump_mlir(const module&) { return {}; }
+
+std::string dump_mlir(const module& m, const std::vector<shape>& inputs)
+{
+    use(m);
+    use(inputs);
+    return {};
+}
+
 // Disabling clang-tidy warning on non-real useage.
 // NOLINTBEGIN(performance-unnecessary-value-param)
 code_object_op