From 4629d76d75c11afd2a964aa93b89fccb64feabae Mon Sep 17 00:00:00 2001
From: Alexander Pivovarov <pivovaa@amazon.com>
Date: Sun, 23 Jun 2019 15:51:22 -0700
Subject: [PATCH] Add mod supoort in relay.build

---
 python/tvm/relay/build_module.py              | 22 ++++++++++++++-----
 tests/python/frontend/caffe2/test_forward.py  |  2 +-
 tests/python/frontend/coreml/test_forward.py  |  2 +-
 tests/python/frontend/darknet/test_forward.py |  2 +-
 tests/python/frontend/keras/test_forward.py   |  2 +-
 tests/python/frontend/mxnet/test_forward.py   |  2 +-
 tests/python/frontend/onnx/test_forward.py    |  2 +-
 .../frontend/tensorflow/test_forward.py       |  4 ++--
 tests/python/frontend/tflite/test_forward.py  |  4 +---
 tutorials/frontend/deploy_model_on_android.py |  2 +-
 tutorials/frontend/deploy_ssd_gluoncv.py      |  2 +-
 tutorials/frontend/from_caffe2.py             |  2 +-
 tutorials/frontend/from_coreml.py             |  2 +-
 tutorials/frontend/from_darknet.py            |  2 +-
 tutorials/frontend/from_tensorflow.py         |  2 +-
 tutorials/frontend/from_tflite.py             |  2 +-
 16 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/python/tvm/relay/build_module.py b/python/tvm/relay/build_module.py
index 1aa4d5ae57c46..4cfdc26f14741 100644
--- a/python/tvm/relay/build_module.py
+++ b/python/tvm/relay/build_module.py
@@ -19,6 +19,7 @@
 from a Relay expression.
 """
 import numpy as np
+import warnings
 
 from tvm import expr as tvm_expr
 from .. import nd as _nd, target as _target, autotvm
@@ -27,6 +28,7 @@
 from . import ir_pass
 from . import ty as _ty
 from . import expr as _expr
+from .module import Module as _Module
 from .backend import interpreter as _interpreter
 from .backend.vm import VMExecutor
 
@@ -137,14 +139,14 @@ def get_params(self):
         return ret
 
 
-def build(func, target=None, target_host=None, params=None):
+def build(mod, target=None, target_host=None, params=None):
     """Helper function that builds a Relay function to run on TVM graph
     runtime.
 
     Parameters
     ----------
-    func: relay.Function
-        The function to build.
+    mod : relay.Module
+        The module to build. Using relay.Function is deprecated.
 
     target : str, :any:`tvm.target.Target`, or dict of str(i.e. device/context
     name) to str/tvm.target.Target, optional
