diff --git a/CHANGELOG.md b/CHANGELOG.md index 66525c60007..17b847f862a 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/src/driver/verify.cpp b/src/driver/verify.cpp index 9f3ce8f79e0..d94dcf4a030 100644 --- a/src/driver/verify.cpp +++ b/src/driver/verify.cpp @@ -36,7 +36,6 @@ #include #include #include -#include #include #include @@ -100,8 +99,7 @@ static std::vector 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); diff --git a/src/include/migraphx/instruction_traversal.hpp b/src/include/migraphx/instruction_traversal.hpp index 57ba5ecf10c..96a4a8e3510 100644 --- a/src/include/migraphx/instruction_traversal.hpp +++ b/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 { + // 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 diff --git a/src/simplify_qdq.cpp b/src/simplify_qdq.cpp index 829843b5917..a03caf16061 100644 --- a/src/simplify_qdq.cpp +++ b/src/simplify_qdq.cpp @@ -23,6 +23,7 @@ */ #include #include +#include #include #include #include @@ -38,6 +39,11 @@ #include #include #include +#include +#include +#include +#include +#include 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>> slices; - auto axis = any_cast(cat_ins->get_operator()).axis; + auto axis = cat_ins->normalized_operator().to_value().at("axis").to(); 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 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,27 +561,90 @@ 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 qdq_skip_ops = { + "pack_fp4", "unpack_fp4", "broadcast", "slice", "reshape", "reshape_lazy", "pad"}; + +constexpr std::array qdq_replay_ops = { + "broadcast", "slice", "reshape", "reshape_lazy"}; + +constexpr std::array 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 +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 ops; +}; + struct remove_qdq_pairs { + static std::optional 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 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 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"]; @@ -562,11 +652,21 @@ struct remove_qdq_pairs { 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 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 { diff --git a/test/simplify_qdq_test.cpp b/test/simplify_qdq_test.cpp index e123499fde7..63f01defa08 100644 --- a/test/simplify_qdq_test.cpp +++ b/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 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}};