zhiics commented on a change in pull request #4996: [relay][external codegen]
outline and inline lifted functions for external codegen
URL: https://github.com/apache/incubator-tvm/pull/4996#discussion_r389194638
##########
File path: tests/python/relay/test_pass_partition_graph.py
##########
@@ -427,10 +465,110 @@ def test_extern_dnnl_mobilenet():
(1, 1000), ref_res.asnumpy(), tol=1e-5, params=params)
+def test_function_lifting():
+ def partition():
+ data = relay.var("data", relay.TensorType((1, 3, 224, 224), "float32"))
+ weight = relay.var("weight", relay.TensorType((16, 3, 3, 3),
"float32"))
+ bn_gamma = relay.var("bn_gamma", relay.TensorType((16, ), "float32"))
+ bn_beta = relay.var("bn_beta", relay.TensorType((16, ), "float32"))
+ bn_mmean = relay.var("bn_mean", relay.TensorType((16, ), "float32"))
+ bn_mvar = relay.var("bn_var", relay.TensorType((16, ), "float32"))
+
+ 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)
+
+ func = relay.Function([data, weight, bn_gamma, bn_beta, bn_mmean,
+ bn_mvar], bn_output.astuple())
+ mod = tvm.IRModule()
+ mod["main"] = func
+ op_list = ["nn.batch_norm", "nn.conv2d"]
+ mod = WhiteListAnnotator(op_list, "test_compiler")(mod)
+
+ opt_pass = transform.Sequential([
+ transform.InferType(),
+ transform.PartitionGraph(),
+ transform.SimplifyInference(),
+ transform.FoldConstant(),
+ transform.AlterOpLayout()
Review comment:
Let's keep the previous one as it tests the lifting. We can create a
separate one to test batch_norm with inline. There are a few other tests for
outlining and inlining tests in this file as well. I think this should be
enough.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
