diff --git a/hpvm/projects/onnx_frontend/frontend/codegen_hpvm.py b/hpvm/projects/onnx_frontend/frontend/codegen_hpvm.py index fbdf0aea083ed34ae0f89b51426747a97272380b..68b8dedb6938124d161cbab8af2bbe176a8c0114 100644 --- a/hpvm/projects/onnx_frontend/frontend/codegen_hpvm.py +++ b/hpvm/projects/onnx_frontend/frontend/codegen_hpvm.py @@ -50,10 +50,10 @@ class CodeGen: # Aux functions ################################################ - def _allocate_varname(self) -> str: - varname = f"var_{self.var_count}" + def _inc_var_count(self) -> int: + var_n = self.var_count self.var_count += 1 - return varname + return var_n @classmethod def emit_weights(cls, weights: List[WeightTensor]) -> List[dict]: @@ -103,11 +103,11 @@ class HpvmCodeGen(CodeGen): assert len(inputs) == 1 self.variables[node] = self.variables[inputs[0]] continue - varname = self._allocate_varname() - self.variables[node] = varname, False # not root-node arg + var_idx = self._inc_var_count() + self.variables[node] = var_idx, False # not root-node arg node_envs.append( { - "name": varname, + "idx": var_idx, "input_size": len(inputs), "edges": self._emit_hpvm_node_edges(inputs), "call_name": func_name, diff --git a/hpvm/projects/onnx_frontend/frontend/codegen_tensor.py b/hpvm/projects/onnx_frontend/frontend/codegen_tensor.py index 44058b0ac454cca150c19fc568fde4b9d726d199..3c511803274534f8a88c9d62d049df918ec67010 100644 --- a/hpvm/projects/onnx_frontend/frontend/codegen_tensor.py +++ b/hpvm/projects/onnx_frontend/frontend/codegen_tensor.py @@ -36,7 +36,7 @@ class TensorCodeGen(CodeGen): assert len(inputs) == 1 self.variables[node] = self.variables[inputs[0]] continue - varname = self._allocate_varname() + varname = f"var_{self._inc_var_count()}" self.variables[node] = varname input_args = [self.variables[n] for n in inputs] + extra_args graph_code.append( diff --git a/hpvm/projects/onnx_frontend/frontend/template_hpvm.cpp.in b/hpvm/projects/onnx_frontend/frontend/template_hpvm.cpp.in index 7f6785c41a623012787b68c3d8a438024979ad14..5634a3b969cec31e61b27bf3a1e91c62bff7214a 100644 --- a/hpvm/projects/onnx_frontend/frontend/template_hpvm.cpp.in +++ b/hpvm/projects/onnx_frontend/frontend/template_hpvm.cpp.in @@ -4,7 +4,7 @@ #include <tensorUtils.h> {% for node in nodes %} -void {{node.name}}_node( +void var_{{node.idx}}_node( {%- for n in range(node.input_size) -%} void *t{{n}}, size_t bytes_t{{n}}{{", " if not loop.last}} {%- endfor %}) { @@ -12,6 +12,7 @@ void *t{{n}}, size_t bytes_t{{n}}{{", " if not loop.last}} __hpvm__attributes({{node.input_size}}, {% for n in range(node.input_size) -%} t{{n}}{{", " if not loop.last}} {%- endfor %}, 0); + __hpvm__node_id({{node.idx + 1}}); void *r = {{node.call_name}}({% for n in range(node.input_size) -%} t{{n}}{{", " if not loop.last}} {%- endfor %}{{", " if node.call_args}}{{node.call_args|join(", ")}}); @@ -29,14 +30,14 @@ void *{{n}}, size_t {{n}}_bytes{{", " if not loop.last}} {%- endfor %}, 0); {% for node in nodes %} - void* {{node.name}} = __hpvm__createNodeND(0, {{node.name}}_node); + void* var_{{node.idx}} = __hpvm__createNodeND(0, var_{{node.idx}}_node); {% for edge in node.edges %} {% if edge.is_bindin %} - __hpvm__bindIn({{node.name}}, {{edge.input_idx * 2 + 1}}, {{edge.edge_idx * 2}}, 0); - __hpvm__bindIn({{node.name}}, {{edge.input_idx * 2 + 2}}, {{edge.edge_idx * 2 + 1}}, 0); + __hpvm__bindIn(var_{{node.idx}}, {{edge.input_idx * 2 + 1}}, {{edge.edge_idx * 2}}, 0); + __hpvm__bindIn(var_{{node.idx}}, {{edge.input_idx * 2 + 2}}, {{edge.edge_idx * 2 + 1}}, 0); {% else %} - __hpvm__edge({{edge.input_node}}, {{node.name}}, 1, 0, {{edge.edge_idx * 2}}, 0); - __hpvm__edge({{edge.input_node}}, {{node.name}}, 1, 1, {{edge.edge_idx * 2 + 1}}, 0); + __hpvm__edge({{edge.input_node}}, var_{{node.idx}}, 1, 0, {{edge.edge_idx * 2}}, 0); + __hpvm__edge({{edge.input_node}}, var_{{node.idx}}, 1, 1, {{edge.edge_idx * 2 + 1}}, 0); {% endif %} {% endfor %}