// Copyright Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

#include <rocRoller/KernelGraph/Transforms/AddPrefetch.hpp>
#include <rocRoller/KernelGraph/Transforms/AddPrefetch_detail.hpp>

#include <rocRoller/CommandSolution.hpp>
#include <rocRoller/Expression.hpp>
#include <rocRoller/KernelOptions_detail.hpp>

#include <rocRoller/KernelGraph/ControlGraph/ControlGraph.hpp>
#include <rocRoller/KernelGraph/ControlGraph/LastRWTracer.hpp>
#include <rocRoller/KernelGraph/ControlGraph/Operation.hpp>
#include <rocRoller/KernelGraph/ControlToCoordinateMapper.hpp>
#include <rocRoller/KernelGraph/CoordinateGraph/Dimension.hpp>
#include <rocRoller/KernelGraph/Transforms/Simplify.hpp>
#include <rocRoller/KernelGraph/Utils.hpp>
#include <rocRoller/KernelGraph/Visitors.hpp>
#include <rocRoller/Operations/Command.hpp>
#include <rocRoller/Operations/Operations.hpp>

/**
@class AddPrefetch
@brief Add prefetching for load operations.

# Overview

Prefetching is automatically applied to tiles loaded through LDS
within an unrolled ForLoop.

Prefetching is controlled by:

1. `bool prefetch` - Enable the prefetch transformation.

2. `int prefetchInFlight` - How many loads are put in-flight.

3. `bool prefetchMixMemOps` - Should global loads be inter-mixed with
    math operations, or isolated?

4. `int prefetchLDSFactor` - Ratio of how many sub-tiles are
   prefetched from LDS before the next unroll segment.

## Single prefetch

Focusing on how the two sets of buffers are operated on, we can view
the control flow of loads vertically as follows

    Buf0      Buf1
    ============== for loop preamble
     PG        PG
     CL
    -------------- barrier
     PL
    ============== unrolled for loop begin (counter += 2)
    -=-=-=-=-=-=-= unroll u=0 segment
     CV
     OPR
     PG        CL
    -------------- barrier
               PL
    -=-=-=-=-=-=-= unroll u=1 segment
               CV
               OPR
     CL        PG
    -------------- barrier
     PL
    ============== for loop end

where

- `PG` denotes _prefetch global_: issuing global to vgpr loads
- `CL` denotes _commit to lds_: vgpr to lds stores
- `PL` denotes _prefetch lds_: issuing lds to vgpr loads
- `CV` denotes _commit to vgpr_: waiting on lds to vgpr loads
- `OPR` denotes _operating on vgpr_: doing math (hopefully many cycles)

The order of these must be: `PG CL PL CV OPR`.

Note that:

1. Within the for-loop, there is always a barrier between a `CL` and a
   `PL`.

2. Global prefetches are not stalled by barriers.

3. Within a unrolled segment; `CV` and `OPR` can be mixed.  The
   rocRoller scheduler will make sure, eg, appropriate wait-counts are
   inserted to enforce dependencies.

4. The `PL` may be a "partial" prefetch.  In this case, the remaining
   loads from LDS into VGPRs happends in the immediately following
   `CV`.

## Two prefetches, unrolled

With two-prefetches, we obtain

    Buf0      Buf1      Buf2
    ======================== for loop preamble
     PG        PG        PG
     CL
    ------------------------ barrier
     PL
    ======================== unrolled for loop begin (counter += 3)
    -=-=-=-=-=-=-=-=-=-=-=-= unroll u=0 segment
     CV
     OPR
     PG        CL
    ------------------------ barrier
               PL
    -=-=-=-=-=-=-=-=-=-=-=-= unroll u=1 segment
               CV
               OPR
               PG        CL
    ------------------------ barrier
                         PL
    -=-=-=-=-=-=-=-=-=-=-=-= unroll u=2 segment
                         CV
                         OPR
     CL                  PG
    ------------------------ barrier
     PL
    ======================== for loop end

In the single pre-fetch mode, we have $U=2$ unrolled segments.  In the
two-prefetch mode, we have $U=3$ unrolled segments.  In both cases,
the scheduling of operations in unrolled segment $u$ are:

    -=-=-=-=-=-=-=-=-=-=-=-=-=- unroll u segment
    CV  [u]
    OPR [u]
    PG  [u]     CL [(u+1) % U]
    --------------------------- barrier
                PL [(u+1) % U]

*/

namespace rocRoller
{
    namespace KernelGraph
    {
        namespace AddPrefetchDetail
        {
            namespace CF = rocRoller::KernelGraph::ControlGraph;
            namespace CT = rocRoller::KernelGraph::CoordinateGraph;

            using GD = rocRoller::Graph::Direction;
            using namespace ControlGraph;
            using namespace CoordinateGraph;

            std::map<int, int> findPrefetch(KernelGraph const& kgraph)
            {
                std::map<int, int> rv;

                auto candidates = kgraph.control.getNodes<LoadTiled>();
                for(auto const& candidate : candidates)
                {
                    auto [user, direction] = getOperationTarget(candidate, kgraph);
                    if(!kgraph.coordinates.get<User>(user).has_value())
                        continue;
                    auto [required, path]   = findRequiredCoordinates(user, direction, kgraph);
                    auto forLoopCoordinates = filterCoordinates<ForLoop>(required, kgraph);
                    auto unrollCoordinates  = filterCoordinates<Unroll>(required, kgraph);

                    auto maybeForLoop = findContainingOperation<ForLoopOp>(candidate, kgraph);

                    if(maybeForLoop.has_value())
                    {
                        if(rv.contains(*maybeForLoop))
                            continue;

                        // TODO: Only do the K-Loop for now
                        auto fl = kgraph.control.get<ForLoopOp>(*maybeForLoop);
                        if(fl->loopName != rocRoller::KLOOP)
                            continue;

                        // Prefetching requires LDS for double-buffering. Check if there are
                        // any StoreLDSTile operations in the loop body. Without LDS stores,
                        // prefetching doesn't make sense (e.g., BufferToVGPR load path).
                        auto isBodyPredicate = kgraph.control.isElemType<Body>();
                        auto isStoreLDSTile  = kgraph.control.isElemType<StoreLDSTile>();
                        auto bodyEdges
                            = filter(isBodyPredicate,
                                     kgraph.control.getNeighbours<GD::Downstream>(*maybeForLoop))
                                  .to<std::vector>();
                        auto storeLDSTileNodes
                            = kgraph.control.findNodes(bodyEdges, isStoreLDSTile, GD::Downstream);
                        if(storeLDSTileNodes.empty())
                        {
                            Log::debug("KernelGraph::AddPrefetch(): ForLoop {} has no LDS stores, "
                                       "skipping prefetch",
                                       *maybeForLoop);
                            continue;
                        }

                        auto forLoopCoord     = getForLoopCoords(*maybeForLoop, kgraph).first;
                        auto maybeUnrollCoord = findUnrollNeighbour(kgraph, forLoopCoord);
                        if(forLoopCoordinates.contains(forLoopCoord)
                           && maybeUnrollCoord.has_value())
                        {
                            auto myUnroll
                                = getUnrollValueForOp(kgraph, *maybeUnrollCoord, candidate);

                            if(myUnroll > 0)
                            {
                                Dimension unroll
                                    = kgraph.coordinates.get<Unroll>(*maybeUnrollCoord).value();
                                auto unrollSize = getUnsignedInt(evaluate(getSize(unroll)));

                                Log::debug("KernelGraph::AddPrefetch(): ForLoop {} is a prefetch "
                                           "candidate: "
                                           "{} {} ({})",
                                           *maybeForLoop,
                                           *maybeUnrollCoord,
                                           unrollSize,
                                           candidate);

                                rv[*maybeForLoop] = unrollSize;
                            }
                        }
                    }
                }

                return rv;
            }

