def __init__(self, main_model: ModuleStruct):
super(MatcherLauncher).__init__()
self.main_model = main_model
self._global_context = GlobalContext()
self._graph_inputs = self._global_context.onnx_graph_info.get(
"graph_inputs")
self._graph_outputs = self._global_context.onnx_graph_info.get(
"graph_outputs")
def public_module_shared_weight_statement_generation(
public_module: ModuleStruct):
"""Return the statement of declaration of shared weights in its public module."""
statements = []
for passthrough_w_onnx_name, passthrough_w_var_name in public_module.shared_weights_collection.items(
):
parameter_statement = GlobalContext(
).repeated_weights_declaration.get(passthrough_w_onnx_name)
declare_statement = f"self.{passthrough_w_var_name} = {parameter_statement}"
statements.append(declare_statement)
return statements
def main_model_special_process_inputs(main_model: ModuleStruct):
"""Call in preprocess"""
# allocate main model construct x
prec_edges = main_model.external_precursor_nodes_names
graph_inputs = GlobalContext().onnx_graph_info.get('graph_inputs')
inputs = dict()
for edge in graph_inputs:
if not edge in inputs and edge in prec_edges:
regular_edge = MatcherHelper.regular_edge_name(edge)
inputs[edge] = regular_edge
main_model.inputs_register = inputs
def __init__(self, sid, module_name, nodes, dataloader, merged_modules):
global MODULE_NAME_MGR, MODULE_NAME_INDEXING
self.sid = sid
if module_name in MODULE_NAME_MGR:
self.fake_module_name = MODULE_NAME_MGR[module_name]
else:
self.fake_module_name = f"Module{MODULE_NAME_INDEXING}"
MODULE_NAME_INDEXING += 1
MODULE_NAME_MGR[module_name] = self.fake_module_name
self.module_name = module_name
if module_name not in GlobalContext().known_module_name:
GlobalContext().known_module_name[
self.fake_module_name] = module_name
self.nodes = {
nd: dataloader.nodes_dict.get(nd) or merged_modules.get(nd)
for nd in nodes
}
self.heads, self.inputs, self.tails, self.outputs = get_area_heads_and_tails(
self.nodes, dataloader)
self.start_rank, self.end_rank = get_area_rank(self.heads, self.tails,
dataloader)
def __init__(self):
"""Init the generator."""
# define MUST have params
self._node_struct_collections = OrderedDict()
self._module_struct_collections = OrderedDict()
self._module_depth_max = 0
self._module_depth_min = 0
# define intermediate var. during conversion
self._module_map = OrderedDict()
self._global_context = GlobalContext()
self._global_context.node_struct_collections = self._node_struct_collections
self._repeated_submodules = set()
def check_node_has_shared_weight(node: NodeStruct):
"""
Check the node has shared weight and return all of them.
Args:
node (NodeStruct): NodeStruct instance.
Returns:
list, a list of shared weight onnx names
"""
shared_weight_names = []
for shared_weight_name, repeated_node_list in GlobalContext(
).repeated_weights.items():
if node.onnx_name in repeated_node_list:
shared_weight_names.append(shared_weight_name)
return shared_weight_names
def _add_known_module_name(search_path):
"""
Add known module name to GlobalContext.
Args:
search_path (SearchPath): Search path.
"""
ctx = GlobalContext()
if search_path.pattern.known_module_name:
ctx.known_module_name[
search_path.pattern.
module_name] = search_path.pattern.known_module_name
for it in search_path.recursion_path:
if it.pattern.known_module_name:
ctx.known_module_name[
it.pattern.module_name] = it.pattern.known_module_name
return ctx
def __init__(self,
onnx_model,
model_path: str,
input_nodes: dict,
output_nodes: list,
infer_shape=True):
log_console.info("Onnx simplifying begins.")
onnx_sim = OnnxSimplify()
onnx_model_sim = onnx_sim.run_onnx_simplify(onnx_model, model_path,
input_nodes)
log_console.info("Onnx simplifying is finished.")
self.model = onnx_model_sim
self.model_path = model_path
self.graph = onnx_model_sim.graph
self.nodes = onnx_model_sim.graph.node
self.input_nodes = input_nodes
self.output_nodes = output_nodes
# args for init
self._is_infer_shape = infer_shape
self._global_context = GlobalContext()
# params parsed in init
self.inferred_model = None
self._nodes_dict = OrderedDict() # {node_name: OnnxNode} NO INPUT NODE
self.tensors_dict = {} # {tensor_name: OnnxTensor}
self.value_info_dict = {} # Not contains input and output nodes
# Record the weight names used many times.
self.repeated_weight = dict()
self.node_output_shape_dict = OrderedDict() # {node_name: [int]}
# Key is edge of ONNX ir graph, value is the corresponding precursor node.
self.output_name_to_node_name = dict()
# Define dynamic nodes to be evaluated with onnxruntime
self.dynamic_resize_node = list()
self.dynamic_reshape_node = list()
self.eliminated_nodes = list()
# Validate init params
self._check_user_provided_info()
self.initialize()
def _recursive_form_module(self):
"""Main routine in generator to build modules from bottom to top."""
# 1. List repeated submodules
repeated_submodules = self._list_repeated_submodules()
# 2. List reused parameters
formal_parameters = self._list_formal_parameters(repeated_submodules)
# 3. Build base subdmodules and set in/ext params translation
for module_struct in self.module_structs.values():
if module_struct.pattern_id == -1: # is main module
continue
formal_args = formal_parameters.get(module_struct.pattern_id)
module_struct.update_args_translation_list(formal_args)
# 4. Form parent modules
md_collection_len = len(self.module_structs.keys())
len_changes = True
while len_changes:
self._add_submodule_to_parent()
new_len = len(self.module_structs.keys())
if md_collection_len != new_len:
md_collection_len = new_len
else:
len_changes = False
GlobalContext().build_struct_finished = True
# 5. Update all translated args from module map
self._update_all_modules_args_translator()
# 6. Update all nodes and moudles input/output
self.build_outputs_connection()
self.module_structs.get('[]').allocate_construct_header_x()
self.module_structs.get('[]').collect_returns()
matcher = MatcherLauncher(self.module_structs.get('[]'))
matcher.matching_process()
for nd_struct in self.node_structs.values():
if nd_struct.fragment.metadata.get("operation") == "Split":
self._split_op_procs(nd_struct)
def _generate_from_module_struct(self, md_struct, repeated_submodules):
"""
Generate the code of current Module Struct, collecting data from submodules.
Args:
md_struct (ModuleStruct): The ModuleStruct which generates codes.
repeated_submodules (dict): The dict contains all submodules which use repeatedly.
Can get this dict from generator.
"""
# Define Module header code line below
class_name = md_struct.class_name
# define a class declaration
self.new_line = f"class {class_name}(nn.Cell):"
# Get all formal args from nodes
module_def_args = ['self']
if md_struct.args_translator.actual_args:
for actual in md_struct.args_translator.actual_args.keys():
module_def_args.append(actual)
if md_struct.args_translator.formal_args:
for formal in md_struct.args_translator.formal_args.keys():
module_def_args.append(formal)
# set passthrough weights for shared weights, no need for main model
if md_struct.identifier != []:
module_def_args = SharedWeightHelper.add_shared_weights_in_init_statement(md_struct, module_def_args)
# For code line in init & construct blocks
init_lines = list()
cons_lines = list()
for (_, struct) in md_struct.get_generate_order():
if isinstance(struct, NodeStruct): # Generate code line for Node.
_ = struct.code_line_in_init()
_ = struct.code_line_in_construct()
init_str, cons_str = struct.fragment.fragment()
init_str = [f"{SECOND_LEVEL_INDENT}{x}" for x in init_str]
cons_str = [f"{SECOND_LEVEL_INDENT}{x}" for x in cons_str]
code_line_construct = cons_str
init_lines += init_str
cons_lines += cons_str
else: # is ModuleStruct
# check if this instance generated CodeStruct
if GlobalContext().code_structs.get(struct.pattern_id) is None:
CodeStruct(struct, repeated_submodules)
code_line_init = struct.code_line_in_init()
code_line_construct = struct.code_line_in_construct(inputs=struct.matched_inputs)
init_lines.append(f"{SECOND_LEVEL_INDENT}{' = '.join(code_line_init)}")
cons_lines.append(f"{SECOND_LEVEL_INDENT}{' = '.join(code_line_construct)}")
# define header of init block
self.new_line = f"{FIRST_LEVEL_INDENT}def __init__({', '.join(module_def_args)}):"
self.new_line = f"{SECOND_LEVEL_INDENT}super({class_name}, self).__init__()"
#add shared weights declaration in init code part
if md_struct.identifier == []:
passthrough_w_declaration = SharedWeightHelper.public_module_shared_weight_statement_generation(md_struct)
for s in passthrough_w_declaration:
self.new_line = f"{SECOND_LEVEL_INDENT}{s}"
# add init code lines to code line list.
self.code_line_list += init_lines
self.new_line = f"{NEW_LINE * 2}"
# define header of construct block
inputs = ['self'] + list(md_struct.inputs_register.values())
self.new_line = f"{FIRST_LEVEL_INDENT}def construct({', '.join(inputs)}):"
# add construct code lines to code line list.
self.code_line_list += cons_lines
# define returns
returns = []
# take opt_var_name to return_list
for output_edge in md_struct.outputs_register.keys():
opt_var_name = md_struct.internal_outputs_collection.get(output_edge)
if opt_var_name is None:
raise ValueError(f"Module {md_struct.identifier} has an output {output_edge} has unknown opt_var_name.")
returns.append(opt_var_name)
self.new_line = f"{SECOND_LEVEL_INDENT}return {', '.join(returns)}"
self.new_line = f"{NEW_LINE * 2}"
GlobalContext().code_structs[md_struct.pattern_id] = self
def __init__(self, nd_struct_list, init_as_parent=False, parent_base=None):
"""Init. a module by NodeStructs."""
self.pattern_id = -1 # pattern num, -1 as Main module
self.pattern_uid = -1 # unique module id for this pattern
self.parent_id = None # parent's pattern num
self.parent_uid = None # parent's pattern module unique id
self.initialized = False
self.identifier = None
self.module_name = None
self.scope_depth = None
self.head_nd_struct = None
self.head_nd_struct_index = None
self.tail_nd_struct = None
self.tail_nd_struct_index = None
self._node_structs = list()
self._module_structs = list()
self._fragment = None
self._args_translator = None
self._parent_module_struct = None
# only store original formal args name, not global
self._nodes_structs_formal_args_list = list()
# define other settings here
self._node_args_translation_list = list()
self._var_name_mgr = LocalVarNameMgr()
self.construct_header_x = OrderedDict(
) # key is header x, value is precursors onnx name
self.inputs_in_construct_header = OrderedDict(
) # key is precursors onnx name, value is x in parent construct
# key is node's onnx name(output provider), value is (provider_succ_name, opt_var_name)
self.outputs_collection = dict()
self.matched_inputs = list(
) # Matched inputs will can be directly used by code line generation
# key is ext. succ node onnx name, value is local opt_var
self.external_successor_local_returns_map = OrderedDict()
# Define outputs manager, note this will be assigned later by Generator.
self.outputs_manager = None
self._global_context = GlobalContext()
# Define a dict to store the reference for quick searching
self.rapid_reference = dict()
# new vars for matcher
self.inputs_register = OrderedDict() # reg by sub
self.outputs_register = OrderedDict() # reg by sub
self.internal_outputs_collection = dict() # reg by sub
# new vars for shared weights
self.shared_weights_collection = dict() # reg by sub
self.shared_weights_counter = 0 # updated by sub
if init_as_parent and (parent_base is not None):
self.reset_as_parent_passed_in(parent_base)
else:
# start initialization
if not self.initialized:
self._init_module(nd_struct_list)
else:
self._update_module(nd_struct_list)
# assign this module reference to node
for (_, nd_struct) in nd_struct_list:
nd_struct.parent_module_struct = self