[rocroller] Build Basic-Block Dependence DAGs

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

@@ -0,0 +1,33 @@
+// Copyright Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <rocRoller/KernelGraph/ControlGraph/ControlGraph.hpp>
+
+namespace rocRoller::KernelGraph::DataDependenceDAG
+{
+    /**
+     * Builds the data dependence graph for the given kernel graph,
+     * specifically based on its control graph portion. It explicitly
+     * represents only the data dependences(flow, anti, and output) between
+     * the control graph nodes at the basic-block level.
+     *
+     * The data dependence graph has no cycles in it, therefore it
+     * is termed as data dependence DAG (directed acyclic graph).
+     * Since it involves multiple basic-blocks, the data dependence graph
+     * may not be a connected graph at the moment. In the future,
+     * we can make it as a connected graph by transforming it into a
+     * program dependence graph (which includes both data and control
+     * dependence information).
+     *
+     * It describes the dependency relationship between the control graph nodes,
+     * so it makes sense to use the existing `rocRoller::KernelGraph::ControlGraph`
+     * structure for its represenatation.
+     *
+     * The dependences are indicated using `Sequence` edges at each basic-block
+     * level i.e. between the nodes sharing the same immediate body-parent
+     * in `graph.control`.
+     */
+    ControlGraph::ControlGraph ConstructDataDependenceDAG(KernelGraph const& graph);
+}

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

@@ -0,0 +1,55 @@
+// 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:
+        /**
+         * Initializes the data dependence DAG structure (`m_dependenceDAG`)
+         * with only the control nodes from the given kernel graph.
+         */
+        DataDependenceDAGDetail(KernelGraph const& graph);
+
+        /**
+         * Returns the body-parent for the given node in the control graph.
+         */
+        int getBodyParent(int control);
+        /**
+         * Adds a dependence edge(represented via `Sequence`) between the given
+         * source and destination nodes in the dependence DAG (`m_dependenceDAG),
+         * if both the nodes have the same body-parent.
+         *
+         * Returns true if a dependence edge is added or exists between the
+         * given source and destination nodes, otherwise returns false.
+         */
+        bool addDependenceEdge(int source, int dest);
+        /**
+         * Makes necessary updates to `m_latestWriteToCoord`, `m_latestReadsToCoord`
+         * and `m_dependenceDAG` structures for the given `ReadWriteRecord`.
+         */
+        void processReadWriteRecord(ControlFlowRWTracer::ReadWriteRecord const& record);
+        /**
+         * Builds the data dependence DAG by populating the `m_dependenceDAG` structure
+         * with required dependence edges based on the trace generated using `ControlFlowRWTracer`
+         * for `m_graph`.
+         */
+        void constructDataDependenceDAG();
+        /**
+         * Returns the data dependence DAG (`m_dependenceDAG`).
+         */
+        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`.
          */

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

@@ -12,6 +12,8 @@ namespace rocRoller::KernelGraph::NodeScheduling
     {
         auto theNodes = graph.control.getNodes().filter(pred);
 
+        // TODO: Use key as pair<nodeID, ControlEdge> to indicate the associated edge
+        // type with the bodyParent.
         std::unordered_map<int, std::vector<int>> rv;
         for(auto node : theNodes)
         {
@@ -125,4 +127,4 @@ namespace rocRoller::KernelGraph::NodeScheduling
         Log::debug("Desired order: \n{}", fmt::join(desiredOrder, "\n"));
         Log::debug("Actual order: \n{}", fmt::join(nodes, "\n"));
     }
-}
+}

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

@@ -394,43 +394,8 @@ namespace rocRoller::KernelGraph
 
     void ControlFlowRWTracer::operator()(ForLoopOp const& op, int tag)
     {
-        //
-        // Don't examine for loop intialize or increment operations.
-        //
-        // Assign operations within loop initialisation operations
-        // are scoped already.
-        //
-        // Assign operations within loop increment operations
-        // typically involve: incrementing loop counters and
-        // offsets.  Loop counters are scoped already.
-        //
-        // Offsets are created by AssignIndexExpressions and are
-        // used in other nodes like LoadTiled.  These
-        // references do not explicitly appear in the graph.
-        //
-        // If we examine loop increment operations and "track" an
-        // offset increment, but don't track it during loads, then
-        // a Deallocate node would be mis-placed.
-        //
-        // A few solutions:
-        //
-        // 1. Don't examine loop increment operations.  They
-        // already appear in Scopes so are deallocated regardless.
-        // Fairly easy but perhaps we miss an opporunity to free
-        // up registers early.
-        //
-        // 2. Teach the tracker how to dig into all nodes.  Very
-        // tedious and not future-proof.
-        //
-        // 3. Expose all references in the graph.  Ideal but we
-        // aren't there yet.
-        //
-
-        // auto init = m_graph.control.getOutputNodeIndices<Initialize>(tag).to<std::set>();
-        // generate(init);
-
-        auto incr = m_graph.control.getOutputNodeIndices<ForLoopIncrement>(tag).to<std::set>();
-        generate(incr);
+        auto init = m_graph.control.getOutputNodeIndices<Initialize>(tag).to<std::set>();
+        generate(init);
 
         CollectDataFlowExpressionVisitor visitor;
         visitor.call(op.condition);
@@ -441,6 +406,9 @@ namespace rocRoller::KernelGraph
 
         auto body = m_graph.control.getOutputNodeIndices<Body>(tag).to<std::set>();
         generate(body);
+
+        auto incr = m_graph.control.getOutputNodeIndices<ForLoopIncrement>(tag).to<std::set>();
+        generate(incr);
     }
 
     void ControlFlowRWTracer::operator()(Kernel const& op, int tag)

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,173 @@
+// 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)
+            {
+                // 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));
+                }
+            }
+
+            ControlGraph::ControlGraph DataDependenceDAGDetail::getDataDependenceDAG()
+            {
+                return m_dependenceDAG;
+            }
+
+            void DataDependenceDAGDetail::constructDataDependenceDAG()
+            {
+                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();
+            }
+
+            bool DataDependenceDAGDetail::addDependenceEdge(int source, int dest)
+            {
+                AssertFatal(source != dest, ShowValue(source), ShowValue(dest));
+
+                if(getBodyParent(source) != getBodyParent(dest))
+                    return false;
+
+                if(!m_dependenceDAG.findEdge(source, dest).has_value())
+                {
+                    m_dependenceDAG.addElement(ControlGraph::Sequence(), {source}, {dest});
+                }
+
+                return true;
+            }
+
+            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) or WR(flow dep) edge
+                    auto depAdded = addDependenceEdge(writeIter->second, record.control);
+
+                    //TODO: Check the order and ensure that writeIter->second
+                    //      happens before record.control.
+                    //if(depAdded)
+                    //{
+                    //    auto order = m_graph.control.compareNodes(
+                    //        UseCacheIfAvailable, writeIter->second, record.control);
+                    //    AssertFatal(order == ControlGraph::NodeOrdering::LeftFirst
+                    //                    || order == ControlGraph::NodeOrdering::RightInBodyOfLeft,
+                    //                ShowValue(order),
+                    //                ShowValue(writeIter->second),
+                    //                ShowValue(record.control),
+                    //                ShowValue(record.coordinate),
+                    //                ShowValue(record.rw));
+                    //}
+                }
+
+                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
+                        auto depAdded = addDependenceEdge(readControl, record.control);
+
+                        //TODO: Check the order and ensure that writeIter->second
+                        //      happens before record.control.
+                        //if(depAdded)
+                        //{
+                        //    auto order = m_graph.control.compareNodes(
+                        //        UseCacheIfAvailable, readControl, record.control);
+                        //    AssertFatal(order == ControlGraph::NodeOrdering::LeftFirst
+                        //                    || order
+                        //                           == ControlGraph::NodeOrdering::RightInBodyOfLeft,
+                        //                ShowValue(order),
+                        //                ShowValue(readControl),
+                        //                ShowValue(record.control),
+                        //                ShowValue(record.coordinate),
+                        //                ShowValue(record.rw));
+                        //}
+                    }
+
+                    // 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);
+                }
+            }
+        }
+    }
+}