            std::optional<int> getLoadForExchange(int exchangeTag, KernelGraph const& graph)
            {
                auto isLoadPredicate = [&](int operation) -> bool {
                    auto maybeLoad = graph.control.get<LoadTiled>(operation);
                    return maybeLoad.has_value();
                };

                auto findConnections = [&](int coordinate) -> std::optional<int> {
                    for(auto c : graph.mapper.getCoordinateConnections(coordinate))
                        if(isLoadPredicate(c.control))
                            return c.control;
                    return {};
                };

                auto exchangeTileTag = graph.mapper.get<MacroTile>(exchangeTag);
                for(auto const edge :
                    graph.coordinates.getNeighbours<GD::Downstream>(exchangeTileTag))
                {
                    auto maybeIndex = graph.coordinates.get<Index>(edge);
                    if(!maybeIndex)
                        continue;
                    auto indexTileTag
                        = only(graph.coordinates.getNeighbours<GD::Downstream>(edge)).value();
                    auto tag = findConnections(indexTileTag);
                    if(tag.has_value())
                        return tag;
                }
                return {};
            }

            bool isLoadForExchange(int loadTag, KernelGraph const& graph)
            {
                if(not(graph.control.get<LoadTiled>(loadTag)
                       or graph.control.get<LoadLDSTile>(loadTag)))
                    return false;

                auto isExchangePredicate = [&](int operation) -> bool {
                    auto maybeExchange = graph.control.get<Exchange>(operation);
                    return maybeExchange.has_value();
                };

                auto checkConnections = [&](int coordinate) -> bool {
                    for(auto c : graph.mapper.getCoordinateConnections(coordinate))
                        if(isExchangePredicate(c.control))
                            return true;
                    return false;
                };

                auto tileTag = graph.mapper.get<MacroTile>(loadTag);
                if(checkConnections(tileTag))
                    return true;

                for(auto edge : graph.coordinates.getNeighbours<GD::Upstream>(tileTag))
                {
                    auto maybeIndex = graph.coordinates.get<Index>(edge);
                    if(!maybeIndex)
                        continue;
                    auto indexTileTag
                        = only(graph.coordinates.getNeighbours<GD::Upstream>(edge)).value();
                    if(checkConnections(indexTileTag))
                        return true;
                }
                return false;
            }

            bool IsDirectLoadToVGPR(KernelGraph const& k, int loadTag)
            {
                const auto maybeLoad = k.control.get<LoadTiled>(loadTag);
                if(!maybeLoad)
                    return false;

                auto macroTileTag = k.mapper.get<MacroTile>(loadTag);

                const auto maybeMacroTile = k.coordinates.get<MacroTile>(macroTileTag);
                if(!maybeMacroTile)
                    return false;
                const auto macroTile = maybeMacroTile.value();

                const auto top = getTopSetCoordinate(k, loadTag);

                auto isStoreLDSTilePredicate = k.control.isElemType<StoreLDSTile>();
                auto storeLDSTileNodes
                    = k.control.findNodes(top, isStoreLDSTilePredicate, GD::Downstream);
                const auto isLoadInLDSPathViaVGPR = std::any_of(
                    storeLDSTileNodes.begin(), storeLDSTileNodes.end(), [&](int storeTag) {
                        return macroTileTag == k.mapper.get<MacroTile>(storeTag);
                    });

                return !isLoadInLDSPathViaVGPR;
            }

            std::vector<int> FindTopSetCoordOfLoadLDSTileOpsInCurrentSegment(
                KernelGraph const& graph, int currentSegmentStart, int currentSegmentEnd)
            {
                auto isLoadLDSTile = graph.control.isElemType<LoadLDSTile>();

                auto edgePredicate = [&](int edgeTag) {
                    for(auto neighbour : graph.control.getNeighbours(edgeTag, GD::Downstream))
                    {
                        if(currentSegmentEnd == neighbour)
                            return false;
                    }
                    return true;
                };

                auto topSetCoordOfLoadLDSTileOps
                    = graph.control
                          .depthFirstVisit(currentSegmentStart, edgePredicate, GD::Downstream)
                          .filter(isLoadLDSTile)
                          .map([&graph](int tag) { return getTopSetCoordinate(graph, tag); })
                          .to<std::vector>();
                AssertFatal(
                    not topSetCoordOfLoadLDSTileOps.empty(),
                    " AddPrefetch: Expected at least one LoadLDSTile in the current segment");

                return topSetCoordOfLoadLDSTileOps;
            }
        }

        namespace CF = rocRoller::KernelGraph::ControlGraph;
        namespace CT = rocRoller::KernelGraph::CoordinateGraph;

        using GD = rocRoller::Graph::Direction;
        using namespace ControlGraph;
        using namespace CoordinateGraph;
        using namespace Expression;
        using namespace Register;
        using namespace AddPrefetchDetail;

        /**
         * Add prefetch transformer.
         */
        struct AddPrefetchVisitor
        {
            AddPrefetchVisitor(CommandParametersPtr params, ContextPtr context)
                : m_params(params)
                , m_context(context)
            {
            }

            void commitForLoop(KernelGraph&                 graph,
                               NaryArgumentColouring const& colouring,
                               int                          forLoop,
                               int                          numUnroll);
            void orderLoadsBeforeMultiplies(KernelGraph& graph, int forLoop, int numUnroll);

            void stage(KernelGraph const& graph);
            void commit(KernelGraph&, NaryArgumentColouring const&);

        private:
            // Keys are: ForLoop tag, Unroll value/segment, LDS tag
            std::map<int, std::map<int, std::map<int, LDSOperationInfo>>> m_info;

            std::map<int, int>                                    m_scopes;
            std::map<int, int>                                    m_prefetchLoops;
            std::map<int, std::map<int, std::unordered_set<int>>> m_prefetchUnrollBodyStarts;
            std::map<int, std::map<int, std::unordered_set<int>>> m_prefetchUnrollBodyEnds;
            std::map<int, std::map<int, std::set<std::tuple<int, int, int>>>>
                                                           m_prefetchFromLDSChains;
            std::map<int, std::map<int, std::vector<int>>> m_loadFromLDSChains;
            std::map<int, std::unordered_set<int>>         m_prefetchDelete;
            std::map<int, int>                             m_exchangeSegment;

            std::map<int, std::map<int, std::vector<int>>> m_deferredToOrder;

            struct LoadTiledInfoComparator
            {
                bool operator()(const LoadTiledInfo& a, const LoadTiledInfo& b) const
                {
                    return a.top < b.top;
                }
            };

            // Global Memory To VGPR loads not part of a to LDS path nor loads for exchange operations
            // ForLoop -> Unroll value -> list of {topSetCoordinate, loadTag}
            std::map<int, std::map<int, std::set<LoadTiledInfo, LoadTiledInfoComparator>>>
                m_directLoads;

            CommandParametersPtr m_params;
            ContextPtr           m_context;
        };

        KernelGraph AddPrefetch::apply(KernelGraph const& original)
        {
            auto graph = original;
            removeRedundantBodyEdges(graph);
            removeRedundantNOPs(graph);

            auto colouring = colourByNaryArgument(graph);

            if(m_params->prefetch)
            {
                auto visitor = AddPrefetchVisitor(m_params, m_context);
                AssertFatal(m_params->unrollK > 1, "KLoop must be unrolled when prefetching.");
                visitor.stage(graph);
                visitor.commit(graph, colouring);
            }

            removeRedundantSequenceEdges(graph);

            return graph;
        }

