Remove concat->reshapes->slice

CHANGELOG.md

@@ -47,6 +47,7 @@ Full documentation for MIGraphX is available at
 * Added a `dyn_slice` operator, `dyn_slice(data, starts, ends)`, whose symbolic `starts`/`ends` attributes describe the run-time bound inputs so a data-dependent slice keeps a symbolic output shape; the axes are an attribute since they must be known when the shape is computed (#5112).
 * Added symbolic normalization of operator attributes holding `sym::expr`, clamping each value against its axis length symbolically (#5148).
 * Added symbolic evaluation of tensor values to preserve ONNX shape-tensor expressions through arithmetic and dynamic shape consumers, including `Reshape`, `Range`, `Slice`, `Expand`, `ConstantOfShape`, and `Trilu` (#5148).
+* Added a `find_slice_reshaped_concat` matcher to `simplify_reshapes` that forwards a slice reading exactly one segment of a concat through intervening reshape/transpose view ops, removing the concat entirely (#5183).
 * Added find_concat_same_broadcast matcher to convert concat of identical broadcasts into a single multibroadcast to reduce hipCopy() (#5179).
 
 ### Changed

src/fuse_attention.cpp

@@ -896,11 +896,9 @@ struct find_kv_cache_attention
 
         auto keys = match::opaque(
             match::skip(match::name(skip_set))(match::name("concat_past_present")).bind("pres_k"));
-        auto keys_transpose = match::opaque(match::name("transpose")(match::arg(0)(keys)));
-        auto k_transpose    = match::opaque(match::skip(match::name(skip_set))(keys_transpose));
-        auto queries = match::name("slice");
-        auto gemm1 =
-            match::opaque(match::name("dot")(match::arg(0)(queries), match::arg(1)(k_transpose)));
+        auto keys_transpose  = match::opaque(match::name("transpose")(match::arg(0)(keys)));
+        auto k_transpose     = match::opaque(match::skip(match::name(skip_set))(keys_transpose));
+        auto gemm1           = match::opaque(match::name("dot")(match::arg(1)(k_transpose)));
         auto gemm1_maybe_cvt = match::opaque(match::skip(match::name("convert"))(gemm1));
         auto scale    = match::opaque(match::name("mul")(match::any_arg(0, 1)(gemm1_maybe_cvt)));
         auto constant = match::opaque(match::is_constant());

src/include/migraphx/shape_transform_descriptor.hpp

@@ -31,6 +31,7 @@
 #include <cstdint>
 #include <iosfwd>
 #include <set>
+#include <utility>
 #include <vector>
 
 namespace migraphx {
@@ -88,6 +89,15 @@ struct MIGRAPHX_EXPORT shape_transform_descriptor
     bool apply_transpose(const std::vector<std::int64_t>& permutation);
     bool apply_broadcast(const std::vector<std::size_t>& out_lens,
                          optional<std::size_t> axis = nullopt);
+    // Restrict the source axis to the range selected by slicing the output
+    // dimensions given by slice_axes/starts/ends. Returns the selected
+    // [start, end) range of the source axis, or nullopt when the slice does
+    // not select one contiguous range of only this axis.
+    optional<std::pair<std::size_t, std::size_t>>
+    slice_axis(std::size_t axis,
+               const std::vector<std::size_t>& slice_axes,
+               const std::vector<std::size_t>& starts,
+               const std::vector<std::size_t>& ends);
     void simplify();
     std::size_t elements() const;
     std::vector<operation> generate(const std::vector<std::size_t>& input_dims = {},

src/shape_transform_descriptor.cpp

@@ -816,7 +816,7 @@ bool shape_transform_descriptor::apply(const std::vector<operation>& ops)
     for(const auto& op : ops)
     {
         auto v = op.to_value();
-        if(contains({"reshape", "squeeze", "unsqueeze", "flatten"}, op.name()))
+        if(contains({"reshape", "reshape_lazy", "squeeze", "unsqueeze", "flatten"}, op.name()))
         {
             dims = compute_dims(op, dims);
             if(not apply_reshape(dims))
@@ -988,6 +988,153 @@ bool shape_transform_descriptor::apply_broadcast(const std::vector<std::size_t>&
     return true;
 }
 
+// The range of an axis subdimension selected by a slice of the output
+struct dimension_slice
+{
+    dimension::sub* sub;
+    std::size_t start;
+    std::size_t end;
+
+    bool full() const { return start == 0 and end == sub->len; }
+    std::size_t size() const { return end - start; }
+};
+
+static bool sub_from_axis(const dimension::sub& s, std::size_t axis)
+{
+    return not s.origin_axis().empty() and s.origin_axis().front() == axis;
+}
+
+// Unit subdimensions carry no element order, so if the wider subdimensions
+// are already in output order, renumber all split indices to output order to
+// avoid generating a gratuitous transpose
+static void renumber_in_output_order(const std::vector<dimension_slice>& dst_slices,
+                                     std::size_t axis)
+{
+    std::vector<dimension::sub*> wider;
+    transform_if(
+        dst_slices.begin(),
+        dst_slices.end(),
+        std::back_inserter(wider),
+        [](const dimension_slice& s) { return s.sub->len > 1; },
+        [](const dimension_slice& s) { return s.sub; });
+    if(not std::is_sorted(
+           wider.begin(), wider.end(), by(std::less<>{}, [](const auto* s) -> const auto& {
+               return s->origin_axis();
+           })))
+        return;
+    for(std::size_t i : range(dst_slices.size()))
+        set_origin_axis(*dst_slices[i].sub, {axis, i});
+}
+
+// Map each subdimension of the axis to the range the slice selects from it.
+// Returns nullopt when a sliced output dimension does not map entirely to a
+// subdimension of the axis.
+static optional<std::vector<dimension_slice>>
+collect_dimension_slices(std::vector<dimension>& dimensions,
+                         std::size_t axis,
+                         const std::vector<std::size_t>& slice_axes,
+                         const std::vector<std::size_t>& starts,
+                         const std::vector<std::size_t>& ends)
+{
+    std::vector<dimension_slice> result;
+    for(auto i : range(dimensions.size()))
+    {
+        auto& dim = dimensions[i];
+        auto it   = std::find(slice_axes.begin(), slice_axes.end(), i);
+        if(it == slice_axes.end())
+        {
+            // Unsliced output axes keep the full range of each subdimension
+            transform_if(
+                dim.subdimensions.begin(),
+                dim.subdimensions.end(),
+                std::back_inserter(result),
+                [&](const auto& s) { return sub_from_axis(s, axis); },
+                [](auto& s) { return dimension_slice{&s, 0, s.len}; });
+            continue;
+        }
+        auto sit =
+            std::find_if(dim.subdimensions.begin(), dim.subdimensions.end(), [&](const auto& s) {
+                return sub_from_axis(s, axis) and s.len == dim.len();
+            });
+        if(sit == dim.subdimensions.end())
+            return nullopt;
+        auto k = std::distance(slice_axes.begin(), it);
+        if(starts[k] >= ends[k])
+            return nullopt;
+        // The remaining subdimensions of the axis have a length of 1, so they
+        // keep their full range
+        transform_if(
+            dim.subdimensions.begin(),
+            dim.subdimensions.end(),
+            std::back_inserter(result),
+            [&](const auto& s) { return sub_from_axis(s, axis); },
+            [&](auto& s) -> dimension_slice {
+                if(&s == &*sit)
+                    return {&s, starts[k], ends[k]};
+                return {&s, 0, s.len};
+            });
+    }
+    return result;
+}
+
+optional<std::pair<std::size_t, std::size_t>>
+shape_transform_descriptor::slice_axis(std::size_t axis,
+                                       const std::vector<std::size_t>& slice_axes,
+                                       const std::vector<std::size_t>& starts,
+                                       const std::vector<std::size_t>& ends)
+{
+    assert(slice_axes.size() == starts.size() and slice_axes.size() == ends.size());
+    auto dst_slices = collect_dimension_slices(dimensions, axis, slice_axes, starts, ends);
+    if(not dst_slices.has_value())
+        return nullopt;
+    // Broadcasted subdimensions do not map back to a source range
+    if(std::any_of(dst_slices->begin(), dst_slices->end(), [](const dimension_slice& s) {
+           return s.sub->has_hidden_axis();
+       }))
+        return nullopt;
+
+    // The subdimensions ordered by their split lineage form a mixed-radix
+    // decomposition of the axis, outermost first
+    auto sub_slices = *dst_slices;
+    std::sort(sub_slices.begin(),
+              sub_slices.end(),
+              by(std::less<>{},
+                 [](const dimension_slice& s) -> const auto& { return s.sub->origin_axis(); }));
+
+    // The selected ranges must form one contiguous range along the axis:
+    // every subdimension outside the innermost restricted one must select a
+    // single index, and everything inside it the full range
+    auto rit = std::find_if(sub_slices.rbegin(), sub_slices.rend(), [](const dimension_slice& s) {
+        return not s.full();
+    });
+    if(rit != sub_slices.rend() and
+       not std::all_of(sub_slices.begin(), std::prev(rit.base()), [](const dimension_slice& s) {
+           return s.size() == 1;
+       }))
+        return nullopt;
+    auto [start, total] = std::accumulate(
+        sub_slices.rbegin(),
+        sub_slices.rend(),
+        std::make_pair(std::size_t{0}, std::size_t{1}),
+        [](auto acc, const dimension_slice& s) {
+            return std::make_pair(acc.first + s.start * acc.second, acc.second * s.sub->len);
+        });
+    auto end = start + transform_accumulate(sub_slices.begin(),
+                                            sub_slices.end(),
+                                            std::size_t{1},
+                                            std::multiplies<>{},
+                                            [](const dimension_slice& s) { return s.size(); });
+    assert(end <= total);
+
+    // Restrict each subdimension to its selected range so the descriptor now
+    // maps the sliced source axis to the sliced output
+    std::for_each(sub_slices.begin(), sub_slices.end(), [](const dimension_slice& s) {
+        s.sub->len = s.size();
+    });
+    renumber_in_output_order(*dst_slices, axis);
+    return std::make_pair(start, end);
+}
+
 // Remove subdimensions of 1
 static void remove_1_sub_dims(std::vector<dimension::sub>& subdimensions)
 {

src/simplify_reshapes.cpp

@@ -816,6 +816,98 @@ struct find_concat_multibroadcasts
     }
 };
 
+// Offsets of each concat input along the concat axis, from 0 up to the total
+// length of the axis
+std::vector<std::size_t> concat_offsets(const std::vector<instruction_ref>& inputs,
+                                        std::size_t axis)
+{
+    std::vector<std::size_t> result = {0};
+    transform_partial_sum(inputs.begin(),
+                          inputs.end(),
+                          std::back_inserter(result),
+                          std::plus<>{},
+                          [&](instruction_ref ins) { return ins->get_shape().lens()[axis]; });
+    return result;
+}
+
+// Forward a slice that reads exactly one segment of a concat through the view
+// ops in between, so the concat can be removed entirely. This shows up when a
+// model repacks tensors (such as qkv) that an op decomposition then re-slices.
+struct find_slice_reshaped_concat
+{
+    static const auto& view_ops()
+    {
+        static const std::unordered_set<std::string> names = {
+            "reshape", "reshape_lazy", "squeeze", "unsqueeze", "flatten", "transpose"};
+        return names;
+    }
+
+    auto matcher() const
+    {
+        // Multi-input slices take runtime starts/ends and are not normalized
+        return match::name("slice")(match::nargs(1),
+                                    match::arg(0)(match::skip(match::name(view_ops()))(
+                                        match::name("concat").bind("concat"))));
+    }
+
+    void apply(module& m, const match::matcher_result& mr) const
+    {
+        auto slice_ins  = mr.result;
+        auto concat_ins = mr.instructions["concat"];
+        if(concat_ins->get_shape().dynamic())
+            return;
+        std::vector<operation> ops;
+        auto x = slice_ins->inputs().front();
+        while(x != concat_ins)
+        {
+            if(not contains(view_ops(), x->name()))
+                return;
+            ops.push_back(x->get_operator());
+            x = x->inputs().front();
+        }
+        // A direct slice of a concat is handled by find_concat_slice
+        if(ops.empty())
+            return;
+        std::reverse(ops.begin(), ops.end());
+
+        const auto& clens = concat_ins->get_shape().lens();
+        std::size_t axis  = any_cast<op::concat>(concat_ins->normalized_operator()).axis;
+        assert(axis < clens.size());
+
+        // Track the element mapping of the view chain with a descriptor
+        auto desc = shape_transform_descriptor::create(clens, ops);
+        if(desc.empty())
+            return;
+
+        // Restrict the concat axis to the range of elements the slice selects
+        auto slice_val = slice_ins->normalized_operator().to_value();
+        auto selected  = desc.slice_axis(axis,
+                                         slice_val["axes"].to_vector<std::size_t>(),
+                                         slice_val["starts"].to_vector<std::size_t>(),
+                                         slice_val["ends"].to_vector<std::size_t>());
+        if(not selected.has_value())
+            return;
+        auto [lo, hi] = *selected;
+
+        // The selected range must be exactly one segment of the concat
+        const auto& inputs = concat_ins->inputs();
+        auto prefix        = concat_offsets(inputs, axis);
+        auto it            = std::find(prefix.begin(), std::prev(prefix.end()), lo);
+        if(it == std::prev(prefix.end()))
+            return;
+        auto idx = std::distance(prefix.begin(), it);
+        if(prefix[idx + 1] != hi)
+            return;
+        auto seg = inputs[idx];
+
+        if(desc.lens() != slice_ins->get_shape().lens())
+            return;
+        desc.simplify();
+        auto y = insert_ops(m, slice_ins, desc.generate(), seg);
+        m.replace_instruction(slice_ins, y);
+    }
+};
+
 struct find_concat_slice
 {
     auto matcher() const
@@ -852,11 +944,7 @@ struct find_concat_slice
         {
             return;
         }
-        std::vector<size_t> prefix_scan = {0};
-        std::transform(
-            inputs.begin(), inputs.end(), std::back_inserter(prefix_scan), [&](const auto& i) {
-                return prefix_scan.back() + i->get_shape().lens()[concat_axis];
-            });
+        auto prefix_scan = concat_offsets(inputs, concat_axis);
         for(const auto& sins : slice_candidates)
         {
             auto sop           = any_cast<op::slice>(sins->get_operator());
@@ -2162,6 +2250,7 @@ void simplify_reshapes::apply(module& m) const
     if(enable_gather_rewrite)
         match::find_matches(m, find_gather{});
     m.repeat_while_changes(depth, [&] {
+        match::find_matches(m, find_slice_reshaped_concat{});
         match::find_matches(m,
                             find_nop_reshapes{},
                             find_flatten{},