Rewrite gather to transpose/reshape/broadcast/slice

src/common_dims.cpp

@@ -97,25 +97,66 @@ static bool compute_common_dim(std::vector<std::size_t>& cd_dims,
     assert(state1.get() < state2.get());
     auto d2    = state2.get();
     auto dims  = state1.dims_for(d2);
-    auto n     = elements(dims);
     auto naxes = distance(dims);
+
     if(naxes == 0)
         return false;
+
+    // Check if state1 has a remainder from previous split
+    bool has_remainder = (state1.rem != 1);
+
+    // Compute the product of dimensions, adjusting for remainder if needed
+    auto n = elements(dims);
+    if(has_remainder and naxes > 0)
+    {
+        n = n / *dims.begin() * (*dims.begin() / state1.rem);
+    }
+
     // If not divisible then we can't compute a common dim
     if((d2 % n) != 0)
         return false;
+
     auto rem = d2 / n;
-    state1.add_multi_axes(naxes, cd_dims.size());
-    state2.add_axes(rem == 1 ? naxes : naxes + 1, cd_dims.size());
+    auto start_pos = cd_dims.size();
 
+    // Add axes mappings
+    if(has_remainder)
+    {
+        // state1: dimension was split, keep axes together
+        state1.add_axes(naxes, start_pos);
+        // state2: axes should include the previous remainder dimension
+        state2.add_axes(rem == 1 ? naxes : naxes + 1, start_pos - 1);
+    }
+    else
+    {
+        // state1: separate axes for each dimension
+        state1.add_multi_axes(naxes, start_pos);
+        // state2: normal axes mapping
+        state2.add_axes(rem == 1 ? naxes : naxes + 1, start_pos);
+    }
+
+    // Add dimensions to cd_dims
+    if(has_remainder and naxes > 0)
+    {
+        // Adjust the first dimension by dividing by the remainder
+        cd_dims.push_back(*dims.begin() / state1.rem);
+        cd_dims.insert(cd_dims.end(), std::next(dims.begin()), dims.end());
+    }
+    else
+    {
+        cd_dims.insert(cd_dims.end(), dims.begin(), dims.end());
+    }
+
+    // Add remainder dimension if needed
+    if(rem != 1)
+        cd_dims.push_back(rem);
+
+    // Update states
     state1.rem = rem;
     state2.rem = 1;
-
-    cd_dims.insert(cd_dims.end(), dims.begin(), dims.end());
-    if(state1.rem != 1)
-        cd_dims.push_back(state1.rem);
-    state1.next(distance(dims));
+    state1.next(naxes);
     state2.next();
+
     return true;
 }
 
@@ -152,6 +193,22 @@ common_dims common_dims::compute(const std::vector<std::size_t>& dims1,
                 return {};
         }
     }
+
+    // Handle case where one state has a remainder that equals the next dimension
+    // In this case, the dimension was already added as a remainder, we just need the axes mapping
+    auto handle_remaining_dimension = [&cd](common_dim_state& state) {
+        if(not state.is_end() and state.rem != 1 and state.get() == 1)
+        {
+            // The remainder already added to cd_dims matches this dimension
+            // Add a single axes mapping
+            state.axes_map->push_back({cd.dims.size() - 1});
+            state.next();
+        }
+    };
+
+    handle_remaining_dimension(state1);
+    handle_remaining_dimension(state2);
+
     assert(elements(dims1) == elements(cd.dims));
     return cd;
 }

src/include/migraphx/fuse_attention.hpp

