From 227f9453f89925809b2a73c0f8afe343e114cd9d Mon Sep 17 00:00:00 2001 From: sepcnt <30561671+sepcnt@users.noreply.github.com> Date: Tue, 15 Sep 2026 13:22:30 +0800 Subject: [PATCH 1/2] [FIX][TIR] Fix fixed-point multiply legalization --- src/target/intrin_rule.cc | 25 +-- .../test_tir_transform_q_multiply_shift.py | 151 ++++++++++++++++++ 2 files changed, 164 insertions(+), 12 deletions(-) create mode 100644 tests/python/tirx-transform/test_tir_transform_q_multiply_shift.py diff --git a/src/target/intrin_rule.cc b/src/target/intrin_rule.cc index 22d2d97d42ce..2fe40e12edf5 100644 --- a/src/target/intrin_rule.cc +++ b/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 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 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()) { + return as_const_int(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 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 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_q_multiply_shift.py b/tests/python/tirx-transform/test_tir_transform_q_multiply_shift.py new file mode 100644 index 000000000000..667e2995141b --- /dev/null +++ b/tests/python/tirx-transform/test_tir_transform_q_multiply_shift.py @@ -0,0 +1,151 @@ +# 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. + +"""Regression tests for the default fixed-point multiply legalization rules.""" + +import pytest +import tvm_ffi + +import tvm +import tvm.testing +from tvm import tirx + + +def _const(value, lanes, dtype="int32"): + value = tirx.const(value, dtype) + return value if lanes == 1 else tirx.Broadcast(value, lanes) + + +def _legalize(name, lanes, *args): + dtype = "int32" if lanes == 1 else f"int32x{lanes}" + call = tirx.call_intrin(dtype, name, *args) + return tvm.ir.Op.get(name).get_attr("default.FLegalize")(call) + + +def _check_value(expr, bindings, expected): + substituted = tvm_ffi.structural_map( + expr, (tirx.Var, lambda var: bindings.get(var, var)), order="post" + ) + actual = tvm.arith.Analyzer().simplify(substituted) + tvm.ir.assert_structural_equal(actual, expected) + + +def _reference(x, y, q, left_shift, right_shift): + total_shift = q + right_shift + return (((x << left_shift) * y) + (1 << (total_shift - 1))) >> total_shift + + +@pytest.mark.parametrize( + ("multiplier", "shift", "q"), + [ + (None, 1, 31), + (None, None, 31), + (12345, None, 31), + (1 << 30, None, 31), + (1 << 30, -1, 31), + (1 << 30, 0, 31), + (1 << 30, 1, 31), + (1 << 30, 2, 31), + (1 << 30, 1, 30), + ], +) +def test_q_multiply_shift(multiplier, shift, q): + lanes = 1 + x, y, s = [tirx.Var(name, "int32") for name in ("x", "y", "s")] + lowered = _legalize( + "tirx.q_multiply_shift", + lanes, + x, + y if multiplier is None else _const(multiplier, lanes), + _const(q, lanes), + s if shift is None else _const(shift, lanes), + ) + if multiplier == 1 << 30 and q == 31 and shift == 1: + tvm.ir.assert_structural_equal(lowered, x) + for x_value, y_value, s_value in [ + (-1001, 1 << 30, -1), + (-3, 12345, 0), + (3, -(1 << 29), 1), + (1001, (1 << 30) + 1, 2), + ]: + effective_y = y_value if multiplier is None else multiplier + effective_s = s_value if shift is None else shift + expected = _reference(x_value, effective_y, q, max(effective_s, 0), max(-effective_s, 0)) + bindings = {x: _const(x_value, lanes), y: _const(y_value, lanes), s: _const(s_value, lanes)} + _check_value(lowered, bindings, _const(expected, lanes)) + + +@pytest.mark.parametrize("flag_dtype", ["int32", "bool"]) +@pytest.mark.parametrize("flag", [0, 1]) +def test_q_multiply_shift_per_axis(flag, flag_dtype): + lanes = 1 + x, y = [tirx.Var(name, "int32") for name in ("x", "y")] + lowered = _legalize( + "tirx.q_multiply_shift_per_axis", + lanes, + x, + y, + _const(2, lanes), + _const(1, lanes), + _const(31, lanes), + _const(flag, lanes, flag_dtype), + _const(1, lanes, flag_dtype), + ) + for x_value in [-1001, -3, 0, 3, 1001]: + expected = _reference(x_value, 1 << 30, 31, 2 if flag else 0, 1) + _check_value( + lowered, {x: _const(x_value, lanes), y: _const(1 << 30, lanes)}, _const(expected, lanes) + ) + + +@pytest.mark.parametrize("lanes", [2, 4]) +@pytest.mark.parametrize("case", ["runtime_y", "runtime_s", "broadcast_y", "zero_exponent"]) +def test_q_multiply_shift_vector(case, lanes): + x, y, s = [tirx.Var(name, f"int32x{lanes}") for name in ("x", "y", "s")] + if case == "broadcast_y": + y = tirx.Broadcast(tirx.Var("multiplier", "int32"), lanes) + elif case != "runtime_y": + y = _const(1 << 30, lanes) + if case != "runtime_s": + s = _const(1, lanes) + lowered = _legalize("tirx.q_multiply_shift", lanes, x, y, _const(31, lanes), s) + tvm.ir.assert_structural_equal(lowered.ty, x.ty) + if case == "zero_exponent": + tvm.ir.assert_structural_equal(lowered, x) + + +@pytest.mark.parametrize("lanes", [2, 4]) +@pytest.mark.parametrize("flag", [0, 1]) +def test_q_multiply_shift_per_axis_vector(flag, lanes): + x, y = [tirx.Var(name, f"int32x{lanes}") for name in ("x", "y")] + args = [x, y, _const(2, lanes), _const(1, lanes), _const(31, lanes)] + lowered = _legalize( + "tirx.q_multiply_shift_per_axis", lanes, *args, _const(flag, lanes), _const(1, lanes) + ) + expected = _legalize( + "tirx.q_multiply_shift_per_axis", + lanes, + *args, + _const(flag, lanes, "bool"), + _const(1, lanes, "bool"), + ) + analyzer = tvm.arith.Analyzer() + tvm.ir.assert_structural_equal(analyzer.simplify(lowered), analyzer.simplify(expected)) + + +if __name__ == "__main__": + tvm.testing.main() From 7636ce46c517ab19a56479aa3effaa6a10cf0052 Mon Sep 17 00:00:00 2001 From: sepcnt <30561671+sepcnt@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:49:59 +0800 Subject: [PATCH 2/2] [TEST][TIR] Minimize fixed-point multiply regression coverage --- .../test_tir_transform_q_multiply_shift.py | 134 +++--------------- 1 file changed, 23 insertions(+), 111 deletions(-) diff --git a/tests/python/tirx-transform/test_tir_transform_q_multiply_shift.py b/tests/python/tirx-transform/test_tir_transform_q_multiply_shift.py index 667e2995141b..99e1d7632aa2 100644 --- a/tests/python/tirx-transform/test_tir_transform_q_multiply_shift.py +++ b/tests/python/tirx-transform/test_tir_transform_q_multiply_shift.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -"""Regression tests for the default fixed-point multiply legalization rules.""" +"""Minimal regressions for fixed-point multiply legalization.""" import pytest import tvm_ffi @@ -25,126 +25,38 @@ from tvm import tirx -def _const(value, lanes, dtype="int32"): - value = tirx.const(value, dtype) - return value if lanes == 1 else tirx.Broadcast(value, lanes) - - -def _legalize(name, lanes, *args): - dtype = "int32" if lanes == 1 else f"int32x{lanes}" - call = tirx.call_intrin(dtype, name, *args) +def _legalize(name, *args): + call = tirx.call_intrin("int32", name, *args) return tvm.ir.Op.get(name).get_attr("default.FLegalize")(call) -def _check_value(expr, bindings, expected): - substituted = tvm_ffi.structural_map( - expr, (tirx.Var, lambda var: bindings.get(var, var)), order="post" - ) - actual = tvm.arith.Analyzer().simplify(substituted) - tvm.ir.assert_structural_equal(actual, expected) - - -def _reference(x, y, q, left_shift, right_shift): - total_shift = q + right_shift - return (((x << left_shift) * y) + (1 << (total_shift - 1))) >> total_shift - - -@pytest.mark.parametrize( - ("multiplier", "shift", "q"), - [ - (None, 1, 31), - (None, None, 31), - (12345, None, 31), - (1 << 30, None, 31), - (1 << 30, -1, 31), - (1 << 30, 0, 31), - (1 << 30, 1, 31), - (1 << 30, 2, 31), - (1 << 30, 1, 30), - ], -) -def test_q_multiply_shift(multiplier, shift, q): - lanes = 1 - x, y, s = [tirx.Var(name, "int32") for name in ("x", "y", "s")] - lowered = _legalize( - "tirx.q_multiply_shift", - lanes, - x, - y if multiplier is None else _const(multiplier, lanes), - _const(q, lanes), - s if shift is None else _const(shift, lanes), +@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" ) - if multiplier == 1 << 30 and q == 31 and shift == 1: - tvm.ir.assert_structural_equal(lowered, x) - for x_value, y_value, s_value in [ - (-1001, 1 << 30, -1), - (-3, 12345, 0), - (3, -(1 << 29), 1), - (1001, (1 << 30) + 1, 2), - ]: - effective_y = y_value if multiplier is None else multiplier - effective_s = s_value if shift is None else shift - expected = _reference(x_value, effective_y, q, max(effective_s, 0), max(-effective_s, 0)) - bindings = {x: _const(x_value, lanes), y: _const(y_value, lanes), s: _const(s_value, lanes)} - _check_value(lowered, bindings, _const(expected, lanes)) + tvm.ir.assert_structural_equal(tvm.arith.Analyzer().simplify(lowered), tirx.const(3, "int32")) -@pytest.mark.parametrize("flag_dtype", ["int32", "bool"]) -@pytest.mark.parametrize("flag", [0, 1]) -def test_q_multiply_shift_per_axis(flag, flag_dtype): - lanes = 1 - x, y = [tirx.Var(name, "int32") for name in ("x", "y")] - lowered = _legalize( - "tirx.q_multiply_shift_per_axis", - lanes, - x, - y, - _const(2, lanes), - _const(1, lanes), - _const(31, lanes), - _const(flag, lanes, flag_dtype), - _const(1, lanes, flag_dtype), - ) - for x_value in [-1001, -3, 0, 3, 1001]: - expected = _reference(x_value, 1 << 30, 31, 2 if flag else 0, 1) - _check_value( - lowered, {x: _const(x_value, lanes), y: _const(1 << 30, lanes)}, _const(expected, lanes) - ) +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) -@pytest.mark.parametrize("lanes", [2, 4]) -@pytest.mark.parametrize("case", ["runtime_y", "runtime_s", "broadcast_y", "zero_exponent"]) -def test_q_multiply_shift_vector(case, lanes): - x, y, s = [tirx.Var(name, f"int32x{lanes}") for name in ("x", "y", "s")] - if case == "broadcast_y": - y = tirx.Broadcast(tirx.Var("multiplier", "int32"), lanes) - elif case != "runtime_y": - y = _const(1 << 30, lanes) - if case != "runtime_s": - s = _const(1, lanes) - lowered = _legalize("tirx.q_multiply_shift", lanes, x, y, _const(31, lanes), s) - tvm.ir.assert_structural_equal(lowered.ty, x.ty) - if case == "zero_exponent": - 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")) -@pytest.mark.parametrize("lanes", [2, 4]) -@pytest.mark.parametrize("flag", [0, 1]) -def test_q_multiply_shift_per_axis_vector(flag, lanes): - x, y = [tirx.Var(name, f"int32x{lanes}") for name in ("x", "y")] - args = [x, y, _const(2, lanes), _const(1, lanes), _const(31, lanes)] - lowered = _legalize( - "tirx.q_multiply_shift_per_axis", lanes, *args, _const(flag, lanes), _const(1, lanes) - ) - expected = _legalize( - "tirx.q_multiply_shift_per_axis", - lanes, - *args, - _const(flag, lanes, "bool"), - _const(1, lanes, "bool"), - ) - analyzer = tvm.arith.Analyzer() - tvm.ir.assert_structural_equal(analyzer.simplify(lowered), analyzer.simplify(expected)) +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__":