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