| # 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 numpy as np |
| import tvm |
| import tvm.testing |
| from tvm import relay |
| from tvm.relay.backend import Executor |
| from tvm.relay import transform |
| from tvm.relay.build_module import bind_params_by_name |
| from tvm.relay.testing import run_infer_type, create_workload |
| |
| |
| def annot_expr(e): |
| """Returns e wrapped with an on_device annotation.""" |
| return relay.op.annotation.on_device(e, tvm.cpu(), constrain_result=True) |
| |
| |
| def run_opt_pass(expr, opt_pass): |
| assert isinstance(opt_pass, tvm.transform.Pass) |
| |
| mod = tvm.IRModule.from_expr(expr) |
| mod = relay.transform.InferType()(mod) |
| mod = opt_pass(mod) |
| entry = mod["main"] |
| return entry if isinstance(expr, relay.Function) else entry.body |
| |
| |
| def test_concatenate_const(): |
| def before(): |
| data = tvm.nd.array(np.array([1.0, 2.0, 3.0])) |
| const = relay.const(data) |
| concat = relay.op.concatenate([const, const], axis=0) |
| func = relay.Function([], concat) |
| return func |
| |
| def expected(): |
| data = tvm.nd.array(np.array([1.0, 2.0, 3.0, 1.0, 2.0, 3.0])) |
| const = relay.const(data) |
| func = relay.Function([], const) |
| return func |
| |
| zz = run_opt_pass(before(), transform.FoldConstant()) |
| zexpected = run_opt_pass(expected(), transform.InferType()) |
| assert tvm.ir.structural_equal(zz, zexpected) |
| |
| |
| def test_fold_const(): |
| c_data = np.array([1, 2, 3]).astype("float32") |
| t = relay.TensorType([1, 2, 3], "float32") |
| |
| def before(): |
| c = relay.const(c_data) |
| x = relay.var("x", t) |
| y = relay.add(c, c) |
| y = relay.multiply(y, relay.const(2, "float32")) |
| y = relay.add(x, y) |
| z = relay.add(y, c) |
| return relay.Function([x], z) |
| |
| def expected(): |
| x = relay.var("x", t) |
| c_folded = (c_data + c_data) * 2 |
| y = relay.add(x, relay.const(c_folded)) |
| z = relay.add(y, relay.const(c_data)) |
| return relay.Function([x], z) |
| |
| # the fold constant should work on any context. |
| with tvm.target.Target("cuda"): |
| zz = run_opt_pass(before(), transform.FoldConstant()) |
| zexpected = run_opt_pass(expected(), transform.InferType()) |
| tvm.ir.assert_structural_equal(zz, zexpected) |
| |
| |
| def test_fold_const_with_on_device(): |
| """Make sure on_device annotations don't get in the way of constant folding""" |
| c_data = np.array([1, 2, 3]).astype("float32") |
| t = relay.TensorType([1, 2, 3], "float32") |
| |
| def before(): |
| c = relay.const(c_data) |
| x = relay.var("x", t) |
| x.virtual_device_ = tvm.cpu() |
| y = relay.add(c, c) |
| y = relay.multiply(y, relay.const(2, "float32")) |
| y = relay.add(x, y) |
| z = relay.add(y, c) |
| f = relay.Function([x], z) |
| f.virtual_device_ = tvm.cpu() |
| return f |
| |
| def expected(): |
| x = relay.var("x", t) |
| x.virtual_device_ = tvm.cpu() |
| c_folded = (c_data + c_data) * 2 |
| y = relay.add(x, relay.const(c_folded)) |
| z = relay.add(y, relay.const(c_data)) |
| f = relay.Function([x], z) |
| f.virtual_device_ = tvm.cpu() |
| return f |
| |
| zz = run_opt_pass(before(), transform.FoldConstant()) |
| zexpected = run_opt_pass(expected(), transform.InferType()) |
| tvm.ir.assert_structural_equal(zz, zexpected) |
| |
| |
| def test_fold_let(): |
| c_data = np.array(1).astype("float32") |
| t = relay.TensorType([1], "float32") |
| |
| def before(): |
| sb = relay.ScopeBuilder() |
| x = relay.var("x", t) |
| t1 = sb.let("t1", relay.const(c_data)) |
| t2 = sb.let("t2", relay.add(t1, t1)) |
| t3 = sb.let("t3", relay.add(t2, x)) |
| sb.ret(t3) |
| return relay.Function([x], sb.get()) |
| |
| def expected(): |
| sb = relay.ScopeBuilder() |
| x = relay.var("x", t) |
| c_folded = c_data + c_data |
| t3 = sb.let("t3", relay.add(relay.const(c_folded), x)) |
| sb.ret(t3) |
| return relay.Function([x], sb.get()) |
| |
| zz = run_opt_pass(before(), transform.FoldConstant()) |
| zexpected = run_opt_pass(expected(), transform.InferType()) |
| tvm.ir.assert_structural_equal(zz, zexpected) |
| |
| |
| def test_fold_let_with_on_device(): |
| """Make sure on_device annotations don't get in the way of constant folding, |
| and inlined constants bring their annotations with them.""" |
| c_data = np.array(1).astype("float32") |
| t = relay.TensorType([1], "float32") |
| |
| def before(): |
| sb = relay.ScopeBuilder() |
| x = relay.var("x", t) |
| x.virtual_device_ = tvm.cpu() |
| t1 = sb.let("t1", annot_expr(relay.const(c_data))) |
| t2 = sb.let("t2", annot_expr(relay.add(t1, t1))) |
| t3 = sb.let("t3", annot_expr(relay.add(t2, x))) |
| sb.ret(t3) |
| f = relay.Function([x], sb.get()) |
| f.virtual_device_ = tvm.cpu() |
| return f |
| |
| def expected(): |
| sb = relay.ScopeBuilder() |
| x = relay.var("x", t) |
| x.virtual_device_ = tvm.cpu() |
| c_folded = c_data + c_data |
| t3 = sb.let("t3", annot_expr(relay.add(annot_expr(relay.const(c_folded)), x))) |
| sb.ret(t3) |
| f = relay.Function([x], sb.get()) |
| f.virtual_device_ = tvm.cpu() |
| return f |
| |
| zz = run_opt_pass(before(), transform.FoldConstant()) |
| zexpected = run_opt_pass(expected(), transform.InferType()) |
| tvm.ir.assert_structural_equal(zz, zexpected) |
| |
| |
| def test_fold_tuple(): |
| c_data = np.array(1).astype("float32") |
| t = relay.TensorType([1], "float32") |
| |
| def before(): |
| c = relay.const(c_data) |
| x = relay.var("x", t) |
| y = relay.Tuple([x, c]) |
| z = relay.add(y[1], c) |
| z = relay.add(z, y[0]) |
| return relay.Function([x], z) |
| |
| def expected(): |
| c = relay.const(c_data + c_data) |
| x = relay.var("x", t) |
| z = relay.add(c, x) |
| return relay.Function([x], z) |
| |
| zz = run_opt_pass(before(), transform.FoldConstant()) |
| zexpected = run_opt_pass(expected(), transform.InferType()) |
| tvm.ir.assert_structural_equal(zz, zexpected) |
| |
| |
| def test_fold_concat(): |
| c_data = np.array([[1, 2, 3]]).astype("float32") |
| |
| def before(): |
| a = relay.const(c_data) |
| b = relay.const(c_data) |
| y = relay.concatenate((a, b), axis=0) |
| return relay.Function([], y) |
| |
| def expected(): |
| y_data = np.concatenate((c_data, c_data), axis=0) |
| y = relay.const(y_data) |
| return relay.Function([], y) |
| |
| zz = run_opt_pass(before(), transform.FoldConstant()) |
| zexpected = run_opt_pass(expected(), transform.InferType()) |
| tvm.ir.assert_structural_equal(zz, zexpected) |
| |
| |
| def test_fold_if(): |
| cond_data = np.array(1).astype("bool") |
| x_data = np.array([[1, 2, 3]]).astype("float32") |
| |
| def before(): |
| a = relay.const(cond_data) |
| x = relay.const(x_data) |
| y = relay.const(x_data) |
| iff = relay.If(a, x + y, x - y) |
| return relay.Function([], iff) |
| |
| def expected(): |
| y_data = x_data + x_data |
| y = relay.const(y_data) |
| return relay.Function([], y) |
| |
| zz = run_opt_pass(before(), transform.FoldConstant()) |
| zexpected = run_opt_pass(expected(), transform.InferType()) |
| tvm.ir.assert_structural_equal(zz, zexpected) |
| |
| cond_data = np.array(0).astype("bool") |
| |
| def before(): |
| a = relay.const(cond_data) |
| x = relay.const(x_data) |
| y = relay.const(x_data) |
| iff = relay.If(a, x + y, x - y) |
| return relay.Function([], iff) |
| |
| def expected(): |
| y_data = x_data - x_data |
| y = relay.const(y_data) |
| return relay.Function([], y) |
| |
| zz = run_opt_pass(before(), transform.FoldConstant()) |
| zexpected = run_opt_pass(expected(), transform.InferType()) |
| tvm.ir.assert_structural_equal(zz, zexpected) |
| |
| |
| def test_fold_shape_of(): |
| c_shape = (8, 9, 10) |
| |
| def before(dtype): |
| x = relay.var("x", shape=c_shape, dtype="float32") |
| y = relay.var("y", shape=c_shape, dtype="float32") |
| z = relay.shape_of(x + y, dtype) |
| return relay.Function([x, y], z) |
| |
| def expected(dtype): |
| x = relay.var("x", shape=c_shape, dtype="float32") |
| y = relay.var("y", shape=c_shape, dtype="float32") |
| z = relay.const(np.array(c_shape).astype(dtype), dtype=dtype) |
| func = relay.Function([x, y], z) |
| return func |
| |
| for dtype in ["int32", "float32"]: |
| zz = run_opt_pass(before(dtype), transform.FoldConstant()) |
| zexpected = run_opt_pass(expected(dtype), transform.InferType()) |
| tvm.ir.assert_structural_equal(zz, zexpected) |
| |
| |
| def test_fold_ndarray_size(): |
| c_shape = (8, 9, 10) |
| |
| def before(dtype): |
| x = relay.var("x", shape=c_shape, dtype="float32") |
| y = relay.var("y", shape=c_shape, dtype="float32") |
| z = relay.ndarray_size(x + y, dtype) |
| return relay.Function([x, y], z) |
| |
| def expected(dtype): |
| x = relay.var("x", shape=c_shape, dtype="float32") |
| y = relay.var("y", shape=c_shape, dtype="float32") |
| z = relay.const(np.size(np.zeros(c_shape)), dtype=dtype) |
| func = relay.Function([x, y], z) |
| mod = tvm.IRModule.from_expr(func) |
| return mod["main"] |
| |
| for dtype in ["int32", "float32"]: |
| zz = run_opt_pass(before(dtype), transform.FoldConstant()) |
| zexpected = run_opt_pass(expected(dtype), transform.InferType()) |
| tvm.ir.assert_structural_equal(zz, zexpected) |
| |
| |
| def test_fold_batch_norm(): |
| def expected(): |
| data = relay.var("data", relay.TensorType((1, 3, 224, 224), "float32")) |
| weight = relay.const(np.zeros((16, 3, 3, 3))) |
| bias = relay.const(np.zeros((16, 1, 1))) |
| conv = relay.nn.conv2d( |
| data=data, weight=weight, kernel_size=(3, 3), channels=16, padding=(1, 1) |
| ) |
| add = relay.add(conv, bias) |
| return relay.Function(relay.analysis.free_vars(add), add) |
| |
| remove_bn_pass = tvm.transform.Sequential( |
| [ |
| relay.transform.InferType(), |
| relay.transform.SimplifyInference(), |
| relay.transform.FoldConstant(), |
| relay.transform.FoldScaleAxis(), |
| ] |
| ) |
| |
| data = relay.var("data", relay.TensorType((1, 3, 224, 224), "float32")) |
| weight = relay.var("weight") |
| bn_gamma = relay.var("bn_gamma") |
| bn_beta = relay.var("bn_beta") |
| bn_mmean = relay.var("bn_mean") |
| bn_mvar = relay.var("bn_var") |
| |
| conv = relay.nn.conv2d( |
| data=data, weight=weight, kernel_size=(3, 3), channels=16, padding=(1, 1) |
| ) |
| bn_output = relay.nn.batch_norm(conv, bn_gamma, bn_beta, bn_mmean, bn_mvar) |
| |
| def initializer(_, param): |
| param = np.zeros(param.shape) |
| |
| mod, params = create_workload(bn_output[0], initializer) |
| mod["main"] = bind_params_by_name(mod["main"], params) |
| |
| with tvm.transform.PassContext(opt_level=3): |
| mod = remove_bn_pass(mod) |
| |
| expect = run_infer_type(expected()) |
| tvm.ir.assert_structural_equal(mod["main"], expect) |
| |
| |
| def test_fold_dropout(): |
| def before(): |
| # A constant graph to fire fold constant |
| data = relay.const(np.arange(10).astype(np.float32)) |
| dropout = relay.nn.dropout(data) |
| add = dropout + relay.const(1.0) |
| return relay.Function(relay.analysis.free_vars(add), add) |
| |
| passes = tvm.transform.Sequential( |
| [ |
| relay.transform.InferType(), |
| relay.transform.FoldConstant(), |
| ] |
| ) |
| |
| before_mod = tvm.IRModule.from_expr(before()) |
| |
| with tvm.transform.PassContext(opt_level=3): |
| after_mod = passes(before_mod) |
| |
| tvm.ir.assert_structural_equal(run_infer_type(before_mod["main"]), after_mod["main"]) |
| |
| |
| def test_fold_qnn_const(): |
| def before(): |
| # QNN op with 2 constant arguments. |
| add = relay.qnn.op.add( |
| relay.const(np.ones((2, 3), dtype="uint8"), dtype="uint8"), |
| relay.const(np.ones((2, 3), dtype="uint8"), dtype="uint8"), |
| lhs_scale=relay.const(2.0), |
| lhs_zero_point=relay.const(0), |
| rhs_scale=relay.const(2.0), |
| rhs_zero_point=relay.const(0), |
| output_scale=relay.const(1.0), |
| output_zero_point=relay.const(0), |
| ) |
| # QNN op with 1 constant and 1 non-constant arguments. |
| a = relay.var("a", shape=[2, 3], dtype="float32") |
| dense = relay.qnn.op.dense( |
| relay.qnn.op.quantize(a, relay.const(1.0), relay.const(0)), |
| add, |
| input_zero_point=relay.const(0), |
| kernel_zero_point=relay.const(0), |
| input_scale=relay.const(2.0), |
| kernel_scale=relay.const(2.0), |
| units=None, |
| ) |
| # QNN op with 2 non-constant arguments. |
| b = relay.var("b", shape=[2], dtype="float32") |
| bias = relay.qnn.op.add( |
| dense, |
| relay.qnn.op.quantize(b, relay.const(1.0), relay.const(0), out_dtype="int32"), |
| lhs_scale=relay.const(2.0), |
| lhs_zero_point=relay.const(0), |
| rhs_scale=relay.const(2.0), |
| rhs_zero_point=relay.const(0), |
| output_scale=relay.const(1.0), |
| output_zero_point=relay.const(0), |
| ) |
| return relay.Function([a, b], bias) |
| |
| def expected(): |
| a = relay.var("a", shape=[2, 3], dtype="float32") |
| dense = relay.qnn.op.dense( |
| relay.qnn.op.quantize(a, relay.const(1.0), relay.const(0)), |
| relay.const(np.array([[4, 4, 4], [4, 4, 4]], dtype="uint8"), dtype="uint8"), |
| input_zero_point=relay.const(0), |
| kernel_zero_point=relay.const(0), |
| input_scale=relay.const(2.0), |
| kernel_scale=relay.const(2.0), |
| units=None, |
| ) |
| b = relay.var("b", shape=[2], dtype="float32") |
| bias = relay.qnn.op.add( |
| dense, |
| relay.qnn.op.quantize(b, relay.const(1.0), relay.const(0), out_dtype="int32"), |
| lhs_scale=relay.const(2.0), |
| lhs_zero_point=relay.const(0), |
| rhs_scale=relay.const(2.0), |
| rhs_zero_point=relay.const(0), |
| output_scale=relay.const(1.0), |
| output_zero_point=relay.const(0), |
| ) |
| return relay.Function([a, b], bias) |
| |
| # Nothing changed after applying FoldConstant |
| a = run_opt_pass(before(), transform.FoldConstant()) |
| b = run_opt_pass(before(), transform.InferType()) |
| tvm.ir.assert_structural_equal(a, b) |
| |
| # Fold QNN constants |
| a = run_opt_pass(before(), transform.FoldConstant(fold_qnn=True)) |
| b = run_opt_pass(expected(), transform.InferType()) |
| tvm.ir.assert_structural_equal(a, b) |
| |
| |
| def test_fold_quantize(): |
| t = relay.TensorType([1, 2, 3], "int8") |
| |
| def before(): |
| data = tvm.nd.array(np.array([1.0, 2.0, 3.0], dtype="float32")) |
| const_fp = relay.const(data, dtype="float32") |
| const_i8 = relay.qnn.op.quantize( |
| const_fp, output_scale=relay.const(0.5), output_zero_point=relay.const(0) |
| ) |
| x = relay.var("x", t) |
| sub = relay.op.subtract(x, const_i8) |
| func = relay.Function([x], sub) |
| return func |
| |
| def expected(): |
| data = tvm.nd.array(np.array([2, 4, 6], dtype="int8")) |
| const_i8 = relay.const(data, dtype="int8") |
| x = relay.var("x", t) |
| sub = relay.op.subtract(x, const_i8) |
| func = relay.Function([x], sub) |
| return func |
| |
| # Nothing changed after applying FoldConstant |
| a = run_opt_pass(before(), transform.FoldConstant()) |
| b = run_opt_pass(before(), transform.InferType()) |
| tvm.ir.assert_structural_equal(a, b) |
| |
| # Fold QNN constants |
| a = run_opt_pass(before(), transform.FoldConstant(fold_qnn=True)) |
| b = run_opt_pass(expected(), transform.InferType()) |
| tvm.ir.assert_structural_equal(a, b) |
| |
| |
| def test_fold_qnn_conv2d_qnn_mul(): |
| def before(): |
| dtype = "uint8" |
| op0 = relay.qnn.op.conv2d( |
| relay.const(np.ones((1, 1, 2, 2), dtype=dtype), dtype=dtype), |
| relay.const(np.ones((1, 1, 2, 2), dtype=dtype), dtype=dtype), |
| input_zero_point=relay.const(0, "int32"), |
| kernel_zero_point=relay.const(0, "int32"), |
| input_scale=relay.const(1.0, "float32"), |
| kernel_scale=relay.const(1.0, "float32"), |
| kernel_size=(2, 2), |
| channels=1, |
| ) |
| op = relay.qnn.op.mul( |
| op0, |
| relay.const(np.array([10], dtype="int32"), dtype="int32"), |
| relay.const(1.0, dtype="float32"), |
| relay.const(0, dtype="int32"), |
| relay.const(1.0, dtype="float32"), |
| relay.const(0, dtype="int32"), |
| relay.const(1.0, dtype="float32"), |
| relay.const(0, dtype="int32"), |
| ) |
| func = relay.Function([], op) |
| return func |
| |
| def expected(): |
| data = relay.const(np.array([[[[40]]]], dtype="int32"), dtype="int32") |
| func = relay.Function([], data) |
| return func |
| |
| # Nothing changed after applying FoldConstant |
| a = run_opt_pass(before(), transform.FoldConstant()) |
| b = run_opt_pass(before(), transform.InferType()) |
| tvm.ir.assert_structural_equal(a, b) |
| |
| # Fold QNN constants |
| a = run_opt_pass(before(), transform.FoldConstant(fold_qnn=True)) |
| b = run_opt_pass(expected(), transform.InferType()) |
| tvm.ir.assert_structural_equal(a, b) |
| |
| |
| def test_fold_requantize(): |
| def before(): |
| data = tvm.nd.array(np.array([1, 2, 3], dtype="int8")) |
| const_i8 = relay.const(data, dtype="int8") |
| op = relay.qnn.op.requantize( |
| const_i8, |
| input_scale=relay.const(2.0, dtype="float32"), |
| input_zero_point=relay.const(1, dtype="int32"), |
| output_scale=relay.const(1.0, dtype="float32"), |
| output_zero_point=relay.const(1, dtype="int32"), |
| ) |
| x = relay.var("x", relay.TensorType([3], "int8")) |
| add = relay.op.add(op, x) |
| func = relay.Function([x], add) |
| return func |
| |
| def expected(): |
| data = tvm.nd.array(np.array([1, 3, 5], dtype="int8")) |
| const_i8 = relay.const(data, dtype="int8") |
| x = relay.var("x", relay.TensorType([3], "int8")) |
| add = relay.op.add(const_i8, x) |
| func = relay.Function([x], add) |
| return func |
| |
| # Nothing changed after applying FoldConstant |
| a = run_opt_pass(before(), transform.FoldConstant()) |
| b = run_opt_pass(before(), transform.InferType()) |
| tvm.ir.assert_structural_equal(a, b) |
| |
| # Fold QNN constants |
| a = run_opt_pass(before(), transform.FoldConstant(fold_qnn=True)) |
| b = run_opt_pass(expected(), transform.InferType()) |
| tvm.ir.assert_structural_equal(a, b) |
| |
| |
| def test_pass_link_params(): |
| """ |
| This test checks ensures that proper executor is passed to interpreter instance |
| The test will fail if FoldConstant does not override the executor due to "int8" |
| is not supported in ScheduleBuilder |
| """ |
| |
| def expr(): |
| z = relay.const(10, dtype="int8") |
| return relay.cast(z, dtype="int32") |
| |
| mod = tvm.IRModule.from_expr(expr()) |
| mod = tvm.relay.transform.InferType()(mod) |
| # Add executor with link-params |
| mod = mod.with_attr("executor", Executor("aot", {"link-params": True})) |
| mod = tvm.relay.transform.FoldConstant()(mod) |
| |
| |
| if __name__ == "__main__": |
| tvm.testing.main() |
