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
351 lines
12 KiB
C++
351 lines
12 KiB
C++
#ifndef DLRM_LDS_MAPPING_UTIL_HPP
|
|
#define DLRM_LDS_MAPPING_UTIL_HPP
|
|
|
|
#pragma GCC diagnostic push
|
|
#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
|
|
#include <rocwmma/rocwmma.hpp>
|
|
#include <rocwmma/rocwmma_transforms.hpp>
|
|
#pragma GCC diagnostic pop
|
|
|
|
#include <rocwmma/internal/io_config.hpp>
|
|
#include <rocwmma/internal/mapping_util.hpp>
|
|
#include <rocwmma/internal/utils.hpp>
|
|
|
|
namespace rocwmma
|
|
{
|
|
class LdsKW
|
|
{
|
|
};
|
|
class LdsKH
|
|
{
|
|
};
|
|
class LdsKHC
|
|
{
|
|
};
|
|
class LdsRF
|
|
{
|
|
};
|
|
|
|
template <>
|
|
constexpr const char* dataTypeToString<LdsKW>()
|
|
{
|
|
return "LdsKW";
|
|
}
|
|
|
|
template <>
|
|
constexpr const char* dataTypeToString<LdsKH>()
|
|
{
|
|
return "LdsKH";
|
|
}
|
|
|
|
template <>
|
|
constexpr const char* dataTypeToString<LdsKHC>()
|
|
{
|
|
return "LdsKHC";
|
|
}
|
|
|
|
template <>
|
|
constexpr const char* dataTypeToString<LdsRF>()
|
|
{
|
|
return "LdsRF";
|
|
}
|
|
|
|
template <uint32_t BlockM,
|
|
uint32_t BlockN,
|
|
uint32_t BlockK,
|
|
typename DataT,
|
|
typename LayoutA,
|
|
typename LayoutB,
|
|
typename LayoutLds,
|
|
typename MappingLds,
|
|
uint32_t TBlockX,
|
|
uint32_t TBlockY,
|
|
uint32_t BlocksX,
|
|
uint32_t BlocksY>
|
|
struct LdsMappingUtil;
|
|
|
|
template <uint32_t BlockM,
|
|
uint32_t BlockN,
|
|
uint32_t BlockK,
|
|
typename DataT,
|
|
typename LayoutA,
|
|
typename LayoutB,
|
|
typename LayoutLds,
|
|
uint32_t TBlockX,
|
|
uint32_t TBlockY,
|
|
uint32_t BlocksX,
|
|
uint32_t BlocksY>
|
|
struct LdsMappingUtil<BlockM,
|
|
BlockN,
|
|
BlockK,
|
|
DataT,
|
|
LayoutA,
|
|
LayoutB,
|
|
LayoutLds,
|
|
LdsRF,
|
|
TBlockX,
|
|
TBlockY,
|
|
BlocksX,
|
|
BlocksY>
|
|
{
|
|
/* LdsRF (Register file)
|
|
* Lds usage is mapped as a vertical register file of width 64.
|
|
* This geometry is compatible with both A and B fragments, as
|
|
* fragments do not support partial registers, regardless of
|
|
* DataLayout.
|
|
*
|
|
* There is no need for data transpose as A and B should already
|
|
* be aligned.
|
|
*
|
|
* All that really needs to happen is we interpret the fragments
|
|
* as registers and provide 1-1 mapping into LDS memory with width
|
|
* of 64. The matrix_b type of LDS fragments is conducive to this
|
|
* mapping as the current config writes rows of fragment data, such
|
|
* as a vertical register file would need.
|
|
*
|
|
* In this case, block dimensions for A and B cannot be extended by
|
|
* respective BlocksX and BlocksY, due to the LDS geometry which does
|
|
* not easily map to sub-blocks.
|
|
*
|
|
* Local Layout (LDS): e.g. BlockMNK = (32, 32, 8)
|
|
*
|
|
* Elements 0.....31 32.....64
|
|
* _______________
|
|
* Reg0 | C0 | C4 |
|
|
* BlockA Reg1 | C1 | C5 |
|
|
* (0) Reg2 | C2 | C6 |
|
|
* Reg3 | C3 | C7 |
|
|
* ... ... ...
|
|
* _______________
|
|
* Reg0 | C0 | C4 |
|
|
* BlockA Reg1 | C1 | C5 |
|
|
* (X-1) Reg2 | C2 | C6 |
|
|
* Reg3 | C3 | C7 |
|
|
*
|
|
* ... ... ...
|
|
* _______________
|
|
* Reg0 | R0 | R4 |
|
|
* BlockB Reg1 | R1 | R5 |
|
|
* (0) Reg2 | R2 | R6 |
|
|
* Reg3 | R3 | R7 |
|
|
* ... ... ...
|
|
* _______________
|
|
* Reg0 | R0 | R4 |
|
|
* BlockB Reg1 | R1 | R5 |
|
|
* (Y-1) Reg2 | R2 | R6 |
|
|
* Reg3 | R3 | R7 |
|
|
*/
|
|
|
|
static constexpr uint32_t registerFileWidth = Constants::AMDGCN_WAVE_SIZE;
|
|
|
|
// Cooperative fragments
|
|
using GlobalReadFragA = fragment<matrix_a, BlockM, BlockN, BlockK, DataT, LayoutA>;
|
|
|
|
using GlobalReadFragB = fragment<matrix_b, BlockM, BlockN, BlockK, DataT, LayoutB>;
|
|
|
|
// Local Write
|
|
// Vertical register file fulfilled by matrix_b with BlockN = 64
|
|
using LocalWriteFragA
|
|
= apply_data_layout_t<apply_register_file_t<GlobalReadFragA>, LayoutLds>;
|
|
using LocalWriteFragB
|
|
= apply_data_layout_t<apply_register_file_t<GlobalReadFragB>, LayoutLds>;
|
|
|
|
// Sanity checks
|
|
static_assert(GlobalReadFragA::size() == LocalWriteFragA::size(),
|
|
"Incompatible register count");
|
|
static_assert(GlobalReadFragB::size() == LocalWriteFragB::size(),
|
|
"Incompatible register count");
|
|
static_assert(LocalWriteFragA::size() * registerFileWidth == BlockM * BlockK,
|
|
"Incompatible element count");
|
|
static_assert(LocalWriteFragB::size() * registerFileWidth == BlockN * BlockK,
|
|
"Incompatible element count");
|
|
|
|
// Final fragments are same as global
|
|
using FragA = fragment<matrix_a, BlockM, BlockN, BlockK, DataT, LayoutA>;
|
|
using FragB = fragment<matrix_b, BlockM, BlockN, BlockK, DataT, LayoutB>;
|
|
|
|
// Local read same as write
|
|
using LocalReadFragA = apply_data_layout_t<apply_register_file_t<FragA>, LayoutLds>;
|
|
using LocalReadFragB = apply_data_layout_t<apply_register_file_t<FragB>, LayoutLds>;
|
|
|
|
private:
|
|
// Calculate offsets based on DataLayout of LDS, A frags and B frags.
|
|
// All of them will have the same workgroupDim and waveCoord values.
|
|
using LdsOffsets = GetMappingUtil_t<LocalWriteFragA>;
|
|
using GlobalAOffsets = GetMappingUtil_t<GlobalReadFragA>;
|
|
using GlobalBOffsets = GetMappingUtil_t<GlobalReadFragB>;
|
|
|
|
template <typename FragT>
|
|
using IOTraits = typename GetIOConfig_t<FragT>::IOTraits;
|
|
|
|
public:
|
|
__device__ static inline auto ldsWidth()
|
|
{
|
|
return registerFileWidth;
|
|
}
|
|
|
|
__device__ static inline auto ldsHeight()
|
|
{
|
|
auto workgroupDim = LdsOffsets::workgroupDim();
|
|
return LocalWriteFragA::size() * BlocksX * get<0>(workgroupDim)
|
|
+ LocalWriteFragB::size() * BlocksY * get<1>(workgroupDim);
|
|
}
|
|
|
|
__device__ static inline uint32_t ld()
|
|
{
|
|
return std::is_same<LayoutLds, row_major>::value ? ldsWidth() : ldsHeight();
|
|
}
|
|
|
|
__device__ static inline uint32_t baseOffsetA()
|
|
{
|
|
return 0;
|
|
}
|
|
|
|
__device__ static inline uint32_t baseOffsetB()
|
|
{
|
|
auto matrixCoord = make_coord2d(
|
|
LocalWriteFragA::size() * BlocksX * get<0>(LdsOffsets::workgroupDim()), 0);
|
|
return LdsOffsets::dataOffset(matrixCoord, ld());
|
|
}
|
|
|
|
__device__ static inline uint32_t waveOffsetA()
|
|
{
|
|
auto matrixCoord = make_coord2d(
|
|
LocalWriteFragA::size() * BlocksX * get<0>(LdsOffsets::waveCoord()), 0);
|
|
return LdsOffsets::dataOffset(matrixCoord, ld());
|
|
}
|
|
|
|
__device__ static inline uint32_t waveOffsetB()
|
|
{
|
|
auto matrixCoord = make_coord2d(
|
|
LocalWriteFragB::size() * BlocksY * get<1>(LdsOffsets::waveCoord()), 0);
|
|
return LdsOffsets::dataOffset(matrixCoord, ld());
|
|
}
|
|
|
|
__device__ static inline uint32_t blockOffsetA(uint32_t blockX)
|
|
{
|
|
auto matrixCoord = make_coord2d(LocalWriteFragA::size() * blockX, 0);
|
|
return LdsOffsets::dataOffset(matrixCoord, ld());
|
|
}
|
|
|
|
__device__ static inline uint32_t blockOffsetB(uint32_t blockY)
|
|
{
|
|
auto matrixCoord = make_coord2d(LocalWriteFragB::size() * blockY, 0);
|
|
return LdsOffsets::dataOffset(matrixCoord, ld());
|
|
}
|
|
|
|
__device__ static inline void
|
|
prefetchCoopGlobalA(DataT* baseLds, DataT const* baseA, uint32_t lda)
|
|
{
|
|
// Cooperative transfer
|
|
for(int32_t i = 0; i < BlocksX; ++i)
|
|
{
|
|
// Issue global read
|
|
GlobalReadFragA fetchA;
|
|
load_matrix_sync(fetchA,
|
|
baseA
|
|
+ GlobalAOffsets::dataOffset(make_coord2d(BlockM * i, 0), lda),
|
|
lda);
|
|
|
|
// Issue local store
|
|
store_matrix_sync(baseLds + baseOffsetA() + waveOffsetA() + blockOffsetA(i),
|
|
reinterpret_cast<LocalWriteFragA&>(fetchA),
|
|
ld());
|
|
}
|
|
}
|
|
|
|
__device__ static inline void
|
|
prefetchCoopGlobalB(DataT* baseLds, DataT const* baseB, uint32_t ldb)
|
|
{
|
|
// Cooperative transfer
|
|
for(int32_t i = 0; i < BlocksY; ++i)
|
|
{
|
|
// Issue global read
|
|
GlobalReadFragB fetchB;
|
|
load_matrix_sync(fetchB,
|
|
baseB
|
|
+ GlobalBOffsets::dataOffset(make_coord2d(0, BlockN * i), ldb),
|
|
ldb);
|
|
|
|
// Issue local store
|
|
store_matrix_sync(baseLds + baseOffsetB() + waveOffsetB() + blockOffsetB(i),
|
|
reinterpret_cast<LocalWriteFragB&>(fetchB),
|
|
ld());
|
|
}
|
|
}
|
|
|
|
__device__ static inline void
|
|
prefetchGlobalA(DataT* baseLds, DataT const* baseA, uint32_t lda)
|
|
{
|
|
for(int32_t i = 0; i < BlocksX; ++i)
|
|
{
|
|
// Issue global read
|
|
GlobalReadFragA fetchA;
|
|
load_matrix_sync(fetchA,
|
|
baseA
|
|
+ GlobalAOffsets::dataOffset(make_coord2d(BlockM * i, 0), lda),
|
|
lda);
|
|
|
|
// Issue local store
|
|
store_matrix_sync(baseLds + baseOffsetA() + waveOffsetA() + blockOffsetA(i),
|
|
reinterpret_cast<LocalWriteFragA&>(fetchA),
|
|
ld());
|
|
}
|
|
}
|
|
|
|
__device__ static inline void
|
|
prefetchGlobalB(DataT* baseLds, DataT const* baseB, uint32_t ldb)
|
|
{
|
|
for(int32_t i = 0; i < BlocksY; ++i)
|
|
{
|
|
// Issue global read
|
|
GlobalReadFragB fetchB;
|
|
load_matrix_sync(fetchB,
|
|
baseB
|
|
+ GlobalBOffsets::dataOffset(make_coord2d(0, BlockN * i), ldb),
|
|
ldb);
|
|
|
|
// Issue local store
|
|
store_matrix_sync(baseLds + baseOffsetB() + waveOffsetB() + blockOffsetB(i),
|
|
reinterpret_cast<LocalWriteFragB&>(fetchB),
|
|
ld());
|
|
}
|
|
}
|
|
|
|
__device__ static inline void
|
|
prefetchLocalA(FragA& fragA, DataT const* baseLds, uint32_t blockX)
|
|
{
|
|
load_matrix_sync(reinterpret_cast<LocalReadFragA&>(fragA),
|
|
baseLds + baseOffsetA() + waveOffsetA() + blockOffsetA(blockX),
|
|
ld());
|
|
}
|
|
|
|
__device__ static inline void
|
|
prefetchLocalB(FragB& fragB, DataT const* baseLds, uint32_t blockY)
|
|
{
|
|
load_matrix_sync(reinterpret_cast<LocalReadFragB&>(fragB),
|
|
baseLds + baseOffsetB() + waveOffsetB() + blockOffsetB(blockY),
|
|
ld());
|
|
}
|
|
|
|
__device__ static inline void prefetchLocalA(FragA* fragsA, DataT const* baseLds)
|
|
{
|
|
for(int i = 0; i < BlocksX; i++)
|
|
{
|
|
prefetchLocalA(fragsA[i], baseLds, i);
|
|
}
|
|
}
|
|
|
|
__device__ static inline void prefetchLocalB(FragB* fragsB, DataT const* baseLds)
|
|
{
|
|
for(int i = 0; i < BlocksY; i++)
|
|
{
|
|
prefetchLocalB(fragsB[i], baseLds, i);
|
|
}
|
|
}
|
|
};
|
|
|
|
} // namespace rocwmma
|
|
|
|
#endif // DLRM_LDS_MAPPING_UTIL_HPP
|