Add tag context manager to AXI master to reuse per-ID processing components

cocotbext/axi/axi_master.py

@@ -49,6 +49,92 @@ AxiReadRespCmd = namedtuple("AxiReadRespCmd", ["address", "length", "size", "cyc
 AxiReadResp = namedtuple("AxiReadResp", ["address", "data", "resp", "user"])
 
 
+class TagContext:
+    def __init__(self, manager):
+        self.current_tag = 0
+        self._cmd_queue = Queue()
+        self._current_cmd = None
+        self._resp_queue = Queue()
+        self._cr = None
+        self._manager = manager
+
+    async def get_resp(self):
+        return await self._resp_queue.get()
+
+    def get_resp_nowait(self):
+        return self._resp_queue.get_nowait()
+
+    def _start(self):
+        if self._cr is None:
+            self._cr = cocotb.fork(self._process_queue())
+
+    def _flush(self):
+        flushed_cmds = []
+        if self._cr is not None:
+            self._cr.kill()
+            self._cr = None
+        self._manager._set_idle(self)
+        if self._current_cmd is not None:
+            flushed_cmds.append(self._current_cmd)
+            self._current_cmd = None
+        while not self._cmd_queue.empty():
+            flushed_cmds.append(self._cmd_queue.get_nowait())
+        while not self._resp_queue.empty():
+            self._resp_queue.get_nowait()
+        return flushed_cmds
+
+    async def _process_queue(self):
+        while True:
+            cmd = await self._cmd_queue.get()
+            self._current_cmd = cmd
+            await self._manager._process(self, cmd)
+            self._current_cmd = None
+
+            if self._cmd_queue.empty() and self._resp_queue.empty():
+                self._manager._set_idle(self)
+
+
+class TagContextManager:
+    def __init__(self, process):
+        self._context_list = []
+        self._context_idle_list = []
+        self._context_mapping = {}
+        self._process = process
+
+    def _get_context(self, tag):
+        if tag in self._context_mapping:
+            return self._context_mapping[tag]
+        elif self._context_idle_list:
+            context = self._context_idle_list.pop()
+        else:
+            context = TagContext(self)
+            self._context_list.append(context)
+        context._start()
+        context.current_tag = tag
+        self._context_mapping[tag] = context
+        return context
+
+    def start_cmd(self, tag, cmd):
+        context = self._get_context(tag)
+        context._cmd_queue.put_nowait(cmd)
+
+    def put_resp(self, tag, resp):
+        context = self._get_context(tag)
+        context._resp_queue.put_nowait(resp)
+
+    def _set_idle(self, context):
+        if context.current_tag in self._context_mapping:
+            del self._context_mapping[context.current_tag]
+            self._context_idle_list.append(context)
+        context.current_tag = None
+
+    def flush(self):
+        flushed_cmds = []
+        for c in self._context_list:
+            flushed_cmds.extend(c._flush())
+        return flushed_cmds
+
+
 class AxiMasterWrite(Reset):
     def __init__(self, bus, clock, reset=None, reset_active_level=True, max_burst_len=256):
         self.log = logging.getLogger(f"cocotb.{bus.aw._entity._name}.{bus.aw._name}")
@@ -72,9 +158,7 @@ class AxiMasterWrite(Reset):
         self.cur_id = 0
         self.active_id = Counter()
 
-        self.int_write_resp_command_queue = [Queue() for k in range(self.id_count)]
+        self.tag_context_manager = TagContextManager(self._process_write_resp_id)
-        self.current_write_resp_command = [None for k in range(self.id_count)]
-        self.int_write_resp_queue_list = [Queue() for k in range(self.id_count)]
 
         self.in_flight_operations = 0
         self._idle = Event()
@@ -104,7 +188,6 @@ class AxiMasterWrite(Reset):
 
         self._process_write_cr = None
         self._process_write_resp_cr = None
-        self._process_write_resp_id_cr = None
 
         self._init_reset(reset, reset_active_level)
 
@@ -207,10 +290,6 @@ class AxiMasterWrite(Reset):
             if self._process_write_resp_cr is not None:
                 self._process_write_resp_cr.kill()
                 self._process_write_resp_cr = None
-            if self._process_write_resp_id_cr is not None:
-                for cr in self._process_write_resp_id_cr:
-                    cr.kill()
-                self._process_write_resp_id_cr = None
 
             self.aw_channel.clear()
             self.w_channel.clear()
@@ -230,21 +309,9 @@ class AxiMasterWrite(Reset):
                 self.current_write_command = None
                 flush_cmd(cmd)
 
-            for q in self.int_write_resp_command_queue:
+            for cmd in self.tag_context_manager.flush():
-                while not q.empty():
-                    cmd = q.get_nowait()
                 flush_cmd(cmd)
 
-            for k in range(len(self.current_write_resp_command)):
-                if self.current_write_resp_command[k]:
-                    cmd = self.current_write_resp_command[k]
-                    self.current_write_resp_command[k] = None
-                    flush_cmd(cmd)
-
-            for q in self.int_write_resp_queue_list:
-                while not q.empty():
-                    q.get_nowait()
-
             self.cur_id = 0
             self.active_id = Counter()
 
@@ -256,8 +323,6 @@ class AxiMasterWrite(Reset):
                 self._process_write_cr = cocotb.fork(self._process_write())
             if self._process_write_resp_cr is None:
                 self._process_write_resp_cr = cocotb.fork(self._process_write_resp())
-            if self._process_write_resp_id_cr is None:
-                self._process_write_resp_id_cr = [cocotb.fork(self._process_write_resp_id(i)) for i in range(self.id_count)]
 
     async def _process_write(self):
         while True:
@@ -361,7 +426,7 @@ class AxiMasterWrite(Reset):
                 cycle_offset = (cycle_offset + num_bytes) % self.byte_width
 
             resp_cmd = AxiWriteRespCmd(cmd.address, len(cmd.data), cmd.size, cycles, cmd.prot, burst_list, cmd.event)
-            await self.int_write_resp_command_queue[awid].put(resp_cmd)
+            self.tag_context_manager.start_cmd(awid, resp_cmd)
 
             self.current_write_command = None
 
@@ -374,18 +439,16 @@ class AxiMasterWrite(Reset):
             if self.active_id[bid] <= 0:
                 raise Exception(f"Unexpected burst ID {bid}")
 
-            await self.int_write_resp_queue_list[bid].put(b)
+            self.tag_context_manager.put_resp(bid, b)
 
-    async def _process_write_resp_id(self, bid):
+    async def _process_write_resp_id(self, context, cmd):
-        while True:
+        bid = context.current_tag
-            cmd = await self.int_write_resp_command_queue[bid].get()
-            self.current_write_resp_command[bid] = cmd
 
         resp = AxiResp.OKAY
         user = []
 
         for burst_length in cmd.burst_list:
-                b = await self.int_write_resp_queue_list[bid].get()
+            b = await context.get_resp()
 
             burst_resp = AxiResp(b.bresp)
             burst_user = int(b.buser)
@@ -410,8 +473,6 @@ class AxiMasterWrite(Reset):
 
         cmd.event.set(write_resp)
 
-            self.current_write_resp_command[bid] = None
-
         self.in_flight_operations -= 1
 
         if self.in_flight_operations == 0:
@@ -439,9 +500,7 @@ class AxiMasterRead(Reset):
         self.cur_id = 0
         self.active_id = Counter()
 
-        self.int_read_resp_command_queue = [Queue() for k in range(self.id_count)]
+        self.tag_context_manager = TagContextManager(self._process_read_resp_id)
-        self.current_read_resp_command = [None for k in range(self.id_count)]
-        self.int_read_resp_queue_list = [Queue() for k in range(self.id_count)]
 
         self.in_flight_operations = 0
         self._idle = Event()
@@ -469,7 +528,6 @@ class AxiMasterRead(Reset):
 
         self._process_read_cr = None
         self._process_read_resp_cr = None
-        self._process_read_resp_id_cr = None
 
         self._init_reset(reset, reset_active_level)
 
@@ -567,10 +625,6 @@ class AxiMasterRead(Reset):
             if self._process_read_resp_cr is not None:
                 self._process_read_resp_cr.kill()
                 self._process_read_resp_cr = None
-            if self._process_read_resp_id_cr is not None:
-                for cr in self._process_read_resp_id_cr:
-                    cr.kill()
-                self._process_read_resp_id_cr = None
 
             self.ar_channel.clear()
             self.r_channel.clear()
@@ -589,21 +643,9 @@ class AxiMasterRead(Reset):
                 self.current_read_command = None
                 flush_cmd(cmd)
 
-            for q in self.int_read_resp_command_queue:
+            for cmd in self.tag_context_manager.flush():
-                while not q.empty():
-                    cmd = q.get_nowait()
                 flush_cmd(cmd)
 
-            for k in range(len(self.current_read_resp_command)):
-                if self.current_read_resp_command[k]:
-                    cmd = self.current_read_resp_command[k]
-                    self.current_read_resp_command[k] = None
-                    flush_cmd(cmd)
-
-            for q in self.int_read_resp_queue_list:
-                while not q.empty():
-                    q.get_nowait()
-
             self.cur_id = 0
             self.active_id = Counter()
 
@@ -615,8 +657,6 @@ class AxiMasterRead(Reset):
                 self._process_read_cr = cocotb.fork(self._process_read())
             if self._process_read_resp_cr is None:
                 self._process_read_resp_cr = cocotb.fork(self._process_read_resp())
-            if self._process_read_resp_id_cr is None:
-                self._process_read_resp_id_cr = [cocotb.fork(self._process_read_resp_id(i)) for i in range(self.id_count)]
 
     async def _process_read(self):
         while True:
@@ -679,7 +719,7 @@ class AxiMasterRead(Reset):
                 cur_addr += num_bytes
 
             resp_cmd = AxiReadRespCmd(cmd.address, cmd.length, cmd.size, cycles, cmd.prot, burst_list, cmd.event)
-            await self.int_read_resp_command_queue[arid].put(resp_cmd)
+            self.tag_context_manager.start_cmd(arid, resp_cmd)
 
             self.current_read_command = None
 
@@ -702,14 +742,12 @@ class AxiMasterRead(Reset):
             cur_rid = rid
 
             if int(r.rlast):
-                await self.int_read_resp_queue_list[rid].put(burst)
+                self.tag_context_manager.put_resp(rid, burst)
                 burst = []
                 cur_rid = None
 
-    async def _process_read_resp_id(self, rid):
+    async def _process_read_resp_id(self, context, cmd):
-        while True:
+        rid = context.current_tag
-            cmd = await self.int_read_resp_command_queue[rid].get()
-            self.current_read_resp_command[rid] = cmd
 
         num_bytes = 2**cmd.size
 
@@ -727,7 +765,7 @@ class AxiMasterRead(Reset):
         first = True
 
         for burst_length in cmd.burst_list:
-                burst = await self.int_read_resp_queue_list[rid].get()
+            burst = await context.get_resp()
 
             if len(burst) != burst_length:
                 raise Exception(f"Burst length incorrect (ID {rid}, expected {burst_length}, got {len(burst)}")
@@ -769,8 +807,6 @@ class AxiMasterRead(Reset):
 
         cmd.event.set(read_resp)
 
-            self.current_read_resp_command[rid] = None
-
         self.in_flight_operations -= 1
 
         if self.in_flight_operations == 0: