Fix MLIR conv-pointwise-layout fusion splitting

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).

src/targets/gpu/jit/mlir.cpp

@@ -135,6 +135,26 @@
     return *it;
 }
 
+static optional<instruction_ref> find_layout_tail_split(instruction_ref pointwise_ins)
+{
+    auto output_path_range = get_output_path(pointwise_ins);
+    std::vector<instruction_ref> 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<mlir_compiler>
 {
     std::vector<std::string> names() const { return {"gpu::mlir_op"}; }
@@ -199,6 +219,28 @@
         }
     }
 
+    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 @@
             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<module_with_inputs, 2> 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<code_object_op>(
+                    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<mlir_code_object> 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<module_with_inputs, 2> 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<mlir_code_object> cops = {cop1, mlir_code_object{cop2}};
             return insert(cops, mod_splits, ins, split_ins);
         }
@@ -345,8 +409,21 @@
                     insert_mlir(m, ins, any_cast<code_object_op>(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<instruction_ref> pw_inputs_updated;
                 std::transform(pw_inputs.begin(),
                                pw_inputs.end(),
@@ -361,7 +438,17 @@
                                });
                 auto pw_ins =
                     insert_mlir(m, ins, any_cast<code_object_op>(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<code_object_op>(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<code_object_op>(ops[2]), copy_input, ins->inputs().back());
+                return m.replace_instruction(ins, copy_ins);
             }};
     }
 

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 <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_conv_add_reshape_lazy_transpose
+    : verify_program<test_conv_add_reshape_lazy_transpose>
+{
+    // 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"; }
+};

test/verify/test_conv_add_tune.cpp

@@ -78,3 +78,4 @@ template struct test_conv_add_tune<migraphx::shape::fp8e4m3fnuz_type>;
 template struct test_conv_add_tune<migraphx::shape::fp8e5m2fnuz_type>;
 template struct test_conv_add_tune<migraphx::shape::fp8e4m3fn_type>;
 template struct test_conv_add_tune<migraphx::shape::fp8e5m2_type>;
+

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 <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_conv_relu_reshape_lazy_transpose
+    : verify_program<test_conv_relu_reshape_lazy_transpose>
+{
+    // 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"; }
+};