Topological initialization

src/PIM/Common/LabeledList.hpp

@@ -0,0 +1,318 @@
+#pragma once
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/ilist_node.h"
+#include "llvm/ADT/simple_ilist.h"
+
+#include <cassert>
+#include <iterator>
+#include <limits>
+#include <type_traits>
+
+namespace onnx_mlir {
+
+template <typename NodeT>
+class LabeledList;
+
+template <typename NodeT>
+class LabeledListNode : public llvm::ilist_node<NodeT> {
+  friend class LabeledList<NodeT>;
+
+public:
+  using Label = uint64_t;
+
+  LabeledListNode() = default;
+  LabeledListNode(const LabeledListNode&) = delete;
+  LabeledListNode(LabeledListNode&&) = default;
+  LabeledListNode& operator=(LabeledListNode&&) = delete;
+
+  ~LabeledListNode() { assert(owner_ == nullptr && "destroying a linked LabeledListNode"); }
+
+  bool isLinked() const { return owner_ != nullptr; }
+  Label getOrderLabel() const { return label; }
+
+  friend bool operator<(const LabeledListNode& lft, const LabeledListNode& rgt){
+    return lft.label < rgt.label;
+  }
+
+private:
+  const void* owner_ = nullptr;
+  Label label = 0;
+};
+
+template <typename NodeT>
+class LabeledList {
+
+  using Label = typename NodeT::Label;
+
+  static constexpr Label kLowerSentinel = 0;
+  static constexpr Label kUpperSentinel = std::numeric_limits<Label>::max();
+  static constexpr Label kRelabelGap = 2;
+
+public:
+  using List = llvm::simple_ilist<NodeT>;
+  using Iterator = typename List::iterator;
+  using RIterator = typename List::reverse_iterator;
+  using ConstIterator = typename List::const_iterator;
+
+  LabeledList() = default;
+  LabeledList(const LabeledList&) = delete;
+  LabeledList& operator=(const LabeledList&) = delete;
+  LabeledList(LabeledList&&) = delete;
+  LabeledList& operator=(LabeledList&&) = delete;
+
+  ~LabeledList() { clear(); }
+
+  bool empty() const { return size_ == 0; }
+  size_t size() const { return size_; }
+
+  NodeT* front() { return empty() ? nullptr : &nodes_.front(); }
+  const NodeT* front() const { return empty() ? nullptr : &nodes_.front(); }
+
+  NodeT* back() { return empty() ? nullptr : &nodes_.back(); }
+  const NodeT* back() const { return empty() ? nullptr : &nodes_.back(); }
+
+  static NodeT* previous(NodeT* node) {
+    if (!node || !owner(node))
+      return nullptr;
+    auto* list = owner(node);
+    auto it = node->getIterator();
+    if (it == list->nodes_.begin())
+      return nullptr;
+    return *std::prev(it);
+  }
+
+  static const NodeT* previous(const NodeT* node) {
+    if (!node || !owner(node))
+      return nullptr;
+    const auto* list = owner(node);
+    auto it = node->getIterator();
+    if (it == list->nodes_.begin())
+      return nullptr;
+    return *std::prev(it);
+  }
+
+  static NodeT* next(NodeT* node) {
+    if (!node || !owner(node))
+      return nullptr;
+    auto* list = owner(node);
+    auto it = std::next(node->getIterator());
+    if (it == list->nodes_.end())
+      return nullptr;
+    return *it;
+  }
+
+  static const NodeT* next(const NodeT* node) {
+    if (!node || !owner(node))
+      return nullptr;
+    const auto* list = owner(node);
+    auto it = std::next(node->getIterator());
+    if (it == list->nodes_.end())
+      return nullptr;
+    return *it;
+  }
+
+  bool contains(const NodeT* node) const { return node && node->owner_ == this; }
+
+  Label getOrderLabel(const NodeT* node) const {
+    assert(contains(node) && "node must belong to this list");
+    return node->label_;
+  }
+
+  bool comesBefore(const NodeT* lhs, const NodeT* rhs) const {
+    assert(contains(lhs) && contains(rhs) && "nodes must belong to this list");
+    return lhs->label_ < rhs->label_;
+  }
+
+  void pushFront(NodeT* node) { insertBefore(front(), node); }
+
+  void pushBack(NodeT* node) { insertBefore(nullptr, node); }
+
+  void insertBefore(NodeT* nextNode, NodeT* node) {
+    assert(node && "cannot insert a null node");
+    assert(!node->owner_ && "node is already linked");
+    assert(nextNode == nullptr || contains(nextNode));
+
+    Iterator nextIt = nextNode ? getIteratorFor(nextNode) : nodes_.end();
+    nodes_.insert(nextIt, *node);
+    node->owner_ = this;
+    ++size_;
+    assignLabel(getIteratorFor(node));
+  }
+
+  void insertAfter(NodeT* prevNode, NodeT* node) {
+    assert(prevNode == nullptr || contains(prevNode));
+    if (prevNode == nullptr)
+      insertBefore(front(), node);
+    else
+      insertBefore(next(prevNode), node);
+  }
+
+  void remove(NodeT* node) {
+    assert(contains(node) && "node must belong to this list");
+    nodes_.remove(*node);
+    node->owner_ = nullptr;
+    node->label_ = 0;
+    --size_;
+  }
+
+  void moveBefore(NodeT* node, NodeT* nextNode) {
+    assert(contains(node) && "node must belong to this list");
+    assert(nextNode == nullptr || contains(nextNode));
+
+    Iterator nodeIt = getIteratorFor(node);
+    Iterator nextIt = nextNode ? getIteratorFor(nextNode) : nodes_.end();
+    if (nodeIt == nextIt || std::next(nodeIt) == nextIt)
+      return;
+
+    nodes_.splice(nextIt, nodes_, nodeIt);
+    assignLabel(getIteratorFor(node));
+  }
+
+  void moveAfter(NodeT* node, NodeT* prevNode) {
+    assert(contains(node) && "node must belong to this list");
+    assert(prevNode == nullptr || contains(prevNode));
+
+    Iterator nextIt = prevNode ? std::next(getIteratorFor(prevNode)) : nodes_.begin();
+    if (getIteratorFor(node) == nextIt)
+      return;
+    moveBefore(node, nextIt == nodes_.end() ? nullptr : &*nextIt);
+  }
+
+  void clear() {
+    while (!nodes_.empty()) {
+      NodeT* node = &nodes_.front();
+      node->owner_ = nullptr;
+      node->label = 0;
+      nodes_.remove(*node);
+    }
+    size_ = 0;
+  }
+
+  Iterator begin() { return nodes_.begin(); }
+
+  Iterator end() { return nodes_.end(); }
+
+  RIterator rbegin() { return nodes_.rbegin(); }
+  RIterator rend() { return nodes_.rend(); }
+
+private:
+  static const LabeledList* owner(const NodeT* node) { return node->owner_; }
+  static LabeledList* owner(NodeT* node) { return node->owner_; }
+
+  static Label lowerLabel(const NodeT* node) { return node ? node->label : kLowerSentinel; }
+  static Label upperLabel(const NodeT* node) { return node ? node->label : kUpperSentinel; }
+
+  static Label labelGap(Label lower, Label upper) {
+    assert(lower < upper && "labels must be strictly ordered");
+    return upper - lower;
+  }
+
+  static bool hasMidpoint(Label lower, Label upper) { return labelGap(lower, upper) > 1; }
+
+  static bool hasRelabelSlack(Label lower, Label upper, size_t nodeCount) {
+    Label gap = labelGap(lower, upper);
+    return gap / static_cast<Label>(nodeCount + 1) >= kRelabelGap;
+  }
+
+  Iterator getIteratorFor(NodeT* node) { return node->getIterator(); }
+  ConstIterator getiteratorFor(const NodeT* node) const { return node->getIterator(); }
+
+  NodeT* previousNode(Iterator it) {
+    if (it == nodes_.begin())
+      return nullptr;
+    return &*std::prev(it);
+  }
+
+  const NodeT* previousNode(ConstIterator it) const {
+    if (it == nodes_.begin())
+      return nullptr;
+    return &*std::prev(it);
+  }
+
+  NodeT* nextNode(Iterator it) {
+    ++it;
+    if (it == nodes_.end())
+      return nullptr;
+    return &*it;
+  }
+
+  const NodeT* nextNode(ConstIterator it) const {
+    ++it;
+    if (it == nodes_.end())
+      return nullptr;
+    return &*it;
+  }
+
+  void assignLabel(Iterator it) {
+    Label lower = lowerLabel(previousNode(it));
+    Label upper = upperLabel(nextNode(it));
+    if (hasMidpoint(lower, upper)) {
+      (*it).label = lower + static_cast<Label>(labelGap(lower, upper) / 2);
+      return;
+    }
+
+    relabelAround(it);
+  }
+
+  void relabelAround(Iterator center) {
+    size_t targetCount = 1;
+    while (true) {
+      Iterator left = center;
+      Iterator right = center;
+      size_t actualCount = 1;
+      expandWindow(center, targetCount, left, right, actualCount);
+
+      Label lower = lowerLabel(previousNode(left));
+      Label upper = upperLabel(nextNode(right));
+      if (hasRelabelSlack(lower, upper, actualCount)) {
+        relabelWindow(left, actualCount, lower, upper);
+        return;
+      }
+
+      if (left == nodes_.begin() && nextNode(right) == nullptr) {
+        assert(hasRelabelSlack(lower, upper, actualCount) && "label space exhausted");
+        relabelWindow(left, actualCount, lower, upper);
+        return;
+      }
+
+      targetCount *= 2;
+    }
+  }
+
+  void expandWindow(Iterator center, size_t targetCount, Iterator& left, Iterator& right, size_t& actualCount) {
+    left = center;
+    right = center;
+    actualCount = 1;
+
+    while (actualCount < targetCount && (left != nodes_.begin() || nextNode(right) != nullptr)) {
+      if (left != nodes_.begin()) {
+        --left;
+        ++actualCount;
+        if (actualCount == targetCount)
+          break;
+      }
+      if (nextNode(right) != nullptr) {
+        ++right;
+        ++actualCount;
+      }
+    }
+  }
+
+  void relabelWindow(Iterator left, size_t nodeCount, Label lower, Label upper) {
+    assert(nodeCount > 0 && "relabel window must not be empty");
+    Label step = labelGap(lower, upper) / static_cast<Label>(nodeCount + 1);
+    assert(step >= 1 && "relabel step must be positive");
+
+    Iterator it = left;
+    for (size_t index = 1; index <= nodeCount; ++index) {
+      (*it).label = lower + step * index;
+      ++it;
+    }
+  }
+
+  List nodes_;
+  size_t size_ = 0;
+};
+
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/DCPAnalysis.cpp

@@ -54,7 +54,6 @@ DCPAnalysisResult DCPAnalysis::runAnalysis() {
       }
     }
   }
