module.py 9.45 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
"""Container of compiled functions of TVM."""
18
from __future__ import absolute_import as _abs
19

20
import struct
21
from collections import namedtuple
Hu Shiwen committed
22

23
from ._ffi.function import ModuleBase, _set_class_module
24
from ._ffi.function import _init_api
25
from ._ffi.libinfo import find_include_path
26
from .contrib import cc as _cc, tar as _tar, util as _util
27

28
ProfileResult = namedtuple("ProfileResult", ["mean", "results"])
29

Hu Shiwen committed
30

31 32
class Module(ModuleBase):
    """Module container of all TVM generated functions"""
Hu Shiwen committed
33

34 35 36 37 38 39 40 41 42 43 44 45 46 47 48
    def __repr__(self):
        return "Module(%s, %x)" % (self.type_key, self.handle.value)

    @property
    def type_key(self):
        """Get type key of the module."""
        return _GetTypeKey(self)

    def get_source(self, fmt=""):
        """Get source code from module, if available.

        Parameters
        ----------
        fmt : str, optional
            The specified format.
49 50 51 52 53

        Returns
        -------
        source : str
            The result source code.
54 55 56 57 58 59 60 61 62
        """
        return _GetSource(self, fmt)

    @property
    def imported_modules(self):
        """Get imported modules

        Returns
        ----------
63
        modules : list of Module
64 65
            The module
        """
66
        nmod = _ImportsSize(self)
67 68 69 70 71
        return [_GetImport(self, i) for i in range(nmod)]

    def save(self, file_name, fmt=""):
        """Save the module to file.

72 73 74
        This do not save the dependent device modules.
        See also export_shared

75 76 77 78 79 80
        Parameters
        ----------
        file_name : str
            The name of the file.
        fmt : str
            The format of the file.
81 82 83 84

        See Also
        --------
        Module.export_library : export the module to shared library.
85 86 87
        """
        _SaveToFile(self, file_name, fmt)

Tianqi Chen committed
88 89 90 91
    def export_library(self,
                       file_name,
                       fcompile=None,
                       **kwargs):
92 93 94 95 96 97 98 99 100
        """Export the module and its imported device code one library.

        This function only works on host llvm modules.
        It will pack all the imported modules

        Parameters
        ----------
        file_name : str
            The name of the shared library.
Tianqi Chen committed
101

102
        fcompile : function(target, file_list, kwargs), optional
Tianqi Chen committed
103
            Compilation function to use create dynamic library.
104 105
            If fcompile has attribute object_format, will compile host library
            to that format. Otherwise, will use default format "o".
Tianqi Chen committed
106

107
        kwargs : dict, optional
Tianqi Chen committed
108
            Additional arguments passed to fcompile
109
        """
110 111 112 113
        from pathlib import Path
        if isinstance(file_name, Path):
            file_name = str(file_name)

114 115 116 117 118 119
        if self.type_key == "stackvm":
            if not file_name.endswith(".stackvm"):
                raise ValueError("Module[%s]: can only be saved as stackvm format."
                                 "did you build with LLVM enabled?" % self.type_key)
            self.save(file_name)
            return
120

121 122
        if not (self.type_key == "llvm" or self.type_key == "c"):
            raise ValueError("Module[%s]: Only llvm and c support export shared" % self.type_key)
123
        temp = _util.tempdir()
124 125 126
        if fcompile is not None and hasattr(fcompile, "object_format"):
            object_format = fcompile.object_format
        else:
127 128 129 130 131
            if self.type_key == "llvm":
                object_format = "o"
            else:
                assert self.type_key == "c"
                object_format = "cc"
132
        path_obj = temp.relpath("lib." + object_format)
133 134
        self.save(path_obj)
        files = [path_obj]
135
        is_system_lib = self.type_key == "llvm" and self.get_function("__tvm_is_system_module")()
136 137 138
        if self.imported_modules:
            path_cc = temp.relpath("devc.cc")
            with open(path_cc, "w") as f:
139
                f.write(_PackImportsToC(self, is_system_lib))
140
            files.append(path_cc)
141 142 143 144 145
        if not fcompile:
            if file_name.endswith(".tar"):
                fcompile = _tar.tar
            else:
                fcompile = _cc.create_shared
146 147
        if self.type_key == "c":
            kwargs.update({'options': ["-I" + path for path in find_include_path()]})
Tianqi Chen committed
148
        fcompile(file_name, files, **kwargs)
149

150
    def time_evaluator(self, func_name, ctx, number=10, repeat=1, min_repeat_ms=0):
151 152 153 154 155 156 157 158 159 160 161
        """Get an evaluator that measures time cost of running function.

        Parameters
        ----------
        func_name: str
            The name of the function in the module.

        ctx: TVMContext
            The context we should run this function on.

        number: int
162 163
            The number of times to run this function for taking average.
            We call these runs as one `repeat` of measurement.
164 165

        repeat: int, optional
166 167 168 169 170 171 172 173 174 175 176 177 178
            The number of times to repeat the measurement.
            In total, the function will be invoked (1 + number x repeat) times,
            where the first one is warm up and will be discarded.
            The returned result contains `repeat` costs,
            each of which is an average of `number` costs.

        min_repeat_ms: int, optional
            The minimum duration of one `repeat` in milliseconds.
            By default, one `repeat` contains `number` runs. If this parameter is set,
            the parameters `number` will be dynamically adjusted to meet the
            minimum duration requirement of one `repeat`.
            i.e., When the run time of one `repeat` falls below this time, the `number` parameter
            will be automatically increased.
179 180 181

        Note
        ----
182
        The function will be invoked  (1 + number x repeat) times,
183 184 185 186 187
        with the first call discarded in case there is lazy initialization.

        Returns
        -------
        ftimer : Function
188 189
            The function that takes same argument as func and returns a ProfileResult.
            The ProfileResult reports `repeat` time costs in seconds.
190 191
        """
        try:
192
            feval = _RPCTimeEvaluator(
193
                self, func_name, ctx.device_type, ctx.device_id, number, repeat, min_repeat_ms)
Hu Shiwen committed
194

195 196 197
            def evaluator(*args):
                """Internal wrapped evaluator."""
                # Wrap feval so we can add more stats in future.
198 199 200 201 202
                blob = feval(*args)
                fmt = "@" + ("d" * repeat)
                results = struct.unpack(fmt, blob)
                mean = sum(results) / float(repeat)
                return ProfileResult(mean=mean, results=results)
Hu Shiwen committed
203

204
            return evaluator
205 206 207
        except NameError:
            raise NameError("time_evaluate is only supported when RPC is enabled")

208

209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229
def system_lib():
    """Get system-wide library module singleton.

    System lib is a global module that contains self register functions in startup.
    Unlike normal dso modules which need to be loaded explicitly.
    It is useful in environments where dynamic loading api like dlopen is banned.

    To build system lib function, simply specify target option ```llvm --system-lib```
    The system lib will be available as long as the result code is linked by the program.

    The system lib is intended to be linked and loaded during the entire life-cyle of the program.
    If you want dynamic loading features, use dso modules instead.

    Returns
    -------
    module : Module
        The system-wide library module.
    """
    return _GetSystemLib()


230
def load(path, fmt=""):
231
    """Load module from file.
232 233 234 235 236 237 238 239 240

    Parameters
    ----------
    path : str
        The path to the module file.

    fmt : str, optional
        The format of the file, if not specified
        it will be inferred from suffix of the file.
241 242 243 244 245

    Returns
    -------
    module : Module
        The loaded module
246 247 248 249 250

    Note
    ----
    This function will automatically call
    cc.create_shared if the path is in format .o or .tar
251
    """
252 253 254 255 256 257
    # High level handling for .o and .tar file.
    # We support this to be consistent with RPC module load.
    if path.endswith(".o"):
        _cc.create_shared(path + ".so", path)
        path += ".so"
    elif path.endswith(".tar"):
258
        tar_temp = _util.tempdir(custom_path=path.replace('.tar', ''))
259 260 261 262 263
        _tar.untar(path, tar_temp.temp_dir)
        files = [tar_temp.relpath(x) for x in tar_temp.listdir()]
        _cc.create_shared(path + ".so", files)
        path += ".so"
    # Redirect to the load API
264 265
    return _LoadFromFile(path, fmt)

266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289

def enabled(target):
    """Whether module runtime is enabled for target

    Parameters
    ----------
    target : str
        The target device type.

    Returns
    -------
    enabled : bool
        Whether runtime is enabled.

    Examples
    --------
    The following code checks if gpu is enabled.

    >>> tvm.module.enabled("gpu")
    """
    return _Enabled(target)


_init_api("tvm.module")
290
_set_class_module(Module)