diff --git a/src/simplify_algebra.cpp b/src/simplify_algebra.cpp index b054d007edf..3ce1f4d915a 100644 --- a/src/simplify_algebra.cpp +++ b/src/simplify_algebra.cpp @@ -1095,11 +1095,9 @@ 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. +// 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 { auto matcher() const { return match::name("concat")(match::same_inputs()); } @@ -1119,18 +1117,25 @@ 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; + const auto& out_lens = ins->get_shape().lens(); - auto out_lens = lens; - out_lens[axis] = inputs.size(); - assert(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. + if(lens[axis] == 1) + { + m.replace_instruction(ins, make_op("multibroadcast", {{"out_lens", out_lens}}), x); + return; + } - m.replace_instruction(ins, make_op("multibroadcast", {{"out_lens", out_lens}}), x); + // 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(); + assert(bcast_lens[axis] == 1); + 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); }