diff --git a/aidge_core/export_utils/export_registry.py b/aidge_core/export_utils/export_registry.py index eabc6eb6b656b0c0b0c8381b665d6a6d8f3805ff..fd24008a6de6c58c1e78f088e086817e2a769373 100644 --- a/aidge_core/export_utils/export_registry.py +++ b/aidge_core/export_utils/export_registry.py @@ -28,6 +28,10 @@ class ExportLib(aidge_core.OperatorImpl): # key: Path where static file is # Value: Path where to copy the file relative to the export root static_files: Dict[str, str] = {} + # Custom main generation jinja file + main_jinja_path = None + # Main memory section + memory_section = None # PRIVATE # Registry of exportNode, class level dictionary, shared across all ExportLib _cls_export_node_registry = {} diff --git a/aidge_core/export_utils/scheduler_export.py b/aidge_core/export_utils/scheduler_export.py index 6829832feda7af7b2c808df4cd430fc77b37b3cb..df0b4a385327e4bdccd6fe4de46043d151658dbd 100644 --- a/aidge_core/export_utils/scheduler_export.py +++ b/aidge_core/export_utils/scheduler_export.py @@ -6,7 +6,7 @@ from aidge_core.export_utils import ExportLib, generate_file, copy_file from typing import List, Tuple -def scheduler_export(scheduler, export_folder_path: str, export_lib: ExportLib = None, memory_manager=None, memory_manager_args=None) -> None: +def scheduler_export(scheduler, export_folder_path: str, export_lib: ExportLib = None, memory_manager=None, memory_manager_args=None, labels=False) -> None: graphview = scheduler.graph_view() export_folder = Path().absolute() / export_folder_path @@ -114,6 +114,7 @@ def scheduler_export(scheduler, export_folder_path: str, export_lib: ExportLib = headers=set(list_configs), actions=list_actions, mem_ctype=inputs_dtype[0], # Legacy behavior ... + mem_section=export_lib.mem_section, peak_mem=peak_mem, inputs_name=inputs_name, inputs_dtype=inputs_dtype, @@ -136,14 +137,20 @@ def scheduler_export(scheduler, export_folder_path: str, export_lib: ExportLib = if len(outputs_name) != len(outputs_dtype) or len(outputs_name) != len(outputs_size): raise RuntimeError("FATAL: Output args list does not have the same length this is an internal bug.") + if export_lib is not None and export_lib.main_jinja_path is not None: + main_jinja_path = export_lib.main_jinja_path + else : + main_jinja_path = str(ROOT / "templates" / "main.jinja") + generate_file( str(export_folder / "main.cpp"), - str(ROOT / "templates" / "main.jinja"), + main_jinja_path, func_name=func_name, inputs_name=inputs_name, outputs_name=outputs_name, outputs_dtype=outputs_dtype, - outputs_size=outputs_size + outputs_size=outputs_size, + labels=labels ) if export_lib is not None: diff --git a/aidge_core/export_utils/templates/forward.jinja b/aidge_core/export_utils/templates/forward.jinja index aec5867b132e0ece9a8a39d9cf5daadeb25ea24a..fde4b2a1392c4ada353af06246951e26c6236df6 100644 --- a/aidge_core/export_utils/templates/forward.jinja +++ b/aidge_core/export_utils/templates/forward.jinja @@ -14,7 +14,11 @@ {%- endfor %} // Memory block +{%- if mem_section == None %} static {{mem_ctype}} mem[{{peak_mem}}]; +{%- else %} +static {{mem_ctype}} mem[{{peak_mem}}] __attribute__((section("{{ mem_section }}"))); +{%- endif %} {# Forward function #} {#- Support multiple inputs with different datatypes and multiple outputs with different datatypes -#}