Move pointwise gen to a seperate header

src/targets/gpu/compile_gen.cpp

@@ -25,6 +25,12 @@
 #include <migraphx/shape.hpp>
 #include <migraphx/permutation.hpp>
 #include <migraphx/stringutils.hpp>
+#include <migraphx/module.hpp>
+#include <migraphx/dead_code_elimination.hpp>
+#include <migraphx/eliminate_common_subexpression.hpp>
+#include <migraphx/cpp_generator.hpp>
+#include <migraphx/pass_manager.hpp>
+#include <migraphx/instruction.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -122,6 +128,47 @@ std::string make_transformer_args(std::vector<std::string> transformers)
     return join_strings(std::move(transformers), ", ");
 }
 
+std::string generate_pointwise(const module& pm, const std::string& name)
+{
+    module m = pm;
+    run_passes(m, {eliminate_common_subexpression{}, dead_code_elimination{}});
+    cpp_generator g;
+    g.fmap([](const std::string& fname) { return "migraphx::" + fname; });
+    g.add_point_op("where", "${function:where}(${0}, ${1}, ${2})");
+    g.add_point_op("prelu", "${function:where}(${0} < 0, ${0} * ${1}, ${0})");
+    g.add_point_op("sign",
+                   "${function:where}(${0} > 0, 1, ${function:where}(${0} < 0, -1, 0))");
+    g.add_point_op("equal", "migraphx::abs(${0} == ${1})");
+    g.add_point_op("less", "migraphx::abs(${0} < ${1})");
+    g.add_point_op("greater", "migraphx::abs(${0} > ${1})");
+    g.add_point_op("not", "migraphx::abs(not ${0})");
+    // Add explict conversions
+    g.fresult([](const shape& s) {
+        return "migraphx::convert<" + shape::cpp_type(s.type()) + ">";
+    });
+    g.create_function(
+        g.generate_module(m).set_attributes({"__device__"}).set_generic_types(m).set_name(name));
+    return g.str();
+}
+
+static std::vector<std::string> get_op_names(const module& m)
+{
+    std::vector<std::string> result;
+    for(auto& ins : m)
+    {
+        if(starts_with(ins.name(), "@"))
+            continue;
+        result.push_back(ins.name());
+    }
+    return result;
+}
+
+std::string generate_name_from_ops(const module& m)
+{
+    auto op_names      = get_op_names(m);
+    return join_strings(op_names, "_");
+}
+
 } // namespace gen
 } // namespace gpu
 } // namespace MIGRAPHX_INLINE_NS

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

@@ -25,6 +25,7 @@
 #define MIGRAPHX_GUARD_GPU_COMPILE_GEN_HPP
 
 #include <migraphx/config.hpp>
+#include <migraphx/module_ref.hpp>
 #include <string>
 #include <unordered_map>
 #include <vector>
@@ -62,6 +63,10 @@ std::string make_transformer_args(Ts... xs)
     return make_transformer_args({xs.str()...});
 }
 
+std::string generate_pointwise(const module& pm, const std::string& name);
+
+std::string generate_name_from_ops(const module& m);
+
 } // namespace gen
 } // namespace gpu
 } // namespace MIGRAPHX_INLINE_NS

src/targets/gpu/jit/pointwise.cpp

@@ -65,18 +65,6 @@ __global__ void ${kernel}(${params})
 
 )__migraphx__";
 
-static std::vector<std::string> get_op_names(const module& m)
-{
-    std::vector<std::string> result;
-    for(auto& ins : m)
-    {
-        if(starts_with(ins.name(), "@"))
-            continue;
-        result.push_back(ins.name());
-    }
-    return result;
-}
-
 struct pointwise_compiler : compiler<pointwise_compiler>
 {
     std::vector<std::string> names() const { return {"pointwise", "contiguous"}; }
@@ -127,31 +115,13 @@ struct pointwise_compiler : compiler<pointwise_compiler>
         {
             assert(not ins->module_inputs().empty());
             auto* pm = ins->module_inputs().front();
-            run_passes(*pm, {eliminate_common_subexpression{}, dead_code_elimination{}});
-            cpp_generator g;
-            g.fmap([](const std::string& fname) { return "migraphx::" + fname; });
-            g.add_point_op("where", "${function:where}(${0}, ${1}, ${2})");
-            g.add_point_op("prelu", "${function:where}(${0} < 0, ${0} * ${1}, ${0})");
-            g.add_point_op("sign",
-                           "${function:where}(${0} > 0, 1, ${function:where}(${0} < 0, -1, 0))");
-            g.add_point_op("equal", "migraphx::abs(${0} == ${1})");
-            g.add_point_op("less", "migraphx::abs(${0} < ${1})");
-            g.add_point_op("greater", "migraphx::abs(${0} > ${1})");
-            g.add_point_op("not", "migraphx::abs(not ${0})");
-            // Add explict conversions
-            g.fresult([](const shape& s) {
-                return "migraphx::convert<" + shape::cpp_type(s.type()) + ">";
-            });
-            auto name = g.create_function(
-                g.generate_module(*pm).set_attributes({"__device__"}).set_generic_types(*pm));
-            std::string lambda = "MIGRAPHX_LIFT(" + name + ")";
-            auto op_names      = get_op_names(*pm);
-            op_names.push_back("kernel");
-            auto op_name_string = join_strings(op_names, "_");
+            auto pf = generate_pointwise(*pm, "inner_pointwise");
+            std::string lambda = "MIGRAPHX_LIFT(inner_pointwise)";
+            auto kernel_name = generate_name_from_ops(*pm) + "_kernel";
             return replace(compile_op(
                 ctx,
                 to_shapes(ins->inputs()),
-                {{"lambda", lambda}, {"preamble", g.str()}, {"kernel", op_name_string}}));
+                {{"lambda", lambda}, {"preamble", pf}, {"kernel", kernel_name}}));
         }
     }
 };