Move instruction to end before removing them

src/dead_code_elimination.cpp

@@ -8,26 +8,24 @@ namespace migraph {
 
 void dead_code_elimination::apply(program& p) const
 {
-    for(auto i : iterator_for(p))
+    auto last = std::prev(p.end());
+    for(auto ins : iterator_for(p))
     {
-        // Skip over instructions that may have been removed
-        if(!p.has_instruction(i))
-            continue;
         // Skip the last instruction
-        if(i == std::prev(p.end()))
+        if(ins == last)
             break;
-        fix([&](auto self, auto ins) {
-            assert(p.has_instruction(ins));
-            if(ins->output.empty())
+        fix([&](auto self, auto leaf) {
+            assert(p.has_instruction(leaf));
+            if(leaf->output.empty())
             {
-                std::cout << p << std::endl;
-                auto args = ins->arguments;
-                p.remove_instruction(ins);
+                auto args = leaf->arguments;
+                p.move_instruction(leaf, p.end());
                 for(auto arg : args)
                     self(arg);
             }
-        })(i);
+        })(ins);
     }
+    p.remove_instructions(std::next(last), p.end());
 }
 
 } // namespace migraph

src/include/migraph/program.hpp

@@ -53,6 +53,9 @@ struct program
     replace_instruction(instruction_ref ins, operation op, std::vector<instruction_ref> args);
 
     instruction_ref remove_instruction(instruction_ref ins);
+    instruction_ref remove_instructions(instruction_ref first, instruction_ref last);
+    
+    instruction_ref move_instruction(instruction_ref src, instruction_ref dst);
 
     template <class... Ts>
     instruction_ref add_literal(Ts&&... xs)

src/program.cpp

@@ -60,6 +60,23 @@ instruction_ref program::remove_instruction(instruction_ref ins)
     return impl->instructions.erase(ins);
 }
 
+instruction_ref program::remove_instructions(instruction_ref first, instruction_ref last)
+{
+    assert(has_instruction(first));
+    assert(has_instruction(last));
+    std::for_each(first, last, [&](instruction& ins) {
+        ins.clear_arguments();
+    });
+    assert(std::all_of(first, last, [&](instruction& ins){ return ins.output.empty(); }));
+    return impl->instructions.erase(first, last);
+}
+
+instruction_ref program::move_instruction(instruction_ref src, instruction_ref dst)
+{
+    impl->instructions.splice(dst, impl->instructions, src);
+    return src;
+}
+
 instruction_ref program::add_literal(literal l)
 {
     impl->instructions.emplace_front(std::move(l));

test/CMakeLists.txt

@@ -52,6 +52,8 @@ function(add_test_command NAME EXE)
                 if(NOT RESULT EQUAL 0)
                     # TODO: check for core files based on pid when setting /proc/sys/kernel/core_uses_pid
                     if(EXISTS ${TEST_DIR}/core)
+                        set(\$ENV{UBSAN_OPTIONS} print_stacktrace=1)
+                        set(\$ENV{ASAN_OPTIONS} print_stacktrace=1)
                         execute_process(COMMAND ${MIGRAPH_GDB} $<TARGET_FILE:${EXE}> ${TEST_DIR}/core -batch -ex bt)
                     endif()
                     message(FATAL_ERROR \"Test failed\")

test/dead_code_elimination_test.cpp

@@ -27,7 +27,7 @@ void simple_test()
     EXPECT(result != migraph::literal{4});
 }
 
-void duplicate_test()
+void duplicate_test1()
 {
     migraph::program p;
 
@@ -43,8 +43,26 @@ void duplicate_test()
     EXPECT(result != migraph::literal{4});
 }
 
+void duplicate_test2()
+{
+    migraph::program p;
+
+    auto one = p.add_literal(1);
+    auto two = p.add_literal(2);
+    p.add_instruction(sum_op{}, one, two);
+    p.add_instruction(minus_op{}, one, two);
+    p.add_instruction(sum_op{}, one, two);
+    auto count = std::distance(p.begin(), p.end());
+    p.compile(dce_target{});
+    EXPECT(std::distance(p.begin(), p.end()) == (count - 2));
+    auto result = p.eval({});
+    EXPECT(result == migraph::literal{3});
+    EXPECT(result != migraph::literal{4});
+}
+
 int main()
 {
     simple_test();
-    duplicate_test();
+    duplicate_test1();
+    duplicate_test2();
 }