Method apply

src/memory_coloring.cpp:300–369

}

void memory_coloring::apply(module& m) const
{
    const std::size_t alignment = find_max_alignment(m, allocation_op);
    auto conflict_table         = build_conflict_table(m, allocation_op);
    auto as                     = allocation_segment::build(m, conflict_table, alignment);

    // All allocations should have a segment
    assert(std::all_of(conflict_table.begin(), conflict_table.end(), [&](auto&& pp) {
        return as.get_segment(pp.first);
    }));

    // Adjacent allocations should not have overlapping segments
    assert(std::none_of(conflict_table.begin(), conflict_table.end(), [&](auto&& pp) {
        auto* x = as.get_segment(pp.first);
        return std::any_of(pp.second.begin(), pp.second.end(), [&](auto ins) {
            auto* y = as.get_segment(ins);
            assert(x and y);
            return is_overlap(*x, *y);
        });
    }));

    // Print out segments
    if(enabled(MIGRAPHX_DEBUG_MEMORY_COLORING{}))
    {
        for(auto&& pp : conflict_table)
        {
            std::cout << "------- conflict -------" << std::endl;
            auto s1 = as.ins2segment.at(pp.first);
            std::cout << s1.first << ", " << s1.second << ": ";
            m.debug_print(pp.first);
            for(auto ins : pp.second)
            {
                auto s2 = as.ins2segment.at(ins);
                std::cout << s2.first << ", " << s2.second << ": ";
                m.debug_print(ins);
            }
        }
    }

    // Total memory
    std::size_t n = as.max() * alignment;

    // Replace allocations
    auto mem = m.add_parameter("scratch", shape{shape::int8_type, {n}});
    for(auto&& [ins, seg] : as.ins2segment)
    {
        assert(ins->name() == allocation_op);
        auto s             = ins->get_shape();
        std::size_t offset = seg.first * alignment;
        assert(offset < n);
        m.replace_instruction(
            ins, make_op("load", {{"shape", to_value(s)}, {"offset", offset}}), mem);
    }

    // Replace zero allocation
    for(auto ins : iterator_for(m))
    {
        if(ins->name() != allocation_op)