Files
TheRock/rocm-libraries/shared/rocroller/test/unit/GEMMStreamKTest.cpp
T
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

380 lines
16 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 "GEMMTestBase.hpp"
#include <rocRoller/GPUArchitecture/GPUCapability.hpp>
namespace GEMMTests
{
using namespace rocRoller;
namespace SolutionParams = rocRoller::Parameters::Solution;
// ProblemConfig: (dataTypeAB, macM, macN, macK, m, n, k, numWGs)
using ProblemConfig = std::tuple<rocRoller::DataType, int, int, int, int, int, int, int>;
// ========================================================================
// GEMMStreamKMultipleFixupsTestSuite
// ========================================================================
// Params: ProblemConfig, StreamKMode, loadPathA, loadPathB, storePath
class GEMMStreamKMultipleFixupsTestSuite
: public BaseGEMMContextFixture<std::tuple<ProblemConfig,
StreamKMode,
SolutionParams::LoadPath, /* loadPathA */
SolutionParams::LoadPath, /* loadPathB */
SolutionParams::StorePath /* storePath */>>
{
};
TEST_P(GEMMStreamKMultipleFixupsTestSuite, GPU_GEMM_StreamK_MultipleFixups)
{
if(m_context->targetArchitecture().HasCapability(GPUCapability::HasWMMA))
{
GTEST_SKIP() << "Skipping GEMMStreamKMultipleFixupsTestSuite on architecture "
<< m_context->targetArchitecture().target().toString();
}
auto [problemConfig, mode, loadPathA, loadPathB, storePath] = std::get<1>(GetParam());
auto [dataTypeAB, macM, macN, macK, m, n, k, numWGs] = problemConfig;
GEMMProblem gemm;
gemm.macM = macM;
gemm.macN = macN;
gemm.macK = macK;
gemm.m = m;
gemm.n = n;
gemm.k = k;
gemm.numWGs = numWGs;
if(dataTypeAB == DataType::Half)
{
gemm.waveK = 8;
}
gemm.workgroupSizeX = 128;
gemm.workgroupSizeY = 2;
// assert that the number of output tiles is smaller than number of WGs
// which means there is not enough data-parallel tiles, and has to split
// K dimension into multiple tiles
ASSERT_GE(gemm.numWGs, gemm.m * gemm.n / gemm.macM / gemm.macN);
gemm.streamK = mode;
gemm.loadPathA = loadPathA;
gemm.loadPathB = loadPathB;
gemm.storePath = storePath;
switch(dataTypeAB)
{
case DataType::Half:
basicGEMM<Half>(gemm, false, false, 100);
break;
case DataType::Float:
basicGEMM<float>(gemm, false, false, 100);
break;
default:
Throw<FatalError>(fmt::format("Unexpected data type: {}. ", toString(dataTypeAB)));
}
}
INSTANTIATE_TEST_SUITE_P(
GEMMStreamKTest,
GEMMStreamKMultipleFixupsTestSuite,
::testing::Combine(
currentGPUISA(),
::testing::Combine(
::testing::Values(
// ProblemConfig: (dataTypeAB, macM, macN, macK, m, n, k, numWGs)
ProblemConfig{rocRoller::DataType::Half, 128, 128, 16, 128, 256, 15936, 128},
ProblemConfig{rocRoller::DataType::Float,
64,
64,
64,
256,
256,
16384,
256}), /* problemConfig */
::testing::Values(
StreamKMode::Standard, StreamKMode::TwoTile, StreamKMode::TwoTileDPFirst),
::testing::Values(SolutionParams::LoadPath::BufferToLDSViaVGPR,
SolutionParams::LoadPath::BufferToVGPR,
SolutionParams::LoadPath::GlobalToVGPR,
SolutionParams::LoadPath::GlobalToLDSViaVGPR), /* loadPathA */
::testing::Values(SolutionParams::LoadPath::BufferToLDSViaVGPR,
SolutionParams::LoadPath::BufferToVGPR,
SolutionParams::LoadPath::GlobalToVGPR,
SolutionParams::LoadPath::GlobalToLDSViaVGPR), /* loadPathB */
::testing::Values(
SolutionParams::StorePath::VGPRToGlobalMemoryWithBuffer,
SolutionParams::StorePath::VGPRToGlobalMemoryViaLDSWithBuffer) /* storePath */
)));
// ========================================================================
// GEMMStreamKWGMTestSuite
// ========================================================================
// Params: workgroupMapping dim, workgroupMapping value, workgroupRemapXCC, StreamKMode, loadPathA, loadPathB
class GEMMStreamKWGMTestSuite
: public BaseGEMMContextFixture<std::tuple<int, /* workgroupMapping dim */
int, /* workgroupMapping value */
bool, /* workgroupRemapXCC */
StreamKMode,
SolutionParams::LoadPath, /* loadPathA */
SolutionParams::LoadPath /* loadPathB */>>
{
};
TEST_P(GEMMStreamKWGMTestSuite, GPU_GEMM_StreamK_WorkgroupMapping)
{
if(m_context->targetArchitecture().HasCapability(GPUCapability::HasWMMA))
{
GTEST_SKIP() << "Skipping StreamKWGMTestGPU on architecture "
<< m_context->targetArchitecture().target().toString();
}
GEMMProblem gemm;
hipDeviceProp_t deviceProperties;
ASSERT_THAT(hipGetDeviceProperties(&deviceProperties, 0), HasHipSuccess(0));
gemm.numWGs = deviceProperties.multiProcessorCount;
gemm.m = gemm.macM * 8;
gemm.n = gemm.macN * gemm.numWGs / 2 + gemm.macN * 2;
ASSERT_GE(gemm.m * gemm.n / gemm.macM / gemm.macN, gemm.numWGs);
gemm.k = gemm.macK * 8;
std::tie(gemm.workgroupMappingDim,
gemm.workgroupMappingValue,
gemm.workgroupRemapXCC,
gemm.streamK,
gemm.loadPathA,
gemm.loadPathB)
= std::get<1>(GetParam());
basicGEMM<float>(gemm);
}
INSTANTIATE_TEST_SUITE_P(
GEMMStreamKTest,
GEMMStreamKWGMTestSuite,
::testing::Combine(
currentGPUISA(),
::testing::Combine(::testing::Values(0, 1), /* workgroupMapping dim */
::testing::Values(1, 2, 6), /* workgroupMapping value */
::testing::Values(true, false), /* remapWorkgroupXCC */
::testing::Values(StreamKMode::Standard,
StreamKMode::TwoTile,
StreamKMode::TwoTileDPFirst),
::testing::Values(SolutionParams::LoadPath::BufferToLDSViaVGPR,
SolutionParams::LoadPath::GlobalToLDSViaVGPR),
::testing::Values(SolutionParams::LoadPath::BufferToLDSViaVGPR,
SolutionParams::LoadPath::GlobalToLDSViaVGPR))));
// ========================================================================
// GEMMStreamKTestSuite
// ========================================================================
// PrefetchConfig: (prefetchInFlight, prefetchLDSFactor)
using PrefetchConfig = std::tuple<int, int>;
// Params: typeAB, unrollK, loadPathA, loadPathB, storePath, mode, betaZero, prefetchConfig
class GEMMStreamKTestSuite : public BaseGEMMContextFixture<std::tuple<rocRoller::DataType,
int,
SolutionParams::LoadPath,
SolutionParams::LoadPath,
SolutionParams::StorePath,
rocRoller::StreamKMode,
bool,
PrefetchConfig>>
{
};
TEST_P(GEMMStreamKTestSuite, GPU_GEMM_StreamK_Parameterized)
{
if(m_context->targetArchitecture().HasCapability(GPUCapability::HasWMMA))
{
GTEST_SKIP() << "Skipping StreamKTestGPU on architecture "
<< m_context->targetArchitecture().target().toString();
}
auto [typeAB, unrollK, loadPathA, loadPathB, storePath, mode, betaZero, prefetchConfig]
= std::get<1>(GetParam());
auto [prefetchInFlight, prefetchLDSFactor] = prefetchConfig;
GEMMProblem gemm;
if(typeAB == DataType::Half)
{
gemm.waveK = 8;
gemm.macK = 16;
}
hipDeviceProp_t deviceProperties;
ASSERT_THAT(hipGetDeviceProperties(&deviceProperties, 0), HasHipSuccess(0));
gemm.numWGs = deviceProperties.multiProcessorCount;
gemm.m = gemm.macM * 8;
gemm.n = gemm.macN * gemm.numWGs / 2 + gemm.macN * 2;
ASSERT_GE(gemm.m * gemm.n / gemm.macM / gemm.macN, gemm.numWGs);
gemm.streamK = mode;
gemm.k = gemm.macK * 8;
gemm.loadPathA = loadPathA;
gemm.loadPathB = loadPathB;
gemm.storePath = storePath;
gemm.unrollK = unrollK;
gemm.transA = "T";
gemm.transB = "N";
if(betaZero)
gemm.beta = 0;
// Enable prefetch if prefetchInFlight > 0
if(prefetchInFlight > 0)
{
gemm.prefetch = true;
gemm.prefetchInFlight = prefetchInFlight;
gemm.prefetchLDSFactor = prefetchLDSFactor;
}
switch(typeAB)
{
case DataType::Half:
basicGEMM<Half>(gemm);
break;
case DataType::Float:
basicGEMM<float>(gemm);
break;
default:
Throw<FatalError>(fmt::format("Unexpected data type: {}. ", toString(typeAB)));
}
}
INSTANTIATE_TEST_SUITE_P(
GEMMTest,
GEMMStreamKMultipleFixupsTestSuite,
::testing::Combine(
currentGPUISA(),
::testing::Combine(
::testing::Values(
// ProblemConfig: (dataTypeAB, macM, macN, macK, m, n, k, numWGs)
ProblemConfig{rocRoller::DataType::Half, 128, 128, 16, 128, 256, 15936, 128},
ProblemConfig{rocRoller::DataType::Float, 128, 128, 64, 128, 128, 4096, 128},
ProblemConfig{rocRoller::DataType::Float,
64,
64,
64,
256,
256,
16384,
256}), /* problemConfig */
::testing::Values(
StreamKMode::Standard, StreamKMode::TwoTile, StreamKMode::TwoTileDPFirst),
::testing::Values(SolutionParams::LoadPath::BufferToLDSViaVGPR,
SolutionParams::LoadPath::BufferToVGPR,
SolutionParams::LoadPath::GlobalToVGPR,
SolutionParams::LoadPath::GlobalToLDSViaVGPR), /* loadPathA */
::testing::Values(SolutionParams::LoadPath::BufferToLDSViaVGPR,
SolutionParams::LoadPath::BufferToVGPR,
SolutionParams::LoadPath::GlobalToVGPR,
SolutionParams::LoadPath::GlobalToLDSViaVGPR), /* loadPathB */
::testing::Values(
SolutionParams::StorePath::VGPRToGlobalMemoryViaLDSWithBuffer,
SolutionParams::StorePath::VGPRToGlobalMemoryWithBuffer) /* storePath */
)));
using StreamKParamGenerator
= ::testing::internal::ParamGenerator<GEMMStreamKTestSuite::ParamType>;
static auto FilterValidStreamKParams(StreamKParamGenerator&& inputParamGenerator)
{
using LP = SolutionParams::LoadPath;
using DT = rocRoller::DataType;
using SM = rocRoller::StreamKMode;
std::vector<GEMMStreamKTestSuite::ParamType> filtered;
for(auto const& inputParam : inputParamGenerator)
{
auto const& params = std::get<1>(inputParam);
auto const& typeAB = std::get<0>(params);
auto const& unrollK = std::get<1>(params);
auto const& loadPathA = std::get<2>(params);
auto const& loadPathB = std::get<3>(params);
auto const& storePath = std::get<4>(params);
auto const& mode = std::get<5>(params);
auto const& [prefetchInFlight, prefetchLDSFactor] = std::get<7>(params);
// Prefetch requires LDS
if((prefetchInFlight > 0) and (loadPathA == LP::BufferToVGPR)
and (loadPathB == LP::BufferToVGPR))
{
continue;
}
// Prefetch requires unrollK > 1
if(prefetchInFlight > 0 and unrollK < 2)
{
continue;
}
// Runs out of LDS
if((typeAB == DT::Float) and (loadPathA == LP::BufferToLDSViaVGPR)
and (loadPathB == LP::BufferToLDSViaVGPR) and (SolutionParams::IsLDSStore(storePath))
and (mode != SM::Standard))
{
continue;
}
// Runs out of VGPRs: Float + both LDS paths + non-Standard StreamK + prefetch
if((typeAB == DT::Float) and (loadPathA == LP::BufferToLDSViaVGPR)
and (loadPathB == LP::BufferToLDSViaVGPR) and (mode != SM::Standard)
and (prefetchInFlight > 0))
{
continue;
}
filtered.push_back(inputParam);
}
return ::testing::ValuesIn(filtered);
}
INSTANTIATE_TEST_SUITE_P(
GEMMStreamKTest,
GEMMStreamKTestSuite,
FilterValidStreamKParams(::testing::Combine(
currentGPUISA(),
::testing::Combine(
::testing::Values(rocRoller::DataType::Float, rocRoller::DataType::Half),
::testing::Values(1, 2), // unrollK
::testing::Values(SolutionParams::LoadPath::BufferToVGPR,
SolutionParams::LoadPath::BufferToLDSViaVGPR),
::testing::Values(SolutionParams::LoadPath::BufferToVGPR,
SolutionParams::LoadPath::BufferToLDSViaVGPR),
::testing::Values(
SolutionParams::StorePath::VGPRToGlobalMemoryViaLDSWithBuffer,
SolutionParams::StorePath::VGPRToGlobalMemoryWithBuffer), // storePath
::testing::Values(rocRoller::StreamKMode::Standard,
rocRoller::StreamKMode::TwoTile,
rocRoller::StreamKMode::TwoTileDPFirst),
::testing::Values(true, false), // betaZero
::testing::Values(PrefetchConfig{0, 0},
PrefetchConfig{1, 0},
PrefetchConfig{1, 2},
PrefetchConfig{2, 0},
PrefetchConfig{2, 2})))));
} // namespace GEMMTests