        void AddPrefetchVisitor::commit(KernelGraph& k, NaryArgumentColouring const& colouring)
        {
            Log::debug("KernelGraph::AddPrefetch()::commit()");

            for(auto [forLoop, numUnroll] : m_prefetchLoops)
            {
                commitForLoop(k, colouring, forLoop, numUnroll);
            }

            removeRedundantSequenceEdges(k);
            removeRedundantBodyEdges(k);
        }

        /**
        * @brief Order loads before Multiplies
        */
        void AddPrefetchVisitor::orderLoadsBeforeMultiplies(KernelGraph& graph, int forLoop, int u)
        {
            auto starts = m_prefetchUnrollBodyStarts[forLoop][u];

            auto isLoadPredicate = [&graph](int x) {
                return graph.control.get<LoadTiled>(x).has_value()
                       || graph.control.get<LoadLDSTile>(x).has_value();
            };

            auto isMultiplyPredicate = graph.control.isElemType<Multiply>();

            auto isLHSRHSPredicate = [](Connections::ConnectionSpec const& spec) -> bool {
                auto nary = std::visit(
                    rocRoller::overloaded{[](Connections::JustNaryArgument arg) {
                                              return std::optional<NaryArgument>{arg.argument};
                                          },
                                          [](Connections::TypeAndNaryArgument arg) {
                                              return std::optional<NaryArgument>{arg.argument};
                                          },
                                          [](auto x) { return std::optional<NaryArgument>{}; }},
                    spec);
                if(!nary.has_value())
                    return false;
                return *nary == NaryArgument::LHS || *nary == NaryArgument::RHS
                       || *nary == NaryArgument::LHS_SCALE || *nary == NaryArgument::RHS_SCALE;
            };

            std::map<int, int> loadMap{};
            for(auto loadTag : graph.control.findNodes(starts, isLoadPredicate))
            {
                auto tileTag     = graph.mapper.get<MacroTile>(loadTag);
                loadMap[tileTag] = getTopSetCoordinate(graph, loadTag);
            }

            auto isExchangePredicate = graph.control.isElemType<Exchange>();

            for(auto exchangeTag : graph.control.findNodes(starts, isExchangePredicate))
            {
                for(auto conn : graph.mapper.getConnections(exchangeTag))
                {
                    auto coord = only(graph.coordinates.getOutputNodeIndices(conn.coordinate,
                                                                             CT::isEdge<CT::Index>))
                                     .value_or(conn.coordinate);
                    if(not loadMap.contains(coord))
                        continue;

                    Log::debug("Adding load-before-exchange Sequence edge from {} to {} for {}",
                               loadMap[coord],
                               exchangeTag,
                               toString(conn.connection));

                    graph.control.addElement(Sequence(), {loadMap[coord]}, {exchangeTag});
                }
                m_prefetchUnrollBodyStarts[forLoop][u].erase(exchangeTag);
            }

            for(auto exchangeTag : graph.control.findNodes(starts, isExchangePredicate))
            {
                auto destTileTag = graph.mapper.get(exchangeTag, NaryArgument::DEST);

                auto pred = m_context->kernelOptions()->scaleSkipPermlane
                                    != rocRoller::ScaleSkipPermlaneMode::None
                                ? CT::isEdge<Segment>
                                : CT::isEdge<Index>;

                auto tileTags
                    = graph.coordinates.getInputNodeIndices(destTileTag, pred).to<std::vector>();
                AssertFatal(!tileTags.empty(), "swizzle indexed tiles not found");
                for(auto tileTag : tileTags)
                    loadMap[tileTag] = getTopSetCoordinate(graph, exchangeTag);
            }

            for(auto multiplyTag : graph.control.findNodes(starts, isMultiplyPredicate))
            {
                for(auto conn : graph.mapper.getConnections(multiplyTag))
                {
                    // Multiply arguments may have been loaded by LDS
                    // prefetching, or may be single-scales from a
                    // LoadSGPR operations.
                    //
                    // These arguments won't be in the loadMap, and we
                    // don't need to order them here.
                    if(not loadMap.contains(conn.coordinate))
                        continue;
                    if(not isLHSRHSPredicate(conn.connection))
                        continue;

                    Log::debug("Adding load-before-multiply Sequence edge from {} to {} for {}",
                               loadMap[conn.coordinate],
                               multiplyTag,
                               toString(conn.connection));

                    graph.control.addElement(Sequence(), {loadMap[conn.coordinate]}, {multiplyTag});
                }
                m_prefetchUnrollBodyStarts[forLoop][u].erase(multiplyTag);
            }
        }

