| # 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 pytest |
| import os |
| import time |
| import numpy as np |
| import tvm |
| import tvm.testing |
| from tvm import relay |
| from tvm.relay import transform, build_module |
| from tvm.relay.testing import run_opt_pass |
| from tvm.contrib import graph_executor, pipeline_executor, pipeline_executor_build |
| from tvm._ffi import get_global_func |
| from tvm.contrib import cc as _cc |
| |
| |
| def graph_split(expr, split_conf, params=None): |
| """Splitting the graph into a list of subgraphs""" |
| |
| def get_dep_var(sub_var_dep): |
| return [var for var in sub_var_dep[len(sub_var_dep) - 1]["ref_nodes"]] |
| |
| def parse_dependency(value, snode_dep, new_input_idx): |
| new_args = [] |
| need_update = False |
| for var in value.args: |
| is_free_var = False |
| for dep in snode_dep[:-1]: |
| if var in dep["nodes"]: |
| # Mark the previous subgraph node as a dependency. |
| dep["nodes"][var] += 1 |
| dep["ref_nodes"][var] = dep["nodes"][var] |
| # The var of this call is a free_var |
| is_free_var = True |
| # if the var of this call is a free_var, recreate it and give it a fixed input name. |
| if is_free_var: |
| need_update = True |
| new_args.append(relay.var(f"data_n_{new_input_idx}", var.checked_type)) |
| new_input_idx += 1 |
| else: |
| new_args.append(var) |
| # if the 'tvm.relay.expr.Call' has a free_var, recreate it with new name as 'data_n_*'. |
| if need_update: |
| value = tvm.relay.expr.Call( |
| value.op, new_args, value.attrs, value.type_args, value.span |
| ) |
| return value, snode_dep, new_input_idx |
| |
| def merge_constant_expr(constant_expr, expr): |
| # merge constant express with a express |
| if not isinstance(constant_expr.body, tvm.relay.expr.Let): |
| return tvm.relay.expr.Let(constant_expr.var, constant_expr.value, expr) |
| |
| return tvm.relay.expr.Let( |
| constant_expr.var, constant_expr.value, merge_constant_expr(constant_expr.body, expr) |
| ) |
| |
| def _recursion(anf, pipeline_mods, split_conf, constant_expr): |
| # Enumurate all operators of compute graph, then split the compute graph into a group of |
| # subgraph. |
| nonlocal operator_index_map |
| nonlocal new_input_idx |
| nonlocal snode_dep |
| cur_node_dep = snode_dep[len(snode_dep) - 1] |
| if isinstance(anf, tvm.relay.Function): |
| return tvm.relay.Function( |
| anf.params, |
| _recursion(anf.body, pipeline_mods, split_conf, constant_expr), |
| anf.ret_type, |
| anf.type_params, |
| anf.attrs, |
| ) |
| if isinstance(anf, tvm.relay.expr.Let): |
| value = anf.value |
| # record the constant expr to make sure all sugraphs can find correct constant. |
| if isinstance(value, tvm.relay.expr.Constant): |
| if not constant_expr: |
| constant_expr = tvm.relay.expr.Let(anf.var, value, anf.var) |
| else: |
| constant_expr = tvm.relay.expr.Let(anf.var, value, constant_expr) |
| if isinstance(value, tvm.relay.expr.Call): |
| new_args = [] |
| # build current var list |
| cur_node_dep["nodes"][anf.var] = 0 |
| # Get the dependency information of the nodes. |
| value, snode_dep, new_input_idx = parse_dependency(value, snode_dep, new_input_idx) |
| if isinstance(value.op, tvm.ir.Op): |
| if value.op.name in operator_index_map: |
| operator_index_map[value.op.name] += 1 |
| else: |
| operator_index_map[value.op.name] = 0 |
| split_operator_name = split_conf[0]["op_name"] if split_conf else "" |
| split_operator_index = split_conf[0]["op_index"] if split_conf else "" |
| # if a operator name and repeating count in the network match with the values |
| # of the 'split configuration', then this place is where we should do the |
| # graph splitting. |
| if ( |
| split_conf |
| and split_operator_name in operator_index_map |
| and operator_index_map[split_operator_name] >= split_operator_index |
| ): |
| # Do graph splitting. |
| split_conf.pop(0) |
| snode_dep.append({"nodes": {}, "ref_nodes": {}}) |
| ann = _recursion( |
| anf.body, |
| pipeline_mods, |
| split_conf, |
| constant_expr, |
| ) |
| snode_dep.pop() |
| dep_vars = get_dep_var(snode_dep) |
| # When the nodes of the current subgraph are the depedency node of another |
| # subgraph, we need to set them as the output of current subgraph. |
| body = relay.Tuple(dep_vars) if len(dep_vars) > 1 else anf.var |
| # when the operator of current subgraph uses previous subgraph constant |
| # as the argument of a "relay.expr.call", such constant may become a free |
| # varaible if the constant does not exist in the current subgraph. |
| # merge the previous constant with current subgraph to avoid such issue. |
| if constant_expr: |
| ann = merge_constant_expr(constant_expr, ann) |
| ann = run_opt_pass(ann, transform.ToGraphNormalForm()) |
| mod = tvm.IRModule.from_expr(ann) |
| pipeline_mods.insert(0, mod) |
| # Return the last node of the current subgraph. |
| return tvm.relay.expr.Let(anf.var, value, body) |
| return tvm.relay.expr.Let( |
| anf.var, |
| value, |
| _recursion(anf.body, pipeline_mods, split_conf, constant_expr), |
| ) |
| else: |
| return anf |
| |
| snode_dep = [{"nodes": {}, "ref_nodes": {}}] |
| pipeline_mods = [] |
| operator_index_map = {} |
| # Used to tracking new input which caused by graph splitting. |
| new_input_idx = 0 |
| constant_expr = None |
| subgraph_split_conf = split_conf.copy() |
| # Binding the parameters. |
| if params: |
| expr = build_module.bind_params_by_name(expr, params) |
| anf = run_opt_pass(expr, transform.ToANormalForm()) |
| anf = run_opt_pass(anf, transform.InferType()) |
| ann = _recursion( |
| anf, |
| pipeline_mods, |
| subgraph_split_conf, |
| constant_expr, |
| ) |
| ann = run_opt_pass(ann.body, transform.ToGraphNormalForm()) |
| mod = tvm.IRModule.from_expr(ann) |
| pipeline_mods.insert(0, mod) |
| return pipeline_mods |
| |
| |
| def get_network(): |
| # Get a list of modules representing subgraphs. |
| mods = [] |
| dshape = (3, 3) |
| data = relay.var("data_0", relay.TensorType(dshape, "float32")) |
| data21 = relay.var("data_1", relay.TensorType(dshape, "float32")) |
| data_net1_output_1 = relay.var("data_0", relay.TensorType(dshape, "float32")) |
| data_net1_output_2 = relay.var("data_1", relay.TensorType(dshape, "float32")) |
| data_net2_output_1 = relay.var("data_0", relay.TensorType(dshape, "float32")) |
| mvalue1 = np.full((1), 1).astype("float32") |
| mvalue2 = np.full((1), 2).astype("float32") |
| mvalue3 = np.full((1), 3).astype("float32") |
| mv1 = relay.Constant(tvm.nd.array(mvalue1)) |
| mv2 = relay.Constant(tvm.nd.array(mvalue2)) |
| mv3 = relay.Constant(tvm.nd.array(mvalue3)) |
| # There are three outputs in the first model. |
| net1_output1 = relay.add(data, mv1) |
| net1_output2 = relay.subtract(data, mv2) |
| net1_output3 = relay.concatenate((net1_output1, net1_output2), axis=0) |
| (net1_output3, _) = relay.split(net1_output3, indices_or_sections=2, axis=0) |
| net1_output3 = relay.add(net1_output3, mv2) |
| # The second model uses the output named net1_output3 of the first model as the first input, |
| # the second input of the second model is data21. |
| net2 = relay.add(net1_output3, mv2) |
| net2 = relay.add(net2, data21) |
| net2_output = relay.add(net2, mv3) |
| # The third model uses the output named net2_output of the second model as the first input |
| # and uses the output named net1_output2 of the first model as the second input. |
| net3 = relay.multiply(net2_output, mv3) |
| net3 = relay.add(net3, net1_output2) |
| return tvm.IRModule.from_expr(relay.Function([data, data21], relay.Tuple([net3]))), dshape |
| |
| |
| def get_split_mod(): |
| mod, dshape = get_network() |
| split_conf = [{"op_name": "add", "op_index": 1}, {"op_name": "add", "op_index": 4}] |
| mods = graph_split(mod["main"], split_conf) |
| return mods, dshape |
| |
| |
| def get_mannual_mod(): |
| # Get a list of modules representing subgraphs. |
| mods = [] |
| dshape = (3, 3) |
| data = relay.var("data_0", relay.TensorType(dshape, "float32")) |
| data21 = relay.var("data_1", relay.TensorType(dshape, "float32")) |
| data_net1_output_1 = relay.var("data_0", relay.TensorType(dshape, "float32")) |
| data_net1_output_2 = relay.var("data_1", relay.TensorType(dshape, "float32")) |
| data_net2_output_1 = relay.var("data_0", relay.TensorType(dshape, "float32")) |
| mvalue1 = np.full((1), 1).astype("float32") |
| mvalue2 = np.full((1), 2).astype("float32") |
| mvalue3 = np.full((1), 3).astype("float32") |
| mv1 = relay.Constant(tvm.nd.array(mvalue1)) |
| mv2 = relay.Constant(tvm.nd.array(mvalue2)) |
| mv3 = relay.Constant(tvm.nd.array(mvalue3)) |
| |
| # There are three outputs in the first model. |
| |
| net1_output1 = relay.add(data, mv1) |
| net1_output2 = relay.subtract(data, mv2) |
| net1_output3 = relay.multiply(data, mv3) |
| |
| # The second model use output named net1_output1 of the first model as the first input, |
| # the second input of the second model is data21. |
| net2 = relay.add(data_net1_output_1, mv2) |
| net2 = relay.add(net2, data21) |
| net2_output = relay.add(net2, mv3) |
| |
| # The third model use the output named net2_output of the second model as the first input |
| # and use the output named net1_output2 of the first model as the second input. |
| net3 = relay.multiply(data_net2_output_1, mv3) |
| net3 = relay.add(net3, data_net1_output_2) |
| |
| mods.append( |
| tvm.IRModule.from_expr( |
| relay.Function([data], relay.Tuple([net1_output1, net1_output2, net1_output3])) |
| ) |
| ) |
| mods.append(tvm.IRModule.from_expr(relay.Function([data_net1_output_1, data21], net2_output))) |
| mods.append( |
| tvm.IRModule.from_expr(relay.Function([data_net1_output_2, data_net2_output_1], net3)) |
| ) |
| |
| return mods, dshape |
| |
| |
| def get_manual_conf(mods, target): |
| # This function is used to generate manual pipeline configuration. |
| mod_config = {} |
| # The third output is the final output, the second output is for mod3, the first output |
| # is for mod2 input. |
| pipe_config1 = { |
| "mod_idx": 0, |
| "cpu_affinity": "0", |
| "output": [ |
| {"output_idx": 0, "dependencies": [{"mod_idx": 1, "input_name": "data_n_0"}]}, |
| {"output_idx": 1, "dependencies": [{"mod_idx": 2, "input_name": "data_n_2"}]}, |
| ], |
| } |
| mod_config[mods[0]] = { |
| "pipeline": pipe_config1, |
| "target_host": None, |
| "mod_name": "default", |
| "build": None, |
| "params": None, |
| "target": target[0], |
| "fcompile": _cc.create_shared, |
| "dev": target[1], |
| } |
| |
| pipe_config2 = { |
| "mod_idx": 1, |
| "cpu_affinity": "0", |
| "output": [ |
| {"output_idx": 0, "dependencies": [{"mod_idx": 2, "input_name": "data_n_1"}]}, |
| ], |
| } |
| mod_config[mods[1]] = { |
| "pipeline": pipe_config2, |
| "target_host": None, |
| "mod_name": "default", |
| "build": None, |
| "params": None, |
| "target": "llvm", |
| "fcompile": None, |
| "dev": tvm.cpu(0), |
| } |
| |
| pipe_config3 = { |
| "mod_idx": 2, |
| "cpu_affinity": "0", |
| "output": [{"output_idx": 0, "dependencies": [{"global_output_index": 0}]}], |
| } |
| mod_config[mods[2]] = { |
| "pipeline": pipe_config3, |
| "target_host": None, |
| "mod_name": "default", |
| "build": None, |
| "params": None, |
| "target": "llvm", |
| "fcompile": None, |
| "dev": tvm.cpu(0), |
| } |
| return mod_config |
| |
| |
| def recreate_parameters(mod): |
| # Get the binding parameters from a module, then create the same parameters with different data. |
| # This function is used to test the "parameter" connection. |
| with tvm.transform.PassContext(opt_level=3): |
| lib = relay.build(mod, "llvm") |
| |
| mod_customized_params = {} |
| for key, value in lib.params.items(): |
| new_value = value.numpy() + np.full(value.shape, 10).astype(value.dtype) |
| mod_customized_params[key] = tvm.nd.array(new_value) |
| return mod_customized_params, mod |
| |
| |
| def run_modules( |
| mod_configs, |
| dev, |
| target, |
| global_input_name, |
| global_input_data, |
| mod_set_input, |
| input_name, |
| input_data, |
| params_mod=None, |
| params=None, |
| ): |
| # Running modules in serialized model. The returnning data are used to verify the pipeline |
| # executor result. |
| mod_input = {} |
| final_output = {} |
| idx = 0 |
| for mod in mod_configs: |
| with tvm.transform.PassContext(opt_level=3): |
| lib = relay.build(mod, target) |
| |
| m = graph_executor.GraphModule(lib["default"](dev)) |
| # Getting the input data then setting the input data into the module. |
| if idx in mod_input: |
| for input in mod_input[idx]: |
| input = mod_input[idx][input] |
| m.set_input(input["index"], input["data"]) |
| else: |
| m.set_input(global_input_name, global_input_data) |
| |
| # Setting the "input_data" into the module. |
| if mod == mod_set_input: |
| m.set_input(input_name, input_data) |
| # If the module is "params_mod" then setting the parameters to this module. |
| if params_mod == mod: |
| m.set_input(None, None, **params) |
| |
| m.run() |
| n = m.get_num_outputs() |
| # Setting current output data as the input of next module. |
| mconfig = mod_configs[mod] |
| for output in mconfig["pipeline"]["output"]: |
| output_data = m.get_output(output["output_idx"]).numpy() |
| for dep in output["dependencies"]: |
| is_global = False |
| if "global_output_index" in dep: |
| is_global = True |
| name = dep["global_output_index"] |
| else: |
| mod_idx = dep["mod_idx"] |
| name = dep["input_name"] |
| if is_global: |
| final_output[name] = output_data |
| else: |
| if mod_idx in mod_input: |
| mod_input[mod_idx][name] = {"index": name, "data": output_data} |
| else: |
| mod_input[mod_idx] = {name: {"index": name, "data": output_data}} |
| idx = idx + 1 |
| |
| return final_output |
| |
| |
| def reset_cpu_affinity(affinity): |
| # Restore the CPU affinity into the default value. |
| config_threadpool = get_global_func("runtime.config_threadpool") |
| config_threadpool(-2, 0) |
| os.sched_setaffinity(0, affinity) |
| |
| |
| def test_pipe_runtime_error_check(): |
| # This function is used to trigger runtime error by applying wrong logic. |
| if pipeline_executor_build.pipeline_executor_build_enabled(): |
| # Get three pipeline modules here. |
| (mod1, mod2, mod3), dshape = get_split_mod() |
| |
| # The input or output name is illegal and expects a runtime error. |
| pipe_error = pipeline_executor_build.PipelineConfig() |
| with pytest.raises(RuntimeError): |
| pipe_error[mod1]["output"][9] |
| |
| with pytest.raises(RuntimeError): |
| pipe_error[mod1]["input"]["data_9"] |
| |
| # The module connection will cause a cycle in DAG and expects runtime error. |
| with pytest.raises(RuntimeError): |
| pipe_error[mod1]["output"][0].connect(pipe_error[mod2]["input"]["data_0"]) |
| pipe_error[mod2]["output"][0].connect(pipe_error[mod1]["input"]["data_0"]) |
| |
| # The module connection is illegal and expects runtime error. |
| |
| with pytest.raises(RuntimeError): |
| pipe_error[mod1]["output"][0].connect(pipe_error[mod1]["input"]["data_0"]) |
| |
| with pytest.raises(RuntimeError): |
| pipe_error[mod1]["input"]["data_0"].connect(pipe_error[mod1]["input"]["data_0"]) |
| |
| with pytest.raises(RuntimeError): |
| pipe_error[mod1]["input"]["data_0"].connect(pipe_error[mod2]["input"]["data_0"]) |
| |
| with pytest.raises(RuntimeError): |
| pipe_error[mod1]["output"][0].connect(pipe_error["input"]["data_0"]) |
| |
| with pytest.raises(RuntimeError): |
| pipe_error["input"]["data_0"].connect(pipe_error[mod1]["output"][0]) |
| |
| with pytest.raises(RuntimeError): |
| pipe_error["output"]["0"].connect(pipe_error[mod1]["output"][0]) |
| |
| # Create pipeline executor to check the executor runtime errors. |
| pipe_config = pipeline_executor_build.PipelineConfig() |
| pipe_config[mod1].target = "llvm" |
| pipe_config[mod1].dev = tvm.cpu(0) |
| pipe_config["param_group"]["param_0"].connect(pipe_config[mod1]["param"]) |
| pipe_config[mod1]["output"][0].connect(pipe_config["output"]["0"]) |
| # Build and create a pipeline module. |
| with tvm.transform.PassContext(opt_level=3): |
| pipeline_mod_factory = pipeline_executor_build.build(pipe_config) |
| pipeline_module = pipeline_executor.PipelineModule(pipeline_mod_factory) |
| customized_parameters, _ = recreate_parameters(mod1) |
| |
| # Checking the pipeline executor runtime errors. |
| with pytest.raises(RuntimeError): |
| pipeline_module.set_params("param_0", None) |
| |
| with pytest.raises(RuntimeError): |
| pipeline_module.set_params("param_1", customized_parameters) |
| |
| |
| def test_pipeline(): |
| if pipeline_executor_build.pipeline_executor_build_enabled(): |
| target_list = tvm.testing.enabled_targets() |
| for target in target_list: |
| affinity = os.sched_getaffinity(0) |
| # Get the three pipeline modules here. |
| (mod1, mod2, mod3), dshape = get_split_mod() |
| |
| # Prepare batch data for pipeline computation. |
| datas = [] |
| for i in range(5): |
| datas.append(np.full(dshape, 3 + i).astype("float32")) |
| |
| pipe_config = pipeline_executor_build.PipelineConfig() |
| |
| customized_parameters, customized_parameters_mod = recreate_parameters(mod1) |
| assert customized_parameters_mod == mod1 |
| # The global parameters group named "param_0" will be connected to "mod1" as parameters. |
| pipe_config["param_group"]["param_0"].connect(pipe_config[mod1]["param"]) |
| # The pipeline input named "data_a" will be connected to a input named "data_0" |
| # of mod1. |
| pipe_config["input"]["data_a"].connect(pipe_config[mod1]["input"]["data_0"]) |
| |
| # The pipeline Input named "data_b" will be connected to a input named "data_1" |
| # of mod2. |
| pipe_config["input"]["data_b"].connect(pipe_config[mod2]["input"]["data_1"]) |
| |
| # The mod1 output[0] will be connected to a input named "data_n_0" of mod2. |
| pipe_config[mod1]["output"][0].connect(pipe_config[mod2]["input"]["data_n_0"]) |
| |
| # The mod1 output[1] will be connected to a input named "data_n_2" of mod3. |
| pipe_config[mod1]["output"][1].connect(pipe_config[mod3]["input"]["data_n_2"]) |
| |
| # The mod2 output[2] will be connected to a input named "data_n_1" of mod3. |
| pipe_config[mod2]["output"][0].connect(pipe_config[mod3]["input"]["data_n_1"]) |
| |
| # The mod3 output[0] will be connected to pipeline output[0]. |
| pipe_config[mod3]["output"][0].connect(pipe_config["output"]["0"]) |
| # Print configuration (print(pipe_config)), the result looks like following. |
| # |
| # Params |
| # |param_0: mod0:param |
| # |
| # Inputs |
| # |data_a: mod0:data_0 |
| # |data_b: mod1:data_1 |
| # |
| # output |
| # |output(0) : mod2.output(0) |
| # |
| # connections |
| # |mod0.output(0)-> mod1.data_n_0 |
| # |mod0.output(1)-> mod2.data_n_2 |
| # |mod1.output(0)-> mod2.data_n_1 |
| |
| # Set other parameters. |
| pipe_config[mod1].target = target[0] |
| pipe_config[mod1].dev = target[1] |
| pipe_config[mod1].cpu_affinity = "0" |
| pipe_config[mod1].fcompile = _cc.create_shared |
| |
| pipe_config[mod2].target = "llvm" |
| pipe_config[mod2].dev = tvm.cpu(0) |
| pipe_config[mod2].cpu_affinity = "0" |
| |
| pipe_config[mod3].target = "llvm" |
| pipe_config[mod3].dev = tvm.cpu(0) |
| pipe_config[mod3].cpu_affinity = "0" |
| # Checking the configuration of modules dependency. |
| mconfig = pipe_config.get_config() |
| assert mconfig["module_connection"] == get_manual_conf([mod1, mod2, mod3], target) |
| |
| # Build and create a pipeline module. |
| with tvm.transform.PassContext(opt_level=3): |
| pipeline_mod_factory = pipeline_executor_build.build(pipe_config) |
| |
| # Export the parameter configuration to a file. |
| directory_path = tvm.contrib.utils.tempdir().temp_dir |
| # If the directory does not exist, create it. |
| if not os.path.exists(directory_path): |
| os.makedirs(directory_path) |
| config_file_name = pipeline_mod_factory.export_library(directory_path) |
| |
| # Use the output of build to create and initialize PipelineModule. |
| pipeline_module = pipeline_executor.PipelineModule(pipeline_mod_factory) |
| assert pipeline_module |
| |
| # Use the import function to create and initialize PipelineModule. |
| pipeline_module_test = pipeline_executor.PipelineModule.load_library(config_file_name) |
| assert pipeline_module_test.num_outputs == 1 |
| |
| input_map = pipeline_module_test.get_input_pipeline_map("data_b") |
| assert input_map[0] == "1" and input_map[1] == "data_1" |
| input_map = pipeline_module_test.get_input_pipeline_map("data_a") |
| assert input_map[0] == "0" and input_map[1] == "data_0" |
| module_index = pipeline_module_test.get_params_group_pipeline_map("param_0") |
| assert module_index == 0 |
| # Using the parameters group name to set parameters. |
| pipeline_module_test.set_params("param_0", customized_parameters) |
| normal_outputs = [] |
| for round in range(0, len(datas)): |
| data = datas[round] |
| # Getting the result without setting customized parameters. |
| wrong_output = run_modules( |
| mconfig["module_connection"], |
| tvm.cpu(), |
| "llvm", |
| "data_0", |
| data, |
| mod2, |
| "data_1", |
| data, |
| ) |
| # Getting the result with setting customized parameters. |
| normal_output = run_modules( |
| mconfig["module_connection"], |
| tvm.cpu(), |
| "llvm", |
| "data_0", |
| data, |
| mod2, |
| "data_1", |
| data, |
| customized_parameters_mod, |
| customized_parameters, |
| ) |
| # Appending the normal output into the list in order to do future correctness |
| # checking. |
| normal_outputs.append(normal_output) |
| # Setting the input data into the pipeline executor. |
| pipeline_module_test.set_input("data_a", tvm.nd.array(data)) |
| pipeline_module_test.set_input("data_b", tvm.nd.array(data)) |
| input_map = pipeline_module_test.get_input_pipeline_map("data_a") |
| # Checking whether the input setting of the first runtime is successful. |
| # The input of the rest of runtime will go into a queue and we can not check |
| # these input data here. |
| if input_map[0] == "0": |
| input_data = pipeline_module_test.get_input("data_a") |
| tvm.testing.assert_allclose(data, input_data.numpy()) |
| |
| assert pipeline_module_test.num_inputs == 2 |
| # Running the pipeline executor in the pipeline mode. |
| pipeline_module_test.run() |
| |
| for k in range(0, len(datas)): |
| statistic_time = 0 |
| outputs = pipeline_module_test.get_output() |
| while len(outputs) == 0: |
| outputs = pipeline_module_test.get_output() |
| statistic_time = statistic_time + 1 |
| # Setting the timeout to 10 seconds. |
| assert statistic_time < 5 |
| time.sleep(1) |
| |
| for i in range(len(outputs)): |
| tvm.testing.assert_allclose(normal_outputs[k][i], outputs[i].numpy()) |
| assert not (normal_output[i] == wrong_output[i]).all() |
| |
| assert pipeline_module_test.num_executing_pipeline == round + 1 |
| |
| # Reset the cpu affinity after a test. |
| reset_cpu_affinity(affinity) |
| |
| |
| if __name__ == "__main__": |
| tvm.testing.main() |