diff --git a/CHANGELOG.md b/CHANGELOG.md index 66525c60007..edbc8b93963 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -83,6 +83,7 @@ Full documentation for MIGraphX is available at * Fixed `QLinearConv` parsing for models with a bias and per-tensor weight quantization, which previously threw `same_dims: dequantizelinear: Dimensions do not match` (e.g. `resnet50_int8`); the bias scale is now broadcast to the bias shape before dequantizing. * Fixed the GPU problem cache failing to find entries after reload for pooling operator, resulting in redundant re-benchmarking when using a saved `MIGRAPHX_PROBLEM_CACHE`. * Fixed `slice_concat_gather` matcher and interaction between same table and cross table gather fusions(#5038). +* Fixed `gpu::mlir_op` compilation failures for convolution and pointwise fusions followed by layout operations by splitting them into MLIR, pointwise, and layout-copy kernels when needed. ### Optimized * Optimized flash decoding recombination in `fuse_attention` to use the exp-normalize form (#5090). diff --git a/src/targets/gpu/jit/mlir.cpp b/src/targets/gpu/jit/mlir.cpp index 7bc73a80220..e74c2a8537c 100644 --- a/src/targets/gpu/jit/mlir.cpp +++ b/src/targets/gpu/jit/mlir.cpp @@ -135,6 +135,26 @@ static instruction_ref find_final_split(instruction_ref split_ins) return *it; } +static optional find_layout_tail_split(instruction_ref pointwise_ins) +{ + auto output_path_range = get_output_path(pointwise_ins); + std::vector output_path(output_path_range.begin(), output_path_range.end()); + if(output_path.size() < 2) + return nullopt; + auto is_layout = [](instruction_ref ins) { + return contains({"flatten", "reshape", "reshape_lazy", "squeeze", "transpose", "unsqueeze"}, + ins->name()); + }; + auto it = std::find_if(std::next(output_path.begin()), output_path.end(), is_layout); + if(it == output_path.end()) + return nullopt; + if(not std::all_of(it, output_path.end(), [&](instruction_ref ins) { + return is_layout(ins) or ins->name() == "@return"; + })) + return nullopt; + return *std::prev(it); +} + struct mlir_compiler : compiler { std::vector names() const { return {"gpu::mlir_op"}; } @@ -199,6 +219,28 @@ struct mlir_compiler : compiler } } + mlir_code_object + compile_mlir_part(context& ctx, const module_with_inputs& mwi, const value& solution) const + { + auto input_shapes = to_shapes(mwi.inputs); + input_shapes.push_back(mwi.mod.get_output_shapes().front()); + return compile_mlir(ctx, mwi.mod, input_shapes, solution); + } + + code_object_op compile_pointwise_part(context& ctx, module_with_inputs& mwi) const + { + auto input_shapes = to_shapes(mwi.inputs); + if(mwi.mod.get_output_shapes().size() == 1) + { + input_shapes.push_back(mwi.mod.get_output_shapes().front()); + } + else + { + input_shapes.push_back(shape{mwi.mod.get_output_shapes()}); + } + return compile_pointwise_module(ctx, input_shapes, &mwi.mod); + } + compiler_replace compile(context& ctx, instruction_ref ins, const operation&, const value& solution) const { @@ -219,15 +261,37 @@ struct mlir_compiler : compiler auto input_args = ins->inputs(); // remove alloc buffer input_args.pop_back(); - auto split_ins = find_final_split(gemm_like_ins); + auto tail_split = find_layout_tail_split(pointwise_ins); + auto split_ins = + tail_split.has_value() ? gemm_like_ins : find_final_split(gemm_like_ins); std::array mod_splits = smod->split(input_args, {split_ins}); - auto dot_mlir_inputs = to_shapes(mod_splits[0].inputs); - // add alloc for the gemm output - dot_mlir_inputs.push_back(mod_splits[0].mod.get_output_shapes().front()); - mlir_code_object cop1 = compile_mlir(ctx, mod_splits[0].mod, dot_mlir_inputs, solution); - auto pw_shapes = to_shapes(mod_splits[1].inputs); - pw_shapes.push_back(ins->get_shape()); - auto cop2 = compile_pointwise_module(ctx, pw_shapes, &mod_splits[1].mod); + if(not is_module_fusible(mod_splits[0].mod, ctx, solution)) + { + split_ins = gemm_like_ins; + mod_splits = smod->split(input_args, {split_ins}); + } + if(tail_split.has_value()) + { + auto mod_splits3 = smod->split(input_args, {split_ins}, {tail_split.value()}); + auto copy_input_shape = mod_splits3[2].mod.get_output_shapes().front(); + auto copy_cop = any_cast( + gpu::compile_op("hip::copy", + ctx, + {copy_input_shape, ins->inputs().back()->get_shape()}, + {{"lambda", "[](auto x) { return make_tuple(x); }"}, + {"kernel", "hip_copy_kernel"}})); + std::vector cops = { + compile_mlir_part(ctx, mod_splits3[0], solution), + mlir_code_object{compile_pointwise_part(ctx, mod_splits3[1])}, + mlir_code_object{copy_cop}}; + std::array mods = {std::move(mod_splits3[0]), + std::move(mod_splits3[1])}; + return insert(cops, mods, ins, split_ins); + } + + auto cop1 = compile_mlir_part(ctx, mod_splits[0], solution); + auto cop2 = compile_pointwise_part(ctx, mod_splits[1]); + assert(cop2.expected_inputs.back() == ins->get_shape()); std::vector cops = {cop1, mlir_code_object{cop2}}; return insert(cops, mod_splits, ins, split_ins); } @@ -345,8 +409,21 @@ struct mlir_compiler : compiler insert_mlir(m, ins, any_cast(ops[0]), dot_inputs_updated); auto pwm = mods[1]; pwm.replace(split_ins, mlir_ins); + const bool has_copy_tail = ops.size() == 3; auto pw_inputs = pwm.inputs; - pw_inputs.push_back(ins->inputs().back()); + if(has_copy_tail) + { + auto pw_alloc = m.insert_instruction( + ins, + migraphx::make_op( + "hip::allocate", + {{"shape", to_value(mods[1].mod.get_output_shapes().front())}})); + pw_inputs.push_back(pw_alloc); + } + else + { + pw_inputs.push_back(ins->inputs().back()); + } std::vector pw_inputs_updated; std::transform(pw_inputs.begin(), pw_inputs.end(), @@ -361,7 +438,17 @@ struct mlir_compiler : compiler }); auto pw_ins = insert_mlir(m, ins, any_cast(ops[1]), pw_inputs_updated); - return m.replace_instruction(ins, pw_ins); + if(not has_copy_tail) + return m.replace_instruction(ins, pw_ins); + + auto copy_input_shape = any_cast(ops[2]).expected_inputs.front(); + auto copy_input = m.insert_instruction( + ins, + migraphx::make_op("as_shape", {{"shape", to_value(copy_input_shape)}}), + pw_ins); + auto copy_ins = m.insert_instruction( + ins, any_cast(ops[2]), copy_input, ins->inputs().back()); + return m.replace_instruction(ins, copy_ins); }}; } diff --git a/test/verify/test_conv_add_reshape_lazy_transpose.cpp b/test/verify/test_conv_add_reshape_lazy_transpose.cpp new file mode 100644 index 00000000000..3d0b10b0be8 --- /dev/null +++ b/test/verify/test_conv_add_reshape_lazy_transpose.cpp @@ -0,0 +1,64 @@ +/* + * The MIT License (MIT) + * + * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "verify_program.hpp" +#include +#include +#include + +struct test_conv_add_reshape_lazy_transpose + : verify_program +{ + // This covers split-k perfConfigs that reject conv+pointwise+layout fusion and require + // compiling the fused MLIR op as conv, pointwise, and layout-copy kernels. + migraphx::program create_program() const + { + migraphx::program p; + auto* mm = p.get_main_module(); + auto input = mm->add_parameter("x", {migraphx::shape::half_type, {1, 256, 16, 16}}); + auto weight = mm->add_literal( + migraphx::generate_literal({migraphx::shape::half_type, {1, 256, 3, 2}}, 1)); + auto y = mm->add_parameter("y", {migraphx::shape::half_type, {1, 1, 8, 8}}); + + auto conv = mm->add_instruction( + migraphx::make_op("convolution", {{"padding", {1, 1, 1, 0}}, {"stride", {2, 2}}}), + input, + weight); + auto add = mm->add_instruction(migraphx::make_op("add"), conv, y); + + auto reshape = mm->add_instruction( + migraphx::make_op("reshape_lazy", {{"dims", {1, 1, 4, 2, 8}}}), add); + mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 3, 2, 4}}}), + reshape); + return p; + } + + // Turn on Exhaustive-tune to enable split-k perf-configs from MLIR + migraphx::compile_options get_compile_options() const + { + return migraphx::compile_options{.exhaustive_tune = true}; + } + + std::string section() const { return "conv"; } +}; diff --git a/test/verify/test_conv_add_tune.cpp b/test/verify/test_conv_add_tune.cpp index e208b2a872c..b17b4265718 100644 --- a/test/verify/test_conv_add_tune.cpp +++ b/test/verify/test_conv_add_tune.cpp @@ -78,3 +78,4 @@ template struct test_conv_add_tune; template struct test_conv_add_tune; template struct test_conv_add_tune; template struct test_conv_add_tune; + diff --git a/test/verify/test_conv_relu_reshape_lazy_transpose.cpp b/test/verify/test_conv_relu_reshape_lazy_transpose.cpp new file mode 100644 index 00000000000..9ac74465224 --- /dev/null +++ b/test/verify/test_conv_relu_reshape_lazy_transpose.cpp @@ -0,0 +1,62 @@ +/* + * The MIT License (MIT) + * + * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "verify_program.hpp" +#include +#include +#include + +struct test_conv_relu_reshape_lazy_transpose + : verify_program +{ + // This covers split-k perfConfigs that reject conv+unary-pointwise+layout fusion. + migraphx::program create_program() const + { + migraphx::program p; + auto* mm = p.get_main_module(); + auto input = mm->add_parameter("x", {migraphx::shape::half_type, {1, 256, 16, 16}}); + auto weight = mm->add_literal( + migraphx::generate_literal({migraphx::shape::half_type, {1, 256, 3, 2}}, 1)); + + auto conv = mm->add_instruction( + migraphx::make_op("convolution", {{"padding", {1, 1, 1, 0}}, {"stride", {2, 2}}}), + input, + weight); + auto relu = mm->add_instruction(migraphx::make_op("relu"), conv); + + auto reshape = mm->add_instruction( + migraphx::make_op("reshape_lazy", {{"dims", {1, 1, 4, 2, 8}}}), relu); + mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 3, 2, 4}}}), + reshape); + return p; + } + + // Turn on Exhaustive-tune to enable split-k perf-configs from MLIR + migraphx::compile_options get_compile_options() const + { + return migraphx::compile_options{.exhaustive_tune = true}; + } + + std::string section() const { return "conv"; } +};