diff --git a/aidge_core/mem_info.py b/aidge_core/mem_info.py
index c7ca85bbd73bd205850b19616e53fda210749a80..53ff73cef2d8d2f17428b5fc3f9917e94c7dce8a 100644
--- a/aidge_core/mem_info.py
+++ b/aidge_core/mem_info.py
@@ -5,6 +5,9 @@ from pathlib import Path
import aidge_core
from typing import Tuple, List
+import matplotlib.pyplot as plt
+import aidge_core.mem_info
+import numpy as np
# Default memory management, which can be used for development
def compute_default_mem_info(scheduler: aidge_core.Scheduler) -> Tuple[int, List]:
@@ -41,20 +44,78 @@ def compute_default_mem_info(scheduler: aidge_core.Scheduler) -> Tuple[int, List
mem_info[node] = [] # No meminfo for producer
return mem_size, mem_info
+# TODO remove
+# def _gnuplot_installed():
+# try:
+# # Run gnuplot with the --version flag and capture the output
+# subprocess.run(["gnuplot", "--version"])
+# return True
+# except FileNotFoundError:
+# aidge_core.Log.warn("Gnuplot is not installed.")
+# return False
+# except subprocess.CalledProcessError:
+# aidge_core.Log.warn("Gnuplot command found but failed to run.")
+# return False
+
+def log_meminfo(mem_manager:aidge_core.MemoryManager, path: Path, diplay_names:bool):
+ """Generate a graph representing the memory allocation of each ouputs.
+
+ Block with the smae color correspond to the same memory plane.
+
+ :param mem_manager: Memory manager to log
+ :type mem_manager: aidge_core.memory_manager
+ :param path: Path where to save the figure
+ :type path: Path
+ :param diplay_names: If True Node names are diplayed alongside their block
+ :type diplay_names: bool
+ """
-def _gnuplot_installed():
- try:
- # Run gnuplot with the --version flag and capture the output
- subprocess.run(["gnuplot", "--version"])
- return True
- except FileNotFoundError:
- aidge_core.Log.warn("Gnuplot is not installed.")
- return False
- except subprocess.CalledProcessError:
- aidge_core.Log.warn("Gnuplot command found but failed to run.")
- return False
-
-def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder: Path = None, wrapping: bool = False) -> Tuple[int, List[dict]]:
+ max_lifetime = mem_manager.get_max_lifetime()
+
+ # peak_usage in kwords
+ peak_usage = mem_manager.get_peak_usage() / 1024
+
+ # Set figure size 1920x1080 px
+ plt.figure(figsize=(19.20, 10.80))
+ # Same color for each planes
+ colors = plt.cm.viridis(np.linspace(0, 1, len(mem_manager.get_planes()) + 1))
+ color_id = 1
+ for node, planes in mem_manager.get_planes().items():
+ for plane in planes:
+ cont_offset = plane.get_contiguous_offset()
+ cont_size = plane.get_contiguous_size()
+ allocated = plane.mem_space.allocated
+ released = plane.mem_space.released
+ is_released = released >= 0 and not plane.mem_space.dependencies
+ x_start = allocated
+ y_start = cont_offset / 1024.0
+ y_end = (cont_offset + cont_size) / 1024.0
+ x_end = max_lifetime if not is_released else released
+
+ plt.fill_betweenx(
+ [y_start, y_end],
+ x_start,
+ x_end + 1,
+ color=colors[color_id % len(colors)]
+ )
+
+ if diplay_names:
+ # Rotation for lisibility!
+ plt.text(x_end,y_end, node.name(), rotation=45)
+ color_id += 1
+
+ plt.axhline(y=peak_usage, color='red', linestyle='--')
+ plt.text(0, peak_usage, f'Peak usage = {peak_usage} KWords', color='red')
+ plt.xlabel("Time")
+ plt.ylabel("Memory usage (KWords)")
+ plt.title("Memory Usage Over Time")
+ plt.grid(True)
+ plt.savefig(path)
+ plt.close()
+ aidge_core.Log.notice(f"Generated memory management info at: {path}")
+
+
+def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder: Path = None, wrapping: bool = False, display_names: bool=True) -> Tuple[int, List[dict]]:
"""Generates optimized memory information for a computation graph managed by a scheduler.
This function analyzes the memory usage of a computation graph, determining the memory peak
@@ -70,6 +131,8 @@ def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder
:param wrapping: Boolean flag to enable or disable wrap-around buffer optimization.
Defaults to `False`.
:type wrapping: bool, optional
+ :param diplay_names: If True Node names are diplayed in the memory plot alongside their block, defaults to False
+ :type diplay_names: bool, optional
:return: A tuple containing the peak memory size and a list of memory information for each
scheduled node. The memory information for each node includes details such as size,
offset, stride, length, count, and optional wrap-around details.
@@ -88,18 +151,21 @@ def generate_optimized_memory_info(scheduler: aidge_core.Scheduler, stats_folder
nodes_at_input = [n[0] for n in scheduler.graph_view().inputs()]
if stats_folder is not None:
- if _gnuplot_installed():
- # Use gnuplot to generate the log
- os.makedirs(str(Path(stats_folder) / "graph"), exist_ok=True)
- mem_manager.log("memory_info")
- os.chmod("memory_info_plot.gnu", 0o777)
- os.system("./memory_info_plot.gnu")
- shutil.move("memory_info", str(Path(stats_folder) / "graph" / "memory_info"))
- shutil.move("memory_info_plot.png", str(
- Path(stats_folder) / "graph" / "memory_info_plot.png"))
- os.remove("memory_info_plot.gnu")
- else:
- aidge_core.Log.warn("Warning: gnuplot is not installed, could not generate stat folder.")
+ log_meminfo(mem_manager, Path(stats_folder) / "memory_info.png", display_names)
+ # TODO remove
+ # if _gnuplot_installed():
+ # # Use gnuplot to generate the log
+ # os.makedirs(str(Path(stats_folder) / "graph"), exist_ok=True)
+ # mem_manager.log("memory_info")
+ # os.chmod("memory_info_plot.gnu", 0o777)
+ # os.system("./memory_info_plot.gnu")
+ # shutil.move("memory_info", str(Path(stats_folder) / "graph" / "memory_info"))
+ # shutil.move("memory_info_plot.png", str(
+ # Path(stats_folder) / "graph" / "memory_info_plot.png"))
+ # os.remove("memory_info_plot.gnu")
+ # else:
+ # aidge_core.Log.warn("Warning: gnuplot is not installed, could not generate stat folder.")
+
# In the export, we currently use an unified memory buffer whose size
# is determined by the memory peak usage
mem_size = mem_manager.get_peak_usage()