/******************************************************************************* * * 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 #include #pragma GCC diagnostic pop namespace rocwmma { namespace GlobalMapping { namespace detail { template // 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; using MfmaFragB = fragment; using MfmaFragC = fragment; using MfmaFragD = fragment; using MfmaFragAcc = fragment; // 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; // Projection of C coordinate in direction of A template __device__ constexpr static inline auto projCoordA(CoordC const& coordC); // Projection of C coordinate in direction of B template __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 struct BlockLevelMapping : public detail::MappingBase { /* * 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; // Global wave tile R/W be in sections of MFMA sized fragments using GRFragA = fragment>; using GRFragB = fragment>; 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 struct WaveLevelMapping : public detail::MappingBase { /* * 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; // 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>; using GRFragB = fragment>; 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 struct WorkgroupLevelMapping : public detail::MappingBase { // 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; 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; using WaveTileB = fragment; using WaveTileC = fragment; using WaveTileD = fragment; using WaveTileAcc = fragment; // We can guess layouts are interleaved by testing A's MatrixLayout. static constexpr bool IsInterleaved = layout_traits< typename GetIOConfig_t::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; using MfmaFragB = conditional_t; using MfmaFragC = conditional_t; using MfmaFragD = conditional_t; using MfmaFragAcc = conditional_t; // 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; using MfmaBuffB = conditional_t; using MfmaBuffC = conditional_t; using MfmaBuffD = conditional_t; using MfmaBuffAcc = conditional_t; // Global reads for A/B are macro tiles // Global C/D: // Interleaved layouts wave tile // Non-interleaved layouts block tile using GRFragA = fragment>; using GRFragB = fragment>; using GRFragC = conditional_t; using GWFragD = conditional_t; // 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; using GWBuffD = conditional_t; // 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