Add output during tracing

src/include/migraphx/program.hpp

@@ -9,6 +9,7 @@
 #include <migraphx/instruction_ref.hpp>
 #include <migraphx/target.hpp>
 #include <migraphx/tracer.hpp>
+#include <migraphx/env.hpp>
 #include <migraphx/config.hpp>
 #include <algorithm>
 #include <iostream>
@@ -16,6 +17,9 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
+MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_COMPILE)
+MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_EVAL)
+
 struct program_impl;
 
 const operation& get_operation(instruction_ref ins);
@@ -107,6 +111,8 @@ struct program
 
     void dry_run(parameter_map params) const;
 
+    void annotate(std::ostream& os, std::function<void(instruction_ref)> a) const;
+
     friend std::ostream& operator<<(std::ostream& os, const program& p);
     friend bool operator==(const program& x, const program& y);
     friend bool operator!=(const program& x, const program& y) { return !(x == y); }

src/program.cpp

@@ -15,9 +15,6 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
-MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_COMPILE)
-MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_EVAL)
-
 struct program_impl
 {
     // A list is used to keep references to an instruction stable
@@ -532,6 +529,13 @@ void program::dry_run(std::unordered_map<std::string, argument> params) const
     generic_eval(*this, ctx, std::move(params), [](auto&&...) { return argument{}; });
 }
 
+void program::annotate(std::ostream& os, std::function<void(instruction_ref)> a) const
+{
+    print_program(os, *this, [&](auto ins, auto&&) {
+        a(ins);
+    });
+}
+
 bool operator==(const program& x, const program& y) { return to_string(x) == to_string(y); }
 
 std::ostream& operator<<(std::ostream& os, const program& p)

src/schedule.cpp

@@ -224,6 +224,17 @@ void schedule::apply(program& p) const
             self(i);
     })(last);
 
+    if(enabled(MIGRAPHX_TRACE_COMPILE{}))
+    {
+        p.annotate(std::cout, [&](auto ins) {
+            std::cout << ":";
+            std::cout << " weight=" << si.weights.at(ins);
+            if (si.has_stream(ins))
+                std::cout << " stream=" << si.get_stream(ins);
+        });
+        std::cout << std::endl;
+    }
+
     // Schedule instructions
     for(auto ins : iterator_for(p))
     {

test/schedule_test.cpp

@@ -120,14 +120,13 @@ bool check_conflicts(migraphx::program& p, migraphx::instruction_ref x, migraphx
     return false;
 }
 
-void check_conflicts(migraphx::program& p,
-                     std::vector<std::vector<migraphx::instruction_ref>> conflicts)
+void check_conflicts(migraphx::program& p, std::vector<std::vector<migraphx::instruction_ref>> conflicts)
 {
     migraphx::dfor(conflicts.size(), conflicts.size())([&](auto i, auto j) {
-        if(i == j)
+        if (i == j)
             return;
-        for(auto ins1 : conflicts[i])
-            for(auto ins2 : conflicts[j])
+        for(auto ins1:conflicts[i])
+            for(auto ins2:conflicts[j])
                 CHECK(check_conflicts(p, ins1, ins2));
     });
 }
@@ -149,12 +148,11 @@ std::vector<std::size_t> get_wait_for(migraphx::instruction_ref ins)
     return wf;
 }
 
-template <class T>
-std::vector<migraphx::instruction_ref>
-chain(migraphx::program& p, std::size_t n, T x, migraphx::instruction_ref input)
+template<class T>
+std::vector<migraphx::instruction_ref> chain(migraphx::program& p, std::size_t n, T x, migraphx::instruction_ref input)
 {
     std::vector<migraphx::instruction_ref> result;
-    for(std::size_t i = 0; i < n; i++)
+    for(std::size_t i = 0;i < n;i++)
     {
         result.push_back(p.add_instruction(x, input));
         input = result.back();
@@ -207,7 +205,7 @@ TEST_CASE(two_weights)
     p.compile(schedule_target{&stream});
     EXPECT(stream.count(one) == 0);
     EXPECT(stream.at(i1) == 1);
-    for(auto ins : c1)
+    for(auto ins:c1)
         EXPECT(stream.at(ins) == 0);
     EXPECT(stream.at(binary) == 0);
     EXPECT(get_wait_for(binary) == get_wait_for(stream[binary], {stream[c1.back()], stream[i1]}));
@@ -223,20 +221,15 @@ TEST_CASE(four_weights)
     auto c2 = chain(p, 3, unary_op{}, one);
     auto c3 = chain(p, 2, unary_op{}, one);
     auto i1   = p.add_instruction(unary_op{}, one);
-    auto binary = p.add_instruction(nary_op{}, i1, c1.back());
+    auto binary = p.add_instruction(nary_op{}, i1, c1.back(), c2.back(), c3.back());
     p.compile(schedule_target{&stream});
     EXPECT(stream.count(one) == 0);
     EXPECT(stream.at(i1) == 3);
-    for(auto ins : c1)
-        EXPECT(stream.at(ins) == 0);
-    for(auto ins : c2)
-        EXPECT(stream.at(ins) == 1);
-    for(auto ins : c3)
-        EXPECT(stream.at(ins) == 2);
+    for(auto ins:c1) EXPECT(stream.at(ins) == 0);
+    for(auto ins:c2) EXPECT(stream.at(ins) == 1);
+    for(auto ins:c3) EXPECT(stream.at(ins) == 2);
     EXPECT(stream.at(binary) == 0);
-    EXPECT(get_wait_for(binary) ==
-           get_wait_for(stream[binary],
-                        {stream[c1.back()], stream[c2.back()], stream[c3.back()], stream[i1]}));
+    EXPECT(get_wait_for(binary) == get_wait_for(stream[binary], {stream[c1.back()], stream[c2.back()], stream[c3.back()], stream[i1]}));
     check_conflicts(p, {c1, c2, c3, {i1}});
 }