Fuse Split-Reduce with MLIR #3319
Conversation
Yeah. They are a bit tricky to write. Let me add a one/two. I have verify test otherwise. |
| auto mlir_output_with_attrs = | ||
| migraphx::interpolate_string(mlir_output, {{"attrs", get_attrs()}}); | ||
| CHECK(encode(s) == encode(mlir_output_with_attrs)); | ||
| // EXPECT(verify_mlir(m)); |
There was a problem hiding this comment.
verify is failing. Therefore disabling it for now. Could be an issue with rocMLIR.
Added tests |
| migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {2, 32, 10, 64, 64}}}), b); | ||
| auto fused = | ||
| add_mlir(p2, | ||
| "mlir_main:pointwise0_main:split_reduce0", |
There was a problem hiding this comment.
minor: why do we lose "convolution" in the name the MLIR instruction?
There was a problem hiding this comment.
Names are constructed from modules that are fused. convolution or dot would appear as mlir_op[op="" attribute.
AMDMIGraphX/src/targets/gpu/fuse_mlir.cpp
Line 140 in 05b2ff4
src/targets/gpu/mlir.cpp
Outdated
| for(const auto i : iterator_for(m)) | ||
| { | ||
| if(starts_with(i->name(), "@")) | ||
| { | ||
| continue; | ||
| } | ||
| problem_config += " " + i->name(); | ||
| } | ||
| tc.problem = problem_config; |
There was a problem hiding this comment.
I had these changes to experiment and work around problem_cache issue. Reverting these
src/targets/gpu/fuse_mlir.cpp
Outdated
| { | ||
| param_map_2[skip_input] = skip_input; | ||
| } | ||
| return main_mod.fuse(*sub_pm, inputs, ¶m_map_2).front(); |
There was a problem hiding this comment.
Actually, fuse is a poor choice here. Thats why you need to skip the parameters in the param map. Also, it doesnt insert instruction at pos. Instead we should add a insert_inline method that can insert the instructions correctly:
std::vector<instruction_ref>
module::insert_inline(instruction_ref ins,
const module& m,
const std::vector<instruction_ref>& inputs,
std::unordered_map<instruction_ref, instruction_ref>* map_ins,
module::inserter insert)
{
std::unordered_map<instruction_ref, instruction_ref> default_map_ins;
if(map_ins == nullptr)
map_ins = &default_map_ins;
auto param_map = m.get_ins_param_map(inputs, true);
map_ins.insert(param_map.begin(), param_map.end());
return this->insert_instructions(ins, &m, map_ins, std::move(insert));
}Then you can do main_mod.insert_inline(pos, *sub_pm, inputs, ¶m_map_2).front(), and you can skip the skip_input changes.
This build is not recommended to merge 🔴 |
🔴bert_large_uncased_fp16: FAILED: MIGraphX is not within tolerance - check verbose output |
| @@ -50,6 +50,8 @@ struct MIGRAPHX_GPU_EXPORT mlir_code_object | |||
| std::vector<value> prefill_values = {}; | |||
| }; | |||
|
|
|||
| MIGRAPHX_GPU_EXPORT bool is_reduce(const instruction& ins); | |||
There was a problem hiding this comment.
I dont see this used outside of mlir.cpp. I think it can be removed from the header.
There was a problem hiding this comment.
It is being used on both fuse_mlir.cpp and mlir.cpp
Part of #3212
Depends on #3097 #3299 and ROCm/rocMLIR#1590