Fuse Split-Reduce with MLIR #3319
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -43,6 +43,7 @@ namespace gpu { | |
|
|
||
| MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_ENABLE_EXTRA_MLIR); | ||
| MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_ENABLE_MLIR_INPUT_FUSION); | ||
| MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_ENABLE_MLIR_REDUCE_FUSION); | ||
| MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_DISABLE_MLIR); | ||
| /** | ||
| * @brief Declares a new MIGraphX environment variable which forces to generate | ||
|
|
@@ -386,13 +387,59 @@ bool is_pointwise_op_supported_by_mlir(const instruction& i) | |
| return false; | ||
| } | ||
|
|
||
| bool is_reduce_op_supported_by_mlir(const instruction& i) | ||
| { | ||
| using type_t = shape::type_t; | ||
| const auto& name = i.name(); | ||
| const auto result_type = i.get_shape().type(); | ||
| const std::initializer_list<type_t> allowed_types = { | ||
| type_t::float_type, type_t::half_type, type_t::fp8e4m3fnuz_type}; | ||
| // Preliminary type check. | ||
| if(not contains(allowed_types, result_type)) | ||
| { | ||
| return false; | ||
| } | ||
| const std::initializer_list<std::string> reduce_ops = {"reduce_mean", "reduce_sum"}; | ||
| return contains(reduce_ops, i.name()); | ||
| } | ||
|
|
||
| // A separate function so we can remove operators that are supported by mlir | ||
| // but not supported for an input fusion. | ||
| bool is_pointwise_op_supported_by_mlir_for_input(const instruction& i) | ||
| { | ||
| return is_pointwise_op_supported_by_mlir(i); | ||
| } | ||
|
|
||
| MIGRAPHX_PRED_MATCHER(mlir_split_reduce, instruction_ref ins) | ||
| { | ||
| if(ins->name() != "split_fused_reduce") | ||
| return false; | ||
| auto* mod_arg = ins->module_inputs().front(); | ||
| auto supported_reshapes = reshaper_names(); | ||
| supported_reshapes.erase("slice"); | ||
| std::unordered_set<std::string> builtins = {"@param", "@literal", "@return"}; | ||
| for(const auto i : iterator_for(*mod_arg)) | ||
| { | ||
| if(is_reduce(*i)) | ||
| { | ||
| if(not is_reduce_op_supported_by_mlir(*i)) | ||
| return false; | ||
| } | ||
| else if(i->name() == "pointwise") | ||
| { | ||
| if(not std::all_of(i->module_inputs().front()->begin(), | ||
| i->module_inputs().front()->end(), | ||
| &is_pointwise_op_supported_by_mlir)) | ||
| return false; | ||
| } | ||
| else if(not contains(reshaper_names(), i->name()) and not contains(builtins, i->name())) | ||
| { | ||
| return false; | ||
| } | ||
| } | ||
| return true; | ||
| } | ||
|
|
||
| MIGRAPHX_PRED_MATCHER(mlir_pointwise, instruction_ref ins) | ||
| { | ||
| if(ins->name() != "pointwise") | ||
|
|
@@ -423,6 +470,100 @@ std::vector<instruction_ref> mlir_contiguous(module_pass_manager& mpm, | |
| return result; | ||
| } | ||
|
|
||
| struct find_mlir_split_reduce | ||
| { | ||
| mlir_mode conv_mode = mlir_mode::none; | ||
| mlir_mode dot_mode = mlir_mode::none; | ||
| auto matcher() const | ||
| { | ||
| auto dot_or_conv = match::name("gpu::mlir_op"); | ||
| // TODO: Handle reshapes inbetween | ||
| return mlir_split_reduce()(match::any_of[match::inputs()](dot_or_conv.bind("gemm"))); | ||
| } | ||
|
|
||
| void apply(module_pass_manager& mpm, const match::matcher_result& r) const | ||
| { | ||
| auto reduce_ins = r.result; | ||
| auto gemm_ins = r.instructions["gemm"]; | ||
| assert(gemm_ins->get_shape().sub_shapes().empty()); | ||
| auto* rm = reduce_ins->module_inputs().front(); | ||
| auto names = rm->get_parameter_names(); | ||
| std::sort(names.begin(), names.end()); | ||
| module_ref gemm_old_mm = gemm_ins->module_inputs().front(); | ||
| module_ref mm = | ||
| mpm.create_module(gemm_old_mm->name() + "_split_fused_reduce", *gemm_old_mm); | ||
| // remove last return instruction | ||
| if(std::prev(mm->end())->name() == "@return") | ||
| { | ||
| mm->remove_instruction(std::prev(mm->end())); | ||
| } | ||
| mm->set_bypass(); | ||
| std::unordered_map<instruction_ref, instruction_ref> param_map; | ||
| param_map[gemm_ins] = std::prev(mm->end()); | ||
| bool gemm_has_multi_outs = gemm_ins->outputs().size() > 1; | ||
| auto return_vals = | ||
| mm->fuse(*rm, | ||
| reduce_ins->inputs(), | ||
| ¶m_map, | ||
| [&](module& main_mod, | ||
| instruction_ref pos, | ||
| const operation& op, | ||
| const std::vector<instruction_ref>& inputs, | ||
| const std::vector<module_ref>& mod_args) { | ||
| if(op.name() == "pointwise") | ||
| { | ||
| for(const auto& skip_param : inputs) | ||
| { | ||
| if(not contains(param_map, skip_param)) | ||
| { | ||
| param_map[skip_param] = | ||
| skip_param; // skip adding parameter for inputs of | ||
| // pointwise inside split_fused_reduce | ||
| } | ||
| } | ||
|
There was a problem hiding this comment. Why is this needed? If its not in the There was a problem hiding this comment. This is fusing pointwise module arg inside the fused_mlir_module. Note that pointwise module is submodule to Inputs to pointwise would be available in Line 1047 in ae2b026 But they won't be passed down when "inserter" is invoked. Line 260 in ae2b026 There was a problem hiding this comment. I see thats because you are reusing the outer param_map. I am not sure that is needed though. You could just use another param_map: auto param_map_2 = create_param_map_with_literals(&main_mod, sub_pm, op.compute_shape(to_shapes(inputs), mod_args));
return main_mod.fuse(*sub_pm, inputs, ¶m_map_2).front();There was a problem hiding this comment. adding params to same There was a problem hiding this comment. I also need to add There was a problem hiding this comment.
That should never happen. Since the pointwise module is used to generate a single GPU kernel(or function in the case of fused_reduce) it should never access instructions from the parent scope.
Sure, you can use auto param_map_2 = sub_pm->get_ins_param_map(inputs, true);
auto literal_param_map = create_param_map_with_literals(&main_mod, sub_pm, op.compute_shape(to_shapes(inputs), mod_args));
param_map_2.insert(literal_param_map.begin(), literal_param_map.end());
return main_mod.fuse(*sub_pm, inputs, ¶m_map_2).front();There was a problem hiding this comment.
No, I can't it will create map from Line 1004 in ae2b026 I need map from |
||
| auto* sub_pm = mod_args.front(); | ||
| auto param_map_2 = create_param_map_with_literals( | ||
| &main_mod, sub_pm, op.compute_shape(to_shapes(inputs), mod_args)); | ||
| param_map.insert(param_map_2.begin(), param_map_2.end()); | ||
| return main_mod.fuse(*sub_pm, inputs, ¶m_map).front(); | ||
| } | ||
| return main_mod.insert_instruction(pos, op, inputs, mod_args); | ||
| }); | ||
| if(gemm_has_multi_outs) | ||
| { | ||
| return_vals.insert(return_vals.end(), param_map[gemm_ins]); | ||
| } | ||
| mm->add_return(return_vals); | ||
| std::vector<instruction_ref> inputs; | ||
| std::copy_if(reduce_ins->inputs().begin(), | ||
| reduce_ins->inputs().end(), | ||
| std::back_inserter(inputs), | ||
| [&](auto input) { return input != gemm_ins; }); | ||
| inputs.insert(inputs.end(), gemm_ins->inputs().begin(), gemm_ins->inputs().end()); | ||
| if(gemm_has_multi_outs) | ||
| { | ||
| auto fused_ins = mpm.get_module().insert_instruction( | ||
| reduce_ins, mlir_op{gemm_ins->get_operator()}, mlir_contiguous(mpm, inputs), {mm}); | ||
| auto dot_ins = mpm.get_module().insert_instruction( | ||
| reduce_ins, | ||
| migraphx::make_op("get_tuple_elem", {{"index", return_vals.size() - 1}}), | ||
| fused_ins); | ||
|
|
||
| mpm.get_module().replace_instruction(gemm_ins, dot_ins); | ||
| for(const auto outs : reduce_ins->outputs()) | ||
| { | ||
| assert(outs->get_operator().name() == "get_tuple_elem"); | ||
| mpm.get_module().replace_instruction(outs, outs->get_operator(), fused_ins); | ||
| } | ||
| } | ||
| else | ||
| { | ||
| mpm.get_module().replace_instruction( | ||
| reduce_ins, mlir_op{gemm_ins->get_operator()}, mlir_contiguous(mpm, inputs), {mm}); | ||
| } | ||
| } | ||
| }; | ||
|
|
||
| struct find_mlir_fused_ops | ||
| { | ||
| mlir_mode conv_mode = mlir_mode::none; | ||
|
|
@@ -714,15 +855,25 @@ void fuse_mlir::apply(module_pass_manager& mpm) const | |
| mpm, | ||
| find_mlir_fused_ops{.conv_mode = get_mode("fused_convolution", mlir_mode::fast), | ||
| .dot_mode = get_mode("fused_dot", mlir_mode::fast)}); | ||
|
|
||
| match::find_matches( | ||
| mpm, | ||
| find_mlir_standalone_convolution_op{get_mode("convolution", mlir_mode::fast)}, | ||
| find_mlir_standalone_dot_op{get_mode("dot", mlir_mode::fast)}); | ||
|
|
||
| mpm.run_pass(dead_code_elimination{}); | ||
| if(enabled(MIGRAPHX_ENABLE_MLIR_REDUCE_FUSION{})) | ||
| { | ||
| match::find_matches( | ||
| mpm, | ||
| find_mlir_split_reduce{.conv_mode = get_mode("fused_convolution", mlir_mode::fast), | ||
| .dot_mode = get_mode("fused_dot", mlir_mode::fast)}); | ||
| } | ||
|
|
||
| if(enabled(MIGRAPHX_ENABLE_MLIR_INPUT_FUSION{})) | ||
| { | ||
| match::find_matches(mpm, find_pointwise_mlir{}); | ||
| } | ||
| #else | ||
| (void)mpm; | ||
| #endif | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why is this flipped?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are cases where same input instruction is mapped to multiple parameters.
e.g.
split_fused_reduce(x, y, x)In those cases, having mapping from
input--> paramwould de-duplicate it and only add single parameter.Later
AMDMIGraphX/src/module.cpp
Line 238 in ae2b026
Here it won't find parameter in the
map_insand would try to add it and fails