<filename>python/tvm/contrib/msir/core/ir/graph.py
import os
from importlib.machinery import SourceFileLoader

import tvm
from tvm import relay
from tvm.contrib.msir.core import _ffi_api
from tvm.contrib.msir.core import utils as msir_utils


class MSIRGraphModule(object):
    def __init__(self, module):
        self.module = module
        self.name = module["get_name"]()
        self._to_str = module["to_string"]
        self._to_abstract = module["to_abstract"]
        self._to_prototxt = module["to_prototxt"]

    def __str__(self):
        return self._to_str()

    def abstract(self):
        return self._to_abstract()

    def visualize(self,path=None):
        path=path or os.path.join(msir_utils.get_cache_dir(), self.name+".prototxt")
        with open(path,"w") as f:
            f.write(self._to_prototxt())


class MSIRGraphRuntimeModule(object):
    def __init__(self, module, 
        is_train = False,
        need_prune = False,
        need_quantize = False,
        need_collect = False,
        runnable = True):
        self.module = module.imported_modules[1]
        self.target = self.module["get_target"]()
        assert self.module.type_key == "msir_" + self.target, "Moduel key {} mismatch with msir_{}".format( \
            self.module.type_key, self.target)
        self._config_runtime = self.module["config_runtime"]
        self._get_weight_datas = self.module["get_weight_datas"]
        self._get_weight_infos = self.module["get_weight_infos"]
        self.config_runtime(is_train, need_prune, need_quantize, need_collect, runnable)

    def config_runtime(self,
        is_train = False,
        need_prune = False,
        need_quantize = False,
        need_collect = False,
        runnable = True):
        self._config_runtime(msir_utils.get_version(self.target), \
            is_train, need_prune, need_quantize, need_collect, runnable)

    def get_graph(self):
        return MSIRGraphModule(self.module["get_graph"]())

    def get_source(self, fmt="python"):
        if fmt == "python":
            return self.module.get_source(fmt)
        raise Exception("format {} is not support".format(fmt))

    def load_source(self, fmt="python", path=None):
        if fmt == "python":
            path=path or os.path.join(msir_utils.get_build_dir(), self.get_graph().name+".py")
            with open(path, "w") as f:
                f.write(self.get_source(fmt))
            foo = SourceFileLoader(path.replace(".py", ""), path).load_module()
            return getattr(foo, self.get_graph().name)
        raise Exception("format {} is not support".format(fmt))

    def load_weights(self):
        valid_weights = {}
        weight_datas = self._get_weight_datas()
        weight_infos = self._get_weight_infos()
        # TODO : transfrom weights based on optype and ref
        for k, (optype, ref, layout) in weight_infos.items():
            assert k in weight_datas, "Can not find data for " + str(k)
            valid_weights[k] = weight_datas[k].asnumpy()
            if ref == "bias":
                valid_weights[k] = valid_weights[k].reshape(valid_weights[k].size)
        return valid_weights


def build_from_relay(relay_module, entry_name="main", graph_name="msir"):
    module = _ffi_api.BuildFromRelay(relay_module, entry_name, graph_name)
    return MSIRGraphModule(module)


def visualize_relay(relay_module, entry_name="main", graph_name="msir", path=None):
    graph=build_from_relay(relay_module, entry_name, graph_name)
    graph.visualize(path)
<filename>python/tvm/contrib/msir/core/ir/graph.py
import os
from importlib.machinery import SourceFileLoader

import tvm
from tvm import relay
from tvm.contrib.msir.core import _ffi_api
from tvm.contrib.msir.core import utils as msir_utils


class MSIRGraphModule(object):
    def __init__(self, module):
        self.module = module
        self.name = module["get_name"]()
        self._to_str = module["to_string"]
        self._to_abstract = module["to_abstract"]
        self._to_prototxt = module["to_prototxt"]

    def __str__(self):
        return self._to_str()

    def abstract(self):
        return self._to_abstract()

    def visualize(self,path=None):
        path=path or os.path.join(msir_utils.get_cache_dir(), self.name+".prototxt")
        with open(path,"w") as f:
            f.write(self._to_prototxt())


class MSIRGraphRuntimeModule(object):
    def __init__(self, module, 
        is_train = False,
        need_prune = False,
        need_quantize = False,
        need_collect = False,
        runnable = True):
        self.module = module.imported_modules[1]
        self.target = self.module["get_target"]()
        assert self.module.type_key == "msir_" + self.target, "Moduel key {} mismatch with msir_{}".format( \
            self.module.type_key, self.target)
        self._config_runtime = self.module["config_runtime"]
        self._get_weight_datas = self.module["get_weight_datas"]
        self._get_weight_infos = self.module["get_weight_infos"]
        self.config_runtime(is_train, need_prune, need_quantize, need_collect, runnable)

    def config_runtime(self,
        is_train = False,
        need_prune = False,
        need_quantize = False,
        need_collect = False,
        runnable = True):
        self._config_runtime(msir_utils.get_version(self.target), \
            is_train, need_prune, need_quantize, need_collect, runnable)

    def get_graph(self):
        return MSIRGraphModule(self.module["get_graph"]())

    def get_source(self, fmt="python"):
        if fmt == "python":
            return self.module.get_source(fmt)
        raise Exception("format {} is not support".format(fmt))

    def load_source(self, fmt="python", path=None):
        if fmt == "python":
            path=path or os.path.join(msir_utils.get_build_dir(), self.get_graph().name+".py")
            with open(path, "w") as f:
                f.write(self.get_source(fmt))
            foo = SourceFileLoader(path.replace(".py", ""), path).load_module()
            return getattr(foo, self.get_graph().name)
        raise Exception("format {} is not support".format(fmt))

    def load_weights(self):
        valid_weights = {}
        weight_datas = self._get_weight_datas()
        weight_infos = self._get_weight_infos()
        # TODO : transfrom weights based on optype and ref
        for k, (optype, ref, layout) in weight_infos.items():
            assert k in weight_datas, "Can not find data for " + str(k)
            valid_weights[k] = weight_datas[k].asnumpy()
            if ref == "bias":
                valid_weights[k] = valid_weights[k].reshape(valid_weights[k].size)
        return valid_weights


def build_from_relay(relay_module, entry_name="main", graph_name="msir"):
    module = _ffi_api.BuildFromRelay(relay_module, entry_name, graph_name)
    return MSIRGraphModule(module)


def visualize_relay(relay_module, entry_name="main", graph_name="msir", path=None):
    graph=build_from_relay(relay_module, entry_name, graph_name)
    graph.visualize(path)
# TODO : transfrom weights based on optype and ref