Method apply

src/split_reduce.cpp:176–242

}

void split_reduce::apply(module_pass_manager& mpm) const
{
    for(auto ins : iterator_for(mpm.get_module()))
    {
        if(ins->name() != "fused_reduce")
            continue;
        auto* rm = ins->module_inputs().front();
        if(get_reduce_size(rm) < split_size)
            continue;
        splitter s{rm};
        auto splits = s.find_splits();
        if(splits.empty())
            continue;
        // Only use split reduce with float for now
        // TODO: Support other data types
        if(not std::all_of(splits.begin(), splits.end(), [](instruction_ref split) {
               return contains({shape::float_type, shape::half_type}, split->get_shape().type());
           }))
            continue;
        auto v    = ins->get_operator().to_value();
        auto axes = v["axes"].to_vector<std::int64_t>();

        auto alive = s.find_alive(splits);

        std::array<module::with_inputs, 2> mods;
        if(not alive.empty())
        {
            auto mods3 = rm->split(ins->inputs(), alive, splits);
            auto r     = insert_module_inline(mpm.get_module(), ins, mods3[0]);
            mods3[1].replace(alive, r);
            mods3[2].replace(alive, r);
            mods = {std::move(mods3[1]), std::move(mods3[2])};
        }
        else
        {
            mods = rm->split(ins->inputs(), splits);
        }

        auto* splitm = mpm.create_module(rm->name() + "_split", std::move(mods[0].mod));
        splitm->set_bypass();

        // Insert split reduce
        auto split_reduce = mpm.get_module().insert_instruction(
            ins,
            make_op("split_fused_reduce", {{"axes", axes}, {"assign", assign_op(splits)}}),
            mods[0].inputs,
            {splitm});

        std::vector<instruction_ref> split_reduce_each;
        if(splits.size() == 1)
        {
            split_reduce_each = {split_reduce};
        }
        else
        {
            transform(range(splits.size()), std::back_inserter(split_reduce_each), [&](auto i) {
                return mpm.get_module().insert_instruction(
                    ins, make_op("get_tuple_elem", {{"index", i}}), split_reduce);