Add initial scheduler with topo sort

src/CMakeLists.txt

@@ -17,6 +17,7 @@ add_library(migraphx
     instruction.cpp
     program.cpp
     shape.cpp
+    schedule.cpp
     simplify_algebra.cpp
     simplify_reshapes.cpp
     opt/memory_coloring.cpp

src/include/migraphx/schedule.hpp

+#ifndef MIGRAPHX_GUARD_RTGLIB_SCHEDULE_HPP
+#define MIGRAPHX_GUARD_RTGLIB_SCHEDULE_HPP
+
+#include <string>
+#include <migraphx/instruction_ref.hpp>
+#include <migraphx/schedule_model.hpp>
+#include <migraphx/config.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+struct program;
+
+/**
+ * Schedule instructions for concurrent execution
+ */
+struct schedule
+{
+    schedule_model model{};
+    std::string name() const { return "schedule"; }
+    void apply(program& p) const;
+};
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/schedule_model.hpp

+#ifndef MIGRAPHX_GUARD_SCHEDULE_MODEL_HPP
+#define MIGRAPHX_GUARD_SCHEDULE_MODEL_HPP
+
+#include <cassert>
+#include <string>
+#include <functional>
+#include <memory>
+#include <type_traits>
+#include <utility>
+
+#include <migraphx/config.hpp>
+#include <migraphx/instruction_ref.hpp>
+#include <vector>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+struct program;
+struct operation;
+
+#ifdef DOXYGEN
+
+/// An interface for target-dependent model for the scheduler
+struct schedule_model
+{
+    /// Get the number of concurrent instruction allowed
+    std::size_t concurrency() const;
+    /// Schedule a concurrent instruction
+    void schedule_instruction(program& p, instruction_ref ins, std::size_t n) const;
+    // Insert necessary waits before an instruction
+    void wait(program& p,
+              instruction_ref ins,
+              std::size_t wait_on,
+              const std::vector<std::size_t>& wait_for) const;
+    /// Compute weights for an operation
+    std::size_t weight(const operation& op) const;
+};
+
+#else
+
+/*
+ * Type-erased interface for:
+ *
+ * struct schedule_model
+ * {
+ *      std::size_t concurrency() const;
+ *      void schedule_instruction(program& p,instruction_ref ins,std::size_t n) const;
+ *      void wait(program& p,instruction_ref ins,std::size_t wait_on,const std::vector<std::size_t>&
+ * wait_for) const; std::size_t weight(const operation& op) const;
+ * };
+ *
+ */
+
+struct schedule_model
+{
+    // Constructors
+    schedule_model() = default;
+
+    template <typename PrivateDetailTypeErasedT>
+    schedule_model(PrivateDetailTypeErasedT value)
+        : private_detail_te_handle_mem_var(
+              std::make_shared<private_detail_te_handle_type<
+                  typename std::remove_reference<PrivateDetailTypeErasedT>::type>>(
+                  std::forward<PrivateDetailTypeErasedT>(value)))
+    {
+    }
+
+    // Assignment
+    template <typename PrivateDetailTypeErasedT>
+    schedule_model& operator=(PrivateDetailTypeErasedT value)
+    {
+        if(private_detail_te_handle_mem_var.unique())
+            *private_detail_te_handle_mem_var = std::forward<PrivateDetailTypeErasedT>(value);
+        else if(!private_detail_te_handle_mem_var)
+            private_detail_te_handle_mem_var = std::make_shared<PrivateDetailTypeErasedT>(
+                std::forward<PrivateDetailTypeErasedT>(value));
+        return *this;
+    }
+
+    // Cast
+    template <typename PrivateDetailTypeErasedT>
+    PrivateDetailTypeErasedT* any_cast()
+    {
+        return private_detail_te_get_handle().type() == typeid(PrivateDetailTypeErasedT)
+                   ? std::addressof(static_cast<private_detail_te_handle_type<
+                                        typename std::remove_cv<PrivateDetailTypeErasedT>::type>&>(
+                                        private_detail_te_get_handle())
+                                        .private_detail_te_value)
+                   : nullptr;
+    }
+
+    template <typename PrivateDetailTypeErasedT>
+    const typename std::remove_cv<PrivateDetailTypeErasedT>::type* any_cast() const
+    {
+        return private_detail_te_get_handle().type() == typeid(PrivateDetailTypeErasedT)
+                   ? std::addressof(static_cast<const private_detail_te_handle_type<
+                                        typename std::remove_cv<PrivateDetailTypeErasedT>::type>&>(
+                                        private_detail_te_get_handle())
+                                        .private_detail_te_value)
+                   : nullptr;
+    }
+
+    const std::type_info& type_id() const
+    {
+        if(private_detail_te_handle_empty())
+            return typeid(std::nullptr_t);
+        else
+            return private_detail_te_get_handle().type();
+    }
+
+    std::size_t concurrency() const
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        return (*this).private_detail_te_get_handle().concurrency();
+    }
+
+    void schedule_instruction(program& p, instruction_ref ins, std::size_t n) const
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        (*this).private_detail_te_get_handle().schedule_instruction(
+            p, std::move(ins), std::move(n));
+    }
+
+    void wait(program& p,
+              instruction_ref ins,
+              std::size_t wait_on,
+              const std::vector<std::size_t>& wait_for) const
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        (*this).private_detail_te_get_handle().wait(
+            p, std::move(ins), std::move(wait_on), wait_for);
+    }
+
+    std::size_t weight(const operation& op) const
+    {
+        assert((*this).private_detail_te_handle_mem_var);
+        return (*this).private_detail_te_get_handle().weight(op);
+    }
+
+    friend bool is_shared(const schedule_model& private_detail_x,
+                          const schedule_model& private_detail_y)
+    {
+        return private_detail_x.private_detail_te_handle_mem_var ==
+               private_detail_y.private_detail_te_handle_mem_var;
+    }
+
+    private:
+    struct private_detail_te_handle_base_type
+    {
+        virtual ~private_detail_te_handle_base_type() {}
+        virtual std::shared_ptr<private_detail_te_handle_base_type> clone() const = 0;
+        virtual const std::type_info& type() const                                = 0;
+
+        virtual std::size_t concurrency() const                                                 = 0;
+        virtual void schedule_instruction(program& p, instruction_ref ins, std::size_t n) const = 0;
+        virtual void wait(program& p,
+                          instruction_ref ins,
+                          std::size_t wait_on,
+                          const std::vector<std::size_t>& wait_for) const                       = 0;
+        virtual std::size_t weight(const operation& op) const                                   = 0;
+    };
+
+    template <typename PrivateDetailTypeErasedT>
+    struct private_detail_te_handle_type : private_detail_te_handle_base_type
+    {
+        template <typename PrivateDetailTypeErasedU = PrivateDetailTypeErasedT>
+        private_detail_te_handle_type(
+            PrivateDetailTypeErasedT value,
+            typename std::enable_if<std::is_reference<PrivateDetailTypeErasedU>::value>::type* =
+                nullptr)
+            : private_detail_te_value(value)
+        {
+        }
+
+        template <typename PrivateDetailTypeErasedU = PrivateDetailTypeErasedT>
+        private_detail_te_handle_type(
+            PrivateDetailTypeErasedT value,
+            typename std::enable_if<!std::is_reference<PrivateDetailTypeErasedU>::value,
+                                    int>::type* = nullptr) noexcept
+            : private_detail_te_value(std::move(value))
+        {
+        }
+
+        std::shared_ptr<private_detail_te_handle_base_type> clone() const override
+        {
+            return std::make_shared<private_detail_te_handle_type>(private_detail_te_value);
+        }
+
+        const std::type_info& type() const override { return typeid(private_detail_te_value); }
+
+        std::size_t concurrency() const override { return private_detail_te_value.concurrency(); }
+
+        void schedule_instruction(program& p, instruction_ref ins, std::size_t n) const override
+        {
+
+            private_detail_te_value.schedule_instruction(p, std::move(ins), std::move(n));
+        }
+
+        void wait(program& p,
+                  instruction_ref ins,
+                  std::size_t wait_on,
+                  const std::vector<std::size_t>& wait_for) const override
+        {
+
+            private_detail_te_value.wait(p, std::move(ins), std::move(wait_on), wait_for);
+        }
+
+        std::size_t weight(const operation& op) const override
+        {
+
+            return private_detail_te_value.weight(op);
+        }
+
+        PrivateDetailTypeErasedT private_detail_te_value;
+    };
+
+    template <typename PrivateDetailTypeErasedT>
+    struct private_detail_te_handle_type<std::reference_wrapper<PrivateDetailTypeErasedT>>
+        : private_detail_te_handle_type<PrivateDetailTypeErasedT&>
+    {
+        private_detail_te_handle_type(std::reference_wrapper<PrivateDetailTypeErasedT> ref)
+            : private_detail_te_handle_type<PrivateDetailTypeErasedT&>(ref.get())
+        {
+        }
+    };
+
+    bool private_detail_te_handle_empty() const
+    {
+        return private_detail_te_handle_mem_var == nullptr;
+    }
+
+    const private_detail_te_handle_base_type& private_detail_te_get_handle() const
+    {
+        assert(private_detail_te_handle_mem_var != nullptr);
+        return *private_detail_te_handle_mem_var;
+    }
+
+    private_detail_te_handle_base_type& private_detail_te_get_handle()
+    {
+        assert(private_detail_te_handle_mem_var != nullptr);
+        if(!private_detail_te_handle_mem_var.unique())
+            private_detail_te_handle_mem_var = private_detail_te_handle_mem_var->clone();
+        return *private_detail_te_handle_mem_var;
+    }
+
+    std::shared_ptr<private_detail_te_handle_base_type> private_detail_te_handle_mem_var;
+};
+
+template <typename ValueType>
+inline const ValueType* any_cast(const schedule_model* x)
+{
+    return x->any_cast<ValueType>();
+}
+
+template <typename ValueType>
+inline ValueType* any_cast(schedule_model* x)
+{
+    return x->any_cast<ValueType>();
+}
+
+template <typename ValueType>
+inline ValueType& any_cast(schedule_model& x)
+{
+    auto* y = x.any_cast<typename std::remove_reference<ValueType>::type>();
+    if(y == nullptr)
+        throw std::bad_cast();
+    return *y;
+}
+
+template <typename ValueType>
+inline const ValueType& any_cast(const schedule_model& x)
+{
+    const auto* y = x.any_cast<typename std::remove_reference<ValueType>::type>();
+    if(y == nullptr)
+        throw std::bad_cast();
+    return *y;
+}
+
+#endif
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/target.hpp

