[rocroller] Build Basic-Block Dependence DAGs

shared/rocroller/lib/include/rocRoller/KernelGraph/ControlGraph/DataDependenceDAG.hpp

@@ -0,0 +1,11 @@
+// Copyright Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <rocRoller/KernelGraph/ControlGraph/ControlGraph.hpp>
+
+namespace rocRoller::KernelGraph::DataDependenceDAG
+{
+    ControlGraph::ControlGraph ConstructDataDependenceDAG(KernelGraph const& graph);
+}

shared/rocroller/lib/include/rocRoller/KernelGraph/ControlGraph/DataDependenceDAG_detail.hpp

@@ -0,0 +1,28 @@
+// Copyright Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <rocRoller/KernelGraph/ControlGraph/ControlFlowRWTracer.hpp>
+
+namespace rocRoller::KernelGraph::DataDependenceDAG::Detail
+{
+    class DataDependenceDAGDetail
+    {
+    public:
+        DataDependenceDAGDetail(KernelGraph const& graph);
+
+        int  getBodyParent(int control);
+        void addDependenceEdge(int sourceControl, int destControl);
+        void processReadWriteRecord(ControlFlowRWTracer::ReadWriteRecord const& record);
+        void constructDataDependenceDAG();
+        ControlGraph::ControlGraph getDataDependenceDAG();
+
+    private:
+        KernelGraph const&                     m_graph;
+        ControlGraph::ControlGraph             m_dependenceDAG;
+        std::unordered_map<int, int>           m_bodyParentCache;
+        std::map<int, int>                     m_latestWriteToCoord;
+        std::map<int, std::unordered_set<int>> m_latestReadsToCoord;
+    };
+}

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

@@ -32,7 +32,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`.
          */
@@ -46,6 +46,15 @@ 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.
+     *
+     * The dependences are represented at each basic block level i.e. between the
+     * nodes sharing the same immediate body-parent in `graph.control`.
+     * The flow, anti, and output data dependences between such nodes are included.
+         */
+        ControlGraph::ControlGraph ConstructDataDependenceDAG(KernelGraph const& graph);
     }
 }
 

shared/rocroller/lib/source/KernelGraph/ControlGraph/CMakeLists.txt

@@ -4,5 +4,6 @@
 target_sources(rocroller
     PRIVATE
         "${CMAKE_CURRENT_SOURCE_DIR}/ControlGraph.cpp"
+        "${CMAKE_CURRENT_SOURCE_DIR}/DataDependenceDAG.cpp"
         "${CMAKE_CURRENT_SOURCE_DIR}/Operation.cpp"
 )

shared/rocroller/lib/source/KernelGraph/ControlGraph/DataDependenceDAG.cpp

@@ -0,0 +1,146 @@
+// Copyright Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include <rocRoller/KernelGraph/ControlGraph/DataDependenceDAG.hpp>
+#include <rocRoller/KernelGraph/ControlGraph/DataDependenceDAG_detail.hpp>
+#include <rocRoller/KernelGraph/Utils.hpp>
+
+namespace rocRoller::KernelGraph
+{
+    namespace DataDependenceDAG
+    {
+        using ReadWrite = ControlFlowRWTracer::ReadWrite;
+
+        ControlGraph::ControlGraph ConstructDataDependenceDAG(KernelGraph const& graph)
+        {
+            using namespace Detail;
+
+            DataDependenceDAGDetail obj(graph);
+            obj.constructDataDependenceDAG();
+            return obj.getDataDependenceDAG();
+        }
+
+        namespace Detail
+        {
+            DataDependenceDAGDetail::DataDependenceDAGDetail(KernelGraph const& graph)
+                : m_graph(graph)
+            {
+            }
+
+            ControlGraph::ControlGraph DataDependenceDAGDetail::getDataDependenceDAG()
+            {
+                return m_dependenceDAG;
+            }
+
+            void DataDependenceDAGDetail::constructDataDependenceDAG()
+            {
+                // Insert all control graph nodes into the data dependence DAG
+                for(auto node : m_graph.control.getNodes())
+                {
+                    m_dependenceDAG.setElement(node, m_graph.control.getElement(node));
+                }
+
+                auto tracer  = ControlFlowRWTracer(m_graph);
+                auto records = tracer.coordinatesReadWrite();
+
+                // This assumes that the trace is ordered and records for the
+                // same control operation are consecutive.
+                std::unordered_set<int> seen;
+                for(auto iter = records.begin(); iter != records.end();)
+                {
+                    auto currentControl = iter->control;
+
+                    AssertFatal(seen.find(currentControl) == seen.end(),
+                                "The records for the same control operation are not consecutive.",
+                                ShowValue(currentControl));
+
+                    for(; iter != records.end() && iter->control == currentControl; ++iter)
+                    {
+                        processReadWriteRecord(*iter);
+                    }
+
+                    seen.insert(currentControl);
+                }
+            }
+
+            int DataDependenceDAGDetail::getBodyParent(int control)
+            {
+                if(auto iter = m_bodyParentCache.find(control); iter != m_bodyParentCache.end())
+                    return iter->second;
+
+                auto topSetCoordinate = getTopSetCoordinate(m_graph, control);
+                auto bodyParent       = bodyParents(topSetCoordinate, m_graph).take(1).only();
+                AssertFatal(bodyParent.has_value(),
+                            "Control node has no body parent",
+                            ShowValue(control),
+                            ShowValue(topSetCoordinate));
+
+                m_bodyParentCache.emplace(control, bodyParent.value());
+                return bodyParent.value();
+            }
+
+            void DataDependenceDAGDetail::addDependenceEdge(int sourceControl, int destControl)
+            {
+                AssertFatal(
+                    sourceControl != destControl, ShowValue(sourceControl), ShowValue(destControl));
+
+                auto sourceBodyParent = getBodyParent(sourceControl);
+                auto destBodyParent   = getBodyParent(destControl);
+
+                if(sourceBodyParent != destBodyParent)
+                    return;
+
+                if(!m_dependenceDAG.findEdge(sourceControl, destControl).has_value())
+                {
+                    m_dependenceDAG.addElement(
+                        ControlGraph::Sequence(), {sourceControl}, {destControl});
+                }
+            }
+
+            void DataDependenceDAGDetail::processReadWriteRecord(
+                ControlFlowRWTracer::ReadWriteRecord const& record)
+            {
+                AssertFatal(record.rw != ReadWrite::Count,
+                            ShowValue(record.control),
+                            ShowValue(record.coordinate),
+                            ShowValue(record.rw));
+
+                if(auto writeIter = m_latestWriteToCoord.find(record.coordinate);
+                   writeIter != m_latestWriteToCoord.end())
+                {
+                    AssertFatal(writeIter->second != record.control,
+                                ShowValue(writeIter->second),
+                                ShowValue(record.control),
+                                ShowValue(record.coordinate),
+                                ShowValue(record.rw));
+
+                    // adds WW(output dep) and WR(flow dep) edges
+                    addDependenceEdge(writeIter->second, record.control);
+                }
+
+                if(record.rw == ReadWrite::WRITE || record.rw == ReadWrite::READWRITE)
+                {
+                    for(auto const readControl : m_latestReadsToCoord[record.coordinate])
+                    {
+                        if(readControl == record.control)
+                            continue;
+
+                        // adds RW(anti dep) edges
+                        addDependenceEdge(readControl, record.control);
+                    }
+
+                    // Since the current control node writes into this coord,
+                    // the latest reads info needs to be reset.
+                    m_latestReadsToCoord[record.coordinate].clear();
+                    // update the latest write to coord
+                    m_latestWriteToCoord[record.coordinate] = record.control;
+                }
+
+                if(record.rw == ReadWrite::READ || record.rw == ReadWrite::READWRITE)
+                {
+                    m_latestReadsToCoord[record.coordinate].insert(record.control);
+                }
+            }
+        }
+    }
+}

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

@@ -6,6 +6,10 @@
 #include <rocRoller/KernelGraph/Transforms/Simplify.hpp>
 #include <rocRoller/KernelGraph/Utils.hpp>
 
+#include <map>
+#include <unordered_map>
+#include <unordered_set>
+
 namespace rocRoller::KernelGraph::NodeScheduling
 {
     ControlGraph::ControlGraph createSubGraph(KernelGraph const&      graph,
@@ -51,5 +55,4 @@ namespace rocRoller::KernelGraph::NodeScheduling
 
         return subGraph;
     }
-
 }

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

@@ -3,8 +3,8 @@
 
 #include <rocRoller/KernelGraph/Transforms/RemoveImplicitScheduling.hpp>
 
+#include <rocRoller/KernelGraph/ControlGraph/DataDependenceDAG.hpp>
 #include <rocRoller/KernelGraph/NodeSchedulingUtils.hpp>
-
 #include <rocRoller/KernelGraph/Utils.hpp>
 
 namespace rocRoller::KernelGraph
@@ -25,54 +25,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;
@@ -99,16 +51,26 @@ namespace rocRoller::KernelGraph
 
             Log::debug("Got connections.");
 
+            auto dependenceDAG = DataDependenceDAG::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())
+                        {
+                            graph.control.addElement(ControlGraph::Sequence(), {nodeA}, {nodeB});
+                        }
+                    }
+                    else if(seqEdge.has_value())
                     {
                         auto upstream = connectionsToKeep.at(nodeB);
                         auto order
@@ -118,8 +80,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
@@ -134,6 +96,7 @@ namespace rocRoller::KernelGraph
     {
         auto rv = original;
 
+        // grouped by immediate body-parent
         auto groupedMultiplyNodes = NodeScheduling::getGroupedNodes<ControlGraph::Multiply>(rv);
 
         for(auto& [parent, nodes] : groupedMultiplyNodes)

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

@@ -1123,7 +1123,8 @@ namespace rocRoller
                 else
                     break;
 
-                AssertFatal(graph.mapper.get<CT::Unroll>(tag) > 0,
+                AssertFatal(graph.mapper.get<CT::Unroll>(tag) > 0
+                                || graph.mapper.get<CT::ForLoop>(tag) > 0,
                             "SetCoordinate needs Unroll dimension");
             }
             return tag;

shared/rocroller/test/catch/RemoveImplicitSchedulingTest.cpp

@@ -139,7 +139,9 @@ namespace RemoveImplicitSchedulingTest
                     CHECK(order == graph.control.compareNodes(UpdateCache, node1, node2));
                 }
             }
-            CHECK(newlyUnorderedPairs > 0);
+
+            if(dataTypeAB == DataType::FP4)
+                CHECK(newlyUnorderedPairs > 0);
         }
     }