From f90215f75b9613e67ab31d9bb666adca66678759 Mon Sep 17 00:00:00 2001 From: Paul Date: Mon, 12 Oct 2026 17:06:33 -0500 Subject: [PATCH 1/6] Fuse concat swapped halves --- src/simplify_algebra.cpp | 69 +++++++++ src/targets/gpu/fuse_concat_past_present.cpp | 50 +++++-- test/gpu/fuse_concat_past_present.cpp | 83 +++++++++++ test/simplify_algebra_test.cpp | 142 +++++++++++++++++++ 4 files changed, 332 insertions(+), 12 deletions(-) diff --git a/src/simplify_algebra.cpp b/src/simplify_algebra.cpp index 1191d663ca2..26e1f506fec 100644 --- a/src/simplify_algebra.cpp +++ b/src/simplify_algebra.cpp @@ -1212,6 +1212,74 @@ struct find_concat_same_broadcast } }; +// Matches a concat of the two equal halves of one tensor in swapped order +// along the concat axis (a rotate-half), with both slices identical on +// every other axis. +MIGRAPHX_PRED_MATCHER(concat_of_swapped_halves, instruction_ref ins) +{ + if(ins->name() != "concat" or ins->get_shape().dynamic()) + return false; + const auto& inputs = ins->inputs(); + if(inputs.size() != 2) + return false; + if(not all_of(inputs, [](instruction_ref i) { return i->name() == "slice"; })) + return false; + if(inputs[0]->inputs().front() != inputs[1]->inputs().front()) + return false; + auto axis = any_cast(ins->normalized_operator()).axis; + auto hi = any_cast(inputs[0]->normalized_operator()); + auto lo = any_cast(inputs[1]->normalized_operator()); + if(hi.axes != lo.axes) + return false; + auto it = std::find(hi.axes.begin(), hi.axes.end(), axis); + if(it == hi.axes.end()) + return false; + auto i = std::distance(hi.axes.begin(), it); + if(hi.starts[i] != lo.ends[i] or hi.ends[i] - hi.starts[i] != lo.ends[i] - lo.starts[i]) + return false; + // The other axes must slice the same range + auto starts = lo.starts; + auto ends = lo.ends; + starts[i] = hi.starts[i]; + ends[i] = hi.ends[i]; + return std::tie(hi.starts, hi.ends) == std::tie(starts, ends); +} + +// Replace a swapped-halves concat with a `where` over the split axis, so the +// data movement becomes a pointwise op over views that can fuse into its +// consumers (such as the rotary prologue of attention), with no copy kernel. +struct find_concat_swapped_halves +{ + auto matcher() const { return concat_of_swapped_halves(); } + + void apply(module& m, const match::matcher_result& r) const + { + auto ins = r.result; + auto axis = any_cast(ins->normalized_operator()).axis; + const auto& out_lens = ins->get_shape().lens(); + // Split the concat axis into {2, half} + auto lens = out_lens; + lens.insert(lens.begin() + axis, 2); + lens[axis + 1] /= 2; + std::vector mask_lens(lens.size(), 1); + mask_lens[axis] = 2; + auto mask = m.add_literal(literal{shape{shape::bool_type, mask_lens}, {true, false}}); + std::vector inputs = { + m.insert_instruction(ins, make_op("multibroadcast", {{"out_lens", lens}}), mask)}; + std::transform(ins->inputs().begin(), + ins->inputs().end(), + std::back_inserter(inputs), + [&](instruction_ref half) { + auto unsqueezed = m.insert_instruction( + ins, make_op("unsqueeze", {{"axes", {axis}}}), half); + return m.insert_instruction( + ins, make_op("multibroadcast", {{"out_lens", lens}}), unsqueezed); + }); + auto w = m.insert_instruction(ins, make_op("where"), inputs); + m.replace_instruction(ins, make_op("reshape", {{"dims", out_lens}}), w); + } +}; + struct find_concat_conv { auto matcher() const @@ -2935,6 +3003,7 @@ void simplify_algebra::apply(module& m) const find_conv_concat_split_fuse{}, find_concat_same_input{}, find_concat_same_broadcast{}, + find_concat_swapped_halves{}, find_concat_op{}, find_split_concat{}, find_splits{}, diff --git a/src/targets/gpu/fuse_concat_past_present.cpp b/src/targets/gpu/fuse_concat_past_present.cpp index c7941e0cb72..14fbd18ae2c 100644 --- a/src/targets/gpu/fuse_concat_past_present.cpp +++ b/src/targets/gpu/fuse_concat_past_present.cpp @@ -29,6 +29,8 @@ #include #include #include +#include +#include #include #include #include @@ -102,21 +104,38 @@ struct find_concat_past_present match::arg(0)(match::skip(view)(producer))); } - // The producer must write exactly the elements of the cache slot, in the - // same memory order - static bool fills_slot(instruction_ref producer, instruction_ref cur, instruction_ref cache) + // The cache slot the append writes, in the layout of cur + static shape slot_shape(instruction_ref cur, instruction_ref cache) + { + const auto& ks = cache->get_shape(); + auto lens = ks.lens(); + lens[seq_axis] = cur->get_shape().lens()[seq_axis]; + return {ks.type(), lens, ks.strides()}; + } + + // The producer must write exactly the elements of the cache slot in the + // same memory order. The views between it and the concat only re-index a + // standard tensor, so the slot is re-indexed back to the producer dims + // as a lazy reshape; nullopt when it cant be. + static optional + producer_slot(instruction_ref producer, instruction_ref cur, instruction_ref cache) { const auto& ps = producer->get_shape(); const auto& cs = cur->get_shape(); const auto& ks = cache->get_shape(); - if(ps != cs or not ps.standard() or not ks.standard()) - return false; + if(not ps.standard() or not cs.standard() or not ks.standard()) + return nullopt; + if(ps.type() != cs.type() or ps.elements() != cs.elements()) + return nullopt; if(cs.ndim() != 4 or ks.ndim() != 4) - return false; - auto seq = cs.lens()[seq_axis]; - auto slot = ks.lens(); - slot[seq_axis] = seq; - return cs.lens() == slot and seq <= ks.lens()[seq_axis]; + return nullopt; + auto slot = slot_shape(cur, cache); + if(cs.lens() != slot.lens() or slot.lens()[seq_axis] > ks.lens()[seq_axis]) + return nullopt; + auto reshaped = reshape_dims(slot, ps.lens(), {.lazy = true}); + if(not reshaped.has_value()) + return nullopt; + return reshaped->to_static(); } // One host copy of the sequence length serves every concat that reads it @@ -137,7 +156,8 @@ struct find_concat_past_present auto cur = ins->inputs()[0]; auto slk = ins->inputs()[1]; auto cache = ins->inputs()[2]; - if(not fills_slot(producer, cur, cache)) + auto pslot = producer_slot(producer, cur, cache); + if(not pslot.has_value()) return; assert(not producer->inputs().empty()); auto alloc = producer->inputs().back(); @@ -174,8 +194,14 @@ struct find_concat_past_present m.move_instruction(producer, ins); std::for_each(chain.rbegin(), chain.rend(), [&](auto x) { m.move_instruction(x, ins); }); + assert(view->get_shape() == slot_shape(cur, cache)); + // The producer writes its own dims, so give it the slot re-indexed as a view + if(view->get_shape() != *pslot) + view = + m.insert_instruction(ins, make_op("reshape_lazy", {{"dims", pslot->lens()}}), view); + assert(view->get_shape() == *pslot); auto v = producer->get_operator().to_value(); - v["output_shape"] = to_value(view->get_shape()); + v["output_shape"] = to_value(*pslot); auto inputs = producer->inputs(); inputs.back() = view; m.replace_instruction( diff --git a/test/gpu/fuse_concat_past_present.cpp b/test/gpu/fuse_concat_past_present.cpp index 20b2f7f0a85..015d27bb632 100644 --- a/test/gpu/fuse_concat_past_present.cpp +++ b/test/gpu/fuse_concat_past_present.cpp @@ -150,6 +150,89 @@ TEST_CASE(skip_multi_use_producer) EXPECT(p1.sort() == p2.sort()); } +// pointwise over the present shape split as {b, h, seq, 2, head/2} -> +// reshape_lazy -> concat_past_present; the producer writes the slot through +// the reshape +TEST_CASE(fuse_decode_reshaped_producer) +{ + migraphx::shape split_shape{migraphx::shape::half_type, {1, 2, 1, 2, 2}}; + migraphx::program p1; + { + auto* mm = p1.get_main_module(); + auto x = mm->add_parameter("x", split_shape); + auto y = mm->add_parameter("y", split_shape); + auto slk = mm->add_parameter("slk", index_shape()); + auto cache = mm->add_parameter("cache", cache_shape()); + auto* pm = create_pointwise_module(p1, "main:pointwise0", {x, y}, single_pointwise("mul")); + auto alloc = mm->add_instruction( + migraphx::make_op("allocate", {{"shape", migraphx::to_value(split_shape)}})); + auto pw = mm->add_instruction(make_precompile_op("pointwise"), {x, y, alloc}, {pm}); + auto rsp = mm->add_instruction( + migraphx::make_op("reshape_lazy", {{"dims", present_shape(1).lens()}}), pw); + auto cpp = mm->add_instruction( + make_precompile_op(migraphx::make_op("concat_past_present", {{"kv_num_heads", 2}}), + cache_shape()), + rsp, + slk, + cache); + mm->add_return({cpp}); + } + run_pass(p1); + + migraphx::program p2; + { + auto* mm = p2.get_main_module(); + auto x = mm->add_parameter("x", split_shape); + auto y = mm->add_parameter("y", split_shape); + auto slk = mm->add_parameter("slk", index_shape()); + auto cache = mm->add_parameter("cache", cache_shape()); + auto* pm = create_pointwise_module(p2, "main:pointwise0", {x, y}, single_pointwise("mul")); + auto scalar = mm->add_instruction(migraphx::make_op("hip::load_scalar"), slk); + auto slot = + mm->add_instruction(migraphx::make_op("gpu::slice_at", {{"axis", 2}}), cache, scalar); + auto view = mm->add_instruction( + migraphx::make_op("reshape_lazy", {{"dims", split_shape.lens()}}), slot); + migraphx::shape split_slot{migraphx::shape::half_type, {1, 2, 1, 2, 2}, {64, 32, 4, 2, 1}}; + auto pw = mm->add_instruction( + make_precompile_op(migraphx::make_op("pointwise"), split_slot), {x, y, view}, {pm}); + auto dep = mm->add_instruction(migraphx::make_op("identity"), cache, pw); + mm->add_return({dep}); + } + EXPECT(p1.sort() == p2.sort()); +} + +// A transpose that reorders non-unit dims changes the memory order, so the +// producer cannot write the slot directly +TEST_CASE(skip_transposed_producer) +{ + migraphx::shape transposed_shape{migraphx::shape::half_type, {1, 4, 1, 2}}; + migraphx::program p1; + { + auto* mm = p1.get_main_module(); + auto x = mm->add_parameter("x", transposed_shape); + auto y = mm->add_parameter("y", transposed_shape); + auto slk = mm->add_parameter("slk", index_shape()); + auto cache = mm->add_parameter("cache", cache_shape()); + auto* pm = create_pointwise_module(p1, "main:pointwise0", {x, y}, single_pointwise("mul")); + auto alloc = mm->add_instruction( + migraphx::make_op("allocate", {{"shape", migraphx::to_value(transposed_shape)}})); + auto pw = mm->add_instruction(make_precompile_op("pointwise"), {x, y, alloc}, {pm}); + auto trp = mm->add_instruction( + migraphx::make_op("transpose", {{"permutation", {0, 3, 2, 1}}}), pw); + auto cpp = mm->add_instruction( + make_precompile_op(migraphx::make_op("concat_past_present", {{"kv_num_heads", 2}}), + cache_shape()), + trp, + slk, + cache); + mm->add_return({cpp}); + } + migraphx::program p2 = p1; + run_pass(p1); + + EXPECT(p1.sort() == p2.sort()); +} + // Unlowered decode append: mul(x, y) -> concat_past_present(mul, slk, cache) static migraphx::program make_decode_program() { diff --git a/test/simplify_algebra_test.cpp b/test/simplify_algebra_test.cpp index 2583ac73497..ef67cdfa26d 100644 --- a/test/simplify_algebra_test.cpp +++ b/test/simplify_algebra_test.cpp @@ -6284,6 +6284,148 @@ TEST_CASE(simplify_concat_same_broadcast_nonunit_axis_dim) EXPECT(m1 == m2); } +TEST_CASE(simplify_concat_swapped_halves) +{ + migraphx::shape s{migraphx::shape::float_type, {2, 6}}; + migraphx::module m1; + { + auto x = m1.add_parameter("x", s); + auto hi = m1.add_instruction( + migraphx::make_op("slice", {{"axes", {1}}, {"starts", {3}}, {"ends", {6}}}), x); + auto lo = m1.add_instruction( + migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {3}}}), x); + auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), hi, lo); + m1.add_return({concat}); + } + run_pass(m1); + + // The unsqueezed slices simplify further into slices of a split reshape + migraphx::module m2; + { + auto x = m2.add_parameter("x", s); + auto split = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2, 3}}}), x); + auto mask = m2.add_literal( + migraphx::literal{{migraphx::shape::bool_type, {1, 2, 1}}, {true, false}}); + auto mask_b = m2.add_instruction( + migraphx::make_op("multibroadcast", {{"out_lens", {2, 2, 3}}}), mask); + auto hi = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1}}, {"ends", {2}}}), split); + auto hi_b = + m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2, 2, 3}}}), hi); + auto lo = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {1}}}), split); + auto lo_b = + m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2, 2, 3}}}), lo); + auto w = m2.add_instruction(migraphx::make_op("where"), mask_b, hi_b, lo_b); + auto reshape = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 6}}}), w); + m2.add_return({reshape}); + } + EXPECT(m1 == m2); +} + +TEST_CASE(simplify_concat_swapped_halves_multi_axis) +{ + migraphx::shape s{migraphx::shape::float_type, {1, 4, 6}}; + migraphx::module m1; + { + auto x = m1.add_parameter("x", s); + auto hi = m1.add_instruction( + migraphx::make_op("slice", {{"axes", {1, 2}}, {"starts", {0, 3}}, {"ends", {2, 6}}}), + x); + auto lo = m1.add_instruction( + migraphx::make_op("slice", {{"axes", {1, 2}}, {"starts", {0, 0}}, {"ends", {2, 3}}}), + x); + auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 2}}), hi, lo); + m1.add_return({concat}); + } + run_pass(m1); + + migraphx::module m2; + { + auto x = m2.add_parameter("x", s); + auto hi = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {1, 2}}, {"starts", {0, 3}}, {"ends", {2, 6}}}), + x); + auto lo = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {1, 2}}, {"starts", {0, 0}}, {"ends", {2, 3}}}), + x); + auto mask = m2.add_literal( + migraphx::literal{{migraphx::shape::bool_type, {1, 1, 2, 1}}, {true, false}}); + auto mask_b = m2.add_instruction( + migraphx::make_op("multibroadcast", {{"out_lens", {1, 2, 2, 3}}}), mask); + auto hi_u = m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {2}}}), hi); + auto hi_b = m2.add_instruction( + migraphx::make_op("multibroadcast", {{"out_lens", {1, 2, 2, 3}}}), hi_u); + auto lo_u = m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {2}}}), lo); + auto lo_b = m2.add_instruction( + migraphx::make_op("multibroadcast", {{"out_lens", {1, 2, 2, 3}}}), lo_u); + auto w = m2.add_instruction(migraphx::make_op("where"), mask_b, hi_b, lo_b); + auto reshape = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 2, 6}}}), w); + m2.add_return({reshape}); + } + EXPECT(m1 == m2); +} + +TEST_CASE(simplify_concat_swapped_halves_unequal) +{ + migraphx::shape s{migraphx::shape::float_type, {2, 6}}; + migraphx::module m1; + { + auto x = m1.add_parameter("x", s); + auto hi = m1.add_instruction( + migraphx::make_op("slice", {{"axes", {1}}, {"starts", {4}}, {"ends", {6}}}), x); + auto lo = m1.add_instruction( + migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {4}}}), x); + auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), hi, lo); + m1.add_return({concat}); + } + auto m2 = m1; + run_pass(m1); + + EXPECT(m1 == m2); +} + +TEST_CASE(simplify_concat_swapped_halves_different_inputs) +{ + migraphx::shape s{migraphx::shape::float_type, {2, 6}}; + migraphx::module m1; + { + auto x = m1.add_parameter("x", s); + auto y = m1.add_parameter("y", s); + auto hi = m1.add_instruction( + migraphx::make_op("slice", {{"axes", {1}}, {"starts", {3}}, {"ends", {6}}}), y); + auto lo = m1.add_instruction( + migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {3}}}), x); + auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), hi, lo); + m1.add_return({concat}); + } + auto m2 = m1; + run_pass(m1); + + EXPECT(m1 == m2); +} + +TEST_CASE(simplify_concat_swapped_halves_other_axis_differs) +{ + migraphx::shape s{migraphx::shape::float_type, {1, 4, 6}}; + migraphx::module m1; + { + auto x = m1.add_parameter("x", s); + auto hi = m1.add_instruction( + migraphx::make_op("slice", {{"axes", {1, 2}}, {"starts", {0, 3}}, {"ends", {2, 6}}}), + x); + auto lo = m1.add_instruction( + migraphx::make_op("slice", {{"axes", {1, 2}}, {"starts", {2, 0}}, {"ends", {4, 3}}}), + x); + auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 2}}), hi, lo); + m1.add_return({concat}); + } + auto m2 = m1; + run_pass(m1); + + EXPECT(m1 == m2); +} + TEST_CASE(simplify_concat_unsqueeze) { migraphx::module m1; From 2d85fff03db8cbab898dee3f9471ddbdc12e8964 Mon Sep 17 00:00:00 2001 From: Paul Date: Mon, 12 Oct 2026 17:45:43 -0500 Subject: [PATCH 2/6] Handle view ops as well --- src/eliminate_concat.cpp | 90 +++++++-------- src/include/migraphx/eliminate_concat.hpp | 12 ++ src/targets/gpu/fuse_concat_past_present.cpp | 64 +++++------ test/eliminate_concat_test.cpp | 51 +++++++++ test/gpu/fuse_concat_past_present.cpp | 112 ++++++++++++++++++- 5 files changed, 245 insertions(+), 84 deletions(-) diff --git a/src/eliminate_concat.cpp b/src/eliminate_concat.cpp index 123dd183678..3410b0fc3f4 100644 --- a/src/eliminate_concat.cpp +++ b/src/eliminate_concat.cpp @@ -40,6 +40,47 @@ namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { +optional, shape>> invert_view_chain( + const std::vector& chain, const shape& input_shape, const shape& view_shape) +{ + std::vector fwd; + std::transform(chain.begin(), chain.end(), std::back_inserter(fwd), [](const operation& op) { + if(op.name() == "reshape_lazy") + return make_op("reshape", {{"dims", op.to_value().at("dims")}}); + return op; + }); + auto inv = shape_transform_descriptor::create(input_shape.lens(), fwd).invert(); + if(inv.empty()) + return std::nullopt; + auto result = inv.generate(); + // Reshapes need to be lazy so the view is aliased instead of copied at runtime + std::transform(result.begin(), result.end(), result.begin(), [](const operation& op) { + if(op.name() != "reshape") + return op; + return make_op("reshape_lazy", op.to_value()); + }); + try + { + auto s = std::accumulate( + result.begin(), result.end(), view_shape, [](const shape& acc, const operation& op) { + std::vector inputs = {acc}; + if(op.output_alias(inputs).empty()) + MIGRAPHX_THROW("Not a view operator: " + op.name()); + auto next = op.compute_shape(inputs); + if(next.elements() != acc.elements()) + MIGRAPHX_THROW("Not an element-preserving view: " + op.name()); + return next; + }); + if(s.lens() != input_shape.lens()) + return std::nullopt; + return std::make_pair(std::move(result), s); + } + catch(const migraphx::exception&) + { + return std::nullopt; + } +} + namespace { const std::unordered_set& invertible_view_names() @@ -102,55 +143,6 @@ struct concat_optimizer return is_allocation(alias) and alias->get_shape().lens() == producer->get_shape().lens(); } - // Build the view ops that turn a slice of the super buffer (shaped like - // the concat input) into the producer's output shape by inverting the - // view chain. Returns the ops and the resulting view shape; nullopt when - // the inverse does not exist or would require a data movement. - static std::optional, shape>> invert_view_chain( - const std::vector& chain, const shape& producer_shape, const shape& slice_shape) - { - std::vector fwd; - std::transform( - chain.begin(), chain.end(), std::back_inserter(fwd), [](const operation& op) { - if(op.name() == "reshape_lazy") - return make_op("reshape", {{"dims", op.to_value().at("dims")}}); - return op; - }); - auto inv = shape_transform_descriptor::create(producer_shape.lens(), fwd).invert(); - if(inv.empty()) - return std::nullopt; - auto result = inv.generate(); - // Reshapes need to be lazy so the slice is aliased instead of copied at runtime - std::transform(result.begin(), result.end(), result.begin(), [](const operation& op) { - if(op.name() != "reshape") - return op; - return make_op("reshape_lazy", op.to_value()); - }); - try - { - auto s = std::accumulate(result.begin(), - result.end(), - slice_shape, - [](const shape& acc, const operation& op) { - std::vector inputs = {acc}; - if(op.output_alias(inputs).empty()) - MIGRAPHX_THROW("Not a view operator: " + op.name()); - auto next = op.compute_shape(inputs); - if(next.elements() != acc.elements()) - MIGRAPHX_THROW("Not an element-preserving view: " + - op.name()); - return next; - }); - if(s.lens() != producer_shape.lens()) - return std::nullopt; - return std::make_pair(std::move(result), s); - } - catch(const migraphx::exception&) - { - return std::nullopt; - } - } - struct plan { instruction_ref producer; diff --git a/src/include/migraphx/eliminate_concat.hpp b/src/include/migraphx/eliminate_concat.hpp index 6ee92cdb902..f1e02887f30 100644 --- a/src/include/migraphx/eliminate_concat.hpp +++ b/src/include/migraphx/eliminate_concat.hpp @@ -25,15 +25,27 @@ #define MIGRAPHX_GUARD_RTGLIB_ELIMINATE_CONCAT_HPP #include +#include +#include #include #include #include +#include +#include +#include namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { struct module; +// Build the view ops that re-index a view shaped like the output of a chain +// of view ops back into the shape of its input by inverting the chain. +// Returns the ops and the resulting view shape; nullopt when the inverse +// does not exist or would need a copy. +MIGRAPHX_EXPORT optional, shape>> invert_view_chain( + const std::vector& chain, const shape& input_shape, const shape& view_shape); + /** * Remove concat operators by having each operator write to different chunk of memory. */ diff --git a/src/targets/gpu/fuse_concat_past_present.cpp b/src/targets/gpu/fuse_concat_past_present.cpp index 14fbd18ae2c..294cdbc9d20 100644 --- a/src/targets/gpu/fuse_concat_past_present.cpp +++ b/src/targets/gpu/fuse_concat_past_present.cpp @@ -24,17 +24,17 @@ #include #include #include +#include #include #include #include #include #include -#include -#include #include #include #include #include +#include #include namespace migraphx { @@ -113,29 +113,17 @@ struct find_concat_past_present return {ks.type(), lens, ks.strides()}; } - // The producer must write exactly the elements of the cache slot in the - // same memory order. The views between it and the concat only re-index a - // standard tensor, so the slot is re-indexed back to the producer dims - // as a lazy reshape; nullopt when it cant be. - static optional - producer_slot(instruction_ref producer, instruction_ref cur, instruction_ref cache) + // The concat input must cover exactly one cache slot + static bool fills_slot(instruction_ref cur, instruction_ref cache) { - const auto& ps = producer->get_shape(); const auto& cs = cur->get_shape(); const auto& ks = cache->get_shape(); - if(not ps.standard() or not cs.standard() or not ks.standard()) - return nullopt; - if(ps.type() != cs.type() or ps.elements() != cs.elements()) - return nullopt; + if(not ks.standard() or cs.type() != ks.type()) + return false; if(cs.ndim() != 4 or ks.ndim() != 4) - return nullopt; - auto slot = slot_shape(cur, cache); - if(cs.lens() != slot.lens() or slot.lens()[seq_axis] > ks.lens()[seq_axis]) - return nullopt; - auto reshaped = reshape_dims(slot, ps.lens(), {.lazy = true}); - if(not reshaped.has_value()) - return nullopt; - return reshaped->to_static(); + return false; + return cs.lens() == slot_shape(cur, cache).lens() and + cs.lens()[seq_axis] <= ks.lens()[seq_axis]; } // One host copy of the sequence length serves every concat that reads it @@ -156,14 +144,27 @@ struct find_concat_past_present auto cur = ins->inputs()[0]; auto slk = ins->inputs()[1]; auto cache = ins->inputs()[2]; - auto pslot = producer_slot(producer, cur, cache); - if(not pslot.has_value()) + if(not fills_slot(cur, cache)) return; assert(not producer->inputs().empty()); auto alloc = producer->inputs().back(); if(alloc->name() != "allocate" or alloc->outputs().size() != 1) return; + // The producer writes in its own dims and layout, so the slot is + // re-indexed back through the views between it and the concat + auto path = get_input_path(cur); + std::vector chain(path.begin(), + std::find(path.begin(), path.end(), producer)); + std::vector chain_ops; + std::transform(chain.rbegin(), + chain.rend(), + std::back_inserter(chain_ops), + [](instruction_ref x) { return x->get_operator(); }); + auto inverse = invert_view_chain(chain_ops, producer->get_shape(), slot_shape(cur, cache)); + if(not inverse.has_value()) + return; + auto seq = cur->get_shape().lens()[seq_axis]; instruction_ref view; if(seq == 1) @@ -188,20 +189,19 @@ struct find_concat_past_present // The view is defined at the concat, so the producer and its view chain // move down to it; both are single-use so nothing in between needs them - auto path = get_input_path(cur); - std::vector chain(path.begin(), - std::find(path.begin(), path.end(), producer)); m.move_instruction(producer, ins); std::for_each(chain.rbegin(), chain.rend(), [&](auto x) { m.move_instruction(x, ins); }); assert(view->get_shape() == slot_shape(cur, cache)); - // The producer writes its own dims, so give it the slot re-indexed as a view - if(view->get_shape() != *pslot) - view = - m.insert_instruction(ins, make_op("reshape_lazy", {{"dims", pslot->lens()}}), view); - assert(view->get_shape() == *pslot); + view = std::accumulate(inverse->first.begin(), + inverse->first.end(), + view, + [&](instruction_ref start, const operation& op) { + return m.insert_instruction(ins, op, start); + }); + assert(view->get_shape() == inverse->second); auto v = producer->get_operator().to_value(); - v["output_shape"] = to_value(*pslot); + v["output_shape"] = to_value(inverse->second); auto inputs = producer->inputs(); inputs.back() = view; m.replace_instruction( diff --git a/test/eliminate_concat_test.cpp b/test/eliminate_concat_test.cpp index 1b69c01dc8d..cd41f32e9ef 100644 --- a/test/eliminate_concat_test.cpp +++ b/test/eliminate_concat_test.cpp @@ -831,4 +831,55 @@ TEST_CASE(input_used_non_packed_unsupported) EXPECT(m1.sort() == m2.sort()); } +TEST_CASE(invert_view_chain_reshape) +{ + // A strided view of {2, 4} re-indexed to the {2, 2, 2} input of a reshape + migraphx::shape view{migraphx::shape::float_type, {2, 4}, {8, 1}}; + auto inverse = + migraphx::invert_view_chain({migraphx::make_op("reshape_lazy", {{"dims", {2, 4}}})}, + {migraphx::shape::float_type, {2, 2, 2}}, + view); + EXPECT(inverse.has_value()); + EXPECT(inverse->first == std::vector{ + migraphx::make_op("reshape_lazy", {{"dims", {2, 2, 2}}})}); + EXPECT(inverse->second == migraphx::shape{migraphx::shape::float_type, {2, 2, 2}, {8, 2, 1}}); +} + +TEST_CASE(invert_view_chain_transpose_reshape) +{ + // {3, 2, 4} -> transpose -> {2, 3, 4} -> reshape -> {2, 12}, inverted on a strided view + migraphx::shape view{migraphx::shape::float_type, {2, 12}, {24, 1}}; + auto inverse = + migraphx::invert_view_chain({migraphx::make_op("transpose", {{"permutation", {1, 0, 2}}}), + migraphx::make_op("reshape_lazy", {{"dims", {2, 12}}})}, + {migraphx::shape::float_type, {3, 2, 4}}, + view); + EXPECT(inverse.has_value()); + EXPECT(inverse->first == std::vector{ + migraphx::make_op("reshape_lazy", {{"dims", {2, 3, 4}}}), + migraphx::make_op("transpose", {{"permutation", {1, 0, 2}}})}); + EXPECT(inverse->second == migraphx::shape{migraphx::shape::float_type, {3, 2, 4}, {4, 24, 1}}); +} + +TEST_CASE(invert_view_chain_not_lazy) +{ + // The view strides cannot merge the two axes, so the reshape would copy + migraphx::shape view{migraphx::shape::float_type, {2, 4}, {8, 1}}; + EXPECT(not migraphx::invert_view_chain({migraphx::make_op("reshape_lazy", {{"dims", {2, 4}}})}, + {migraphx::shape::float_type, {8}}, + view) + .has_value()); +} + +TEST_CASE(invert_view_chain_broadcast) +{ + // A broadcast duplicates elements so it has no inverse + migraphx::shape view{migraphx::shape::float_type, {2, 4}}; + EXPECT(not migraphx::invert_view_chain( + {migraphx::make_op("multibroadcast", {{"out_lens", {2, 4}}})}, + {migraphx::shape::float_type, {2, 1}}, + view) + .has_value()); +} + int main(int argc, const char* argv[]) { test::run(argc, argv); } diff --git a/test/gpu/fuse_concat_past_present.cpp b/test/gpu/fuse_concat_past_present.cpp index 015d27bb632..22b2217fef6 100644 --- a/test/gpu/fuse_concat_past_present.cpp +++ b/test/gpu/fuse_concat_past_present.cpp @@ -201,9 +201,9 @@ TEST_CASE(fuse_decode_reshaped_producer) EXPECT(p1.sort() == p2.sort()); } -// A transpose that reorders non-unit dims changes the memory order, so the -// producer cannot write the slot directly -TEST_CASE(skip_transposed_producer) +// pointwise in a transposed layout -> transpose -> concat_past_present; the +// slot is transposed back so the producer writes it in its own dims +TEST_CASE(fuse_transposed_producer) { migraphx::shape transposed_shape{migraphx::shape::half_type, {1, 4, 1, 2}}; migraphx::program p1; @@ -227,6 +227,112 @@ TEST_CASE(skip_transposed_producer) cache); mm->add_return({cpp}); } + run_pass(p1); + + migraphx::program p2; + { + auto* mm = p2.get_main_module(); + auto x = mm->add_parameter("x", transposed_shape); + auto y = mm->add_parameter("y", transposed_shape); + auto slk = mm->add_parameter("slk", index_shape()); + auto cache = mm->add_parameter("cache", cache_shape()); + auto* pm = create_pointwise_module(p2, "main:pointwise0", {x, y}, single_pointwise("mul")); + auto scalar = mm->add_instruction(migraphx::make_op("hip::load_scalar"), slk); + auto slot = + mm->add_instruction(migraphx::make_op("gpu::slice_at", {{"axis", 2}}), cache, scalar); + auto view = mm->add_instruction( + migraphx::make_op("transpose", {{"permutation", {0, 3, 2, 1}}}), slot); + migraphx::shape transposed_slot{migraphx::shape::half_type, {1, 4, 1, 2}, {64, 1, 4, 32}}; + auto pw = + mm->add_instruction(make_precompile_op(migraphx::make_op("pointwise"), transposed_slot), + {x, y, view}, + {pm}); + auto dep = mm->add_instruction(migraphx::make_op("identity"), cache, pw); + mm->add_return({dep}); + } + EXPECT(p1.sort() == p2.sort()); +} + +// Prefill rotary: the pointwise follows the seq-major qkv layout split as +// {b, seq, heads, 2, head/2} -> transpose -> reshape_lazy -> concat_past_present +TEST_CASE(fuse_prefill_transposed_producer) +{ + migraphx::shape seq_major{migraphx::shape::half_type, {1, 4, 2, 2, 2}}; + migraphx::program p1; + { + auto* mm = p1.get_main_module(); + auto x = mm->add_parameter("x", seq_major); + auto y = mm->add_parameter("y", seq_major); + auto slk = mm->add_parameter("slk", index_shape()); + auto cache = mm->add_parameter("cache", cache_shape()); + auto* pm = create_pointwise_module(p1, "main:pointwise0", {x, y}, single_pointwise("mul")); + auto alloc = mm->add_instruction( + migraphx::make_op("allocate", {{"shape", migraphx::to_value(seq_major)}})); + auto pw = mm->add_instruction(make_precompile_op("pointwise"), {x, y, alloc}, {pm}); + auto trp = mm->add_instruction( + migraphx::make_op("transpose", {{"permutation", {0, 2, 1, 3, 4}}}), pw); + auto rsp = mm->add_instruction( + migraphx::make_op("reshape_lazy", {{"dims", present_shape(4).lens()}}), trp); + auto cpp = mm->add_instruction( + make_precompile_op(migraphx::make_op("concat_past_present", {{"kv_num_heads", 2}}), + cache_shape()), + rsp, + slk, + cache); + mm->add_return({cpp}); + } + run_pass(p1); + + migraphx::program p2; + { + auto* mm = p2.get_main_module(); + auto x = mm->add_parameter("x", seq_major); + auto y = mm->add_parameter("y", seq_major); + mm->add_parameter("slk", index_shape()); + auto cache = mm->add_parameter("cache", cache_shape()); + auto* pm = create_pointwise_module(p2, "main:pointwise0", {x, y}, single_pointwise("mul")); + auto slot = mm->add_instruction( + migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {4}}}), cache); + auto split = mm->add_instruction( + migraphx::make_op("reshape_lazy", {{"dims", {1, 2, 4, 2, 2}}}), slot); + auto view = mm->add_instruction( + migraphx::make_op("transpose", {{"permutation", {0, 2, 1, 3, 4}}}), split); + migraphx::shape seq_major_slot{ + migraphx::shape::half_type, {1, 4, 2, 2, 2}, {64, 4, 32, 2, 1}}; + auto pw = mm->add_instruction( + make_precompile_op(migraphx::make_op("pointwise"), seq_major_slot), {x, y, view}, {pm}); + auto dep = mm->add_instruction(migraphx::make_op("identity"), cache, pw); + mm->add_return({dep}); + } + EXPECT(p1.sort() == p2.sort()); +} + +// The slot strides cannot merge the head and head-size axes into one, so the +// producer cannot write it through the reshape +TEST_CASE(skip_unmergeable_producer) +{ + migraphx::shape flat_shape{migraphx::shape::half_type, {1, 8}}; + migraphx::program p1; + { + auto* mm = p1.get_main_module(); + auto x = mm->add_parameter("x", flat_shape); + auto y = mm->add_parameter("y", flat_shape); + auto slk = mm->add_parameter("slk", index_shape()); + auto cache = mm->add_parameter("cache", cache_shape()); + auto* pm = create_pointwise_module(p1, "main:pointwise0", {x, y}, single_pointwise("mul")); + auto alloc = mm->add_instruction( + migraphx::make_op("allocate", {{"shape", migraphx::to_value(flat_shape)}})); + auto pw = mm->add_instruction(make_precompile_op("pointwise"), {x, y, alloc}, {pm}); + auto rsp = mm->add_instruction( + migraphx::make_op("reshape_lazy", {{"dims", present_shape(1).lens()}}), pw); + auto cpp = mm->add_instruction( + make_precompile_op(migraphx::make_op("concat_past_present", {{"kv_num_heads", 2}}), + cache_shape()), + rsp, + slk, + cache); + mm->add_return({cpp}); + } migraphx::program p2 = p1; run_pass(p1); From f73aad1ba3bd4ec51b9a828394e725631354af1b Mon Sep 17 00:00:00 2001 From: Paul Fultz II Date: Thu, 15 Oct 2026 18:11:30 -0500 Subject: [PATCH 3/6] Fix instruction insertion for producer in fuse_concat Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/targets/gpu/fuse_concat_past_present.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/targets/gpu/fuse_concat_past_present.cpp b/src/targets/gpu/fuse_concat_past_present.cpp index 294cdbc9d20..c39580c9fa9 100644 --- a/src/targets/gpu/fuse_concat_past_present.cpp +++ b/src/targets/gpu/fuse_concat_past_present.cpp @@ -197,7 +197,7 @@ struct find_concat_past_present inverse->first.end(), view, [&](instruction_ref start, const operation& op) { - return m.insert_instruction(ins, op, start); + return m.insert_instruction(producer, op, start); }); assert(view->get_shape() == inverse->second); auto v = producer->get_operator().to_value(); From 4242af994b0e6afbc3bdd5193a049479e8206935 Mon Sep 17 00:00:00 2001 From: Paul Date: Sat, 17 Oct 2026 09:20:09 -0500 Subject: [PATCH 4/6] Only for dnnl ops --- src/targets/cpu/dnnl.cpp | 16 +++- src/targets/cpu/include/migraphx/cpu/dnnl.hpp | 3 + src/targets/cpu/lowering.cpp | 80 ++++++++++++------- 3 files changed, 68 insertions(+), 31 deletions(-) diff --git a/src/targets/cpu/dnnl.cpp b/src/targets/cpu/dnnl.cpp index b8f4baffa7d..e03afc9f632 100644 --- a/src/targets/cpu/dnnl.cpp +++ b/src/targets/cpu/dnnl.cpp @@ -22,6 +22,7 @@ * THE SOFTWARE. */ #include +#include #if defined(__GNUC__) && __GNUC__ <= 5 namespace std { @@ -57,7 +58,7 @@ dnnl_context& get_dnnl_context() #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wswitch-enum" #endif -dnnl::memory::data_type to_dnnl_memory_data_type(shape::type_t t) +static std::optional find_dnnl_memory_data_type(shape::type_t t) { using dt = dnnl::memory::data_type; using st = shape::type_t; @@ -68,15 +69,24 @@ dnnl::memory::data_type to_dnnl_memory_data_type(shape::type_t t) case st::int32_type: return dt::s32; case st::int8_type: return dt::s8; case st::uint8_type: return dt::u8; - case st::fp8e4m3fnuz_type: MIGRAPHX_THROW("fp8e4m3fnuz unsupported in DNNL"); case st::bf16_type: return dt::bf16; - default: MIGRAPHX_THROW("Unsupported data type"); + default: return std::nullopt; } } #ifdef __clang__ #pragma clang diagnostic pop #endif +bool is_dnnl_supported_type(shape::type_t t) { return find_dnnl_memory_data_type(t).has_value(); } + +dnnl::memory::data_type to_dnnl_memory_data_type(shape::type_t t) +{ + auto dt = find_dnnl_memory_data_type(t); + if(not dt.has_value()) + MIGRAPHX_THROW("Unsupported data type in DNNL: " + shape::name(t)); + return *dt; +} + dnnl::memory::format_tag to_dnnl_memory_format_tag(std::size_t n) { switch(n) diff --git a/src/targets/cpu/include/migraphx/cpu/dnnl.hpp b/src/targets/cpu/include/migraphx/cpu/dnnl.hpp index 9256454dfbe..26c4ad9b7f8 100644 --- a/src/targets/cpu/include/migraphx/cpu/dnnl.hpp +++ b/src/targets/cpu/include/migraphx/cpu/dnnl.hpp @@ -61,6 +61,9 @@ dnnl_context& get_dnnl_context(); dnnl::memory::data_type to_dnnl_memory_data_type(shape::type_t t); +// oneDNN has no bool, int64, double or fp8 memory types +bool is_dnnl_supported_type(shape::type_t t); + dnnl::memory::format_tag to_dnnl_memory_format_tag(std::size_t n); template diff --git a/src/targets/cpu/lowering.cpp b/src/targets/cpu/lowering.cpp index fdea6202a4b..06719732a1e 100644 --- a/src/targets/cpu/lowering.cpp +++ b/src/targets/cpu/lowering.cpp @@ -49,6 +49,7 @@ #include #include #include +#include #include #include #include @@ -303,6 +304,17 @@ struct cpu_rnn_var_sl_last_output }; MIGRAPHX_REGISTER_OP(cpu_rnn_var_sl_last_output) +static bool has_dnnl_type(instruction_ref ins) +{ + return is_dnnl_supported_type(ins->get_shape().type()); +} + +// oneDNN cannot run an op whose output or inputs use a type it lacks (bool, int64, fp8, ...) +static bool has_dnnl_types(instruction_ref ins, const std::vector& inputs) +{ + return has_dnnl_type(ins) and std::all_of(inputs.begin(), inputs.end(), &has_dnnl_type); +} + struct cpu_apply { module* modl; @@ -319,6 +331,13 @@ struct cpu_apply }); } + void extend_dnnl_op(const std::string& op_name, const std::string& dnnl_name) + { + apply_map.emplace(op_name, [=](instruction_ref ins) { + return replace_dnnl(ins, make_op(dnnl_name, ins->get_operator().to_value())); + }); + } + void extend_dnnl_algos(const std::string& dnnl_name, const std::vector>& algos) { @@ -332,7 +351,7 @@ struct cpu_apply return ins; v["algo"] = algo; auto op = make_op(dnnl_name, v); - return replace(ins, op); + return replace_dnnl(ins, op); }); } } @@ -347,8 +366,7 @@ struct cpu_apply bind_inputs.end(), std::back_inserter(inputs), [&](const auto& s) { return r.instructions[s]; }); - inputs.push_back(this->insert_allocation(ins, ins->get_shape())); - modl->replace_instruction(ins, op, inputs); + this->replace_dnnl(ins, op, inputs); }); } @@ -381,17 +399,17 @@ struct cpu_apply {"reduce_min", "reduction_min"}, {"reduce_sum", "reduction_sum"}, }); - extend_op("concat", "dnnl::concat"); - extend_op("contiguous", "dnnl::reorder"); - extend_op("convolution", "dnnl::convolution"); + extend_dnnl_op("concat", "dnnl::concat"); + extend_dnnl_op("contiguous", "dnnl::reorder"); + extend_dnnl_op("convolution", "dnnl::convolution"); #ifndef MIGRAPHX_ENABLE_ZENDNN - extend_op("dot", "dnnl::dot"); + extend_dnnl_op("dot", "dnnl::dot"); #endif extend_op("erf", "cpu::erf"); extend_op("gather", "cpu::gather"); - extend_op("logsoftmax", "dnnl::logsoftmax"); - extend_op("lrn", "dnnl::lrn"); - extend_op("softmax", "dnnl::softmax"); + extend_dnnl_op("logsoftmax", "dnnl::logsoftmax"); + extend_dnnl_op("lrn", "dnnl::lrn"); + extend_dnnl_op("softmax", "dnnl::softmax"); extend_op("im2col", "cpu::im2col", false); extend_op("pad", "cpu::pad", false); @@ -413,12 +431,6 @@ struct cpu_apply // Apply these operators first so the inputs can be const folded for(auto it : iterator_for(*modl)) { - // skip lowering if input has fp8 as one of the inputs since oneDNN doesn't have fp8 - // supported yet. - if(std::any_of(it->inputs().begin(), it->inputs().end(), [](const auto& i) { - return contains(fp8_types{}.get(), i->get_shape().type()); - })) - continue; if(it->name() == "pow") { apply_pow(it); @@ -426,12 +438,6 @@ struct cpu_apply } for(auto it : iterator_for(*modl)) { - // skip lowering if input has fp8 as one of the inputs since oneDNN doesn't have fp8 - // supported yet. - if(std::any_of(it->inputs().begin(), it->inputs().end(), [](const auto& i) { - return contains(fp8_types{}.get(), i->get_shape().type()); - })) - continue; if(it->name() == "pooling") { apply_pooling(it); @@ -458,10 +464,11 @@ struct cpu_apply auto beta = read_scalar(ins->inputs()[1]); if(beta.empty()) return ins; - return replace(ins, - make_op("dnnl::eltwise", - {{"algo", "eltwise_pow"}, {"alpha", 1.0}, {"beta", beta.front()}}), - {ins->inputs().front()}); + return replace_dnnl( + ins, + make_op("dnnl::eltwise", + {{"algo", "eltwise_pow"}, {"alpha", 1.0}, {"beta", beta.front()}}), + {ins->inputs().front()}); } // TODO: update lowering to run the reference @@ -481,7 +488,7 @@ struct cpu_apply if(has_op("dnnl::pooling") and ins->get_shape().type() == shape::type_t::float_type and not v["ceil_mode"].to() and v["mode"].to() != op::pooling_mode::lpnorm) - return replace(ins, make_op("dnnl::pooling", op.to_value())); + return replace_dnnl(ins, make_op("dnnl::pooling", op.to_value())); return ins; } /* @@ -490,6 +497,8 @@ struct cpu_apply */ instruction_ref apply_reshape(instruction_ref ins) const { + if(not has_dnnl_types(ins, ins->inputs())) + return ins; std::vector before_contiguous_args = ins->inputs(); auto before_alloc = insert_allocation(ins, before_contiguous_args.front()->get_shape().as_standard()); @@ -518,7 +527,7 @@ struct cpu_apply auto v = op.to_value(); if(has_op("dnnl::convolution_backwards") and v["group"].to() == 1) { - return replace(ins, make_op("dnnl::convolution_backwards", op.to_value())); + return replace_dnnl(ins, make_op("dnnl::convolution_backwards", op.to_value())); } return ins; } @@ -548,6 +557,21 @@ struct cpu_apply return modl->replace_instruction(ins, op, inputs); } + // Leave the op for the reference implementation when oneDNN lacks one of its types + instruction_ref replace_dnnl(instruction_ref ins, const operation& op) const + { + return replace_dnnl(ins, op, ins->inputs()); + } + + instruction_ref replace_dnnl(instruction_ref ins, + const operation& op, + const std::vector& inputs) const + { + if(not has_dnnl_types(ins, inputs)) + return ins; + return replace(ins, op, inputs); + } + instruction_ref insert_allocation(instruction_ref ins, const shape& s) const { return modl->insert_instruction(ins, make_op("allocate", {{"shape", to_value(s)}})); From a6121690318b63e5b0f542767cedeebcbecdf791 Mon Sep 17 00:00:00 2001 From: Paul Date: Sat, 17 Oct 2026 10:33:13 -0500 Subject: [PATCH 5/6] Update year --- src/targets/cpu/dnnl.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/targets/cpu/dnnl.cpp b/src/targets/cpu/dnnl.cpp index e03afc9f632..593d16bf533 100644 --- a/src/targets/cpu/dnnl.cpp +++ b/src/targets/cpu/dnnl.cpp @@ -1,7 +1,7 @@ /* * The MIT License (MIT) * - * Copyright (c) 2015-2025 Advanced Micro Devices, Inc. All rights reserved. + * 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 From a454aabe86c3c5d053830fd90d2027739fb8eaf7 Mon Sep 17 00:00:00 2001 From: Paul Date: Sat, 17 Oct 2026 17:12:21 -0500 Subject: [PATCH 6/6] Handle reshape_lazy case --- src/targets/gpu/fuse_concat_past_present.cpp | 16 +++--- test/gpu/fuse_concat_past_present.cpp | 52 ++++++++++++++++++++ 2 files changed, 59 insertions(+), 9 deletions(-) diff --git a/src/targets/gpu/fuse_concat_past_present.cpp b/src/targets/gpu/fuse_concat_past_present.cpp index c39580c9fa9..7db2be02731 100644 --- a/src/targets/gpu/fuse_concat_past_present.cpp +++ b/src/targets/gpu/fuse_concat_past_present.cpp @@ -187,26 +187,24 @@ struct find_concat_past_present cache); } - // The view is defined at the concat, so the producer and its view chain - // move down to it; both are single-use so nothing in between needs them - m.move_instruction(producer, ins); - std::for_each(chain.rbegin(), chain.rend(), [&](auto x) { m.move_instruction(x, ins); }); - + // A retargeted copy of the producer is inserted at the concat rather + // than rewriting it in place: the new strides could make the original + // view chain fail shape recomputation before it becomes dead assert(view->get_shape() == slot_shape(cur, cache)); view = std::accumulate(inverse->first.begin(), inverse->first.end(), view, [&](instruction_ref start, const operation& op) { - return m.insert_instruction(producer, op, start); + return m.insert_instruction(ins, op, start); }); assert(view->get_shape() == inverse->second); auto v = producer->get_operator().to_value(); v["output_shape"] = to_value(inverse->second); auto inputs = producer->inputs(); inputs.back() = view; - m.replace_instruction( - producer, make_op("gpu::precompile_op", v), inputs, producer->module_inputs()); - m.replace_instruction(ins, make_op("identity"), {cache, producer}); + auto fused = m.insert_instruction( + ins, make_op("gpu::precompile_op", v), inputs, producer->module_inputs()); + m.replace_instruction(ins, make_op("identity"), {cache, fused}); } }; diff --git a/test/gpu/fuse_concat_past_present.cpp b/test/gpu/fuse_concat_past_present.cpp index 22b2217fef6..1a96c74020b 100644 --- a/test/gpu/fuse_concat_past_present.cpp +++ b/test/gpu/fuse_concat_past_present.cpp @@ -201,6 +201,58 @@ TEST_CASE(fuse_decode_reshaped_producer) EXPECT(p1.sort() == p2.sort()); } +// pointwise over {h, 2, head/2} -> reshape_lazy -> concat_past_present; the +// slot gives the producer a non-packed leading stride that the original +// reshape cannot alias, so the retargeted producer must not recompute it +TEST_CASE(fuse_decode_reshaped_producer_strided) +{ + migraphx::shape split_shape{migraphx::shape::half_type, {2, 2, 2}}; + migraphx::program p1; + { + auto* mm = p1.get_main_module(); + auto x = mm->add_parameter("x", split_shape); + auto y = mm->add_parameter("y", split_shape); + auto slk = mm->add_parameter("slk", index_shape()); + auto cache = mm->add_parameter("cache", cache_shape()); + auto* pm = create_pointwise_module(p1, "main:pointwise0", {x, y}, single_pointwise("mul")); + auto alloc = mm->add_instruction( + migraphx::make_op("allocate", {{"shape", migraphx::to_value(split_shape)}})); + auto pw = mm->add_instruction(make_precompile_op("pointwise"), {x, y, alloc}, {pm}); + auto rsp = mm->add_instruction( + migraphx::make_op("reshape_lazy", {{"dims", present_shape(1).lens()}}), pw); + auto cpp = mm->add_instruction( + make_precompile_op(migraphx::make_op("concat_past_present", {{"kv_num_heads", 2}}), + cache_shape()), + rsp, + slk, + cache); + mm->add_return({cpp}); + } + run_pass(p1); + + migraphx::program p2; + { + auto* mm = p2.get_main_module(); + auto x = mm->add_parameter("x", split_shape); + auto y = mm->add_parameter("y", split_shape); + auto slk = mm->add_parameter("slk", index_shape()); + auto cache = mm->add_parameter("cache", cache_shape()); + auto* pm = create_pointwise_module(p2, "main:pointwise0", {x, y}, single_pointwise("mul")); + auto scalar = mm->add_instruction(migraphx::make_op("hip::load_scalar"), slk); + auto slot = + mm->add_instruction(migraphx::make_op("gpu::slice_at", {{"axis", 2}}), cache, scalar); + auto split = mm->add_instruction( + migraphx::make_op("reshape_lazy", {{"dims", {1, 2, 1, 2, 2}}}), slot); + auto view = mm->add_instruction(migraphx::make_op("squeeze", {{"axes", {0, 2}}}), split); + migraphx::shape split_slot{migraphx::shape::half_type, {2, 2, 2}, {32, 2, 1}}; + auto pw = mm->add_instruction( + make_precompile_op(migraphx::make_op("pointwise"), split_slot), {x, y, view}, {pm}); + auto dep = mm->add_instruction(migraphx::make_op("identity"), cache, pw); + mm->add_return({dep}); + } + EXPECT(p1.sort() == p2.sort()); +} + // pointwise in a transposed layout -> transpose -> concat_past_present; the // slot is transposed back so the producer writes it in its own dims TEST_CASE(fuse_transposed_producer)