Fix hang with CSE pass when using submodules (#1050)

* add submodule test
* remove for loop
* simplify reshape test

src/eliminate_common_subexpression.cpp

@@ -32,7 +32,12 @@ void cse_range(module& p, Range&& r)
                 continue;
             p.replace_instruction(ins, eq);
             processed_ins.emplace(ins);
-            auto outputs = eq->outputs();
+            std::vector<instruction_ref> outputs;
+            std::copy_if(eq->outputs().begin(),
+                         eq->outputs().end(),
+                         std::back_inserter(outputs),
+                         [&](auto x) { return p.has_instruction(x); });
+
             std::sort(outputs.begin(), outputs.end(), [&](auto x, auto y) {
                 return std::distance(eq, x) < std::distance(eq, y);
             });

test/eliminate_common_subexpression_test.cpp

@@ -6,6 +6,12 @@
 
 #include <test.hpp>
 
+void run_pass(migraphx::program& p)
+{
+    migraphx::run_passes(
+        p, {migraphx::eliminate_common_subexpression{}, migraphx::dead_code_elimination{}});
+}
+
 void run_pass(migraphx::module& m)
 {
     migraphx::run_passes(
@@ -142,4 +148,59 @@ TEST_CASE(cse_test_literal)
     EXPECT(m1 == m2);
 }
 
+TEST_CASE(cse_test_submodule)
+{
+    migraphx::shape si{migraphx::shape::int64_type};
+    migraphx::shape s{migraphx::shape::int64_type, {1}};
+    migraphx::shape sc{migraphx::shape::bool_type};
+
+    auto create_program = [&](bool remove_literal = false) {
+        migraphx::program p;
+        std::vector<bool> vc    = {true};
+        std::vector<int64_t> vd = {3};
+        auto* mm                = p.get_main_module();
+
+        auto in_cond = mm->add_parameter("ccond", sc);
+        auto in_val  = mm->add_parameter("val", s);
+        auto b0      = mm->add_literal(migraphx::literal(sc, vc));
+        auto b1      = b0;
+        if(not(remove_literal))
+            b1 = mm->add_literal(migraphx::literal(sc, vc));
+
+        auto* body1 = p.create_module("loop_module1");
+        body1->add_parameter("#loop_module_in_1", sc);
+        auto in_v1 = body1->add_parameter("#loop_module_in_2", s);
+        auto l1    = body1->add_literal(migraphx::literal(si, vd));
+        auto ad1   = body1->add_instruction(migraphx::make_op("add"), l1, l1);
+        auto val1  = body1->add_instruction(migraphx::make_op("add"), in_v1, ad1);
+        auto cond1 = body1->add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::bool_type}}), b0);
+        auto cond2 = body1->add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::bool_type}}), b1);
+        body1->add_return({cond1, cond2, val1, val1});
+
+        auto* body2 = p.create_module("loop_module2");
+        body2->add_parameter("#loop_module_in_1", sc);
+        auto in_v2 = body2->add_parameter("#loop_module_in_2", s);
+        auto l2    = body2->add_literal(migraphx::literal(si, vd));
+        auto ad2   = body2->add_instruction(migraphx::make_op("add"), l2, l2);
+        auto val2  = body2->add_instruction(migraphx::make_op("add"), in_v2, ad2);
+        auto cond3 = body2->add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::bool_type}}), b1);
+        body2->add_return({cond3, val2, val2});
+
+        auto loop1 = mm->add_instruction(
+            migraphx::make_op("loop", {{"max_iterations", 1}}), {in_cond, in_val}, {body1});
+        auto loop2 = mm->add_instruction(
+            migraphx::make_op("loop", {{"max_iterations", 1}}), {in_cond, in_val}, {body2});
+
+        mm->add_return({loop1, loop2});
+
+        return p;
+    };
+    auto p = create_program();
+    run_pass(p);
+    EXPECT(p == create_program(true));
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }