From 1948155fb22d9a2806483cef28beaf45bfe59bcc Mon Sep 17 00:00:00 2001
From: Mahyar Samani <msamani@ucdavis.edu>
Date: Wed, 30 Oct 2024 12:37:22 -0700
Subject: [PATCH] stdlib: AbstractMemorySystem.get_mem_interfaces.

This change adds a new method to AbstractMemorySystem to allow
getting its objects of the class MemInterface. This is useful
when certain other classes require a list of MemInterface objects
to create physical memory. In addition, ChanneledMemory and
HighBandwidthMemory implement this function.
---
 .../gem5/components/memory/abstract_memory_system.py      | 7 +++++++
 src/python/gem5/components/memory/hbm.py                  | 8 ++++++++
 src/python/gem5/components/memory/memory.py               | 4 ++++
 3 files changed, 19 insertions(+)

diff --git a/src/python/gem5/components/memory/abstract_memory_system.py b/src/python/gem5/components/memory/abstract_memory_system.py
index 6d24e724b6..67d9187d61 100644
--- a/src/python/gem5/components/memory/abstract_memory_system.py
+++ b/src/python/gem5/components/memory/abstract_memory_system.py
@@ -37,6 +37,7 @@ from typing import (
 from m5.objects import (
     AddrRange,
     MemCtrl,
+    MemInterface,
     Port,
     Root,
     SubSystem,
@@ -79,6 +80,12 @@ class AbstractMemorySystem(SubSystem):
         """Get all of the memory controllers in this memory system."""
         raise NotImplementedError
 
+    @abstractmethod
+    def get_mem_interfaces(self) -> List[MemInterface]:
+        """Get all memory interfaces in this memory system.
+        Useful when creating physical memory objects."""
+        raise NotImplementedError
+
     @abstractmethod
     def get_size(self) -> int:
         """Returns the total size of the memory system."""
diff --git a/src/python/gem5/components/memory/hbm.py b/src/python/gem5/components/memory/hbm.py
index 1e1c83e2bd..4774634d34 100644
--- a/src/python/gem5/components/memory/hbm.py
+++ b/src/python/gem5/components/memory/hbm.py
@@ -29,6 +29,7 @@
 
 from math import log
 from typing import (
+    List,
     Optional,
     Sequence,
     Tuple,
@@ -40,6 +41,7 @@ from m5.objects import (
     AddrRange,
     DRAMInterface,
     HBMCtrl,
+    MemInterface,
     Port,
 )
 
@@ -172,6 +174,12 @@ class HighBandwidthMemory(ChanneledMemory):
             (addr_ranges[i], ctrl.port) for i, ctrl in enumerate(self.mem_ctrl)
         ]
 
+    @overrides(ChanneledMemory)
+    def get_mem_interfaces(self) -> List[MemInterface]:
+        return [ctrl.dram for ctrl in self.get_memory_controllers()] + [
+            ctrl.dram_2 for ctrl in self.get_memory_controllers()
+        ]
+
 
 def HBM2Stack(
     size: Optional[str] = "4GiB",
diff --git a/src/python/gem5/components/memory/memory.py b/src/python/gem5/components/memory/memory.py
index 2d922c8e83..9b9c501bc8 100644
--- a/src/python/gem5/components/memory/memory.py
+++ b/src/python/gem5/components/memory/memory.py
@@ -185,6 +185,10 @@ class ChanneledMemory(AbstractMemorySystem):
     def get_memory_controllers(self) -> List[MemCtrl]:
         return [ctrl for ctrl in self.mem_ctrl]
 
+    @overrides(AbstractMemorySystem)
+    def get_mem_interfaces(self) -> List[DRAMInterface]:
+        return self._dram
+
     @overrides(AbstractMemorySystem)
     def get_size(self) -> int:
         return self._size