[rocroller] Build Basic-Block Dependence DAGs

shared/rocroller/lib/include/rocRoller/KernelGraph/NodeSchedulingUtils.hpp

@@ -56,7 +56,7 @@ namespace rocRoller::KernelGraph
          * The sub-graph is created by adding the given nodes to a new control graph, and then
          * adding edges between the nodes based on the order of the nodes in the original
          * control graph.
-         * 
+         *
          * rv.compare(cacheMode, a, b) should always return the same result as the original
          * control graph as long as a and b are both in `nodes`.
          */
@@ -70,6 +70,11 @@ namespace rocRoller::KernelGraph
          * `graph.control`.
          */
         void orderNodes(KernelGraph const& graph, std::vector<int>& nodes, auto const& comp);
+
+        /**
+     * Builds the data dependence graph for the given kernel graph
+     */
+        ControlGraph::ControlGraph constructDataDependenceDAG(KernelGraph const& graph);
     }
 }
 

shared/rocroller/lib/source/KernelGraph/NodeSchedulingUtils.cpp

@@ -75,4 +75,127 @@ namespace rocRoller::KernelGraph::NodeScheduling
         return subGraph;
     }
 
+    ControlGraph::ControlGraph constructDataDependenceDAG(KernelGraph const& graph)
+    {
+        ControlGraph::ControlGraph dependenceDAG;
+
+        // Insert all control graph nodes into the data dependence DAG
+        for(auto node : graph.control.getNodes())
+        {
+            dependenceDAG.setElement(node, graph.control.getElement(node));
+        }
+
+        // Create tracer and get all the trace records from it.
+        auto tracer  = ControlFlowRWTracer(graph);
+        auto records = tracer.coordinatesReadWrite();
+
+        std::map<int, int>                     latestWriteToCoord;
+        std::map<int, std::unordered_set<int>> latestReadsToCoord;
+
+        // This assumes that the trace is ordered and records for the
+        // same control operation are consecutive.
+        auto it = records.begin();
+        while(it != records.end())
+        {
+            int currentControl = it->control;
+            while(it != records.end() && it->control == currentControl)
+            {
+                auto coord = it->coordinate;
+
+                // adds WW(output dep) and WR(flow dep) edges
+                if(latestWriteToCoord.contains(coord))
+                {
+                    AssertFatal(latestWriteToCoord[coord] != currentControl,
+                                ShowValue(it->control),
+                                ShowValue(it->coordinate),
+                                ShowValue(it->rw));
+
+                    auto sourceNodeBodyParent
+                        = bodyParents(getTopSetCoordinate(graph, latestWriteToCoord[coord]), graph)
+                              .take(1)
+                              .only();
+                    AssertFatal(sourceNodeBodyParent.has_value(),
+                                "Source node has no body parent",
+                                ShowValue(latestWriteToCoord[coord]));
+
+                    auto destNodeBodyParent
+                        = bodyParents(getTopSetCoordinate(graph, currentControl), graph)
+                              .take(1)
+                              .only();
+                    AssertFatal(destNodeBodyParent.has_value(),
+                                "Dest node has no body parent",
+                                ShowValue(currentControl));
+
+                    if(sourceNodeBodyParent.value() == destNodeBodyParent.value())
+                    {
+                        auto order = graph.control.compareNodes(
+                            UseCacheIfAvailable, latestWriteToCoord[coord], currentControl);
+                        AssertFatal(order == ControlGraph::NodeOrdering::LeftFirst
+                                        || order == ControlGraph::NodeOrdering::RightInBodyOfLeft,
+                                    ShowValue(latestWriteToCoord[coord]),
+                                    ShowValue(currentControl),
+                                    ShowValue(order));
+
+                        dependenceDAG.addElement(ControlGraph::Sequence(),
+                                                 {latestWriteToCoord[coord]},
+                                                 {currentControl});
+                    }
+                }
+
+                if(it->rw == ControlFlowRWTracer::WRITE || it->rw == ControlFlowRWTracer::READWRITE)
+                {
+                    // adds RW(anti dep) edges
+                    for(auto const read : latestReadsToCoord[coord])
+                    {
+                        if(read == currentControl)
+                            continue;
+
+                        auto sourceNodeBodyParent
+                            = bodyParents(getTopSetCoordinate(graph, read), graph).take(1).only();
+                        AssertFatal(sourceNodeBodyParent.has_value(),
+                                    "Source node has no body parent",
+                                    ShowValue(read));
+
+                        auto destNodeBodyParent
+                            = bodyParents(getTopSetCoordinate(graph, currentControl), graph)
+                                  .take(1)
+                                  .only();
+                        AssertFatal(destNodeBodyParent.has_value(),
+                                    "Dest node has no body parent",
+                                    ShowValue(currentControl));
+
+                        if(sourceNodeBodyParent.value() == destNodeBodyParent.value())
+                        {
+                            auto order = graph.control.compareNodes(
+                                UseCacheIfAvailable, read, currentControl);
+                            AssertFatal(order == ControlGraph::NodeOrdering::LeftFirst
+                                            || order
+                                                   == ControlGraph::NodeOrdering::RightInBodyOfLeft,
+                                        ShowValue(read),
+                                        ShowValue(currentControl),
+                                        ShowValue(order));
+
+                            dependenceDAG.addElement(
+                                ControlGraph::Sequence(), {read}, {currentControl});
+                        }
+                    }
+
+                    // Since the current control node writes into this coord,
+                    // the latest reads info needs to be reset.
+                    latestReadsToCoord[coord].clear();
+
+                    // update the latest write to coord
+                    latestWriteToCoord[coord] = currentControl;
+                }
+
+                if(it->rw == ControlFlowRWTracer::READ || it->rw == ControlFlowRWTracer::READWRITE)
+                {
+                    latestReadsToCoord[coord].insert(currentControl);
+                }
+
+                it++;
+            }
+        }
+        return dependenceDAG;
+    }
 }

shared/rocroller/lib/source/KernelGraph/Transformations/RemoveImplicitScheduling.cpp

@@ -48,54 +48,6 @@ namespace rocRoller::KernelGraph
                 return std::nullopt;
             };
 
-            auto loop = getLoopOp(nodes.front()).value_or(-1);
-
-            auto colouring = colourByUnrollValue(graph, -1);
-
-            Log::debug(toString(colouring));
-
-            using Colour = std::set<std::pair<int, int>>;
-
-            std::map<Colour, std::vector<int>> reverse;
-
-            for(auto node : nodes)
-            {
-                auto const& opColour = colouring.operationColour.at(node);
-
-                Colour key(opColour.begin(), opColour.end());
-
-                reverse[key].push_back(node);
-            }
-
-            for(auto& [key, keyOps] : reverse)
-            {
-                std::ranges::sort(keyOps, TopologicalCompare(graph));
-            }
-
-            std::set<int> edgesToKeep;
-
-            for(auto& [key, keyOps] : reverse)
-            {
-                for(int idx = 0; idx + 1 < keyOps.size(); idx++)
-                {
-                    auto thisEdge = graph.control.findEdge(keyOps[idx], keyOps[idx + 1]);
-
-                    if(!thisEdge)
-                    {
-                        AssertFatal(graph.control.compareNodes(
-                                        UseCacheIfAvailable, keyOps[idx], keyOps[idx + 1])
-                                    == ControlGraph::NodeOrdering::LeftFirst);
-
-                        thisEdge = graph.control.addElement(
-                            ControlGraph::Sequence(), {keyOps[idx]}, {keyOps[idx + 1]});
-                    }
-
-                    edgesToKeep.insert(*thisEdge);
-                }
-            }
-
-            Log::debug("Keeping edges ({})", fmt::join(edgesToKeep, ", "));
-
             auto notMultiply = [&graph](int idx) {
                 if(graph.control.getElementType(idx) != Graph::ElementType::Node)
                     return false;
@@ -122,16 +74,30 @@ namespace rocRoller::KernelGraph
 
             Log::debug("Got connections.");
 
+            auto dependenceDAG = NodeScheduling::constructDataDependenceDAG(graph);
+
             for(auto nodeA : nodes)
             {
                 for(auto nodeB : nodes)
                 {
                     if(nodeA == nodeB)
                         continue;
 
-                    auto thisEdge = graph.control.findEdge(nodeA, nodeB);
+                    auto depEdge = dependenceDAG.findEdge(nodeA, nodeB);
+                    auto seqEdge = graph.control.findEdge(nodeA, nodeB);
 
-                    if(thisEdge.has_value() && !edgesToKeep.contains(*thisEdge))
+                    if(depEdge.has_value())
+                    {
+                        if(!seqEdge.has_value())
+                        {
+                            AssertFatal(
+                                graph.control.compareNodes(UseCacheIfAvailable, nodeA, nodeB)
+                                == ControlGraph::NodeOrdering::LeftFirst);
+
+                            graph.control.addElement(ControlGraph::Sequence(), {nodeA}, {nodeB});
+                        }
+                    }
+                    else if(seqEdge.has_value())
                     {
                         auto upstream = connectionsToKeep.at(nodeB);
                         auto order
@@ -141,8 +107,8 @@ namespace rocRoller::KernelGraph
                                     ShowValue(order),
                                     ShowValue(upstream),
                                     ShowValue(nodeB),
-                                    ShowValue(*thisEdge));
-                        graph.control.deleteElement(*thisEdge);
+                                    ShowValue(seqEdge.value()));
+                        graph.control.deleteElement(seqEdge.value());
                         if(order == ControlGraph::NodeOrdering::LeftFirst)
                             graph.control.chain<ControlGraph::Sequence>(upstream, nodeB);
                         else
@@ -157,6 +123,7 @@ namespace rocRoller::KernelGraph
     {
         auto rv = original;
 
+        // grouped by immediate body-parent
         auto groupedMultiplyNodes = NodeScheduling::getGroupedNodes<ControlGraph::Multiply>(rv);
 
         for(auto& [parent, nodes] : groupedMultiplyNodes)