module.py 7.3 KB
Newer Older
1
"""Container of compiled functions of TVM."""
2
from __future__ import absolute_import as _abs
3

4
import struct
5
from collections import namedtuple
Hu Shiwen committed
6

7
from ._ffi.function import ModuleBase, _set_class_module
8
from ._ffi.function import _init_api
9
from .contrib import cc as _cc, tar as _tar, util as _util
10

11
ProfileResult = namedtuple("ProfileResult", ["mean", "results"])
12

Hu Shiwen committed
13

14 15
class Module(ModuleBase):
    """Module container of all TVM generated functions"""
Hu Shiwen committed
16

17 18 19 20 21 22 23 24 25 26 27 28 29 30 31
    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.
32 33 34 35 36

        Returns
        -------
        source : str
            The result source code.
37 38 39 40 41 42 43 44 45
        """
        return _GetSource(self, fmt)

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

        Returns
        ----------
46
        modules : list of Module
47 48
            The module
        """
49
        nmod = _ImportsSize(self)
50 51 52 53 54
        return [_GetImport(self, i) for i in range(nmod)]

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

55 56 57
        This do not save the dependent device modules.
        See also export_shared

58 59 60 61 62 63
        Parameters
        ----------
        file_name : str
            The name of the file.
        fmt : str
            The format of the file.
64 65 66 67

        See Also
        --------
        Module.export_library : export the module to shared library.
68 69 70
        """
        _SaveToFile(self, file_name, fmt)

Tianqi Chen committed
71 72 73 74
    def export_library(self,
                       file_name,
                       fcompile=None,
                       **kwargs):
75 76 77 78 79 80 81 82 83
        """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
84

85
        fcompile : function(target, file_list, kwargs), optional
Tianqi Chen committed
86
            Compilation function to use create dynamic library.
87 88
            If fcompile has attribute object_format, will compile host library
            to that format. Otherwise, will use default format "o".
Tianqi Chen committed
89 90 91

        kwargs : dict, optiona;
            Additional arguments passed to fcompile
92
        """
93 94 95 96
        if self.type_key == "stacktvm":
            raise ValueError("Module[%s]: export_library requires llvm module,"
                             " did you build with LLVM enabled?" % self.type_key)

97
        if self.type_key != "llvm":
98
            raise ValueError("Module[%s]: Only llvm support export shared" % self.type_key)
99
        temp = _util.tempdir()
100 101 102 103 104
        if fcompile is not None and hasattr(fcompile, "object_format"):
            object_format = fcompile.object_format
        else:
            object_format = "o"
        path_obj = temp.relpath("lib." + object_format)
105 106
        self.save(path_obj)
        files = [path_obj]
Tianqi Chen committed
107
        is_system_lib = self.get_function("__tvm_is_system_module")()
108 109 110
        if self.imported_modules:
            path_cc = temp.relpath("devc.cc")
            with open(path_cc, "w") as f:
111
                f.write(_PackImportsToC(self, is_system_lib))
112
            files.append(path_cc)
113 114 115 116 117
        if not fcompile:
            if file_name.endswith(".tar"):
                fcompile = _tar.tar
            else:
                fcompile = _cc.create_shared
Tianqi Chen committed
118
        fcompile(file_name, files, **kwargs)
119

120
    def time_evaluator(self, func_name, ctx, number, repeat=1):
121 122 123 124 125 126 127 128 129 130 131
        """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
132 133 134 135 136
            The number of steps used in measuring each time interval

        repeat: int, optional
            Number of times to run the timer measurement
            If repeat equals 3, then we will get 3 numbers in the ProfileResult.
137 138 139

        Note
        ----
140
        The function will be invoked  repeat * number + 1 times,
141 142 143 144 145 146 147 148 149
        with the first call discarded in case there is lazy initialization.

        Returns
        -------
        ftimer : Function
            The function that takes same argument as func
            and return a float representing seconds per function call.
        """
        try:
150
            feval = _RPCTimeEvaluator(
151
                self, func_name, ctx.device_type, ctx.device_id, number, repeat)
Hu Shiwen committed
152

153 154 155
            def evaluator(*args):
                """Internal wrapped evaluator."""
                # Wrap feval so we can add more stats in future.
156 157 158 159 160
                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
161

162
            return evaluator
163 164 165
        except NameError:
            raise NameError("time_evaluate is only supported when RPC is enabled")

166

167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187
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()


188
def load(path, fmt=""):
189
    """Load module from file.
190 191 192 193 194 195 196 197 198

    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.
199 200 201 202 203

    Returns
    -------
    module : Module
        The loaded module
204 205 206 207 208

    Note
    ----
    This function will automatically call
    cc.create_shared if the path is in format .o or .tar
209
    """
210 211 212 213 214 215 216 217 218 219 220 221
    # 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"):
        tar_temp = _util.tempdir()
        _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
222 223
    return _LoadFromFile(path, fmt)

224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247

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")
248
_set_class_module(Module)