Correctly add module

src/include/migraphx/matcher.hpp

@@ -326,8 +326,8 @@ match::matcher_result find_match(module& modl, M&& m)
 MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_MATCHES)
 
 /// Find matches for an instruction in the module
-template <class... Ms>
-void find_matches(module& mod, instruction_ref ins, Ms&&... ms)
+template <class Mod, class... Ms>
+void find_matches(Mod& mod, instruction_ref ins, Ms&&... ms)
 {
 #if !defined(__GNUC__) || defined(__clang__) || __GNUC__ > 5
     const
@@ -340,13 +340,13 @@ void find_matches(module& mod, instruction_ref ins, Ms&&... ms)
                 return;
             if(trace > 1)
                 std::cout << "Match: " << get_type_name(m) << std::endl;
-            auto r = match_instruction(mod, ins, m.matcher());
-            if(r.result == mod.end())
+            auto r = match_instruction(get_module(mod), ins, m.matcher());
+            if(r.result == get_module(mod).end())
                 return;
             if(trace > 0)
             {
                 std::cout << "Matched by " << get_type_name(m) << std::endl;
-                mod.debug_print(ins);
+                get_module(mod).debug_print(ins);
             }
             m.apply(mod, r);
             match = true;
@@ -355,10 +355,10 @@ void find_matches(module& mod, instruction_ref ins, Ms&&... ms)
 }
 
 /// Find matches in a module
-template <class... Ms>
-void find_matches(module& mod, Ms&&... ms)
+template <class Mod, class... Ms>
+void find_matches(Mod& mod, Ms&&... ms)
 {
-    for(auto ins : iterator_for(mod))
+    for(auto ins : iterator_for(get_module(mod)))
     {
         find_matches(mod, ins, ms...);
     }

src/include/migraphx/module.hpp

@@ -178,6 +178,11 @@ struct module
     std::unique_ptr<module_impl> impl;
 };
 
+inline module& get_module(module& m)
+{
+    return m;
+}
+
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx
 

src/include/migraphx/pass_manager.hpp

@@ -21,6 +21,8 @@ struct module_pass_manager
     virtual ~module_pass_manager() {}
 };
 
+module& get_module(module_pass_manager& mpm);
+
 void run_passes(module& mod, const std::vector<pass>& passes, tracer trace = tracer{});
 void run_passes(program& prog, const std::vector<pass>& passes, tracer trace = tracer{});
 

src/targets/gpu/fuse_mlir.cpp

@@ -47,7 +47,7 @@ struct find_conv_pointwise
         return match::name("pointwise")(match::any_of[match::inputs()](convolution.bind("x")));
     }
 
-    void apply(module& m, match::matcher_result r) const
+    void apply(module_pass_manager& mpm, match::matcher_result r) const
     {
         auto ins      = r.result;
         auto conv_ins = r.instructions["convolution"];
@@ -61,13 +61,14 @@ struct find_conv_pointwise
            }))
             return;
         std::sort(names.begin(), names.end());
-        module mm{};
+        module_ref mm = mpm.create_module("mlir_" + pm->name());
+        mm->set_bypass();
         std::unordered_map<instruction_ref, instruction_ref> param_map;
-        auto x    = mm.add_parameter("x" + std::to_string(names.size()),
+        auto x    = mm->add_parameter("x" + std::to_string(names.size()),
                                   conv_ins->inputs().at(0)->get_shape());
-        auto w    = mm.add_parameter("x" + std::to_string(names.size() + 1),
+        auto w    = mm->add_parameter("x" + std::to_string(names.size() + 1),
                                   conv_ins->inputs().at(1)->get_shape());
-        auto conv = mm.add_instruction(conv_ins->get_operator(), {x, w});
+        auto conv = mm->add_instruction(conv_ins->get_operator(), {x, w});
         std::transform(names.begin(),
                        names.end(),
                        ins->inputs().begin(),
@@ -76,9 +77,9 @@ struct find_conv_pointwise
                            if(input == x_ins)
                                return std::make_pair(pm->get_parameter(name), conv);
                            return std::make_pair(pm->get_parameter(name),
-                                                 mm.add_parameter(name, input->get_shape()));
+                                                 mm->add_parameter(name, input->get_shape()));
                        });
-        mm.add_return(mm.insert_module_instructions(mm.end(), pm, param_map));
+        mm->add_return(mm->insert_module_instructions(mm->end(), pm, param_map));
 
         std::vector<instruction_ref> inputs;
         std::copy_if(ins->inputs().begin(),
@@ -86,13 +87,13 @@ struct find_conv_pointwise
                      std::back_inserter(inputs),
                      [&](auto input) { return input != conv_ins; });
         inputs.insert(inputs.end(), conv_ins->inputs().begin(), conv_ins->inputs().end());
-        m.replace_instruction(
-            ins, mlir_conv{conv_ins->get_operator()}, inputs, ins->module_inputs());
+        mpm.get_module().replace_instruction(
+            ins, mlir_conv{conv_ins->get_operator()}, inputs, {mm});
     }
 };
 } // namespace
 
-void fuse_mlir::apply(module& m) const { match::find_matches(m, find_conv_pointwise{}); }
+void fuse_mlir::apply(module_pass_manager& mpm) const { match::find_matches(mpm, find_conv_pointwise{}); }
 
 } // namespace gpu
 

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

@@ -7,7 +7,7 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
-struct module;
+struct module_pass_manager;
 
 namespace gpu {
 
@@ -15,7 +15,7 @@ struct fuse_mlir
 {
     context* ctx = nullptr;
     std::string name() const { return "gpu::fuse_mlir"; }
-    void apply(module& m) const;
+    void apply(module_pass_manager& m) const;
 };
 
 } // namespace gpu