Rewrite memory manager log_info in Python.

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()