More general memory management

aidge_export_cpp/export.py

@@ -19,7 +19,9 @@ def export(export_folder_name: str,
            inputs_tensor: aidge_core.Tensor = None,
            labels: aidge_core.Tensor = None,
            dev_mode: bool = False,
-           aidge_cmp: bool = False):
+           aidge_cmp: bool = False,
+           memory_manager = generate_optimized_memory_info,
+           memory_manager_args = {}):
     
     """ Export an aidge_core.Scheduler to C++ code
     
@@ -68,13 +70,16 @@ def export(export_folder_name: str,
     - Copy all needed kernels
     """
 
+    if "stats_folder" not in memory_manager_args:
+        memory_manager_args["stats_folder"] = f"{export_folder_name}/stats"
+
     scheduler_export(scheduler,
                      export_folder_name,
                      ExportLibCpp,
-                     memory_manager=generate_optimized_memory_info,
-                     memory_manager_args={
-                         "stats_folder": f"{export_folder_name}/stats"},
+                     memory_manager=memory_manager,
+                     memory_manager_args=memory_manager_args,
                      dev_mode=dev_mode)
+    graphview.save(f"{export_folder_name}/graph")
     
     # Generate main file
     generate_main_cpp(export_folder_name, graphview, labels=labels, inputs_tensor=inputs_tensor)

aidge_export_cpp/operators/Producer.py

@@ -2,7 +2,7 @@ import os
 from pathlib import Path
 import numpy as np
 import aidge_core
-from aidge_core.export_utils import ExportNode, generate_file
+from aidge_core.export_utils import ExportNodeCpp, generate_file
 from aidge_export_cpp import ROOT
 from aidge_export_cpp import ExportLibCpp
 
@@ -43,22 +43,34 @@ def export_params(name: str,
     )
 
 @ExportLibCpp.register("Producer", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.any)))
-class ProducerCPP(ExportNode):
+class ProducerCPP(ExportNodeCpp):
     def __init__(self, node, mem_info):
         super().__init__(node, mem_info)
         self.values = np.array(self.operator.get_output(0))
         self.ignore = node.attributes().has_attr("ignore")
 
     def export(self, export_folder: Path):
-        if not self.ignore :
-            header_path = f"include/parameters/{self.attributes['name']}.h"
-            export_params(
-                self.attributes['out_name'][0],
-                self.values.reshape(-1),
-                str(export_folder / header_path))
-            return [header_path]
-        return []
+        if self.ignore:
+            return []
+
+        path_to_definition = f"{self.config_path}/{self.attributes['name']}.{self.config_extension}"
+
+        try:
+            aidge_core.export_utils.code_generation.generate_file(
+                str(export_folder / path_to_definition),
+                str(ROOT / "templates" / "configuration" / "producer_config.jinja"),
+                **self.attributes
+            )
+        except Exception as e:
+            raise RuntimeError(f"Error when creating config file for {self.node.name()}[{self.node.type()}].") from e
+
+        header_path = f"include/parameters/{self.attributes['name']}.h"
+        export_params(
+            self.attributes['out_name'][0],
+            self.values.reshape(-1),
+            str(export_folder / header_path))
+        return [path_to_definition, header_path]
 
     def forward(self):
         # A Producer does nothing during forward
-        return []
\ No newline at end of file
+        return []

aidge_export_cpp/templates/configuration/_meminfo.jinja

 // MEMINFO CONF
+{% for inidx in range(nb_in) -%}
+{# Specify a default memory layout for standalone input tensors -#}
+{% if not in_node[inidx] %}
+#define {{ in_name[inidx]|upper }}_MEM_SIZE {{ in_size[inidx] }}
+#define {{ in_name[inidx]|upper }}_MEM_OFFSET 0
+#define {{ in_name[inidx]|upper }}_MEM_STRIDE {{ in_size[inidx] }}
+#define {{ in_name[inidx]|upper }}_MEM_LENGTH 1
+#define {{ in_name[inidx]|upper }}_MEM_CONT_SIZE {{ in_size[inidx] }}
+#define {{ in_name[inidx]|upper }}_MEM_CONT_OFFSET 0
+#define {{ in_name[inidx]|upper }}_MEM_WRAP_OFFSET 0
+#define {{ in_name[inidx]|upper }}_MEM_WRAP_SIZE 0
+{% endif %}
+{% endfor %}
+
 {% for outidx in range(nb_out) -%}
-#define {{ out_name[outidx]|upper }}_SIZE {{ mem_info_size[outidx]}}
-#define {{ out_name[outidx]|upper }}_OFFSET {{ mem_info_offset[outidx]}}
-#define {{ out_name[outidx]|upper }}_STRIDE {{ mem_info_stride[outidx]}}
-#define {{ out_name[outidx]|upper }}_LENGTH {{ mem_info_length[outidx]}}
-#define {{ out_name[outidx]|upper }}_CONT_SIZE {{ mem_info_cont_size[outidx]}}
-#define {{ out_name[outidx]|upper }}_CONT_OFFSET {{ mem_info_cont_offset[outidx]}}
-#define {{ out_name[outidx]|upper }}_WRAP_OFFSET {{ mem_info_wrap_offset[outidx]}}
-#define {{ out_name[outidx]|upper }}_WRAP_SIZE {{ mem_info_wrap_size[outidx]}}
+#define {{ out_name[outidx]|upper }}_MEM_SIZE {{ mem_info_size[outidx]}}
+#define {{ out_name[outidx]|upper }}_MEM_OFFSET {{ mem_info_offset[outidx]}}
+#define {{ out_name[outidx]|upper }}_MEM_STRIDE {{ mem_info_stride[outidx]}}
+#define {{ out_name[outidx]|upper }}_MEM_LENGTH {{ mem_info_length[outidx]}}
+#define {{ out_name[outidx]|upper }}_MEM_CONT_SIZE {{ mem_info_cont_size[outidx]}}
+#define {{ out_name[outidx]|upper }}_MEM_CONT_OFFSET {{ mem_info_cont_offset[outidx]}}
+#define {{ out_name[outidx]|upper }}_MEM_WRAP_OFFSET {{ mem_info_wrap_offset[outidx]}}
+#define {{ out_name[outidx]|upper }}_MEM_WRAP_SIZE {{ mem_info_wrap_size[outidx]}}
 {% endfor %}

aidge_export_cpp/templates/configuration/producer_config.jinja

+{#- For name header -#}
+#ifndef {{ name|upper }}_LAYER_H
+#define {{ name|upper }}_LAYER_H
+
+{# For layer configuration -#}
+{% include "./_meminfo.jinja" %}
+
+#endif /* {{ name|upper }}_LAYER_H */
\ No newline at end of file

aidge_export_cpp/templates/kernel_forward/_mem_offset.jinja

 {%- for outidx in range(nb_out) %}
-{{out_cdtype[outidx]}}* {{out_name[outidx]}} = ({{out_cdtype[outidx]}}*) (mem + {{out_name[outidx]|upper}}_OFFSET);
+{{out_cdtype[outidx]}}* {{out_name[outidx]}} = ({{out_cdtype[outidx]}}*) (mem + {{out_name[outidx]|upper}}_MEM_OFFSET);
 {%- endfor %}

aidge_export_cpp/unit_tests/test_export.py

@@ -154,6 +154,7 @@ class test_operator_export(unittest.TestCase):
         aidge_core.adapt_to_backend(graph_view)
         graph_view.forward_dims(dims=in_dims)
         graph_view.save(export_folder + "/graph")
+    
         scheduler = aidge_core.SequentialScheduler(graph_view)
         scheduler.generate_scheduling()