        void AddPrefetchVisitor::commitForLoop(KernelGraph&                 graph,
                                               NaryArgumentColouring const& colouring,
                                               int                          forLoop,
                                               int                          numUnroll)
        {
            auto logger = rocRoller::Log::getLogger();
            logger->debug("KernelGraph::AddPrefetch()::commitForLoop({})", forLoop);

            AssertFatal(isOperation<ForLoopOp>(graph.control.getElement(forLoop)));

            auto forLoopCoord = getForLoopCoords(forLoop, graph).first;
            auto unrollCoord  = findUnrollNeighbour(graph, forLoopCoord).value();

            std::vector<NaryArgument> argumentOrder = {NaryArgument::RHS,
                                                       NaryArgument::RHS_SCALE,
                                                       NaryArgument::LHS,
                                                       NaryArgument::LHS_SCALE};

            auto sortBy = [&](auto& container, auto const& order, auto getter) {
                std::stable_sort(
                    container.begin(), container.end(), [&](const auto& a, const auto& b) {
                        auto itA = std::find(
                            order.begin(), order.end(), colouring.coordinateColour.at(getter(a)));
                        AssertFatal(itA != order.end(), ShowValue(getter(a)));
                        auto itB = std::find(
                            order.begin(), order.end(), colouring.coordinateColour.at(getter(b)));
                        AssertFatal(itB != order.end(), ShowValue(getter(b)));
                        return itA < itB;
                    });
            };

            //
            // Delete connecting edges
            //
            for(auto tag : m_prefetchDelete[forLoop])
            {
                if(graph.control.exists(tag))
                    graph.control.deleteElement(tag);
            }

            // At this point, each of the unrolled loop bodies are
            // detached and isolated from the rest of the graph.

            std::map<int, std::vector<LDSOperationInfo>> loadsByUnroll;
            for(int u = 0; u < numUnroll; u++)
            {
                for(auto [target, info] : m_info[forLoop][u])
                    loadsByUnroll[u].push_back(info);
                sortBy(loadsByUnroll[u], argumentOrder, [](const auto& info) { return info.user; });
            }

            AssertFatal(loadsByUnroll.size() == numUnroll,
                        ShowValue(loadsByUnroll.size()),
                        ShowValue(numUnroll));

            std::map<int, std::vector<LoadTiledInfo>> directLoadsByUnroll;
            for(int u = 0; u < numUnroll; u++)
            {
                for(auto load : m_directLoads[forLoop][u])
                    directLoadsByUnroll[u].push_back(load);
                // TODO: Order by argument type?
            }

            const auto numDirectLoads = directLoadsByUnroll.size();
            AssertFatal(numDirectLoads == 0 or numDirectLoads == numUnroll,
                        ShowValue(numDirectLoads),
                        ShowValue(numUnroll));

            //
            // Add Scope above the ForLoop
            //
            if(!m_scopes.contains(forLoop))
            {
                m_scopes[forLoop]
                    = replaceWith(graph, forLoop, graph.control.addElement(Scope()), false);
            }
            auto scope = m_scopes[forLoop];

            //
            // Prefetch before ForLoop
            //
            auto preNOP = graph.control.addElement(NOP());
            graph.control.addElement(Sequence(), {preNOP}, {forLoop});

            std::vector<int> preChain;

            int numInFlight = m_params->prefetchInFlight;

            // Loads first
            for(int u = 0; u < numInFlight; ++u)
            {
                // Direct loads before global prefetch loads in pre-loop
                if(directLoadsByUnroll.contains(u))
                {
                    for(auto& load : directLoadsByUnroll[u])
                    {
                        logger->debug("  prefetch: pre-loop direct load: unroll {} load {} "
                                      "of user {} with top {}.",
                                      u,
                                      load.tag,
                                      load.user,
                                      load.top);
                        auto loadChain = duplicateChain(graph, {load.top});
                        preChain.push_back(loadChain);
                    }
                }
                for(auto load : loadsByUnroll[u])
                {
                    logger->debug(
                        "  prefetch: pre-loop global load: unroll {} user {}", u, load.user);
                    auto loadChain = duplicateChain(graph, {load.globalChain});
                    preChain.push_back(loadChain);
                }
            }

            // StoreLDS next
            auto storeLDScounter   = 0;
            auto direct2LDSCounter = 0;
            for(auto load : loadsByUnroll[0])
            {
                logger->debug("  prefetch: pre-loop commit lds: unroll {} user {}", 0, load.user);
                auto storeChain = duplicateChain(graph, {load.ldsChain});
                preChain.push_back(storeChain);

                storeLDScounter++;

                auto op = std::get<Operation>(graph.control.getElement(load.globalOperation));
                if(std::holds_alternative<LoadTiled>(op))
                {
                    auto macroTile = graph.coordinates.get<MacroTile>(
                        graph.mapper.get<MacroTile>(load.globalOperation));
                    if(macroTile && macroTile->memoryType == MemoryType::WAVE_Direct2LDS)
                    {
                        direct2LDSCounter++;
                    }
                }
            }

            auto prefetchDirect2LDS
                = ((direct2LDSCounter > 0) && (storeLDScounter == direct2LDSCounter));

            graph.control.addElement(Body(), {scope}, {preChain[0]});
            for(uint i = 1; i < preChain.size(); ++i)
            {
                graph.control.addElement(Sequence(), {preChain[i - 1]}, {preChain[i]});
            }

            auto addLDSPrefetchChains = [&](int u, int pre, int post, bool duplicate) -> int {
                std::vector<int> prefetchChain;

                std::map<std::tuple<NaryArgument, int>, std::vector<int>> ldsByArgAndSmallK;
                std::map<NaryArgument, std::set<int>>                     ldsPrefetchValues;
                for(auto [target, smallk, chain] : m_prefetchFromLDSChains[forLoop][u])
                {
                    int dchain = duplicate ? duplicateChain(graph, {chain}) : chain;
                    AssertFatal(colouring.operationColour.contains(chain), ShowValue(chain));
                    auto argument = colouring.operationColour.at(chain);
                    ldsByArgAndSmallK[{argument, smallk}].push_back(dchain);
                    ldsPrefetchValues[argument].insert(smallk);
                }

                /*
                 */
                auto mixDataAndScale = [&](auto dataArgument, auto scaleArgument) {
                    // Count unique lds-prefetch unroll (small-k) values
                    int numDataPrefetchValues  = ldsPrefetchValues[dataArgument].size();
                    int numScalePrefetchValues = ldsPrefetchValues[scaleArgument].size();

                    // Get number of data/scale chains (grows with number of jammed wavetiles)
                    size_t numDataChains = 0;
                    if(numDataPrefetchValues > 0)
                    {
                        int minDataK  = *ldsPrefetchValues[dataArgument].begin();
                        numDataChains = ldsByArgAndSmallK.at({dataArgument, minDataK}).size();
                    }

                    size_t numScaleChains = 0;
                    if(numScalePrefetchValues > 0)
                    {
                        int minScaleK  = *ldsPrefetchValues[scaleArgument].begin();
                        numScaleChains = ldsByArgAndSmallK.at({scaleArgument, minScaleK}).size();
                    }

                    if(numScaleChains > 0)
                    {
                        auto numDataPerScale = numDataChains / numScaleChains;

                        for(int idxScale = 0; idxScale < numScaleChains; ++idxScale)
                        {
                            for(int idxData = 0; idxData < numDataPerScale; ++idxData)
                            {
                                for(int k = 0; k < numDataPrefetchValues; ++k)
                                    prefetchChain.push_back(
                                        ldsByArgAndSmallK.at({dataArgument, k})
                                            .at(idxScale * numDataPerScale + idxData));
                            }

                            for(int k = 0; k < numScalePrefetchValues; ++k)
                                prefetchChain.push_back(
                                    ldsByArgAndSmallK.at({scaleArgument, k}).at(idxScale));
                        }
                    }
                    else
                    {
                        for(int idxData = 0; idxData < numDataChains; ++idxData)
                        {
                            for(int k = 0; k < numDataPrefetchValues; ++k)
                                prefetchChain.push_back(
                                    ldsByArgAndSmallK.at({dataArgument, k}).at(idxData));
                        }
                    }
                };
                mixDataAndScale(NaryArgument::RHS, NaryArgument::RHS_SCALE);
                mixDataAndScale(NaryArgument::LHS, NaryArgument::LHS_SCALE);

                AssertFatal(!prefetchChain.empty());

                logger->debug("  prefetch: lds prefetch: ordering {} to {} (top)",
                              pre,
                              prefetchChain.front());
                graph.control.addElement(Sequence(), {pre}, {prefetchChain.front()});
                for(uint i = 1; i < prefetchChain.size(); ++i)
                {
                    logger->debug("  prefetch: lds prefetch: ordering {} to {} (chain)",
                                  prefetchChain[i - 1],
                                  prefetchChain[i]);
                    graph.control.addElement(
                        Sequence(), {prefetchChain[i - 1]}, {prefetchChain[i]});
                }
                logger->debug("  prefetch: lds prefetch: ordering {} to {} (bottom)",
                              prefetchChain.back(),
                              post);
                graph.control.addElement(Sequence(), {prefetchChain.back()}, {post});

                return prefetchChain.front();
            };

            if(!m_prefetchFromLDSChains[forLoop].empty())
                addLDSPrefetchChains(0, preChain.back(), preNOP, true);
            else
                graph.control.addElement(Sequence(), {preChain.back()}, {preNOP});

            //
            // ForLoop body
            //

            // Update SetCoordinates for LoadTile operations
            for(uint u = 0; u < numUnroll; ++u)
            {
                auto prefetchGlobalU   = (u + numInFlight) % numUnroll;
                auto prefetchCoordExpr = literal(u + numInFlight);

                for(auto load : loadsByUnroll[prefetchGlobalU])
                {
                    logger->debug("  prefetch: in-loop: set coordinate: load {} user {} expr {}",
                                  prefetchGlobalU,
                                  load.user,
                                  toString(prefetchCoordExpr));

                    auto setPrefetchCoord = SetCoordinate(prefetchCoordExpr);

                    auto maybeSetCoordinate = graph.control.get<SetCoordinate>(load.globalChain);
                    auto loadUnrollCoord    = graph.mapper.get<Unroll>(load.globalChain);
                    if(maybeSetCoordinate && loadUnrollCoord == unrollCoord)
                    {
                        graph.control.setElement(load.globalChain, setPrefetchCoord);
                    }
                    else
                    {
                        Throw<FatalError>("Mismatched SetCoordinate node above LoadTile.");
                    }
                }
            }

            // Update SetCoordinates for direct loads operations
            for(uint u = 0; u < numUnroll; ++u)
            {
                auto prefetchGlobalU = (u + numInFlight) % numUnroll;
                for(auto load : m_directLoads[forLoop][prefetchGlobalU])
                {
                    // for each SetCoordinate in chain
                    for(auto setCoordTag : graph.control.findNodes(
                            load.top, graph.control.isElemType<SetCoordinate>(), GD::Downstream))
                    {
                        auto setCoordOp = graph.control.getNode<SetCoordinate>(setCoordTag);
                        auto setCoordUnrollCoord = graph.mapper.get<Unroll>(setCoordTag);

                        if(setCoordUnrollCoord == unrollCoord)
                        {
                            auto prefetchCoordExpr = literal(u + numInFlight);

                            logger->debug(
                                "  prefetch: in-loop: set coordinate: load {} user {} expr {}",
                                u,
                                load.user,
                                toString(prefetchCoordExpr));

                            graph.control.setElement(setCoordTag, SetCoordinate(prefetchCoordExpr));
                        }
                    }
                }
            }

            // Build Unroll segment boundaries
            std::vector<int> segmentBoundaries = {forLoop};
            for(uint u = 0; u < numUnroll; ++u)
                segmentBoundaries.push_back(graph.control.addElement(NOP()));

            auto separateMemOps = !m_params->prefetchMixMemOps;

            // Unrolled loop over prefetch segments
            for(uint u = 0; u < numUnroll; ++u)
            {
                logger->debug("  prefetch: in-loop: segment {}", u);

                // Connect the segment to the preceding segment boundary
                //
                // Note that the first boundary is the forLoop, and
                // the remaining boundaries are NOPs.  Therefore
                // segmentBoundaries[u] is the "preceding" boundary.
                for(auto tag : m_prefetchUnrollBodyStarts[forLoop][u])
                {
                    if(u == 0)
                        graph.control.addElement(Body(), {segmentBoundaries[u]}, {tag});
                    else
                    {
                        auto descOfSegmentStart
                            = graph.control
                                  .depthFirstVisit(
                                      tag, graph.control.isElemType<Sequence>(), GD::Downstream)
                                  .to<std::set>();

                        if(!descOfSegmentStart.contains(segmentBoundaries[u]))
                        {
                            graph.control.addElement(Sequence(), {segmentBoundaries[u]}, {tag});
                        }
                    }
                }

                auto toOrder = filter(graph.control.isElemType<LoadLDSTile>(),
                                      graph.control.depthFirstVisit(segmentBoundaries[u],
                                                                    Graph::Direction::Downstream))
                                   .to<std::vector>();
                orderMemoryNodes(graph, toOrder, false);

                auto globalPrefetchU = (u + numInFlight) % numUnroll;
                auto ldsPrefetchU    = (u + 1) % numUnroll;

                auto nop = separateMemOps ? graph.control.addElement(NOP()) : -1;

                auto previousBoundaryStart = separateMemOps ? nop : segmentBoundaries[u];

                // Order direct loads before global prefetch loads within this segment
                // and chain them sequentially
                if(directLoadsByUnroll.contains(globalPrefetchU))
                {
                    auto        directLoadPrefetchU = (u + numInFlight) % numUnroll;
                    const auto& directLoads         = directLoadsByUnroll[directLoadPrefetchU];

                    const auto firstLoadTopSetCoord = directLoads[0].top;

                    if(separateMemOps)
                    {
                        graph.control.addElement(
                            Sequence(), {previousBoundaryStart}, {firstLoadTopSetCoord});
                    }
                    else
                    {
                        auto edge = (u == 0) ? ControlEdge{Body()} : ControlEdge{Sequence()};
                        graph.control.addElement(
                            edge, {previousBoundaryStart}, {firstLoadTopSetCoord});
                    }

                    for(size_t i = 1; i < directLoads.size(); ++i)
                    {
                        graph.control.addElement(
                            Sequence(), {directLoads[i - 1].top}, {directLoads[i].top});
                    }

                    previousBoundaryStart = directLoads[directLoads.size() - 1].top;
                }

                // Issue global loads
                auto globalLoads = loadsByUnroll[globalPrefetchU];
                logger->debug("  prefetch: in-loop: issue global loads {}",
                              globalLoads[0].globalChain);
                if(separateMemOps)
                {
                    graph.control.addElement(
                        Sequence(), {previousBoundaryStart}, {globalLoads[0].globalChain});
                }
                else
                {
                    auto edge = (u == 0) ? ControlEdge{Body()} : ControlEdge{Sequence()};
                    graph.control.addElement(
                        edge, {previousBoundaryStart}, {globalLoads[0].globalChain});
                }

                logger->debug("  prefetch: in-loop: global load {} user {}",
                              globalPrefetchU,
                              globalLoads[0].user);

                for(int i = 1; i < globalLoads.size(); i++)
                {
                    graph.control.addElement(
                        Sequence(), {globalLoads[i - 1].globalChain}, {globalLoads[i].globalChain});

                    logger->debug("  prefetch: in-loop: global load {} user {}",
                                  globalPrefetchU,
                                  globalLoads[i].user);
                }

                if(separateMemOps)
                {
                    graph.control.addElement(Sequence(),
                                             {globalLoads[globalLoads.size() - 1].globalChain},
                                             {segmentBoundaries[u + 1]});
                }

                // Commit in-flight to LDS
                auto globalStores = loadsByUnroll[ldsPrefetchU];

                logger->debug("  prefetch: in-loop: commit lds {} user {}",
                              ldsPrefetchU,
                              globalStores[0].user);

                for(int i = 1; i < globalStores.size(); i++)
                {
                    graph.control.addElement(
                        Sequence(), {globalStores[i - 1].ldsChain}, {globalStores[i].ldsChain});

                    logger->debug("  prefetch: in-loop: connecting: {} to {}",
                                  globalStores[i - 1].ldsChain,
                                  globalStores[i].ldsChain);

                    logger->debug("  prefetch: in-loop: commit lds {} user {}",
                                  ldsPrefetchU,
                                  globalStores[i].user);
                }

                logger->debug("  prefetch: in-loop: ordering {} to {}",
                              globalLoads[globalLoads.size() - 1].globalChain,
                              globalStores[0].ldsChain);
                graph.control.addElement(Sequence(),
                                         {globalLoads[globalLoads.size() - 1].globalChain},
                                         {globalStores[0].ldsChain});

                // Prefetch from LDS
                int firstPrefetchFromLDS = -1;
                if(m_prefetchFromLDSChains[forLoop].contains(ldsPrefetchU))
                {
                    auto predecessorOfLoadLDSTileChain
                        = globalStores[globalStores.size() - 1].ldsChain;

                    // When prefetching Direct2LDS ops and mixing of memory operations is enabled,
                    // Direct2LDS and LoadLDSTile ops can be made parallel by connecting LoadLDSTile
                    // chains to the current segment boundary. Otherwise, LoadLDSTile operations
                    // are sequenced after Direct2LDS/StoreLDSTile.
                    if(prefetchDirect2LDS && m_params->prefetchMixMemOps)
                    {
                        auto predecessorBySequenceEdge
                            = only(graph.control.getInputNodeIndices<Sequence>(
                                globalLoads[0].globalChain));
                        auto maybeForLoop = only(
                            graph.control.getInputNodeIndices<Body>(globalLoads[0].globalChain));
                        AssertFatal(
                            maybeForLoop || predecessorBySequenceEdge,
                            "Expected prefetch global load to be connected to either the for "
                            "loop body or by a single sequence edge to a node, that is either the "
                            "segment boundary nop or a direct global load.");

                        predecessorOfLoadLDSTileChain = segmentBoundaries[u];
                        // In segment 0, make LoadLDSTile parallel with
                        // ForLoop --Body->  NOP --Sequence-> LDSPrefetchChain.
                        if(u == 0)
                        {
                            const auto nop = graph.control.addElement(NOP());
                            graph.control.addElement(Body(), {segmentBoundaries[u]}, {nop});
                            predecessorOfLoadLDSTileChain = nop;
                        }

                        graph.control.addElement(Sequence(),
                                                 {globalStores[globalStores.size() - 1].ldsChain},
                                                 {segmentBoundaries[u + 1]});
                    }

                    firstPrefetchFromLDS = addLDSPrefetchChains(ldsPrefetchU,
                                                                predecessorOfLoadLDSTileChain,
                                                                segmentBoundaries[u + 1],
                                                                false);
                }
                else
                {
                    graph.control.addElement(Sequence(),
                                             {globalStores[globalStores.size() - 1].ldsChain},
                                             {segmentBoundaries[u + 1]});
                }

                // To ensure proper memory ordering, the last
                // load-from-lds chain of the current segment must
                // finish before the first prefetch-from-lds chain
                // starts (which prefetchs for the next segment; but
                // appears in the current segment)
                if(firstPrefetchFromLDS != -1)
                {
                    int lastLoadFromLDS = -1;
                    for(auto tagA : m_loadFromLDSChains[forLoop][u])
                    {
                        bool isLast = true;
                        for(auto tagB : m_loadFromLDSChains[forLoop][u])
                        {
                            if(tagA == tagB)
                                continue;
                            if(graph.control.compareNodes(rocRoller::UpdateCache, tagA, tagB)
                               == NodeOrdering::LeftFirst)
                            {
                                isLast = false;
                                break;
                            }
                        }
                        if(isLast)
                        {
                            lastLoadFromLDS = tagA;
                            break;
                        }
                    }

                    if(lastLoadFromLDS != -1)
                    {
                        Log::debug("  prefetch: in-loop: lds ordering {} to {}",
                                   lastLoadFromLDS,
                                   firstPrefetchFromLDS);
                        graph.control.addElement(
                            Sequence(), {lastLoadFromLDS}, {firstPrefetchFromLDS});
                    }
                }

                orderLoadsBeforeMultiplies(graph, forLoop, u);

                auto loads
                    = filter(graph.control.isElemType<LoadTiled>(),
                             graph.control.depthFirstVisit(m_prefetchUnrollBodyStarts[forLoop][u],
                                                           Graph::Direction::Downstream))
                          .to<std::vector>();
                toOrder.clear();
                std::transform(std::begin(loads),
                               std::end(loads),
                               std::back_inserter(toOrder),
                               [&graph](int tag) { return getTopSetCoordinate(graph, tag); });
                orderMemoryNodes(graph, loads, false);

                // Connect the segment to the proceeding segment boundary
                if(separateMemOps)
                {
                    for(auto tag : m_prefetchUnrollBodyEnds[forLoop][u])
                        graph.control.addElement(Sequence(), {tag}, {nop});
                }
                else
                {
                    for(auto tag : m_prefetchUnrollBodyEnds[forLoop][u])
                        graph.control.addElement(Sequence(), {tag}, {segmentBoundaries[u + 1]});
                }

                // When mixing memory operations and not prefetching LoadLDSTile ops,
                // ensure that all LoadLDSTile operations in the current are scheduled
                // before any StoreLDSTile part of prefetch load-store chains that
                // are not Direct2LDS.
                if(m_params->prefetchMixMemOps && m_params->prefetchLDSFactor == 0)
                {
                    const auto currentSegmentStart = segmentBoundaries[u];
                    const auto currentSegmentEnd   = segmentBoundaries[u + 1];

                    std::vector<int> nonDirect2LDSStoreLDSTileTops;
                    for(const auto storeLDSTileOp : globalStores)
                    {
                        auto macroTile = graph.coordinates.get<MacroTile>(
                            graph.mapper.get<MacroTile>(storeLDSTileOp.globalOperation));
                        AssertFatal(macroTile,
                                    "Expected a MacroTile coordinate for StoreLDSTile operation.");

                        if(macroTile->memoryType != MemoryType::WAVE_Direct2LDS)
                        {
                            nonDirect2LDSStoreLDSTileTops.push_back(storeLDSTileOp.ldsChain);
                        }
                    }

                    for(const auto topSetCoordOfLoadLDSTile :
                        FindTopSetCoordOfLoadLDSTileOpsInCurrentSegment(
                            graph, currentSegmentStart, currentSegmentEnd))
                    {
                        for(const auto storeLDSTileOp : nonDirect2LDSStoreLDSTileTops)
                        {
                            graph.control.addElement(
                                Sequence(), {topSetCoordOfLoadLDSTile}, {storeLDSTileOp});
                        }
                    }
                }
            }

            //
            // Make exchange scale loads happen first!
            //
            {
                auto isExchangePredicate = graph.control.isElemType<Exchange>();

                auto bodies = graph.control.getOutputNodeIndices<Body>(forLoop).to<std::vector>();
                auto exchanges
                    = graph.control.findNodes(bodies, isExchangePredicate).to<std::vector>();

                std::map<int, int> scaleLoadU;

                for(auto const exchangeTag : exchanges)
                {
                    auto prefetchGlobalU
                        = (m_exchangeSegment[exchangeTag] + numInFlight) % numUnroll;

                    auto loadTag = getLoadForExchange(exchangeTag, graph);

                    // When loading pre-swizzled scales from LDS, no
                    // LoadTiled node exists.
                    if(!loadTag)
                    {
                        Log::debug("No matching load operation found for Exchange({}); assuming it "
                                   "is pre-swizzled from LDS.",
                                   exchangeTag);
                        continue;
                    }

                    auto const search = scaleLoadU.find(loadTag.value());
                    if(search == scaleLoadU.end() || search->second > prefetchGlobalU)
                        scaleLoadU[loadTag.value()] = prefetchGlobalU;
                }

                for(auto const [loadTag, u] : scaleLoadU)
                {
                    std::unordered_set<int> orderBeforeTags;
                    for(auto const info : loadsByUnroll[u])
                        orderBeforeTags.insert(info.globalChain);

                    auto topOp = getTopSetCoordinate(graph, loadTag);
                    for(auto const orderBeforeTag : orderBeforeTags)
                        graph.control.addElement(Sequence(), {topOp}, {orderBeforeTag});
                }
            }
        }

        void updateExchangeColouring(std::map<int, int>&    operationUnroll,
                                     KernelGraph const&     graph,
                                     UnrollColouring const& colouring,
                                     int                    forLoop,
                                     int                    unrollCoord)
        {
            auto isMultiplyPredicate = graph.control.isElemType<Multiply>();

            auto bodies = graph.control.getOutputNodeIndices<Body>(forLoop).to<std::vector>();
            auto multiplyTags
                = graph.control.findNodes(bodies, isMultiplyPredicate).to<std::unordered_set>();

            for(auto multiplyTag : multiplyTags)
            {
                auto lhsExchange
                    = getExchangeForMultiply(graph, multiplyTag, NaryArgument::LHS_SCALE);
                auto rhsExchange
                    = getExchangeForMultiply(graph, multiplyTag, NaryArgument::RHS_SCALE);

                if(lhsExchange)
                    operationUnroll[*lhsExchange] = operationUnroll[multiplyTag];
                if(rhsExchange)
                    operationUnroll[*rhsExchange] = operationUnroll[multiplyTag];
            }
        }

        void AddPrefetchVisitor::stage(KernelGraph const& k)
        {
            m_prefetchLoops = findPrefetch(k);

            auto colouring       = colourByUnrollValue(k);
            auto isBodyPredicate = k.control.isElemType<Body>();

            std::map<int, int> unrollCoordSizes;
            {
                for(auto unrollTag : k.coordinates.getNodes<Unroll>())
                {
                    auto unroll                 = k.coordinates.get<Unroll>(unrollTag).value();
                    unrollCoordSizes[unrollTag] = getUnsignedInt(evaluate(unroll.size));
                }
            }

            // Map: Operation (in loop body) to Unroll coordinate value
            std::map<int, int> operationUnroll;

            for(auto [forLoop, numUnroll] : m_prefetchLoops)
            {
                auto forLoopCoord     = getForLoopCoords(forLoop, k).first;
                auto maybeUnrollCoord = findUnrollNeighbour(k, forLoopCoord);
                AssertFatal(maybeUnrollCoord, "Prefetch with no unroll coordinate.");
                auto unrollCoord = *maybeUnrollCoord;

                Log::debug("AddPrefetch::stage: Unroll coordinate is {}", unrollCoord);

                for(auto [opTag, opColours] : colouring.operationColour)
                {
                    if(opColours.contains(unrollCoord))
                        operationUnroll[opTag] = opColours[unrollCoord];
                }

                updateExchangeColouring(operationUnroll, k, colouring, forLoop, unrollCoord);
            }

            auto alreadySeen = std::unordered_set<int>();

            //
            // Find global loads, LDS stores, and remaining (which
            // will be top of unroll bodies)
            //
            for(auto [forLoop, numUnroll] : m_prefetchLoops)
            {
                auto bodyEdges
                    = filter(isBodyPredicate, k.control.getNeighbours<GD::Downstream>(forLoop))
                          .to<std::vector>();

                // Find LDS stores and mark them for detachment
                auto isStoreLDSTile = k.control.isElemType<StoreLDSTile>();
                for(auto storeLDSTag :
                    k.control.findNodes(bodyEdges, isStoreLDSTile, GD::Downstream))
                {
                    auto top = getTopSetCoordinate(k, storeLDSTag);
                    for(auto edge : k.control.getNeighbours(top, GD::Upstream))
                        m_prefetchDelete[forLoop].insert(edge);
                    for(auto edge : k.control.getNeighbours(top, GD::Downstream))
                    {
                        if(!isBodyPredicate(edge))
                            m_prefetchDelete[forLoop].insert(edge);
                    }

                    auto target = getLDSOperationTarget(k, storeLDSTag);

                    m_info[forLoop][operationUnroll[storeLDSTag]][target].ldsOperation
                        = storeLDSTag;
                    m_info[forLoop][operationUnroll[storeLDSTag]][target].ldsChain = top;

                    Log::debug("AddPrefetch::stage: LDS store operation {} top {} target {}",
                               storeLDSTag,
                               top,
                               target);

                    alreadySeen.insert(top);
                }

                // Find global loads and detach them
                auto isLoadTiled = k.control.isElemType<LoadTiled>();
                for(auto loadTag : k.control.findNodes(bodyEdges, isLoadTiled, GD::Downstream))
                {
                    // If there isn't an info entry yet, then there
                    // isn't a matching StoreLDSTile operation.  In
                    // this case, LDS isn't being used for this User
                    // coordinate; don't try pre-fetching it unless
                    // it is a direct load (Global Memory -> VGPR).
                    auto user = k.mapper.get<User>(loadTag);
                    if(!m_info[forLoop][operationUnroll[loadTag]].contains(user))
                    {
                        const auto isDirectLoad   = IsDirectLoadToVGPR(k, loadTag);
                        const auto isExchangeLoad = isLoadForExchange(loadTag, k);
                        auto       ok = m_params->prefetchScale && (isExchangeLoad || isDirectLoad);
                        if(m_params->prefetchMixMemOps && !ok)
                        {
                            Throw<FatalError>(
                                "AddPrefetch: An unexpected direct load (not through LDS) was "
                                "detected, "
                                "and memory-operation mixing is enabled.  The AddPrefetch pass "
                                "can not continue.");
                        }

                        if(isDirectLoad && not(isExchangeLoad))
                        {
                            if(!operationUnroll.contains(loadTag))
                                continue;

                            auto top = getTopSetCoordinate(k, loadTag);

                            if(alreadySeen.contains(top))
                                continue;

                            auto unrollU = operationUnroll[loadTag];

                            Log::debug("AddPrefetch::stage: Direct load operation {} "
                                       "top {} unroll {}.",
                                       loadTag,
                                       top,
                                       unrollU);

                            m_directLoads[forLoop][unrollU].insert(
                                {.user = user, .tag = loadTag, .top = top});

                            for(auto edge : k.control.getNeighbours(top, GD::Upstream))
                                m_prefetchDelete[forLoop].insert(edge);
                            for(auto edge : k.control.getNeighbours(top, GD::Downstream))
                            {
                                if(!isBodyPredicate(edge))
                                    m_prefetchDelete[forLoop].insert(edge);
                            }

                            alreadySeen.insert(top);
                        }
                        else
                        {
                            Log::debug(
                                "AddPrefetch::stage: Skipping global non-LDS load operation {}",
                                loadTag);
                        }

                        continue;
                    }

                    auto top = getTopSetCoordinate(k, loadTag);
                    for(auto edge : k.control.getNeighbours(top, GD::Upstream))
                        m_prefetchDelete[forLoop].insert(edge);
                    for(auto edge : k.control.getNeighbours(top, GD::Downstream))
                    {
                        if(!isBodyPredicate(edge))
                            m_prefetchDelete[forLoop].insert(edge);
                    }

                    m_info[forLoop][operationUnroll[loadTag]][user].user            = user;
                    m_info[forLoop][operationUnroll[loadTag]][user].globalOperation = loadTag;
                    m_info[forLoop][operationUnroll[loadTag]][user].globalChain     = top;

                    Log::debug("AddPrefetch::stage: Global load operation {} top {} user {}",
                               loadTag,
                               top,
                               user);

                    alreadySeen.insert(top);
                }
            }

            //
            // Find separator edges and mark for deletion
            //
            auto inOperationUnroll = [&](int x) -> bool { return operationUnroll.contains(x); };
            for(auto [forLoop, numUnroll] : m_prefetchLoops)
            {
                for(auto bodyTop : k.control.getOutputNodeIndices<Body>(forLoop))
                {
                    for(auto bodyElem : filter(inOperationUnroll,
                                               k.control.depthFirstVisit(bodyTop, GD::Downstream)))
                    {
                        for(auto edge : filter(k.control.isElemType<Sequence>(),
                                               k.control.getNeighbours<GD::Downstream>(bodyElem)))
                        {
                            auto otherElem
                                = only(k.control.getNeighbours<GD::Downstream>(edge)).value();

                            if(operationUnroll.contains(otherElem)
                               && operationUnroll[otherElem] != operationUnroll[bodyElem])
                            {
                                m_prefetchDelete[forLoop].insert(edge);
                            }
                        }
                    }
                }
            }

            // Find things that haven't been seen yet (so aren't
            // global loads or LDS stores) and that aren't in a
            // cluster, and mark them as the body-starts
            for(auto [forLoop, numUnroll] : m_prefetchLoops)
            {
                auto headless = [&](int x) {
                    if(alreadySeen.contains(x))
                        return false;
                    for(auto incomingEdge : k.control.getNeighbours<GD::Upstream>(x))
                        if(!m_prefetchDelete[forLoop].contains(incomingEdge))
                            return false;
                    if(!operationUnroll.contains(x))
                        return false;
                    return true;
                };

                auto bodyEdges
                    = filter(isBodyPredicate, k.control.getNeighbours<GD::Downstream>(forLoop))
                          .to<std::vector>();

                for(auto start : k.control.findNodes(bodyEdges, headless, GD::Downstream))
                    m_prefetchUnrollBodyStarts[forLoop][operationUnroll[start]].insert(start);
            }

            //
            // Within each segment, determine which LoadLDSTile
            // operations should be prefetched.
            //
            // That is, some LoadLDSTile operation from the next
            // segment can be moved into (prefetched) the current
            // segment.
            //
            int splitLDSPrefetchFactor = m_params->prefetchLDSFactor;
            for(auto [forLoop, numUnroll] : m_prefetchLoops)
            {
                auto forLoopCoord     = getForLoopCoords(forLoop, k).first;
                auto maybeUnrollCoord = findUnrollNeighbour(k, forLoopCoord);
                AssertFatal(maybeUnrollCoord, "Prefetch with no unroll coordinate.");
                auto prefetchUnrollCoord = *maybeUnrollCoord;

                auto bodyEdges
                    = filter(isBodyPredicate, k.control.getNeighbours<GD::Downstream>(forLoop))
                          .to<std::vector>();
                auto isLoadLDSTile = k.control.isElemType<LoadLDSTile>();

                int prefetchLDSUnrollCoord = -1;
                {
                    // LDS prefetch colouring is based on the "small k", not jammed coordinates.
                    for(auto loadLDSTileTag :
                        k.control.findNodes(bodyEdges, isLoadLDSTile, GD::Downstream))
                    {
                        auto colour = colouring.operationColour[loadLDSTileTag];
                        for(auto kv : colour)
                        {
                            auto unrollCoord = kv.first;
                            if(unrollCoord == prefetchUnrollCoord)
                                continue;
                            auto incoming
                                = k.coordinates.getInputNodeIndices(unrollCoord, CT::isEdge<Split>)
                                      .to<std::vector>();
                            auto incomingJammed
                                = filterCoordinates<JammedWaveTileNumber>(incoming, k);
                            if(incomingJammed.empty())
                            {
                                prefetchLDSUnrollCoord = unrollCoord;
                            }
                        }
                    }
                }
                AssertFatal(prefetchLDSUnrollCoord != -1, "Can not find LDS prefetch coordinate.");

                Log::debug("AddPrefetch::stage: LDS prefetch: factor {} Unroll "
                           "coordinate {} size {}",
                           splitLDSPrefetchFactor,
                           prefetchLDSUnrollCoord,
                           unrollCoordSizes[prefetchLDSUnrollCoord]);

                for(auto loadLDSTileTag :
                    k.control.findNodes(bodyEdges, isLoadLDSTile, GD::Downstream))
                {
                    auto colour = colouring.operationColour[loadLDSTileTag];
                    auto u      = colour[prefetchUnrollCoord];

                    auto prefetchLDSUnrollValue = colour[prefetchLDSUnrollCoord];

                    auto loadLDSTileChain = getTopSetCoordinate(k, loadLDSTileTag);

                    if(splitLDSPrefetchFactor == 0)
                    {
                        // Keep load in same segment
                        m_loadFromLDSChains[forLoop][u].push_back(loadLDSTileChain);
                        continue;
                    }

                    if(splitLDSPrefetchFactor > 0)
                    {
                        if(prefetchLDSUnrollValue
                           >= unrollCoordSizes[prefetchLDSUnrollCoord] / splitLDSPrefetchFactor)
                        {
                            // Keep load in same segment
                            Log::debug(
                                "AddPrefetch::stage: LDS load operation {} keep (unroll value {})",
                                loadLDSTileTag,
                                prefetchLDSUnrollValue);
                            m_loadFromLDSChains[forLoop][u].push_back(loadLDSTileChain);
                            continue;
                        }
                    }

                    // Move load to previous segment (it is prefetchable)
                    Log::debug("AddPrefetch::stage: LDS load operation {} move (unroll value {})",
                               loadLDSTileTag,
                               prefetchLDSUnrollValue);

                    auto [target, direction] = getOperationTarget(loadLDSTileTag, k);
                    m_prefetchFromLDSChains[forLoop][u].insert(
                        {target, prefetchLDSUnrollValue, loadLDSTileChain});
                    m_prefetchUnrollBodyStarts[forLoop][u].erase(loadLDSTileChain);

                    for(auto inEdge : k.control.getNeighbours<GD::Upstream>(loadLDSTileChain))
                    {
                        m_prefetchDelete[forLoop].insert(inEdge);
                    }

                    for(auto outEdge :
                        filter(k.control.isElemType<Sequence>(),
                               k.control.getNeighbours<GD::Downstream>(loadLDSTileChain)))
                    {
                        if(m_prefetchDelete[forLoop].contains(outEdge))
                            continue;

                        auto outNode
                            = only(k.control.getNeighbours<GD::Downstream>(outEdge)).value();

                        m_prefetchUnrollBodyStarts[forLoop][u].insert(outNode);
                        m_prefetchDelete[forLoop].insert(outEdge);
                    }
                }
            }

            // Mark exchanges
            for(auto [forLoop, numUnroll] : m_prefetchLoops)
            {
                auto isExchangePredicate = k.control.isElemType<Exchange>();
                auto bodies = k.control.getOutputNodeIndices<Body>(forLoop).to<std::vector>();
                auto exchangeTags
                    = k.control.findNodes(bodies, isExchangePredicate).to<std::unordered_set>();
                for(auto exchangeTag : exchangeTags)
                    m_exchangeSegment[exchangeTag] = operationUnroll[exchangeTag];
            }

            //
            // Find tail end of each unrolled loop body
            //
            for(auto [forLoop, numUnroll] : m_prefetchLoops)
            {
                for(auto u = 0; u < numUnroll; ++u)
                {
                    auto starts = m_prefetchUnrollBodyStarts[forLoop][u];

                    int  fLoop                   = forLoop;
                    auto onlyFollowSequenceEdges = [&](int x) -> bool {
                        auto isSequence = CF::isEdge<Sequence>(k.control.getElement(x));
                        auto willDelete = m_prefetchDelete[fLoop].contains(x);
                        return isSequence && !willDelete;
                    };

                    for(auto op :
                        k.control.depthFirstVisit(starts, onlyFollowSequenceEdges, GD::Downstream))
                    {
                        auto outgoing = k.control.getNeighbours(op, GD::Downstream);
                        for(auto tag : m_prefetchDelete[forLoop])
                        {
                            std::erase(outgoing, tag);
                        }
                        if(outgoing.empty())
                        {
                            m_prefetchUnrollBodyEnds[forLoop][u].insert(op);
                        }
                    }
                }
            }
        }

        ConstraintStatus AcceptablePrefetchNodes(const KernelGraph& k)
        {
            TIMER(t, "Constraint::AcceptablePrefetchNodes");

            ConstraintStatus retval;

            for(auto [forLoop, numUnroll] : findPrefetch(k))
            {
                for(auto bodyTop : k.control.getOutputNodeIndices<Body>(forLoop))
                {
                    for(auto bodyElem : k.control.depthFirstVisit(bodyTop, GD::Downstream))
                    {
                        if(k.control.get<Body>(bodyElem) || k.control.get<Sequence>(bodyElem)
                           || k.control.get<SetCoordinate>(bodyElem)
                           || k.control.get<LoadTiled>(bodyElem)
                           || k.control.get<LoadLDSTile>(bodyElem)
                           || k.control.get<StoreLDSTile>(bodyElem)
                           || k.control.get<Exchange>(bodyElem) || k.control.get<Multiply>(bodyElem)
                           || k.control.get<NOP>(bodyElem))
                        {
                            continue;
                        }

                        auto op = k.control.getNode(bodyElem);
                        retval.combine(false,
                                       concatenate("Found unsupported node ",
                                                   toString(op),
                                                   " (",
                                                   bodyElem,
                                                   ") in forloop ",
                                                   forLoop,
                                                   "."));
                    }
                }
            }

            return retval;
        }
    }
}