Test for waits

test/include/test.hpp

@@ -36,6 +36,19 @@ inline std::ostream& operator<<(std::ostream& s, std::nullptr_t)
     return s;
 }
 
+template<class T>
+inline std::ostream& operator<<(std::ostream& s, const std::vector<T>& v)
+{
+    char delim = '{';
+    for(auto&& x:v) 
+    {
+        s << delim << " " << x;
+        delim = ',';
+    }
+    s << " }";
+    return s;
+}
+
 template <class T, class U, class Operator>
 struct expression
 {

test/schedule_test.cpp

@@ -45,6 +45,23 @@ struct binary_op
     }
 };
 
+struct wait_event
+{
+    std::vector<std::size_t> wait_for;
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return migraphx::pack(f(self.wait_for, "wait_for"));
+    }
+    std::string name() const { return "wait_event"; }
+    migraphx::shape compute_shape(const std::vector<migraphx::shape>&) const { return {}; }
+
+    migraphx::argument compute(migraphx::context&, const migraphx::shape&, const std::vector<migraphx::argument>&) const
+    {
+        return {};
+    }
+};
+
 using instruction_map = std::unordered_map<migraphx::instruction_ref, std::size_t>;
 
 struct schedule_model_test
@@ -52,7 +69,7 @@ struct schedule_model_test
     instruction_map* ins2stream;
     std::size_t concurrency() const { return 4; }
     void
-    schedule_instruction(migraphx::program& p, migraphx::instruction_ref ins, std::size_t n) const
+    schedule_instruction(migraphx::program&, migraphx::instruction_ref ins, std::size_t n) const
     {
         (*ins2stream)[ins] = n;
     }
@@ -61,6 +78,8 @@ struct schedule_model_test
               std::size_t wait_on,
               const std::vector<std::size_t>& wait_for) const
     {
+        (*ins2stream)[ins] = wait_on;
+        p.insert_instruction(ins, wait_event{wait_for});
     }
     std::size_t weight(const migraphx::operation& op) const
     {
@@ -98,18 +117,36 @@ bool check_conflicts(migraphx::program& p, migraphx::instruction_ref x, migraphx
     return false;
 }
 
+std::vector<std::size_t> get_wait_for(std::size_t wait_on, std::vector<std::size_t> wait_for)
+{
+    wait_for.erase(std::find(wait_for.begin(), wait_for.end(), wait_on));
+    std::sort(wait_for.begin(), wait_for.end());
+    return wait_for;
+}
+
+std::vector<std::size_t> get_wait_for(migraphx::instruction_ref ins)
+{
+    auto wait_ins = std::prev(ins);
+    if (wait_ins->name() != "wait_event")
+        return {};
+    auto wf = migraphx::any_cast<wait_event>(wait_ins->get_operator()).wait_for;
+    std::sort(wf.begin(), wf.end());
+    return wf;
+}
+
 TEST_CASE(test1)
 {
-    instruction_map ins2stream;
+    instruction_map stream;
     migraphx::program p;
     auto one    = p.add_literal(1);
     auto two    = p.add_literal(2);
     auto onep   = p.add_instruction(unary_op{}, one);
     auto twop   = p.add_instruction(unary_op{}, two);
     auto binary = p.add_instruction(binary_op{}, onep, twop);
-    p.compile(schedule_target{&ins2stream});
-    EXPECT(ins2stream.at(onep) != ins2stream.at(twop));
-    EXPECT(ins2stream.at(binary) == 0);
+    p.compile(schedule_target{&stream});
+    EXPECT(stream.at(onep) != stream.at(twop));
+    EXPECT(stream.at(binary) == 0);
+    EXPECT(get_wait_for(binary) == get_wait_for(stream[binary], {stream[onep], stream[twop]}));
     EXPECT(check_conflicts(p, onep, twop));
 }