Fuse the rotate-half concat into attention and let the K rotary write the KV cache directly

src/eliminate_concat.cpp

@@ -40,6 +40,47 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
+optional<std::pair<std::vector<operation>, shape>> invert_view_chain(
+    const std::vector<operation>& chain, const shape& input_shape, const shape& view_shape)
+{
+    std::vector<operation> 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<shape> 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<std::string>& 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<std::pair<std::vector<operation>, shape>> invert_view_chain(
-        const std::vector<operation>& chain, const shape& producer_shape, const shape& slice_shape)
-    {
-        std::vector<operation> 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<shape> 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;

src/include/migraphx/eliminate_concat.hpp

@@ -25,15 +25,27 @@
 #define MIGRAPHX_GUARD_RTGLIB_ELIMINATE_CONCAT_HPP
 
 #include <string>
+#include <utility>
+#include <vector>
 #include <migraphx/instruction_ref.hpp>
 #include <migraphx/concat_opt.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/operation.hpp>
+#include <migraphx/optional.hpp>
+#include <migraphx/shape.hpp>
 
 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<std::pair<std::vector<operation>, shape>> invert_view_chain(
+    const std::vector<operation>& chain, const shape& input_shape, const shape& view_shape);
+
 /**
  * Remove concat operators by having each operator write to different chunk of memory.
  */

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<op::concat>(ins->normalized_operator()).axis;
+    auto hi   = any_cast<op::slice>(inputs[0]->normalized_operator());
+    auto lo   = any_cast<op::slice>(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<op::concat>(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<std::size_t> 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<instruction_ref> 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{},

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
@@ -22,6 +22,7 @@
  * THE SOFTWARE.
  */
 #include <migraphx/cpu/dnnl.hpp>
+#include <optional>
 
 #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<dnnl::memory::data_type> 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)

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 <class R>