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

src/target/intrin_rule.cc

@@ -280,25 +280,25 @@ 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 int_node->value;
+      // Probe scalar or broadcast constants without rejecting runtime values.
+      auto get_int_value = [](const PrimExpr& node) {
+        if (const auto* broadcast = node.as<prim::BroadcastNode>()) {
+          return as_const_int(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 int_node->value;
+        return as_const_int(node);
       };
       // 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 int64_t* y_value = get_int_value(y);
+      const int64_t* q_value = get_int_value(q);
+      const int64_t* s_value = get_int_value(s);
+      if (y_value && *y_value == (1 << 30) && q_value && *q_value == 31 && s_value) {
         PrimExpr exp = s - 1;
-        int exp_val = get_int_value(s) - 1;
+        int64_t exp_val = *s_value - 1;
+        if (exp_val == 0) return x;
         if (exp_val > 0) {
-          // power of 2 is greater than 0, apply left shift.
+          // 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.
@@ -339,6 +339,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_q_multiply_shift.py

@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Minimal regressions for fixed-point multiply legalization."""
+
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm import tirx
+
+
+def _legalize(name, *args):
+    call = tirx.call_intrin("int32", name, *args)
+    return tvm.ir.Op.get(name).get_attr("default.FLegalize")(call)
+
+
+@pytest.mark.parametrize("runtime_arg", ["multiplier", "shift"])
+def test_q_multiply_shift_runtime_argument(runtime_arg):
+    value = tirx.Var("value", "int32")
+    multiplier = value if runtime_arg == "multiplier" else tirx.const(1 << 30, "int32")
+    shift = value if runtime_arg == "shift" else tirx.const(1, "int32")
+    lowered = _legalize("tirx.q_multiply_shift", 3, multiplier, 31, shift)
+    replacement = tirx.const(1 << 30 if runtime_arg == "multiplier" else 1, "int32")
+    lowered = tvm_ffi.structural_map(
+        lowered, (tirx.Var, lambda var: replacement if var.same_as(value) else var), order="post"
+    )
+    tvm.ir.assert_structural_equal(tvm.arith.Analyzer().simplify(lowered), tirx.const(3, "int32"))
+
+
+def test_q_multiply_shift_zero_exponent():
+    x = tirx.Var("x", "int32")
+    lowered = _legalize("tirx.q_multiply_shift", x, 1 << 30, 31, 1)
+    tvm.ir.assert_structural_equal(lowered, x)
+
+
+def test_q_multiply_shift_non_q31():
+    lowered = _legalize("tirx.q_multiply_shift", 3, 1 << 30, 30, 2)
+    tvm.ir.assert_structural_equal(tvm.arith.Analyzer().simplify(lowered), tirx.const(12, "int32"))
+
+
+def test_q_multiply_shift_per_axis_integer_flag():
+    lowered = _legalize("tirx.q_multiply_shift_per_axis", 3, 1 << 30, 2, 1, 31, 1, 1)
+    tvm.ir.assert_structural_equal(tvm.arith.Analyzer().simplify(lowered), tirx.const(3, "int32"))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()