tvm.contrib.graph_executor

Minimum graph executor that executes graph containing TVM PackedFunc.

class tvm.contrib.graph_executor.GraphModule(module)

Wrapper runtime module.

This is a thin wrapper of the underlying TVM module. you can also directly call set_input, run, and get_output of underlying module functions

Parameters

moduletvm.runtime.Module

The internal tvm module that holds the actual graph functions.

Attributes

moduletvm.runtime.Module

The internal tvm module that holds the actual graph functions.

Examples

import tvm
from tvm import relay
from tvm.contrib import graph_executor

# build the library using graph executor
lib = relay.build(...)
lib.export_library("compiled_lib.so")
# load it back as a runtime
lib: tvm.runtime.Module = tvm.runtime.load_module("compiled_lib.so")
# Call the library factory function for default and create
# a new runtime.Module, wrap with graph module.
gmod = graph_executor.GraphModule(lib["default"](dev))
# use the graph module.
gmod.set_input("x", data)
gmod.run()

__getitem__(key)

Get internal module function

Parameters

keystr

The key to the module.

benchmark(device, func_name='run', repeat=5, number=5, min_repeat_ms=None, limit_zero_time_iterations=100, end_to_end=False, cooldown_interval_ms=0, repeats_to_cooldown=1, **kwargs)

Calculate runtime of a function by repeatedly calling it.

Use this function to get an accurate measurement of the runtime of a function. The function is run multiple times in order to account for variability in measurements, processor speed or other external factors. Mean, median, standard deviation, min and max runtime are all reported. On GPUs, CUDA and ROCm specifically, special on-device timers are used so that synchonization and data transfer operations are not counted towards the runtime. This allows for fair comparison of runtimes across different functions and models. The end_to_end flag switches this behavior to include data transfer operations in the runtime.

The benchmarking loop looks approximately like so:

for r in range(repeat):
    time_start = now()
    for n in range(number):
        func_name()
    time_end = now()
    total_times.append((time_end - time_start)/number)

Parameters

func_namestr

The function to benchmark. This is ignored if end_to_end is true.

repeatint

Number of times to run the outer loop of the timing code (see above). The output will contain repeat number of datapoints.

numberint

Number of times to run the inner loop of the timing code. This inner loop is run in between the timer starting and stopping. In order to amortize any timing overhead, number should be increased when the runtime of the function is small (less than a 1/10 of a millisecond).

min_repeat_msOptional[int]

If set, the inner loop will be run until it takes longer than min_repeat_ms milliseconds. This can be used to ensure that the function is run enough to get an accurate measurement.

limit_zero_time_iterationsOptional[int]

The maximum number of repeats when measured time is equal to 0. It helps to avoid hanging during measurements.

end_to_endbool

If set, include time to transfer input tensors to the device and time to transfer returned tensors in the total runtime. This will give accurate timings for end to end workloads.

cooldown_interval_ms: Optional[int]

The cooldown interval in milliseconds between the number of repeats defined by repeats_to_cooldown.

repeats_to_cooldown: Optional[int]

The number of repeats before the cooldown is activated.

kwargsDict[str, Object]

Named arguments to the function. These are cached before running timing code, so that data transfer costs are not counted in the runtime.

Returns

timing_resultsBenchmarkResult

Runtimes of the function. Use .mean to access the mean runtime, use .results to access the individual runtimes (in seconds).

debug_get_output(node, out)

Run graph up to node and get the output to out

Parameters

nodeint / str

The node index or name

outNDArray

The output array container

get_input(index, out=None)

Get index-th input to out

Parameters

indexint

The input index

outNDArray

The output array container

get_input_index(name)

Get inputs index via input name.

Parameters

namestr

The input key name

Returns

index: int

The input index. -1 will be returned if the given input name is not found.

get_input_info()

Return the 'shape' and 'dtype' dictionaries of the graph.

备注

We can't simply get the input tensors from a TVM graph because weight tensors are treated equivalently. Therefore, to find the input tensors we look at the 'arg_nodes' in the graph (which are either weights or inputs) and check which ones don't appear in the params (where the weights are stored). These nodes are therefore inferred to be input tensors.

Returns

shape_dictMap

Shape dictionary - {input_name: tuple}.

dtype_dictMap

dtype dictionary - {input_name: dtype}.

get_num_inputs()

Get the number of inputs to the graph

Returns

countint

The number of inputs.

get_num_outputs()

Get the number of outputs from the graph

Returns

countint

The number of outputs.

get_output(index, out=None)

Get index-th output to out

Parameters

indexint

The output index

outNDArray

The output array container

load_params(params_bytes)

Load parameters from serialized byte array of parameter dict.

Parameters

params_bytesbytearray

The serialized parameter dict.

run(**input_dict)

Run forward execution of the graph

Parameters

input_dict: dict of str to NDArray

List of input values to be feed to

set_input(key=None, value=None, **params)

Set inputs to the module via kwargs

Parameters

keyint or str

The input key

valuethe input value.

The input value

paramsdict of str to NDArray

Additional arguments

set_input_zero_copy(key=None, value=None, **params)

Set inputs to the module via kwargs with zero memory copy

Parameters

keyint or str

The input key

valuethe input value in DLPack

The input value

paramsdict of str to NDArray

Additional arguments

set_output_zero_copy(key, value)

Set outputs to the module with zero memory copy

Parameters

keyint or str

The output key

valuethe output value in DLPack

The output value

share_params(other, params_bytes)

Share parameters from pre-existing GraphExecutor instance.

Parameters

other: GraphExecutor

The parent GraphExecutor from which this instance should share it's parameters.

params_bytesbytearray

The serialized parameter dict (used only for the parameter names).

tvm.contrib.graph_executor.create(graph_json_str, libmod, device)

Create a runtime executor module given a graph and module.

Parameters

graph_json_strstr

The graph to be deployed in json format output by json graph. The graph can contain operator(tvm_op) that points to the name of PackedFunc in the libmod.

libmodtvm.runtime.Module

The module of the corresponding function

deviceDevice or list of Device

The device to deploy the module. It can be local or remote when there is only one Device. Otherwise, the first device in the list will be used as this purpose. All device should be given for heterogeneous execution.

Returns

graph_moduleGraphModule

Runtime graph module that can be used to execute the graph.

Note

See also tvm.contrib.graph_executor.GraphModule for examples to directly construct a GraphModule from an exported relay compiled library.

tvm.contrib.graph_executor.get_device(libmod, device)

Parse and validate all the device(s).

Parameters

libmodtvm.runtime.Module

The module of the corresponding function

device : Device or list of Device

Returns

device : list of Device num_rpc_dev : Number of rpc devices device_type_id : List of device type and device id