@@ -37,7 +37,7 @@ struct module_pass_manager;
 
 struct MIGRAPHX_EXPORT fuse_attention
 {
-    bool attn_enabled = false;
+    bool attn_enabled                         = false;
     bool flash_decoding_enabled               = false;
     std::size_t flash_decoding_num_splits     = 0;
     std::size_t flash_decoding_threshold      = 32;

src/include/migraphx/shape_transform_descriptor.hpp

@@ -27,6 +27,7 @@
 
 #include <migraphx/config.hpp>
 #include <migraphx/optional.hpp>
+#include <migraphx/shape.hpp>
 #include <cstdint>
 #include <iosfwd>
 #include <set>
@@ -106,6 +107,8 @@ struct MIGRAPHX_EXPORT shape_transform_descriptor
     std::vector<std::vector<std::size_t>> common_axes_map_from_src() const;
     std::vector<std::vector<std::size_t>> common_axes_map_from_dst() const;
 
+    std::vector<std::size_t> get_dst_axes_from_src(std::size_t axis) const;
+
     bool empty() const;
     std::vector<std::size_t> lens() const;
 
@@ -158,6 +161,10 @@ struct MIGRAPHX_EXPORT shape_transform_descriptor
 MIGRAPHX_EXPORT std::vector<operation>
 optimize_shape_transforms(const std::vector<std::size_t>& dims, const std::vector<operation>& ops);
 
+// Generate the shape transforms for strided view
+MIGRAPHX_EXPORT optional<std::vector<operation>>
+generate_shape_transforms_for(shape s, const std::vector<std::size_t>& idims, std::int64_t offset);
+
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx
 #endif // MIGRAPHX_GUARD_MIGRAPHX_SHAPE_TRANSFORM_DESCRIPTOR_HPP

src/include/migraphx/simplify_reshapes.hpp

@@ -40,6 +40,7 @@ struct MIGRAPHX_EXPORT simplify_reshapes
 {
     size_t depth = 4;
     bool enable_op_shape_transform_op = false;
+    bool enable_gather_rewrite        = false;
     std::string name() const { return "simplify_reshapes"; }
     void apply(module& m) const;
 };

src/module.cpp

@@ -1320,7 +1320,7 @@ module::print_py(std::ostream& os,
             if(ins->name() == "@literal")
             {
                 os << mname << ".add_literal(";
-                if(ins->get_shape().elements() < 10)
+                if(ins->get_shape().elements() < 1024)
                 {
                     os << "migraphx.create_argument(";
                     print_py_shape(os, ins->get_shape());

src/pass_manager.cpp

@@ -182,7 +182,7 @@ struct module_pm : module_pass_manager
         catch(const std::exception& e)
         {
             std::cerr << "Error " << p.name() << ": " << e.what() << std::endl;
-            auto clk = std::chrono::steady_clock::now().time_since_epoch().count();
+            auto clk         = std::chrono::steady_clock::now().time_since_epoch().count();
             fs::path dirname = fs::temp_directory_path() / "migraphx";
             fs::create_directories(dirname);
             std::string base = p.name() + std::to_string(clk) + ".mxr";

src/shape_transform_descriptor.cpp

@@ -1870,6 +1870,28 @@
     return result;
 }
 
+std::vector<std::size_t> shape_transform_descriptor::get_dst_axes_from_src(std::size_t axis) const
+{
+    std::vector<std::size_t> result;
+    for(auto i : range(dimensions.size()))
+    {
+        const auto& d = dimensions[i];
+        auto it       = std::find_if(d.subdimensions.begin(), d.subdimensions.end(), [&](auto& s) {
+            if(s.axis.empty())
+                return false;
+            return s.axis.front() == axis;
+        });
+        if(it == d.subdimensions.end())
+            continue;
+        // If it maps to a subdimesion then exit as there isnt a clear mapping
+        if(d.len() != it->len)
+            return {};
+        result.push_back(i);
+    }
+    // TODO: Put it in the correct order if there is multiple axes
+    return result;
+}
+
 bool shape_transform_descriptor::empty() const { return dimensions.empty(); }
 
 std::vector<std::size_t> shape_transform_descriptor::lens() const
@@ -2011,5 +2033,209 @@
     return sd.generate();
 }
 
+// Replace broadcasted dimensions with size 1, and set the stride to the previous stride
+static shape unbroadcast(const shape& s)
+{
+    std::vector<std::size_t> lens    = s.lens();
+    std::vector<std::size_t> strides = s.strides();
+    auto stride_it                   = std::find_if(
+        s.strides().begin(), s.strides().end(), [](auto stride) { return stride != 0; });
+    std::size_t prev_stride = stride_it == s.strides().end() ? 1 : *stride_it;
+    for(std::size_t i = 0; i < lens.size(); ++i)
+    {
+        if(strides[i] == 0)
+        {
+            lens[i]    = 1;
+            strides[i] = prev_stride;
+        }
+        else
+        {
+            prev_stride = strides[i];
+        }
+    }
+    return {s.type(), lens, strides};
+}
+
+static std::size_t adjust_strided_shape(shape& s, std::size_t n)
+{
+    auto lens    = s.lens();
+    auto strides = s.strides();
+
+    // Insert a dim of 1 so it can be used to handle steps
+    if(std::none_of(strides.begin(), strides.end(), [](auto stride) { return stride == 1; }) and
+       std::any_of(strides.begin(), strides.end(), [](auto stride) { return stride != 0; }))
+    {
+        lens.push_back(1);
+        strides.push_back(1);
+    }
+
+    auto last_axis      = std::max_element(strides.begin(), strides.end()) - strides.begin();
+    auto total_elements = std::max<std::size_t>(1, strides[last_axis] * lens[last_axis]);
+    // Add a dim of 1 to the front so it can handle extra elements
+    auto extra = n / total_elements;
+    if(extra > 1)
+    {
+        strides.insert(strides.begin(), total_elements);
+        lens.insert(lens.begin(), 1);
+    }
+    s = shape(s.type(), lens, strides);
+    return std::max<std::size_t>(1, extra);
+}
+
+// Generate the shape transforms for strided view
+optional<std::vector<operation>>
+generate_shape_transforms_for(shape s, const std::vector<std::size_t>& idims, std::int64_t offset)
+{
+    std::vector<operation> result;
+    if(s.lens().empty())
+        return std::nullopt;
+    std::size_t ielements =
+        std::accumulate(idims.begin(), idims.end(), std::size_t(1), std::multiplies<>());
+    auto extra = adjust_strided_shape(s, ielements);
+    // TODO: Improve handling of multiple dimensions, for now just reshape to 1 dimension
+    if(idims.size() != 1)
+    {
+        result.push_back(make_op("reshape", {{"dims", {ielements}}}));
+        auto ops = generate_shape_transforms_for(s, {ielements}, offset);
+        if(not ops)
+            return std::nullopt;
+        result.insert(result.end(), ops->begin(), ops->end());
+        return result;
+    }
+    auto pre_broadcast = unbroadcast(s);
+    auto perm          = find_permutation(pre_broadcast);
+    auto iperm         = invert_permutation(perm);
+    auto pre_transpose = reorder_shape(pre_broadcast, perm);
+
+    std::vector<std::size_t> start_lens;
+    std::adjacent_difference(pre_transpose.strides().begin(),
+                             pre_transpose.strides().end(),
+                             std::back_inserter(start_lens),
+                             [](auto y, auto x) -> std::size_t {
+                                 assert(x >= y);
+                                 assert(y != 0);
+                                 if((x % y) != 0)
+                                     return 0;
+                                 return x / y;
+                             });
+    if(std::any_of(start_lens.begin(), start_lens.end(), [](auto len) { return len == 0; }))
+        return std::nullopt;
+    start_lens.front() = extra > 1 ? extra : pre_transpose.lens().front();
+
+    std::size_t nelements =
+        std::accumulate(start_lens.begin(), start_lens.end(), std::size_t(1), std::multiplies<>());
+
+    if(nelements < pre_transpose.elements() * extra)
+        return std::nullopt;
+
+    std::vector<std::size_t> start_mask(start_lens.size(), 0);
+    if(offset != 0)
+    {
+        shape start_shape{shape::float_type, start_lens};
+        auto idx = start_shape.multi(offset);
+
+        std::vector<std::size_t> overhead;
+        std::transform(start_lens.begin(),
+                       start_lens.end(),
+                       pre_transpose.lens().begin(),
+                       std::back_inserter(overhead),
+                       [](auto start_len, auto len) { return start_len - len; });
+        if(std::equal(
+               idx.begin(), idx.end(), overhead.begin(), overhead.end(), [](auto i, auto over) {
+                   return i <= over;
+               }))
+        {
+            start_mask = reorder_dims(idx, iperm);
+            offset     = 0;
+        }
+    }
+
+    std::vector<std::size_t> pre_slice_mask;
+    std::transform(start_lens.begin(),
+                   start_lens.end(),
+                   pre_transpose.lens().begin(),
+                   std::back_inserter(pre_slice_mask),
+                   [](auto start_len, auto len) -> std::size_t {
+                       if(start_len == len)
+                           return 0;
+                       return len;
+                   });
+    auto slice_mask = reorder_dims(pre_slice_mask, iperm);
+
+    std::vector<std::size_t> blens = reorder_dims(start_lens, iperm);
+    std::transform(s.lens().begin(),
+                   s.lens().end(),
+                   blens.begin(),
+                   blens.begin(),
+                   [](auto len, auto blen) -> std::size_t {
+                       if(blen == 1)
+                           return len;
+                       return blen;
+                   });
+
+    std::vector<operation> ops;
+    ops.push_back(make_op("multibroadcast", {{"out_lens", blens}}));
+    ops.push_back(make_op("transpose", {{"permutation", invert_permutation(perm)}}));
+    ops.push_back(make_op("reshape", {{"dims", start_lens}}));
+    std::reverse(ops.begin(), ops.end());
+
+    auto desc = shape_transform_descriptor::create({nelements}, ops);
+
+    auto end = offset + nelements;
+    if(offset != 0 or nelements != ielements)
+    {
+
+        // If the end is out of bounds broadcast it to pad it
+        if(end > ielements)
+        {
+            result.push_back(make_op("broadcast", {{"axis", 1}, {"out_lens", {2, ielements}}}));
+            result.push_back(make_op("reshape", {{"dims", {2 * ielements}}}));
+        }
+        result.push_back(make_op("slice", {{"axes", {0}}, {"starts", {offset}}, {"ends", {end}}}));
+    }
+
+    auto opt_ops = desc.generate();
+    result.insert(result.end(), opt_ops.begin(), opt_ops.end());
+
+    std::vector<std::size_t> axes;
+    std::transform(slice_mask.begin(),
+                   slice_mask.end(),
+                   range(slice_mask.size()).begin(),
+                   join_back_inserter(axes),
+                   [](std::size_t mask, std::size_t idx) -> std::vector<std::size_t> {
+                       if(mask > 0)
+                           return {idx};
+                       return {};
+                   });
+
+    if(not axes.empty())
+    {
+        std::vector<std::size_t> starts;
+        std::transform(slice_mask.begin(),
+                       slice_mask.end(),
+                       start_mask.begin(),
+                       join_back_inserter(starts),
+                       [](std::size_t mask, std::size_t start) -> std::vector<std::size_t> {
+                           if(mask == 0)
+                               return {};
+                           return {start};
+                       });
+        std::vector<std::size_t> ends;
+        std::transform(slice_mask.begin(),
+                       slice_mask.end(),
+                       s.lens().begin(),
+                       join_back_inserter(ends),
+                       [](std::size_t mask, std::size_t len) -> std::vector<std::size_t> {
+                           if(mask == 0)
+                               return {};
+                           return {len};
+                       });
+        std::transform(ends.begin(), ends.end(), starts.begin(), ends.begin(), std::plus<>{});
+
+        result.push_back(make_op("slice", {{"axes", axes}, {"starts", starts}, {"ends", ends}}));
+    }
+    return result;
+}
+
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx