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
1809 lines
84 KiB
C++
1809 lines
84 KiB
C++
// Copyright Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
/**
|
|
* AddStreamK -- stream accumulation tiles.
|
|
*
|
|
* The usual example is matrix-matrix multiply:
|
|
*
|
|
* D = A B
|
|
*
|
|
* where A and B are tiled so that A has M x K tiles, and B has K x N
|
|
* tiles. Tiles in the accumulation (K) dimension will be streamed.
|
|
*
|
|
* The flattened M * N * K global tile-space will be distributed
|
|
* evenly among the WGs.
|
|
*
|
|
* Each WG needs to iterate over its portion of the flattened global
|
|
* tile-space.
|
|
*
|
|
* To facilitate accumulation initialization and prefetching etc, we
|
|
* use two loops to accomplish this: an outer forTileIdx loop and an
|
|
* inner forAccumIdx loop.
|
|
*
|
|
* The inner forAccumIdx loop iterates over the accumulation tiles
|
|
* one-by-one. The outer forTileIdx loop iterates over the local
|
|
* tile-space. Its increment is: however many tiles the inner
|
|
* forAccumIdx loop processed.
|
|
*
|
|
* When the inner forAccumIdx loop advances, it will be advancing to a
|
|
* new K tile, but will remain within the same M/N tile. When the
|
|
* outer forTileIdx loop advances, it will be advancing to a new M/N
|
|
* tile.
|
|
*
|
|
* The local combined index
|
|
*
|
|
* wgTile = forTileIdx + forAccumIdx
|
|
*
|
|
* is the current WGs local tile in its portion of the global
|
|
* tile-space. Then
|
|
*
|
|
* tile = tilesPerWG * wg + wgTile
|
|
*
|
|
* is the global tile that the WG is processing. Given the global
|
|
* tile, the M/N/K tile coordinates are
|
|
*
|
|
* m = (tile / numTilesK) / numTilesN;
|
|
* n = (tile / numTilesK) % numTilesN;
|
|
* k = tile % numTilesK;
|
|
*
|
|
* In the "two-tile" version, we unevenly split (sunder) the global
|
|
* tile space into two partitions: the "streaming" partition and the
|
|
* "data-parallel" partition.
|
|
*
|
|
* In the data-parallel partition, we want each WG to process whole
|
|
* output tiles so that: synchronization between WGs is not necessary,
|
|
* and cache locality may be improved. Processing whole output tiles
|
|
* implies that the number of global tiles processed in the
|
|
* data-parallel partition must be a multiple of WG * K.
|
|
*
|
|
* To minimize quantization inefficiency over the WGs in the streaming
|
|
* partition, we want each WG to process at least one output tile, but
|
|
* no more than two.
|
|
*
|
|
* In the data-parallel section, each WG will process
|
|
*
|
|
* N_DP = (((M * N * K) // (WG * K) - 1) * WG * K
|
|
*
|
|
* tiles; where // is "floor division". Note:
|
|
*
|
|
* 1. This a multiple of WG * K (no synchronization)
|
|
* 2. The "- 1" pushes enough tiles over to the streaming partition
|
|
* to minimize quantization inefficiency over the WGs.
|
|
*
|
|
* Simplifying, we obtain
|
|
*
|
|
* N_DP = ((M * N * K) // (WG * K) - 1) * WG * K
|
|
* = ((M * N) // WG - 1) * WG * K
|
|
*
|
|
* This means that
|
|
*
|
|
* N_SK = M * N * K - N_DP
|
|
* = K * (M * N - ((M * N) // WG - 1) * WG)
|
|
* = K * (M * N - ((M * N) // WG * WG - WG))
|
|
* = K * (M * N - (M * N) // WG * WG + WG)
|
|
* = K * ((M * N) % WG + WG)
|
|
*
|
|
* tiles will be processed in the streaming partition.
|
|
*/
|
|
|
|
#include <map>
|
|
#include <optional>
|
|
#include <unordered_set>
|
|
#include <vector>
|
|
|
|
#include <rocRoller/Expression.hpp>
|
|
#include <rocRoller/ExpressionTransformations.hpp>
|
|
#include <rocRoller/KernelGraph/ControlGraph/ControlFlowRWTracer.hpp>
|
|
#include <rocRoller/KernelGraph/ControlGraph/Operation.hpp>
|
|
#include <rocRoller/KernelGraph/KernelGraph.hpp>
|
|
#include <rocRoller/KernelGraph/Transforms/AddStreamK.hpp>
|
|
#include <rocRoller/KernelGraph/Utils.hpp>
|
|
#include <rocRoller/Utilities/Logging.hpp>
|
|
|
|
namespace rocRoller
|
|
{
|
|
namespace KernelGraph
|
|
{
|
|
namespace Expression = rocRoller::Expression;
|
|
namespace CG = rocRoller::KernelGraph::CoordinateGraph;
|
|
|
|
using namespace ControlGraph;
|
|
using namespace CoordinateGraph;
|
|
using namespace Expression;
|
|
using namespace Register;
|
|
|
|
using GD = rocRoller::Graph::Direction;
|
|
|
|
/*
|
|
* Transform info.
|
|
*/
|
|
|
|
/// Basic info about the incoming loops that need to be transformed.
|
|
struct LoopInfo
|
|
{
|
|
/// Sub-dimensions of dangling `MacroTileNumber`s that should be included in the streaming construct.
|
|
std::vector<int> dimensionIndices;
|
|
/// Control tag of incoming top-loop
|
|
int topLoopOp;
|
|
/// Control tag of incoming accumulator-loop
|
|
int accumulatorLoopOp;
|
|
/// Control tag of XLoop
|
|
int xLoop = -1;
|
|
/// Control tag of YLoop
|
|
int yLoop = -1;
|
|
/// Length of XLoop
|
|
unsigned int xLoopSize = 1;
|
|
/// Length of YLoop
|
|
unsigned int yLoopSize = 1;
|
|
};
|
|
|
|
/// Info about accumulator.
|
|
struct AccumulatorInfo
|
|
{
|
|
/// Coordinate dimension of the accumulator-loop (the ForLoop)
|
|
int accumulatorCoord;
|
|
/// Coordinate dimension of tile into which the accumulator-loop accumulates
|
|
int accumulatorTile;
|
|
/// DataType of accumulator tile
|
|
VariableType accumulatorVarType;
|
|
/// Set of control tags, after the accumulator-loop, that use the accumulator-tile.
|
|
std::unordered_set<int> usesAccumulatorTile;
|
|
|
|
/// Staged MacroTileNumber coordinates
|
|
//
|
|
// Mapping: dimension -> set of MacroTileNumber coordinates
|
|
std::map<int, std::unordered_set<int>> tileNumberCoords;
|
|
};
|
|
|
|
/// Info about the SK and DP tile spaces.
|
|
struct StreamKInfo
|
|
{
|
|
int localTileSpaceSK; /// Streaming tile space
|
|
std::optional<int> localTileSpaceDP; /// Data-parallel tile space (TwoTile only)
|
|
std::optional<int> selector; /// SK vs DP selector for Sunder (TwoTile only)
|
|
};
|
|
|
|
/// Info about the "number of tiles" kernel arguments.
|
|
struct ArgumentInfo
|
|
{
|
|
// Kernel arguments
|
|
std::vector<Expression::ExpressionPtr> numTiles, numTileArgExprs;
|
|
Expression::ExpressionPtr numSKTiles, numDPTiles;
|
|
Expression::ExpressionPtr numWGs, numSKTilesPerWG, numDPTilesPerWG;
|
|
};
|
|
|
|
//
|
|
// Helpers
|
|
//
|
|
|
|
auto makeFindLoopPredicate(KernelGraph const& graph, std::string const& loopName)
|
|
{
|
|
auto findLoopPredicate = [&](int tag) -> bool {
|
|
auto maybeForLoop = graph.control.get<ForLoopOp>(tag);
|
|
if(!maybeForLoop)
|
|
return false;
|
|
if(maybeForLoop->loopName == loopName)
|
|
return true;
|
|
return false;
|
|
};
|
|
return findLoopPredicate;
|
|
}
|
|
|
|
template <typename... Args>
|
|
ExpressionPtr minimum(ExpressionPtr x, ExpressionPtr y, Args... remaining)
|
|
{
|
|
if constexpr(sizeof...(remaining) > 0)
|
|
return minimum(minimum(x, y), remaining...);
|
|
|
|
return std::make_shared<Expression::Expression>(Expression::Conditional{x < y, x, y});
|
|
}
|
|
|
|
// x - y, but if x < y, return 0.
|
|
ExpressionPtr conditionalSubtract(ExpressionPtr x, ExpressionPtr y, DataType dataType)
|
|
{
|
|
return std::make_shared<Expression::Expression>(
|
|
Expression::Conditional{x > y, x - y, literal(0, dataType)});
|
|
}
|
|
|
|
/**
|
|
* Create Assign node to add partially accumulated tiles.
|
|
*
|
|
* Returns tag of new Assign node.
|
|
*/
|
|
int addFixup(KernelGraph& graph,
|
|
int partialMacTileTag,
|
|
int destMacTileTag,
|
|
DataType dataType,
|
|
uint numVGPRs)
|
|
{
|
|
auto lhsExpr = std::make_shared<Expression::Expression>(
|
|
Expression::DataFlowTag{destMacTileTag, Register::Type::Vector, dataType});
|
|
auto rhsExpr = std::make_shared<Expression::Expression>(
|
|
Expression::DataFlowTag{partialMacTileTag, Register::Type::Vector, dataType});
|
|
|
|
auto addExpr = lhsExpr + rhsExpr;
|
|
auto fixupTag
|
|
= graph.control.addElement(Assign{Register::Type::Vector, addExpr, numVGPRs});
|
|
|
|
graph.coordinates.addElement(DataFlow(), {partialMacTileTag}, {destMacTileTag});
|
|
|
|
graph.mapper.connect(fixupTag, destMacTileTag, NaryArgument::DEST);
|
|
|
|
return fixupTag;
|
|
}
|
|
|
|
/*
|
|
* Scratch tile info.
|
|
*/
|
|
|
|
struct ScratchInfo
|
|
{
|
|
int store; //< Source internal tile for store operation
|
|
int load; //< Destination internal tile for load operation
|
|
int tile; //< Scratch tile
|
|
int setPlusOne; //< SetCoordinate operation; must be inserted above the Load
|
|
};
|
|
|
|
struct SendInfo
|
|
{
|
|
int preWaitZero; //< WaitZero before conditional
|
|
int sendCond; //< Conditional send-block
|
|
};
|
|
|
|
struct RecvInfo
|
|
{
|
|
int preWaitZero;
|
|
int receiveCond;
|
|
int setPlusOne;
|
|
};
|
|
|
|
/**
|
|
* Create coordinate transforms for exchanging partially
|
|
* accumulated tiles through scratch space.
|
|
*/
|
|
ScratchInfo loadStoreMacroTileSCRATCH(KernelGraph& graph,
|
|
std::vector<DeferredConnection>& storeConnections,
|
|
std::vector<DeferredConnection>& loadConnections,
|
|
int macTileTag,
|
|
VariableType varType,
|
|
LoopInfo const& loopInfo,
|
|
ArgumentInfo const& argInfo,
|
|
int forReceiveTileLoopCoord,
|
|
CommandParametersPtr params,
|
|
ContextPtr context)
|
|
{
|
|
auto macTile = graph.coordinates.getNode<MacroTile>(macTileTag);
|
|
|
|
auto numWorkgroupsX = argInfo.numWGs;
|
|
auto numWorkgroupsY = literal(1u);
|
|
|
|
auto sizeX = simplify(numWorkgroupsX * literal(static_cast<uint>(macTile.sizes[0])));
|
|
auto sizeY = simplify(numWorkgroupsY * literal(static_cast<uint>(macTile.sizes[1])));
|
|
|
|
auto strideX = sizeY;
|
|
auto strideY = literal(1u);
|
|
|
|
auto globalScratch = newScratchCoordinate(
|
|
simplify(sizeX * sizeY), varType, Operations::ScratchPolicy::None, context);
|
|
auto globalScratchTag = graph.coordinates.addElement(globalScratch);
|
|
|
|
std::vector<unsigned int> jammedSizes = {loopInfo.xLoopSize, loopInfo.yLoopSize};
|
|
|
|
// Store
|
|
auto jammedStoreScratchTileTag = createInternalTile(
|
|
graph, varType, macTileTag, jammedSizes, false, params, context);
|
|
graph.coordinates.addElement(View(), {jammedStoreScratchTileTag}, {macTileTag});
|
|
|
|
auto jammedStoreScratchTile
|
|
= *graph.coordinates.get<MacroTile>(jammedStoreScratchTileTag);
|
|
jammedStoreScratchTile.layoutType = LayoutType::SCRATCH;
|
|
graph.coordinates.setElement(jammedStoreScratchTileTag, jammedStoreScratchTile);
|
|
|
|
std::vector<int> storeSubDimensions
|
|
= {graph.coordinates.addElement(SubDimension(0, sizeX, strideX)),
|
|
graph.coordinates.addElement(SubDimension(1, sizeY, strideY))};
|
|
graph.coordinates.addElement(
|
|
Join(), storeSubDimensions, std::vector<int>{globalScratchTag});
|
|
|
|
{
|
|
auto [nMacX, iMacX, iMacY] = addStore1DMacroTileCT(
|
|
graph, storeConnections, macTileTag, storeSubDimensions);
|
|
|
|
addStoreThreadTileCT(graph,
|
|
storeConnections,
|
|
jammedStoreScratchTileTag,
|
|
iMacX,
|
|
iMacY,
|
|
context->kernel()->workgroupSize(),
|
|
jammedSizes,
|
|
true);
|
|
|
|
graph.coordinates.addElement(
|
|
DataFlow(), {jammedStoreScratchTileTag}, {globalScratchTag});
|
|
}
|
|
|
|
storeConnections.push_back(DC<MacroTile>(jammedStoreScratchTileTag));
|
|
storeConnections.push_back(DC<User>(globalScratchTag));
|
|
|
|
// Load
|
|
auto jammedLoadScratchTileTag = createInternalTile(
|
|
graph, varType, macTileTag, jammedSizes, false, params, context);
|
|
auto jammedLoadScratchTile
|
|
= *graph.coordinates.get<MacroTile>(jammedLoadScratchTileTag);
|
|
jammedLoadScratchTile.layoutType = LayoutType::SCRATCH;
|
|
graph.coordinates.setElement(jammedLoadScratchTileTag, jammedLoadScratchTile);
|
|
|
|
auto loadScratchTileTag
|
|
= createInternalTile(graph, varType, macTileTag, params, context);
|
|
auto loadScratchTile = *graph.coordinates.get<MacroTile>(loadScratchTileTag);
|
|
loadScratchTile.layoutType = LayoutType::SCRATCH;
|
|
graph.coordinates.setElement(loadScratchTileTag, loadScratchTile);
|
|
|
|
graph.coordinates.addElement(View(), {jammedLoadScratchTileTag}, {loadScratchTileTag});
|
|
|
|
std::vector<int> loadSubDimensions
|
|
= {graph.coordinates.addElement(SubDimension(0, sizeX, strideX)),
|
|
graph.coordinates.addElement(SubDimension(1, sizeY, strideY))};
|
|
graph.coordinates.addElement(
|
|
Split(), std::vector<int>{globalScratchTag}, loadSubDimensions);
|
|
|
|
{
|
|
auto [nMacX, iMacX, iMacY]
|
|
= addLoad1DMacroTileCT(graph, loadConnections, macTileTag, loadSubDimensions);
|
|
|
|
addLoadThreadTileCT(graph,
|
|
loadConnections,
|
|
jammedLoadScratchTileTag,
|
|
iMacX,
|
|
iMacY,
|
|
context->kernel()->workgroupSize(),
|
|
jammedSizes,
|
|
true);
|
|
|
|
graph.coordinates.addElement(DataFlow(), {loadScratchTileTag}, {globalScratchTag});
|
|
}
|
|
|
|
// Find the dangling MacroTileNumber and add one to it...
|
|
auto danglingMacroTileNumber = [&graph](int tag) -> bool {
|
|
auto maybeTileNumber = graph.coordinates.get<MacroTileNumber>(tag);
|
|
if(!maybeTileNumber)
|
|
return false;
|
|
if(!std::empty(graph.coordinates.getNeighbours<GD::Downstream>(tag)))
|
|
return false;
|
|
if(maybeTileNumber->dim != 0)
|
|
return false;
|
|
return true;
|
|
};
|
|
|
|
auto nextTileNumTag
|
|
= graph.coordinates.findNodes(globalScratchTag, danglingMacroTileNumber)
|
|
.take(1)
|
|
.only()
|
|
.value();
|
|
|
|
auto one = Expression::literal(1u);
|
|
auto plusOneTag = graph.coordinates.addElement(Linear(one, one));
|
|
auto setPlusOneTag = graph.control.addElement(SetCoordinate(one));
|
|
graph.mapper.connect<Linear>(setPlusOneTag, plusOneTag);
|
|
|
|
auto tileNumTag = graph.coordinates.addElement(
|
|
*graph.coordinates.get<MacroTileNumber>(nextTileNumTag)); // Copy existing
|
|
graph.coordinates.addElement(
|
|
Split(), {nextTileNumTag}, {tileNumTag, plusOneTag, forReceiveTileLoopCoord});
|
|
|
|
loadConnections.push_back(DC<MacroTile>(loadScratchTileTag));
|
|
loadConnections.push_back(DC<User>(globalScratchTag));
|
|
|
|
return {jammedStoreScratchTileTag, loadScratchTileTag, globalScratchTag, setPlusOneTag};
|
|
}
|
|
|
|
/**
|
|
* Create send-tile block, which is roughly:
|
|
*
|
|
* WaitZero()
|
|
* if sendTileExpr:
|
|
* StoreTile()
|
|
* WaitZero()
|
|
* Barrier()
|
|
* if wave0:
|
|
* flag = Assign(SGPR, 1u);
|
|
* StoreSGPR(flag)
|
|
* WaitZero()
|
|
*/
|
|
SendInfo sendTile(KernelGraph& graph,
|
|
ExpressionPtr sendTileExpr,
|
|
std::vector<DeferredConnection> const& storeConnections,
|
|
int flagsScratchTag,
|
|
DataType scratchDataType,
|
|
VariableType numTilesVarType,
|
|
LoopInfo const& loopInfo,
|
|
ContextPtr context)
|
|
{
|
|
auto one = Expression::literal(1u);
|
|
auto zero = Expression::literal(0, numTilesVarType);
|
|
auto DF = [numTilesVarType](int tag) {
|
|
return std::make_shared<Expression::Expression>(
|
|
Expression::DataFlowTag{tag, Register::Type::Scalar, numTilesVarType});
|
|
};
|
|
|
|
auto workgroup = graph.coordinates.addElement(Workgroup(0));
|
|
graph.coordinates.addElement(PassThrough(), {workgroup}, {flagsScratchTag});
|
|
|
|
// Store tile
|
|
int forX = cloneForLoop(graph, loopInfo.xLoop);
|
|
int forY = cloneForLoop(graph, loopInfo.yLoop);
|
|
|
|
auto store = StoreTiled(scratchDataType);
|
|
|
|
// TODO: Improve setting of arch-specific buffer options
|
|
if(!(context->targetArchitecture().target().isCDNA1GPU()
|
|
|| context->targetArchitecture().target().isCDNA2GPU()))
|
|
{
|
|
store.bufOpts = {.sc1 = true};
|
|
}
|
|
|
|
auto storeTileTag = graph.control.addElement(store);
|
|
for(auto const& c : storeConnections)
|
|
graph.mapper.connect(storeTileTag, c.coordinate, c.connectionSpec);
|
|
|
|
graph.control.addElement(Body(), {forX}, {forY});
|
|
graph.control.addElement(Body(), {forY}, {storeTileTag});
|
|
|
|
auto jammedX = graph.mapper.get<JammedWaveTileNumber>(storeTileTag, 0);
|
|
if(jammedX != -1)
|
|
graph.coordinates.addElement(
|
|
PassThrough(), {graph.mapper.get<ForLoop>(forX)}, {jammedX});
|
|
auto jammedY = graph.mapper.get<JammedWaveTileNumber>(storeTileTag, 1);
|
|
if(jammedY != -1)
|
|
graph.coordinates.addElement(
|
|
PassThrough(), {graph.mapper.get<ForLoop>(forY)}, {jammedY});
|
|
|
|
auto sendTileTag = graph.control.addElement(
|
|
ConditionalOp{sendTileExpr, ConditionalMode::Branch, "Send Tile"});
|
|
auto barrierTag = graph.control.addElement(Barrier());
|
|
auto waitZeroTag = graph.control.addElement(WaitZero());
|
|
|
|
// Store flag
|
|
auto flagRegister = graph.coordinates.addElement(VGPR());
|
|
|
|
auto assignFlagTag = graph.control.addElement(Assign{Register::Type::Scalar, one});
|
|
graph.mapper.connect(assignFlagTag, flagRegister, NaryArgument::DEST);
|
|
|
|
// TODO: Improve setting of arch-specific buffer options
|
|
BufferInstructionOptions bufOpts{.glc = true};
|
|
|
|
auto storeFlagTag = graph.control.addElement(StoreSGPR(DataType::UInt32, bufOpts));
|
|
graph.mapper.connect<User>(storeFlagTag, flagsScratchTag);
|
|
graph.mapper.connect<VGPR>(storeFlagTag, flagRegister);
|
|
|
|
// Create workitem coordinate and expression for wave 0 check
|
|
// Only workitem 0 (wave 0) should write to the flag
|
|
auto workitemTag = graph.coordinates.addElement(Workitem(0));
|
|
auto workitemDF = std::make_shared<Expression::Expression>(
|
|
Expression::DataFlowTag{workitemTag, Register::Type::Vector, DataType::UInt32});
|
|
auto isWave0Expr = (workitemDF == Expression::literal(0u));
|
|
auto wave0FlagStoreTag = graph.control.addElement(
|
|
ConditionalOp{isWave0Expr, ConditionalMode::Branch, "Wave0 Store Flag"});
|
|
|
|
// Add to control
|
|
auto preWaitZeroTag = graph.control.addElement(WaitZero());
|
|
|
|
graph.control.addElement(Sequence(), {preWaitZeroTag}, {sendTileTag});
|
|
graph.control.addElement(Body(), {sendTileTag}, {forX});
|
|
graph.control.chain<Sequence>(forX, waitZeroTag, barrierTag, wave0FlagStoreTag);
|
|
graph.control.addElement(Body(), {wave0FlagStoreTag}, {assignFlagTag});
|
|
auto waitAfterStoreFlagTag = graph.control.addElement(WaitZero());
|
|
graph.control.chain<Sequence>(assignFlagTag, storeFlagTag, waitAfterStoreFlagTag);
|
|
|
|
return {preWaitZeroTag, sendTileTag};
|
|
}
|
|
|
|
/**
|
|
* Create receive-tile block, which is roughly:
|
|
*
|
|
* WaitZero()
|
|
* if receiveTileExpr:
|
|
* for each partial result (loop created externally)
|
|
* nextWG = WG + 1 + i
|
|
* Assert nextWG < numScratch
|
|
* do:
|
|
* LoadSGPR(flag[nextWG])
|
|
* while flag[nextWG] == 0
|
|
* Barrier()
|
|
* if wave0:
|
|
* Assign(flag[nextWG] = 0)
|
|
* StoreSGPR(flag[nextWG])
|
|
* WaitZero()
|
|
* partiallyAccumulatedTile = LoadTiled()
|
|
* fullyAccumulatedTile = Assign(localPartiallyAccumulatedTile)
|
|
* fullyAccumulatedTile = Assign(fullyAccumulatedTile + partiallyAccumulatedTile)
|
|
* localPartiallyAccumulatedTile = Assign(fullyAccumulatedTile)
|
|
* WaitZero()
|
|
*
|
|
* Note this also update all subsequent references to
|
|
* localPartiallyAccumulatedTile to fullyAccumulatedTile.
|
|
*/
|
|
RecvInfo receiveTile(KernelGraph& graph,
|
|
ExpressionPtr receiveTileExpr,
|
|
int scratchTileTag,
|
|
std::vector<DeferredConnection> const& loadConnections,
|
|
int flagsScratchTag,
|
|
int accumulatorTileTag,
|
|
std::unordered_set<int> const& usesAccumulatorTile,
|
|
DataType dataType,
|
|
LoopInfo const& loopInfo,
|
|
ArgumentInfo const& argInfo,
|
|
int forReceiveTileLoopOp,
|
|
int forReceiveTileLoopCoord,
|
|
CommandParametersPtr params,
|
|
ContextPtr context)
|
|
{
|
|
auto DF = [](int tag) {
|
|
return std::make_shared<Expression::Expression>(
|
|
Expression::DataFlowTag{tag, Register::Type::Scalar, DataType::UInt32});
|
|
};
|
|
|
|
auto one = Expression::literal(1u);
|
|
auto zero = Expression::literal(0u);
|
|
|
|
auto workgroup = graph.coordinates.addElement(Workgroup(0));
|
|
|
|
// Read tile
|
|
auto receiveTileTag = graph.control.addElement(
|
|
ConditionalOp{receiveTileExpr, ConditionalMode::Branch, "Receive Tile"});
|
|
|
|
// Read flag
|
|
auto plusOneTag = graph.coordinates.addElement(Linear(one, one));
|
|
auto setPlusOneTag = graph.control.addElement(SetCoordinate(one));
|
|
graph.mapper.connect<Linear>(setPlusOneTag, plusOneTag);
|
|
|
|
auto nextWorkgroupTag = graph.coordinates.addElement(Linear(nullptr, one));
|
|
graph.coordinates.addElement(
|
|
Split(), {nextWorkgroupTag}, {workgroup, plusOneTag, forReceiveTileLoopCoord});
|
|
|
|
// TODO: Improve setting of arch-specific buffer options
|
|
BufferInstructionOptions bufOpts{.glc = true};
|
|
|
|
auto flagRegister = graph.coordinates.addElement(VGPR());
|
|
auto loadFlagTag = graph.control.addElement(LoadSGPR(DataType::UInt32, bufOpts));
|
|
|
|
auto numScratch = argInfo.numWGs;
|
|
|
|
auto boundsCheckTag = graph.control.addElement(
|
|
AssertOp{"Bounds Check", (DF(nextWorkgroupTag) < numScratch)});
|
|
|
|
graph.mapper.connect<User>(loadFlagTag, flagsScratchTag);
|
|
graph.mapper.connect<VGPR>(loadFlagTag, flagRegister);
|
|
|
|
// TODO: See if the ControlFlowRWTracer can discover this dependency without this explicit connection.
|
|
// This is currently needed for InlineIncrements to work correctly.
|
|
graph.mapper.connect<ForLoop>(loadFlagTag, forReceiveTileLoopCoord);
|
|
|
|
graph.coordinates.addElement(PassThrough(), {flagsScratchTag}, {nextWorkgroupTag});
|
|
|
|
auto doWhileTag = graph.control.addElement(
|
|
DoWhileOp{(DF(flagRegister) == zero), "Global sync spin loop"});
|
|
|
|
// The coordinate graph for load, store, and reset flags:
|
|
//
|
|
// Workgroup
|
|
// |
|
|
// PassThrough
|
|
// |
|
|
// v
|
|
// flagsScratchTag <-Duplicate-- resetFlagsScratchTag
|
|
// | ^
|
|
// PassThrough PassThrough
|
|
// | |
|
|
// v resetNextWorkgroupTag
|
|
// nextWorkgroupTag ^
|
|
// | Join
|
|
// Split / | \
|
|
// / | \ / | \
|
|
// v v v / | \
|
|
// Workgroup plusOne forReceiveTileLoop
|
|
//
|
|
// Note: nextWorkgroupTag and resetNextWorkgroupTag both connect to the same
|
|
// neighbors (Workgroup, plusOne, forReceiveTileLoop) but in opposite directions
|
|
// (Split vs Join). This allows loading flags from WG+1+i and resetting flags
|
|
// at the same index.
|
|
|
|
// Create coordinate to indicate flag index to reset
|
|
auto resetNextWorkgroupTag = graph.coordinates.addElement(Linear(nullptr, one));
|
|
graph.coordinates.addElement(
|
|
Join(), {workgroup, plusOneTag, forReceiveTileLoopCoord}, {resetNextWorkgroupTag});
|
|
|
|
// Duplicate flags scratch coordinate for reset operation
|
|
auto resetFlagsScratchTag
|
|
= graph.coordinates.addElement(*graph.coordinates.get<User>(flagsScratchTag));
|
|
graph.coordinates.addElement(Duplicate(), {resetFlagsScratchTag}, {flagsScratchTag});
|
|
graph.coordinates.addElement(
|
|
PassThrough(), {resetNextWorkgroupTag}, {resetFlagsScratchTag});
|
|
|
|
// Reset flag operations
|
|
auto assignResetFlagTag
|
|
= graph.control.addElement(Assign{Register::Type::Scalar, zero});
|
|
graph.mapper.connect(assignResetFlagTag, flagRegister, NaryArgument::DEST);
|
|
|
|
auto resetFlagTag = graph.control.addElement(StoreSGPR(DataType::UInt32, bufOpts));
|
|
graph.mapper.connect<User>(resetFlagTag, resetFlagsScratchTag);
|
|
graph.mapper.connect<VGPR>(resetFlagTag, flagRegister);
|
|
|
|
// TODO: See if the ControlFlowRWTracer can discover this dependency without this explicit connection.
|
|
// This is currently needed for InlineIncrements to work correctly.
|
|
graph.mapper.connect<ForLoop>(resetFlagTag, forReceiveTileLoopCoord);
|
|
|
|
// Create workitem coordinate and expression for wave 0 check
|
|
// Only workitem 0 (wave 0) should write to the flag
|
|
auto workitemTag = graph.coordinates.addElement(Workitem(0));
|
|
auto workitemDF = std::make_shared<Expression::Expression>(
|
|
Expression::DataFlowTag{workitemTag, Register::Type::Vector, DataType::UInt32});
|
|
auto isWave0Expr = (workitemDF == Expression::literal(0u));
|
|
auto wave0ResetFlagTag = graph.control.addElement(
|
|
ConditionalOp{isWave0Expr, ConditionalMode::Branch, "Wave0 Reset Flag"});
|
|
|
|
auto barrierBeforeResetTag = graph.control.addElement(Barrier());
|
|
|
|
auto accumulatorTile = graph.coordinates.get<MacroTile>(accumulatorTileTag);
|
|
uint numRegisters = accumulatorTile->elements()
|
|
/ (product(context->kernel()->workgroupSize()) * loopInfo.xLoopSize
|
|
* loopInfo.yLoopSize);
|
|
|
|
// Read tile
|
|
int loadAddForX = cloneForLoop(graph, loopInfo.xLoop);
|
|
int loadAddForY = cloneForLoop(graph, loopInfo.yLoop);
|
|
|
|
auto loadTileTag = graph.control.addElement(LoadTiled(dataType));
|
|
for(auto const& c : loadConnections)
|
|
graph.mapper.connect(loadTileTag, c.coordinate, c.connectionSpec);
|
|
|
|
graph.control.addElement(Body(), {loadAddForX}, {loadAddForY});
|
|
graph.control.addElement(Body(), {loadAddForY}, {loadTileTag});
|
|
|
|
auto jammedX = graph.mapper.get<JammedWaveTileNumber>(loadTileTag, 0);
|
|
if(jammedX != -1)
|
|
graph.coordinates.addElement(
|
|
PassThrough(), {jammedX}, {graph.mapper.get<ForLoop>(loadAddForX)});
|
|
auto jammedY = graph.mapper.get<JammedWaveTileNumber>(loadTileTag, 1);
|
|
if(jammedY != -1)
|
|
graph.coordinates.addElement(
|
|
PassThrough(), {jammedY}, {graph.mapper.get<ForLoop>(loadAddForY)});
|
|
|
|
// Assign accumulator tile to temporal vgpr tile
|
|
auto fullyAccumulatedTileTag = graph.coordinates.addElement(MacroTile());
|
|
auto localAccumulatorTileExpr = std::make_shared<Expression::Expression>(
|
|
Expression::DataFlowTag{accumulatorTileTag, Register::Type::Vector, dataType});
|
|
auto assignAccTileOp = graph.control.addElement(
|
|
Assign{Register::Type::Vector, localAccumulatorTileExpr, numRegisters});
|
|
graph.mapper.connect(assignAccTileOp, fullyAccumulatedTileTag, NaryArgument::DEST);
|
|
graph.control.addElement(Sequence(), {loadTileTag}, {assignAccTileOp});
|
|
|
|
// Fixup
|
|
auto fixupTag
|
|
= addFixup(graph, scratchTileTag, fullyAccumulatedTileTag, dataType, numRegisters);
|
|
graph.control.addElement(Sequence(), {assignAccTileOp}, {fixupTag});
|
|
|
|
// Assign fixup result to accumulator tile
|
|
auto fullyAccumulatedTileExpr = std::make_shared<Expression::Expression>(
|
|
Expression::DataFlowTag{fullyAccumulatedTileTag, Register::Type::Vector, dataType});
|
|
auto assignFixupResult = graph.control.addElement(
|
|
Assign{Register::Type::Accumulator, fullyAccumulatedTileExpr, numRegisters});
|
|
graph.mapper.connect(assignFixupResult, accumulatorTileTag, NaryArgument::DEST);
|
|
graph.control.addElement(Sequence(), {fixupTag}, {assignFixupResult});
|
|
|
|
// Add to control
|
|
auto preWaitZeroTag = graph.control.addElement(WaitZero());
|
|
auto postWaitZeroTag = graph.control.addElement(WaitZero());
|
|
|
|
graph.control.chain<Sequence>(preWaitZeroTag, receiveTileTag);
|
|
|
|
graph.control.addElement(Body(), {receiveTileTag}, {forReceiveTileLoopOp});
|
|
graph.control.addElement(Body(), {forReceiveTileLoopOp}, {boundsCheckTag});
|
|
graph.control.addElement(Sequence(), {boundsCheckTag}, {doWhileTag});
|
|
graph.control.addElement(Body(), {doWhileTag}, {loadFlagTag});
|
|
|
|
graph.control.chain<Sequence>(doWhileTag, barrierBeforeResetTag, wave0ResetFlagTag);
|
|
graph.control.addElement(Body(), {wave0ResetFlagTag}, {assignResetFlagTag});
|
|
auto waitAfterRestFlagStoreTag = graph.control.addElement(WaitZero());
|
|
graph.control.chain<Sequence>(
|
|
assignResetFlagTag, resetFlagTag, waitAfterRestFlagStoreTag);
|
|
graph.control.chain<Sequence>(wave0ResetFlagTag, loadAddForX, postWaitZeroTag);
|
|
|
|
return {preWaitZeroTag, receiveTileTag, setPlusOneTag};
|
|
}
|
|
|
|
/**
|
|
* Creates a ForLoop with a non-constant increment.
|
|
*/
|
|
int conditionalFor(KernelGraph& graph,
|
|
int incrementCoord,
|
|
ExpressionPtr conditionExpr,
|
|
ExpressionPtr incrementExpr,
|
|
std::string const& loopName,
|
|
VariableType varType)
|
|
{
|
|
auto forLoop = graph.control.addElement(ForLoopOp{conditionExpr, loopName});
|
|
graph.mapper.connect(forLoop, incrementCoord, NaryArgument::DEST);
|
|
|
|
auto initialize = graph.control.addElement(
|
|
Assign{Register::Type::Scalar, Expression::literal(0, varType)});
|
|
graph.mapper.connect(initialize, incrementCoord, NaryArgument::DEST);
|
|
|
|
auto increment
|
|
= graph.control.addElement(Assign{Register::Type::Scalar, incrementExpr});
|
|
graph.mapper.connect(increment, incrementCoord, NaryArgument::DEST);
|
|
|
|
graph.control.addElement(Initialize(), {forLoop}, {initialize});
|
|
graph.control.addElement(ForLoopIncrement(), {forLoop}, {increment});
|
|
|
|
return forLoop;
|
|
}
|
|
|
|
/**
|
|
* Add chain of SetCoordinate and Assign nodes to initialize coordinates.
|
|
*
|
|
* @p setCoordinateGroups is a list of {coordinates, value} pairs.
|
|
* Each group's coordinates get SetCoordinate nodes with the group's
|
|
* value. All SetCoordinate nodes are chained via Body edges from @p top.
|
|
*
|
|
* @p assigns get Assign nodes with @p assignValue, chained via
|
|
* Sequence edges after the SetCoordinate chain.
|
|
*/
|
|
int initializeCoordinates(
|
|
KernelGraph& graph,
|
|
int top,
|
|
std::initializer_list<std::pair<std::initializer_list<int>, ExpressionPtr>>
|
|
setCoordinateGroups,
|
|
std::initializer_list<int> assigns,
|
|
ExpressionPtr assignValue)
|
|
{
|
|
std::vector<int> setChain;
|
|
for(auto const& [coordinates, value] : setCoordinateGroups)
|
|
{
|
|
for(auto coordinate : coordinates)
|
|
{
|
|
auto setCoordinate = graph.control.addElement(SetCoordinate(value));
|
|
graph.mapper.connect(setCoordinate, coordinate, NaryArgument::DEST);
|
|
setChain.push_back(setCoordinate);
|
|
}
|
|
}
|
|
|
|
std::vector<int> assignChain;
|
|
for(auto tag : assigns)
|
|
{
|
|
auto def = graph.control.addElement(Assign{Register::Type::Scalar, assignValue});
|
|
graph.mapper.connect(def, tag, NaryArgument::DEST);
|
|
assignChain.push_back(def);
|
|
}
|
|
|
|
graph.control.addElement(Body(), {top}, {setChain.front()});
|
|
for(auto i = 1; i < setChain.size(); ++i)
|
|
{
|
|
graph.control.addElement(Body(), {setChain[i - 1]}, {setChain[i]});
|
|
}
|
|
|
|
graph.control.addElement(Body(), {setChain.back()}, {assignChain.front()});
|
|
for(auto i = 1; i < assignChain.size(); ++i)
|
|
{
|
|
graph.control.addElement(Sequence(), {assignChain[i - 1]}, {assignChain[i]});
|
|
}
|
|
|
|
auto nop = graph.control.addElement(NOP());
|
|
graph.control.addElement(Sequence(), {assignChain.back()}, {nop});
|
|
|
|
return nop;
|
|
}
|
|
|
|
//
|
|
// AddStreamK methods
|
|
//
|
|
|
|
std::string AddStreamK::name() const
|
|
{
|
|
return concatenate("AddStreamK");
|
|
}
|
|
|
|
AddStreamK::AddStreamK(ContextPtr context,
|
|
CommandParametersPtr params,
|
|
std::string const& topLoop,
|
|
std::string const& accumulatorLoop,
|
|
ExpressionPtr numWGs)
|
|
: m_context(context)
|
|
, m_params(params)
|
|
, m_dimensionIndices(params->loopOverOutputTilesDimensions)
|
|
, m_topLoop(topLoop)
|
|
, m_accumulatorLoop(accumulatorLoop)
|
|
, m_numWGs(numWGs)
|
|
{
|
|
}
|
|
|
|
/**
|
|
* Add coordinate transforms that go between:
|
|
* 1. the global tile space and the SK tile space,
|
|
* 2. the global tile space and the DP tile space, and
|
|
* 3. the SK/DP tile spaces and the various MacroTileNumber coordinates.
|
|
*/
|
|
StreamKInfo addTileSpaceCT(KernelGraph& graph,
|
|
bool forward,
|
|
LoopInfo const& loopInfo,
|
|
AccumulatorInfo const& accumInfo,
|
|
ArgumentInfo const& argInfo,
|
|
ExpressionPtr numTotalTiles,
|
|
StreamKConfig streamKMode)
|
|
{
|
|
// Create forward/reverse tile-numbers for each dimension
|
|
// and attach to all staged tile-number coordinates
|
|
std::vector<int> tileNumbers;
|
|
for(auto d : loopInfo.dimensionIndices)
|
|
{
|
|
// Skip if no dangling MacroTileNumber at this dimension
|
|
if(accumInfo.tileNumberCoords.at(d).empty())
|
|
continue;
|
|
|
|
auto tileNumber = graph.coordinates.addElement(
|
|
MacroTileNumber(d, argInfo.numTiles[d], nullptr));
|
|
|
|
for(auto tileNumTag : accumInfo.tileNumberCoords.at(d))
|
|
{
|
|
if(forward
|
|
&& std::empty(graph.coordinates.getNeighbours<GD::Downstream>(tileNumTag)))
|
|
graph.coordinates.addElement(PassThrough(), {tileNumTag}, {tileNumber});
|
|
if(!forward
|
|
&& std::empty(graph.coordinates.getNeighbours<GD::Upstream>(tileNumTag)))
|
|
graph.coordinates.addElement(PassThrough(), {tileNumber}, {tileNumTag});
|
|
}
|
|
|
|
tileNumbers.push_back(tileNumber);
|
|
}
|
|
|
|
// Create forward/reverse accumulator tile-numbers
|
|
//
|
|
// Appending means: accumulating dimension is fastest moving
|
|
//
|
|
tileNumbers.push_back(accumInfo.accumulatorCoord);
|
|
|
|
// Create forward/reverse flattened tile-space
|
|
auto tileSpace = graph.coordinates.addElement(Linear(numTotalTiles, nullptr));
|
|
auto tileSpaceSK = graph.coordinates.addElement(Linear(argInfo.numSKTiles, nullptr));
|
|
|
|
auto WGs = graph.coordinates.addElement(Linear(argInfo.numWGs, nullptr));
|
|
auto workgroup = graph.coordinates.addElement(Workgroup());
|
|
|
|
auto localTileSpaceSK
|
|
= graph.coordinates.addElement(Linear(argInfo.numSKTilesPerWG, nullptr));
|
|
|
|
std::optional<int> localTileSpaceDP;
|
|
std::optional<int> selector;
|
|
|
|
if(streamKMode.isTwoTileMode())
|
|
{
|
|
bool dpFirst = streamKMode == StreamKMode::TwoTileDPFirst;
|
|
|
|
auto tileSpaceDP
|
|
= graph.coordinates.addElement(Linear(argInfo.numDPTiles, nullptr));
|
|
auto selectorCoord = graph.coordinates.addElement(Linear(literal(2u), nullptr));
|
|
auto localDP
|
|
= graph.coordinates.addElement(Linear(argInfo.numDPTilesPerWG, nullptr));
|
|
|
|
if(forward)
|
|
{
|
|
graph.coordinates.addElement(
|
|
Flatten(), tileNumbers, std::vector<int>{tileSpace});
|
|
if(dpFirst)
|
|
{
|
|
graph.coordinates.addElement(
|
|
Sunder(), {tileSpace}, {tileSpaceDP, tileSpaceSK, selectorCoord});
|
|
}
|
|
else
|
|
{
|
|
graph.coordinates.addElement(
|
|
Sunder(), {tileSpace}, {tileSpaceSK, tileSpaceDP, selectorCoord});
|
|
}
|
|
graph.coordinates.addElement(PassThrough(), {WGs}, {workgroup});
|
|
graph.coordinates.addElement(Tile(), {tileSpaceSK}, {WGs, localTileSpaceSK});
|
|
graph.coordinates.addElement(Tile(), {tileSpaceDP}, {WGs, localDP});
|
|
}
|
|
else
|
|
{
|
|
if(dpFirst)
|
|
{
|
|
graph.coordinates.addElement(
|
|
Sunder(), {tileSpaceDP, tileSpaceSK, selectorCoord}, {tileSpace});
|
|
}
|
|
else
|
|
{
|
|
graph.coordinates.addElement(
|
|
Sunder(), {tileSpaceSK, tileSpaceDP, selectorCoord}, {tileSpace});
|
|
}
|
|
graph.coordinates.addElement(Tile(), std::vector<int>{tileSpace}, tileNumbers);
|
|
graph.coordinates.addElement(PassThrough(), {workgroup}, {WGs});
|
|
graph.coordinates.addElement(Flatten(), {WGs, localTileSpaceSK}, {tileSpaceSK});
|
|
graph.coordinates.addElement(Flatten(), {WGs, localDP}, {tileSpaceDP});
|
|
}
|
|
|
|
localTileSpaceDP = localDP;
|
|
selector = selectorCoord;
|
|
}
|
|
else
|
|
{
|
|
// Standard mode: entire tile space is SK, no Sunder needed.
|
|
if(forward)
|
|
{
|
|
graph.coordinates.addElement(
|
|
Flatten(), tileNumbers, std::vector<int>{tileSpace});
|
|
graph.coordinates.addElement(PassThrough(), {tileSpace}, {tileSpaceSK});
|
|
graph.coordinates.addElement(PassThrough(), {WGs}, {workgroup});
|
|
graph.coordinates.addElement(Tile(), {tileSpaceSK}, {WGs, localTileSpaceSK});
|
|
}
|
|
else
|
|
{
|
|
graph.coordinates.addElement(PassThrough(), {tileSpaceSK}, {tileSpace});
|
|
graph.coordinates.addElement(Tile(), std::vector<int>{tileSpace}, tileNumbers);
|
|
graph.coordinates.addElement(PassThrough(), {workgroup}, {WGs});
|
|
graph.coordinates.addElement(Flatten(), {WGs, localTileSpaceSK}, {tileSpaceSK});
|
|
}
|
|
}
|
|
|
|
return {localTileSpaceSK, localTileSpaceDP, selector};
|
|
}
|
|
|
|
//
|
|
// Stage
|
|
//
|
|
// Look for all leaf MacroTileNumbers with matching
|
|
// sub-dimension.
|
|
//
|
|
// Matches are: tile->dim in m_dimensions.
|
|
//
|
|
AccumulatorInfo stage(KernelGraph const& graph, LoopInfo const& loopInfo)
|
|
{
|
|
AccumulatorInfo accumInfo;
|
|
|
|
accumInfo.accumulatorCoord = getForLoopCoords(loopInfo.accumulatorLoopOp, graph).first;
|
|
Log::debug(" accumulator loop coord: {}", accumInfo.accumulatorCoord);
|
|
|
|
// Find accumulator tile: look above accumulator-loop for an assign statement
|
|
accumInfo.accumulatorTile = -1;
|
|
auto maybeAccumulatorInit = only(graph.control.findElements([&](int tag) -> bool {
|
|
auto maybeAssign = graph.control.get<Assign>(tag);
|
|
if(!maybeAssign)
|
|
return false;
|
|
auto nextTag = only(graph.control.getOutputNodeIndices<Sequence>(tag));
|
|
return nextTag.value_or(-1) == loopInfo.accumulatorLoopOp;
|
|
}));
|
|
if(maybeAccumulatorInit)
|
|
{
|
|
auto init = graph.control.get<Assign>(*maybeAccumulatorInit);
|
|
accumInfo.accumulatorVarType = resultVariableType(init->expression);
|
|
|
|
auto dst = graph.mapper.get(*maybeAccumulatorInit, NaryArgument::DEST);
|
|
auto maybeAccumTile = graph.coordinates.get<MacroTile>(dst);
|
|
if(maybeAccumTile)
|
|
accumInfo.accumulatorTile = dst;
|
|
}
|
|
|
|
ControlFlowRWTracer tracer(graph);
|
|
for(auto m : tracer.coordinatesReadWrite(accumInfo.accumulatorTile))
|
|
{
|
|
if(graph.control.compareNodes(
|
|
rocRoller::UseCacheIfAvailable, loopInfo.accumulatorLoopOp, m.control)
|
|
== NodeOrdering::LeftFirst)
|
|
{
|
|
accumInfo.usesAccumulatorTile.insert(m.control);
|
|
}
|
|
}
|
|
|
|
// Find all dangling MacroTileNumber dimensions associated
|
|
// with the requested dimensions
|
|
for(auto dimension : loopInfo.dimensionIndices)
|
|
{
|
|
auto danglingTileNumberPredicate = [&](int tag) {
|
|
auto maybeTileNumber = graph.coordinates.get<MacroTileNumber>(tag);
|
|
if(!maybeTileNumber)
|
|
return false;
|
|
if(maybeTileNumber->dim != dimension)
|
|
return false;
|
|
if(std::empty(graph.coordinates.getNeighbours<GD::Upstream>(tag))
|
|
|| std::empty(graph.coordinates.getNeighbours<GD::Downstream>(tag)))
|
|
return true;
|
|
return false;
|
|
};
|
|
|
|
// NOTE: a dimension might not have any dangling MacroTileNumber.
|
|
// This happens when workgroupMapping (RemapOutputTiles)
|
|
// applied as M & N dimensions get flattened int a new MacroTileNumber.
|
|
accumInfo.tileNumberCoords[dimension]
|
|
= graph.coordinates.findElements(danglingTileNumberPredicate)
|
|
.to<std::unordered_set>();
|
|
}
|
|
|
|
return accumInfo;
|
|
}
|
|
|
|
//
|
|
// Commit
|
|
//
|
|
void commit(ContextPtr context,
|
|
CommandParametersPtr params,
|
|
KernelGraph& graph,
|
|
LoopInfo const& loopInfo,
|
|
AccumulatorInfo const& accumInfo,
|
|
ArgumentInfo const& argInfo)
|
|
{
|
|
//
|
|
// Find the epilogue operations; detach them.
|
|
//
|
|
// They will be re-attached later.
|
|
//
|
|
std::vector<int> epilogueOperations;
|
|
for(auto tag : graph.control.getNeighbours<GD::Downstream>(loopInfo.topLoopOp))
|
|
{
|
|
auto maybeSequence = graph.control.get<Sequence>(tag);
|
|
if(maybeSequence)
|
|
{
|
|
epilogueOperations.push_back(
|
|
only(graph.control.getNeighbours<GD::Downstream>(tag)).value());
|
|
graph.control.deleteElement(tag);
|
|
}
|
|
}
|
|
|
|
//
|
|
// Create new Scope and insert it above the top-loop
|
|
//
|
|
auto scope = graph.control.addElement(Scope());
|
|
replaceWith(graph, loopInfo.topLoopOp, scope, false);
|
|
|
|
//
|
|
// Compute size of global and local tile-spaces
|
|
//
|
|
auto numAccumTiles = argInfo.numTileArgExprs.back();
|
|
auto numTotalTiles = numAccumTiles;
|
|
for(auto d : loopInfo.dimensionIndices)
|
|
numTotalTiles = numTotalTiles * argInfo.numTiles.at(d);
|
|
numTotalTiles = simplify(numTotalTiles);
|
|
|
|
auto numTilesVarType = resultType(numTotalTiles).varType;
|
|
auto one = Expression::literal(1, numTilesVarType);
|
|
auto zero = Expression::literal(0, numTilesVarType);
|
|
|
|
Log::debug(" numAccumTiles: {}", toString(numAccumTiles));
|
|
Log::debug(" numTotalTiles: {}", toString(numTotalTiles));
|
|
Log::debug(" varType: {}", toString(numTilesVarType));
|
|
|
|
//
|
|
// Helper
|
|
//
|
|
auto DF = [numTilesVarType](int tag) {
|
|
return std::make_shared<Expression::Expression>(
|
|
Expression::DataFlowTag{tag, Register::Type::Scalar, numTilesVarType});
|
|
};
|
|
|
|
auto wg = graph.coordinates.addElement(Workgroup(0));
|
|
auto wgExpr = std::make_shared<Expression::Expression>(
|
|
Expression::DataFlowTag{wg, Register::Type::Scalar, DataType::UInt32});
|
|
|
|
//
|
|
// Duplicate operations from the top-loop down for DP section of two-tile StreamK.
|
|
//
|
|
// TODO This seems overly aggressive. There should be a
|
|
// way of looking at the DataFlow graph to figure out
|
|
// where to stop.
|
|
int dpTopLoop, dpAccumLoop;
|
|
if(params->streamK.isTwoTileMode())
|
|
{
|
|
// Keep the accumulator tile shared between SK and DP sections
|
|
// so that the fix-up pass uses the correct registers.
|
|
auto accumTile = accumInfo.accumulatorTile;
|
|
auto reindexer = std::make_shared<GraphReindexer>();
|
|
dpTopLoop
|
|
= only(duplicateControlNodes(graph,
|
|
reindexer,
|
|
{loopInfo.topLoopOp},
|
|
[accumTile](int x) { return x == accumTile; }))
|
|
.value();
|
|
dpAccumLoop = reindexer->control[loopInfo.accumulatorLoopOp];
|
|
|
|
// Duplicate the epilogue. We enable duplication here
|
|
// to help the deallocation pass reduce register
|
|
// pressure (shortens lifetimes)
|
|
auto epilogueDuplicates = duplicateControlNodes(
|
|
graph, nullptr, epilogueOperations, [](int x) { return false; });
|
|
for(auto const& epilogueDuplicate : epilogueDuplicates)
|
|
graph.control.addElement(Sequence(), {dpTopLoop}, {epilogueDuplicate});
|
|
}
|
|
|
|
//
|
|
// Add forward/reverse tile-space coordinate transforms
|
|
//
|
|
|
|
auto forwardInfo = addTileSpaceCT(
|
|
graph, true, loopInfo, accumInfo, argInfo, numTotalTiles, params->streamK);
|
|
auto reverseInfo = addTileSpaceCT(
|
|
graph, false, loopInfo, accumInfo, argInfo, numTotalTiles, params->streamK);
|
|
|
|
// Hijack the old accumulator loop. Remove the DataFlow edge
|
|
// from the accumulator increment counter to the original
|
|
// accumulator ForLoop coordinate.
|
|
auto forAccumIncr = only(graph.coordinates.getInputNodeIndices(
|
|
accumInfo.accumulatorCoord, CG::isEdge<CG::DataFlow>))
|
|
.value();
|
|
{
|
|
auto dataflow
|
|
= only(graph.coordinates.getNeighbours<GD::Downstream>(forAccumIncr)).value();
|
|
graph.coordinates.deleteElement(dataflow);
|
|
}
|
|
|
|
// Create new forward/reverse accum ForLoop dims, shared by both
|
|
// SK and DP (both iterate over the same number of accum tiles).
|
|
auto forwardForAccumIdx = graph.coordinates.addElement(ForLoop(numAccumTiles, one));
|
|
auto reverseForAccumIdx = graph.coordinates.addElement(ForLoop(numAccumTiles, one));
|
|
|
|
// Wire the accumulator increment counter to the new dims.
|
|
graph.coordinates.addElement(DataFlow(), {forAccumIncr}, {forwardForAccumIdx});
|
|
graph.coordinates.addElement(DataFlow(), {forAccumIncr}, {reverseForAccumIdx});
|
|
|
|
// Identify edges let index computation traverse from the new
|
|
// ForLoop dims back to the original coordinate for offsets/strides.
|
|
graph.coordinates.addElement(
|
|
Identify(), {forwardForAccumIdx}, {accumInfo.accumulatorCoord});
|
|
graph.coordinates.addElement(
|
|
Identify(), {reverseForAccumIdx}, {accumInfo.accumulatorCoord});
|
|
|
|
graph.mapper.connect<ForLoop>(loopInfo.accumulatorLoopOp, forwardForAccumIdx);
|
|
|
|
// Create tile ForLoop dims and connect them to the local tile
|
|
// spaces via Split/Join. SK and DP get separate tile dims
|
|
// but share the accum dims created above.
|
|
|
|
auto skForTileIncr = graph.coordinates.addElement(Linear(argInfo.numSKTilesPerWG, one));
|
|
auto skForwardForTileIdx
|
|
= graph.coordinates.addElement(ForLoop(argInfo.numSKTilesPerWG, one));
|
|
auto skReverseForTileIdx
|
|
= graph.coordinates.addElement(ForLoop(argInfo.numSKTilesPerWG, one));
|
|
|
|
Log::debug(
|
|
" SK tile ForLoops: fwd {} rev {}", skForwardForTileIdx, skReverseForTileIdx);
|
|
Log::debug(
|
|
" Shared accum ForLoops: fwd {} rev {}", forwardForAccumIdx, reverseForAccumIdx);
|
|
|
|
graph.coordinates.addElement(
|
|
Split(), {forwardInfo.localTileSpaceSK}, {skForwardForTileIdx, forwardForAccumIdx});
|
|
graph.coordinates.addElement(
|
|
Join(), {skReverseForTileIdx, reverseForAccumIdx}, {reverseInfo.localTileSpaceSK});
|
|
|
|
graph.coordinates.addElement(DataFlow(), {skForTileIncr}, {skForwardForTileIdx});
|
|
graph.coordinates.addElement(DataFlow(), {skForTileIncr}, {skReverseForTileIdx});
|
|
|
|
// DP tile ForLoop dims (TwoTile only)
|
|
std::optional<int> dpForTileIncr;
|
|
std::optional<int> dpForwardForTileIdx;
|
|
std::optional<int> dpReverseForTileIdx;
|
|
|
|
if(params->streamK.isTwoTileMode())
|
|
{
|
|
dpForTileIncr = graph.coordinates.addElement(Linear(argInfo.numDPTilesPerWG, one));
|
|
|
|
dpForwardForTileIdx
|
|
= graph.coordinates.addElement(ForLoop(argInfo.numDPTilesPerWG, one));
|
|
dpReverseForTileIdx
|
|
= graph.coordinates.addElement(ForLoop(argInfo.numDPTilesPerWG, one));
|
|
|
|
Log::debug(" DP tile ForLoops: fwd {} rev {}",
|
|
dpForwardForTileIdx.value(),
|
|
dpReverseForTileIdx.value());
|
|
|
|
graph.coordinates.addElement(Split(),
|
|
{forwardInfo.localTileSpaceDP.value()},
|
|
{dpForwardForTileIdx.value(), forwardForAccumIdx});
|
|
graph.coordinates.addElement(Join(),
|
|
{dpReverseForTileIdx.value(), reverseForAccumIdx},
|
|
{reverseInfo.localTileSpaceDP.value()});
|
|
|
|
graph.coordinates.addElement(
|
|
DataFlow(), {dpForTileIncr.value()}, {dpForwardForTileIdx.value()});
|
|
graph.coordinates.addElement(
|
|
DataFlow(), {dpForTileIncr.value()}, {dpReverseForTileIdx.value()});
|
|
}
|
|
|
|
// The original accumulator coordinate no longer has a DataFlow
|
|
// input. Replace it with a Linear so that loop-related passes
|
|
// don't treat it as an active ForLoop iterator.
|
|
{
|
|
auto accumulatorCoordSize
|
|
= getSize(graph.coordinates.getNode(accumInfo.accumulatorCoord));
|
|
auto linear = Linear(accumulatorCoordSize, nullptr);
|
|
graph.coordinates.setElement(accumInfo.accumulatorCoord, linear);
|
|
}
|
|
|
|
// Pre-compute and save (into SGPRs):
|
|
//
|
|
// 1. Number of accumulation tiles processed in the accumulator loop
|
|
// 2. The first accumulator tile index processed
|
|
// 3. The last accumulator tile index processed
|
|
//
|
|
// The matching Assign operations are created below
|
|
auto currentTile = graph.coordinates.addElement(VGPR());
|
|
auto numAccumTilesProcessed = graph.coordinates.addElement(VGPR());
|
|
auto firstAccumTile = graph.coordinates.addElement(VGPR());
|
|
auto lastAccumTile = graph.coordinates.addElement(VGPR());
|
|
|
|
//
|
|
// Create local-tile loop
|
|
//
|
|
int forTileSKOp;
|
|
{
|
|
auto numSKTilesPerWGBound = argInfo.numSKTilesPerWG;
|
|
auto numSKTilesBound = conditionalSubtract(
|
|
argInfo.numSKTiles, argInfo.numSKTilesPerWG * wgExpr, numTilesVarType.dataType);
|
|
auto forTileSKConditionExpr
|
|
= DF(skForTileIncr) < minimum(numSKTilesPerWGBound, numSKTilesBound);
|
|
|
|
auto forTileSKIncrementExpr = DF(skForTileIncr) + DF(numAccumTilesProcessed);
|
|
|
|
forTileSKOp = conditionalFor(graph,
|
|
skForTileIncr,
|
|
forTileSKConditionExpr,
|
|
forTileSKIncrementExpr,
|
|
"SKStreamTileLoop",
|
|
numTilesVarType);
|
|
}
|
|
|
|
int assignCurrentTile;
|
|
{
|
|
assignCurrentTile = graph.control.addElement(Assign{
|
|
Register::Type::Scalar, argInfo.numSKTilesPerWG * wgExpr + DF(skForTileIncr)});
|
|
graph.mapper.connect(assignCurrentTile, currentTile, NaryArgument::DEST);
|
|
}
|
|
|
|
// Compute the starting K tile index within the current output tile.
|
|
// firstAccumTile = currentTile % numAccumTiles gives which K tile (0 to numAccumTiles-1)
|
|
// we begin processing within this (M, N) output tile.
|
|
int assignFirstAccumTile;
|
|
{
|
|
auto firstAccumTileExpr = DF(currentTile) % numAccumTiles;
|
|
assignFirstAccumTile
|
|
= graph.control.addElement(Assign{Register::Type::Scalar, firstAccumTileExpr});
|
|
graph.mapper.connect(assignFirstAccumTile, firstAccumTile, NaryArgument::DEST);
|
|
}
|
|
|
|
int assignNumAccumTilesProcessed;
|
|
{
|
|
// Compute how many K tiles to process in this iteration.
|
|
// sameNonAccumTileBound = numAccumTiles - firstAccumTile is the number of K tiles
|
|
// remaining in the current output tile. We take the minimum of this and other bounds.
|
|
auto sameNonAccumTileBound = conditionalSubtract(
|
|
numAccumTiles, DF(firstAccumTile), numTilesVarType.dataType);
|
|
auto numSKTilesPerWGBound = conditionalSubtract(
|
|
argInfo.numSKTilesPerWG, DF(skForTileIncr), numTilesVarType.dataType);
|
|
auto numSKTilesExprBound = conditionalSubtract(
|
|
argInfo.numSKTiles, DF(currentTile), numTilesVarType.dataType);
|
|
|
|
assignNumAccumTilesProcessed = graph.control.addElement(Assign{
|
|
Register::Type::Scalar,
|
|
minimum(sameNonAccumTileBound, numSKTilesPerWGBound, numSKTilesExprBound)});
|
|
graph.mapper.connect(
|
|
assignNumAccumTilesProcessed, numAccumTilesProcessed, NaryArgument::DEST);
|
|
}
|
|
|
|
int assignLastAccumTile;
|
|
{
|
|
// Compute the last K tile index processed within this output tile.
|
|
// Since numAccumTilesProcessed <= numAccumTiles - firstAccumTile (bounded by
|
|
// sameNonAccumTileBound to stay within the same output tile), the result is
|
|
// guaranteed to be in range [0, numAccumTiles - 1].
|
|
auto lastAccumTileExpr = DF(firstAccumTile) + DF(numAccumTilesProcessed) - one;
|
|
assignLastAccumTile
|
|
= graph.control.addElement(Assign{Register::Type::Scalar, lastAccumTileExpr});
|
|
graph.mapper.connect(assignLastAccumTile, lastAccumTile, NaryArgument::DEST);
|
|
}
|
|
|
|
auto forAccumOp = *graph.control.get<ForLoopOp>(loopInfo.accumulatorLoopOp);
|
|
forAccumOp.condition = DF(forAccumIncr) < DF(numAccumTilesProcessed);
|
|
graph.control.setElement(loopInfo.accumulatorLoopOp, forAccumOp);
|
|
|
|
//
|
|
// After the accumulator loop, we might need to send and/or
|
|
// receive partially accumulated tiles.
|
|
//
|
|
ScratchInfo scratchTileInfo;
|
|
SendInfo sendInfo;
|
|
RecvInfo receiveInfo;
|
|
int postAccumulationCond;
|
|
if(accumInfo.accumulatorTile != -1)
|
|
{
|
|
auto remainAccumTiles = numAccumTiles - DF(lastAccumTile) - one;
|
|
|
|
// For loop that receives and accumulates partial results from other workgroups.
|
|
// The loop iterates ceil(remainAccumTiles / numSKTilesPerWG) times.
|
|
// Using the equivalence: i < ceil(a/b) iff i*b < a (for a >= 0, b > 0),
|
|
// we express the loop condition as a multiplication to avoid runtime division.
|
|
auto iteratorCoord = graph.coordinates.addElement(Linear(numAccumTiles, one));
|
|
auto forReceiveTileLoopCoord
|
|
= graph.coordinates.addElement(ForLoop(numAccumTiles, one));
|
|
graph.coordinates.addElement(
|
|
DataFlow(), {iteratorCoord}, {forReceiveTileLoopCoord});
|
|
|
|
auto iterator = std::make_shared<Expression::Expression>(Expression::DataFlowTag{
|
|
iteratorCoord, Register::Type::Scalar, numTilesVarType});
|
|
|
|
// Loop condition: iterates while iterator * numSKTilesPerWG < remainAccumTiles
|
|
// This is equivalent to iterator < ceil(remainAccumTiles / numSKTilesPerWG)
|
|
auto loopCondition = iterator * argInfo.numSKTilesPerWG < remainAccumTiles;
|
|
|
|
auto forReceiveTileLoopOp
|
|
= graph.control.addElement(ForLoopOp{loopCondition, rocRoller::RECEIVE});
|
|
graph.mapper.connect(forReceiveTileLoopOp, iteratorCoord, NaryArgument::DEST);
|
|
graph.mapper.connect<ForLoop>(forReceiveTileLoopOp, forReceiveTileLoopCoord);
|
|
|
|
auto initialAssign = graph.control.addElement(
|
|
Assign{Register::Type::Scalar, Expression::literal(0, numTilesVarType)});
|
|
graph.mapper.connect(initialAssign, iteratorCoord, NaryArgument::DEST);
|
|
|
|
auto incrementAssign
|
|
= graph.control.addElement(Assign{Register::Type::Scalar, iterator + one});
|
|
graph.mapper.connect(incrementAssign, iteratorCoord, NaryArgument::DEST);
|
|
|
|
graph.control.addElement(Initialize(), {forReceiveTileLoopOp}, {initialAssign});
|
|
graph.control.addElement(
|
|
ForLoopIncrement(), {forReceiveTileLoopOp}, {incrementAssign});
|
|
|
|
// Create scratch space for flags
|
|
auto flagsScratch
|
|
= newScratchCoordinate(argInfo.numWGs,
|
|
DataType::UInt32,
|
|
Operations::ScratchPolicy::ZeroedBeforeAndAfter,
|
|
context);
|
|
auto flagsScratchTag = graph.coordinates.addElement(flagsScratch);
|
|
|
|
// Create scratch space for partially accumulated tiles
|
|
std::vector<DeferredConnection> storeConnections, loadConnections;
|
|
scratchTileInfo = loadStoreMacroTileSCRATCH(graph,
|
|
storeConnections,
|
|
loadConnections,
|
|
accumInfo.accumulatorTile,
|
|
accumInfo.accumulatorVarType,
|
|
loopInfo,
|
|
argInfo,
|
|
forReceiveTileLoopCoord,
|
|
params,
|
|
context);
|
|
|
|
// Add send and receive
|
|
auto hasFirstAccumTile = DF(firstAccumTile) == zero;
|
|
auto doesntHaveFirstAccumTile = DF(firstAccumTile) != zero;
|
|
auto doesntHaveLastAccumTile = DF(lastAccumTile) < (numAccumTiles - one);
|
|
|
|
sendInfo = sendTile(graph,
|
|
doesntHaveFirstAccumTile,
|
|
storeConnections,
|
|
flagsScratchTag,
|
|
accumInfo.accumulatorVarType.dataType,
|
|
numTilesVarType,
|
|
loopInfo,
|
|
context);
|
|
|
|
receiveInfo = receiveTile(graph,
|
|
hasFirstAccumTile && doesntHaveLastAccumTile,
|
|
scratchTileInfo.load,
|
|
loadConnections,
|
|
flagsScratchTag,
|
|
accumInfo.accumulatorTile,
|
|
accumInfo.usesAccumulatorTile,
|
|
accumInfo.accumulatorVarType.dataType,
|
|
loopInfo,
|
|
argInfo,
|
|
forReceiveTileLoopOp,
|
|
forReceiveTileLoopCoord,
|
|
params,
|
|
context);
|
|
|
|
postAccumulationCond = graph.control.addElement(ConditionalOp{
|
|
hasFirstAccumTile, ConditionalMode::Branch, "Post-accumulation Condition"});
|
|
}
|
|
else
|
|
{
|
|
scratchTileInfo.setPlusOne = graph.control.addElement(Scope());
|
|
sendInfo.preWaitZero = graph.control.addElement(NOP());
|
|
sendInfo.sendCond = graph.control.addElement(NOP());
|
|
receiveInfo.preWaitZero = graph.control.addElement(NOP());
|
|
receiveInfo.receiveCond = graph.control.addElement(NOP());
|
|
receiveInfo.setPlusOne = graph.control.addElement(Scope());
|
|
postAccumulationCond = graph.control.addElement(Scope());
|
|
|
|
graph.control.addElement(Sequence(), {sendInfo.preWaitZero}, {sendInfo.sendCond});
|
|
graph.control.addElement(
|
|
Sequence(), {receiveInfo.preWaitZero}, {receiveInfo.receiveCond});
|
|
}
|
|
graph.control.addElement(Sequence(), {receiveInfo.receiveCond}, {postAccumulationCond});
|
|
|
|
// Add local-tile loop and accumulator loop
|
|
if(params->streamK.isTwoTileMode())
|
|
{
|
|
bool dpFirst = params->streamK == StreamKMode::TwoTileDPFirst;
|
|
int scopeSK = graph.control.addElement(Scope());
|
|
int scopeDP = graph.control.addElement(Scope());
|
|
|
|
// Connect scopes in the appropriate order.
|
|
if(dpFirst)
|
|
{
|
|
graph.control.addElement(Sequence(), {scope}, {scopeDP});
|
|
graph.control.addElement(Sequence(), {scopeDP}, {scopeSK});
|
|
}
|
|
else
|
|
{
|
|
graph.control.addElement(Sequence(), {scope}, {scopeSK});
|
|
graph.control.addElement(Sequence(), {scopeSK}, {scopeDP});
|
|
}
|
|
|
|
auto skSelectorValue = dpFirst ? one : zero;
|
|
auto dpSelectorValue = dpFirst ? zero : one;
|
|
|
|
// Initialize SK section.
|
|
// The DP tile dims are also initialized to zero here so
|
|
// that the reverse Sunder can be fully traversed. The
|
|
// selector ensures only the SK branch is evaluated at
|
|
// runtime.
|
|
auto skInit = initializeCoordinates(
|
|
graph,
|
|
scopeSK,
|
|
{{{skForwardForTileIdx,
|
|
skReverseForTileIdx,
|
|
dpForwardForTileIdx.value(),
|
|
dpReverseForTileIdx.value(),
|
|
forwardForAccumIdx,
|
|
reverseForAccumIdx},
|
|
zero},
|
|
{{forwardInfo.selector.value(), reverseInfo.selector.value()},
|
|
skSelectorValue}},
|
|
{skForTileIncr,
|
|
forAccumIncr,
|
|
currentTile,
|
|
numAccumTilesProcessed,
|
|
firstAccumTile,
|
|
lastAccumTile},
|
|
zero);
|
|
|
|
// Chain SK-specific setup inside scopeSK
|
|
graph.control.chain<Sequence>(skInit, receiveInfo.setPlusOne);
|
|
graph.control.chain<Body>(
|
|
receiveInfo.setPlusOne, scratchTileInfo.setPlusOne, forTileSKOp);
|
|
graph.control.chain<Body>(forTileSKOp, assignCurrentTile);
|
|
graph.control.chain<Sequence>(assignCurrentTile,
|
|
assignFirstAccumTile,
|
|
assignNumAccumTilesProcessed,
|
|
assignLastAccumTile,
|
|
loopInfo.topLoopOp,
|
|
sendInfo.preWaitZero);
|
|
|
|
graph.control.chain<Sequence>(sendInfo.sendCond, receiveInfo.preWaitZero);
|
|
|
|
// Initialize DP section.
|
|
// The SK tile dims are also initialized to zero here so
|
|
// that the reverse Sunder can be fully traversed. The
|
|
// selector ensures only the DP branch is evaluated at
|
|
// runtime.
|
|
auto dpInit = initializeCoordinates(
|
|
graph,
|
|
scopeDP,
|
|
{{{skForwardForTileIdx,
|
|
skReverseForTileIdx,
|
|
dpForwardForTileIdx.value(),
|
|
dpReverseForTileIdx.value(),
|
|
forwardForAccumIdx,
|
|
reverseForAccumIdx},
|
|
zero},
|
|
{{forwardInfo.selector.value(), reverseInfo.selector.value()},
|
|
dpSelectorValue}},
|
|
{dpForTileIncr.value(), forAccumIncr},
|
|
zero);
|
|
|
|
// Create DP tile loop (uses its own dpForTileIncr counter).
|
|
int forTileDPOp;
|
|
{
|
|
auto numDPTilesPerWGBound = argInfo.numDPTilesPerWG;
|
|
auto numDPTilesBound = conditionalSubtract(argInfo.numDPTiles,
|
|
argInfo.numDPTilesPerWG * wgExpr,
|
|
numTilesVarType.dataType);
|
|
auto forTileDPConditionExpr = DF(dpForTileIncr.value())
|
|
< minimum(numDPTilesPerWGBound, numDPTilesBound);
|
|
|
|
auto forTileDPIncrementExpr = DF(dpForTileIncr.value()) + DF(forAccumIncr);
|
|
|
|
forTileDPOp = conditionalFor(graph,
|
|
dpForTileIncr.value(),
|
|
forTileDPConditionExpr,
|
|
forTileDPIncrementExpr,
|
|
"DPStreamTileLoop",
|
|
numTilesVarType);
|
|
}
|
|
|
|
graph.control.addElement(Sequence(), {dpInit}, {forTileDPOp});
|
|
graph.control.addElement(Body(), {forTileDPOp}, {dpTopLoop});
|
|
|
|
auto forAccumOp = *graph.control.get<ForLoopOp>(dpAccumLoop);
|
|
forAccumOp.condition = DF(forAccumIncr) < numAccumTiles;
|
|
graph.control.setElement(dpAccumLoop, forAccumOp);
|
|
|
|
graph.mapper.connect<ForLoop>(dpAccumLoop, forwardForAccumIdx);
|
|
}
|
|
else
|
|
{
|
|
// Standard mode -- no selector, no DP dims
|
|
auto standardInit = initializeCoordinates(graph,
|
|
scope,
|
|
{{{skForwardForTileIdx,
|
|
skReverseForTileIdx,
|
|
forwardForAccumIdx,
|
|
reverseForAccumIdx},
|
|
zero}},
|
|
{skForTileIncr,
|
|
forAccumIncr,
|
|
currentTile,
|
|
numAccumTilesProcessed,
|
|
firstAccumTile,
|
|
lastAccumTile},
|
|
zero);
|
|
|
|
graph.control.chain<Sequence>(standardInit, receiveInfo.setPlusOne);
|
|
graph.control.chain<Body>(
|
|
receiveInfo.setPlusOne, scratchTileInfo.setPlusOne, forTileSKOp);
|
|
graph.control.chain<Body>(forTileSKOp, assignCurrentTile);
|
|
graph.control.chain<Sequence>(assignCurrentTile,
|
|
assignFirstAccumTile,
|
|
assignNumAccumTilesProcessed,
|
|
assignLastAccumTile,
|
|
loopInfo.topLoopOp,
|
|
sendInfo.preWaitZero);
|
|
|
|
graph.control.chain<Sequence>(sendInfo.sendCond, receiveInfo.preWaitZero);
|
|
}
|
|
|
|
//
|
|
// Now can re-attach epilogueOperations. We defer
|
|
// this to now so that code above can duplicate the
|
|
// for-loops.
|
|
//
|
|
// Make sure receive happens before other operations after
|
|
// the accumulator loop.
|
|
//
|
|
for(auto tag : epilogueOperations)
|
|
{
|
|
graph.control.addElement(Body(), {postAccumulationCond}, {tag});
|
|
}
|
|
}
|
|
|
|
ArgumentInfo setupArguments(ContextPtr context,
|
|
CommandParametersPtr params,
|
|
ExpressionPtr numWGs,
|
|
KernelGraph const& graph,
|
|
LoopInfo const& loopInfo,
|
|
AccumulatorInfo const& accumInfo)
|
|
{
|
|
ArgumentInfo argInfo;
|
|
|
|
auto k = context->kernel();
|
|
|
|
argInfo.numTileArgExprs.resize(loopInfo.dimensionIndices.size() + 1);
|
|
|
|
auto numWGsDT = DataType::UInt32;
|
|
auto numTilesDT = DataType::UInt32;
|
|
auto one = Expression::literal(1, numTilesDT);
|
|
auto zero = Expression::literal(0, numTilesDT);
|
|
|
|
// On entry, numWGs is an Expression that either:
|
|
// 1. Pulls a value from a CommandArgument
|
|
// 2. Is a literal (for testing)
|
|
|
|
argInfo.numWGs = k->addArgument(
|
|
{rocRoller::NUMWGS, numWGsDT, DataDirection::ReadOnly, convert(numWGsDT, numWGs)});
|
|
|
|
// Fill number-of-tiles using MacroTileNumber sizes
|
|
// from load operations (store operations are missing
|
|
// that info).
|
|
for(auto dimension : loopInfo.dimensionIndices)
|
|
{
|
|
for(auto tileNumberTag : accumInfo.tileNumberCoords.at(dimension))
|
|
{
|
|
if(!argInfo.numTileArgExprs[dimension])
|
|
{
|
|
auto macTileNumber = graph.coordinates.get<MacroTileNumber>(tileNumberTag);
|
|
if(macTileNumber->size)
|
|
argInfo.numTileArgExprs[dimension]
|
|
= convert(numTilesDT, macTileNumber->size);
|
|
}
|
|
}
|
|
|
|
argInfo.numTileArgExprs.back() = convert(
|
|
numTilesDT, graph.coordinates.get<ForLoop>(accumInfo.accumulatorCoord)->size);
|
|
|
|
for(auto tileNumberTag : accumInfo.tileNumberCoords.at(dimension))
|
|
{
|
|
AssertFatal(argInfo.numTileArgExprs[dimension]);
|
|
Log::debug(" dimension: {} coord: {} size: {}",
|
|
dimension,
|
|
tileNumberTag,
|
|
toString(argInfo.numTileArgExprs[dimension]));
|
|
}
|
|
}
|
|
|
|
// Make kernel arguments:
|
|
//
|
|
// numWGs: Value
|
|
// numTiles0: Computed on host: M / macM
|
|
// numTiles1: Computed on host: N / macN
|
|
// numTilesAcc: Computed on host: K / macK
|
|
// numSKTilesPerWG: Computed on host:
|
|
//
|
|
// for basic StreamK: (numTiles0 * numTiles1 * numTilesAcc + numWGs - 1) / numWGs
|
|
// for 2-tile StreamK: ((numTilesAcc * ((numTiles0 * numTiles1) % numWGs + numWGs) + numWGs - 1) / numWGs
|
|
//
|
|
// numDPTilesPerWG: Computed on host:
|
|
//
|
|
// for basic StreamK: 0
|
|
// for 2-tile StreamK: ((numTiles0 * numTiles1) / numWGs - 1) * numTilesAcc
|
|
//
|
|
for(auto d : loopInfo.dimensionIndices)
|
|
{
|
|
if(argInfo.numTileArgExprs[d] == nullptr)
|
|
continue;
|
|
|
|
argInfo.numTiles.push_back(k->addArgument({concatenate("numTiles", d),
|
|
numTilesDT,
|
|
DataDirection::ReadOnly,
|
|
argInfo.numTileArgExprs[d]}));
|
|
}
|
|
|
|
argInfo.numTiles.push_back(k->addArgument({"numTilesAcc",
|
|
numTilesDT,
|
|
DataDirection::ReadOnly,
|
|
argInfo.numTileArgExprs.back()}));
|
|
|
|
ExpressionPtr numSKTilesArgExpr, numDPTilesArgExpr;
|
|
ExpressionPtr numSKTilesPerWGArgExpr, numDPTilesPerWGArgExpr;
|
|
{
|
|
ExpressionPtr numAccTiles = argInfo.numTileArgExprs.back();
|
|
ExpressionPtr numNonAccTiles = nullptr;
|
|
for(auto d : loopInfo.dimensionIndices)
|
|
{
|
|
if(argInfo.numTileArgExprs[d] == nullptr)
|
|
continue;
|
|
|
|
numNonAccTiles = numNonAccTiles ? numNonAccTiles * argInfo.numTileArgExprs[d]
|
|
: argInfo.numTileArgExprs[d];
|
|
}
|
|
|
|
if(params->streamK.isTwoTileMode())
|
|
{
|
|
auto enoughNonAccTilesExpr = numNonAccTiles > numWGs - one;
|
|
|
|
numSKTilesArgExpr
|
|
= conditional(enoughNonAccTilesExpr,
|
|
(numNonAccTiles % numWGs + numWGs) * numAccTiles,
|
|
numNonAccTiles * numAccTiles);
|
|
|
|
numSKTilesPerWGArgExpr = (numSKTilesArgExpr + numWGs - one) / numWGs;
|
|
numDPTilesArgExpr
|
|
= conditional(enoughNonAccTilesExpr,
|
|
(numNonAccTiles / numWGs - one) * numWGs * numAccTiles,
|
|
zero);
|
|
numDPTilesPerWGArgExpr = conditional(
|
|
enoughNonAccTilesExpr, (numNonAccTiles / numWGs - one) * numAccTiles, zero);
|
|
}
|
|
else
|
|
{
|
|
numSKTilesArgExpr = numNonAccTiles * numAccTiles;
|
|
numSKTilesPerWGArgExpr = (numSKTilesArgExpr + numWGs - one) / numWGs;
|
|
numDPTilesArgExpr = zero;
|
|
numDPTilesPerWGArgExpr = zero;
|
|
}
|
|
}
|
|
|
|
argInfo.numSKTiles = k->addArgument(
|
|
{"numSKTiles", numTilesDT, DataDirection::ReadOnly, numSKTilesArgExpr});
|
|
|
|
argInfo.numSKTilesPerWG = k->addArgument(
|
|
{"numSKTilesPerWG", numTilesDT, DataDirection::ReadOnly, numSKTilesPerWGArgExpr});
|
|
|
|
if(params->streamK.isTwoTileMode())
|
|
{
|
|
argInfo.numDPTiles = k->addArgument(
|
|
{"numDPTiles", numTilesDT, DataDirection::ReadOnly, numDPTilesArgExpr});
|
|
|
|
argInfo.numDPTilesPerWG = k->addArgument({"numDPTilesPerWG",
|
|
numTilesDT,
|
|
DataDirection::ReadOnly,
|
|
numDPTilesPerWGArgExpr});
|
|
}
|
|
else
|
|
{
|
|
argInfo.numDPTiles = zero;
|
|
argInfo.numDPTilesPerWG = zero;
|
|
}
|
|
|
|
return argInfo;
|
|
}
|
|
|
|
KernelGraph AddStreamK::apply(KernelGraph const& original)
|
|
{
|
|
TIMER(t, "KernelGraph::AddStreamK");
|
|
|
|
LoopInfo loopInfo;
|
|
|
|
loopInfo.dimensionIndices = m_dimensionIndices;
|
|
|
|
auto kernel = *original.control.roots().begin();
|
|
|
|
// Find the loop control nodes
|
|
auto maybeTopLoopOp
|
|
= original.control.findNodes(kernel, makeFindLoopPredicate(original, m_topLoop))
|
|
.take(1)
|
|
.only();
|
|
if(!maybeTopLoopOp)
|
|
{
|
|
rocRoller::Log::warn("Unable to find ForLoop '{}' during AddStreamK pass. "
|
|
"AddStreamK transform skipped.",
|
|
m_topLoop);
|
|
return original;
|
|
}
|
|
loopInfo.topLoopOp = *maybeTopLoopOp;
|
|
|
|
auto maybeAccumLoopOp = only(
|
|
original.control.findElements(makeFindLoopPredicate(original, m_accumulatorLoop)));
|
|
if(!maybeAccumLoopOp)
|
|
{
|
|
rocRoller::Log::warn("Unable to find ForLoop '{}' during AddStreamK pass. "
|
|
"AddStreamK transform skipped.",
|
|
m_accumulatorLoop);
|
|
return original;
|
|
}
|
|
loopInfo.accumulatorLoopOp = *maybeAccumLoopOp;
|
|
|
|
auto loopLength = [&](int tag) {
|
|
auto dimTag = original.mapper.get(tag, NaryArgument::DEST);
|
|
auto sizeExpr = getSize(original.coordinates.getNode(dimTag));
|
|
|
|
AssertFatal(
|
|
Expression::evaluationTimes(sizeExpr)[Expression::EvaluationTime::Translate],
|
|
"Invalid length for x or y loop");
|
|
auto length = Expression::evaluate(sizeExpr);
|
|
AssertFatal(isInteger(length), "For loop length should be an integer");
|
|
|
|
return getUnsignedInt(length);
|
|
};
|
|
|
|
auto maybeXLoopOp
|
|
= original.control
|
|
.findNodes(kernel, makeFindLoopPredicate(original, rocRoller::XLOOP))
|
|
.take(1)
|
|
.only();
|
|
if(maybeXLoopOp)
|
|
{
|
|
loopInfo.xLoop = *maybeXLoopOp;
|
|
loopInfo.xLoopSize = loopLength(loopInfo.xLoop);
|
|
}
|
|
|
|
auto maybeYLoopOp
|
|
= original.control
|
|
.findNodes(kernel, makeFindLoopPredicate(original, rocRoller::YLOOP))
|
|
.take(1)
|
|
.only();
|
|
if(maybeYLoopOp)
|
|
{
|
|
loopInfo.yLoop = *maybeYLoopOp;
|
|
loopInfo.yLoopSize = loopLength(loopInfo.yLoop);
|
|
}
|
|
|
|
auto graph = original;
|
|
auto accumInfo = stage(graph, loopInfo);
|
|
|
|
auto argInfo
|
|
= setupArguments(m_context, m_params, m_numWGs, graph, loopInfo, accumInfo);
|
|
|
|
commit(m_context, m_params, graph, loopInfo, accumInfo, argInfo);
|
|
|
|
return graph;
|
|
}
|
|
}
|
|
}
|