[FIX][TIR] Fix fixed-point multiply legalization

src/target/intrin_rule.cc

@@ -282,25 +282,24 @@ TVM_REGISTER_OP("tirx.q_multiply_shift")
       PrimExpr q = call->args[2].as_or_throw<PrimExpr>();
       PrimExpr s = call->args[3].as_or_throw<PrimExpr>();
 
-      // Lambda function to extract the int value from PrimExpr
-      auto get_int_value = [](const PrimExpr node) {
-        if (auto int_node = node.as<IntImmNode>()) {
-          return static_cast<int64_t>(int_node->value);
+      // Probe scalar or broadcast constants without rejecting runtime values.
+      auto get_int_imm = [](PrimExpr node) {
+        if (const auto* broadcast = node.as<prim::BroadcastNode>()) {
+          node = broadcast->value;
         }
-        auto broadcast_node = node.as<prim::BroadcastNode>();
-        TVM_FFI_ICHECK(broadcast_node != nullptr);
-        auto int_node = broadcast_node->value.as<IntImmNode>();
-        TVM_FFI_ICHECK(int_node != nullptr);
-        return static_cast<int64_t>(int_node->value);
+        return node.as<IntImmNode>();
       };
       // Power of 2 is determined by the fixed_point_multiplier == 1 << 30. In case of power of
       // 2, fixed point multiplier will represent a float value of 0.5. In fixed point, this is
       // represented by 1 << 30.
-      if (get_int_value(y) == (1 << 30)) {
+      const auto* y_value = get_int_imm(y);
+      const auto* q_value = get_int_imm(q);
+      const auto* s_value = get_int_imm(s);
+      if (y_value && y_value->value == (1 << 30) && q_value && q_value->value == 31 && s_value) {
+        if (s_value->value == 1) return x;
         PrimExpr exp = s - 1;
-        int exp_val = get_int_value(s) - 1;
-        if (exp_val > 0) {
-          // power of 2 is greater than 0, apply left shift.
+        if (s_value->value > 1) {
+          // A positive exponent only needs a left shift.
           return x << exp;
         } else {
           // power of 2 is less than 0, round and then apply right shift.
@@ -341,6 +340,7 @@ TVM_REGISTER_OP("tirx.q_multiply_shift_per_axis")
       PrimExpr right_shift = call->args[3].as_or_throw<PrimExpr>();
       PrimExpr q = call->args[4].as_or_throw<PrimExpr>();
       PrimExpr is_lshift_required = call->args[5].as_or_throw<PrimExpr>();
+      is_lshift_required = is_lshift_required != MakeConst(is_lshift_required.ty(), 0);
       // Note, 7th argument is "is_rshift_required" flag, but we don't need that here.
       // PrimExpr is_rshift_required = call->args[6];
 

tests/python/tirx-transform/test_tir_transform_lower_intrin.py

@@ -316,6 +316,44 @@ def test_lower_floordiv_overflow_checks():
     check_value(res, [x], data_check5, lambda a: (min(max(a, -10), 32758)) // 3)
 
 
+@pytest.mark.parametrize("runtime_arg", [1, 3], ids=["multiplier", "shift"])
+def test_q_multiply_shift_runtime_argument(runtime_arg):
+    value = tvm.tirx.Var("value", "int32")
+    args = [tvm.tirx.const(arg, "int32") for arg in [3, 1 << 30, 31, 1]]
+    replacement = args[runtime_arg]
+    args[runtime_arg] = value
+    lowered = lower_intrin([value], tvm.tirx.call_intrin("int32", "tirx.q_multiply_shift", *args))
+    lowered = tvm_ffi.structural_map(
+        lowered, (tvm.tirx.Var, lambda var: replacement if var.same_as(value) else var)
+    )
+    tvm.ir.assert_structural_equal(tvm.sym.Analyzer().simplify(lowered), tvm.tirx.const(3, "int32"))
+
+
+def test_q_multiply_shift_zero_exponent():
+    x = tvm.tirx.Var("x", "int32")
+    lowered = lower_intrin(
+        [x], tvm.tirx.call_intrin("int32", "tirx.q_multiply_shift", x, 1 << 30, 31, 1)
+    )
+    tvm.ir.assert_structural_equal(lowered, x)
+
+
+def test_q_multiply_shift_non_q31():
+    lowered = lower_intrin(
+        [], tvm.tirx.call_intrin("int32", "tirx.q_multiply_shift", 3, 1 << 30, 30, 2)
+    )
+    tvm.ir.assert_structural_equal(
+        tvm.sym.Analyzer().simplify(lowered), tvm.tirx.const(12, "int32")
+    )
+
+
+def test_q_multiply_shift_per_axis_integer_flag():
+    lowered = lower_intrin(
+        [],
+        tvm.tirx.call_intrin("int32", "tirx.q_multiply_shift_per_axis", 3, 1 << 30, 2, 1, 31, 1, 1),
+    )
+    tvm.ir.assert_structural_equal(tvm.sym.Analyzer().simplify(lowered), tvm.tirx.const(3, "int32"))
+
+
 if __name__ == "__main__":
     test_lower_floordiv()
     test_lower_floormod()