rebias uint8 to int8 on models with mixed datatypes

src/propagate_precision.cpp

@@ -142,6 +142,16 @@ static std::optional<instruction_ref> get_next_input(instruction_ref ins)
     return nullopt;
 }
 
+// Promotion converts an op's inputs to the target type, which only preserves the value when the
+// inputs share the target's category. quantizelinear mixes a float scale with integral operands,
+// so promoting it to int would truncate the scale (e.g. 0.05 -> 0) and zero the result.
+static bool inputs_match_category(instruction_ref ins, precision target)
+{
+    return std::all_of(ins->inputs().begin(), ins->inputs().end(), [&](instruction_ref input) {
+        return same_category(precision{input->get_shape().type()}, target);
+    });
+}
+
 // Find all adjacent instructions that could be upgraded with higher precision
 // by traversing the inputs from a convert
 
@@ -158,6 +168,8 @@ static std::unordered_set<instruction_ref> find_adjacent_inputs(instruction_ref
         // Stop when crossing a type category boundary (e.g., int or fp8 to float)
         if(not same_category(precision{ins->get_shape().type()}, target))
             return;
+        if(not inputs_match_category(ins, target))
+            return;
         auto next = get_next_input(ins);
         if(not next.has_value())
             return;
@@ -184,6 +196,8 @@ static std::unordered_set<instruction_ref> find_adjacent_outputs(instruction_ref
             // Stop when crossing a type category boundary (e.g., int or fp8 to float)
             if(not same_category(precision{output->get_shape().type()}, target))
                 continue;
+            if(not inputs_match_category(output, target))
+                continue;
             auto next = get_next_input(output);
             if(not next.has_value())
                 continue;

src/simplify_qdq.cpp

@@ -81,6 +81,27 @@ auto propagate_quantized_ins(module& m,
     return input_ins;
 }
 
+// Rebias a uint8 operand and its zero point to int8 by subtracting 128 from both: (q - 128) -
+// (zp - 128) == q - zp, so the dequantized value is unchanged. Done via int32 to avoid wrap-around.
+instruction_ref
+rebias_uint8_to_int8(module& m, instruction_ref pos, instruction_ref qdata, instruction_ref& zp)
+{
+    if(qdata->get_shape().type() != migraphx::shape::uint8_type)
+        return qdata;
+    auto subtract_128 = [&](instruction_ref x) {
+        auto x_i32 = m.insert_instruction(
+            pos, make_op("convert", {{"target_type", migraphx::shape::int32_type}}), x);
+        auto lit   = m.add_literal(literal{shape{migraphx::shape::int32_type}, {128}});
+        auto lit_b = m.insert_instruction(
+            pos, make_op("multibroadcast", {{"out_lens", x->get_shape().lens()}}), lit);
+        auto diff = m.insert_instruction(pos, make_op("sub"), x_i32, lit_b);
+        return m.insert_instruction(
+            pos, make_op("convert", {{"target_type", migraphx::shape::int8_type}}), diff);
+    };
+    zp = subtract_128(zp);
+    return subtract_128(qdata);
+}
+
 struct match_find_quantizable_ops
 {
     static bool
@@ -156,6 +177,7 @@ struct match_find_quantizable_ops
 
         std::set<migraphx::shape::type_t> supported_types = fp8_types{}.get();
         supported_types.insert(migraphx::shape::int8_type);
+        supported_types.insert(migraphx::shape::uint8_type);
         auto in1 = dq1->inputs().front();
         auto in2 = dq2->inputs().front();
         if(not contains(supported_types, in1->get_shape().type()) or
@@ -164,6 +186,10 @@ struct match_find_quantizable_ops
             return;
         }
 
+        // quant_convolution/quant_dot need both operands to share a type, so rebias uint8 to int8.
+        qop_args.at(0) = rebias_uint8_to_int8(m, qop, qop_args.at(0), zp1);
+        qop_args.at(1) = rebias_uint8_to_int8(m, qop, qop_args.at(1), zp2);
+
         instruction_ref dq;
         instruction_ref out_scale;
         instruction_ref out_zp;
@@ -279,6 +305,11 @@ struct match_find_quantizable_ops
             }
         }
 
+        // The correction inherits a transposed layout from its broadcast zero-point operand, which
+        // folds to a transposed literal that the fused MLIR dequant kernel misreads as standard.
+        if(out_zp->get_shape().packed() and not out_zp->get_shape().standard())
+            out_zp = m.insert_instruction(qop, make_op("contiguous"), out_zp);
+
         dq = m.insert_instruction(qop, make_op("dequantizelinear"), dq, out_scale, out_zp);
         if(is_fp16_model)
         {
@@ -562,11 +593,18 @@ struct remove_qdq_pairs
         {
             return;
         }
+        // Bypassing the pair drops the ops in between, which is only valid when they round-trip to
+        // identity. Skip if a shape-altering op (e.g. reshape) changed the lengths.
+        auto q_input = q_ins->inputs().front();
+        if(q_input->get_shape().lens() != dq_ins->get_shape().lens())
+        {
+            return;
+        }
         // 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, q_input);
         }
     }
 };

