derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
1631 lines
71 KiB
C++
1631 lines
71 KiB
C++
// 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;
|
|
}
|
|
}
|
|
}
|