Files
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
Flatten subprojects into a single base monorepo
2026-09-02 16:34:47 -07:00

442 lines
20 KiB
C++

// Copyright Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#ifdef ROCROLLER_USE_HIP
#include <hip/hip_ext.h>
#include <hip/hip_runtime.h>
#endif /* ROCROLLER_USE_HIP */
#include <algorithm>
#include <random>
#include <rocRoller/AssemblyKernel.hpp>
#include <rocRoller/CodeGen/ArgumentLoader.hpp>
#include <rocRoller/CommandSolution.hpp>
#include <rocRoller/DataTypes/DataTypes.hpp>
#include <rocRoller/Expression.hpp>
#include <rocRoller/ExpressionTransformations.hpp>
#include <rocRoller/KernelGraph/KernelGraph.hpp>
#include <rocRoller/KernelOptions.hpp>
#include <rocRoller/Operations/Command.hpp>
#include <rocRoller/Scheduling/Observers/FileWritingObserver.hpp>
#include <rocRoller/TensorDescriptor.hpp>
#include <rocRoller/Utilities/Error.hpp>
#include <rocRoller/Utilities/Logging.hpp>
#include <rocRoller/Utilities/Timer.hpp>
#include "GPUContextFixture.hpp"
#include "GenericContextFixture.hpp"
#include "SourceMatcher.hpp"
#include "Utilities.hpp"
#include <common/GEMMProblem.hpp>
#include <common/mxDataGen.hpp>
namespace GEMMDriverTest
{
using namespace rocRoller;
namespace SolutionParams = rocRoller::Parameters::Solution;
// ========================================================================
// GEMMFusionTestSuite
// ========================================================================
struct GEMMFusionTestSuite : public CurrentGPUContextFixture
{
template <typename T>
void basicGEMMRelu(ContextPtr& m_context,
const GEMMProblem& gemm,
bool debuggable = false,
bool setIdentity = false,
int numIters = 1)
{
REQUIRE_ARCH_CAP(GPUCapability::HasMFMA);
// D (MxN) = alpha * A (MxK) X B (KxN) + beta * C (MxN)
int M = gemm.m;
int N = gemm.n;
int K = gemm.k;
float alpha = gemm.alpha;
float beta = gemm.beta;
T reluAlpha = 0.9;
AssertFatal(M % gemm.macM == 0, "MacroTile size mismatch (M)");
AssertFatal(N % gemm.macN == 0, "MacroTile size mismatch (N)");
if(gemm.unrollK > 0)
{
AssertFatal(K % (gemm.macK * gemm.unrollK) == 0,
"MacroTile size mismatch (K unroll)");
}
AssertFatal(gemm.workgroupSizeX % gemm.wavefrontSize == 0,
"Workgroup Size X must be multiply of wave front size");
uint wavetilePerWavefrontM
= gemm.wavefrontSize * gemm.macM / gemm.waveM / gemm.workgroupSizeX;
uint wavetilePerWavefrontN = gemm.macN / gemm.waveN / gemm.workgroupSizeY;
AssertFatal(gemm.macM % (gemm.waveM * wavetilePerWavefrontM) == 0,
"WaveTile size mismatch (M)");
AssertFatal(gemm.macN % (gemm.waveN * wavetilePerWavefrontN) == 0,
"WaveTile size mismatch (N)");
uint workgroupSizeX = gemm.workgroupSizeX * gemm.workgroupSizeY;
uint workgroupSizeY = 1;
uint numWorkgroupX;
uint numWorkgroupY;
if(gemm.loopOverTiles > 0)
{
// multiple output macro tiles per workgroup
numWorkgroupX = M * N / gemm.macM / gemm.macN / 2;
numWorkgroupY = 1;
}
else if(gemm.streamK)
{
numWorkgroupX = gemm.numWGs;
numWorkgroupY = 1;
}
else
{
// one output macro tile per workgroup
numWorkgroupX = M / gemm.macM;
numWorkgroupY = N / gemm.macN;
}
auto NX = std::make_shared<Expression::Expression>(numWorkgroupX * workgroupSizeX);
auto NY = std::make_shared<Expression::Expression>(numWorkgroupY * workgroupSizeY);
auto NZ = std::make_shared<Expression::Expression>(1u);
// Host data
std::vector<T> hostA;
std::vector<T> hostB;
std::vector<T> hostC;
auto dataType = TypeInfo<T>::Var.dataType;
TensorDescriptor descA(dataType, {size_t(M), size_t(K)}, gemm.transA);
TensorDescriptor descB(dataType, {size_t(K), size_t(N)}, gemm.transB);
TensorDescriptor descC(dataType, {size_t(M), size_t(N)}, "N");
TensorDescriptor descRelu(dataType, {size_t(M), size_t(N)}, "N");
auto seed = 31415u;
DGenInput(seed, hostA, descA, hostB, descB, hostC, descC);
if(setIdentity)
{
SetIdentityMatrix(hostA, K, M);
SetIdentityMatrix(hostB, N, K);
std::fill(hostC.begin(), hostC.end(), static_cast<T>(0.0));
}
std::shared_ptr<T> deviceA = make_shared_device(hostA);
std::shared_ptr<T> deviceB = make_shared_device(hostB);
std::shared_ptr<T> deviceC = make_shared_device(hostC);
std::shared_ptr<T> deviceD = make_shared_device<T>(M * N, 0.0);
auto command = std::make_shared<Command>();
std::vector<size_t> oneStridesN
= gemm.literalStrides ? std::vector<size_t>({(size_t)1}) : std::vector<size_t>({});
std::vector<size_t> oneStridesT = gemm.literalStrides
? std::vector<size_t>({(size_t)0, (size_t)1})
: std::vector<size_t>({});
auto tagTensorA = command->addOperation(rocRoller::Operations::Tensor(
2, dataType, {}, gemm.transA == "N" ? oneStridesN : oneStridesT)); // A
auto tagLoadA = command->addOperation(rocRoller::Operations::T_Load_Tiled(tagTensorA));
auto tagTensorB = command->addOperation(rocRoller::Operations::Tensor(
2, dataType, {}, gemm.transB == "N" ? oneStridesN : oneStridesT)); // B
auto tagLoadB = command->addOperation(rocRoller::Operations::T_Load_Tiled(tagTensorB));
auto tagTensorC = command->addOperation(
rocRoller::Operations::Tensor(2, dataType, {}, oneStridesN)); // C
auto tagLoadC = command->addOperation(rocRoller::Operations::T_Load_Tiled(tagTensorC));
auto tagScalarAlpha
= command->addOperation(rocRoller::Operations::Scalar(DataType::Float)); // alpha
auto tagLoadAlpha
= command->addOperation(rocRoller::Operations::T_Load_Scalar(tagScalarAlpha));
auto tagScalarBeta
= command->addOperation(rocRoller::Operations::Scalar(DataType::Float)); // beta
auto tagLoadBeta
= command->addOperation(rocRoller::Operations::T_Load_Scalar(tagScalarBeta));
auto tagScalarReluAlpha = command->addOperation(
rocRoller::Operations::Scalar(DataType::Float)); // leaky relu alpha
auto tagLoadReluAlpha
= command->addOperation(rocRoller::Operations::T_Load_Scalar(tagScalarReluAlpha));
auto tagLiteralZero
= command->addOperation(rocRoller::Operations::Literal(0.0f)); // zero
auto tagAB
= command->addOperation(rocRoller::Operations::T_Mul(tagLoadA, tagLoadB)); // A * B
rocRoller::Operations::T_Execute execute(command->getNextTag());
auto tagBetaC
= execute.addXOp(rocRoller::Operations::E_Mul(tagLoadBeta, tagLoadC)); // beta * C
auto tagAlphaAB = execute.addXOp(
rocRoller::Operations::E_Mul(tagLoadAlpha, tagAB)); // alpha * (A * B)
Operations::OperationTag tagD;
if(gemm.betaInFma)
{
tagD = execute.addXOp(rocRoller::Operations::E_Add(
tagBetaC, tagAlphaAB)); // beta * C + alpha * (A * B)
}
else
{
tagD = execute.addXOp(rocRoller::Operations::E_Add(
tagAlphaAB, tagBetaC)); // alpha * (A * B) + beta * C
}
auto tagDGtZero = execute.addXOp(
rocRoller::Operations::E_GreaterThan(tagD, tagLiteralZero)); // D > 0
auto tagDReluAlpha = execute.addXOp(
rocRoller::Operations::E_Mul(tagD, tagLoadReluAlpha)); // D * reluAlpha
auto tagRelu = execute.addXOp(rocRoller::Operations::E_Conditional(
tagDGtZero, tagD, tagDReluAlpha)); // D > 0 ? D : D * reluAlpha
command->addOperation(std::move(execute));
auto tagTensorRelu = command->addOperation(
rocRoller::Operations::Tensor(2, dataType, {}, oneStridesN)); // E
command->addOperation(rocRoller::Operations::T_Store_Tiled(tagRelu, tagTensorRelu));
Operations::OperationTag tagScratch[static_cast<int>(Operations::ScratchPolicy::Count)];
for(int i = 0; i < static_cast<int>(Operations::ScratchPolicy::Count); ++i)
{
auto policy = static_cast<Operations::ScratchPolicy>(i);
tagScratch[i] = command->allocateTag();
command->addOperation(rocRoller::Operations::Scratch(tagScratch[i], policy));
command->allocateArgument(
VariableType(DataType::UInt32, PointerType::PointerGlobal),
tagScratch[i],
ArgumentType::Value,
DataDirection::ReadWrite,
getScratchName(policy));
}
auto params = std::make_shared<CommandParameters>();
params->setManualKernelDimension(2);
// TODO: Calculate these values internally based on workgroup sizes.
params->setWaveTilesPerWavefront(wavetilePerWavefrontM, wavetilePerWavefrontN);
params->fuseLoops = gemm.fuseLoops;
params->tailLoops = gemm.tailLoops;
params->allowAmbiguousMemoryNodes = gemm.allowAmbiguousMemoryNodes;
params->unrollK = gemm.unrollK;
params->packMultipleElementsInto1VGPR = gemm.packMultipleElementsInto1VGPR;
params->prefetch = gemm.prefetch;
params->prefetchInFlight = gemm.prefetchInFlight;
params->prefetchLDSFactor = gemm.prefetchLDSFactor;
params->prefetchMixMemOps = gemm.prefetchMixMemOps;
params->transposeMemoryAccess.set(LayoutType::MATRIX_A, gemm.transA == "T");
params->transposeMemoryAccess.set(LayoutType::MATRIX_B, gemm.transB == "T");
if(gemm.loopOverTiles > 0)
{
params->loopOverOutputTilesDimensions = {0, 1};
params->loopOverOutputTilesCoordSizes
= {static_cast<uint>(M / gemm.macM), static_cast<uint>(N / gemm.macN)};
params->loopOverOutputTilesIteratedTiles = 2;
}
if(gemm.streamK)
{
REQUIRE_ARCH_CAP(GPUCapability::ArchAccUnifiedRegs);
AssertFatal(
numWorkgroupY == 1,
"Current scratch space implementation assumes that the kernel is launched "
"with numWorkgroupY == 1");
params->loopOverOutputTilesDimensions = {0, 1};
params->streamK = gemm.streamK;
}
auto macTileA = KernelGraph::CoordinateGraph::MacroTile(
{gemm.macM, gemm.macK},
LayoutType::MATRIX_A,
{gemm.waveM, gemm.waveN, gemm.waveK, gemm.waveB},
GetMemoryType(gemm.loadPathA));
auto macTileB = KernelGraph::CoordinateGraph::MacroTile(
{gemm.macK, gemm.macN},
LayoutType::MATRIX_B,
{gemm.waveM, gemm.waveN, gemm.waveK, gemm.waveB},
GetMemoryType(gemm.loadPathB));
auto macTileC = KernelGraph::CoordinateGraph::MacroTile(
{gemm.macM, gemm.macN},
LayoutType::MATRIX_ACCUMULATOR,
{gemm.waveM, gemm.waveN, gemm.waveK, gemm.waveB});
auto macTileD = KernelGraph::CoordinateGraph::MacroTile(
{gemm.macM, gemm.macN},
LayoutType::MATRIX_ACCUMULATOR,
{gemm.waveM, gemm.waveN, gemm.waveK, gemm.waveB},
(gemm.storePath == SolutionParams::StorePath::VGPRToGlobalMemoryViaLDSWithBuffer)
? MemoryType::WAVE_LDS
: MemoryType::WAVE);
params->setDimensionInfo(tagLoadA, macTileA);
params->setDimensionInfo(tagLoadB, macTileB);
params->setDimensionInfo(tagLoadC, macTileC);
params->setDimensionInfo(tagRelu, macTileD);
params->setManualWorkgroupSize({workgroupSizeX, workgroupSizeY, 1});
rocRoller::Log::getLogger()->debug(
"GEMM workgroup sizes {} {} {}", workgroupSizeX, workgroupSizeY, 1);
rocRoller::Log::getLogger()->debug(
"GEMM workitem counts {} {} {}", toString(NX), toString(NY), toString(NZ));
params->setManualWavefrontCount(
{static_cast<uint>(gemm.macM / gemm.waveM / wavetilePerWavefrontM),
static_cast<uint>(gemm.macN / gemm.waveN / wavetilePerWavefrontN)});
CommandKernel commandKernel(command, testKernelName());
commandKernel.setContext(m_context);
commandKernel.setCommandParameters(params);
commandKernel.generateKernel();
CommandArguments commandArgs = command->createArguments();
setCommandTensorArg(commandArgs, tagTensorA, descA, deviceA.get());
setCommandTensorArg(commandArgs, tagTensorB, descB, deviceB.get());
setCommandTensorArg(commandArgs, tagTensorC, descC, deviceC.get());
setCommandTensorArg(commandArgs, tagTensorRelu, descRelu, deviceD.get());
commandArgs.setArgument(tagTensorA, ArgumentType::Value, deviceA.get());
commandArgs.setArgument(tagTensorB, ArgumentType::Value, deviceB.get());
commandArgs.setArgument(tagTensorC, ArgumentType::Value, deviceC.get());
commandArgs.setArgument(tagTensorRelu, ArgumentType::Value, deviceD.get());
commandArgs.setArgument(tagScalarAlpha, ArgumentType::Value, alpha);
commandArgs.setArgument(tagScalarBeta, ArgumentType::Value, beta);
commandArgs.setArgument(
tagScalarReluAlpha, ArgumentType::Value, static_cast<T>(reluAlpha));
// Create scratch space
if(gemm.streamK)
{
commandArgs.setArgument(command->getNextTag(), ArgumentType::Value, gemm.numWGs);
}
std::shared_ptr<uint8_t>
deviceScratch[static_cast<size_t>(Operations::ScratchPolicy::Count)];
size_t scratchSpaceRequired[static_cast<size_t>(Operations::ScratchPolicy::Count)];
for(size_t i = 0; i < static_cast<size_t>(Operations::ScratchPolicy::Count); ++i)
{
scratchSpaceRequired[i] = commandKernel.scratchSpaceRequired(
static_cast<Operations::ScratchPolicy>(i), commandArgs.runtimeArguments());
if(scratchSpaceRequired[i] > 0)
{
deviceScratch[i] = make_shared_device<uint8_t>(scratchSpaceRequired[i], 0);
commandArgs.setArgument(
command->getNextTag(), ArgumentType::Value, deviceScratch[i].get());
}
}
// Host result
std::vector<T> h_result(M * N, 0.0);
rocRoller::CPUMM(h_result,
hostC,
hostA,
hostB,
M,
N,
K,
alpha,
beta,
gemm.transA == "T",
gemm.transB == "T");
// Host leaky relu
for(size_t i = 0; i < M; i++)
{
for(size_t j = 0; j < N; j++)
{
auto b = h_result[i * N + j];
h_result[i * N + j]
= b > 0 ? static_cast<T>(b)
: static_cast<T>(static_cast<T>(reluAlpha) * static_cast<T>(b));
}
}
// Device result
std::vector<T> d_result(M * N);
for(int iteration = 0; iteration < numIters; ++iteration)
{
ASSERT_THAT(hipMemset(deviceD.get(), 0, M * N * sizeof(T)), HasHipSuccess(0));
for(size_t i = 0; i < static_cast<size_t>(Operations::ScratchPolicy::Count); ++i)
{
if(scratchSpaceRequired[i] > 0)
{
ASSERT_THAT(hipMemset(deviceScratch[i].get(), 0, scratchSpaceRequired[i]),
HasHipSuccess(0));
}
}
commandKernel.launchKernel(commandArgs.runtimeArguments());
m_context = commandKernel.getContext();
ASSERT_THAT(
hipMemcpy(
d_result.data(), deviceD.get(), M * N * sizeof(T), hipMemcpyDeviceToHost),
HasHipSuccess(0));
// Verify ZeroedBeforeAndAfter scratch is all zeros after kernel
auto zeroedIdx
= static_cast<size_t>(Operations::ScratchPolicy::ZeroedBeforeAndAfter);
if(scratchSpaceRequired[zeroedIdx] > 0)
{
std::vector<uint8_t> zeroedResult(scratchSpaceRequired[zeroedIdx]);
ASSERT_THAT(hipMemcpy(zeroedResult.data(),
deviceScratch[zeroedIdx].get(),
scratchSpaceRequired[zeroedIdx],
hipMemcpyDeviceToHost),
HasHipSuccess(0));
EXPECT_TRUE(std::all_of(
zeroedResult.begin(), zeroedResult.end(), [](uint8_t v) { return v == 0; }))
<< "ZeroedBeforeAndAfter scratch should be all zeros after kernel "
"execution";
}
auto tol = gemmAcceptableError<T, T, T>(
M, N, K, m_context->targetArchitecture().target());
auto res = compare(d_result, h_result, tol);
Log::info("RNorm is {}", res.relativeNormL2);
if(debuggable && !res.ok)
{
for(size_t i = 0; i < M; i++)
{
for(size_t j = 0; j < N; j++)
{
auto const a = d_result[i * N + j];
auto b = h_result[i * N + j];
if((a - b) * (a - b) / (b * b) > 100.0 * tol.relativeL2Tolerance)
{
std::cout << std::setw(8) << i << std::setw(8) << j << std::setw(16)
<< std::scientific << a << std::setw(16)
<< std::scientific << b << std::setw(16)
<< std::scientific << a - b << std::endl;
}
}
}
}
ASSERT_TRUE(res.ok) << res.message();
}
}
};
TEST_F(GEMMFusionTestSuite, GPU_GEMM_Fusion_Relu)
{
GEMMProblem gemm;
basicGEMMRelu<float>(m_context, gemm);
}
}