Fix moving shape transforms through a shared reduction and fold unsqueeze into multibroadcast

CHANGELOG.md

@@ -16,6 +16,7 @@ Full documentation for MIGraphX is available at
 
 ### Resolved issues
 
+* Fixed `simplify_reshapes` generating invalid inverse shape transforms when moving an element-expanding shape-transform chain through a shared reduction (#5360).
 * Fixed non-standard GPU output layouts being exposed to integrations that require packed tensors (#5345).
 * Fixed mixed `uint8` and `int8` quantization (#5075).
 * Fixed flash-decoding rebuild for fused `@literal`/`@outline` and extra score-shaped `@param` inputs (#5114).

src/include/migraphx/shape_transform_descriptor.hpp

@@ -104,6 +104,8 @@ struct MIGRAPHX_EXPORT shape_transform_descriptor
                                     bool no_broadcast                          = false) const;
 
     std::set<std::size_t> find_broadcasted_axes() const;
+    // Returns true if a dimension is broadcasted without a source axis, ie a new dimension
+    bool has_axisless_broadcast() const;
     bool has_broadcast() const;
     // Returns true if any non-trivial axes are reordered, requiring a transpose
     bool is_transposed() const;
@@ -117,6 +119,8 @@ struct MIGRAPHX_EXPORT shape_transform_descriptor
     shape_transform_descriptor to_common_from_src() const;
     shape_transform_descriptor to_common_from_dst() const;
     shape_transform_descriptor to_dst_from_common() const;
+    // Transform the common dimensions back into the source dimensions. A dimension broadcasted
+    // from no source axis is squeezed away, so it is hidden and must be 1 to rebase.
     shape_transform_descriptor to_src_from_common() const;
 
     // Transform the destination dimensions back into the original dimensions. This is empty when

src/shape_transform_descriptor.cpp

@@ -215,6 +215,23 @@ static std::size_t visible_len(const std::vector<dimension::sub*>& subs)
         });
 }
 
+static bool has_no_axes(const dimension& d)
+{
+    return std::all_of(d.subdimensions.begin(), d.subdimensions.end(), [](const dimension::sub& s) {
+        return s.axis.empty() and s.hidden_axis.empty();
+    });
+}
+
+// Length of the dimension before any broadcast is applied
+static std::size_t input_len(const dimension& d)
+{
+    return transform_accumulate(d.subdimensions.begin(),
+                                d.subdimensions.end(),
+                                std::size_t{1},
+                                std::multiplies<>{},
+                                [](const dimension::sub& s) { return s.axis.empty() ? 1 : s.len; });
+}
+
 static std::vector<std::size_t> compute_dims(const operation& op,
                                              const std::vector<std::size_t>& idims)
 {
@@ -742,53 +759,103 @@ struct rebase_ambiguity_resolver
     static const std::size_t last_axis_split = std::numeric_limits<std::size_t>::max();
 };
 
-shape_transform_descriptor shape_transform_descriptor::rebase(const std::vector<std::size_t>& dims,
-                                                              bool broadcast) const
+// A hidden broadcast of an axis that is rebased to another size cant stay
+// attached to that axis. When a dimension of 1 without an axis precedes it,
+// the axes shift down into that dimension and the broadcast becomes its own
+// axis-less dimension, as if the axis had been unsqueezed and broadcasted.
+static void detach_unabsorbed_broadcasts(std::vector<dimension>& dimensions,
+                                         const std::vector<std::size_t>& dims)
 {
-    auto result   = *this;
-    auto axes_map = rebase_ambiguity_resolver{result, dims}.resolve();
-    for(auto& [axis, subs] : axes_map)
+    auto axes_map = group_axes(dimensions);
+    for(std::size_t h : range(dimensions.size()))
     {
+        auto& subs = dimensions[h].subdimensions;
+        if(subs.size() != 1)
+            continue;
+        auto& s = subs.front();
+        if(not s.has_hidden_axis() or s.hidden_axis.size() != 1)
+            continue;
+        auto axis = s.hidden_axis.front();
         assert(axis < dims.size());
-        auto dim       = dims[axis];
-        if(dim == len(subs))
-        {
-            if(not broadcast)
-            {
-                for(auto* sub : subs)
-                    sub->expose();
-            }
-        }
-        else if(dim == 1)
+        if(axes_map.at(axis).size() != 1)
+            continue;
+        auto dim = dims[axis];
+        if(dim == 0 or dim == 1 or dim == s.len)
+            continue;
+        // Search backwards through dimensions of 1 for one without an axis
+        auto rit = std::find_if(std::make_reverse_iterator(dimensions.begin() + h),
+                                dimensions.rend(),
+                                [](const dimension& d) { return d.len() != 1 or has_no_axes(d); });
+        if(rit == dimensions.rend() or rit->len() != 1)
+            continue;
+        dimension bdim{{dimension::sub{s.len, {}}}};
+        s.expose();
+        s.len = 1;
+        dimensions.erase(std::prev(rit.base()));
+        dimensions.insert(dimensions.begin() + h, bdim);
+    }
+}
+
+// Adjust the subdimensions of an axis to its new dim. Returns false when the
+// axis cant be rebased.
+static bool rebase_axis(const std::vector<dimension::sub*>& subs, std::size_t dim, bool broadcast)
+{
+    if(dim == len(subs))
+    {
+        if(not broadcast)
         {
             for(auto* sub : subs)
-            {
-                if(not sub->has_hidden_axis())
-                    sub->len = 1;
-            }
+                sub->expose();
         }
-        else if(subs.size() == 1)
+    }
+    else if(dim == 1)
+    {
+        for(auto* sub : subs)
         {
-            subs.front()->len = dim;
-            if(broadcast)
-                subs.front()->hide();
-            else
-                subs.front()->expose();
+            if(not sub->has_hidden_axis())
+                sub->len = 1;
         }
-        else if(dim == visible_len(subs))
+    }
+    else if(subs.size() == 1)
+    {
+        // A hidden axis of 1 has no broadcast to absorb a different dim
+        if(not broadcast and subs.front()->has_hidden_axis() and subs.front()->len == 1)
+            return false;
+        subs.front()->len = dim;
+        if(broadcast)
+            subs.front()->hide();
+        else
+            subs.front()->expose();
+    }
+    else if(dim == visible_len(subs))
+    {
+        for(auto* sub : subs)
         {
-            for(auto* sub : subs)
+            if(sub->has_hidden_axis())
             {
-                if(sub->has_hidden_axis())
-                {
-                    sub->expose();
-                    sub->len = 1;
-                }
+                sub->expose();
+                sub->len = 1;
             }
         }
-        else
-            return {};
     }
+    else
+        return false;
+    return true;
+}
+
+shape_transform_descriptor shape_transform_descriptor::rebase(const std::vector<std::size_t>& dims,
+                                                              bool broadcast) const
+{
+    auto result = *this;
+    if(broadcast)
+        detach_unabsorbed_broadcasts(result.dimensions, dims);
+    auto axes_map = rebase_ambiguity_resolver{result, dims}.resolve();
+    if(not std::all_of(axes_map.begin(), axes_map.end(), [&](const auto& p) {
+           const auto& [axis, subs] = p;
+           assert(axis < dims.size());
+           return rebase_axis(subs, dims[axis], broadcast);
+       }))
+        return {};
     for(auto& dim : result.dimensions)
         remove_empty_sub_dims(dim.subdimensions);
     if(broadcast and not is_broadcast_only(dimensions, result.dimensions))
@@ -1704,12 +1771,6 @@ struct operation_list
 
 } // namespace
 
-static bool has_no_axes(const dimension& d)
-{
-    return std::all_of(d.subdimensions.begin(), d.subdimensions.end(), [](const dimension::sub& s) {
-        return s.axis.empty() and s.hidden_axis.empty();
-    });
-}
 static bool has_axes(const dimension& d)
 {
     return std::any_of(d.subdimensions.begin(), d.subdimensions.end(), [](const dimension::sub& s) {
@@ -1793,6 +1854,14 @@ shape_transform_descriptor::generate(const std::vector<std::size_t>& input_dims,
             else
             {
                 result.push_back(make_op("multibroadcast", {{"out_lens", out_lens}}));
+                // Broadcasting aligns from the right, so a dimension without an
+                // axis that only follows input dimensions of 1 needs no unsqueeze
+                auto first_non1 =
+                    std::find_if(new_dims.begin(), new_dims.end(), [](const dimension& d) {
+                        return input_len(d) != 1;
+                    });
+                new_dims.erase(std::remove_if(new_dims.begin(), first_non1, &has_no_axes),
+                               first_non1);
             }
         }
         // If all the dimensions have no axes then there isnt anthing else to do
@@ -1860,6 +1929,13 @@ std::set<std::size_t> shape_transform_descriptor::find_broadcasted_axes() const
     return result;
 }
 
+bool shape_transform_descriptor::has_axisless_broadcast() const
+{
+    auto subs = get_all_subdimensions(dimensions);
+    return std::any_of(subs.begin(), subs.end(), [](const dimension::sub& s) {
+        return s.origin_axis().empty() and s.len != 1;
+    });
+}
 bool shape_transform_descriptor::has_broadcast() const
 {
     return std::any_of(dimensions.begin(), dimensions.end(), [&](const dimension& d) {
@@ -1990,6 +2066,15 @@ shape_transform_descriptor shape_transform_descriptor::to_src_from_common() cons
         return {subdimensions};
     });
     result.simplify();
+    // A broadcasted dimension with no source axis is squeezed away, so hide
+    // its axis to require it to stay 1 when rebased
+    for_each_subdimension(result.dimensions, [&](dimension::sub& s) {
+        if(s.axis.size() != 1)
+            return;
+        const auto& common_sub = subs[s.axis.front()];
+        if(common_sub.origin_axis().empty() and common_sub.len != 1)
+            s.hide();
+    });
     return result;
 }
 

src/simplify_algebra.cpp

@@ -1019,13 +1019,17 @@ struct find_concat_op
         }
         else if(op.name() == "multibroadcast")
         {
-            shape bshape = (*start)->get_shape();
-            auto input   = (*start)->inputs()[0];
-            if(iaxis >= bshape.strides().size() or bshape.strides()[iaxis] == 0)
+            auto input_rank = (*start)->inputs().front()->get_shape().ndim();
+            // The inputs are concatenated, so the axis cant be broadcasted and
+            // the inputs need the same rank for it to map to the same input axis
+            if(std::any_of(start, last, [&](instruction_ref i) {
+                   const auto& s = i->get_shape();
+                   return iaxis >= s.ndim() or s.strides()[iaxis] == 0 or
+                          i->inputs().front()->get_shape().ndim() != input_rank;
+               }))
                 return {start, last};
             op.from_value({{"out_lens", get_output_lens(start, last, iaxis)}});
-            auto delta = bshape.lens().size() - input->get_shape().lens().size();
-            iaxis -= delta;
+            iaxis -= (*start)->get_shape().ndim() - input_rank;
         }
         else if(op.name() == "unsqueeze")
         {

src/simplify_reshapes.cpp

@@ -277,11 +277,17 @@ struct find_op_shape_transform_op
         // shape_transform_descriptor doesnt handle scalars for now
         if(input_ins->get_shape().scalar() or x_ins->get_shape().scalar())
             return false;
-        // If its just a broadcast then skip
+        // A chain of only broadcasts is skipped, unless it broadcasts a new
+        // axis of a reduction, which is the same as an unsqueeze and broadcast
         if(not any_input_of(input_ins, x_ins, [](instruction_ref x) {
                return not contains({"multibroadcast", "broadcast", "contiguous"}, x->name());
            }))
-            return false;
+        {
+            if(not is_reduce(x_ins))
+                return false;
+            auto desc = make_descriptor(x_ins, chain_ops(x_ins, input_ins), input_ins);
+            return not desc.empty() and desc.has_axisless_broadcast();
+        }
         // A chain that changes the element count is only valid for a reduction
         return is_reduce(x_ins) or
                input_ins->get_shape().elements() == x_ins->get_shape().elements();
@@ -292,12 +298,40 @@ struct find_op_shape_transform_op
         return is_reduce(ins) or ins->get_operator().attributes().contains("pointwise");
     }
 
+    // True when x_ins is only consumed through the chain ending at ins, so
+    // it is dead once ins is rewritten
+    static bool
+    is_private_chain(instruction_ref x_ins, instruction_ref input_ins, instruction_ref ins)
+    {
+        if(not all_of(input_ins->outputs(), [&](instruction_ref out) { return out == ins; }))
+            return false;
+        if(x_ins->outputs().size() != 1)
+            return false;
+        return not any_input_of(input_ins->inputs().front(), x_ins, [](instruction_ref i) {
+            return i->outputs().size() != 1;
+        });
+    }
+
     static bool is_reduce(instruction_ref ins)
     {
         return starts_with(ins->name(), "reduce_") or ins->name() == "argmin" or
                ins->name() == "argmax";
     }
 
+    // The shape transform operators from x_ins to input_ins in order
+    static std::vector<operation> chain_ops(instruction_ref x_ins, instruction_ref input_ins)
+    {
+        std::vector<operation> ops;
+        auto next_ins = input_ins;
+        while(next_ins != x_ins)
+        {
+            ops.push_back(next_ins->get_operator());
+            next_ins = next_ins->inputs().front();
+        }
+        std::reverse(ops.begin(), ops.end());
+        return ops;
+    }
+
     template <class F>
     static instruction_ref find_input_if(instruction_ref start, instruction_ref last, F f)
     {
@@ -480,23 +514,20 @@ struct find_op_shape_transform_op
         auto x_ins     = r.instructions["x"];
         auto input_ins = r.instructions["input"];
 
-        std::vector<operation> ops;
-        auto next_ins = input_ins;
-        while(next_ins != x_ins)
-        {
-            ops.push_back(next_ins->get_operator());
-            next_ins = next_ins->inputs().front();
-        }
-        assert(next_ins == x_ins);
-        std::reverse(ops.begin(), ops.end());
-
-        auto desc = make_descriptor(x_ins, ops, input_ins);
+        auto desc = make_descriptor(x_ins, chain_ops(x_ins, input_ins), input_ins);
         if(desc.empty())
             return;
 
         if(not is_valid(x_ins, desc))
             return;
 
+        // x_ins is restored by rebasing the common dims onto its transform, which
+        // fails when the chain broadcasts a new axis, so then it can only be removed
+        auto src_desc        = desc.to_src_from_common();
+        const bool restore_x = not src_desc.rebase(desc.common_dims()).empty();
+        if(not restore_x and not is_private_chain(x_ins, input_ins, ins))
+            return;
+
         // ins is remapped via the dst map; bail if its argmin/argmax axis splits
         if(is_reduce(ins) and ins->get_operator().to_value().contains("axis") and
            not argmax_axis_unsplit(ins, desc.common_axes_map_from_dst()))
@@ -524,7 +555,13 @@ struct find_op_shape_transform_op
                        x_inputs.begin(),
                        reshape_input(x_ins, desc.to_common_from_src()));
         auto new_input_ins = insert(m, x_ins, x_inputs, desc.common_axes_map_from_src());
-        auto new_x_ins     = reshape_input(x_ins, desc.to_src_from_common())(new_input_ins);
+        // Replace old x_ins just in case it is used more than once
+        if(restore_x)
+        {
+            auto new_x_ins = reshape_input(x_ins, src_desc)(new_input_ins);
+            assert(x_ins->get_shape().lens() == new_x_ins->get_shape().lens());
+            m.replace_instruction(x_ins, new_x_ins);
+        }
         if(new_input_ins->get_shape().elements() != input_ins->get_shape().elements())
         {
             auto cdims    = desc.common_dims();
@@ -537,9 +574,6 @@ struct find_op_shape_transform_op
                 return new_input_ins;
             return reshape_input(ins, desc.to_common_from_dst(), true)(input);
         });
-        // Replace old x_ins just in case it is used more than once
-        assert(x_ins->get_shape().lens() == new_x_ins->get_shape().lens());
-        m.replace_instruction(x_ins, new_x_ins);
         // Replace final instruction
         auto pw   = insert(m, ins, inputs, desc.common_axes_map_from_dst());
         auto rins = reshape_input(ins, desc.to_dst_from_common())(pw);

test/fuse_reduce.cpp

@@ -1978,7 +1978,7 @@ TEST_CASE(pointwise_reshapes_reduce_shadowed_broadcast)
         auto w   = mm->add_parameter("w", ws);
         auto pw0 = add_pointwise(p2, "main:pointwise0", {x}, single_pointwise("sqrt"));
         auto wu  = mm->add_instruction(migraphx::make_op("unsqueeze", {{"axes", {0}}}), w);
-        auto xr  = mm->add_instruction(migraphx::make_op("reshape", {{"dims", {1, 1, 2, 4}}}), pw0);
+        auto xr  = mm->add_instruction(migraphx::make_op("reshape", {{"dims", {1, 2, 4}}}), pw0);
         auto xb  = mm->add_instruction(
             migraphx::make_op("multibroadcast", {{"out_lens", {1, 4, 2, 4}}}), xr);
         auto rsum = add_reduce(