Horizontally fuse cross-embedding gather operators #4599
Conversation
| { | ||
| if(finder.is_candidate(ins)) | ||
| groups[finder.group_key(ins)].push_back(ins); | ||
| } |
There was a problem hiding this comment.
I think we should use group_by instead otherwise this wont find groups of independent instructions:
std::vector<instruction_ref> groups;
std::copy_if(...);
std::unordered_map<instruction_ref, std::size_t> pos;
std::adjacent_difference(...);
auto pred = [&](instruction_ref x, instruciton_ref y) {
if(a == b)
return true;
if(x->get_operator() != y->get_operator())
return false;
if(finder.group_key(a) == finder.group_key(b))
return false;
if(pos.at(a) < pos.at(b))
return reaches(a, b);
return reaches(b, a);
};
auto each = [&](auto start, auto last) {
auto n = std::distance(start, last);
if(n < finder.min_group_size())
return;
...
};
group_by(groups.begin(), groups.end(), each, pred);
| return false; | ||
| } | ||
| return true; | ||
| }); |
There was a problem hiding this comment.
This should be checked in the group_key not here.
| auto emb = g->inputs().at(0); | ||
| emb_inputs.push_back(emb); | ||
| emb_offsets.push_back(cumulative_offset); | ||
| cumulative_offset += emb->get_shape().lens().front(); |
There was a problem hiding this comment.
Actually this should be a transform_partial_sum.
| for(auto g : gathers) | ||
| { | ||
| auto emb = g->inputs().at(0); | ||
| emb_inputs.push_back(emb); |
There was a problem hiding this comment.
I dont think this vector is needed thought, It can be easily computed directly.
| return {}; | ||
|
|
||
| // Insert after the last gather in the group (already sorted by position) | ||
| auto insert_pt = std::next(gathers.back()); |
There was a problem hiding this comment.
The insert point should be passed into the fuse function.
| continue; | ||
|
|
||
| assert(replacements.size() == group.size()); | ||
| for(std::size_t i = 0; i < group.size(); i++) |
There was a problem hiding this comment.
Use our for_each algorithm that takes two sequence(like std::transform).
| std::vector<instruction_ref> adjusted_idx_inputs; | ||
| std::vector<std::size_t> idx_sizes; | ||
|
|
||
| for(std::size_t i = 0; i < gathers.size(); i++) |
There was a problem hiding this comment.
Use for_each that takes two sequences.
| {"ends", std::vector<int64_t>{end}}}), | ||
| batched_gather)); | ||
|
|
||
| slice_offset += idx_size; |
There was a problem hiding this comment.
Compute start using partial_sum.
| { | ||
| m.replace_instruction(group[i], replacements[i]); | ||
| } | ||
| } |
There was a problem hiding this comment.
Before the replacements we need to call move_output_instructions_after so the outputs come after inserted operator. This is in fuse_pointwise. It needs to be moved to the module class and unit tests added. I can try and an open a PR for that. It also needs to be fixed because it doesnt work with submodules. For now you will need to add a check like we doin fuse_pointwise.
|
|
||
| fuse_horizontal_ops(m, gather_horizontal_fusion{}); | ||
|
|
||
| dead_code_elimination{}.apply(m); |
There was a problem hiding this comment.
Not necessary, just run this pass afterwards in the target passes.
| run_pass(m1); | ||
| EXPECT(m1 == m2); | ||
| } | ||
|
|
There was a problem hiding this comment.
We need more tests with more complicated graphs as well such as using the first gather before the second gather.
2 participants
Motivation
Technical Details
Changelog Category
Add a
CHANGELOG.mdentry for any option other thanNot Applicable