# 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. import tvm import numpy as np def lower_intrin(stmt): """wrapper to call transformation in stmt""" lower_expr = isinstance(stmt, tvm.expr.Expr) stmt = tvm.stmt.Evaluate(stmt) if lower_expr else stmt stmt = tvm.ir_pass.CanonicalSimplify(stmt) stmt = tvm.ir_pass._LowerIntrinStmt(stmt, "llvm") return stmt.value if lower_expr else stmt.body def check_value(expr, vx, vy, data, fref): n = len(data) A = tvm.placeholder((n,), name="A", dtype=expr.dtype) B = tvm.placeholder((n,), name="B", dtype=expr.dtype) def make_binds(i): x = expr x = tvm.expr.Let(vx, A[i], x) x = tvm.expr.Let(vy, B[i], x) return x C = tvm.compute((n,), make_binds) s = tvm.create_schedule([C.op]) if not tvm.module.enabled("llvm"): return f = tvm.build(s, [A, B, C], "llvm") a = tvm.nd.array(np.array([x for x, y in data], dtype=expr.dtype)) b = tvm.nd.array(np.array([y for x, y in data], dtype=expr.dtype)) c = tvm.nd.array(np.zeros(len(data), dtype=expr.dtype)) f(a, b, c) cref = np.array([fref(x, y) for x, y in data]) np.testing.assert_equal(c.asnumpy(), cref) def get_ref_data(): """Get reference data for every pairs""" import itertools x = range(-10, 10) y = list(range(-10, 10)) y.remove(0) return list(itertools.product(x, y)) def test_lower_floordiv(): data = get_ref_data() for dtype in ["int32", "int64", "int16"]: x = tvm.var("x", dtype=dtype) y = tvm.var("y", dtype=dtype) zero = tvm.const(0, dtype) # no constraints res = lower_intrin(tvm.floordiv(x, y)) check_value(res, x, y, data, lambda a, b: a // b) # rhs >= 0 res = lower_intrin(tvm.expr.Select(y >= 0, tvm.floordiv(x, y), zero)) check_value(res, x, y, data, lambda a, b: a // b if b > 0 else 0) # involves max res = lower_intrin(tvm.expr.Select(y >= 0, tvm.max(tvm.floordiv(x, y), zero), zero)) check_value(res, x, y, data, lambda a, b: max(a // b, 0) if b > 0 else 0) # lhs >= 0 res = lower_intrin(tvm.expr.Select(tvm.all(y >= 0, x >= 0), tvm.floordiv(x, y), zero)) check_value(res, x, y, data, lambda a, b: a // b if b > 0 and a >= 0 else 0) # const power of two res = lower_intrin(tvm.floordiv(x, tvm.const(8, dtype=dtype))) check_value(res, x, y, [(a, b) for a, b in data if b == 8], lambda a, b: a // b) def test_lower_floormod(): data = get_ref_data() for dtype in ["int32", "int64", "int16"]: x = tvm.var("x", dtype=dtype) y = tvm.var("y", dtype=dtype) zero = tvm.const(0, dtype) # no constraints res = lower_intrin(tvm.floormod(x, y)) check_value(res, x, y, data, lambda a, b: a % b) # rhs >= 0 res = lower_intrin(tvm.expr.Select(y >= 0, tvm.floormod(x, y), zero)) check_value(res, x, y, data, lambda a, b: a % b if b > 0 else 0) # lhs >= 0 res = lower_intrin(tvm.expr.Select(tvm.all(y >= 0, x >= 0), tvm.floormod(x, y), zero)) check_value(res, x, y, data, lambda a, b: a % b if b > 0 and a >= 0 else 0) # const power of two res = lower_intrin(tvm.floormod(x, tvm.const(8, dtype=dtype))) check_value(res, x, y, [(a, b) for a, b in data if b == 8], lambda a, b: a % b) if __name__ == "__main__": test_lower_floordiv() test_lower_floormod()