From 05eb21f0ad47cf47be2a8b87cc8a6908a01d1f7c Mon Sep 17 00:00:00 2001 From: Olivier BICHLER <olivier.bichler@cea.fr> Date: Thu, 26 Jun 2025 15:54:58 +0200 Subject: [PATCH] More general memory management --- aidge_export_cpp/export.py | 13 ++++--- aidge_export_cpp/operators/Producer.py | 34 +++++++++++++------ .../templates/configuration/_meminfo.jinja | 30 +++++++++++----- .../configuration/producer_config.jinja | 8 +++++ .../kernel_forward/_mem_offset.jinja | 2 +- aidge_export_cpp/unit_tests/test_export.py | 1 + 6 files changed, 64 insertions(+), 24 deletions(-) create mode 100644 aidge_export_cpp/templates/configuration/producer_config.jinja diff --git a/aidge_export_cpp/export.py b/aidge_export_cpp/export.py index adeab0d..7effde9 100644 --- a/aidge_export_cpp/export.py +++ b/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) diff --git a/aidge_export_cpp/operators/Producer.py b/aidge_export_cpp/operators/Producer.py index 59979fe..bfae74f 100644 --- a/aidge_export_cpp/operators/Producer.py +++ b/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 [] diff --git a/aidge_export_cpp/templates/configuration/_meminfo.jinja b/aidge_export_cpp/templates/configuration/_meminfo.jinja index 15d1542..ad6e76f 100644 --- a/aidge_export_cpp/templates/configuration/_meminfo.jinja +++ b/aidge_export_cpp/templates/configuration/_meminfo.jinja @@ -1,11 +1,25 @@ // 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 %} diff --git a/aidge_export_cpp/templates/configuration/producer_config.jinja b/aidge_export_cpp/templates/configuration/producer_config.jinja new file mode 100644 index 0000000..ad0660f --- /dev/null +++ b/aidge_export_cpp/templates/configuration/producer_config.jinja @@ -0,0 +1,8 @@ +{#- 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 diff --git a/aidge_export_cpp/templates/kernel_forward/_mem_offset.jinja b/aidge_export_cpp/templates/kernel_forward/_mem_offset.jinja index a32b16b..e385318 100644 --- a/aidge_export_cpp/templates/kernel_forward/_mem_offset.jinja +++ b/aidge_export_cpp/templates/kernel_forward/_mem_offset.jinja @@ -1,3 +1,3 @@ {%- 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 %} diff --git a/aidge_export_cpp/unit_tests/test_export.py b/aidge_export_cpp/unit_tests/test_export.py index f0fe59a..871e701 100644 --- a/aidge_export_cpp/unit_tests/test_export.py +++ b/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() -- GitLab