graph_runtime.py 8.21 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# 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.
17
"""Minimum graph runtime that executes graph containing TVM PackedFunc."""
18 19
import numpy as np

20 21
from .._ffi.base import string_types
from .._ffi.function import get_global_func
Zhi committed
22
from .._ffi.runtime_ctypes import TVMContext
23
from ..rpc import base as rpc_base
24 25 26 27 28 29 30 31 32 33 34

def create(graph_json_str, libmod, ctx):
    """Create a runtime executor module given a graph and module.
    Parameters
    ----------
    graph_json_str : str or graph class
        The graph to be deployed in json format output by nnvm graph.
        The graph can only contain one operator(tvm_op) that
        points to the name of PackedFunc in the libmod.
    libmod : tvm.Module
        The module of the corresponding function
Zhi committed
35 36 37 38 39
    ctx : TVMContext or list of TVMContext
        The context to deploy the module. It can be local or remote when there
        is only one TVMContext. Otherwise, the first context in the list will
        be used as this purpose. All context should be given for heterogeneous
        execution.
40 41 42 43 44 45 46 47 48 49
    Returns
    -------
    graph_module : GraphModule
        Runtime graph module that can be used to execute the graph.
    """
    if not isinstance(graph_json_str, string_types):
        try:
            graph_json_str = graph_json_str._tvm_graph_json()
        except AttributeError:
            raise ValueError("Type %s is not supported" % type(graph_json_str))
50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74

    ctx, num_rpc_ctx, device_type_id = get_device_ctx(libmod, ctx)

    if num_rpc_ctx == len(ctx):
        hmod = rpc_base._ModuleHandle(libmod)
        fcreate = ctx[0]._rpc_sess.get_function("tvm.graph_runtime.remote_create")
        return GraphModule(fcreate(graph_json_str, hmod, *device_type_id))

    fcreate = get_global_func("tvm.graph_runtime.create")
    return GraphModule(fcreate(graph_json_str, libmod, *device_type_id))

def get_device_ctx(libmod, ctx):
    """Parse and validate all the device context(s).
    Parameters
    ----------
    libmod : tvm.Module
        The module of the corresponding function
    ctx : TVMContext or list of TVMContext
    Returns
    -------
    ctx : list of TVMContext
    num_rpc_ctx : Number of rpc contexts
    device_type_id : List of device type and device id
    """

Zhi committed
75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
    if isinstance(ctx, TVMContext):
        ctx = [ctx]
    elif not isinstance(ctx, (list, tuple)):
        raise ValueError("ctx has to be the type of TVMContext or a list of "
                         "TVMCTVMContext")
    for cur_ctx in ctx:
        if not isinstance(cur_ctx, TVMContext):
            raise ValueError("ctx has to be the type of TVMContext or a list "
                             "of TVMContext")

    # device_type_id[0], device_type_id[1] are used as the primary/fallback
    # context type and id. All other ones are used as device context for
    # heterogeneous execution.
    num_rpc_ctx = 0
    device_type_id = []
    for cur_ctx in ctx:
        device_type = cur_ctx.device_type
        if device_type >= rpc_base.RPC_SESS_MASK:
            assert libmod.type_key == "rpc"
            assert rpc_base._SessTableIndex(
                libmod) == cur_ctx._rpc_sess._tbl_index
            num_rpc_ctx += 1
            device_type = cur_ctx.device_type % rpc_base.RPC_SESS_MASK
        device_type_id.append(device_type)
        device_type_id.append(cur_ctx.device_id)

    if 0 < num_rpc_ctx < len(ctx):
        raise ValueError("Either all or none of the contexts should be rpc.")
103
    return ctx, num_rpc_ctx, device_type_id
104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122


class GraphModule(object):
    """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
    ----------
    module : Module
        The interal tvm module that holds the actual graph functions.

    Attributes
    ----------
    module : Module
        The interal tvm module that holds the actual graph functions.
    """
Zhi committed
123 124

    def __init__(self, module):
125 126 127 128
        self.module = module
        self._set_input = module["set_input"]
        self._run = module["run"]
        self._get_output = module["get_output"]
129
        self._get_input = module["get_input"]
130
        self._get_num_outputs = module["get_num_outputs"]
131
        self._load_params = module["load_params"]
132
        self._share_params = module["share_params"]
133 134 135 136 137 138 139 140 141 142 143 144 145 146 147

    def set_input(self, key=None, value=None, **params):
        """Set inputs to the module via kwargs

        Parameters
        ----------
        key : int or str
           The input key

        value : the input value.
           The input key

        params : dict of str to NDArray
           Additonal arguments
        """
148
        if key is not None:
Zhi committed
149
            self._get_input(key).copyfrom(value)
150 151 152 153 154 155

        if params:
            # upload big arrays first to avoid memory issue in rpc mode
            keys = list(params.keys())
            keys.sort(key=lambda x: -np.prod(params[x].shape))
            for k in keys:
Zhi committed
156
                self._get_input(k).copyfrom(params[k])
157 158 159 160 161 162 163 164 165 166 167 168 169

    def run(self, **input_dict):
        """Run forward execution of the graph

        Parameters
        ----------
        input_dict: dict of str to NDArray
            List of input values to be feed to
        """
        if input_dict:
            self.set_input(**input_dict)
        self._run()

170 171 172 173 174 175 176 177 178 179 180
    def get_num_outputs(self):
        """Get the number of outputs from the graph

        Returns
        -------
        count : int
            The number of outputs.
        """
        return self._get_num_outputs()

    def get_input(self, index, out=None):
181 182 183 184 185 186 187 188 189 190
        """Get index-th input to out

        Parameters
        ----------
        index : int
            The input index

        out : NDArray
            The output array container
        """
191 192 193
        if out:
            self._get_input(index).copyto(out)
            return out
194

195 196 197
        return self._get_input(index)

    def get_output(self, index, out=None):
198 199 200 201 202
        """Get index-th output to out

        Parameters
        ----------
        index : int
203
            The output index
204 205 206 207

        out : NDArray
            The output array container
        """
208 209 210 211 212
        if out:
            self._get_output(index, out)
            return out

        return self._get_output(index)
213

214 215 216 217 218 219 220 221 222 223 224
    def debug_get_output(self, node, out):
        """Run graph upto node and get the output to out

        Parameters
        ----------
        node : int / str
            The node index or name

        out : NDArray
            The output array container
        """
225 226
        raise NotImplementedError(
            "Please use debugger.debug_runtime as graph_runtime instead.")
227

228 229 230 231 232 233 234 235 236 237
    def load_params(self, params_bytes):
        """Load parameters from serialized byte array of parameter dict.

        Parameters
        ----------
        params_bytes : bytearray
            The serialized parameter dict.
        """
        self._load_params(bytearray(params_bytes))

238 239 240 241 242 243 244 245 246 247 248 249 250
    def share_params(self, other, params_bytes):
        """Share parameters from pre-existing GraphRuntime instance.

        Parameters
        ----------
        other: GraphRuntime
            The parent GraphRuntime from which this instance should share
            it's parameters.
        params_bytes : bytearray
            The serialized parameter dict (used only for the parameter names).
        """
        self._share_params(other.module, bytearray(params_bytes))

251 252 253 254 255 256 257 258 259
    def __getitem__(self, key):
        """Get internal module function

        Parameters
        ----------
        key : str
            The key to the module.
        """
        return self.module[key]