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
400 lines
15 KiB
C++
400 lines
15 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 GEMM_DRIVER_IMPL_HPP
|
|
#define GEMM_DRIVER_IMPL_HPP
|
|
|
|
#pragma GCC diagnostic push
|
|
#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
|
|
#include <rocwmma/rocwmma.hpp>
|
|
#include <rocwmma/rocwmma_transforms.hpp>
|
|
#pragma GCC diagnostic pop
|
|
|
|
namespace rocwmma
|
|
{
|
|
|
|
namespace CooperativeGemm
|
|
{
|
|
|
|
#define GemmDriverT typename GlobalMapping, typename LdsMapping
|
|
|
|
#define GemmDriverT_impl GlobalMapping, LdsMapping
|
|
|
|
template <GemmDriverT>
|
|
template <uint32_t BlocksX>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::globalReadCoopA(
|
|
GRFragA (&grFragsA)[BlocksX], GetDataType_t<GRFragA> const* gAddrA, uint32_t lda)
|
|
{
|
|
auto blockOffset = MappingUtil<GRFragA>::dataOffset(GlobalMapping::blockOffsetA(), lda);
|
|
#pragma unroll
|
|
for(int i = 0; i < BlocksX; i++)
|
|
{
|
|
globalReadCoopA(grFragsA[i], gAddrA + i * blockOffset, lda);
|
|
}
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::globalReadCoopA(
|
|
GRFragA& grFragA, GetDataType_t<GRFragA> const* gAddrA, uint32_t lda)
|
|
{
|
|
rocwmma::load_matrix_sync(grFragA, gAddrA, lda);
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <uint32_t BlocksY>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::globalReadCoopB(
|
|
GRFragB (&grFragsB)[BlocksY], GetDataType_t<GRFragB> const* gAddrB, uint32_t ldb)
|
|
{
|
|
auto blockOffset = MappingUtil<GRFragB>::dataOffset(GlobalMapping::blockOffsetB(), ldb);
|
|
#pragma unroll
|
|
for(int i = 0; i < BlocksY; i++)
|
|
{
|
|
globalReadCoopB(grFragsB[i], gAddrB + i * blockOffset, ldb);
|
|
}
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::globalReadCoopB(
|
|
GRFragB& grFragB, GetDataType_t<GRFragB> const* gAddrB, uint32_t ldb)
|
|
{
|
|
rocwmma::load_matrix_sync(grFragB, gAddrB, ldb);
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <uint32_t BlocksX>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::localWriteCoopA(
|
|
GetDataType_t<GRFragA>* ldsAddr, GRFragA const (&grFragsA)[BlocksX], uint32_t ldlds)
|
|
{
|
|
auto blockOffset = MappingUtil<LWFragA>::dataOffset(LdsMapping::blockOffsetA(), ldlds);
|
|
#pragma unroll
|
|
for(int i = 0; i < BlocksX; i++)
|
|
{
|
|
localWriteCoopA(ldsAddr + i * blockOffset, grFragsA[i], ldlds);
|
|
}
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::localWriteCoopA(
|
|
GetDataType_t<GRFragA>* ldsAddr, GRFragA const& grFragA, uint32_t ldlds)
|
|
{
|
|
rocwmma::store_matrix_sync(ldsAddr, LdsMapping::formatLWFragA(grFragA), ldlds);
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <uint32_t BlocksY>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::localWriteCoopB(
|
|
GetDataType_t<GRFragB>* ldsAddr, GRFragB const (&grFragsB)[BlocksY], uint32_t ldlds)
|
|
{
|
|
auto blockOffset = MappingUtil<LWFragB>::dataOffset(LdsMapping::blockOffsetB(), ldlds);
|
|
#pragma unroll
|
|
for(int i = 0; i < BlocksY; i++)
|
|
{
|
|
localWriteCoopB(ldsAddr + i * blockOffset, grFragsB[i], ldlds);
|
|
}
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::localWriteCoopB(
|
|
GetDataType_t<GRFragB>* ldsAddr, GRFragB const& grFragB, uint32_t ldlds)
|
|
{
|
|
rocwmma::store_matrix_sync(ldsAddr, LdsMapping::formatLWFragB(grFragB), ldlds);
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::localReadA(
|
|
MfmaFragA& fragsA, GetDataType_t<MfmaFragA> const* ldsAddrA, uint32_t ldlds)
|
|
{
|
|
LRFragA lrFragA;
|
|
rocwmma::load_matrix_sync(lrFragA, ldsAddrA, ldlds);
|
|
fragsA.mStorage = LdsMapping::formatLRFragA(lrFragA).mStorage;
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <uint32_t BlocksX>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::localReadA(
|
|
MfmaFragA (&fragsA)[BlocksX], GetDataType_t<MfmaFragA> const* ldsAddrA, uint32_t ldlds)
|
|
{
|
|
auto blockStep = MappingUtil<LRFragA>::dataOffset(LdsMapping::blockOffsetA(), ldlds);
|
|
#pragma unroll
|
|
for(int i = 0; i < BlocksX; i++)
|
|
{
|
|
localReadA(fragsA[i], ldsAddrA, ldlds);
|
|
ldsAddrA += blockStep;
|
|
}
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::localReadB(
|
|
MfmaFragB& fragsB, GetDataType_t<MfmaFragB> const* ldsAddrB, uint32_t ldlds)
|
|
{
|
|
LRFragB lrFragB;
|
|
rocwmma::load_matrix_sync(lrFragB, ldsAddrB, ldlds);
|
|
fragsB.mStorage = LdsMapping::formatLRFragB(lrFragB).mStorage;
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <uint32_t BlocksY>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::localReadB(
|
|
MfmaFragB (&fragsB)[BlocksY], GetDataType_t<MfmaFragB> const* ldsAddrB, uint32_t ldlds)
|
|
{
|
|
auto blockStep = MappingUtil<LRFragB>::dataOffset(LdsMapping::blockOffsetB(), ldlds);
|
|
#pragma unroll
|
|
for(int i = 0; i < BlocksY; i++)
|
|
{
|
|
localReadB(fragsB[i], ldsAddrB, ldlds);
|
|
ldsAddrB += blockStep;
|
|
}
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <typename FragT>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::fill(FragT& frag,
|
|
GetDataType_t<FragT> value)
|
|
{
|
|
rocwmma::fill_fragment(frag, value);
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <typename FragT, uint32_t BlocksX, uint32_t BlocksY>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::fill(FragT (&frags)[BlocksX][BlocksY],
|
|
GetDataType_t<FragT> value)
|
|
{
|
|
#pragma unroll
|
|
for(int i = 0; i < BlocksX; i++)
|
|
{
|
|
#pragma unroll
|
|
for(int j = 0; j < BlocksY; j++)
|
|
{
|
|
fill(frags[i][j], value);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::mfma(MfmaFragAcc& fragAccOut,
|
|
MfmaFragA const& fragA,
|
|
MfmaFragB const& fragB,
|
|
MfmaFragAcc const& fragAccIn)
|
|
{
|
|
rocwmma::mma_sync(fragAccOut, fragA, fragB, fragAccIn);
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <uint32_t BlocksX, uint32_t BlocksY>
|
|
__device__ inline void
|
|
GemmDriver<GemmDriverT_impl>::mfma(MfmaFragAcc (&fragAccOut)[BlocksX][BlocksY],
|
|
MfmaFragA const (&fragA)[BlocksX],
|
|
MfmaFragB const (&fragB)[BlocksY],
|
|
MfmaFragAcc const (&fragAccIn)[BlocksX][BlocksY])
|
|
{
|
|
#pragma unroll
|
|
for(int i = 0; i < BlocksX; i++)
|
|
{
|
|
#pragma unroll
|
|
for(int j = 0; j < BlocksY; j++)
|
|
{
|
|
mfma(fragAccOut[i][j], fragA[i], fragB[j], fragAccIn[i][j]);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::globalReadC(
|
|
MfmaFragC& fragC, GetDataType_t<MfmaFragC> const* gAddrC, uint32_t ldc)
|
|
{
|
|
rocwmma::load_matrix_sync(fragC, gAddrC, ldc);
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <uint32_t BlocksX, uint32_t BlocksY>
|
|
__device__ inline void
|
|
GemmDriver<GemmDriverT_impl>::globalReadC(MfmaFragC (&fragC)[BlocksX][BlocksY],
|
|
GetDataType_t<MfmaFragC> const* gAddrC,
|
|
uint32_t ldc)
|
|
{
|
|
auto blockStepX
|
|
= MappingUtil<MfmaFragC>::dataOffset(GlobalMapping::blockOffsetA(), ldc);
|
|
auto blockStepY
|
|
= MappingUtil<MfmaFragC>::dataOffset(GlobalMapping::blockOffsetB(), ldc);
|
|
#pragma unroll
|
|
for(int i = 0; i < BlocksX; i++)
|
|
{
|
|
auto offsetY = 0u;
|
|
#pragma unroll
|
|
for(int j = 0; j < BlocksY; j++)
|
|
{
|
|
globalReadC(fragC[i][j], gAddrC + offsetY, ldc);
|
|
|
|
offsetY += blockStepY;
|
|
}
|
|
gAddrC += blockStepX;
|
|
}
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::globalWriteD(
|
|
GetDataType_t<MfmaFragD>* gAddrD, MfmaFragD const& fragD, uint32_t ldd)
|
|
{
|
|
rocwmma::store_matrix_sync(gAddrD, fragD, ldd);
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <uint32_t BlocksX, uint32_t BlocksY>
|
|
__device__ inline void
|
|
GemmDriver<GemmDriverT_impl>::globalWriteD(GetDataType_t<MfmaFragD>* gAddrD,
|
|
MfmaFragD const (&fragsD)[BlocksX][BlocksY],
|
|
uint32_t ldd)
|
|
{
|
|
auto blockStepX
|
|
= MappingUtil<MfmaFragD>::dataOffset(GlobalMapping::blockOffsetA(), ldd);
|
|
auto blockStepY
|
|
= MappingUtil<MfmaFragD>::dataOffset(GlobalMapping::blockOffsetB(), ldd);
|
|
#pragma unroll
|
|
for(int i = 0; i < BlocksX; i++)
|
|
{
|
|
auto offsetY = 0u;
|
|
#pragma unroll
|
|
for(int j = 0; j < BlocksY; j++)
|
|
{
|
|
globalWriteD(gAddrD + offsetY, fragsD[i][j], ldd);
|
|
|
|
offsetY += blockStepY;
|
|
}
|
|
gAddrD += blockStepX;
|
|
}
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
__device__ inline void
|
|
GemmDriver<GemmDriverT_impl>::uniformFma(MfmaFragD& fragD,
|
|
GetDataType_t<MfmaFragAcc> alpha,
|
|
MfmaFragAcc const& fragAcc,
|
|
GetDataType_t<MfmaFragAcc> beta,
|
|
MfmaFragC const& fragC)
|
|
{
|
|
// Split into chunks to not overwhelm the register capacity
|
|
constexpr uint32_t ChunkSize = min(8u, MfmaFragD::size());
|
|
for(int i = 0; i < MfmaFragD::size() / ChunkSize; i++)
|
|
{
|
|
for(int j = 0; j < ChunkSize; j++)
|
|
{
|
|
auto idx = i * ChunkSize + j;
|
|
// Perform computation in ComputeT and cast back to OutputT
|
|
fragD.x[idx] = static_cast<GetDataType_t<MfmaFragD>>(
|
|
alpha * fragAcc.x[idx]
|
|
+ beta * static_cast<GetDataType_t<MfmaFragAcc>>(fragC.x[idx]));
|
|
}
|
|
}
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <uint32_t BlocksX, uint32_t BlocksY>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::uniformFma(
|
|
MfmaFragD (&fragsD)[BlocksX][BlocksY],
|
|
GetDataType_t<MfmaFragAcc> alpha,
|
|
MfmaFragAcc const (&fragsAcc)[BlocksX][BlocksY],
|
|
GetDataType_t<MfmaFragAcc> beta,
|
|
MfmaFragC const (&fragsC)[BlocksX][BlocksY])
|
|
{
|
|
#pragma unroll
|
|
for(int i = 0; i < BlocksX; i++)
|
|
{
|
|
#pragma unroll
|
|
for(int j = 0; j < BlocksY; j++)
|
|
{
|
|
uniformFma(fragsD[i][j], alpha, fragsAcc[i][j], beta, fragsC[i][j]);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::syncWorkgroup()
|
|
{
|
|
rocwmma::synchronize_workgroup();
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <int32_t priority>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::prioritize_wavefront()
|
|
{
|
|
using SetPrio = SetPrio<priority>;
|
|
SetPrio::exec();
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <int32_t mask>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::sched_barrier()
|
|
{
|
|
using SchedBarrier = SchedBarrier<mask>;
|
|
SchedBarrier::exec();
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <int32_t vmcnt, int32_t lgkmcnt>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::mem_barrier()
|
|
{
|
|
using Waitcnt = Waitcnt<vmcnt, lgkmcnt>;
|
|
Waitcnt::exec();
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <int32_t vmcnt>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::vector_mem_barrier()
|
|
{
|
|
using WaitVmcnt = WaitVmcnt<vmcnt>;
|
|
WaitVmcnt::exec();
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
template <int32_t lgkmcnt>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::lds_mem_barrier()
|
|
{
|
|
using WaitLgkmcnt = WaitLgkmcnt<lgkmcnt>;
|
|
WaitLgkmcnt::exec();
|
|
}
|
|
|
|
template <GemmDriverT>
|
|
__device__ inline void GemmDriver<GemmDriverT_impl>::scheduleKLoop()
|
|
{
|
|
// K-loop rough scheduling:
|
|
// 1. GR
|
|
// 2. LR
|
|
// 3. Mma
|
|
// 4. LW
|
|
SchedBarrierVmemRead::exec();
|
|
SchedBarrierDsRead::exec();
|
|
SchedBarrierMma::exec();
|
|
SchedBarrierDsWrite::exec();
|
|
}
|
|
|
|
#undef GemmDriverT
|
|
#undef GemmDriverT_impl
|
|
|
|
} // namespace CooperativeGemm
|
|
|
|
} // namespace rocwmma
|
|
|
|
#endif // GEMM_DRIVER_IMPL_HPP
|