@@ -175,6 +177,17 @@ def build(func, target=None, target_host=None, params=None):
     params : dict
         The parameters of the final graph.
     """
+    if isinstance(mod, _Module):
+        func = mod[mod.entry_func]
+    elif isinstance(mod, _expr.Function):
+        func = mod
+        warnings.warn(
+            "Please use input parameter mod (tvm.relay.module.Module) "
+            "instead of deprecated parameter func (tvm.relay.expr.Function)",
+            DeprecationWarning)
+    else:
+        raise ValueError("Type of input parameter mod must be tvm.relay.module.Module")
+
     target = _update_target(target)
 
     if isinstance(target_host, (str, _target.Target)):
@@ -192,8 +205,7 @@ def build(func, target=None, target_host=None, params=None):
 
     with tophub_context:
         bld_mod = BuildModule()
-        graph_json, mod, params = bld_mod.build(func, target, target_host,
-                                                params)
+        graph_json, mod, params = bld_mod.build(func, target, target_host, params)
     return graph_json, mod, params
 
 
diff --git a/tests/python/frontend/caffe2/test_forward.py b/tests/python/frontend/caffe2/test_forward.py
index 465ac70afc5fd..92258bbc284e5 100644
--- a/tests/python/frontend/caffe2/test_forward.py
+++ b/tests/python/frontend/caffe2/test_forward.py
@@ -43,7 +43,7 @@ def get_tvm_output(model,
     mod, params = relay.frontend.from_caffe2(
         model.init_net, model.predict_net, shape_dict, dtype_dict)
     with relay.build_config(opt_level=3):
-        graph, lib, params = relay.build(mod[mod.entry_func], target, params=params)
+        graph, lib, params = relay.build(mod, target, params=params)
 
     m = graph_runtime.create(graph, lib, ctx)
 
diff --git a/tests/python/frontend/coreml/test_forward.py b/tests/python/frontend/coreml/test_forward.py
index 7f0cdd1a51207..13f987c32be78 100644
--- a/tests/python/frontend/coreml/test_forward.py
+++ b/tests/python/frontend/coreml/test_forward.py
@@ -73,7 +73,7 @@ def run_tvm_graph(coreml_model, target, ctx, input_data, input_name, output_shap
 
     mod, params = relay.frontend.from_coreml(coreml_model, shape_dict)
     with relay.transform.build_config(opt_level=3):
-        graph, lib, params = relay.build(mod[mod.entry_func], target, params=params)
+        graph, lib, params = relay.build(mod, target, params=params)
 
     from tvm.contrib import graph_runtime
     m = graph_runtime.create(graph, lib, ctx)
diff --git a/tests/python/frontend/darknet/test_forward.py b/tests/python/frontend/darknet/test_forward.py
index 06916172f2bf5..ebfbbd320391c 100644
--- a/tests/python/frontend/darknet/test_forward.py
+++ b/tests/python/frontend/darknet/test_forward.py
@@ -55,7 +55,7 @@ def _get_tvm_output(net, data, build_dtype='float32', states=None):
     mod, params = relay.frontend.from_darknet(net, data.shape, dtype)
     target = 'llvm'
     shape_dict = {'data': data.shape}
-    graph, library, params = relay.build(mod[mod.entry_func],
+    graph, library, params = relay.build(mod,
                                          target,
                                          params=params)
 
diff --git a/tests/python/frontend/keras/test_forward.py b/tests/python/frontend/keras/test_forward.py
index 2dfb43b70455c..03117b21e3e23 100644
--- a/tests/python/frontend/keras/test_forward.py
+++ b/tests/python/frontend/keras/test_forward.py
@@ -44,7 +44,7 @@ def get_tvm_output(xs, target, ctx, dtype='float32'):
         shape_dict = {name: x.shape for (name, x) in zip(keras_model.input_names, xs)}
         mod, params = relay.frontend.from_keras(keras_model, shape_dict)
         with relay.transform.build_config(opt_level=2):
-            graph, lib, params = relay.build(mod[mod.entry_func],
+            graph, lib, params = relay.build(mod,
                                              target,
                                              params=params)
         m = graph_runtime.create(graph, lib, ctx)
diff --git a/tests/python/frontend/mxnet/test_forward.py b/tests/python/frontend/mxnet/test_forward.py
index c82dc208fdc42..ffef538a5b5c3 100644
--- a/tests/python/frontend/mxnet/test_forward.py
+++ b/tests/python/frontend/mxnet/test_forward.py
@@ -66,7 +66,7 @@ def get_tvm_output(symbol, x, args, auxs, target, ctx, dtype='float32'):
                                                     arg_params=args,
                                                     aux_params=auxs)
         with relay.build_config(opt_level=3):
-            graph, lib, params = relay.build(mod[mod.entry_func], target, params=params)
+            graph, lib, params = relay.build(mod, target, params=params)
         m = graph_runtime.create(graph, lib, ctx)
         # set inputs
         m.set_input("data", tvm.nd.array(x.astype(dtype)))
diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
index a52e3f0cc16e8..d409960118966 100644
--- a/tests/python/frontend/onnx/test_forward.py
+++ b/tests/python/frontend/onnx/test_forward.py
@@ -47,7 +47,7 @@ def get_tvm_output(graph_def, input_data, target, ctx, output_shape=None, output
 
     mod, params = relay.frontend.from_onnx(graph_def, shape_dict)
     with relay.build_config(opt_level=1):
-        graph, lib, params = relay.build(mod[mod.entry_func],
+        graph, lib, params = relay.build(mod,
                                          target,
                                          params=params)
 
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index 4035d0265ebd2..ce66aab2a7e42 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -64,7 +64,7 @@ def run_tvm_graph(graph_def, input_data, input_node, num_output=1,
                                                  shape=shape_dict,
                                                  outputs=out_names)
     with relay.build_config(opt_level=opt_level):
-        graph, lib, params = relay.build(mod[mod.entry_func], target, target_host, params)
+        graph, lib, params = relay.build(mod, target, target_host, params)
 
     ctx = tvm.context(target, 0)
     from tvm.contrib import graph_runtime
@@ -1436,7 +1436,7 @@ def _get_tvm_graph_module(graph_def):
                       'Model/RNN/RNN/multi_rnn_cell/cell_0/lstm_cell/LSTMBlockCell_h':'float32'}
         target = 'llvm'
         with relay.build_config(opt_level=0):
-            graph, lib, params = relay.build(mod[mod.entry_func],
+            graph, lib, params = relay.build(mod,
                                              target,
                                              params=params)
         from tvm.contrib import graph_runtime
diff --git a/tests/python/frontend/tflite/test_forward.py b/tests/python/frontend/tflite/test_forward.py
index 41147fe9e9bf6..b633c58efd2b5 100644
--- a/tests/python/frontend/tflite/test_forward.py
+++ b/tests/python/frontend/tflite/test_forward.py
@@ -68,9 +68,7 @@ def run_tvm_graph(tflite_model_buf, input_data, input_node, num_output=1, target
                                              shape_dict=shape_dict,
                                              dtype_dict=dtype_dict)
     with relay.build_config(opt_level=3):
-        graph, lib, params = relay.build(mod[mod.entry_func],
-                                         target,
-                                         params=params)
+        graph, lib, params = relay.build(mod, target, params=params)
 
     ctx = tvm.context(target, 0)
     from tvm.contrib import graph_runtime
diff --git a/tutorials/frontend/deploy_model_on_android.py b/tutorials/frontend/deploy_model_on_android.py
index 2e416b7f379ac..bc245f00d9287 100644
--- a/tutorials/frontend/deploy_model_on_android.py
+++ b/tutorials/frontend/deploy_model_on_android.py
@@ -263,7 +263,7 @@ def transform_image(image):
 mod, params = relay.frontend.from_keras(keras_mobilenet_v2, shape_dict)
 
 with relay.build_config(opt_level=3):
-    graph, lib, params = relay.build(mod[mod.entry_func], target=target,
+    graph, lib, params = relay.build(mod, target=target,
                                      target_host=target_host, params=params)
 
 # After `relay.build`, you will get three return values: graph,
diff --git a/tutorials/frontend/deploy_ssd_gluoncv.py b/tutorials/frontend/deploy_ssd_gluoncv.py
index 92b488f8fa07e..ad7c27c1b2729 100644
--- a/tutorials/frontend/deploy_ssd_gluoncv.py
+++ b/tutorials/frontend/deploy_ssd_gluoncv.py
@@ -78,7 +78,7 @@
 def build(target):
     mod, params = relay.frontend.from_mxnet(block, {"data": dshape})
     with relay.build_config(opt_level=3):
-        graph, lib, params = relay.build(mod[mod.entry_func], target, params=params)
+        graph, lib, params = relay.build(mod, target, params=params)
     return graph, lib, params
 
 ######################################################################
diff --git a/tutorials/frontend/from_caffe2.py b/tutorials/frontend/from_caffe2.py
index 082dfafb33e70..aadee02f4b027 100644
--- a/tutorials/frontend/from_caffe2.py
+++ b/tutorials/frontend/from_caffe2.py
@@ -89,7 +89,7 @@ def transform_image(image):
 # target x86 CPU
 target = 'llvm'
 with relay.build_config(opt_level=3):
-    graph, lib, params = relay.build(mod[mod.entry_func], target, params=params)
+    graph, lib, params = relay.build(mod, target, params=params)
 
 ######################################################################
 # Execute on TVM
diff --git a/tutorials/frontend/from_coreml.py b/tutorials/frontend/from_coreml.py
index 7eeefb3f2e5b6..1bb7cbd885aba 100644
--- a/tutorials/frontend/from_coreml.py
+++ b/tutorials/frontend/from_coreml.py
@@ -71,7 +71,7 @@
 mod, params = relay.frontend.from_coreml(mlmodel, shape_dict)
 
 with relay.build_config(opt_level=3):
-    graph, lib, params = relay.build(mod[mod.entry_func],
+    graph, lib, params = relay.build(mod,
                                      target,
                                      params=params)
 
diff --git a/tutorials/frontend/from_darknet.py b/tutorials/frontend/from_darknet.py
index d9014e092ff03..f3eb05f9d20b6 100644
--- a/tutorials/frontend/from_darknet.py
+++ b/tutorials/frontend/from_darknet.py
@@ -95,7 +95,7 @@
 shape = {'data': data.shape}
 print("Compiling the model...")
 with relay.build_config(opt_level=3):
-    graph, lib, params = relay.build(mod[mod.entry_func],
+    graph, lib, params = relay.build(mod,
                                      target=target,
                                      target_host=target_host,
                                      params=params)
diff --git a/tutorials/frontend/from_tensorflow.py b/tutorials/frontend/from_tensorflow.py
index 6603c2da15bc2..34865f021230a 100644
--- a/tutorials/frontend/from_tensorflow.py
+++ b/tutorials/frontend/from_tensorflow.py
@@ -140,7 +140,7 @@
 #   lib: target library which can be deployed on target with TVM runtime.
 
 with relay.build_config(opt_level=3):
-    graph, lib, params = relay.build(mod[mod.entry_func],
+    graph, lib, params = relay.build(mod,
                                      target=target,
                                      target_host=target_host,
                                      params=params)
diff --git a/tutorials/frontend/from_tflite.py b/tutorials/frontend/from_tflite.py
index 5a8525133f7c8..7baa45e3a4345 100644
--- a/tutorials/frontend/from_tflite.py
+++ b/tutorials/frontend/from_tflite.py
@@ -145,7 +145,7 @@ def extract(path):
 # target x86 CPU
 target = "llvm"
 with relay.build_config(opt_level=3):
-    graph, lib, params = relay.build(mod[mod.entry_func], target, params=params)
+    graph, lib, params = relay.build(mod, target, params=params)
 
 ######################################################################
 # Execute on TVM