Files
TheRock/rocm-libraries/projects/rocwmma/test/dlrm/dlrm_kernel_base_impl.hpp
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

541 lines
21 KiB
C++

/*******************************************************************************
*
* MIT License
*
* Copyright (C) 2021-2025 Advanced Micro Devices, Inc. All rights reserved.
*
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
* SOFTWARE.
*
*******************************************************************************/
#ifndef DLRM_KERNEL_BASE_IMPL_HPP
#define DLRM_KERNEL_BASE_IMPL_HPP
#include <cmath>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <tuple>
#include <hip/hip_ext.h>
#include <hip/hip_runtime.h>
#include <hip/hip_runtime_api.h>
#include <gtest/gtest.h>
#include <rocwmma/internal/constants.hpp>
#include <rocwmma/internal/utils.hpp>
#include "../common.hpp"
#include "./common.hpp"
#include "dlrm_kernel_base.hpp"
#include "performance.hpp"
// Library includes
#if ROCWMMA_VALIDATION_TESTS
#include "reference.hpp" // Vanilla CPU kernel
#endif // ROCWMMA_VALIDATION_TESTS
namespace rocwmma
{
template <uint32_t TileSize, typename DataT>
DlrmKernelBase<TileSize, DataT>::DlrmKernelBase()
{
reset();
}
template <uint32_t TileSize, typename DataT>
DlrmKernelBase<TileSize, DataT>::~DlrmKernelBase()
{
}
template <uint32_t TileSize, typename DataT>
uint32_t DlrmKernelBase<TileSize, DataT>::ldsUsage() const
{
return 0;
}
template <uint32_t TileSize, typename DataT>
dim3 DlrmKernelBase<TileSize, DataT>::gridDim() const
{
auto& device = DeviceInfo::instance();
if(passDirection == DlrmDirection_t::Forward)
{
return dim3(ceil_div(mMPadded, TileSize * mTBlockX / device->warpSize()),
ceil_div(mM, TileSize),
mB);
}
else
{
return dim3(ceil_div(mMPadded, TileSize * mTBlockX / device->warpSize()),
ceil_div(mK, TileSize),
mB);
}
}
template <uint32_t TileSize, typename DataT>
dim3 DlrmKernelBase<TileSize, DataT>::blockDim() const
{
return dim3(mTBlockX);
}
template <uint32_t TileSize, typename DataT>
bool DlrmKernelBase<TileSize, DataT>::checkDevice() const
{
auto& deviceInfo = DeviceInfo::instance();
auto deviceArch = deviceInfo->getGcnArch();
// Arch
auto isGfx908 = deviceArch == DeviceInfo::GFX908;
auto isGfx11 = (deviceArch == DeviceInfo::GFX1100) || (deviceArch == DeviceInfo::GFX1101)
|| (deviceArch == DeviceInfo::GFX1102)
|| (deviceArch == DeviceInfo::GFX1103)
|| (deviceArch == DeviceInfo::GFX1150)
|| (deviceArch == DeviceInfo::GFX1151)
|| (deviceArch == DeviceInfo::GFX1152)
|| (deviceArch == DeviceInfo::GFX1153);
auto isGfx12 = (deviceArch == DeviceInfo::GFX1200) || (deviceArch == DeviceInfo::GFX1201)
|| (deviceArch == DeviceInfo::GFX1250);
// Datatypes
auto isF64 = std::is_same<DataT, float64_t>::value;
#if !ROCWMMA_TESTS_NO_HALF
auto isH16 = std::is_same<DataT, hfloat16_t>::value;
#else
auto isH16 = false;
#endif // !ROCWMMA_NO_HALF
auto isF16 = std::is_same<DataT, float16_t>::value || isH16;
auto isBF16 = (std::is_same<DataT, bfloat16_t>::value);
auto isI8 = (std::is_same<DataT, int8_t>::value);
// Block size
auto is16x16 = (TileSize == 16);
// No unsupported devices
bool unsupportedDeviceCheck = !(deviceArch == DeviceInfo::UNSUPPORTED_ARCH);
// gfx908 doesn't support f64
bool gfx908F64Check = !(isGfx908 && isF64);
// gfx11 only supports f16, i8 and bf16 inputs with block size 16
bool gfx11Check = !(isGfx11 && ((!isF16 && !isBF16 && !isI8) || !is16x16));
// gfx12 only supports f16, i8 and bf16 inputs with block size 16
bool gfx12Check = !(isGfx12 && ((!isF16 && !isBF16 && !isI8) || !is16x16));
return unsupportedDeviceCheck && gfx908F64Check && gfx11Check && gfx12Check;
}
template <uint32_t TileSize, typename DataT>
bool DlrmKernelBase<TileSize, DataT>::checkSizes() const
{
return (mM >= TileSize && (mM % TileSize == 0) && mK >= TileSize && (mK % TileSize == 0)
&& (mTBlockX % TileSize == 0));
}
template <uint32_t TileSize, typename DataT>
bool DlrmKernelBase<TileSize, DataT>::checkLds() const
{
return ldsUsage() <= DeviceInfo::instance()->sharedMemSize();
}
template <uint32_t TileSize, typename DataT>
void DlrmKernelBase<TileSize, DataT>::reset()
{
mM = mK = mB = 0;
mMPadded = mKPadded = 0;
mRepeats =
#if ROCWMMA_VALIDATION_TESTS
1;
#else
5;
#endif // ROCWMMA_VALIDATION_TESTS
mRunFlag = true;
mTotalGFlops = mMeasuredTFlopsPerSec = 0.0;
mElapsedTimeMs = 0.0;
mEfficiency = -1;
passDirection = DlrmDirection_t::Forward;
mValidationResult = false;
mMaxRelativeError = 0.0;
}
template <uint32_t TileSize, typename DataT>
HipResource* DlrmKernelBase<TileSize, DataT>::getResource()
{
return DataStorage::instance().get();
}
template <uint32_t TileSize, typename DataT>
std::ostream& DlrmKernelBase<TileSize, DataT>::printHeader(std::ostream& stream) const
{
return stream << "TileSize, "
<< "DataT, "
<< "Direction, "
<< "MatM, MatK, MatB, "
#if ROCWMMA_VALIDATION_TESTS
<< "maxRelativeDiff, "
#endif // ROCWMMA_VALIDATION_TESTS
<< "elapsedMs, "
<< "Problem Size(GFlops), "
<< "TFlops/s, "
<< "Efficiency(%), "
<< "Result" << std::endl;
}
template <uint32_t TileSize, typename DataT>
std::ostream& DlrmKernelBase<TileSize, DataT>::printKernel(std::ostream& stream) const
{
if(!mRunFlag)
{
return stream << TileSize << ", " << dataTypeToString<DataT>() << ", "
<< (passDirection == DlrmDirection_t::Forward ? "Forwards" : "Backwards")
<< ", " << mM << ", " << mK << ", " << mB << ", "
#if ROCWMMA_VALIDATION_TESTS
<< "n/a, "
#endif // ROCWMMA_VALIDATION_TESTS
<< "n/a, n/a, n/a, n/a, SKIPPED" << std::endl;
}
else
{
return stream << TileSize << ", " << dataTypeToString<DataT>() << ", "
<< (passDirection == DlrmDirection_t::Forward ? "Forwards" : "Backwards")
<< ", " << mM << ", " << mK << ", " << mB << ", "
#if ROCWMMA_VALIDATION_TESTS
<< mMaxRelativeError << ", "
#endif // ROCWMMA_VALIDATION_TESTS
<< mElapsedTimeMs << ", " << mTotalGFlops << ", " << mMeasuredTFlopsPerSec
<< ", " << mEfficiency << ", "
#if ROCWMMA_VALIDATION_TESTS
<< (mValidationResult ? "PASSED" : "FAILED")
#else
<< "BENCH"
#endif // ROCWMMA_VALIDATION_TESTS
<< std::endl;
}
}
template <uint32_t TileSize, typename DataT>
void DlrmKernelBase<TileSize, DataT>::setup(ProblemParams const& problem)
{
// Reset the flags in case of multiple runs
mRunFlag = true;
// Format incoming problem parameters
std::tie(mTBlockX, mTBlockY)
= std::tie(static_cast<uint32_t const&>(std::get<0>(problem.threadBlockSize)),
static_cast<uint32_t const&>(std::get<1>(problem.threadBlockSize)));
std::tie(mM, mK, mB)
= std::tie(static_cast<uint32_t const&>(std::get<0>(problem.problemSize)),
static_cast<uint32_t const&>(std::get<1>(problem.problemSize)),
static_cast<uint32_t const&>(std::get<2>(problem.problemSize)));
mMPadded = ceil_div(mM, TileSize) * TileSize;
mKPadded = ceil_div(mK, TileSize) * TileSize;
// Determine whether to run forward or backward pass
passDirection = problem.passDirection;
mRunFlag &= checkDevice();
mRunFlag &= checkSizes();
mRunFlag &= checkLds();
if(mRunFlag)
{
auto& dataInstance = DataStorage::instance();
// Initialize matrix storage
if(passDirection == DlrmDirection_t::Forward)
{
dataInstance->resizeFwdStorage(problem.problemSize);
}
else
{
dataInstance->resizeBwdStorage(problem.problemSize);
}
// Initialize matrix data on device and transfer to host for validation
if(passDirection == DlrmDirection_t::Forward)
{
MatrixUtil<row_major>::fillLaunchKernel(
dataInstance->deviceInput().get(), mM, mK, mB);
#if ROCWMMA_VALIDATION_TESTS
dataInstance->copyDeviceToHostFwdInput();
#endif // ROCWMMA_VALIDATION_TESTS
}
else
{
uint gradSize = ((mM * (mM - 1)) / 2) + mK;
MatrixUtil<row_major>::fillLaunchKernel(
dataInstance->deviceInput().get(), mM, mK, mB);
MatrixUtil<row_major>::fillLaunchKernel(
dataInstance->deviceUpstreamGrad().get(), 1, gradSize, mB);
#if ROCWMMA_VALIDATION_TESTS
dataInstance->copyDeviceToHostBwdInput();
#endif // ROCWMMA_VALIDATION_TESTS
}
}
}
template <uint32_t TileSize, typename DataT>
void DlrmKernelBase<TileSize, DataT>::exec()
{
if(mRunFlag)
{
std::function<void()> dlrmKernel;
if(passDirection == DlrmDirection_t::Forward)
{
if(mM == mMPadded && mK == mKPadded)
{
uint inputBatchOffset = mM * mK;
uint outputBatchOffset = ((mM * (mM - 1)) / 2) + mK;
uint accBatchOffset = mM * mM;
dlrmKernel = [this, inputBatchOffset, outputBatchOffset, accBatchOffset]() {
auto& dataInstance = DataStorage::instance();
hipExtLaunchKernelGGL((this->kernelFwdImpl()),
(this->gridDim()),
(this->blockDim()),
(this->ldsUsage()),
0,
nullptr,
nullptr,
0,
dataInstance->deviceInput().get(),
dataInstance->deviceOutput().get(),
dataInstance->deviceAccFwd().get(),
mM,
mK,
mB,
inputBatchOffset,
outputBatchOffset,
accBatchOffset);
CHECK_HIP_ERROR(hipGetLastError());
};
}
}
else
{
if(mM == mMPadded && mK == mKPadded)
{
auto& dataInstance = DataStorage::instance();
uint inputBatchOffset = mM * mK;
uint upstreamBatchOffset = ((mM * (mM - 1)) / 2) + mK;
uint accBatchOffset = mM * mM;
dlrmKernel = [this, inputBatchOffset, upstreamBatchOffset, accBatchOffset]() {
auto& dataInstance = DataStorage::instance();
auto trilGridDim
= dim3(ceil_div(mM * mM, static_cast<uint32_t>(mTBlockX)), 1, mB);
hipEvent_t syncEvent;
CHECK_HIP_ERROR(hipEventCreate(&syncEvent));
hipExtLaunchKernelGGL((this->kernelTrilImpl()),
trilGridDim,
this->blockDim(),
0,
0,
nullptr,
nullptr,
0,
dataInstance->deviceUpstreamGrad().get(),
dataInstance->deviceAccBwd().get(),
mM,
mK,
mB,
upstreamBatchOffset,
accBatchOffset);
CHECK_HIP_ERROR(hipGetLastError());
CHECK_HIP_ERROR(hipEventRecord(syncEvent));
CHECK_HIP_ERROR(hipEventSynchronize(syncEvent));
hipExtLaunchKernelGGL((this->kernelBwdImpl()),
(this->gridDim()),
(this->blockDim()),
(this->ldsUsage()),
0,
nullptr,
nullptr,
0,
dataInstance->deviceInput().get(),
dataInstance->deviceUpstreamGrad().get(),
dataInstance->deviceGrad().get(),
dataInstance->deviceBottomMlpGrad().get(),
dataInstance->deviceAccBwd().get(),
mM,
mK,
mB,
inputBatchOffset,
upstreamBatchOffset,
accBatchOffset);
CHECK_HIP_ERROR(hipGetLastError());
};
}
}
hipEvent_t startEvent, stopEvent;
CHECK_HIP_ERROR(hipEventCreate(&startEvent));
CHECK_HIP_ERROR(hipEventCreate(&stopEvent));
CHECK_HIP_ERROR(hipEventRecord(startEvent));
for(uint32_t i = 0; i < mRepeats; ++i)
{
dlrmKernel();
}
CHECK_HIP_ERROR(hipEventRecord(stopEvent));
CHECK_HIP_ERROR(hipEventSynchronize(stopEvent));
auto timeMs = 0.0f;
CHECK_HIP_ERROR(hipEventElapsedTime(&timeMs, startEvent, stopEvent));
// Calculate efficiency
auto& deviceInfo = DeviceInfo::instance();
auto devicePeakGFlopsPerSec = deviceInfo->peakGFlopsPerSec<DataT>();
auto outputSize = (passDirection == DlrmDirection_t::Forward) ? mM * mM : mM * mK;
mElapsedTimeMs = float64_t(timeMs);
mTotalGFlops = calculateGFlops(outputSize, mB, mK);
mMeasuredTFlopsPerSec = calculateTFlopsPerSec(outputSize, mB, mK, mElapsedTimeMs)
* static_cast<float64_t>(mRepeats);
mEfficiency = round(mMeasuredTFlopsPerSec / devicePeakGFlopsPerSec * 100000.0);
CHECK_HIP_ERROR(hipEventDestroy(startEvent));
CHECK_HIP_ERROR(hipEventDestroy(stopEvent));
#if ROCWMMA_VALIDATION_TESTS
// Run reference CPU kernel
std::function<void()> cpuKernel;
if(passDirection == DlrmDirection_t::Forward)
{
cpuKernel = [this]() {
auto& dataInstance = DataStorage::instance();
dlrm_fwd_CPU<DataT>(dataInstance->hostInput().get(),
dataInstance->hostOutputRef().get(),
mM,
mK,
mB);
};
}
else
{
cpuKernel = [this]() {
auto& dataInstance = DataStorage::instance();
dlrm_bwd_CPU<DataT>(dataInstance->hostInput().get(),
dataInstance->hostUpstreamGrad().get(),
dataInstance->hostBottomMlpGradRef().get(),
dataInstance->hostGradRef().get(),
mM,
mK,
mB);
};
}
cpuKernel();
#endif // ROCWMMA_VALIDATION_TESTS
}
}
template <uint32_t TileSize, typename DataT>
void DlrmKernelBase<TileSize, DataT>::validateResults()
{
#if ROCWMMA_VALIDATION_TESTS
if(mRunFlag)
{
auto& dataInstance = DataStorage::instance();
if(passDirection == DlrmDirection_t::Forward)
{
uint batchSize = ((mM * (mM - 1)) / 2) + mK;
auto reference = dataInstance->template allocDevice<DataT>(batchSize * mB);
dataInstance->copyData(reference, dataInstance->hostOutputRef(), batchSize * mB);
std::tie(mValidationResult, mMaxRelativeError)
= compareEqualLaunchKernel<DataT, DataT>(dataInstance->deviceOutput().get(),
reference.get(),
1,
batchSize,
mB,
10.0);
EXPECT_TRUE(mValidationResult) << "Max relative error: " << mMaxRelativeError;
}
else
{
// Copy reference output gradient to device
auto reference0 = dataInstance->template allocDevice<DataT>(mM * mK * mB);
dataInstance->copyData(reference0, dataInstance->hostGradRef(), mM * mK * mB);
std::tie(mValidationResult, mMaxRelativeError)
= compareEqualLaunchKernel<DataT, DataT>(
dataInstance->deviceGrad().get(), reference0.get(), mM, mK, mB);
EXPECT_TRUE(mValidationResult) << "Max relative error: " << mMaxRelativeError;
double maxRelativeError = mMaxRelativeError;
// Copy reference bottom mlp gradient to device
auto reference1 = dataInstance->template allocDevice<DataT>(mK * mB);
dataInstance->copyData(reference1, dataInstance->hostBottomMlpGradRef(), mK * mB);
std::tie(mValidationResult, mMaxRelativeError)
= compareEqualLaunchKernel<DataT, DataT>(
dataInstance->deviceBottomMlpGrad().get(), reference1.get(), 1, mK, mB);
EXPECT_TRUE(mValidationResult) << "Max relative error: " << mMaxRelativeError;
mMaxRelativeError
= (maxRelativeError > mMaxRelativeError) ? maxRelativeError : mMaxRelativeError;
}
}
#endif
}
template <uint32_t TileSize, typename DataT>
void DlrmKernelBase<TileSize, DataT>::reportResults()
{
if(!KernelI::sHeaderPrinted)
{
printHeader();
KernelI::sHeaderPrinted = true;
}
printKernel();
}
template <uint32_t TileSize, typename DataT>
void DlrmKernelBase<TileSize, DataT>::tearDown()
{
}
} // namespace rocwmma
#endif // DLRM_KERNEL_BASE_IMPL_HPP