From b963bf063c114c66392f62b404a33df94b9a8ba6 Mon Sep 17 00:00:00 2001
From: Axel Farrugia <axel.farrugia@cea.fr>
Date: Tue, 24 Sep 2024 11:46:35 +0200
Subject: [PATCH] Memory section specification for forward generation
---
aidge_core/export_utils/export_registry.py | 2 ++
aidge_core/export_utils/scheduler_export.py | 1 +
aidge_core/export_utils/templates/forward.jinja | 4 ++++
3 files changed, 7 insertions(+)
diff --git a/aidge_core/export_utils/export_registry.py b/aidge_core/export_utils/export_registry.py
index eabc6eb6b..51711b751 100644
--- a/aidge_core/export_utils/export_registry.py
+++ b/aidge_core/export_utils/export_registry.py
@@ -28,6 +28,8 @@ 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] = {}
+ # 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 6829832fe..13b5cb9ce 100644
--- a/aidge_core/export_utils/scheduler_export.py
+++ b/aidge_core/export_utils/scheduler_export.py
@@ -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,
diff --git a/aidge_core/export_utils/templates/forward.jinja b/aidge_core/export_utils/templates/forward.jinja
index aec5867b1..fde4b2a13 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 -#}
--
GitLab