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
560 lines
27 KiB
C++
560 lines
27 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_GLOBAL_MAPPING_HPP
|
|
#define GEMM_GLOBAL_MAPPING_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 GlobalMapping
|
|
{
|
|
namespace detail
|
|
{
|
|
template <uint32_t BlockM, // MFMA BlockM
|
|
uint32_t BlockN, // MFMA BlockN
|
|
uint32_t BlockK, // MFMA BlockK
|
|
typename InputT, // Type of A / B data
|
|
typename OutputT, // Type of C / D data
|
|
typename ComputeT, // Type of accumulator data
|
|
typename LayoutA, // Data Layout of A (row / col major)
|
|
typename LayoutB, // Data Layout of B (row / col major)
|
|
typename LayoutC, // Data Layout of C (row / col major)
|
|
typename LayoutD, // Data Layout of D (row / col major)
|
|
uint32_t BlocksX, // MFMA blocks per wave in X direction
|
|
uint32_t BlocksY, // MFMA blocks per wave in Y direction
|
|
uint32_t TBlockX = 0, // Thread block X dimension
|
|
uint32_t TBlockY = 0> // Thread block Y dimension
|
|
struct MappingBase
|
|
{
|
|
/*
|
|
* This global mapping aligns each wave with a gemm C/D tile
|
|
* that the wave will compute independently. All calculations
|
|
* are in matrix coordinates (Coord2d).
|
|
*
|
|
* Coordinates are scaled into three different levels:
|
|
* - Macro level: context of work for entire workgroup
|
|
* - Wave level: context of work for each wave
|
|
* - Block level: context of work for each mfma block
|
|
*
|
|
* Waves computing tiles on the same row (matrix A) and waves
|
|
* computing tiles on the same col (matrix B) have opportunities
|
|
* for data sharing. As a result, waves collaboration is limited
|
|
* to tiles in same locality.
|
|
*
|
|
* Block size: (BlockM x BlockN)
|
|
* Wave tile size: (BlocksX * BlockSize.x) x (BlocksY * BlockSize.y)
|
|
* Macro Tile size: (WgDim.x * WaveTileSize.x) x (WgDim.y * WaveTileSize.y)
|
|
*
|
|
* Wave data share A: same row
|
|
* Wave data share B: same col
|
|
*
|
|
* Global C layout & workgroup assignment for WGDim = (2x2):
|
|
*
|
|
* W (X, Y) = wave row X, col Y
|
|
*
|
|
*
|
|
* |--------- Macro Tile Y-------------|
|
|
* |-- Wave Tile Y --|
|
|
* |-BlockN-|
|
|
*
|
|
* BlockN x BlocksY BlockN x BlocksY
|
|
* |<--------------->|<--------------->|
|
|
* _ _ _ _ _ _ ___ ________ ________ ________ ________
|
|
* | | | ^ | | | | |
|
|
* | Wave| BlockM | BlockM | | W | W |
|
|
* | Tile| _|_ x | |__ (0, 0) __|__ (0, 1) __|
|
|
* | X | BlocksX | | | |
|
|
* Macro | | | | | |
|
|
* Tile | _|_ _v_ |________|________|________|________|
|
|
* X | ^ | | | | |
|
|
* | BlockM | | W | W |
|
|
* | x | |__ (1, 0) __|__ (1, 1) __|
|
|
* | BlocksX | | | |
|
|
* | | | | |
|
|
* _|_ _v_ |________|________|________|________|
|
|
*
|
|
*
|
|
* TLDR: Global mapping aligns global offsets for A / B / C / D MFMA blocks.
|
|
* Configures fragment and buffer types for per-wave responsibility of
|
|
* BlocksX * BlocksY MFMA blocks.
|
|
*/
|
|
|
|
// Block fragments that are used as inputs to MFMA operations.
|
|
using MfmaFragA = fragment<matrix_a, BlockM, BlockN, BlockK, InputT, LayoutA>;
|
|
using MfmaFragB = fragment<matrix_b, BlockM, BlockN, BlockK, InputT, LayoutB>;
|
|
using MfmaFragC = fragment<accumulator, BlockM, BlockN, BlockK, OutputT, LayoutC>;
|
|
using MfmaFragD = fragment<accumulator, BlockM, BlockN, BlockK, OutputT, LayoutD>;
|
|
using MfmaFragAcc = fragment<accumulator, BlockM, BlockN, BlockK, ComputeT>;
|
|
|
|
// Mfma fragment buffers required for the gemm driver
|
|
using MfmaBuffA = MfmaFragA[BlocksX];
|
|
using MfmaBuffB = MfmaFragB[BlocksY];
|
|
using MfmaBuffC = MfmaFragC[BlocksX][BlocksY];
|
|
using MfmaBuffD = MfmaFragD[BlocksX][BlocksY];
|
|
using MfmaBuffAcc = MfmaFragAcc[BlocksX][BlocksY];
|
|
|
|
using WaveSpace = typename rocwmma::detail::WaveSpace<TBlockX, TBlockY>;
|
|
|
|
// Projection of C coordinate in direction of A
|
|
template <typename CoordC>
|
|
__device__ constexpr static inline auto projCoordA(CoordC const& coordC);
|
|
|
|
// Projection of C coordinate in direction of B
|
|
template <typename CoordC>
|
|
__device__ constexpr static inline auto projCoordB(CoordC const& coordC);
|
|
|
|
///
|
|
/// Dimensions
|
|
///
|
|
|
|
// Macro tile = tile processed by entire workgroup
|
|
// Wave tile = tile processed by current wave
|
|
// Block size = mfma block size
|
|
__device__ constexpr static inline auto macroTileSizeC();
|
|
__device__ constexpr static inline auto waveTileSizeC();
|
|
__device__ constexpr static inline auto blockSizeC();
|
|
__device__ constexpr static inline auto kDim();
|
|
|
|
///
|
|
/// Offsets
|
|
///
|
|
|
|
// The offset from macro tile to the mfma tiles for current wave
|
|
__device__ constexpr static inline auto waveOffsetA();
|
|
__device__ constexpr static inline auto waveOffsetB();
|
|
__device__ constexpr static inline auto waveOffsetC();
|
|
|
|
// The local offsets between mfma blocks
|
|
__device__ constexpr static inline auto blockOffsetA();
|
|
__device__ constexpr static inline auto blockOffsetB();
|
|
__device__ constexpr static inline auto blockOffsetC();
|
|
|
|
// The matrix offset to the next step in the k dimension
|
|
__device__ constexpr static inline auto kStepOffsetA();
|
|
__device__ constexpr static inline auto kStepOffsetB();
|
|
|
|
///
|
|
/// Global matrix coords
|
|
///
|
|
|
|
// Global matrix coordinate of macro tile for the current workgroup
|
|
__device__ constexpr static inline auto macroTileCoordC();
|
|
|
|
// Global matrix coordinate of wave tile for the current wave
|
|
__device__ constexpr static inline auto waveTileCoordC();
|
|
};
|
|
|
|
} // namespace detail
|
|
|
|
template <uint32_t BlockM,
|
|
uint32_t BlockN,
|
|
uint32_t BlockK,
|
|
typename InputT,
|
|
typename OutputT,
|
|
typename ComputeT,
|
|
typename LayoutA,
|
|
typename LayoutB,
|
|
typename LayoutC,
|
|
typename LayoutD,
|
|
uint32_t BlocksX,
|
|
uint32_t BlocksY,
|
|
uint32_t TBlockX = 0,
|
|
uint32_t TBlockY = 0>
|
|
struct BlockLevelMapping : public detail::MappingBase<BlockM,
|
|
BlockN,
|
|
BlockK,
|
|
InputT,
|
|
OutputT,
|
|
ComputeT,
|
|
LayoutA,
|
|
LayoutB,
|
|
LayoutC,
|
|
LayoutD,
|
|
BlocksX,
|
|
BlocksY,
|
|
TBlockX,
|
|
TBlockY>
|
|
{
|
|
/*
|
|
* This flavour of Global Mapping targets A/B/C/D wave tiles iteratively
|
|
* in MFMA fragment chunks:
|
|
* BlocksX x (BlockM x BlockK) for A
|
|
* BlocksY x (BlockN x BlockK) for B
|
|
* BlocksX x BlocksY x (BlockM x BlockN) for C / D
|
|
*
|
|
* This leads to favourable MFMA layouts directly from Global Reads,
|
|
* but may not be the most efficient in some layouts.
|
|
*/
|
|
|
|
using Base = detail::MappingBase<BlockM,
|
|
BlockN,
|
|
BlockK,
|
|
InputT,
|
|
OutputT,
|
|
ComputeT,
|
|
LayoutA,
|
|
LayoutB,
|
|
LayoutC,
|
|
LayoutD,
|
|
BlocksX,
|
|
BlocksY,
|
|
TBlockX,
|
|
TBlockY>;
|
|
|
|
// Global wave tile R/W be in sections of MFMA sized fragments
|
|
using GRFragA = fragment<matrix_a,
|
|
BlockM,
|
|
BlockN,
|
|
BlockK,
|
|
InputT,
|
|
LayoutA,
|
|
fragment_scheduler::coop_row_slice_2d<TBlockX, TBlockY>>;
|
|
using GRFragB = fragment<matrix_b,
|
|
BlockM,
|
|
BlockN,
|
|
BlockK,
|
|
InputT,
|
|
LayoutB,
|
|
fragment_scheduler::coop_col_slice_2d<TBlockX, TBlockY>>;
|
|
using GRFragC = typename Base::MfmaFragC;
|
|
using GWFragD = typename Base::MfmaFragD;
|
|
|
|
// Global wave tile R/W will use MFMA sized fragment buffers
|
|
using GRBuffA = GRFragA[BlocksX];
|
|
using GRBuffB = GRFragB[BlocksY];
|
|
using GRBuffC = typename Base::MfmaBuffC;
|
|
using GWBuffD = typename Base::MfmaBuffD;
|
|
|
|
// The base global matrix coordinate of the current wave tile.
|
|
__device__ constexpr static inline auto readCoordA()
|
|
{
|
|
return Base::projCoordA(readCoordC());
|
|
}
|
|
__device__ constexpr static inline auto readCoordB()
|
|
{
|
|
return Base::projCoordB(readCoordC());
|
|
}
|
|
__device__ constexpr static inline auto readCoordC()
|
|
{
|
|
return Base::waveTileCoordC();
|
|
}
|
|
__device__ constexpr static inline auto writeCoordD()
|
|
{
|
|
return Base::waveTileCoordC();
|
|
}
|
|
|
|
// Indicate if global read for A and B are wave tiles
|
|
__device__ constexpr static inline auto readABWaveTile()
|
|
{
|
|
return true;
|
|
}
|
|
};
|
|
|
|
template <uint32_t BlockM,
|
|
uint32_t BlockN,
|
|
uint32_t BlockK,
|
|
typename InputT,
|
|
typename OutputT,
|
|
typename ComputeT,
|
|
typename LayoutA,
|
|
typename LayoutB,
|
|
typename LayoutC,
|
|
typename LayoutD,
|
|
uint32_t BlocksX,
|
|
uint32_t BlocksY,
|
|
uint32_t TBlockX = 0,
|
|
uint32_t TBlockY = 0>
|
|
struct WaveLevelMapping : public detail::MappingBase<BlockM,
|
|
BlockN,
|
|
BlockK,
|
|
InputT,
|
|
OutputT,
|
|
ComputeT,
|
|
LayoutA,
|
|
LayoutB,
|
|
LayoutC,
|
|
LayoutD,
|
|
BlocksX,
|
|
BlocksY,
|
|
TBlockX,
|
|
TBlockY>
|
|
{
|
|
/*
|
|
* This flavour of Global Mapping targets A/B as a single wave tile sized fragment.
|
|
* C/D wave tiles are targeted iteratively in MFMA fragment chunks:
|
|
* (BlocksX * BlockM) x BlockK for A
|
|
* (BlocksY * BlockN) x BlockK for B
|
|
* BlocksX x BlocksY x (BlockM x BlockN) for C / D
|
|
*
|
|
* This global read A/B fragments are not MFMA friendly, however when written to LDS
|
|
* smaller MFMA friendly fragment may be read directly from the same LDS layout.
|
|
* Larger GR may be more efficient in certain layouts for data pipelining.
|
|
*
|
|
* This layout can accommodate either compile time or runtime TBlockX/Y params.
|
|
*/
|
|
using Base = detail::MappingBase<BlockM,
|
|
BlockN,
|
|
BlockK,
|
|
InputT,
|
|
OutputT,
|
|
ComputeT,
|
|
LayoutA,
|
|
LayoutB,
|
|
LayoutC,
|
|
LayoutD,
|
|
BlocksX,
|
|
BlocksY,
|
|
TBlockX,
|
|
TBlockY>;
|
|
|
|
// Global reads for A/B are single fragment of wave tile size
|
|
// Global R/W for C/D are MFMA sized fragments
|
|
using GRFragA = fragment<matrix_a,
|
|
BlockM * BlocksX,
|
|
BlockN,
|
|
BlockK,
|
|
InputT,
|
|
LayoutA,
|
|
fragment_scheduler::coop_row_slice_2d<TBlockX, TBlockY>>;
|
|
using GRFragB = fragment<matrix_b,
|
|
BlockM,
|
|
BlockN * BlocksY,
|
|
BlockK,
|
|
InputT,
|
|
LayoutB,
|
|
fragment_scheduler::coop_col_slice_2d<TBlockX, TBlockY>>;
|
|
using GRFragC = typename Base::MfmaFragC;
|
|
using GWFragD = typename Base::MfmaFragD;
|
|
|
|
// Global reads for A/B will have one fragment buffer
|
|
// Global R/W for C/D will use MFMA sized multiple fragment buffer
|
|
using GRBuffA = GRFragA;
|
|
using GRBuffB = GRFragB;
|
|
using GRBuffC = typename Base::MfmaBuffC;
|
|
using GWBuffD = typename Base::MfmaBuffD;
|
|
|
|
// A/B global reads will collaborate on wave tile
|
|
__device__ constexpr static inline auto readCoordA()
|
|
{
|
|
return Base::projCoordA(readCoordC());
|
|
}
|
|
__device__ constexpr static inline auto readCoordB()
|
|
{
|
|
return Base::projCoordB(readCoordC());
|
|
}
|
|
|
|
// C/D global R/W on wave tile
|
|
__device__ constexpr static inline auto readCoordC()
|
|
{
|
|
return Base::waveTileCoordC();
|
|
}
|
|
|
|
__device__ constexpr static inline auto writeCoordD()
|
|
{
|
|
return Base::waveTileCoordC();
|
|
}
|
|
|
|
// Indicate if global read for A and B are wave tiles
|
|
__device__ constexpr static inline auto readABWaveTile()
|
|
{
|
|
return true;
|
|
}
|
|
};
|
|
|
|
template <uint32_t BlockM,
|
|
uint32_t BlockN,
|
|
uint32_t BlockK,
|
|
typename InputT,
|
|
typename OutputT,
|
|
typename ComputeT,
|
|
typename LayoutA,
|
|
typename LayoutB,
|
|
typename LayoutC,
|
|
typename LayoutD,
|
|
uint32_t BlocksX,
|
|
uint32_t BlocksY,
|
|
uint32_t TBlockX,
|
|
uint32_t TBlockY>
|
|
struct WorkgroupLevelMapping : public detail::MappingBase<BlockM,
|
|
BlockN,
|
|
BlockK,
|
|
InputT,
|
|
OutputT,
|
|
ComputeT,
|
|
LayoutA,
|
|
LayoutB,
|
|
LayoutC,
|
|
LayoutD,
|
|
BlocksX,
|
|
BlocksY,
|
|
TBlockX,
|
|
TBlockY>
|
|
{
|
|
|
|
// Must provide valid TBlockX/Y params at compile time.
|
|
static_assert((TBlockX > 0) && (TBlockX % Constants::AMDGCN_WAVE_SIZE == 0),
|
|
"Invalid TBlockX dimension");
|
|
static_assert(TBlockY > 0, "Invalid TBlockY dimension");
|
|
|
|
/*
|
|
* This flavour of Global Mapping targets A/B as a single macro tile sized fragment.
|
|
* C/D wave tiles are targeted iteratively in MFMA fragment chunks:
|
|
* (WgX * BlocksX * BlockM) x BlockK for A
|
|
* (WgY * BlocksY * BlockN) x BlockK for B
|
|
* BlocksX x BlocksY x (BlockM x BlockN) for C / D
|
|
*
|
|
* This global read A/B fragments are not MFMA friendly, however when written to LDS
|
|
* smaller MFMA friendly fragment may be read directly from the same LDS layout.
|
|
* Larger GR may be more efficient in certain layouts for data pipelining.
|
|
*/
|
|
using Base = detail::MappingBase<BlockM,
|
|
BlockN,
|
|
BlockK,
|
|
InputT,
|
|
OutputT,
|
|
ComputeT,
|
|
LayoutA,
|
|
LayoutB,
|
|
LayoutC,
|
|
LayoutD,
|
|
BlocksX,
|
|
BlocksY,
|
|
TBlockX,
|
|
TBlockY>;
|
|
|
|
static constexpr uint32_t WaveTileM = BlocksX * BlockM;
|
|
static constexpr uint32_t WaveTileN = BlocksY * BlockN;
|
|
static constexpr uint32_t WaveTileK = BlockK;
|
|
|
|
static constexpr uint32_t MacroTileM
|
|
= TBlockX / Constants::AMDGCN_WAVE_SIZE * WaveTileM;
|
|
static constexpr uint32_t MacroTileN = TBlockY * WaveTileN;
|
|
static constexpr uint32_t MacroTileK = BlockK;
|
|
|
|
// Wave-tiles
|
|
using WaveTileA = fragment<matrix_a, WaveTileM, WaveTileN, WaveTileK, InputT, LayoutA>;
|
|
using WaveTileB = fragment<matrix_b, WaveTileM, WaveTileN, WaveTileK, InputT, LayoutB>;
|
|
using WaveTileC
|
|
= fragment<accumulator, WaveTileM, WaveTileN, WaveTileK, OutputT, LayoutC>;
|
|
using WaveTileD
|
|
= fragment<accumulator, WaveTileM, WaveTileN, WaveTileK, OutputT, LayoutD>;
|
|
using WaveTileAcc = fragment<accumulator, WaveTileM, WaveTileN, WaveTileK, ComputeT>;
|
|
|
|
// We can guess layouts are interleaved by testing A's MatrixLayout.
|
|
static constexpr bool IsInterleaved = layout_traits<
|
|
typename GetIOConfig_t<WaveTileA>::IOLayout::MatrixLayout>::is_interleaved;
|
|
|
|
// Mma operations
|
|
// Interleaved layouts use wave tile mma inputs
|
|
// Non-interleaved layouts use block tile mma inputs
|
|
using MfmaFragA = conditional_t<IsInterleaved, WaveTileA, typename Base::MfmaFragA>;
|
|
using MfmaFragB = conditional_t<IsInterleaved, WaveTileB, typename Base::MfmaFragB>;
|
|
using MfmaFragC = conditional_t<IsInterleaved, WaveTileC, typename Base::MfmaFragC>;
|
|
using MfmaFragD = conditional_t<IsInterleaved, WaveTileD, typename Base::MfmaFragD>;
|
|
using MfmaFragAcc
|
|
= conditional_t<IsInterleaved, WaveTileAcc, typename Base::MfmaFragAcc>;
|
|
|
|
// Mfma fragment buffers required for the gemm driver.
|
|
// Interleaved layouts mma inputs are buffered in single wave tile
|
|
// Non-interleaved layouts mma inputs are buffered in array of block tiles
|
|
using MfmaBuffA = conditional_t<IsInterleaved, MfmaFragA, typename Base::MfmaBuffA>;
|
|
using MfmaBuffB = conditional_t<IsInterleaved, MfmaFragB, typename Base::MfmaBuffB>;
|
|
using MfmaBuffC = conditional_t<IsInterleaved, MfmaFragC, typename Base::MfmaBuffC>;
|
|
using MfmaBuffD = conditional_t<IsInterleaved, MfmaFragD, typename Base::MfmaBuffD>;
|
|
using MfmaBuffAcc
|
|
= conditional_t<IsInterleaved, MfmaFragAcc, typename Base::MfmaBuffAcc>;
|
|
|
|
// Global reads for A/B are macro tiles
|
|
// Global C/D:
|
|
// Interleaved layouts wave tile
|
|
// Non-interleaved layouts block tile
|
|
using GRFragA = fragment<matrix_a,
|
|
MacroTileM,
|
|
MacroTileN,
|
|
MacroTileK,
|
|
InputT,
|
|
LayoutA,
|
|
fragment_scheduler::coop_row_major_2d<TBlockX, TBlockY>>;
|
|
using GRFragB = fragment<matrix_b,
|
|
MacroTileM,
|
|
MacroTileN,
|
|
MacroTileK,
|
|
InputT,
|
|
LayoutB,
|
|
fragment_scheduler::coop_row_major_2d<TBlockX, TBlockY>>;
|
|
using GRFragC = conditional_t<IsInterleaved, WaveTileC, typename Base::MfmaFragC>;
|
|
using GWFragD = conditional_t<IsInterleaved, WaveTileD, typename Base::MfmaFragD>;
|
|
|
|
// Global reads for A/B will have one macro tile buffer
|
|
// Global R/W:
|
|
// Interleaved layouts buffer will be wave-tile
|
|
// Non-interleaved layouts buffer will be array of block tiles
|
|
using GRBuffA = GRFragA;
|
|
using GRBuffB = GRFragB;
|
|
using GRBuffC = conditional_t<IsInterleaved, GRFragC, typename Base::MfmaBuffC>;
|
|
using GWBuffD = conditional_t<IsInterleaved, GWFragD, typename Base::MfmaBuffD>;
|
|
|
|
// A/B global reads will collaborate on macro tile
|
|
__device__ constexpr static inline auto readCoordA()
|
|
{
|
|
return Base::projCoordA(Base::macroTileCoordC());
|
|
}
|
|
__device__ constexpr static inline auto readCoordB()
|
|
{
|
|
return Base::projCoordB(Base::macroTileCoordC());
|
|
}
|
|
|
|
// C/D global R/W on wave tile
|
|
__device__ constexpr static inline auto readCoordC()
|
|
{
|
|
return Base::waveTileCoordC();
|
|
}
|
|
|
|
__device__ constexpr static inline auto writeCoordD()
|
|
{
|
|
return Base::waveTileCoordC();
|
|
}
|
|
|
|
// Indicate if global read for A and B are wave tiles
|
|
__device__ constexpr static inline auto readABWaveTile()
|
|
{
|
|
return false;
|
|
}
|
|
};
|
|
|
|
} // namespace GlobalMapping
|
|
|
|
} // namespace rocwmma
|
|
|
|
#include "gemm_global_mapping_impl.hpp"
|
|
|
|
#endif // GEMM_GLOBAL_MAPPING_HPP
|