[REFACTOR][RELAY] Replace build_config with PassContext (apache#5698)

apps/android_camera/models/prepare_model.py

@@ -87,7 +87,7 @@ def main(model_str, output_path):
     except FileExistsError:
         pass
     print("building...")
-    with relay.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         graph, lib, params = relay.build(net, target, target_host=target_host, params=params)
     print("dumping lib...")
     lib.export_library(output_path_str + '/' + 'deploy_lib_cpu.so', ndk.create_shared)

apps/benchmark/arm_cpu_imagenet_bench.py

@@ -39,7 +39,7 @@ def evaluate_network(network, target, target_host, repeat):
     net, params, input_shape, output_shape = get_network(network, batch_size=1)
 
     print_progress("%-20s building..." % network)
-    with relay.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         graph, lib, params = relay.build(
             net, target=target, target_host=target_host, params=params)
 

apps/benchmark/gpu_imagenet_bench.py

@@ -33,7 +33,7 @@
 def benchmark(network, target):
     net, params, input_shape, output_shape = get_network(network, batch_size=1)
 
-    with relay.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         graph, lib, params = relay.build(net, target=target, params=params)
 
     # create runtime

apps/benchmark/mobile_gpu_imagenet_bench.py

@@ -38,7 +38,7 @@ def evaluate_network(network, target, target_host, dtype, repeat):
     net, params, input_shape, output_shape = get_network(network, batch_size=1, dtype=dtype)
 
     print_progress("%-20s building..." % network)
-    with relay.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         graph, lib, params = relay.build(
             net, target=target, target_host=target_host, params=params)
 

apps/bundle_deploy/build_model.py

@@ -33,7 +33,7 @@ def build_module(opts):
     func = mod["main"]
     func = relay.Function(func.params, relay.nn.softmax(func.body), None, func.type_params, func.attrs)
 
-    with relay.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         graph, lib, params = relay.build(
             func, 'llvm --system-lib', params=params)
 

apps/sgx/src/build_model.py

@@ -37,7 +37,7 @@ def main():
     net, params = relay.testing.resnet.get_workload(
         layers=18, batch_size=dshape[0], image_shape=dshape[1:])
 
-    with relay.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         graph, lib, params = relay.build(
             net, 'llvm --system-lib', params=params)
 

golang/sample/gen_mobilenet_lib.py

@@ -16,7 +16,7 @@
 # under the License.
 
 import os
-from tvm import relay
+from tvm import relay, transform
 from tvm.contrib.download import download_testdata
 
 
@@ -77,7 +77,7 @@ def extract(path):
 target = 'llvm'
 
 # Build with Relay
-with relay.build_config(opt_level=3):
+with transform.PassContext(opt_level=3):
     graph, lib, params = relay.build_module.build(
         mod, target, params=params)
 

python/tvm/relay/frontend/common.py

@@ -505,7 +505,7 @@ def infer_value(input_val, params, mod=None):
         assert all(var.name_hint in params.keys() for var in analysis.free_vars(
             input_val)), "All inputs to infer must be available in params."
         func = _function.Function(analysis.free_vars(input_val), input_val)
-        with tvm.relay.build_config(opt_level=0):
+        with tvm.transform.PassContext(opt_level=0):
             graph, lib, params = tvm.relay.build(func, target="llvm", params=params)
         ctx = tvm.cpu(0)
         m = graph_runtime.create(graph, lib, ctx)

python/tvm/relay/quantize/_calibrate.py

@@ -28,7 +28,6 @@
 from .. import op as _op
 from .. import expr as _expr
 from .. import analysis as _analysis
-from .. import transform as _transform
 from .. import build_module as _build_module
 from ...contrib import graph_runtime
 from .kl_divergence import _find_scale_by_kl
@@ -45,7 +44,7 @@ def _get_profile_runtime(mod):
         target = 'llvm'
         ctx = tvm.context(target)
 
-    with _transform.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         graph, lib, params = _build_module.build(func, target=target)
     runtime = graph_runtime.create(graph, lib, ctx)
     runtime.set_input(**params)

python/tvm/relay/transform/transform.py

@@ -21,6 +21,7 @@
 import types
 import inspect
 import functools
+import warnings
 
 import tvm.ir
 from tvm import te
@@ -34,7 +35,9 @@ def build_config(opt_level=2,
                  required_pass=None,
                  disabled_pass=None,
                  trace=None):
-    """Configure the build behavior by setting config variables.
+    """Configure the build behavior by setting config variables. This function
+    will be deprecated in TVM v0.7. Instead, we should directly use
+    tvm.transform.PassContext.
 
     Parameters
     ----------
@@ -72,8 +75,9 @@ def build_config(opt_level=2,
     pass_context: PassContext
         The pass context for optimizations.
     """
-    return tvm.ir.transform.PassContext(opt_level, required_pass,
-                                        disabled_pass, trace)
+    warnings.warn("relay.build_config will be deprecated. Please use \
+                  tvm.transform.PassContext directly", DeprecationWarning)
+    return tvm.transform.PassContext(opt_level, required_pass, disabled_pass, trace)
 
 
 @tvm._ffi.register_object("relay.FunctionPass")

rust/frontend/examples/resnet/src/build_resnet.py

@@ -75,8 +75,8 @@ def build(target_dir):
             num_layers=18, batch_size=batch_size, image_shape=image_shape)
 
     # compile the model
-    with relay.build_config(opt_level=opt_level):
-            graph, lib, params = relay.build_module.build(net, target, params=params)
+    with tvm.transform.PassContext(opt_level=opt_level):
+        graph, lib, params = relay.build_module.build(net, target, params=params)
 
     # save the model artifacts
     lib.save(deploy_lib)

src/relay/backend/build_module.cc

@@ -304,9 +304,8 @@ class RelayBuildModule : public runtime::ModuleNode {
     // Handle heterogeneous compilation.
     transform::PassContext pass_ctx = PassContext::Current();
     if (targets_.size() > 1) {
-      Optional<IntImm> opt_fallback_dev =
-          pass_ctx->GetConfig("relay.fallback_device_type",
-                              IntImm(runtime::DataType::Int(32), static_cast<int>(kDLCPU)));
+      Optional<Integer> opt_fallback_dev =
+          pass_ctx->GetConfig("relay.fallback_device_type", Integer(static_cast<int>(kDLCPU)));
       auto fallback_dev = opt_fallback_dev.value();
       CHECK_GT(fallback_dev->value, 0U);
       relay_module = RunDeviceAnnotationPass(relay_module, fallback_dev->value);

tests/cpp/relay_transform_sequential.cc

@@ -70,7 +70,7 @@ TEST(Relay, Sequential) {
   auto mod = IRModule::FromExpr(func);
   auto pass_ctx = relay::transform::PassContext::Create();
   pass_ctx->opt_level = 3;
-  pass_ctx->config.Set("relay.fallback_device_type", IntImm(DataType::Int(32), 1));
+  pass_ctx->config.Set("relay.fallback_device_type", Integer(1));
   {
     tvm::With<relay::transform::PassContext> ctx_scope(pass_ctx);
     tvm::With<tvm::Target> tctx(tvm::Target::Create("llvm"));

tests/python/frontend/caffe2/test_forward.py

@@ -43,7 +43,7 @@ def get_tvm_output(model,
     dtype_dict = {input_names: input_data.dtype}
     mod, params = relay.frontend.from_caffe2(
         model.init_net, model.predict_net, shape_dict, dtype_dict)
-    with relay.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         graph, lib, params = relay.build(mod, target, params=params)
 
     m = graph_runtime.create(graph, lib, ctx)

tests/python/frontend/coreml/test_forward.py

@@ -33,7 +33,7 @@
 
 def get_tvm_output(func, x, params, target, ctx,
                    out_shape=(1, 1000), input_name='image', dtype='float32'):
-    with relay.transform.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         graph, lib, params = relay.build(func, target, params=params)
     m = graph_runtime.create(graph, lib, ctx)
     # set inputs
@@ -76,7 +76,7 @@ def run_tvm_graph(coreml_model, target, ctx, input_data, input_name, output_shap
         dtype_dict = {input_name: input_data.dtype}
 
     mod, params = relay.frontend.from_coreml(coreml_model, shape_dict)
-    with relay.transform.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         graph, lib, params = relay.build(mod, target, params=params)
 
     from tvm.contrib import graph_runtime

tests/python/frontend/keras/test_forward.py

@@ -84,7 +84,7 @@ def get_keras_output(xs, dtype='float32'):
     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, layout=layout)
-        with relay.transform.build_config(opt_level=2):
+        with tvm.transform.PassContext(opt_level=2):
             graph, lib, params = relay.build(mod,
                                              target,
                                              params=params)

tests/python/frontend/mxnet/test_forward.py

@@ -66,7 +66,7 @@ def get_tvm_output(symbol, x, args, auxs, target, ctx, dtype='float32'):
                                                     shape_dict,
                                                     arg_params=args,
                                                     aux_params=auxs)
-        with relay.build_config(opt_level=3):
+        with tvm.transform.PassContext(opt_level=3):
             graph, lib, params = relay.build(mod, target, params=params)
         m = graph_runtime.create(graph, lib, ctx)
         # set inputs

tests/python/frontend/mxnet/test_qnn_ops_utils.py

@@ -15,8 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import tvm
 import numpy as np
+import tvm
 from tvm import relay
 from tvm.contrib import graph_runtime
 from tvm.relay.frontend.mxnet_qnn_op_utils import dequantize_mxnet_min_max, \
@@ -39,7 +39,7 @@ def dequantize_test_driver(in_dtype, quant_args, in_data, verify_output_data):
                                                       in_dtype=in_dtype)
         mod = relay.Function(relay.analysis.free_vars(dequantized_output), dequantized_output)
         mod = tvm.IRModule.from_expr(mod)
-        with relay.build_config(opt_level=3):
+        with tvm.transform.PassContext(opt_level=3):
             graph, lib, params = relay.build(mod, "llvm", params=None)
             rt_mod = graph_runtime.create(graph, lib, ctx=tvm.cpu(0))
             rt_mod.set_input(input_data=in_data)
@@ -93,7 +93,7 @@ def quantize_test_driver(out_dtype, quant_args, in_data, verify_output_data):
                                                         out_dtype=out_dtype)
         mod = relay.Function(relay.analysis.free_vars(quantized_output), quantized_output)
         mod = tvm.IRModule.from_expr(mod)
-        with relay.build_config(opt_level=3):
+        with tvm.transform.PassContext(opt_level=3):
             graph, lib, params = relay.build(mod, "llvm", params=None)
             rt_mod = graph_runtime.create(graph, lib, ctx=tvm.cpu(0))
             rt_mod.set_input(input_data=in_data)

tests/python/frontend/onnx/test_forward.py

@@ -65,7 +65,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, opset=opset)
 
-    with relay.build_config(opt_level=1):
+    with tvm.transform.PassContext(opt_level=1):
         graph, lib, params = relay.build(mod,
                                          target,
                                          params=params)

tests/python/frontend/pytorch/qnn_test.py

@@ -41,7 +41,7 @@ def get_tvm_runtime(script_module, input_name, ishape):
     input_shapes = [(input_name, ishape)]
     mod, params = relay.frontend.from_pytorch(script_module, input_shapes)
 
-    with relay.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         # test on only cpu for now, torch cannot run quant models on cuda
         # also not to make CI too slow
         json, lib, params = relay.build(mod, target="llvm", params=params)

tests/python/frontend/pytorch/test_forward.py

@@ -176,7 +176,7 @@ def verify_model(model_name, input_data=[],
     compiled_input = dict(zip(input_names,
                               [inp.cpu().numpy() for inp in baseline_input]))
 
-    with relay.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         for target, ctx in ctx_list:
             relay_graph, relay_lib, relay_params = relay.build(mod, target=target, params=params)
             relay_model = graph_runtime.create(relay_graph, relay_lib, ctx)
@@ -2294,7 +2294,7 @@ def test_forward_pretrained_bert_base_uncased():
     # ----------------------------
 
     target = 'llvm'
-    with relay.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         relay_graph, relay_lib, relay_params = relay.build(mod, target=target, params=params)
 
     ######################################################################

tests/python/frontend/tensorflow/test_bn_dynamic.py

@@ -50,7 +50,7 @@ def verify_fused_batch_norm(shape):
             continue
         mod, params = relay.frontend.from_tensorflow(constant_graph,
                                                      outputs=['output'])
-        with relay.build_config(opt_level=3):
+        with tvm.transform.PassContext(opt_level=3):
             graph, lib, params = relay.build(mod,
                                              target=device,
                                              params=params)

tests/python/frontend/tensorflow/test_forward.py

@@ -123,7 +123,7 @@ def run_tvm_graph(graph_def, input_data, input_node, num_output=1,
         result = ex.evaluate()(*inputs)
         return vmobj_to_list(result)
     else:
-        with relay.build_config(opt_level=opt_level):
+        with tvm.transform.PassContext(opt_level=opt_level):
             graph, lib, params = relay.build(mod, target, target_host, params)
 
         ctx = tvm.context(target, 0)
@@ -2307,7 +2307,7 @@ def _get_tvm_graph_module(graph_def):
                       'Model/RNN/RNN/multi_rnn_cell/cell_0/lstm_cell/LSTMBlockCell_c': 'float32',
                       'Model/RNN/RNN/multi_rnn_cell/cell_0/lstm_cell/LSTMBlockCell_h': 'float32'}
         target = 'llvm'
-        with relay.build_config(opt_level=0):
+        with tvm.transform.PassContext(opt_level=0):
             graph, lib, params = relay.build(mod,
                                              target,
                                              params=params)

tests/python/frontend/tflite/test_forward.py

@@ -109,7 +109,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):
+    with tvm.transform.PassContext(opt_level=3):
         graph, lib, params = relay.build(mod, target, params=params)
 
     ctx = tvm.context(target, 0)

tests/python/nightly/quantization/test_quantization_accuracy.py

@@ -66,7 +66,7 @@ def get_model(model_name, batch_size, qconfig, target=None, original=False, simu
     mod, params = relay.frontend.from_mxnet(gluon_model, {"data": data_shape})
     net = mod['main']
 
-    with relay.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         qfunc = relay.quantize.prerequisite_optimize(net, params=params)
     logging.debug('original')
     logging.debug(qfunc.astext(show_meta_data=False))
@@ -83,7 +83,7 @@ def get_model(model_name, batch_size, qconfig, target=None, original=False, simu
 
 
 def eval_acc(model, dataset, batch_fn, target=tvm.target.cuda(), ctx=tvm.gpu(), log_interval=100):
-    with relay.build_config(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         graph, lib, params = relay.build(model, target)
     # create runtime module
     m = tvm.contrib.graph_runtime.create(graph, lib, ctx)

tests/python/relay/benchmarking/benchmark_vm.py

@@ -36,7 +36,7 @@ def benchmark_execution(mod,
                         model="unknown"):
     def get_graph_runtime_output(mod, data, params, target, ctx,
                                  dtype='float32', number=2, repeat=20):
-        with relay.build_config(opt_level=3):
+        with tvm.transform.PassContext(opt_level=3):
             graph, lib, params = relay.build(mod, target, params=params)
 
         m = graph_runtime.create(graph, lib, ctx)
@@ -59,7 +59,7 @@ def get_graph_runtime_output(mod, data, params, target, ctx,
 
     def get_vm_output(mod, data, params, target, ctx, dtype='float32',
                       number=2, repeat=20):
-        with relay.build_config(opt_level=3):
+        with tvm.transform.PassContext(opt_level=3):
             exe = vm.compile(mod, target, params=params)
             rly_vm = vm_rt.VirtualMachine(exe)
             rly_vm.init(ctx)

tests/python/relay/test_backend_compile_engine.py

@@ -184,7 +184,7 @@ def test_compile_placeholder_bypass():
     z = relay.var("z", shape=(2, 3))
     result = relay.Tuple([x, relay.op.concatenate([y, z], axis=0)])
     func = relay.Function(relay.analysis.free_vars(result), result)
-    with relay.build_config(opt_level=0):
+    with tvm.transform.PassContext(opt_level=0):
        graph, lib, params = relay.build(tvm.IRModule.from_expr(func), 'llvm')
 
 

tests/python/relay/test_backend_graph_runtime.py

@@ -166,7 +166,7 @@ def unit_numpy(X, W):
     z = unit(rnn_dim)
 
     for target, ctx in ctx_list():
-        with relay.build_config(opt_level=2):
+        with tvm.transform.PassContext(opt_level=2):
             graph, lib, params = relay.build(tvm.IRModule.from_expr(z), target)
             m = graph_runtime.create(graph, lib, ctx)
             m.set_input("X", tvm.nd.array(x.astype(dtype)))

tests/python/relay/test_cpp_build_module.py

@@ -115,7 +115,7 @@ def check_conversion(tgt, ctx):
             X = tvm.nd.array(n * np.random.randn(n).astype(src) - n / 2)
 
             # build
-            with relay.build_config(opt_level=1):
+            with tvm.transform.PassContext(opt_level=1):
                 g_json, mmod, params = relay.build(tvm.IRModule.from_expr(func), tgt)
 
             # test

tests/python/relay/test_external_codegen.py

@@ -49,7 +49,8 @@ def update_lib(lib):
         return lib
 
     def check_vm_result():
-        with relay.build_config(opt_level=3, disabled_pass=["AlterOpLayout"]):
+        with tvm.transform.PassContext(opt_level=3,
+                                       disabled_pass=["AlterOpLayout"]):
             exe = relay.vm.compile(mod, target=target)
         code, lib = exe.save()
         lib = update_lib(lib)
@@ -60,7 +61,8 @@ def check_vm_result():
         tvm.testing.assert_allclose(out.asnumpy(), result, rtol=tol, atol=tol)
 
     def check_graph_runtime_result():
-        with relay.build_config(opt_level=3, disabled_pass=["AlterOpLayout"]):
+        with tvm.transform.PassContext(opt_level=3,
+                                       disabled_pass=["AlterOpLayout"]):
             json, lib, _ = relay.build(mod, target=target)
         lib = update_lib(lib)
         rt_mod = tvm.contrib.graph_runtime.create(json, lib, ctx)

tests/python/relay/test_memory_passes.py

@@ -37,7 +37,7 @@ def check_memory_plan(func, check_fn):
     no_plan_result = ex.evaluate(mod['main'])(*args)
 
     # Compute with memory planning.
-    with relay.build_config(opt_level=1, disabled_pass=["MemoryPlan"]):
+    with tvm.transform.PassContext(opt_level=1, disabled_pass=["MemoryPlan"]):
         plan_result = ex.evaluate(mod['main'])(*args)
 
     # Compute Python result.