@@ -22,10 +22,8 @@ struct target
 {
     /// A unique name used to identify the target
     std::string name() const;
-    /// The transformation passes to be run
     /**
      * @brief The transformation pass to be run during compilation.
-     * @details [long description]
      *
      * @param ctx This is the target-dependent context that is created by `get_context`
      * @return The passes to be ran

src/schedule.cpp

+#include <migraphx/schedule.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/iterator_for.hpp>
+#include <migraphx/functional.hpp>
+#include <migraphx/ranges.hpp>
+#include <unordered_map>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+void schedule::apply(program& p) const
+{
+    // Compute accumulated weights
+    std::unordered_map<instruction_ref, std::size_t> weights;
+    auto last = std::prev(p.end());
+    fix<std::size_t>([&](auto self, auto ins) -> std::size_t {
+        if (weights.count(ins) == 0) 
+        {
+            weights[ins] = std::accumulate(ins->inputs().begin(), ins->inputs().end(), model.weight(ins->get_operator()), [&](std::size_t w, instruction_ref i) {
+                return w + self(i);
+            });
+        }
+        return weights[ins];
+    })(last);
+
+    // Topo sort
+    fix([&](auto self, auto ins) {
+        for(auto i:ins->inputs())
+            p.move_instruction(i, p.begin());
+        for(auto i:ins->inputs())
+            self(i);
+    })(last);
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+

src/targets/gpu/CMakeLists.txt

@@ -62,6 +62,7 @@ add_library(migraphx_gpu
     pad.cpp
     gather.cpp
     lrn.cpp
+    schedule_model.cpp
 )
 set_target_properties(migraphx_gpu PROPERTIES EXPORT_NAME gpu)
 rocm_clang_tidy_check(migraphx_gpu)

src/targets/gpu/include/migraphx/gpu/schedule_model.hpp

+#ifndef MIGRAPHX_GUARD_RTGLIB_GPU_SCHEDULE_MODEL_HPP
+#define MIGRAPHX_GUARD_RTGLIB_GPU_SCHEDULE_MODEL_HPP
+
+#include <migraphx/config.hpp>
+#include <migraphx/instruction_ref.hpp>
+#include <vector>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+struct program;
+struct operation;
+
+namespace gpu {
+
+struct schedule_model
+{
+    std::size_t n = 4;
+    std::size_t concurrency() const;
+    void schedule_instruction(program& p, instruction_ref ins, std::size_t n) const;
+    void wait(program& p, instruction_ref ins, std::size_t wait_on, const std::vector<std::size_t>& wait_for) const;
+    std::size_t weight(const operation& op) const;
+};
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/targets/gpu/schedule_model.cpp

+#include <migraphx/gpu/schedule_model.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/operation.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+std::size_t schedule_model::concurrency() const
+{
+    return n;
+}
+void schedule_model::schedule_instruction(program& p, instruction_ref ins, std::size_t n) const
+{}
+void schedule_model::wait(program& p, instruction_ref ins, std::size_t wait_on, const std::vector<std::size_t>& wait_for) const
+{}
+std::size_t schedule_model::weight(const operation& op) const
+{
+    return 1;
+}
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
\ No newline at end of file

src/targets/gpu/target.cpp

@@ -18,6 +18,8 @@
 #include <migraphx/rewrite_rnn.hpp>
 #include <migraphx/eliminate_concat.hpp>
 #include <migraphx/gpu/concat_gpu_opt.hpp>
+#include <migraphx/gpu/schedule_model.hpp>
+#include <migraphx/schedule.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -51,6 +53,7 @@ std::vector<pass> target::get_passes(migraphx::context& gctx) const
         fuse_ops{&ctx},
         dead_code_elimination{},
         write_literals{&ctx},
+        schedule{gpu::schedule_model{}},
         memory_coloring{"hip::allocate"},
         eliminate_workspace{},
         eliminate_allocation{"hip::allocate"},

tools/include/schedule_model.hpp

+#ifndef MIGRAPHX_GUARD_SCHEDULE_MODEL_HPP
+#define MIGRAPHX_GUARD_SCHEDULE_MODEL_HPP
+
+#include <cassert>
+#include <string>
+#include <functional>
+#include <memory>
+#include <type_traits>
+#include <utility>
+
+#include <migraphx/config.hpp>
+#include <migraphx/instruction_ref.hpp>
+#include <vector>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+struct program;
+struct operation;
+
+#ifdef DOXYGEN
+
+/// An interface for target-dependent model for the scheduler
+struct schedule_model
+{
+    /// Get the number of concurrent instruction allowed
+    std::size_t concurrency() const;
+    /// Schedule a concurrent instruction
+    void schedule_instruction(program& p, instruction_ref ins, std::size_t n) const;
+    // Insert necessary waits before an instruction 
+    void wait(program& p, instruction_ref ins, std::size_t wait_on, const std::vector<std::size_t>& wait_for) const;
+    /// Compute weights for an operation
+    std::size_t weight(const operation& op) const;
+};
+
+#else
+
+<%
+interface('schedule_model',
+    virtual('concurrency', returns='std::size_t', const=True),
+    virtual('schedule_instruction', p='program&', ins='instruction_ref', n='std::size_t', const=True),
+    virtual('wait', p='program&', ins='instruction_ref', wait_on='std::size_t', wait_for='const std::vector<std::size_t>&', const=True),
+    virtual('weight', returns='std::size_t', op='const operation&', const=True)
+)
+%>
+
+#endif
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif