Preserve shape ops when removing QDQ pairs

CHANGELOG.md

@@ -65,6 +65,7 @@ Full documentation for MIGraphX is available at
 
 ### Resolved issues
 
+* Fixed `simplify_qdq` Q/DQ removal to preserve intervening shape operations and avoid rewrites that change output shapes.
 * Fixed the reference `nonzero` operator to handle non-standard input layouts such as transposed or broadcasted tensors.
 * Restored support for the documented flat {min,max,optimals} JSON format in migraphx-driver's --default-dyn-dim and --dyn-input-dim flags (#4926).
 * Fixed ONNX `Where` parsing for dynamic-shape inputs that require broadcasting (including mixed static and dynamic inputs), which previously threw `same_dims: where: Dimensions do not match` (#4925).

src/driver/verify.cpp

@@ -36,7 +36,6 @@
 #include <migraphx/stringutils.hpp>
 #include <migraphx/verify_args.hpp>
 #include <migraphx/simplify_qdq.hpp>
-#include <migraphx/dead_code_elimination.hpp>
 #include <migraphx/logger.hpp>
 #include <utility>
 
@@ -100,8 +99,7 @@ static std::vector<argument> run_ref(program p,
 {
     if(vo.ref_use_double)
     {
-        run_passes(
-            p, {fp_to_double{}, simplify_qdq{.remove_qdq_only = true}, dead_code_elimination{}});
+        run_passes(p, {fp_to_double{}, simplify_qdq{.remove_qdq_only = true}});
     }
     p.compile(migraphx::make_target("ref"), options);
     auto out = p.eval(inputs);

src/include/migraphx/instruction_traversal.hpp

@@ -57,6 +57,16 @@ inline auto get_alias_path(instruction_ref ins)
     });
 }
 
+inline auto get_input_path(instruction_ref ins)
+{
+    return unfold(ins, [](instruction_ref x) -> std::optional<instruction_ref> {
+        // Follow only linear input chains; branches terminate the path.
+        if(x->inputs().size() != 1)
+            return std::nullopt;
+        return x->inputs().front();
+    });
+}
+
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx
 #endif // MIGRAPHX_GUARD_MIGRAPHX_INSTRUCTION_TRAVERSAL_HPP

src/simplify_qdq.cpp

@@ -23,6 +23,7 @@
  */
 #include <migraphx/simplify_qdq.hpp>
 #include <migraphx/instruction.hpp>
+#include <migraphx/instruction_traversal.hpp>
 #include <migraphx/iterator_for.hpp>
 #include <migraphx/make_op.hpp>
 #include <migraphx/program.hpp>
@@ -38,6 +39,11 @@
 #include <migraphx/register_op.hpp>
 #include <migraphx/fp8_types.hpp>
 #include <migraphx/match/dq_helpers.hpp>
+#include <algorithm>
+#include <array>
+#include <numeric>
+#include <optional>
+#include <string_view>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -442,7 +448,7 @@ struct match_concat_qlinear
     auto get_slices(instruction_ref cat_ins) const
     {
         std::vector<std::vector<std::pair<std::string, value>>> slices;
-        auto axis    = any_cast<op::concat>(cat_ins->get_operator()).axis;
+        auto axis    = cat_ins->normalized_operator().to_value().at("axis").to<std::size_t>();
         size_t start = 0;
         for(auto cat_inp : cat_ins->inputs())
         {
@@ -462,11 +468,32 @@ struct match_concat_qlinear
         auto scale = ins->inputs()[1];
         auto zp    = ins->inputs()[2];
 
-        auto slices = get_slices(cat_ins);
+        if(not shape::same_lens(scale->get_shape(), ins->get_shape()) or
+           not shape::same_lens(zp->get_shape(), ins->get_shape()))
+        {
+            return;
+        }
+
+        auto slices     = get_slices(cat_ins);
+        auto cat_inputs = cat_ins->inputs();
+        // Each sliced scale and zero point must match the concat input it quantizes, so
+        // check the shapes before inserting anything.
+        auto slice_matches = [&](instruction_ref cat_inp, const auto& slc) {
+            auto slice_op = make_op("slice", slc);
+            return shape::same_lens(slice_op.compute_shape({scale->get_shape()}),
+                                    cat_inp->get_shape()) and
+                   shape::same_lens(slice_op.compute_shape({zp->get_shape()}),
+                                    cat_inp->get_shape());
+        };
+        if(not std::equal(cat_inputs.begin(), cat_inputs.end(), slices.begin(), slice_matches))
+        {
+            return;
+        }
+
         std::vector<instruction_ref> new_cat_inputs;
         std::transform(
-            cat_ins->inputs().begin(),
-            cat_ins->inputs().end(),
+            cat_inputs.begin(),
+            cat_inputs.end(),
             slices.begin(),
             std::back_inserter(new_cat_inputs),
             [&](auto i, const auto& slc) {
@@ -534,39 +561,112 @@ bool is_any_input_int4(instruction_ref a)
  *  that match the name with one input recursively.
  *  Use this after selected quantizations have been made into real quantized instructions.
  */
+constexpr std::array<std::string_view, 7> qdq_skip_ops = {
+    "pack_fp4", "unpack_fp4", "broadcast", "slice", "reshape", "reshape_lazy", "pad"};
+
+constexpr std::array<std::string_view, 4> qdq_replay_ops = {
+    "broadcast", "slice", "reshape", "reshape_lazy"};
+
+constexpr std::array<std::string_view, 2> qdq_pack_unpack_ops = {"pack_fp4", "unpack_fp4"};
+
+bool is_pack_unpack_op(instruction_ref ins) { return contains(qdq_pack_unpack_ops, ins->name()); }
+
+bool is_replay_op(instruction_ref ins) { return contains(qdq_replay_ops, ins->name()); }
+
+template <class Iterator>
+bool has_pack_unpack_op(Iterator first, Iterator last)
+{
+    return std::any_of(first, last, is_pack_unpack_op);
+}
+
+bool is_supported_qdq_path_op(instruction_ref ins, bool has_pack_unpack)
+{
+    return is_pack_unpack_op(ins) or is_replay_op(ins) or
+           (has_pack_unpack and contains(qdq_skip_ops, ins->name()));
+}
+
+MIGRAPHX_PRED_MATCHER(qdq_skip_op, instruction_ref ins)
+{
+    return contains(qdq_skip_ops, ins->name());
+}
+
+struct qdq_replacement
+{
+    instruction_ref input;
+    std::vector<operation> ops;
+};
+
 struct remove_qdq_pairs
 {
+    static std::optional<qdq_replacement> find_replacement(instruction_ref dq_ins,
+                                                           instruction_ref q_ins)
+    {
+        qdq_replacement replacement{q_ins->inputs().front(), {}};
+        auto input_path = get_input_path(dq_ins->inputs().front());
+        auto q_pos      = std::find(input_path.begin(), input_path.end(), q_ins);
+        if(q_pos == input_path.end())
+            return std::nullopt;
+
+        auto path_to_q           = range(input_path.begin(), q_pos);
+        const bool pack_fp4_path = has_pack_unpack_op(input_path.begin(), q_pos);
+        if(not std::all_of(path_to_q.begin(), path_to_q.end(), [&](auto x) {
+               return is_supported_qdq_path_op(x, pack_fp4_path);
+           }))
+        {
+            return std::nullopt;
+        }
+
+        if(not pack_fp4_path)
+        {
+            std::vector<instruction_ref> replay_instructions;
+            std::copy_if(
+                input_path.begin(), q_pos, std::back_inserter(replay_instructions), is_replay_op);
+            // The input path runs from DQ to Q, so replay it from Q to DQ.
+            std::transform(replay_instructions.rbegin(),
+                           replay_instructions.rend(),
+                           std::back_inserter(replacement.ops),
+                           [](auto ins) { return ins->get_operator(); });
+        }
+
+        auto replacement_shape =
+            std::accumulate(replacement.ops.begin(),
+                            replacement.ops.end(),
+                            replacement.input->get_shape(),
+                            [](const auto& s, const auto& op) { return op.compute_shape({s}); });
+        if(replacement_shape != dq_ins->get_shape())
+            return std::nullopt;
+        return replacement;
+    }
+
     auto matcher() const
     {
-        // clang-format off
-        static const std::unordered_set<std::string> skip_set = {
-            "pack_fp4",
-            "unpack_fp4",
-            "broadcast",
-            "slice",
-            "reshape",
-            "reshape_lazy",
-            "pad",
-        };
-        // clang-format on
-        auto q_ins =
-            match::skip(match::name(skip_set))(match::name("quantizelinear").bind("q_ins"));
+        auto q_ins = match::skip(qdq_skip_op())(match::name("quantizelinear").bind("q_ins"));
         return match::name("dequantizelinear")(match::arg(0)(q_ins));
     }
 
-    auto apply(module&, const match::matcher_result& r) const
+    auto apply(module& m, const match::matcher_result& r) const
     {
         auto dq_ins = r.result;
         auto q_ins  = r.instructions["q_ins"];
         if(not is_same_scale_zero(dq_ins, q_ins))
         {
             return;
         }
+        auto replacement = find_replacement(dq_ins, q_ins);
+        if(not replacement.has_value())
+        {
+            return;
+        }
+        auto replacement_ins = replacement->input;
+        for(const auto& op : replacement->ops)
+        {
+            replacement_ins = m.insert_instruction(dq_ins, op, replacement_ins);
+        }
         // Need to copy outputs since will be modifying dq_ins outputs
         std::vector<instruction_ref> dq_outputs = dq_ins->outputs();
         for(auto out : dq_outputs)
         {
-            instruction::replace_argument(out, dq_ins, q_ins->inputs().front());
+            instruction::replace_argument(out, dq_ins, replacement_ins);
         }
     }
 };
@@ -647,6 +747,7 @@ void simplify_qdq::apply(module& m) const
     if(remove_qdq_only)
     {
         match::find_matches(m, remove_qdq_pairs{});
+        migraphx::run_passes(m, {migraphx::dead_code_elimination{}});
     }
     else
     {

test/simplify_qdq_test.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
@@ -95,6 +95,66 @@ TEST_CASE(remove_qdq)
     EXPECT(m1 == m2);
 }
 
+TEST_CASE(remove_qdq_only_runs_dce)
+{
+    migraphx::shape s{migraphx::shape::float_type, {4, 4}};
+
+    migraphx::module m1;
+    {
+        auto input = m1.add_parameter("input", s);
+        auto scale = m1.add_literal(0.5f);
+        auto zero  = m1.add_literal(std::int8_t{0});
+        auto q     = add_quantize_op(m1, "quantizelinear", input, scale, zero);
+        auto dq    = add_quantize_op(m1, "dequantizelinear", q, scale, zero);
+        m1.add_return({dq});
+    }
+
+    migraphx::module m2;
+    {
+        auto input = m2.add_parameter("input", s);
+        m2.add_return({input});
+    }
+
+    migraphx::run_passes(m1, {migraphx::simplify_qdq{.remove_qdq_only = true}});
+    EXPECT(m1 == m2);
+}
+
+TEST_CASE(remove_qdq_preserve_shape_ops)
+{
+    migraphx::shape input_shape{migraphx::shape::float_type, {4, 4}};
+    migraphx::shape output_shape{migraphx::shape::float_type, {4, 2}};
+
+    migraphx::module m1;
+    {
+        auto a     = m1.add_parameter("a", input_shape);
+        auto b     = m1.add_parameter("b", output_shape);
+        auto scale = m1.add_literal(0.5f);
+        auto zero  = m1.add_literal(std::int8_t{0});
+
+        auto q = add_quantize_op(m1, "quantizelinear", a, scale, zero);
+        auto s = m1.add_instruction(
+            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {2}}}), q);
+        auto r   = m1.add_instruction(migraphx::make_op("reshape", {{"dims", {4, 2}}}), s);
+        auto dq  = add_quantize_op(m1, "dequantizelinear", r, scale, zero);
+        auto add = m1.add_instruction(migraphx::make_op("add"), dq, b);
+        m1.add_return({add});
+    }
+
+    migraphx::module m2;
+    {
+        auto a = m2.add_parameter("a", input_shape);
+        auto b = m2.add_parameter("b", output_shape);
+        auto s = m2.add_instruction(
+            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {2}}}), a);
+        auto r   = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {4, 2}}}), s);
+        auto add = m2.add_instruction(migraphx::make_op("add"), r, b);
+        m2.add_return({add});
+    }
+
+    run_pass(m1);
+    EXPECT(m1 == m2);
+}
+
 TEST_CASE(qdq_different_scales)
 {
     migraphx::shape sh1{migraphx::shape::float_type, {100, 100}};
@@ -1591,6 +1651,58 @@ TEST_CASE(pointwise_concat_quant_per_tensor)
     EXPECT(m1 == m2);
 }
 
+TEST_CASE(pointwise_concat_quant_split_two_pointwise_inputs)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {2, 3}};
+    migraphx::shape s2{migraphx::shape::float_type, {2, 5}};
+
+    migraphx::module m1;
+    {
+        auto i1    = m1.add_parameter("i1", s1);
+        auto i2    = m1.add_parameter("i2", s2);
+        auto scale = m1.add_literal(0.5f);
+        auto zero  = m1.add_literal(std::int8_t{0});
+
+        auto relu1 = m1.add_instruction(migraphx::make_op("relu"), i1);
+        auto relu2 = m1.add_instruction(migraphx::make_op("relu"), i2);
+        auto cat   = m1.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), relu1, relu2);
+        auto q     = add_quantize_op(m1, "quantizelinear", cat, scale, zero);
+        m1.add_return({q});
+    }
+
+    migraphx::module m2;
+    {
+        std::vector<std::size_t> cat_lens{2, 8};
+        auto i1    = m2.add_parameter("i1", s1);
+        auto i2    = m2.add_parameter("i2", s2);
+        auto scale = m2.add_literal(0.5f);
+        auto zero  = m2.add_literal(std::int8_t{0});
+
+        auto relu1    = m2.add_instruction(migraphx::make_op("relu"), i1);
+        auto relu2    = m2.add_instruction(migraphx::make_op("relu"), i2);
+        auto scale_mb = broadcast_scale(m2, scale, cat_lens, 1);
+        auto zero_mb  = broadcast_shift(m2, zero, cat_lens);
+
+        auto sc1 = m2.add_instruction(
+            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {3}}}), scale_mb);
+        auto zp1 = m2.add_instruction(
+            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {3}}}), zero_mb);
+        auto q1 = add_quantize_op(m2, "quantizelinear", relu1, sc1, zp1);
+
+        auto sc2 = m2.add_instruction(
+            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {3}}, {"ends", {8}}}), scale_mb);
+        auto zp2 = m2.add_instruction(
+            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {3}}, {"ends", {8}}}), zero_mb);
+        auto q2 = add_quantize_op(m2, "quantizelinear", relu2, sc2, zp2);
+
+        auto cat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 1}}), q1, q2);
+        m2.add_return({cat});
+    }
+
+    run_pass(m1);
+    EXPECT(m1 == m2);
+}
+
 TEST_CASE(pointwise_concat_quant_per_channel)
 {
     migraphx::shape s1{migraphx::shape::float_type, {1, 4, 28, 28}};