From a924621645cbec0491b6a60a6d3185d30cc655b7 Mon Sep 17 00:00:00 2001 From: pranav jonnalagadda-SJ1 Eng_ML Date: Tue, 25 Jan 2022 17:20:51 -0800 Subject: [PATCH] [Relay][DefuseOps pass] bug fix: To support function body types other than call node --- src/relay/transforms/defuse_ops.cc | 18 ++- tests/python/relay/test_pass_defuse_ops.py | 151 ++++++++++++++++++++- 2 files changed, 157 insertions(+), 12 deletions(-) diff --git a/src/relay/transforms/defuse_ops.cc b/src/relay/transforms/defuse_ops.cc index d7a9bfde57c3..0d97d5a7b75c 100644 --- a/src/relay/transforms/defuse_ops.cc +++ b/src/relay/transforms/defuse_ops.cc @@ -55,17 +55,15 @@ class DefuseOpsMutator : public ExprMutator { if (const auto* call = new_n.as()) { if (const auto* func = call->op.as()) { - if (func->body->IsInstance()) { - std::unordered_map name_to_args; - for (size_t i = 0; i < func->params.size(); ++i) { - const std::string& pname = func->params[i]->name_hint(); - ICHECK(name_to_args.cend() == name_to_args.find(pname)) - << "Found multiple parameters share the same variable name `" << pname - << "` which introduces uncertainty in DefuseOps pass"; - name_to_args[pname] = call->args[i]; - } - return FuncBodyMutator(std::move(name_to_args)).Mutate(func->body); + std::unordered_map name_to_args; + for (size_t i = 0; i < func->params.size(); ++i) { + const std::string& pname = func->params[i]->name_hint(); + ICHECK(name_to_args.cend() == name_to_args.find(pname)) + << "Found multiple parameters share the same variable name `" << pname + << "` which introduces uncertainty in DefuseOps pass"; + name_to_args[pname] = call->args[i]; } + return FuncBodyMutator(std::move(name_to_args)).Mutate(func->body); } } return new_n; diff --git a/tests/python/relay/test_pass_defuse_ops.py b/tests/python/relay/test_pass_defuse_ops.py index 2312b2d9ec47..f123bd582b87 100644 --- a/tests/python/relay/test_pass_defuse_ops.py +++ b/tests/python/relay/test_pass_defuse_ops.py @@ -14,6 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import numpy +import pytest import tvm from tvm import relay from tvm.relay import transform @@ -63,6 +65,151 @@ def before(dshape): assert tvm.ir.structural_equal(x, defused) +def test_defuse_complex(): + """Complex defuse testcase""" + + def fused_conv2d_batch_norm(w): + data = relay.var("data", shape=(1, 224, 224, 3)) + bn_gamma0 = relay.var("bn_gamma0", relay.TensorType((64,), "float32")) + bn_beta0 = relay.var("bn_beta0", relay.TensorType((64,), "float32")) + bn_mmean0 = relay.var("bn_mean0", relay.TensorType((64,), "float32")) + bn_mvar0 = relay.var("bn_var0", relay.TensorType((64,), "float32")) + c0 = relay.nn.conv2d( + data, + w, + strides=(2, 2), + padding=(3, 3, 3, 3), + channels=64, + kernel_size=(7, 7), + data_layout="NHWC", + kernel_layout="OHWI", + out_layout="NHWC", + ) + c1 = relay.nn.batch_norm(c0, bn_gamma0, bn_beta0, bn_mmean0, bn_mvar0, axis=3) + c2 = c1[0] + return relay.Function(relay.analysis.free_vars(c2), c2) + + def fused_conv2d_batch_norm_relu(z): + data2 = relay.var("data2", shape=(1, 56, 56, 64)) + bn_gamma0 = relay.var("bn_gamma0", relay.TensorType((64,), "float32")) + bn_beta0 = relay.var("bn_beta0", relay.TensorType((64,), "float32")) + bn_mmean0 = relay.var("bn_mean0", relay.TensorType((64,), "float32")) + bn_mvar0 = relay.var("bn_var0", relay.TensorType((64,), "float32")) + c0 = relay.nn.conv2d( + data2, + z, + padding=(1, 1, 1, 1), + channels=64, + kernel_size=(3, 3), + data_layout="NHWC", + kernel_layout="OHWI", + out_layout="NHWC", + ) + c1 = relay.nn.batch_norm(c0, bn_gamma0, bn_beta0, bn_mmean0, bn_mvar0, axis=3) + c2 = c1[0] + c3 = relay.nn.relu(data=c2) + return relay.Function(relay.analysis.free_vars(c3), c3) + + def fused_max_pool2d(): + data1 = relay.var("data1", shape=(1, 112, 112, 64)) + a1 = relay.nn.max_pool2d( + data1, + pool_size=(3, 3), + strides=(2, 2), + padding=(1, 1, 1, 1), + layout="NHWC", + out_layout="NHWC", + ) + return relay.Function(relay.analysis.free_vars(a1), a1) + + def fused_add_relu(): + data1 = relay.var("data1", shape=(1, 56, 56, 64)) + data2 = relay.var("data2", shape=(1, 56, 56, 64)) + a0 = relay.add(data1, data2) + a1 = relay.nn.relu(a0) + return relay.Function(relay.analysis.free_vars(a1), a1) + + def before_fused(conv_layer1_weight, conv_layer2_weight): + data = relay.var("data", shape=(1, 3, 224, 224)) + data1 = relay.layout_transform(data, src_layout="NCHW", dst_layout="NHWC") + bn_gamma0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) + bn_beta0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) + bn_mmean0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) + bn_mvar0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) + a0 = fused_conv2d_batch_norm(conv_layer1_weight) + a1 = fused_max_pool2d() + a2 = fused_conv2d_batch_norm_relu(conv_layer2_weight) + a3 = fused_add_relu() + y0 = relay.Call(a0, [data1, bn_gamma0, bn_beta0, bn_mmean0, bn_mvar0]) + y1 = relay.Call(a1, [y0]) + y2 = relay.Call(a2, [y1, bn_gamma0, bn_beta0, bn_mmean0, bn_mvar0]) + y3 = relay.Call(a3, [y1, y2]) + return relay.Function(relay.analysis.free_vars(y3), y3) + + def golden_defused(conv_layer1_weight, conv_layer2_weight): + data = relay.var("data", shape=(1, 3, 224, 224)) + data1 = relay.layout_transform(data, src_layout="NCHW", dst_layout="NHWC") + bn_gamma0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) + bn_beta0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) + bn_mmean0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) + bn_mvar0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) + c0 = relay.nn.conv2d( + data1, + conv_layer1_weight, + strides=(2, 2), + padding=(3, 3, 3, 3), + channels=64, + kernel_size=(7, 7), + data_layout="NHWC", + kernel_layout="OHWI", + out_layout="NHWC", + ) + c1 = relay.nn.batch_norm(c0, bn_gamma0, bn_beta0, bn_mmean0, bn_mvar0, axis=3) + c2 = c1[0] + c3 = relay.nn.max_pool2d( + c2, + pool_size=(3, 3), + strides=(2, 2), + padding=(1, 1, 1, 1), + layout="NHWC", + out_layout="NHWC", + ) + c4 = relay.nn.conv2d( + c3, + conv_layer2_weight, + padding=(1, 1, 1, 1), + channels=64, + kernel_size=(3, 3), + data_layout="NHWC", + kernel_layout="OHWI", + out_layout="NHWC", + ) + c5 = relay.nn.batch_norm(c4, bn_gamma0, bn_beta0, bn_mmean0, bn_mvar0, axis=3) + c6 = c5[0] + c7 = relay.nn.relu(c6) + c8 = relay.add(c3, c7) + c9 = relay.nn.relu(c8) + return relay.Function(relay.analysis.free_vars(c9), c9) + + # creating weight constants for the two convolution layers + # in the input fused model and the golden defused model. + conv_layer1_weight = relay.nn.Constant( + tvm.nd.array(numpy.ndarray(shape=(64, 7, 7, 3), dtype="float32")) + ) + conv_layer2_weight = relay.nn.Constant( + tvm.nd.array(numpy.ndarray(shape=(64, 3, 3, 64), dtype="float32")) + ) + x = before_fused(conv_layer1_weight, conv_layer2_weight) + x = run_opt_pass(x, transform.InferType()) + defused = run_opt_pass(x, transform.DefuseOps()) + + golden1 = golden_defused(conv_layer1_weight, conv_layer2_weight) + golden1 = run_opt_pass(golden1, transform.InferType()) + + assert tvm.ir.structural_equal(defused, golden1), ( + "Actual = \n" + str(defused) + "\nGolden = \n" + str(golden1) + ) + + if __name__ == "__main__": - test_defuse_simple() - test_inception_like() + pytest.main([__file__])