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..99e1d7632aa2 --- /dev/null +++ b/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()