Source code for chop.ir.graph.mase_metadata

import logging

from torch import nn

logger = logging.getLogger(__name__)


def get_module_by_name(model, request_name):
    for name, layer in model.named_modules():
        if name == request_name:
            return layer
    return None




class MaseMetadata:
    """
    The metadata of a Mase node in a Mase graph describes the constraints of the
    node for any static analysis or possible transformation. The metadata has a
    tree structure, e.g.
    - common
      - mase_op -> str : the mase op of the node, e.g. placeholder, linear, relu
      - mase_type -> str : the mase type of the node, e.g. module, builtin_func, module_related_func
      - args -> {}
         - $name : name of the arg
           (if the arg is a tensor)
           - type -> type of the arg, e.g. fixed point or float
           - precision -> format of the type, e.g. (10, 5)
           - shape -> shape of the arg
           (if the arg is not a tensor)
           - value of the arg
      - results -> {}
         - $name : name of the result
           (if the result is a tensor)
           - type -> type of the result, e.g. fixed point or float
           - precision -> format of the type, e.g. (10, 5)
           - size -> size of the result
           (if the result is not a tensor)
           - value of the result
    - software: dict
      - args: dict
        - $name (dict): name of the arg, e.g. data_in_0
          - "stat": {"record": {"data": ..., "count": ...},
                     "variance_online": {"variance": ..., "mean": ..., "count": ...}},
                     "variance_precise": {"variance": ..., "mean": ..., "count": ...},
                     "range_n_sigma": {"min": ..., "max": ..., "count": ...},
                     "range_quantile": {"min": ..., "max": ..., "count": ...},
                     "range_min_max": {"min": ..., "max": ..., "count": ...},
                    }.
      - results: dict
        - $name (dict): name of the result, e.g. data_out_0
          - "stat": {"stat_name": { # stat_values } }
    - hardware
      - is_implicit -> bool : whether the node is mapped on hardware or software annotation only
      - verilog_param -> {} : parameters need for customise the hardware module
      - device_id -> int : the ID of the device where the node is mapped, default = -1
      - toolchain -> str : tool chain for code generation, must be INTERNAL, EXTERNAL or HLS
      - module -> str : the name of the used hardware module
      - interface -> {}
         - name : name of the parameters
           - storage : the hardware interface implemented, must be BRAM
           - transpose : whether the data needs to be transposed before emitting
      - dependence_files -> [] : the dependent files for the generated module
    ...
    """

    # Hardware dict
    known_types = ["fixed", "float", "NA"]
    known_toolchain = ["INTERNAL", "EXTERNAL", "HLS"]
    known_storage = ["BRAM"]



    def __init__(
        self,
        node=None,
        model=None,
    ):
        # Top-level model
        self.model = model
        # The fx node of the module in the fx graph of the model
        self.node = node
        # layers that we have in RTL
        self.internal_layers = {nn.Linear: "linear", nn.ReLU: "relu"}

        self.parameters = {
            "common": {},
            "software": {},
            "hardware": {},
        }


    @property
    def module(self):
        # The target module in the model
        # if it is not a "call_module" node, return None
        if self.node.op == "call_module":
            return get_module_by_name(self.model, self.node.target)
        else:
            return None

    @property
    def graph(self):
        # The fx graph of the model
        return self.model.graph

    def __getitem__(self, key):
        return self.parameters[key]

    def __setitem__(self, key, value):
        self.parameters[key] = value