1 files changed, 24 insertions, 20 deletions

diff --git a/src/mem/ruby/network/garnet/flexible-pipeline/NetworkInterface.cc b/src/mem/ruby/network/garnet/flexible-pipeline/NetworkInterface.cc
index 13bbe2b08..26d2423e8 100644
--- a/src/mem/ruby/network/garnet/flexible-pipeline/NetworkInterface.cc
+++ b/src/mem/ruby/network/garnet/flexible-pipeline/NetworkInterface.cc
@@ -49,13 +49,10 @@ NetworkInterface::NetworkInterface(const Params *p)
     m_virtual_networks  = p->virt_nets;
     m_vc_per_vnet = p->vcs_per_vnet;
     m_num_vcs = m_vc_per_vnet*m_virtual_networks;
-
     m_vc_round_robin = 0;
-    m_ni_buffers.resize(m_num_vcs);
-    inNode_ptr.resize(m_virtual_networks);
-    outNode_ptr.resize(m_virtual_networks);
 
     // instantiating the NI flit buffers
+    m_ni_buffers.resize(m_num_vcs);
     for (int i =0; i < m_num_vcs; i++)
         m_ni_buffers[i] = new flitBuffer();
 
@@ -93,18 +90,20 @@ NetworkInterface::addOutPort(NetworkLink *out_link)
 }
 
 void
-NetworkInterface::addNode(vector<MessageBuffer*>& in,
-    vector<MessageBuffer*>& out)
+NetworkInterface::addNode(map<int, MessageBuffer*>& in,
+    map<int, MessageBuffer*>& out)
 {
-    assert(in.size() == m_virtual_networks);
     inNode_ptr = in;
     outNode_ptr = out;
 
-    // protocol injects messages into the NI
-    for (int j = 0; j < m_virtual_networks; j++) {
-        inNode_ptr[j]->setConsumer(this);
-        inNode_ptr[j]->setReceiver(this);
-        outNode_ptr[j]->setSender(this);
+    for (auto& it: in) {
+        // the protocol injects messages into the NI
+        it.second->setConsumer(this);
+        it.second->setReceiver(this);
+    }
+
+    for (auto& it : out) {
+        it.second->setSender(this);
     }
 }
 
@@ -243,12 +242,14 @@ NetworkInterface::wakeup()
 
     //Checking for messages coming from the protocol
     // can pick up a message/cycle for each virtual net
-    for (int vnet = 0; vnet < m_virtual_networks; vnet++) {
-        while (inNode_ptr[vnet]->isReady()) // Is there a message waiting
-        {
-            msg_ptr = inNode_ptr[vnet]->peekMsgPtr();
+    for (auto it = inNode_ptr.begin(); it != inNode_ptr.end(); ++it) {
+        int vnet = (*it).first;
+        MessageBuffer *b = (*it).second;
+
+        while (b->isReady()) { // Is there a message waiting
+            msg_ptr = b->peekMsgPtr();
             if (flitisizeMessage(msg_ptr, vnet)) {
-                inNode_ptr[vnet]->dequeue();
+                b->dequeue();
             } else {
                 break;
             }
@@ -324,14 +325,17 @@ NetworkInterface::scheduleOutputLink()
 void
 NetworkInterface::checkReschedule()
 {
-    for (int vnet = 0; vnet < m_virtual_networks; vnet++) {
-        if (inNode_ptr[vnet]->isReady()) { // Is there a message waiting
+    for (const auto& it : inNode_ptr) {
+        MessageBuffer *b = it.second;
+
+        while (b->isReady()) { // Is there a message waiting
             scheduleEvent(Cycles(1));
             return;
         }
     }
+
     for (int vc = 0; vc < m_num_vcs; vc++) {
-        if (m_ni_buffers[vc]->isReadyForNext(curCycle())) {
+        if (m_ni_buffers[vc]->isReady(curCycle() + Cycles(1))) {
             scheduleEvent(Cycles(1));
             return;
         }