diff --git a/src/target/intrin_rule.cc b/src/target/intrin_rule.cc index 0d59b1a41fa8..94cee2c9b713 100644 --- a/src/target/intrin_rule.cc +++ b/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 s = call->args[3].as_or_throw(); - // Lambda function to extract the int value from PrimExpr - auto get_int_value = [](const PrimExpr node) { - if (auto int_node = node.as()) { - return static_cast(int_node->value); + // Probe scalar or broadcast constants without rejecting runtime values. + auto get_int_imm = [](PrimExpr node) { + if (const auto* broadcast = node.as()) { + node = broadcast->value; } - auto broadcast_node = node.as(); - TVM_FFI_ICHECK(broadcast_node != nullptr); - auto int_node = broadcast_node->value.as(); - TVM_FFI_ICHECK(int_node != nullptr); - return static_cast(int_node->value); + return node.as(); }; // 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 q = call->args[4].as_or_throw(); PrimExpr is_lshift_required = call->args[5].as_or_throw(); + 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]; diff --git a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py index 98cb476b19ae..a67a7d8e6da8 100644 --- a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py +++ b/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()