# 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. # ruff: noqa: F841 import tvm import tvm.testing from tvm import te def test_cast(): analyzer = tvm.arith.Analyzer() x = tvm.tirx.Var("x", "int8") m = analyzer.modular_set((x * 3).astype("uint32")) assert m.coeff == 3 assert m.base == 0 m = analyzer.modular_set((x * 3 + 1).astype("float32").astype("int32")) assert m.coeff == 3 assert m.base == 1 def test_add_sub(): analyzer = tvm.arith.Analyzer() x, y = tvm.tirx.Var("x", "int64"), tvm.tirx.Var("y", "int64") m = analyzer.modular_set(x * 6 + y * 4) assert m.coeff == 2 assert m.base == 0 analyzer.bind(y, x * 4 + 1) m = analyzer.modular_set(1 - y) assert m.coeff == 4 assert m.base == 0 def test_mul(): analyzer = tvm.arith.Analyzer() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") m = analyzer.modular_set((x * 4 + 2) * (y * 6 + 1)) assert m.coeff == 4 assert m.base == 2 def test_shift_left(): analyzer = tvm.arith.Analyzer() x, y = te.var("x"), te.var("y") m = analyzer.modular_set((x * 4 + 2) << 2) assert m.coeff == 16 assert m.base == 8 def test_floormod(): analyzer = tvm.arith.Analyzer() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") m = analyzer.modular_set(tvm.tirx.floormod(x * 128 + y * 4, 256)) assert m.coeff == 4 assert m.base == 0 def test_div_shift(): analyzer = tvm.arith.Analyzer() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") # not sure if x is non-negative tdiv = tvm.tirx.truncdiv m = analyzer.modular_set(tdiv(x * 4 + 2, 2)) assert m.coeff == 1 assert m.base == 0 # right shift always round down so it is fine m = analyzer.modular_set((x * 4 + 2) >> 1) assert m.coeff == 2 assert m.base == 1 fld = tvm.tirx.floordiv m = analyzer.modular_set(fld(x * 4 + 2, 2)) assert m.coeff == 2 assert m.base == 1 # x is non-negative analyzer.update(x, tvm.arith.ConstIntBound(0, 100)) m = analyzer.modular_set(tdiv(x * 4 + 2, 2)) assert m.coeff == 2 assert m.base == 1 def test_mod(): analyzer = tvm.arith.Analyzer() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") tmod = tvm.tirx.truncmod fmod = tvm.tirx.floormod # not sure if x is non-negative m = analyzer.modular_set(tmod(x * 4 + 1, 4)) assert m.coeff == 1 assert m.base == 0 # no need to be positive if base == 0 m = analyzer.modular_set(tmod(x * 4, 4)) assert m.coeff == 4 assert m.base == 0 # floor mod tests m = analyzer.modular_set(fmod(x * 4 + 3, 2)) assert m.coeff == 2 assert m.base == 1 m = analyzer.modular_set(fmod(x * 4 + 3, 8)) assert m.coeff == 4 assert m.base == 3 # x is non-negative analyzer.update(x, tvm.arith.ConstIntBound(0, 100)) m = analyzer.modular_set(tmod(x * 4 + 3, 2)) assert m.coeff == 2 assert m.base == 1 def test_min_max_select(): analyzer = tvm.arith.Analyzer() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") m = analyzer.modular_set(tvm.tirx.min(x * 3, y * 9)) assert m.coeff == 3 assert m.base == 0 m = analyzer.modular_set(tvm.tirx.max(x * 3 + 1, y * 9 + 4)) assert m.coeff == 3 assert m.base == 1 m = analyzer.modular_set(tvm.tirx.Select(x > 0, x * 3 + 1, y * 9 + 2)) assert m.coeff == 1 assert m.base == 0 def test_mix_index(): a = tvm.tirx.Var("a", "int32") b = tvm.tirx.Var("b", "int32") analyzer = tvm.arith.Analyzer() tdiv = tvm.tirx.truncdiv m = analyzer.modular_set(a * 4 + b * 6 + 7) assert m.coeff == 2 assert m.base == 1 m = analyzer.modular_set((a * 4 + 1) * (b * 8 + 3)) assert m.coeff == 4 assert m.base == 3 m = analyzer.modular_set(tdiv(a * 4 + 1, b * 8 + 3)) assert m.coeff == 1 assert m.base == 0 m = analyzer.modular_set((a * 4 + 1) * tdiv(b * 8, 4)) assert m.coeff == 2 assert m.base == 0 m = analyzer.modular_set((a * 12 + 1) - (b * 3 * 7 + 2)) assert m.coeff == 3 assert m.base == 2 m = analyzer.modular_set(a * 12 + tvm.tirx.min(b * 3 * 7, 2)) assert m.coeff == 1 assert m.base == 0 def test_constraint_scope(): a = tvm.tirx.Var("a", "int32") b = tvm.tirx.Var("b", "int32") analyzer = tvm.arith.Analyzer() tmod = tvm.tirx.truncmod with analyzer.constraint_scope(tmod(b, 4) == 2): m = analyzer.modular_set(b + 1) assert m.coeff == 4 assert m.base == 3 with analyzer.constraint_scope(tmod(a, 2) == 1): m = analyzer.modular_set(b + a * 2) assert m.coeff == 4 assert m.base == 0 m = analyzer.modular_set(b + a * 2) assert m.coeff == 2 assert m.base == 0 m = analyzer.modular_set(b + 1) assert m.coeff == 1 assert m.base == 0 def test_intersect(): a = tvm.tirx.Var("a", "int32") analyzer = tvm.arith.Analyzer() tmod = tvm.tirx.truncmod with analyzer.constraint_scope(tmod(a, 4) == 1): with analyzer.constraint_scope(tmod(a, 3) == 1): m = analyzer.modular_set(a) assert m.coeff == 12 assert m.base == 1 with analyzer.constraint_scope(tmod(a, 3) == 2): with analyzer.constraint_scope(tmod(a, 5) == 3): with analyzer.constraint_scope(tmod(a, 7) == 2): m = analyzer.modular_set(a) assert m.coeff == 105 assert m.base == 23 def test_let(): analyzer = tvm.arith.Analyzer() x = tvm.tirx.Var("x", "int32") y = tvm.tirx.Var("y", "int32") m = analyzer.modular_set(tvm.tirx.Let(x, y * 10, x + 1)) assert m.coeff == 10 assert m.base == 1 def test_bitwise_and(): analyzer = tvm.arith.Analyzer() x = tvm.tirx.Var("x", "int32") y = tvm.tirx.Var("y", "int32") # RHS of bitwise_and is 2^p - 1 m = analyzer.modular_set((x * 16 + y * 4) & 31) assert m.coeff == 4 assert m.base == 0 # arbitrary RHS m = analyzer.modular_set((x * 16 + y * 4) & 17) assert m.coeff == 1 assert m.base == 0 if __name__ == "__main__": tvm.testing.main()