Fix #222: Assertion in dead code elimination

src/dead_code_elimination.cpp

@@ -4,6 +4,7 @@
 #include <migraphx/iterator_for.hpp>
 #include <migraphx/functional.hpp>
 #include <migraphx/ranges.hpp>
+#include <unordered_set>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -50,7 +51,7 @@ void dead_code_elimination::apply(program& p) const
             assert(p.has_instruction(leaf));
             if(leaf->outputs().empty())
             {
-                auto args = leaf->inputs();
+                std::unordered_set<instruction_ref> args(leaf->inputs().begin(), leaf->inputs().end());
                 leaf->clear_arguments();
                 assert(bidistance(p, last, leaf) < 0);
                 assert(leaf != ins);

test/dead_code_elimination_test.cpp

@@ -129,4 +129,55 @@ TEST_CASE(undefined_test)
     EXPECT(result != migraphx::literal{4});
 }
 
+TEST_CASE(duplicate_args1)
+{
+    migraphx::program p;
+
+    auto l0 = p.add_literal(0);
+    auto l3 = p.add_literal(3);
+    p.add_instruction(migraphx::op::add{}, l3, l3);
+    p.add_instruction(migraphx::op::identity{}, l0);
+    auto count = std::distance(p.begin(), p.end());
+    p.compile(dce_target{});
+    EXPECT(std::distance(p.begin(), p.end()) != count);
+    EXPECT(std::distance(p.begin(), p.end()) == 2);
+    auto result = p.eval({});
+    EXPECT(result == migraphx::literal{0});
+}
+
+TEST_CASE(duplicate_args2)
+{
+    migraphx::program p;
+
+    auto l0 = p.add_literal(0);
+    auto l3 = p.add_literal(3);
+    auto sum1 = p.add_instruction(migraphx::op::add{}, l0, l3);
+    p.add_instruction(migraphx::op::add{}, sum1, l3);
+    p.add_instruction(migraphx::op::identity{}, l0);
+    auto count = std::distance(p.begin(), p.end());
+    p.compile(dce_target{});
+    EXPECT(std::distance(p.begin(), p.end()) != count);
+    EXPECT(std::distance(p.begin(), p.end()) == 2);
+    auto result = p.eval({});
+    EXPECT(result == migraphx::literal{0});
+}
+
+TEST_CASE(duplicate_args3)
+{
+    migraphx::program p;
+
+    auto l0 = p.add_literal(0);
+    auto l3 = p.add_literal(3);
+    auto sum1 = p.add_instruction(migraphx::op::add{}, l0, l3);
+    auto sum2 = p.add_instruction(migraphx::op::add{}, l0, sum1);
+    p.add_instruction(migraphx::op::add{}, sum2, l3);
+    p.add_instruction(migraphx::op::identity{}, l0);
+    auto count = std::distance(p.begin(), p.end());
+    p.compile(dce_target{});
+    EXPECT(std::distance(p.begin(), p.end()) != count);
+    EXPECT(std::distance(p.begin(), p.end()) == 2);
+    auto result = p.eval({});
+    EXPECT(result == migraphx::literal{0});
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }