From 27bfbcfbf246de1f1fb1ccbb49893a0d71e1d995 Mon Sep 17 00:00:00 2001 From: Paul Date: Sun, 16 Aug 2026 16:22:23 -0500 Subject: [PATCH 1/2] Update find_concat_same_inputs to handle dimension > 1 --- src/simplify_algebra.cpp | 34 +++++++++------- test/simplify_algebra_test.cpp | 71 ++++++++++++++++++++++++++++++---- 2 files changed, 83 insertions(+), 22 deletions(-) diff --git a/src/simplify_algebra.cpp b/src/simplify_algebra.cpp index b054d007edf..45389528be3 100644 --- a/src/simplify_algebra.cpp +++ b/src/simplify_algebra.cpp @@ -1096,10 +1096,8 @@ struct find_concat_op }; // Collapse `concat(x, x, ..., x)` (N copies of the same instruction) into a -// single `multibroadcast` when the concat axis has length 1 in the source -// tensor. This is the common shape that shows up in MoE / KV-cache / RoPE -// expansion code where a tensor is replicated N times along an axis. The -// rewrite turns an O(output_size) memcpy into a strided view. +// broadcast of x. This is the common shape that shows up in MoE / KV-cache / +// RoPE expansion code where a tensor is replicated N times along an axis. struct find_concat_same_input { auto matcher() const { return match::name("concat")(match::same_inputs()); } @@ -1119,18 +1117,26 @@ struct find_concat_same_input if(axis < 0 or axis >= lens.size()) return; - // Safe (no data movement) case: the concat axis is size 1 in the - // source, so it can be broadcast to N. The general lens[axis] > 1 - // case requires unsqueeze + multGibroadcast + reshape and is left - // to a follow-up matcher. - if(lens[axis] != 1) - return; - - auto out_lens = lens; - out_lens[axis] = inputs.size(); + auto out_lens = lens; + out_lens[axis] *= inputs.size(); assert(out_lens == ins->get_shape().lens()); - m.replace_instruction(ins, make_op("multibroadcast", {{"out_lens", out_lens}}), x); + // The concat axis is size 1 in the source, so replicating it is a + // strided view with no data movement. + if(lens[axis] == 1) + { + m.replace_instruction(ins, make_op("multibroadcast", {{"out_lens", out_lens}}), x); + return; + } + + // General case: tile the axis by unsqueezing a unit dim before it, + // broadcasting that dim to N, then folding it back into the axis. + auto unsqueezed = m.insert_instruction(ins, make_op("unsqueeze", {{"axes", {axis}}}), x); + auto bcast_lens = unsqueezed->get_shape().lens(); + bcast_lens[axis] = inputs.size(); + auto bcast = m.insert_instruction( + ins, make_op("multibroadcast", {{"out_lens", bcast_lens}}), unsqueezed); + m.replace_instruction(ins, make_op("reshape", {{"dims", out_lens}}), bcast); } }; diff --git a/test/simplify_algebra_test.cpp b/test/simplify_algebra_test.cpp index babfa599aa3..47010afb2d5 100644 --- a/test/simplify_algebra_test.cpp +++ b/test/simplify_algebra_test.cpp @@ -3130,11 +3130,15 @@ TEST_CASE(simplify_dot_horiz_same_constant) migraphx::module m2; { - auto input = m2.add_parameter("input", s); - auto a = m2.add_literal(migraphx::generate_literal(s, 0)); - auto concat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 2}}), a, a); - auto dot = m2.add_instruction(migraphx::make_op("dot"), input, concat); - auto x = m2.add_instruction( + auto input = m2.add_parameter("input", s); + auto a = m2.add_literal(migraphx::generate_literal(s, 0)); + auto unsqueeze = m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {2}}}), a); + auto bcast = m2.add_instruction( + migraphx::make_op("multibroadcast", {{"out_lens", {3, 2, 2, 2}}}), unsqueeze); + auto reshape = + m2.add_instruction(migraphx::make_op("reshape", {{"dims", {3, 2, 4}}}), bcast); + auto dot = m2.add_instruction(migraphx::make_op("dot"), input, reshape); + auto x = m2.add_instruction( migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {2}}}), dot); auto y = m2.add_instruction( migraphx::make_op("slice", {{"axes", {2}}, {"starts", {2}}, {"ends", {4}}}), dot); @@ -5743,9 +5747,60 @@ TEST_CASE(simplify_concat_same_input_axis_not_one) migraphx::module m2; { - auto x = m2.add_parameter("x", s); - auto concat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), x, x); - m2.add_return({concat}); + auto x = m2.add_parameter("x", s); + auto unsqueeze = m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {1}}}), x); + auto bcast = m2.add_instruction( + migraphx::make_op("multibroadcast", {{"out_lens", {2, 2, 4, 4}}}), unsqueeze); + auto reshape = + m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 8, 4}}}), bcast); + m2.add_return({reshape}); + } + EXPECT(m1 == m2); +} + +TEST_CASE(simplify_concat_same_input_tile_axis_one_not_last) +{ + auto s = migraphx::shape{migraphx::shape::float_type, {2, 3, 5}}; + migraphx::module m1; + { + auto x = m1.add_parameter("x", s); + auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), x, x, x); + m1.add_return({concat}); + } + run_pass(m1); + + migraphx::module m2; + { + auto x = m2.add_parameter("x", s); + auto unsqueeze = m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {1}}}), x); + auto bcast = m2.add_instruction( + migraphx::make_op("multibroadcast", {{"out_lens", {2, 3, 3, 5}}}), unsqueeze); + auto reshape = + m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 9, 5}}}), bcast); + m2.add_return({reshape}); + } + EXPECT(m1 == m2); +} + +TEST_CASE(simplify_concat_same_input_tile_axis_zero) +{ + auto s = migraphx::shape{migraphx::shape::float_type, {3, 2}}; + migraphx::module m1; + { + auto x = m1.add_parameter("x", s); + auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), x, x, x); + m1.add_return({concat}); + } + run_pass(m1); + + migraphx::module m2; + { + auto x = m2.add_parameter("x", s); + auto unsqueeze = m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {0}}}), x); + auto bcast = m2.add_instruction( + migraphx::make_op("multibroadcast", {{"out_lens", {3, 3, 2}}}), unsqueeze); + auto reshape = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {9, 2}}}), bcast); + m2.add_return({reshape}); } EXPECT(m1 == m2); } From bc737c67294e67222660a09a01c687b06d637a62 Mon Sep 17 00:00:00 2001 From: Paul Date: Sun, 16 Aug 2026 16:32:39 -0500 Subject: [PATCH 2/2] Simplify --- src/simplify_algebra.cpp | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/simplify_algebra.cpp b/src/simplify_algebra.cpp index 45389528be3..3ce1f4d915a 100644 --- a/src/simplify_algebra.cpp +++ b/src/simplify_algebra.cpp @@ -1095,8 +1095,8 @@ struct find_concat_op } }; -// Collapse `concat(x, x, ..., x)` (N copies of the same instruction) into a -// broadcast of x. This is the common shape that shows up in MoE / KV-cache / +// Rewrite `concat(x, x, ..., x)` (N copies of the same instruction) as a +// broadcast-based tiling of x. This pattern shows up in MoE / KV-cache / // RoPE expansion code where a tensor is replicated N times along an axis. struct find_concat_same_input { @@ -1117,9 +1117,7 @@ struct find_concat_same_input if(axis < 0 or axis >= lens.size()) return; - auto out_lens = lens; - out_lens[axis] *= inputs.size(); - assert(out_lens == ins->get_shape().lens()); + const auto& out_lens = ins->get_shape().lens(); // The concat axis is size 1 in the source, so replicating it is a // strided view with no data movement. @@ -1131,8 +1129,9 @@ struct find_concat_same_input // General case: tile the axis by unsqueezing a unit dim before it, // broadcasting that dim to N, then folding it back into the axis. - auto unsqueezed = m.insert_instruction(ins, make_op("unsqueeze", {{"axes", {axis}}}), x); - auto bcast_lens = unsqueezed->get_shape().lens(); + auto unsqueezed = m.insert_instruction(ins, make_op("unsqueeze", {{"axes", {axis}}}), x); + auto bcast_lens = unsqueezed->get_shape().lens(); + assert(bcast_lens[axis] == 1); bcast_lens[axis] = inputs.size(); auto bcast = m.insert_instruction( ins, make_op("multibroadcast", {{"out_lens", bcast_lens}}), unsqueezed);