derek/TheRock
1
0 Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
7502bf64a4a4f3bf0074ac1068e837355abcfcde
TheRock/rocm-libraries/shared/rocroller/client/src/gemm.cpp
T
derek.holloway 7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
Details
CodeQL / Analyze (python) (push) Canceled after 0s
Details
pre-commit / pre-commit (push) Canceled after 0s
Details
Multi-Arch CI / setup (push) Canceled after 0s
Details
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Details
Setup (Multi-Arch) / setup (push) Canceled after 0s
Details
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Details
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
Details
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Details
Multi-Arch CI / CI Summary (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Details
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Details
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
Details
Flatten subprojects into a single base monorepo
2026-09-02 16:34:47 -07:00

2357 lines
96 KiB
C++
Raw Blame History

// Copyright Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include "rocRoller/Serialization/YAML.hpp"
#include <algorithm>
#include <filesystem>
#include <fstream>
#include <string>
#ifdef ROCROLLER_USE_HIP
#include <hip/hip_ext.h>
#include <hip/hip_runtime.h>
#endif /* ROCROLLER_USE_HIP */
#include <rocRoller/AssemblyKernel.hpp>
#include <rocRoller/CommandSolution.hpp>
#include <rocRoller/KernelOptions_detail.hpp>
#include <rocRoller/Operations/CommandArgument_fwd.hpp>
#include <rocRoller/Operations/OperationTag.hpp>
#include <rocRoller/Utilities/HipUtils.hpp>
#include <rocRoller/Utilities/Timer.hpp>
#include <rocRoller/Utilities/Utils.hpp>
#include <rocRoller/Utilities/Version.hpp>
#include <rocRoller/Serialization/KernelGraph.hpp>
#include <common/Utilities.hpp>
#include <common/mxDataGen.hpp>
#include "client/CLI_Utils.hpp"
#include "client/DataParallelGEMMSolution.hpp"
#include "client/GEMMParameters.hpp"
#include "client/GEMMParameters_serialization.hpp"
#include "client/RotatingBuffer.hpp"
#include "client/StreamKGEMMSolution.hpp"
#include <mxDataGenerator/PreSwizzle.hpp>
#include <CLI/CLI.hpp>
using namespace rocRoller;
namespace SolutionParams = rocRoller::Parameters::Solution;
enum ReturnCodes : int
{
OK = 0,
GenerateFailure = 1,
CorrectnessFailure = 2,
SolutionNotSupportedOnArch = 3
};
namespace rocRoller::Client::GEMMClient
{
using GEMMSolutionPtr = std::shared_ptr<Client::GEMMClient::GEMMSolution>;
template <typename A, typename B, typename C, typename D>
std::pair<bool, double>
validate(std::vector<typename PackedTypeOf<A>::type> const& h_A,
std::vector<typename PackedTypeOf<B>::type> const& h_B,
std::vector<C> const& h_C,
std::vector<D> const& h_D,
std::vector<uint8_t> const& h_scaleA,
std::vector<uint8_t> const& h_scaleB,
rocRoller::Client::GEMMClient::ProblemParameters const& problemParams,
GPUArchitecture const& arch)
{
using namespace rocRoller::Client::GEMMClient;
// Host result
std::vector<D> h_result(problemParams.m * problemParams.n, static_cast<D>(0.0));
if(!h_scaleA.empty() || !h_scaleB.empty())
{
rocRoller::ScaledCPUMM(h_result,
h_C,
h_A,
h_B,
h_scaleA,
h_scaleB,
problemParams.m,
problemParams.n,
problemParams.k,
problemParams.alpha,
problemParams.beta,
problemParams.types.transA == TransposeType::T,
problemParams.types.transB == TransposeType::T,
problemParams.types.scaleBlockSize,
problemParams.types.scaleTypeA,
problemParams.types.scaleTypeB);
}
else
{
CPUMM(h_result,
h_C,
h_A,
h_B,
problemParams.m,
problemParams.n,
problemParams.k,
problemParams.alpha,
problemParams.beta,
problemParams.types.transA == TransposeType::T,
problemParams.types.transB == TransposeType::T);
}
auto tol = gemmAcceptableError<A, B, D>(
problemParams.m, problemParams.n, problemParams.k, arch.target());
auto res = compare(h_D, h_result, tol);
Log::debug(res.message());
std::cout << "Result: " << (res.ok ? "Correct" : "Incorrect") << std::endl;
std::cout << "RNorm: " << res.relativeNormL2 << std::endl;
if(!res.ok)
{
std::cerr << "WARNING: Result incorrect. " << res.message() << std::endl;
}
return {res.ok, res.relativeNormL2};
}
// D (MxN) = alpha * A (MxK) X B (KxN) + beta * C (MxN)
template <typename A, typename B, typename C, typename D>
Client::BenchmarkResults GEMM(CommandPtr command,
CommandKernelPtr commandKernel,
GEMMSolutionPtr gemm,
ProblemParameters const& problemParams,
RunParameters const& runParams,
BenchmarkParameters const& benchmarkParams,
GPUArchitecture const& arch)
{
using namespace rocRoller::Client;
using namespace rocRoller::Client::GEMMClient;
// Host Data
std::cout << "Generating input data..." << std::endl;
TensorDescriptor descA(fromString<DataType>(problemParams.types.typeA),
{static_cast<unsigned long>(problemParams.m),
static_cast<unsigned long>(problemParams.k)},
problemParams.types.transA == TransposeType::T ? "T" : "N");
TensorDescriptor descB(fromString<DataType>(problemParams.types.typeB),
{static_cast<unsigned long>(problemParams.k),
static_cast<unsigned long>(problemParams.n)},
problemParams.types.transB == TransposeType::T ? "T" : "N");
TensorDescriptor descC(fromString<DataType>(problemParams.types.typeC),
{static_cast<unsigned long>(problemParams.m),
static_cast<unsigned long>(problemParams.n)},
"N");
using PackedTypeA = typename PackedTypeOf<A>::type;
using PackedTypeB = typename PackedTypeOf<B>::type;
std::vector<PackedTypeA> hostA;
std::vector<PackedTypeB> hostB;
std::vector<C> hostC;
std::vector<D> hostD(problemParams.m * problemParams.n, D{});
std::vector<uint8_t> hostScaleA, hostScaleB;
auto seed = 31415u;
if(problemParams.types.scaleA == Operations::ScaleMode::Separate
|| problemParams.types.scaleB == Operations::ScaleMode::Separate)
{
auto scaleBlockSize = problemParams.types.scaleBlockSize;
AssertFatal(scaleBlockSize > 0, "scaleBlockSize must be set to scale A or B.");
AssertFatal(arch.isSupportedScaleBlockSize(scaleBlockSize),
fmt::format("Architecture {} does not support block scaling (size: {}).",
arch.target().toString(),
scaleBlockSize));
AssertFatal(problemParams.k % scaleBlockSize == 0,
fmt::format("K: {} must be a multiple of the scale block size: {}",
problemParams.k,
scaleBlockSize));
DGenInput(seed,
hostA,
descA,
hostB,
descB,
hostC,
descC,
hostScaleA,
hostScaleB,
problemParams.types.scaleTypeA,
problemParams.types.scaleTypeB,
-1.f,
1.f,
static_cast<uint>(scaleBlockSize),
problemParams.initModeA,
problemParams.initModeB,
problemParams.initModeC);
}
else
{
DGenInput(seed,
hostA,
descA,
hostB,
descB,
hostC,
descC,
-1.f,
1.f,
problemParams.initModeA,
problemParams.initModeB,
problemParams.initModeC);
}
// Pre-tile B on the host when pretileB is set (kernel expects pre-tiled layout)
std::vector<PackedTypeB> hostBForKernel(hostB);
if(not problemParams.types.pretileB.empty() && problemParams.types.pretileB.size() == 2)
{
auto const packing = TypeInfo<PackedTypeB>::ElementBits / TypeInfo<B>::ElementBits;
std::vector<size_t> sizes = descB.sizes();
std::vector<size_t> preTileSize = problemParams.types.pretileB;
if(packing > 1)
{
// Tile sizes are in logical elements; convert sizes and preTileSize to packed
// element space so that product(sizes) == hostB.size() and tiles align.
// Packing is along the first (fast) dimension for B (K).
AssertFatal(sizes[0] % packing == 0,
"pretileB: K dimension must be a multiple of packing factor (",
packing,
") for packed type B.");
AssertFatal(preTileSize[0] % packing == 0,
"pretileB: tile K must be a multiple of packing factor (",
packing,
") for packed type B.");
sizes[0] /= packing;
preTileSize[0] /= packing;
}
hostBForKernel = DGen::preSwizzle(hostB, sizes, {}, preTileSize);
}
// Pre-tile A on the host when pretileA is set (kernel expects pre-tiled layout)
std::vector<PackedTypeA> hostAForKernel(hostA);
if(not problemParams.types.pretileA.empty() && problemParams.types.pretileA.size() == 2)
{
auto const packing = TypeInfo<PackedTypeA>::ElementBits / TypeInfo<A>::ElementBits;
std::vector<size_t> sizes = descA.sizes();
std::vector<size_t> preTileSize = problemParams.types.pretileA;
if(packing > 1)
{
// Tile sizes are in logical elements; convert sizes and preTileSize to packed
// element space. Packing is along the last (fast) dimension for A (K).
AssertFatal(sizes[1] % packing == 0,
"pretileA: K dimension must be a multiple of packing factor (",
packing,
") for packed type A.");
AssertFatal(preTileSize[1] % packing == 0,
"pretileA: tile K must be a multiple of packing factor (",
packing,
") for packed type A.");
sizes[1] /= packing;
preTileSize[1] /= packing;
}
// The preSwizzle helper assumes column-major; so we swap sizes here.
std::vector<size_t> swappedSizes = {sizes[1], sizes[0]};
std::vector<size_t> swappedPreTileSize = {preTileSize[1], preTileSize[0]};
hostAForKernel = DGen::preSwizzle(hostA, swappedSizes, {}, swappedPreTileSize);
}
size_t rotatingSize = benchmarkParams.rotatingBuffSize;
int idx = -1, currDeviceAsicRevisionId = -1;
HIP_CHECK(hipGetDevice(&idx));
HIP_CHECK(
hipDeviceGetAttribute(&currDeviceAsicRevisionId, hipDeviceAttributeAsicRevision, idx));
if(arch.target().isCDNA5GPU() && currDeviceAsicRevisionId == 2)
{
Log::warn("Disabling rotating buffers for arch {} revision {}",
toString(arch.target()),
currDeviceAsicRevisionId);
rotatingSize = 0;
}
RotatingBuffer<PackedTypeA> rotatingA(hostAForKernel, rotatingSize);
RotatingBuffer<PackedTypeB> rotatingB(hostBForKernel, rotatingSize);
RotatingBuffer<C> rotatingC(hostC, rotatingSize);
auto deviceD = make_shared_device<D>(problemParams.m * problemParams.n, D{});
std::shared_ptr<uint8_t> deviceScaleA, deviceScaleB;
AssertFatal(problemParams.types.scaleA == Operations::ScaleMode::None
|| problemParams.types.scaleA == Operations::ScaleMode::SingleScale
|| problemParams.types.scaleA == Operations::ScaleMode::Separate,
"Scale mode not supported!",
ShowValue(problemParams.types.scaleA));
AssertFatal(problemParams.types.scaleB == Operations::ScaleMode::None
|| problemParams.types.scaleB == Operations::ScaleMode::SingleScale
|| problemParams.types.scaleB == Operations::ScaleMode::Separate,
"Scale mode not supported!",
ShowValue(problemParams.types.scaleB));
if(problemParams.types.scaleA == Operations::ScaleMode::Separate)
{
if((problemParams.types.scaleSkipPermlane != ScaleSkipPermlaneMode::None)
|| (not problemParams.types.scalePretileA.empty()))
{
auto descScaleA = descA.withNormalizedDimensions();
{
auto sizes = descScaleA.sizes();
sizes[0] /= problemParams.types.scaleBlockSize;
descScaleA = TensorDescriptor(descScaleA.dataType(), std::move(sizes));
}
std::vector<size_t> preSwizzleSize;
if(problemParams.types.scaleSkipPermlane != ScaleSkipPermlaneMode::None)
{
AssertFatal(problemParams.types.scaleShuffleTileA.size() == 3);
preSwizzleSize = problemParams.types.scaleShuffleTileA;
}
std::vector<size_t> preTileSize;
if(not problemParams.types.scalePretileA.empty())
{
AssertFatal(problemParams.types.transA == TransposeType::T,
"Can only pre-tile A if it is TransposeType::T");
// Note, scalePretileA is M x K (usually something like 256 x 8)
//
// Because we used "withNormalizedDimensions"
// above and A is T, the K dimension becomes the
// leftmost dimension
preTileSize = {problemParams.types.scalePretileA[1],
problemParams.types.scalePretileA[0]};
}
auto tmpScaleA = [&]() {
if(problemParams.types.scaleSkipPermlane
== rocRoller::ScaleSkipPermlaneMode::PreSwizzleScaleGFX950)
return DGen::preSwizzleScalesGFX950(
hostScaleA, {descScaleA.sizes()[1], descScaleA.sizes()[0]});
else
return DGen::preSwizzle(
hostScaleA, descScaleA.sizes(), preSwizzleSize, preTileSize);
}();
deviceScaleA = make_shared_device(tmpScaleA);
}
else
{
deviceScaleA = make_shared_device(hostScaleA);
}
}
if(problemParams.types.scaleB == Operations::ScaleMode::Separate)
{
if((problemParams.types.scaleSkipPermlane != ScaleSkipPermlaneMode::None)
|| (not problemParams.types.scalePretileB.empty()))
{
auto descScaleB = descB.withNormalizedDimensions();
{
auto sizes = descScaleB.sizes();
sizes[0] /= problemParams.types.scaleBlockSize;
descScaleB = TensorDescriptor(descScaleB.dataType(), std::move(sizes));
}
std::vector<size_t> preSwizzleSize;
if(problemParams.types.scaleSkipPermlane != ScaleSkipPermlaneMode::None)
{
AssertFatal(problemParams.types.scaleShuffleTileB.size() == 3);
preSwizzleSize = problemParams.types.scaleShuffleTileB;
}
std::vector<size_t> preTileSize;
if(not problemParams.types.scalePretileB.empty())
{
// Note scalePretileB is K x N (usually something like 8 x 256)
//
// Because we used "withNormalizedDimensions"
// above, and B is N, the K dimension stays as the
// leftmost dimension.
AssertFatal(problemParams.types.transB == TransposeType::N,
"Can only pre-tile B if it is TransposeType::N");
preTileSize = {problemParams.types.scalePretileB[0],
problemParams.types.scalePretileB[1]};
};
auto tmpScaleB = [&]() {
if(problemParams.types.scaleSkipPermlane
== rocRoller::ScaleSkipPermlaneMode::PreSwizzleScaleGFX950)
return DGen::preSwizzleScalesGFX950(
hostScaleB, {descScaleB.sizes()[1], descScaleB.sizes()[0]});
else
return DGen::preSwizzle(
hostScaleB, descScaleB.sizes(), preSwizzleSize, preTileSize);
}();
deviceScaleB = make_shared_device(tmpScaleB);
}
else
{
deviceScaleB = make_shared_device(hostScaleB);
}
}
std::cout << "Generating launch parameters and runtime arguments..." << std::endl;
commandKernel->loadKernel();
auto commandArgs = gemm->commandArguments(command, problemParams, runParams);
auto [aTag, bTag, cTag, dTag] = gemm->getABCDTags();
commandArgs.setArgument(dTag, ArgumentType::Value, (D*)deviceD.get());
if(problemParams.types.scaleA == Operations::ScaleMode::Separate)
{
TensorDescriptor descAScale = TensorDescriptor(
problemParams.types.scaleTypeA,
{static_cast<size_t>(problemParams.m),
static_cast<size_t>(problemParams.k / problemParams.types.scaleBlockSize)},
problemParams.types.transA == TransposeType::T ? "T" : "N");
auto [aScaleTag, bScaleTag] = gemm->getABScaleTags();
setCommandTensorArg(commandArgs, aScaleTag.value(), descAScale, deviceScaleA.get());
}
else if(problemParams.types.scaleA == Operations::ScaleMode::SingleScale)
{
uint8_t scaleValue
= floatToScale(problemParams.types.scaleTypeA, problemParams.scaleValueA);
auto [aScaleTag, bScaleTag] = gemm->getABScaleTags();
commandArgs.setArgument(aScaleTag.value(), ArgumentType::Value, scaleValue);
hostScaleA = {scaleValue};
}
if(problemParams.types.scaleB == Operations::ScaleMode::Separate)
{
TensorDescriptor descBScale = TensorDescriptor(
problemParams.types.scaleTypeB,
{static_cast<size_t>(problemParams.k / problemParams.types.scaleBlockSize),
static_cast<size_t>(problemParams.n)},
problemParams.types.transB == TransposeType::T ? "T" : "N");
auto [aScaleTag, bScaleTag] = gemm->getABScaleTags();
setCommandTensorArg(commandArgs, bScaleTag.value(), descBScale, deviceScaleB.get());
}
else if(problemParams.types.scaleB == Operations::ScaleMode::SingleScale)
{
uint8_t scaleValue
= floatToScale(problemParams.types.scaleTypeB, problemParams.scaleValueB);
auto [aScaleTag, bScaleTag] = gemm->getABScaleTags();
commandArgs.setArgument(bScaleTag.value(), ArgumentType::Value, scaleValue);
hostScaleB = {scaleValue};
}
gemm->validateRunParameters(
command, problemParams, runParams, benchmarkParams, commandKernel);
auto runtimeArgs = commandArgs.runtimeArguments();
// Note: the lifetime of deviceScratch needs to exceed kernel executions
std::shared_ptr<uint8_t>
deviceScratch[static_cast<size_t>(Operations::ScratchPolicy::Count)];
for(int i = 0; i < static_cast<int>(Operations::ScratchPolicy::Count); ++i)
{
auto policy = static_cast<Operations::ScratchPolicy>(i);
auto scratchSpaceRequired = commandKernel->scratchSpaceRequired(policy, runtimeArgs);
if(scratchSpaceRequired > 0)
{
deviceScratch[i] = make_shared_device<uint8_t>(scratchSpaceRequired, 0);
commandArgs.setArgument(
gemm->getScratchTag(policy), ArgumentType::Value, deviceScratch[i].get());
}
}
if(benchmarkParams.visualize)
{
Client::visualize(command, *commandKernel, commandArgs);
}
std::cout << std::endl;
std::cout << "Problem:" << std::endl;
std::cout << problemParams << std::endl;
std::cout << "Launching GPU kernel(s)..." << std::endl;
BenchmarkResults result;
result.runParams = runParams;
result.benchmarkParams = benchmarkParams;
// Benchmark runs
for(int outer = 0; outer < benchmarkParams.numOuter; ++outer)
{
// Warmup runs
for(int i = 0; i < benchmarkParams.numWarmUp; ++i)
{
auto spanA = rotatingA.next();
auto spanB = rotatingB.next();
auto spanC = rotatingC.next();
commandArgs.setArgument(
aTag, ArgumentType::Value, reinterpret_cast<unsigned char*>(spanA.data()));
commandArgs.setArgument(
bTag, ArgumentType::Value, reinterpret_cast<unsigned char*>(spanB.data()));
commandArgs.setArgument(
cTag, ArgumentType::Value, reinterpret_cast<unsigned char*>(spanC.data()));
auto runtimeArgs = commandArgs.runtimeArguments();
commandKernel->launchKernel(runtimeArgs);
}
HIP_TIMER(t_kernel, "GEMM", benchmarkParams.numInner);
for(int inner = 0; inner < benchmarkParams.numInner; ++inner)
{
auto spanA = rotatingA.next();
auto spanB = rotatingB.next();
auto spanC = rotatingC.next();
commandArgs.setArgument(
aTag, ArgumentType::Value, reinterpret_cast<unsigned char*>(spanA.data()));
commandArgs.setArgument(
bTag, ArgumentType::Value, reinterpret_cast<unsigned char*>(spanB.data()));
commandArgs.setArgument(
cTag, ArgumentType::Value, reinterpret_cast<unsigned char*>(spanC.data()));
auto runtimeArgs = commandArgs.runtimeArguments();
commandKernel->launchKernel(runtimeArgs, t_kernel, inner);
}
HIP_SYNC(t_kernel);
t_kernel->sleep(50);
auto nanoseconds = t_kernel->allNanoseconds();
result.kernelExecute.insert(
result.kernelExecute.end(), nanoseconds.begin(), nanoseconds.end());
}
double totalTime = 0;
for(auto ke : result.kernelExecute)
totalTime += static_cast<double>(ke) / 1.e9;
double averageTime = totalTime / (benchmarkParams.numInner * benchmarkParams.numOuter);
std::cout << "Average runtime (s): " << averageTime << std::endl;
std::cout << "Average GFLOPS: "
<< (double)problemParams.m * problemParams.n * problemParams.k * 2.0 / averageTime
* 1.e-9
<< std::endl;
std::cerr << "Average GFLOPS: "
<< (double)problemParams.m * problemParams.n * problemParams.k * 2.0 / averageTime
* 1.e-9
<< std::endl;
result.kernelAssemble = TimerPool::nanoseconds("Assembler::assembleMachineCode");
result.kernelGenerate = TimerPool::nanoseconds("CommandKernel::generateKernel");
if(commandKernel->getContext() && commandKernel->getContext()->kernel())
{
auto assemblyKernel = commandKernel->getContext()->kernel();
result.sgprCount = assemblyKernel->sgpr_count();
result.vgprCount = assemblyKernel->vgpr_count();
result.agprCount = assemblyKernel->agpr_count();
result.ldsBytes = assemblyKernel->group_segment_fixed_size();
}
if(benchmarkParams.check)
{
AssertFatal(hipMemcpy(hostD.data(),
deviceD.get(),
problemParams.m * problemParams.n * sizeof(D),
hipMemcpyDeviceToHost)
== (hipError_t)HIP_SUCCESS);
auto [correct, rnorm] = validate<A, B, C, D>(
hostA, hostB, hostC, hostD, hostScaleA, hostScaleB, problemParams, arch);
// Verify ZeroedBeforeAndAfter scratch is all zeros after kernel
auto zeroedIdx = static_cast<size_t>(Operations::ScratchPolicy::ZeroedBeforeAndAfter);
if(deviceScratch[zeroedIdx])
{
auto zeroedSize = commandKernel->scratchSpaceRequired(
Operations::ScratchPolicy::ZeroedBeforeAndAfter, runtimeArgs);
std::vector<uint8_t> zeroedResult(zeroedSize);
AssertFatal(hipMemcpy(zeroedResult.data(),
deviceScratch[zeroedIdx].get(),
zeroedSize,
hipMemcpyDeviceToHost)
== (hipError_t)HIP_SUCCESS);
AssertFatal(
std::all_of(
zeroedResult.begin(), zeroedResult.end(), [](uint8_t v) { return v == 0; }),
"ZeroedBeforeAndAfter scratch should be all zeros after kernel execution");
}
result.checked = true;
result.correct = correct;
result.rnorm = rnorm;
}
return result;
}
template <typename A, typename C, typename D>
Client::BenchmarkResults GEMMMixed(CommandPtr command,
CommandKernelPtr commandKernel,
GEMMSolutionPtr gemm,
ProblemParameters const& problemParams,
RunParameters const& runParams,
BenchmarkParameters const& benchmarkParams,
GPUArchitecture const& arch,
auto typeB)
{
if(typeB == "fp8")
{
return GEMM<A, FP8, C, D>(
command, commandKernel, gemm, problemParams, runParams, benchmarkParams, arch);
}
else if(typeB == "bf8")
{
return GEMM<A, BF8, C, D>(
command, commandKernel, gemm, problemParams, runParams, benchmarkParams, arch);
}
else if(typeB == "fp6")
{
return GEMM<A, FP6, C, D>(
command, commandKernel, gemm, problemParams, runParams, benchmarkParams, arch);
}
else if(typeB == "bf6")
{
return GEMM<A, BF6, C, D>(
command, commandKernel, gemm, problemParams, runParams, benchmarkParams, arch);
}
else if(typeB == "fp4")
{
return GEMM<A, FP4, C, D>(
command, commandKernel, gemm, problemParams, runParams, benchmarkParams, arch);
}
else
Throw<FatalError>("Invalid type for Mixed GEMM.");
}
template <typename C, typename D>
Client::BenchmarkResults GEMMMixed(CommandPtr command,
CommandKernelPtr commandKernel,
GEMMSolutionPtr gemm,
ProblemParameters const& problemParams,
RunParameters const& runParams,
BenchmarkParameters const& benchmarkParams,
GPUArchitecture const& arch,
auto typeA,
auto typeB)
{
if(typeA == "fp8")
{
return GEMMMixed<FP8, C, D>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeB);
}
else if(typeA == "bf8")
{
return GEMMMixed<BF8, C, D>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeB);
}
else if(typeA == "fp6")
{
return GEMMMixed<FP6, C, D>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeB);
}
else if(typeA == "bf6")
{
return GEMMMixed<BF6, C, D>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeB);
}
else if(typeA == "fp4")
{
return GEMMMixed<FP4, C, D>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeB);
}
else
Throw<FatalError>("Invalid type for Mixed GEMM.");
}
template <typename AB>
Client::BenchmarkResults GEMMUniform(CommandPtr command,
CommandKernelPtr commandKernel,
GEMMSolutionPtr gemm,
ProblemParameters const& problemParams,
RunParameters const& runParams,
BenchmarkParameters const& benchmarkParams,
GPUArchitecture const& arch,
auto typeCD)
{
if(typeCD == "float")
{
return GEMM<AB, AB, float, float>(
command, commandKernel, gemm, problemParams, runParams, benchmarkParams, arch);
}
if(typeCD == "half")
{
return GEMM<AB, AB, Half, Half>(
command, commandKernel, gemm, problemParams, runParams, benchmarkParams, arch);
}
if(typeCD == "bf16")
{
return GEMM<AB, AB, BFloat16, BFloat16>(
command, commandKernel, gemm, problemParams, runParams, benchmarkParams, arch);
}
Throw<FatalError>("Invalid CD type for uniform GEMM.");
}
Client::BenchmarkResults GEMMUniform(CommandPtr command,
CommandKernelPtr commandKernel,
GEMMSolutionPtr gemm,
ProblemParameters const& problemParams,
RunParameters const& runParams,
BenchmarkParameters const& benchmarkParams,
GPUArchitecture const& arch,
auto typeAB,
auto typeCD)
{
if(typeAB == "float")
{
return GEMMUniform<float>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeCD);
}
if(typeAB == "half")
{
return GEMMUniform<Half>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeCD);
}
if(typeAB == "bf16")
{
return GEMMUniform<BFloat16>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeCD);
}
if(typeAB == "fp8")
{
return GEMMUniform<FP8>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeCD);
}
if(typeAB == "bf8")
{
return GEMMUniform<BF8>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeCD);
}
if(typeAB == "fp6")
{
return GEMMUniform<FP6>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeCD);
}
if(typeAB == "bf6")
{
return GEMMUniform<BF6>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeCD);
}
if(typeAB == "fp4")
{
return GEMMUniform<FP4>(command,
commandKernel,
gemm,
problemParams,
runParams,
benchmarkParams,
arch,
typeCD);
}
Throw<FatalError>("Invalid AB type for uniform GEMM.");
}
/*
* Generate an instance of GEMMSolution and call generateSolution.
*
* The kind of GEMMSolution returned is based on the parameters in
* solutionParams.
*/
std::pair<GEMMSolutionPtr, int> createGEMMSolution(rocRoller::ContextPtr context,
SolutionParameters const& solution)
{
GEMMSolutionPtr gemmSolution;
auto const& arch = context->targetArchitecture().target();
if(solution.streamK != StreamKMode::None)
{
if(context->targetArchitecture().HasCapability(GPUCapability::ArchAccUnifiedRegs))
{
gemmSolution = std::make_shared<Client::GEMMClient::StreamKGEMMSolution>(context);
}
else
{
std::cout << "Not running StreamK for " << arch.toString() << std::endl;
return {nullptr, ReturnCodes::SolutionNotSupportedOnArch};
}
}
else
{
gemmSolution = std::make_shared<Client::GEMMClient::DataParallelGEMMSolution>(context);
}
AssertFatal(gemmSolution, "No solution!");
return {gemmSolution, ReturnCodes::OK};
}
struct ArchitectureParameters
{
GPUArchitectureTarget target;
};
struct IOParameters
{
bool doSaveAsm, doSaveCO;
std::string saveAsmPath, loadAsmPath;
std::string saveCOPath, loadCOPath;
std::string resultsPath, timersPath;
std::string saveGraphPath;
};
void writeFile(std::filesystem::path const& filename, std::vector<char> const& x)
{
std::ofstream file(filename, std::ios::out | std::ios::binary);
std::copy(x.cbegin(), x.cend(), std::ostream_iterator<unsigned char>(file));
}
int RunGEMMCLI(bool doGenerate,
bool doValidate,
bool doBenchmark,
bool doInfo,
ArchitectureParameters architecture,
SolutionParameters solution,
ProblemParameters problem,
TypeParameters types,
RunParameters run,
BenchmarkParameters benchmark,
IOParameters io)
{
GEMMSolutionPtr gemm;
CommandPtr command;
CommandKernelPtr commandKernel;
// Changing settings has to go before creating the context :(
if(io.doSaveAsm)
{
if(io.saveAsmPath.empty())
io.saveAsmPath = solution.generateKernelName().shortName + ".s";
Settings::getInstance()->set(Settings::SaveAssembly, true);
Settings::getInstance()->set(Settings::AssemblyFile, std::string(io.saveAsmPath));
}
// When loading from .s or .co, currently need to load
// solution from a .yaml file.
if(!io.loadAsmPath.empty())
{
auto yamlPath = std::filesystem::path{io.loadAsmPath};
yamlPath.replace_extension(".yaml");
solution = Serialization::readYAMLFile<SolutionParameters>(yamlPath);
architecture.target = solution.architecture;
if(solution.version != rocRoller::Version::Git())
{
std::cout << "Warning: this version of rocRoller (" << rocRoller::Version::Git()
<< ") differs from the one that generated the kernel." << std::endl;
}
}
auto const arch = GPUArchitectureLibrary::getInstance()->GetArch(architecture.target);
if(solution.scheduler != "")
{
auto schedulerValue = fromString<Scheduling::SchedulerProcedure>(solution.scheduler);
Settings::getInstance()->set(Settings::Scheduler, schedulerValue);
}
if(solution.schedulerCost != "")
{
auto cost = fromString<Scheduling::CostFunction>(solution.schedulerCost);
Settings::getInstance()->set(Settings::SchedulerCost, cost);
}
auto context = Context::ForTarget(arch,
solution.generateKernelName().shortName,
{{.scaleSkipPermlane = solution.types.scaleSkipPermlane,
.ldsSwizzleMode = solution.ldsBankSwizzle}});
if(doInfo)
{
std::cout << "Loading kernel from: " << io.loadCOPath << std::endl;
try
{
auto elfKernels = AssemblyKernels::fromELF(io.loadCOPath).kernels;
AssertFatal(elfKernels.size() == 1,
"Expected exactly one kernel in ELF file, found ",
elfKernels.size());
auto kernelFromELF = elfKernels.at(0);
auto metadataYaml = kernelFromELF.amdgpu_metadata_yaml();
std::cout << metadataYaml << std::endl;
std::cout << *kernelFromELF.command() << std::endl;
}
catch(const std::exception& e)
{
std::cerr << "Error loading ELF file: " << e.what() << std::endl;
return ReturnCodes::GenerateFailure;
}
return ReturnCodes::OK;
}
bool willRunOnGPU = doValidate || doBenchmark;
if(willRunOnGPU)
{
std::cout << "Setting HIP device to " << benchmark.device << std::endl;
HIP_CHECK(hipSetDevice(benchmark.device));
}
if(willRunOnGPU && solution.streamK != StreamKMode::None)
{
if(run.numWGs == 0)
{
hipDeviceProp_t deviceProperties;
AssertFatal(hipGetDeviceProperties(&deviceProperties, 0)
== (hipError_t)HIP_SUCCESS);
run.numWGs = deviceProperties.multiProcessorCount;
}
}
if(doGenerate)
{
std::cout << "Generating for architecture: "
<< context->targetArchitecture().target().toString() << std::endl;
std::cout << std::endl;
std::cout << "Solution:" << std::endl;
std::cout << solution << std::endl;
std::cout << "Generating: " << solution.generateKernelName().shortName << "..."
<< std::endl;
int reason;
std::tie(gemm, reason) = createGEMMSolution(context, solution);
if(!gemm)
return reason;
command = gemm->makeCommand(solution);
commandKernel = gemm->generateCommandKernel(command, solution);
std::string basePath;
if(io.doSaveAsm)
basePath = io.saveAsmPath;
if(io.doSaveCO)
basePath = io.saveCOPath;
if(io.doSaveAsm || io.doSaveCO)
{
// When saveing ASM, also need code-object so that
// Command (ie, workgroup size, argument mapping etc)
// can be de-serialized later.
auto codeObject = commandKernel->assembleKernel();
std::filesystem::path codeObjectPath{basePath};
codeObjectPath.replace_extension(".co");
{
// Output RunParameters into a YAML file
std::filesystem::path yamlRunParams{basePath};
auto stem = yamlRunParams.stem();
std::filesystem::path yamlRunParamsPath = stem.string() + "_runParameters.yaml";
if(!std::filesystem::exists(yamlRunParamsPath))
{
std::ofstream file(yamlRunParamsPath);
Serialization::writeYAML(file, run);
std::cout << "Wrote: " << yamlRunParamsPath.string() << std::endl;
}
}
std::filesystem::path yamlPath{basePath};
yamlPath.replace_extension(".yaml");
if(!std::filesystem::exists(yamlPath))
{
std::ofstream file(yamlPath);
Serialization::writeYAML(file, solution);
std::cout << "Wrote: " << yamlPath.string() << std::endl;
}
writeFile(codeObjectPath, codeObject);
std::cout << "Wrote: " << codeObjectPath.string() << std::endl;
}
if(io.doSaveAsm)
{
std::filesystem::path assemblyPath{basePath};
assemblyPath.replace_extension(".s");
// We don't explicitly write here (the code-gen does),
// but we still emit a message here.
std::cout << "Wrote: " << assemblyPath.string() << std::endl;
}
if(not io.saveGraphPath.empty())
{
std::ofstream file(io.saveGraphPath);
Serialization::writeYAML(file, commandKernel->getKernelGraph());
std::cout << "Wrote: " << io.saveGraphPath << std::endl;
}
}
else
{
int reason;
std::tie(gemm, reason) = createGEMMSolution(context, solution);
command = gemm->makeCommand(solution);
if(!io.loadAsmPath.empty())
{
// When loading ASM, we also load code-object so that
// Command (ie, workgroup size, argument mapping etc)
// can be de-serialized.
std::filesystem::path codeObjectPath{io.loadAsmPath};
codeObjectPath.replace_extension(".co");
std::cout << "Loading kernel meta-data from: " << codeObjectPath.string()
<< std::endl;
commandKernel = std::make_shared<CommandKernel>();
commandKernel->setContext(context);
auto kernel = commandKernel->loadKernelFromCodeObject(
codeObjectPath, solution.generateKernelName().shortName);
command = kernel->command();
std::cout << "Loading kernel from: " << io.loadAsmPath << std::endl;
commandKernel = std::make_shared<CommandKernel>(
command, solution.generateKernelName().shortName);
commandKernel->setContext(context);
commandKernel->loadKernelFromAssembly(io.loadAsmPath,
solution.generateKernelName().shortName);
}
else if(!io.loadCOPath.empty())
{
std::cout << "Loading kernel from: " << io.loadCOPath << std::endl;
commandKernel = std::make_shared<CommandKernel>();
commandKernel->setContext(context);
auto kernel = commandKernel->loadKernelFromCodeObject(
io.loadCOPath, solution.generateKernelName().shortName);
command = kernel->command();
}
}
if(!gemm || !command || !commandKernel)
return ReturnCodes::GenerateFailure;
if(doValidate || doBenchmark)
{
std::cout << "Running..." << std::endl;
if(doValidate)
{
benchmark.check = true;
benchmark.numWarmUp = 0;
benchmark.numOuter = 1;
benchmark.numInner = 1;
}
auto isF8F6F4 = [](auto dtype) {
return (dtype == "fp8" || dtype == "bf8" || dtype == "fp6" || dtype == "bf6"
|| dtype == "fp4");
};
Client::GEMMClient::Result result;
result.problemParams = problem;
result.solutionParams = solution;
result.benchmarkResults.runParams = run;
result.benchmarkResults.benchmarkParams = benchmark;
if(types.typeA == types.typeB && types.typeC == types.typeD)
{
result.benchmarkResults = GEMMUniform(command,
commandKernel,
gemm,
problem,
run,
benchmark,
arch,
types.typeA,
types.typeC);
}
else if((problem.types.typeA != problem.types.typeB) && isF8F6F4(problem.types.typeA)
&& isF8F6F4(problem.types.typeB))
{
result.benchmarkResults = GEMMMixed<float, float>(command,
commandKernel,
gemm,
problem,
run,
benchmark,
arch,
problem.types.typeA,
problem.types.typeB);
}
else
{
Throw<FatalError>("Unsupported combination of datatypes for GEMM");
}
if(!io.resultsPath.empty())
{
std::ofstream file(io.resultsPath);
Serialization::writeYAML(file, result);
}
if(!result.benchmarkResults.correct)
return ReturnCodes::CorrectnessFailure;
}
// Dump timers
if(!io.timersPath.empty())
{
std::ofstream dfile;
dfile.open(io.timersPath, std::ofstream::out | std::ofstream::trunc);
dfile << rocRoller::TimerPool::CSV();
dfile.close();
}
return ReturnCodes::OK;
}
void overwriteTypesFromSolution(TypeParameters& types, SolutionParameters const& solution)
{
if((types.typeA != solution.types.typeA) || (types.typeB != solution.types.typeB)
|| (types.typeC != solution.types.typeC) || (types.typeD != solution.types.typeD)
|| (types.typeAcc != solution.types.typeAcc))
{
std::cout << "NOTE: Types have been superceded by solution." << std::endl;
}
if((types.transA != solution.types.transA) || (types.transB != solution.types.transB))
{
std::cout << "NOTE: Transposes have been superceded by solution." << std::endl;
}
if((types.scaleA != solution.types.scaleA) || (types.scaleA != solution.types.scaleB))
{
std::cout << "NOTE: MX Scalings have been superceded by solution." << std::endl;
}
if(types.scaleBlockSize != solution.types.scaleBlockSize)
{
std::cout << "NOTE: MX scale block size has been superceded by solution." << std::endl;
}
types = solution.types;
}
}
namespace rocRoller::Client::GEMMClient::CLI
{
constexpr auto SolutionParameterArguments = std::make_tuple(
std::make_pair("--arch", &SolutionParameters::architecture),
std::make_pair("--mac_m", &SolutionParameters::macM),
std::make_pair("--mac_n", &SolutionParameters::macN),
std::make_pair("--mac_k", &SolutionParameters::macK),
std::make_pair("--wave_m", &SolutionParameters::waveM),
std::make_pair("--wave_n", &SolutionParameters::waveN),
std::make_pair("--wave_k", &SolutionParameters::waveK),
std::make_pair("--wave_b", &SolutionParameters::waveB),
std::make_pair("--workgroup_size_x", &SolutionParameters::workgroupSizeX),
std::make_pair("--workgroup_size_y", &SolutionParameters::workgroupSizeY),
std::make_pair("--workgroupMappingDim", &SolutionParameters::workgroupMappingDim),
std::make_pair("--workgroupRemapXCC", &SolutionParameters::workgroupRemapXCC),
std::make_pair("--workgroupRemapXCCValue", &SolutionParameters::workgroupRemapXCCValue),
std::make_pair("--loadScale_A", &SolutionParameters::loadPathAScale),
std::make_pair("--loadScale_B", &SolutionParameters::loadPathBScale),
std::make_pair("--swizzleScale", &SolutionParameters::swizzleScale),
std::make_pair("--sts", &SolutionParameters::swizzleTileSize),
std::make_pair("--prefetchScale", &SolutionParameters::prefetchScale),
std::make_pair("--load_A", &SolutionParameters::loadPathA),
std::make_pair("--load_B", &SolutionParameters::loadPathB),
std::make_pair("--padLDS_A", &SolutionParameters::padLDSA),
std::make_pair("--padLDS_B", &SolutionParameters::padLDSB),
std::make_pair("--store", &SolutionParameters::storePath),
std::make_pair("--prefetch", &SolutionParameters::prefetch),
std::make_pair("--prefetchInFlight", &SolutionParameters::prefetchInFlight),
std::make_pair("--prefetchLDSFactor", &SolutionParameters::prefetchLDSFactor),
std::make_pair("--prefetchMixMemOps", &SolutionParameters::prefetchMixMemOps),
std::make_pair("--betaInFMA", &SolutionParameters::betaInFma),
std::make_pair("--scheduler", &SolutionParameters::scheduler),
std::make_pair("--schedulerCost", &SolutionParameters::schedulerCost),
std::make_pair("--tailLoops", &SolutionParameters::tailLoops),
std::make_pair("--streamK", &SolutionParameters::streamK),
std::make_pair("--ldsBankSwizzle", &SolutionParameters::ldsBankSwizzle),
std::make_pair("--workgroup_cluster_size_x", &SolutionParameters::workgroupClusterSizeX),
std::make_pair("--workgroup_cluster_size_y", &SolutionParameters::workgroupClusterSizeY),
std::make_pair("--workgroup_cluster_size_z", &SolutionParameters::workgroupClusterSizeZ));
template <typename T, typename U>
std::string getSolutionParameterArgumentName(U T::*member_ptr)
{
std::optional<std::string> found_name;
std::apply(
[&](auto&&... args) {
(([&] {
if constexpr(std::is_same_v<decltype(args.second), decltype(member_ptr)>)
{
if(args.second == member_ptr)
{
found_name = args.first;
}
}
}()),
...);
},
SolutionParameterArguments);
AssertFatal(found_name, "Internal error: could not find argument name.");
return found_name.value();
}
void updateSolutionFromArguments(rocRoller::Client::GEMMClient::SolutionParameters& solution,
::CLI::App const& app)
{
using SP = rocRoller::Client::GEMMClient::SolutionParameters;
auto SN = [](auto x) {
return rocRoller::Client::GEMMClient::CLI::getSolutionParameterArgumentName(x);
};
auto update = [&](const std::string& optionName, auto& value) -> bool {
if(app.get_option(optionName)->count())
{
if constexpr(std::is_same_v<std::decay_t<decltype(value)>, std::pair<int, int>>)
{
auto arg = app.get_option(optionName)->as<std::string>();
rocRoller::Client::GEMMClient::CLI::ParseIntPair(arg, value);
}
else
{
value = app.get_option(optionName)->as<std::decay_t<decltype(value)>>();
}
return true;
}
return false;
};
// Architecture
if(app.get_option(SN(&SP::architecture))->count())
{
auto architectureName = app.get_option(SN(&SP::architecture))->as<std::string>();
solution.architecture = GPUArchitectureTarget::fromString(architectureName);
}
// Workgroup tile size
bool wgtsSet = false;
if(app.get_option("--wgts")->count())
{
rocRoller::Client::GEMMClient::MNKTuple mnk{0, 0, 0};
if(!ParseMNK(app.get_option("--wgts")->as<std::string>(), mnk))
Throw<FatalError>("Failed to parse WGTS argument.");
solution.macM = mnk.m;
solution.macN = mnk.n;
solution.macK = mnk.k;
wgtsSet = true;
}
bool macSet = false;
macSet |= update(SN(&SP::macM), solution.macM);
macSet |= update(SN(&SP::macN), solution.macN);
macSet |= update(SN(&SP::macK), solution.macK);
if(wgtsSet && macSet)
{
Throw<FatalError>("Workgroup tile size was overspecified. Please use only --wgts or "
"the --mac_M, --mac_N, and --mac_K arguments; but not both.");
}
// Matrix instruction
if(app.get_option("--mi")->count())
{
auto x = rocRoller::Client::GEMMClient::MNKBTuple{0, 0, 0, 1};
if(!ParseMNKB(app.get_option("--mi")->as<std::string>(), x))
Throw<FatalError>("Failed to parse MI argument.");
solution.waveM = x.m;
solution.waveN = x.n;
solution.waveK = x.k;
solution.waveB = x.b;
}
update(SN(&SP::waveM), solution.waveM);
update(SN(&SP::waveN), solution.waveN);
update(SN(&SP::waveK), solution.waveK);
update(SN(&SP::waveB), solution.waveB);
// Workgroup size
bool wgsXYSet = false;
wgsXYSet |= update(SN(&SP::workgroupSizeX), solution.workgroupSizeX);
wgsXYSet |= update(SN(&SP::workgroupSizeY), solution.workgroupSizeY);
bool wgsSet = false;
if(app.get_option("--wgs")->count())
{
XYTuple xy{0, 0};
if(!ParseXY(app.get_option("--wgs")->as<std::string>(), xy))
Throw<FatalError>("Failed to parse WGS argument.");
solution.workgroupSizeX = xy.x;
solution.workgroupSizeY = xy.y;
wgsSet = true;
}
if(wgsSet && wgsXYSet)
{
Throw<FatalError>(
"Workgroup size was overspecified. Please use only --wgs or "
"the --workgroup_size_x and --workgroup_size_y arguments; but not both.");
}
// Workgroup mapping
update(SN(&SP::workgroupMappingDim), solution.workgroupMappingDim);
update(SN(&SP::workgroupRemapXCC), solution.workgroupRemapXCC);
update(SN(&SP::workgroupRemapXCCValue), solution.workgroupRemapXCCValue);
// LDS
if(app.get_option("--lds")->count())
{
auto arg = app.get_option("--lds")->as<std::string>();
solution.loadPathA = SolutionParams::LoadPath::BufferToVGPR;
if(arg.find('A') != std::string::npos)
solution.loadPathA = SolutionParams::LoadPath::BufferToLDSViaVGPR;
solution.loadPathB = SolutionParams::LoadPath::BufferToVGPR;
if(arg.find('B') != std::string::npos)
solution.loadPathB = SolutionParams::LoadPath::BufferToLDSViaVGPR;
solution.storePath = SolutionParams::StorePath::VGPRToGlobalMemoryWithBuffer;
if(arg.find('D') != std::string::npos)
solution.storePath = SolutionParams::StorePath::VGPRToGlobalMemoryViaLDSWithBuffer;
}
update(SN(&SP::loadPathA), solution.loadPathA);
update(SN(&SP::loadPathB), solution.loadPathB);
update(SN(&SP::storePath), solution.storePath);
if(app.get_option("--d2lds")->count())
{
auto arg = app.get_option("--d2lds")->as<std::string>();
solution.loadPathA = SolutionParams::LoadPath::BufferToVGPR;
if(arg.find('A') != std::string::npos)
solution.loadPathA = SolutionParams::LoadPath::BufferToLDS;
solution.loadPathB = SolutionParams::LoadPath::BufferToVGPR;
if(arg.find('B') != std::string::npos)
solution.loadPathB = SolutionParams::LoadPath::BufferToLDS;
}
update(SN(&SP::loadPathAScale), solution.loadPathAScale);
update(SN(&SP::loadPathBScale), solution.loadPathBScale);
if(app.get_option("--mxlds")->count())
{
auto arg = app.get_option("--mxlds")->as<std::string>();
solution.loadPathAScale = SolutionParams::LoadPath::BufferToVGPR;
if(arg.find('A') != std::string::npos)
solution.loadPathAScale = SolutionParams::LoadPath::BufferToLDSViaVGPR;
solution.loadPathBScale = SolutionParams::LoadPath::BufferToVGPR;
if(arg.find('B') != std::string::npos)
solution.loadPathBScale = SolutionParams::LoadPath::BufferToLDSViaVGPR;
}
if(app.get_option("--mxd2lds")->count())
{
auto arg = app.get_option("--mxd2lds")->as<std::string>();
solution.loadPathAScale = SolutionParams::LoadPath::BufferToVGPR;
if(arg.find('A') != std::string::npos)
solution.loadPathAScale = SolutionParams::LoadPath::BufferToLDS;
solution.loadPathBScale = SolutionParams::LoadPath::BufferToVGPR;
if(arg.find('B') != std::string::npos)
solution.loadPathBScale = SolutionParams::LoadPath::BufferToLDS;
}
update(SN(&SP::padLDSA), solution.padLDSA);
update(SN(&SP::padLDSB), solution.padLDSB);
// Swizzling
update(SN(&SP::swizzleScale), solution.swizzleScale);
update(SN(&SP::swizzleTileSize), solution.swizzleTileSize);
update(SN(&SP::prefetchScale), solution.prefetchScale);
update(SN(&SP::ldsBankSwizzle), solution.ldsBankSwizzle);
// Prefetching
update(SN(&SP::prefetch), solution.prefetch);
update(SN(&SP::prefetchInFlight), solution.prefetchInFlight);
update(SN(&SP::prefetchLDSFactor), solution.prefetchLDSFactor);
update(SN(&SP::prefetchMixMemOps), solution.prefetchMixMemOps);
// Tail loops
update(SN(&SP::tailLoops), solution.tailLoops);
// StreamK
update(SN(&SP::streamK), solution.streamK);
// Other
update(SN(&SP::betaInFma), solution.betaInFma);
update(SN(&SP::scheduler), solution.scheduler);
update(SN(&SP::schedulerCost), solution.schedulerCost);
}
}
/*
* Parse the command line and dispatch.
*/
int main(int argc, const char* argv[])
{
using namespace rocRoller::Client::GEMMClient::CLI;
CLI::App app{"GEMM Driver: D (MxN) = alpha * A (MxK) * B (KxN) + beta * C (MxN)"};
app.footer(Settings::getInstance()->help());
//
// Parameters
//
std::string architectureName;
rocRoller::Client::GEMMClient::ArchitectureParameters architecture;
rocRoller::Client::GEMMClient::SolutionParameters solution{
.macM = 64,
.macN = 64,
.macK = -1,
.waveM = -1,
.waveN = -1,
.waveK = -1,
.waveB = -1,
.workgroupSizeX = -1,
.workgroupSizeY = 2,
.workgroupMappingDim = -1,
.workgroupRemapXCC = false,
.workgroupRemapXCCValue = -1,
.workgroupClusterSizeX = 0,
.workgroupClusterSizeY = 0,
.workgroupClusterSizeZ = 0,
.types = {.scaleA = Operations::ScaleMode::None,
.scaleTypeA = DataType::None,
.scaleB = Operations::ScaleMode::None,
.scaleTypeB = DataType::None,
.scaleBlockSize = -1},
.loadPathAScale = SolutionParams::LoadPath::BufferToVGPR,
.loadPathBScale = SolutionParams::LoadPath::BufferToVGPR,
.swizzleScale = false,
.prefetchScale = false,
.loadPathA = SolutionParams::LoadPath::BufferToLDSViaVGPR,
.loadPathB = SolutionParams::LoadPath::BufferToLDSViaVGPR,
.storePath = SolutionParams::StorePath::VGPRToGlobalMemoryViaLDSWithBuffer,
.padLDSA = {0u, 0u},
.padLDSB = {0u, 0u},
.prefetch = false,
.prefetchInFlight = 0,
.prefetchLDSFactor = 0,
.prefetchMixMemOps = false,
.betaInFma = true,
.scheduler = "Priority",
.tailLoops = true,
.streamK = StreamKMode::None,
.version = rocRoller::Version::Git(),
};
rocRoller::Client::GEMMClient::ProblemParameters problem{
.m = 3072,
.n = 4096,
.k = 4096,
.alpha = 2.0f,
.beta = 0.5f,
.scaleValueA = 1.0f,
.scaleValueB = 1.0f,
.initModeA = DataInitMode(Bounded{}),
.initModeB = DataInitMode(Bounded{}),
.initModeC = DataInitMode(Bounded{}),
};
rocRoller::Client::GEMMClient::TypeParameters types;
rocRoller::Client::RunParameters runParams{
.workgroupMappingValue = -1,
.numWGs = 0,
};
rocRoller::Client::BenchmarkParameters benchmarkParams{
.device = 0,
.numWarmUp = 3,
.numOuter = 5,
.numInner = 2,
.rotatingBuffSize = 32'000'000ull,
.check = true,
.visualize = false,
};
rocRoller::Client::GEMMClient::IOParameters io{
.doSaveAsm = false,
.doSaveCO = false,
.saveAsmPath = "",
.loadAsmPath = "",
.saveCOPath = "",
.loadCOPath = "",
.resultsPath = "",
};
//
// Architecture
//
app.option_defaults()->ignore_case()->group("Architecture parameters");
app.add_option("--arch", architectureName, "GPU architecture name (eg, gfx90a).");
//
// Problem definition
//
app.option_defaults()->ignore_case()->group("Problem parameters");
app.add_option("-M,--M", problem.m, "Tensor size M.");
app.add_option("-N,--N", problem.n, "Tensor size N.");
app.add_option("-K,--K", problem.k, "Tensor size K.");
app.add_option("--alpha", problem.alpha, "Alpha scalar.");
app.add_option("--beta", problem.beta, "Beta scalar.");
app.add_option("--scaleValue_A", problem.scaleValueA, "Single scale value for A.");
app.add_option("--scaleValue_B", problem.scaleValueB, "Single scale value for B.");
app.add_option(
"--initMode_A",
[&problem](auto& args) -> bool { return ParseInitMode(args[0], problem.initModeA); },
"Data initialization mode for A [Bounded | BoundedAlternatingSign | Unbounded | "
"Identity | Ones | Zeros | TrigonometricFromFloat | \"NormalFromFloat(<mean>, "
"<std_dev>)\"]. "
"Default: Bounded.");
app.add_option(
"--initMode_B",
[&problem](auto& args) -> bool { return ParseInitMode(args[0], problem.initModeB); },
"Data initialization mode for B [Bounded | BoundedAlternatingSign | Unbounded | "
"Identity | Ones | Zeros | TrigonometricFromFloat | \"NormalFromFloat(<mean>, "
"<std_dev>)\"]. "
"Default: Bounded.");
app.add_option(
"--initMode_C",
[&problem](auto& args) -> bool { return ParseInitMode(args[0], problem.initModeC); },
"Data initialization mode for C [Bounded | BoundedAlternatingSign | Unbounded | "
"Identity | Ones | Zeros | TrigonometricFromFloat | \"NormalFromFloat(<mean>, "
"<std_dev>)\"]. "
"Default: Bounded.");
//
// Problem types
//
app.option_defaults()->ignore_case()->group("Type parameters");
app.add_option("--type_A",
types.typeA,
"Datatype of A matrix [float | half | bf16 | fp8 | bf8 | fp6 | bf6 | fp4]. "
"Default: float.");
app.add_option("--type_B",
types.typeB,
"Datatype of B matrix [float | half | bf16 | fp8 | bf8 | fp6 | bf6 | fp4]. "
"Default: float.");
app.add_option(
"--type_C", types.typeC, "Datatype of C matrix [float | half | bf16]. Default: float.");
app.add_option(
"--type_D", types.typeD, "Datatype of D matrix [float | half | bf16]. Default: float.");
app.add_option("--type_acc",
types.typeAcc,
"Datatype of accumulation [float | half | bf16]. Default: float");
app.add_option(
"--trans_A",
[&types](auto res) -> bool {
types.transA = fromString<Client::GEMMClient::TransposeType>(res[0]);
return true;
},
"N: A is not to be transposed. T: A is to be transposed.",
"N");
app.add_option(
"--trans_B",
[&types](auto res) -> bool {
types.transB = fromString<Client::GEMMClient::TransposeType>(res[0]);
return true;
},
"N: B is not to be transposed. T: B is to be transposed.",
"N");
app.add_option(
"--scale_A",
[&types](auto res) -> bool {
types.scaleA = fromString<Operations::ScaleMode>(res[0]);
return true;
},
"Enable MX scaling of A matrix [None | Separate | SingleScale].",
"Default: None.");
app.add_option(
"--scaleType_A",
[&types](auto res) -> bool {
types.scaleTypeA = fromString<DataType>(res[0]);
return true;
},
"Type for A matrix scales [None | E8M0].",
"Default: None.");
app.add_option(
"--scale_B",
[&types](auto res) -> bool {
types.scaleB = fromString<Operations::ScaleMode>(res[0]);
return true;
},
"Enable MX scaling of B matrix [None | Separate | SingleScale].",
"Default: None.");
app.add_option(
"--scaleType_B",
[&types](auto res) -> bool {
types.scaleTypeB = fromString<DataType>(res[0]);
return true;
},
"Type for B matrix scales [None | E8M0].",
"Default: None.");
app.add_option("--scaleBlockSize",
types.scaleBlockSize,
"Set MX scaling block size for A and B. (default: 32)");
app.add_option("--scaleSkipPermlane",
types.scaleSkipPermlane,
"Experimental: Skip Permlane instructions for scale data. Options: None, "
"PreSwizzleScale, PreSwizzleScaleGFX950.");
auto pretileAOption = app.add_option("--pretileA", "Pre-tile A matrix. Dimensions are: MxK.");
auto pretileBOption = app.add_option("--pretileB", "Pre-tile B matrix. Dimensions are: KxN.");
bool pretileScale = false;
app.add_flag("--pretileScale", pretileScale, "Experimental: pretile scale data.");
//
// Solution parameters
//
app.option_defaults()->ignore_case()->group("Solution parameters");
using SP = rocRoller::Client::GEMMClient::SolutionParameters;
auto SN = [](auto x) {
return rocRoller::Client::GEMMClient::CLI::getSolutionParameterArgumentName(x);
};
app.add_option(SN(&SP::macM), "(Macro) Tile size M.");
app.add_option(SN(&SP::macN), "(Macro) Tile size N.");
app.add_option(SN(&SP::macK), "(Macro) Tile size K.");
app.add_option("--wgts", "Workgroup tile size (m/n/k tuple).");
app.add_option(SN(&SP::waveM), "(MI) Tile size M.");
app.add_option(SN(&SP::waveN), "(MI) Tile size N.");
app.add_option(SN(&SP::waveK), "(MI) Tile size K.");
app.add_option(SN(&SP::waveB), "(MI) Tile size B.");
app.add_option("--mi", "MI (matrix instruction) to use.");
app.add_option(SN(&SP::workgroupSizeX), "Workgroup size in the x dimension.");
app.add_option(SN(&SP::workgroupSizeY), "Workgroup size in the y dimension.");
app.add_option("--wgs", "Workgroup size (x/y pair).");
app.add_option(SN(&SP::workgroupMappingDim),
"Workgroup mapping dimension (-1, 0, 1). Default: -1")
->check(CLI::IsMember({-1, 0, 1}));
app.add_flag(SN(&SP::workgroupRemapXCC), "Use an XCC-aware workgroup remapping.");
app.add_option(SN(&SP::workgroupRemapXCCValue),
"Force an XCC-aware workgroup remapping value. (Optional)");
app.add_option(
SN(&SP::loadPathA),
"How to load A (BufferToVGPR, BufferToLDSViaVGPR, BufferToLDS). Default: BufferToLDS");
app.add_option(
SN(&SP::loadPathB),
"How to load B (BufferToVGPR, BufferToLDSViaVGPR, BufferToLDS). Default: BufferToLDS");
app.add_option(SN(&SP::storePath),
"How to store D (VGPRToGlobalMemoryWithBuffer, VGPRToGlobal, "
"VGPRToGlobalMemoryViaLDSWithBuffer, "
"VGPRToGlobalMemoryViaLDSWithGlobal). Default: "
"VGPRToGlobalMemoryViaLDSWithBuffer");
app.add_option("--lds", "Use LDS for A/B/D.");
app.add_option("--d2lds", "Use direct-to-LDS for A/B.");
app.add_flag(SN(&SP::betaInFma), "Use beta in FMA instruction instead of alpha.");
app.add_option(SN(&SP::scheduler), "Which scheduler to use.");
app.add_option(SN(&SP::schedulerCost), "Which scheduler cost function to use.");
auto descriptionPadLDSA = fmt::format("Byte padding for A LDS buffer. Passed as a pair: "
"contiguous-bytes,padding-bytes, eg {}=1024,8",
SN(&SP::padLDSA));
app.add_option(SN(&SP::padLDSA), descriptionPadLDSA);
auto descriptionPadLDSB = fmt::format("Byte padding for B LDS buffer. Passed as a pair: "
"contiguous-bytes,padding-bytes, eg {}=1024,8",
SN(&SP::padLDSB));
app.add_option(SN(&SP::padLDSB), descriptionPadLDSB);
app.add_flag(SN(&SP::tailLoops), "Enable tail-loops transformation.");
app.add_flag(SN(&SP::prefetch), "Enable prefetching (UnrollK=2 implied).");
app.add_option(SN(&SP::prefetchInFlight), "Number of prefetches in flight at the same time");
app.add_option(SN(&SP::prefetchLDSFactor),
"Prefetch 1/prefetchLDSFactor of MacroTile from LDS");
app.add_flag(SN(&SP::prefetchMixMemOps),
"Mix global and LDS memory operations during prefetching.");
app.add_option(SN(&SP::streamK),
"StreamK mode (None, Standard, TwoTile, TwoTileDPFirst). Default: None")
->check(CLI::IsMember(rocRoller::enumStrings<StreamKMode>()));
app.add_option(SN(&SP::ldsBankSwizzle), "LDS bank swizzle mode (None, Swizzle). Default: None")
->check(CLI::IsMember(rocRoller::enumStrings<LDSBankSwizzleMode>()));
app.add_option(SN(&SP::loadPathAScale),
"How to load AScale (BufferToVGPR, BufferToLDSViaVGPR, BufferToLDS). Default: "
"BufferToLDSViaVGPR");
app.add_option(SN(&SP::loadPathBScale),
"How to load BScale (BufferToVGPR, BufferToLDSViaVGPR, BufferToLDS). Default: "
"BufferToLDSViaVGPR");
app.add_option("--mxlds", "Use LDS for A/B scales.");
app.add_option("--mxd2lds", "Use direct-to-LDS for A/B scales.");
app.add_flag(SN(&SP::swizzleScale), "Use Swizzle when loading A and B scale.");
app.add_option(SN(&SP::swizzleTileSize),
"Size of swizzle tile in MxK/NxL format. The A scale swizzle-tile is MxK. The "
"B scale swizzle-tile is NxL.");
app.add_flag(SN(&SP::prefetchScale), "Prefetch scale values with using Swizzled scales.");
app.add_option("--workgroupMappingValue",
runParams.workgroupMappingValue,
"Workgroup mapping value. Default: -1")
->check(CLI::IsMember({-1}) | CLI::PositiveNumber);
app.add_option("--workgroup_cluster_size_x",
solution.workgroupClusterSizeX,
"Workgroup cluster size in the x dimension.");
app.add_option("--workgroup_cluster_size_y",
solution.workgroupClusterSizeY,
"Workgroup cluster size in the y dimension.");
app.add_option("--workgroup_cluster_size_z",
solution.workgroupClusterSizeZ,
"Workgroup cluster size in the z dimension.");
//
// Benchmarking options
//
app.option_defaults()->ignore_case()->group("Benchmarking parameters");
app.add_option("--num_warmup", benchmarkParams.numWarmUp, "Number of warm-up runs.");
app.add_option("--num_outer", benchmarkParams.numOuter, "Number of outer runs.");
app.add_option("--num_inner", benchmarkParams.numInner, "Number of inner runs.");
app.add_option("--device", benchmarkParams.device, "GPU device ordinal");
app.add_option("--numWGs",
runParams.numWGs,
"Number of workgroups to use with StreamK algorithm. Defaults to number of WGs "
"present on local device.");
app.add_option(
"--rotating_buff_size", benchmarkParams.rotatingBuffSize, "Rotating Buffer Size.");
//
// Client params and shortcuts
//
app.option_defaults()->ignore_case()->group("Client options and shortcuts");
bool noCheckResult = false;
std::string loadPath, examplePath, exampleProblemPath;
app.add_flag(
"--hgemm",
[&types](auto res) -> bool {
types.typeA = "half";
types.typeB = "half";
types.typeC = "half";
types.typeD = "half";
types.typeAcc = "float";
return true;
},
"Overwrite types to: --type_A=half --type_B=half --type_C=half --type_D=half "
"--type_acc=float.");
app.add_flag("--visualize",
benchmarkParams.visualize,
"Dump out volumes describing memory access patterns.");
app.add_flag("--noCheck", noCheckResult, "Do not verify GEMM results against OpenBLAS.");
app.add_option("--yaml", io.resultsPath, "Save results to file.");
app.add_option("--timers", io.timersPath, "Save timers to CSV file.");
//
// generate sub-command
//
std::string loadConfigPath;
auto generate = app.add_subcommand("generate", "Generate a GEMM solution.")->fallthrough();
auto asmOption = generate->add_option("--asm", io.saveAsmPath, "Save yaml+assembly to files.")
->expected(0, 1);
auto coOption
= generate->add_option("--co", io.saveCOPath, "Save code-object to file.")->expected(0, 1);
auto graphOption
= generate->add_option("--graph", io.saveGraphPath, "Save kernel graph to file.")
->expected(0, 1);
generate
->add_option(
"--config", loadConfigPath, "Load solution generation parameters from YAML file.")
->expected(1, 1);
//
// validate sub-command
//
auto validate
= app.add_subcommand("validate",
"Run and validate a GEMM solution (only runs the solution once).")
->fallthrough();
validate->add_option(
"--load", loadPath, "Load solution from code-object (.co) or assembly (.s) file.");
//
// benchmark sub-command
//
auto benchmark
= app.add_subcommand("benchmark",
"Benchmark a GEMM solution (may run the solution many times).")
->fallthrough();
benchmark->add_option(
"--load", loadPath, "Load solution from code-object (.co) or assembly (.s) file.");
std::string loadRunParamsPath;
benchmark->add_option(
"--loadRunParams", loadRunParamsPath, "Load run parameters from YAML file.");
//
// info sub-command
//
auto info = app.add_subcommand("info", "Dump info about a GEMM solution.")->fallthrough();
info->add_option("load", loadPath, "Load solution from code-object (.co).")->required();
//
// example sub-command
//
auto example = app.add_subcommand("example", "Save example generation parameters to YAML file.")
->fallthrough();
example->add_option("save", examplePath, "Example config path.")->required();
//
// example problem parameters sub-command
//
auto exampleProblem
= app.add_subcommand("exampleProblem", "Save example problem parameters to YAML file.")
->fallthrough();
exampleProblem->add_option("save", exampleProblemPath, "Example problem config path")
->required();
//
// Parse and update/validate problem definition
//
CLI11_PARSE(app, argc, argv);
if(pretileAOption->count() > 0)
{
auto xy = pretileAOption->as<Client::GEMMClient::XYTuple>();
types.pretileA = {static_cast<unsigned long>(xy.x), static_cast<unsigned long>(xy.y)};
}
if(pretileBOption->count() > 0)
{
auto xy = pretileBOption->as<Client::GEMMClient::XYTuple>();
types.pretileB = {static_cast<unsigned long>(xy.x), static_cast<unsigned long>(xy.y)};
}
updateSolutionFromArguments(solution, app);
if(architectureName.empty())
architecture.target = GPUArchitectureLibrary::getInstance()
->GetDefaultHipDeviceArch(benchmarkParams.device)
.target();
else
architecture.target = GPUArchitectureTarget::fromString(architectureName);
solution.architecture = architecture.target;
if(!loadConfigPath.empty())
{
solution = Serialization::readYAMLFile<rocRoller::Client::GEMMClient::SolutionParameters>(
loadConfigPath);
updateSolutionFromArguments(solution, app);
overwriteTypesFromSolution(types, solution);
if(solution.architecture.gfx == GPUArchitectureGFX::UNKNOWN)
solution.architecture = architecture.target;
}
if(!loadPath.empty())
{
auto path = std::filesystem::path(loadPath);
if(path.extension() == ".s" || path.extension() == ".yaml")
{
io.loadAsmPath = path.string();
}
else if(path.extension() == ".co")
{
io.loadCOPath = path.string();
}
else
{
Throw<FatalError>("Extension not supported. Can not load solution from ", loadPath);
}
}
if(!loadRunParamsPath.empty())
{
auto path = std::filesystem::path(loadRunParamsPath);
path.replace_extension(".yaml");
// Load RunParameters from a specified YAML file
runParams = Serialization::readYAMLFile<rocRoller::Client::RunParameters>(path);
}
if(!io.loadAsmPath.empty() || !io.loadCOPath.empty())
{
std::filesystem::path yamlPath;
if(!io.loadAsmPath.empty())
yamlPath = std::filesystem::path{io.loadAsmPath};
if(!io.loadCOPath.empty())
yamlPath = std::filesystem::path{io.loadCOPath};
yamlPath.replace_extension(".yaml");
// YAML file does not have the workgroupMappingValue used to generate the kernel.
// Instead, the workgroupMappingValue specified by users in benchmarking will be used.
solution = Serialization::readYAMLFile<rocRoller::Client::GEMMClient::SolutionParameters>(
yamlPath);
overwriteTypesFromSolution(types, solution);
}
if(types.scaleA != Operations::ScaleMode::None && types.scaleB == Operations::ScaleMode::None)
{
types.scaleB = Operations::ScaleMode::SingleScale;
problem.scaleValueB = 1.0f;
types.scaleTypeB = types.scaleTypeA;
}
if(types.scaleB != Operations::ScaleMode::None && types.scaleA == Operations::ScaleMode::None)
{
types.scaleA = Operations::ScaleMode::SingleScale;
problem.scaleValueA = 1.0f;
types.scaleTypeA = types.scaleTypeB;
}
auto const& arch = GPUArchitectureLibrary::getInstance()->GetArch(solution.architecture);
if(types.scaleBlockSize == -1
&& (types.scaleA == Operations::ScaleMode::Separate
|| types.scaleB == Operations::ScaleMode::Separate))
{
AssertFatal(arch.HasCapability(GPUCapability::HasBlockScaling32),
fmt::format("Architecture {} does not support block scaling.",
arch.target().toString()));
types.scaleBlockSize = arch.GetCapability(GPUCapability::DefaultScaleBlockSize);
}
if(solution.workgroupSizeX == -1)
solution.workgroupSizeX = 2 * arch.GetCapability(GPUCapability::DefaultWavefrontSize);
if(solution.workgroupSizeY == -1)
solution.workgroupSizeY = 2;
const DataType typeA = fromString<DataType>(types.typeA);
const DataType typeB = fromString<DataType>(types.typeB);
const DataType typeC = fromString<DataType>(types.typeC);
const DataType typeD = fromString<DataType>(types.typeD);
const DataType typeAcc = fromString<DataType>(types.typeAcc);
AssertFatal((typeAcc == DataType::Float) || (typeAcc == DataType::Half)
|| (typeAcc == DataType::BFloat16));
// TODO: Reevaluate the relationship between problem and solution params.
problem.workgroupMappingDim = solution.workgroupMappingDim;
benchmarkParams.check = !noCheckResult;
io.doSaveAsm = asmOption->count() > 0;
io.doSaveCO = coOption->count() > 0;
// Set default MI and macK sizes
if(arch.HasCapability(GPUCapability::HasMFMA))
{
if(solution.macK == -1)
solution.macK = 64;
if(solution.waveB == -1)
solution.waveB = 1;
if(typeA == DataType::Float && typeB == DataType::Float && typeC == DataType::Float
&& typeD == DataType::Float)
{
if(solution.waveM == -1)
solution.waveM = 32;
if(solution.waveN == -1)
solution.waveN = 32;
if(solution.waveK == -1)
solution.waveK = 2;
}
else if(typeA == DataType::Half && typeB == DataType::Half)
{
if(solution.waveM == -1)
solution.waveM = 32;
if(solution.waveN == -1)
solution.waveN = 32;
if(solution.waveK == -1)
solution.waveK = 8;
}
else if(typeA == DataType::BFloat16 && typeB == DataType::BFloat16)
{
if(solution.waveM == -1)
solution.waveM = 16;
if(solution.waveN == -1)
solution.waveN = 16;
if(solution.waveK == -1)
solution.waveK = 8;
}
else if((typeA == DataType::FP8 && typeB == DataType::FP8)
|| (typeA == DataType::BF8 && typeB == DataType::BF8))
{
if(solution.waveM == -1)
solution.waveM = 16;
if(solution.waveN == -1)
solution.waveN = 16;
if(solution.waveK == -1)
solution.waveK = 32;
}
}
else if(arch.HasCapability(GPUCapability::HasWMMA))
{
if(solution.waveM == -1)
solution.waveM = 16;
if(solution.waveN == -1)
solution.waveN = 16;
if(solution.waveB == -1)
solution.waveB = 1;
if(typeA == DataType::Float && typeB == DataType::Float)
{
if(solution.macK == -1)
solution.macK = 64;
if(arch.HasCapability(GPUCapability::HasWMMA_f32_16x16x4_f32))
{
if(solution.waveK == -1)
solution.waveK = 4;
}
}
else if((typeA == DataType::Half && typeB == DataType::Half)
|| (typeA == DataType::BFloat16 && typeB == DataType::BFloat16))
{
if(solution.macK == -1)
solution.macK = 64;
if(arch.HasCapability(GPUCapability::HasWMMA_f32_16x16x16_f16))
{
if(solution.waveK == -1)
solution.waveK = 16;
}
else if(arch.HasCapability(GPUCapability::HasWMMA_f32_16x16x32_f16))
{
if(solution.waveK == -1)
solution.waveK = 32;
}
}
else if(isUnpackedF8(typeA) && isUnpackedF8(typeB)
&& types.scaleA == Operations::ScaleMode::None
&& types.scaleB == Operations::ScaleMode::None)
{
if(solution.macK == -1)
solution.macK = 64;
if(arch.HasCapability(GPUCapability::HasWMMA_f32_16x16x16_f8))
{
if(solution.waveK == -1)
solution.waveK = 16;
}
else if(arch.HasCapability(GPUCapability::HasWMMA_f32_16x16x64_f8))
{
if(solution.waveK == -1)
solution.waveK = 64;
}
}
else if(isUnpackedF8F6F4(typeA) && isUnpackedF8F6F4(typeB))
{
if(solution.macK == -1)
solution.macK = 128;
if(arch.HasCapability(GPUCapability::HasWMMA_f8f6f4))
{
if(solution.waveK == -1)
solution.waveK = 128;
}
}
}
else
{
Throw<FatalError>("Unsupported arch for GEMM client: ", arch.target().toString());
}
if(arch.target().isRDNA4GPU())
{
// Override default settings for the `example`, `exampleProblem`, and `generate` subcommands.
if((example->parsed() || exampleProblem->parsed() || generate->parsed())
&& typeA == DataType::Float && typeB == DataType::Float)
{
std::cout << "Warning: A and B types and wave sizes have been overridden for RDNA4."
<< std::endl;
types.typeA = "half";
types.typeB = "half";
types.typeC = "half";
types.typeD = "half";
solution.waveM = 16;
solution.waveN = 16;
solution.waveK = 16;
solution.waveB = 1;
solution.macK = 64;
}
if(solution.prefetch)
{
std::cout << "Warning: disabling prefetching for RDNA4." << std::endl;
solution.prefetch = false;
}
}
if(arch.target().isCDNA5GPU())
{
// TODO: Remove these default sizes
if(problem.m * problem.n + problem.m * problem.k + problem.k * problem.n > 512 * 512 * 3)
{
std::cout << "Warning: M, N, K sizes have been overridden to 512 for CDNA5."
<< std::endl;
problem.m = 512;
problem.n = 512;
problem.k = 512;
}
}
AssertFatal(solution.waveM > 0 && solution.waveN > 0 && solution.waveK > 0
&& solution.waveB > 0,
fmt::format("MI tile sizes must be set greater than zero. "
"waveM: {} waveN: {} waveK: {} waveB: {}",
solution.waveM,
solution.waveN,
solution.waveK,
solution.waveB));
AssertFatal(solution.macK >= solution.waveK && solution.macM >= solution.waveM
&& solution.macN >= solution.waveN,
fmt::format("Macro tile sizes must be greater than or equal to MI tile sizes. "
"macM: {} waveM: {} macN: {} waveN: {} macK: {} waveK: {}",
solution.macM,
solution.waveM,
solution.macN,
solution.waveN,
solution.macK,
solution.waveK));
AssertFatal(solution.macK % solution.waveK == 0,
fmt::format("macK must be a multiple of waveK. macK: {} waveK: {}",
solution.macK,
solution.waveK));
AssertFatal(problem.k >= solution.macK,
fmt::format("K must be >= macK. K: {} macK: {}", problem.k, solution.macK));
AssertFatal(problem.k % solution.macK == 0 || solution.tailLoops,
fmt::format("K must be a multiple of macK (or enable tailLoops). K: {} macK: {}",
problem.k,
solution.macK));
if(types.scaleSkipPermlane == rocRoller::ScaleSkipPermlaneMode::PreSwizzleScaleGFX950)
{
AssertFatal(solution.swizzleTileSize.m == 32 && solution.swizzleTileSize.n == 32
&& solution.swizzleTileSize.k == 8,
"When scaleSkipPermlane is PreSwizzleScaleGFX950, swizzleTileSize must be "
"(m=32, n=32, k=8).");
AssertFatal(pretileScale,
"pretileScale must be true when scaleSkipPermlane is PreSwizzleScaleGFX950.");
}
if(types.scaleSkipPermlane != rocRoller::ScaleSkipPermlaneMode::None)
{
AssertFatal(types.transA == Client::GEMMClient::TransposeType::T, ShowValue(types));
AssertFatal(types.scaleA == Operations::ScaleMode::Separate, ShowValue(types));
size_t kSubtile = solution.waveK / types.scaleBlockSize;
AssertFatal(kSubtile == 2 || kSubtile == 4,
ShowValue(kSubtile),
ShowValue(solution.waveK),
ShowValue(types.scaleBlockSize));
types.scaleShuffleTileA = {static_cast<size_t>(solution.swizzleTileSize.m),
256 / static_cast<size_t>(solution.swizzleTileSize.m),
kSubtile};
}
if(types.scaleSkipPermlane != rocRoller::ScaleSkipPermlaneMode::None)
{
AssertFatal(types.transB == Client::GEMMClient::TransposeType::N, ShowValue(types));
AssertFatal(types.scaleB == Operations::ScaleMode::Separate, ShowValue(types));
size_t kSubtile = solution.waveK / types.scaleBlockSize;
AssertFatal(kSubtile == 2 || kSubtile == 4,
ShowValue(kSubtile),
ShowValue(solution.waveK),
ShowValue(types.scaleBlockSize));
types.scaleShuffleTileB = {static_cast<size_t>(solution.swizzleTileSize.n),
256 / static_cast<size_t>(solution.swizzleTileSize.n),
kSubtile};
}
if(pretileScale)
{
AssertFatal(types.scaleShuffleTileA.size() >= 2 && types.scaleShuffleTileB.size() >= 2,
"scaleShuffleTileA and scaleShuffleTileB must have at least 2 elements when "
"pretileScale is enabled");
types.scalePretileA = {types.scaleShuffleTileA[0], types.scaleShuffleTileA[1]};
types.scalePretileB = {types.scaleShuffleTileB[1], types.scaleShuffleTileB[0]};
}
problem.types = types;
solution.types = types;
//
// Run!
//
if(example->parsed())
{
std::ofstream file(examplePath);
Serialization::writeYAML(file, solution);
return 0;
}
if(exampleProblem->parsed())
{
std::ofstream file(exampleProblemPath);
Serialization::writeYAML(file, problem);
return 0;
}
// If no subcommands were given, default behaviour is: generate
// and benchmark.
bool generateAndBenchmark
= !generate->parsed() && !validate->parsed() && !benchmark->parsed() && !info->parsed();
bool doGenerate = generateAndBenchmark || generate->parsed();
bool doValidate = validate->parsed();
bool doBenchmark = generateAndBenchmark || benchmark->parsed();
bool doInfo = info->parsed();
return rocRoller::Client::GEMMClient::RunGEMMCLI(doGenerate,
doValidate,
doBenchmark,
doInfo,
architecture,
solution,
problem,
types,
runParams,
benchmarkParams,
io);
}
Reference in new issue View Git Blame Copy Permalink