-
   GraphDCP graphDCP(spatWeightedComputes, edges);
   graphDCP.DCP();
   return graphDCP.getResult();

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.cpp

@@ -1,3 +1,6 @@
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+
 #include <algorithm>
 #include <array>
 #include <cassert>
@@ -5,16 +8,17 @@
 #include <cstddef>
 #include <deque>
 #include <fstream>
+#include <iterator>
 #include <vector>
 
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "DCPAnalysis.hpp"
 #include "Graph.hpp"
 #include "Task.hpp"
 #include "UniqueWorklist.hpp"
 #include "Utils.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
 
-std::optional<Edge_pair> addEdge(TaskDCP* parent, TaskDCP* child, Weight_t weight) {
+std::optional<DoubleEdge> addEdge(TaskDCP* parent, TaskDCP* child, Weight_t weight) {
   auto old_child = parent->addChild(child, weight);
   auto old_parent = child->addParent(parent, weight);
   assert(old_child.has_value() == old_parent.has_value() && "The edge must be present in both element");
@@ -43,26 +47,34 @@ int getTranferCost(TaskDCP* parent, TaskDCP* child) {
   return child_position->second;
 }
 
-std::array<std::optional<Edge_pair>, 2> GraphDCP::insertTaskInCPU(CPU cpu, TaskDCP* task, size_t position) {
-  std::array<std::optional<Edge_pair>, 2> ret;
+TaskInsertion GraphDCP::insertTaskInCPU(CPU cpu, TaskDCP* task, size_t position) {
+  TaskInsertion ret;
   task->setCPU(cpu);
   task->setWeight(task->computeWeight(this, cpu));
-  auto& list = mapCPUTasks[cpu];
-  unsigned int total_size = list.size();
-  assert(position <= total_size && "Inserting in a not valid position");
-  auto inserted_point = list.begin();
-  inserted_point = list.insert(std::next(list.begin(), position), task);
+  auto& tasksInCPU = mapCPUTasks[cpu];
+  unsigned int numCPUTask = tasksInCPU.size();
+  assert(position <= numCPUTask && "Inserting in a not valid position");
+  auto insertedPoint = tasksInCPU.insert(std::next(tasksInCPU.begin(), position), task);
+  ret.oldTopologicalPosition = std::next(task->getTopologicalPosition());
+  ret.cpuModified = cpu;
+  ret.taskInserted = task;
+  ret.graph = this;
 
-  if (inserted_point != list.begin()) {
-    auto precedent_point = std::prev(inserted_point, 1);
-    auto old_edge = addEdge(*precedent_point, *inserted_point, 0);
-    ret[0] = old_edge;
+  if (insertedPoint != tasksInCPU.begin()) {
+    auto precedentPoint = std::prev(insertedPoint, 1);
+    auto oldEdge = addEdge(*precedentPoint, *insertedPoint, 0);
+    ret.beforeNode = oldEdge;
+
+    if (*task < **insertedPoint)
+      topologicalMoveAfter(task, *precedentPoint);
   }
 
-  if (std::next(inserted_point) != list.end()) {
-    auto next_point = std::next(inserted_point, 1);
-    auto old_edge = addEdge(*inserted_point, *next_point, 0);
-    ret[1] = old_edge;
+  if (std::next(insertedPoint) != tasksInCPU.end()) {
+    auto nextPoint = std::next(insertedPoint, 1);
+    auto oldEdge = addEdge(*insertedPoint, *nextPoint, 0);
+    ret.afterNode = oldEdge;
+    if (**insertedPoint < *task)
+      topologicalMoveBefore(task, *nextPoint);
   }
   return ret;
 }
@@ -88,7 +100,7 @@ void GraphDCP::removeTaskFromCPU(CPU cpu, TaskDCP* task) {
 std::vector<TaskDCP*> GraphDCP::getRoots() {
   std::vector<TaskDCP*> tmp;
   for (auto& task : nodes)
-    if (!task.hasParent())
+    if (!task.hasParents())
       tmp.push_back(&task);
   return tmp;
 }
@@ -96,47 +108,33 @@ std::vector<TaskDCP*> GraphDCP::getRoots() {
 void GraphDCP::initAEST() {
   UniqueWorkList<std::deque<TaskDCP*>> worklists(getRoots());
 
-  while (!worklists.empty()) {
-    TaskDCP& task = *worklists.front();
-    bool modified = true;
-    while (modified) {
-      modified = false;
-      for (auto& child : task.childs) {
-        if (worklists.allElementContained(
-              child.first->parents.begin(), child.first->parents.end(), [](Edge_t edge) { return edge.first; })) {
-          modified |= worklists.push_back(child.first);
-        }
-      }
-    }
+  auto& worklist = topologicalOrder;
+  int max_dcpl = 0;
+  for (auto& task : worklist) {
     int max_parent_aest = 0;
-    for (auto& parent : task.parents) {
-      max_parent_aest = std::max(
-        parent.first->getAEST() + parent.first->getWeight() + getTranferCost(parent.first, &task), max_parent_aest);
+    for (Edge_t parentEdge : task.parents) {
+      max_parent_aest =
+        std::max(parentEdge.first->getAEST() + parentEdge.first->getWeight() + getTranferCost(parentEdge.first, &task),
+                 max_parent_aest);
     }
     task.setAEST(max_parent_aest);
-    worklists.pop_front();
+    max_dcpl = std::max(max_parent_aest + task.getWeight(), max_dcpl);
   }
+  DCPL = max_dcpl;
 }
 
 int GraphDCP::computeAEST(TaskDCP* task, CPU cpu) {
   int max_parent_aest = 0;
-  for (auto& parent : task->parents) {
+  for (Edge_t parentEdge : task->parents) {
     int transfer_cost = 0;
-    if (!(parent.first->isScheduled() && cpu == *parent.first->getCPU()))
-      transfer_cost = getTranferCost(parent.first, task);
-
-    max_parent_aest = std::max(parent.first->getAEST() + parent.first->getWeight() + transfer_cost, max_parent_aest);
+    if (!(parentEdge.first->isScheduled() && cpu == *parentEdge.first->getCPU()))
+      transfer_cost = getTranferCost(parentEdge.first, task);
+    max_parent_aest =
+      std::max(parentEdge.first->getAEST() + parentEdge.first->getWeight() + transfer_cost, max_parent_aest);
   }
   return max_parent_aest;
 }
 
-int GraphDCP::initDCPL() {
-  int max_aest = 0;
-  for (auto& node : nodes)
-    max_aest = std::max(node.getAEST() + node.getWeight(), max_aest);
-  return max_aest;
-}
-
 int GraphDCP::computeDCPL(TaskDCP* task, CPU cpu) {
   int max_aest = 0;
   for (auto& node : nodes)
@@ -148,59 +146,26 @@ int GraphDCP::computeDCPL(TaskDCP* task, CPU cpu) {
 }
 
 void GraphDCP::initALST() {
-  int dcpl = initDCPL();
-  std::vector<TaskDCP*> roots = getRoots();
-  UniqueWorkList<std::vector<TaskDCP*>> worklists(roots);
-  worklists.reserve(nodes.size());
-  size_t i = 0;
-  while (i != worklists.size()) {
-    bool modified = true;
-    while (modified) {
-      modified = false;
-      for (auto& child : worklists.at(i)->childs) {
-        if (worklists.allElementContained(
-              child.first->parents.begin(), child.first->parents.end(), [](Edge_t edge) { return edge.first; })) {
-          modified |= worklists.push_back(child.first);
-        }
-      }
-    }
-    i++;
-  }
+  int dcpl = getDCPL();
+  auto& worklists = topologicalOrder;
 
-  while (!worklists.empty()) {
-    TaskDCP& node = *worklists.back();
+  for (TaskDCP& node : llvm::reverse(worklists)) {
     int min_alst = INT_MAX;
     if (!node.hasChilds())
       min_alst = dcpl - node.getWeight();
-    for (auto child : node.childs)
-      min_alst = std::min(min_alst, child.first->getALST() - node.getWeight() - getTranferCost(&node, child.first));
+    for (Edge_t childEdge : node.childs)
+      min_alst =
+        std::min(min_alst, childEdge.first->getALST() - node.getWeight() - getTranferCost(&node, childEdge.first));
     node.setALST(min_alst);
-    worklists.pop_back();
   }
 }
 
-std::unordered_map<TaskDCP*, int> GraphDCP::computeALST(TaskDCP* task, CPU cpu) {
+llvm::DenseMap<TaskDCP*, int> GraphDCP::computeALST(TaskDCP* task, CPU cpu) {
   int dcpl = computeDCPL(task, cpu);
-  std::unordered_map<TaskDCP*, int> temp_ALST;
-  UniqueWorkList<std::vector<TaskDCP*>> worklists(getRoots());
-  size_t i = 0;
-  while (i != worklists.size()) {
+  llvm::DenseMap<TaskDCP*, int> temp_ALST;
 
-    bool modified = true;
-    while (modified) {
-      modified = false;
-      for (auto& child : worklists.at(i)->childs) {
-        if (worklists.allElementContained(
-              child.first->parents.begin(), child.first->parents.end(), [](Edge_t edge) { return edge.first; })) {
-          modified |= worklists.push_back(child.first);
-        }
-      }
-    }
-    i++;
-  }
-
-  while (!worklists.empty()) {
-    TaskDCP& node = *worklists.back();
+  auto& worklists = topologicalOrder;
+  for (TaskDCP& node : llvm::reverse(worklists)) {
     int min_alst = INT_MAX;
     if (!node.hasChilds()) {
       if (&node != task)
@@ -209,14 +174,13 @@ std::unordered_map<TaskDCP*, int> GraphDCP::computeALST(TaskDCP* task, CPU cpu)
         min_alst = dcpl - node.computeWeight(this, cpu);
     }
 
-    for (auto child : node.childs) {
-      int transfer_cost = getTranferCost(&node, child.first);
-      if (&node == task && child.first->isScheduled() && cpu == *child.first->getCPU())
+    for (Edge_t childEdge : node.childs) {
+      int transfer_cost = getTranferCost(&node, childEdge.first);
+      if (&node == task && childEdge.first->isScheduled() && cpu == *childEdge.first->getCPU())
         transfer_cost = 0;
-      min_alst = std::min(min_alst, temp_ALST[child.first] - node.getWeight() - transfer_cost);
+      min_alst = std::min(min_alst, temp_ALST[childEdge.first] - node.getWeight() - transfer_cost);
     }
     temp_ALST[&node] = min_alst;
-    worklists.pop_back();
   }
   return temp_ALST;
 }
@@ -253,6 +217,124 @@ TaskDCP* GraphDCP::findCandidate(std::vector<TaskDCP*> nodes) {
   return *best_node;
 }
 
+void GraphDCP::initTopological() {
+  UniqueWorkList<std::vector<TaskDCP*>> worklists(getRoots());
+  long long flag = getUniqueFlag();
+  for (auto root : worklists)
+    root->setFlag(flag);
+
+  size_t i = 0;
+  while (i != worklists.size()) {
+    for (auto& child : worklists.at(i)->childs) {
+      TaskDCP* childTask = child.first;
+      if (std::all_of(childTask->parents.begin(), childTask->parents.end(), [flag](Edge_t edge) {
+            return edge.first->getFlag() == flag;
+          })) {
+        worklists.push_back(childTask);
+        childTask->setFlag(flag);
+      }
+    }
+    i++;
+  }
+
+  for (auto task : worklists)
+    topologicalOrder.pushBack(task);
+}
+
+void GraphDCP::topologicalMoveAfter(TaskDCP* task, TaskDCP* pivotPoint) {
+  auto moveChildAfterMe = [this](TaskDCP* origTask) -> void {
+    auto cmp = [](Edge_t lft, Edge_t rgt) { return *rgt.first < *lft.first; };
+    TaskDCP* insertionPoint = origTask;
+    std::vector<Edge_t>& childEdges = origTask->childs;
+    std::vector<TaskDCP*> worklist;
+    worklist.push_back(origTask);
+    size_t i = 0;
+    while (i < worklist.size()) {
+      auto task = worklist[i];
+      // build min heap Complexity 3N
+      std::make_heap(childEdges.begin(), childEdges.end(), cmp);
+      auto lastPoppedIter = childEdges.end();
+      bool foundChildInOrder = false;
+      while (lastPoppedIter != childEdges.begin()) {
+        std::pop_heap(childEdges.begin(), lastPoppedIter, cmp);
+        lastPoppedIter = std::prev(lastPoppedIter, 1);
+        auto currentChild = (*lastPoppedIter).first;
+        if (*currentChild < *task) {
+          topologicalOrder.moveAfter(currentChild, insertionPoint);
+          insertionPoint = currentChild;
+        }
+        else {
+          foundChildInOrder = true;
+          break;
+        }
+      }
+
+      if (foundChildInOrder)
+        lastPoppedIter++;
+
+      for (auto it = lastPoppedIter; it != childEdges.end(); ++it)
+        if (it->first->hasChilds())
+          worklist.push_back(it->first);
+      i++;
+    }
+  };
+
+  if (!(*task < *pivotPoint)) {
+    return;
+    topologicalOrder.moveAfter(task, pivotPoint);
+    if (task->hasChilds())
+      moveChildAfterMe(task);
+  }
+}
+
+void GraphDCP::topologicalMoveBefore(TaskDCP* task, TaskDCP* pivotPoint) {
+
+  auto moveParentBeforeMe = [this](TaskDCP* origTask) -> void {
+    auto cmp = [](Edge_t lft, Edge_t rgt) { return *lft.first < *rgt.first; };
+    TaskDCP* insertionPoint = origTask;
+    std::vector<Edge_t>& parentEdges = origTask->parents;
+    std::vector<TaskDCP*> worklist;
+    worklist.push_back(origTask);
+    size_t i = 0;
+    while (i < worklist.size()) {
+      auto task = worklist[i];
+      // build max heap Complexity 3N
+      std::make_heap(parentEdges.begin(), parentEdges.end(), cmp);
+      auto lastPoppedIter = parentEdges.end();
+      bool foundParentInOrder = false;
+      while (lastPoppedIter != parentEdges.begin()) {
+        std::pop_heap(parentEdges.begin(), lastPoppedIter, cmp);
+        lastPoppedIter = std::prev(lastPoppedIter, 1);
+        auto currentParent = (*lastPoppedIter).first;
+        if (*currentParent < *task) {
+          topologicalOrder.moveBefore(currentParent, insertionPoint);
+          insertionPoint = currentParent;
+        }
+        else {
+          foundParentInOrder = true;
+          break;
+        }
+      }
+
+      if (foundParentInOrder)
+        lastPoppedIter++;
+
+      for (auto it = lastPoppedIter; it != parentEdges.end(); ++it)
+        if (it->first->hasParents())
+          worklist.push_back(it->first);
+      i++;
+    }
+  };
+
+  if (!(*task < *pivotPoint)) {
+    return;
+    topologicalOrder.moveBefore(task, pivotPoint);
+    if (task->hasParents())
+      moveParentBeforeMe(task);
+  }
+}
+
+
 GraphDCP::FindSlot GraphDCP::findSlot(TaskDCP* candidate, CPU cpu, bool push) {
   int aest_on_cpu = computeAEST(candidate, cpu);
   auto tmp_ALST = computeALST(candidate, cpu);
@@ -353,8 +435,8 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
       }
     }
     else if (candidate->hasChilds()) {
-      auto dcpl = initDCPL();
-      auto old_edges = insertTaskInCPU(current_cpu, candidate, slot.index);
+      auto dcpl = getDCPL();
+      auto taskInsertion = insertTaskInCPU(current_cpu, candidate, slot.index);
       initAEST();
       initALST();
       Edge_t smallest_child {nullptr, 0};
@@ -377,15 +459,9 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
         best_composite = slot.aest + child_slot.aest;
         best_slot = slot;
       }
-      removeTaskFromCPU(current_cpu, candidate);
+      taskInsertion.rollBack();
       initAEST();
       initALST();
-      for (auto opt_edge : old_edges) {
-        if (opt_edge.has_value()) {
-          auto double_edge = *opt_edge;
-          addEdge(double_edge.first.first, double_edge.second.first, double_edge.first.second);
-        }
-      }
     }
   }
   if (best_process == -1) {
@@ -398,6 +474,7 @@ void GraphDCP::selectProcessor(TaskDCP* candidate, bool push) {
 }
 
 void GraphDCP::DCP() {
+  initTopological();
   initAEST();
   initALST();
   to_dot();

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.hpp

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
 
 #include <list>
 #include <optional>
@@ -11,7 +12,7 @@
 #include "Task.hpp"
 #include "Utils.hpp"
 
-std::optional<Edge_pair> addEdge(TaskDCP* parent, TaskDCP* child, Weight_t weight);
+std::optional<DoubleEdge> addEdge(TaskDCP* parent, TaskDCP* child, Weight_t weight);
 void removeEdge(TaskDCP* parent, TaskDCP* child);
 int getTranferCost(TaskDCP* parent, TaskDCP* child);
 
@@ -23,20 +24,32 @@ class GraphDCP {
   };
 
   std::vector<TaskDCP> nodes;
+  onnx_mlir::LabeledList<TaskDCP> topologicalOrder;
   std::unordered_map<CPU, std::list<TaskDCP*>> mapCPUTasks;
   CPU last_cpu = 0;
+  long long flag = 1;
+  int DCPL;
 
-  std::array<std::optional<Edge_pair>, 2> insertTaskInCPU(CPU cpu, TaskDCP* task, size_t position);
+  TaskInsertion insertTaskInCPU(CPU cpu, TaskDCP* task, size_t position);
   void removeTaskFromCPU(CPU cpu, TaskDCP* task);
 
   std::vector<TaskDCP*> getRoots();
 
+  long long getUniqueFlag() {return flag++;};
+  
   void initAEST();
-  int computeAEST(TaskDCP* task, CPU cpu);
   int initDCPL();
-  int computeDCPL(TaskDCP* task, CPU cpu);
   void initALST();
-  std::unordered_map<TaskDCP*, int> computeALST(TaskDCP* task, CPU cpu);
+
+  int computeAEST(TaskDCP* task, CPU cpu);
+  int computeDCPL(TaskDCP* task, CPU cpu);
+  int getDCPL() {return DCPL;};
+
+  void initTopological();
+  void topologicalMoveAfter(TaskDCP* task, TaskDCP * pivotPoint); 
+  void topologicalMoveBefore(TaskDCP* task, TaskDCP * pivotPoint); 
+
+  llvm::DenseMap<TaskDCP*, int> computeALST(TaskDCP* task, CPU cpu);
 
   TaskDCP* findCandidate(std::vector<TaskDCP*> nodes);
   void selectProcessor(TaskDCP* candidate, bool push);
@@ -45,13 +58,17 @@ class GraphDCP {
   FindSlot findSlot(TaskDCP* candidate, CPU cpu, bool push);
   void to_dot();
 
+  friend TaskInsertion;
+
+
 public:
   void DCP();
   GraphDCP(llvm::ArrayRef<onnx_mlir::spatial::SpatWeightedCompute> spatWeightedComputes,
            llvm::ArrayRef<EdgesIndex> edges)
   : nodes(), mapCPUTasks() {
-    for (auto spatWeightedCompute : spatWeightedComputes)
+    for (auto spatWeightedCompute : spatWeightedComputes){
       nodes.emplace_back(spatWeightedCompute);
+    }
 
     for (auto [start, end, weight] : edges)
       makeEdge(start, end, weight);

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Task.cpp

@@ -44,4 +44,17 @@ bool TaskDCP::hasDescendent(TaskDCP* child) {
 }
 
 // TODO fare qualcosa di sensato
-int TaskDCP::computeWeight(GraphDCP* graph, CPU cpu) { return orig_weight; }
+int TaskDCP::computeWeight(GraphDCP* graph, CPU cpu) { return origWeight; }
+
+void TaskInsertion::rollBack() {
+  graph->removeTaskFromCPU(cpuModified, taskInserted);
+  if (beforeNode.has_value()) {
+    auto double_edge = *beforeNode;
+    addEdge(double_edge.first.first, double_edge.second.first, double_edge.first.second);
+  }
+  if (afterNode.has_value()) {
+    auto double_edge = *beforeNode;
+    addEdge(double_edge.first.first, double_edge.second.first, double_edge.first.second);
+  }
+  graph->topologicalOrder.moveBefore( taskInserted,&*oldTopologicalPosition );
+}

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Task.hpp

@@ -2,22 +2,25 @@
 
 #include <cassert>
 #include <cstdint>
+#include <iterator>
+#include <list>
 #include <optional>
 #include <vector>
 
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "Utils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
-std::optional<Edge_pair> addEdge(TaskDCP* parent, TaskDCP* child, Weight_t weight);
+std::optional<DoubleEdge> addEdge(TaskDCP* parent, TaskDCP* child, Weight_t weight);
 void removeEdge(TaskDCP* parent, TaskDCP* child);
 
-class TaskDCP {
+class TaskDCP : public onnx_mlir::LabeledListNode<TaskDCP> {
   onnx_mlir::spatial::SpatWeightedCompute spatWeightedCompute;
   int aest;
   int alst;
   std::optional<CPU> scheduledCPU;
   int weight;
-  int orig_weight;
+  int origWeight;
+  long long flag = 0;
 
   std::optional<Edge_t> addChild(TaskDCP* child, Weight_t weight);
   std::optional<Edge_t> addChild(TaskDCP& child, Weight_t weight) { return addChild(&child, weight); }
@@ -36,12 +39,12 @@ public:
   std::vector<Edge_t> childs;
   TaskDCP() = default;
   TaskDCP(onnx_mlir::spatial::SpatWeightedCompute spatWeightedCompute)
-  : spatWeightedCompute(spatWeightedCompute),
+  : onnx_mlir::LabeledListNode<TaskDCP>(), spatWeightedCompute(spatWeightedCompute),
     aest(0),
     alst(0),
     scheduledCPU(),
     weight(getSpatWeightCompute(spatWeightedCompute)),
-    orig_weight(weight),
+    origWeight(weight),
     parents(),
     childs() {}
 
@@ -55,13 +58,13 @@ public:
     if (isScheduled())
       return weight;
     else
-      return orig_weight;
+      return origWeight;
   }
   void setWeight(int val) { weight = val; }
-  void resetWeight() { weight = orig_weight; }
+  void resetWeight() { weight = origWeight; }
   int computeWeight(GraphDCP* graph, CPU cpu);
 
-  bool hasParent() { return parents.size() != 0; }
+  bool hasParents() { return parents.size() != 0; }
   bool hasChilds() { return childs.size() != 0; }
 
   int getAEST() { return aest; }
@@ -81,7 +84,24 @@ public:
   bool isScheduled() const { return scheduledCPU.has_value(); }
   onnx_mlir::spatial::SpatWeightedCompute getSpatWeightedCompute() { return spatWeightedCompute; }
 
-  friend std::optional<Edge_pair> addEdge(TaskDCP* parent, TaskDCP* child, Weight_t weight);
+  void setFlag(long long val) { flag = val; }
+  long long getFlag() const { return flag; }
+
+  onnx_mlir::LabeledList<TaskDCP>::Iterator getTopologicalPosition() { return getIterator(); }
+
+  friend std::optional<DoubleEdge> addEdge(TaskDCP* parent, TaskDCP* child, Weight_t weight);
   friend void removeEdge(TaskDCP* parent, TaskDCP* child);
   friend int getTranferCost(TaskDCP* parent, TaskDCP* child);
 };
+
+
+struct TaskInsertion {
+  std::optional<DoubleEdge> beforeNode;
+  std::optional<DoubleEdge> afterNode;
+  onnx_mlir::LabeledList<TaskDCP>::Iterator oldTopologicalPosition;
+  CPU cpuModified;
+  TaskDCP* taskInserted;
+  GraphDCP* graph;
+
+  void rollBack(); 
+};

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/UniqueWorklist.hpp

@@ -1,7 +1,9 @@
 #pragma once
 
+#include "llvm/ADT/DenseSet.h"
 #include <cassert>
 #include <type_traits>
+#include <iostream>
 #include <unordered_set>
 
 template <typename T, typename = void>
@@ -15,7 +17,7 @@ class UniqueWorkList {
 
   using V = typename T::value_type;
   T storage;
-  std::unordered_set<V> set;
+  llvm::DenseSet<V> set;
 
 public:
   UniqueWorkList() = default;
@@ -24,7 +26,7 @@ public:
   UniqueWorkList(const arg_ty& from)
   : storage() {
     for (auto& element : from) {
-      if (set.count(element) == 0) {
+      if (!set.contains(element)) {
         storage.push_back(element);
         set.insert(element);
       }
@@ -41,7 +43,7 @@ public:
   V& back() { return storage.back(); }
 
   bool push_back(const V& val) {
-    if (set.count(val) == 0) {
+    if (!set.contains(val)) {
       storage.push_back(val);
       set.insert(val);
       return true;
@@ -61,10 +63,11 @@ public:
 
   void pop_back() { storage.pop_back(); }
 
+
   template <typename Iterator, typename Mapper>
   bool allElementContained(Iterator start, Iterator end, Mapper map) {
     while (start != end) {
-      if (set.count(map(*start)) == 0)
+      if (!set.contains(map(*start)))
         return false;
       std::advance(start, 1);
     }

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Utils.hpp

@@ -6,20 +6,24 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <list>
 #include <utility>
 #include <vector>
 
+#include "src/Accelerators/PIM/Common/LabeledList.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Support/TypeUtilities.hpp"
 
+
 using CPU = int;
 using Weight_t = int;
 class TaskDCP;
 class GraphDCP;
 using Edge_t = std::pair<TaskDCP*, Weight_t>;
-using Edge_pair = std::pair<Edge_t, Edge_t>;
+using DoubleEdge = std::pair<Edge_t, Edge_t>;
 using EdgesIndex = std::tuple<int64_t, int64_t, int64_t>;
 
+
 template <typename T>
 void fastRemove(std::vector<std::pair<T*, Weight_t>>& vector, T* to_remove) {
   auto position =