test/simplify_algebra_test.cpp

@@ -1065,6 +1065,29 @@ TEST_CASE(simplify_concat_unpack_int4)
     EXPECT(m1 == m2);
 }
 
+TEST_CASE(simplify_concat_dequantizelinear)
+{
+    // dequantizelinear must not be pulled back through the concat: that would leave the concat
+    // operating on integer values, which some backends (e.g. rocMLIR) cannot fuse.
+    auto si = migraphx::shape{migraphx::shape::int8_type, {4}};
+    auto sf = migraphx::shape{migraphx::shape::float_type, {4}};
+    migraphx::module m1;
+    {
+        auto x      = m1.add_parameter("x", si);
+        auto y      = m1.add_parameter("y", si);
+        auto scale  = m1.add_literal(migraphx::literal{sf, {0.5f, 0.5f, 0.5f, 0.5f}});
+        auto zp     = m1.add_literal(migraphx::literal{si, {5, 5, 5, 5}});
+        auto dq1    = m1.add_instruction(migraphx::make_op("dequantizelinear"), x, scale, zp);
+        auto dq2    = m1.add_instruction(migraphx::make_op("dequantizelinear"), y, scale, zp);
+        auto concat = m1.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), dq1, dq2);
+        m1.add_return({concat});
+    }
+
+    migraphx::module m2 = m1;
+    run_pass(m1);
+    EXPECT(m1 == m2);
+}
+
 TEST_CASE(simplify_concat_add_relu)
 {
     auto s = migraphx::shape{migraphx::shape::int32_type, {1}};

test/simplify_qdq_test.cpp

@@ -41,6 +41,7 @@ namespace match = migraphx::match;
 
 static bool is_convolution(const migraphx::instruction& ins) { return ins.name() == "convolution"; }
 static bool is_dot(const migraphx::instruction& ins) { return ins.name() == "dot"; }
+static bool is_quant_dot(const migraphx::instruction& ins) { return ins.name() == "quant_dot"; }
 
 static void run_pass(migraphx::module& m)
 {
@@ -346,6 +347,62 @@ TEST_CASE(dot_transposed)
     EXPECT(m1 == m2);
 }
 
+TEST_CASE(dot_uint8_input)
+{
+    // uint8 activation with int8 weight must still become a quant_dot (uint8 rebiased to int8),
+    // not have its Q/DQ stripped.
+    migraphx::shape sh1{migraphx::shape::float_type, {4, 8}};
+    migraphx::shape sh2{migraphx::shape::float_type, {8, 6}};
+
+    migraphx::module m1;
+    {
+        auto t1      = m1.add_parameter("t1", sh1);
+        auto t2      = m1.add_parameter("t2", sh2);
+        auto a_scale = m1.add_literal(0.5f);
+        auto a_zp    = m1.add_literal(std::uint8_t{128});
+        auto w_scale = m1.add_literal(0.25f);
+        auto w_zp    = m1.add_literal(std::int8_t{0});
+
+        auto q1  = add_quantize_op(m1, "quantizelinear", t1, a_scale, a_zp);
+        auto d1  = add_quantize_op(m1, "dequantizelinear", q1, a_scale, a_zp);
+        auto q2  = add_quantize_op(m1, "quantizelinear", t2, w_scale, w_zp);
+        auto d2  = add_quantize_op(m1, "dequantizelinear", q2, w_scale, w_zp);
+        auto dot = m1.add_instruction(migraphx::make_op("dot"), d1, d2);
+        m1.add_return({dot});
+    }
+
+    run_pass(m1);
+    EXPECT(any_of(m1, &is_quant_dot));
+    EXPECT(none_of(m1, &is_dot));
+}
+
+TEST_CASE(qdq_reshape_unquantized_dot)
+{
+    // q -> reshape -> dq feeding an unquantized dot stays put: remove_qdq_pairs must not drop the
+    // reshape, which would leave the dot with mismatched ranks.
+    migraphx::shape xsh{migraphx::shape::float_type, {1, 1024, 1, 1}};
+    migraphx::shape wsh{migraphx::shape::float_type, {1024, 1000}};
+
+    migraphx::module m1;
+    {
+        auto x     = m1.add_parameter("x", xsh);
+        auto w     = m1.add_parameter("w", wsh);
+        auto scale = m1.add_literal(0.5f);
+        auto zero  = m1.add_literal(std::int8_t{1});
+
+        auto q1 = add_quantize_op(m1, "quantizelinear", x, scale, zero);
+        auto rs = m1.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 1024}}}), q1);
+        auto d1 = add_quantize_op(m1, "dequantizelinear", rs, scale, zero);
+        auto dot = m1.add_instruction(migraphx::make_op("dot"), d1, w);
+        m1.add_return({dot});
+    }
+
+    migraphx::module m2 = m1;
+
+    run_pass(m1);
+    EXPECT(m1 == m2);
+}
+
 TEST_CASE(dot_reshaped)
 {
     migraphx::shape sh1{migraphx::shape::float_type, {1280, 1000}};

test/verify/test_uint8_rebias_roundtrip.cpp

@@ -0,0 +1,66 @@
+/*
+ * The MIT License (MIT)
+ *
+ * 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
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#include "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/make_op.hpp>
+
+// Isolates the uint8 -> int8 rebias chain simplify_qdq emits for asymmetric activations:
+// quantizelinear(uint8) -> convert(int32) -> sub(128) -> convert(int8), with no conv or MLIR.
+struct test_uint8_rebias_roundtrip : verify_program<test_uint8_rebias_roundtrip>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+
+        migraphx::shape x_shape{migraphx::shape::float_type, {2, 8}};
+        auto x = mm->add_parameter("x", x_shape);
+
+        auto scale = mm->add_literal(migraphx::literal{migraphx::shape::float_type, {0.05f}});
+        auto zp    = mm->add_literal(migraphx::literal{migraphx::shape::uint8_type, {0}});
+        auto scale_b = mm->add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", x_shape.lens()}}), scale);
+        auto zp_b = mm->add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", x_shape.lens()}}), zp);
+
+        auto q = mm->add_instruction(migraphx::make_op("quantizelinear"), x, scale_b, zp_b);
+
+        // Rebias to int8 by subtracting 128 through int32 (same IR as rebias_uint8_to_int8).
+        auto q_i32 = mm->add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::int32_type}}), q);
+        auto k128 = mm->add_literal(migraphx::literal{migraphx::shape::int32_type, {128}});
+        auto k128_b = mm->add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", x_shape.lens()}}), k128);
+        auto diff = mm->add_instruction(migraphx::make_op("sub"), q_i32, k128_b);
+        auto q_i8 = mm->add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::int8_type}}), diff);
+
+        auto out = mm->add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::float_type}}), q_i8);
+        mm->add_return({out});
+        return p;
+    }
+    std::string section() const { return "conv"; }
+};