/******************************************************************************* * * 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 ROCWMMA_API_HPP #define ROCWMMA_API_HPP #include "internal/accessors.hpp" #include "internal/io_scheduler.hpp" #include "internal/io_traits.hpp" #include "internal/pack_util.hpp" #include "internal/types.hpp" /** * \mainpage * * rocWMMA is a C++ header library for accelerating mixed precision matrix multiply-accumulate operations * leveraging specialized GPU matrix cores on AMD's latest discrete GPUs. 'roc' being an AMD-specific * component belonging to the ROCm ecosystem, and WMMA stands for Wavefront Mixed precision Multiply Accumulate. * * rocWMMA leverages modern C++ techniques. It is templated for modularity and uses meta-programming paradigms to provide opportunities for customization * and compile-time inferences and optimizations. The API is seamless across supported CDNA and RDNA architectures. It is also portable with the Nvidia * nvcuda::wmma library, allowing those users to easily migrate to the AMD platform. * * The API is implemented as GPU device code which empowers users with direct use of GPU matrix cores, right from their kernel code. * Major benefits include kernel-level control which allows authoring flexibility and accessibility to compiler optimization passes in-situ * with other device code. Users can therefore decide when and where kernel run-time launches are required, which is not dictated by the API. * * rocWMMA's API facilitates the decomposition of matrix multiply-accumulate problems into discretized blocks (also known as fragments) and enables * parallelization of block-wise operations across multiple GPU wavefronts. The programmer's perspective is simplified to wavefront handling of fragments, * whereas individual threads are handled internally. This can allow for faster development times and a more seamless experience across multiple architectures. * API functions include data loading and storing, matrix multiply-accumulate and helper transforms that operate on data fragment abstractions. Moreover, data movement * between global and local memory can be done cooperatively amongst the wavefronts in a threadblock to enable data sharing and re-use. Matrix multiply-accumulate * functionality supports mixed precision inputs and outputs with native fixed-precision accumulation. * * Supporting code is required for GPU device management and kernel invocation. The kernel code samples and tests provided are built and launched via * the Heterogeneous-Compute Interface for Portability (HIP) ecosystem within ROCm. * * This library is an ongoing Work-In-Progress (WIP). * * For more documentation, please visit https://rocm.docs.amd.com/projects/rocWMMA/en/latest/index.html. * */ namespace rocwmma { //! @defgroup Rocwmma rocWMMA Public API //! //! @brief rocWMMA objects and API function definitions. //! @{ //! @struct row_major //! @brief Meta-tag indicating 2D in-memory data layout as row major. struct row_major { }; //! @struct col_major //! @brief Meta-tag indicating 2D in-memory data layout as column major. struct col_major { }; //! @struct matrix_a //! @brief Meta-tag indicating data context is input Matrix A. struct matrix_a { }; //! @struct matrix_b //! @brief Meta-tag indicating data context is input Matrix B. struct matrix_b { }; //! @struct accumulator //! @brief Meta-tag indicating data context is Accumulator (also used as Matrix C / D). struct accumulator { }; //! @struct layout_t //! @brief Runtime data layout tags //! @var mem_row_major //! @var mem_col_major enum layout_t : uint32_t { mem_row_major, mem_col_major }; namespace fragment_scheduler { //! @typedef default_schedule //! @brief The default fragment scheduler; each wave operates independently. using default_schedule = IOScheduler::Default; //! @typedef coop_row_major_2d //! @brief A cooperative scheduling strategy where each wave in the 2d threadblock //! will contribute to the fragment operation in row_major grid order. //! All waves are scheduled in row_major order. //! E.g. (TBlockX, TBlockY) => 2x2 waves //! w0 = (0, 0), w1 = (0, 1), //! w2 = (1, 0), w3 = (1, 1) //! @tparam TBlockX the size of the thread-block in the X dimension //! @tparam TBlockY the size of the thread-block in the Y dimension template using coop_row_major_2d = IOScheduler::RowMajor2d; //! @typedef coop_col_major_2d //! @brief A cooperative scheduling strategy where each wave in the 2d threadblock //! will contribute to the fragment operation in col_major grid order. //! All waves are scheduled in row_major order. //! E.g. (TBlockX, TBlockY) => 2x2 waves //! w0 = (0, 0), w2 = (0, 1), //! w1 = (1, 0), w3 = (1, 1) //! @tparam TBlockX the size of the thread-block in the X dimension //! @tparam TBlockY the size of the thread-block in the Y dimension template using coop_col_major_2d = IOScheduler::ColMajor2d; //! @typedef coop_row_slice_2d //! @brief A cooperative scheduling strategy where each row of waves //! in the 2d threadblock will contribute to the fragment operation. //! Waves are partitioned into rows. Only waves in the same row //! participate together. //! E.g. (TBlockX, TBlockY) = 2x2 waves //! RowSlice0: w0 = (0, 0), w1 = (0, 1) //! RowSlice1: w0 = (1, 0), w1 = (1, 1) //! @tparam TBlockX the size of the thread-block in the X dimension //! @tparam TBlockY the size of the thread-block in the Y dimension template using coop_row_slice_2d = IOScheduler::RowSlice2d; //! @typedef coop_col_slice_2d //! @brief A cooperative scheduling strategy where each col of waves //! in the 2d threadblock will contribute to the fragment operation. //! Waves are partitioned into cols. Only waves in the same col //! participate together. //! E.g. (TBlockX, TBlockY) = 2x2 waves //! ColSlice0: ColSlice1: //! w0 = (0, 0), w0 = (0, 1), //! w1 = (1, 0) w1 = (1, 1) //! @tparam TBlockX the size of the thread-block in the X dimension //! @tparam TBlockY the size of the thread-block in the Y dimension template using coop_col_slice_2d = IOScheduler::ColSlice2d; //! @typedef single //! @brief A cooperative scheduling strategy where only one wave in //! the thread block will participate. //! @tparam TBlockX the size of the thread-block in the X dimension //! @tparam TBlockY the size of the thread-block in the Y dimension //! @tparam WaveIdx the index of the wave which will participate template using single = IOScheduler::Single; } // namespace fragment_scheduler //! @class fragment //! @brief rocWMMA fragment class. This is the primary object used in block-wise decomposition of the matrix multiply-accumulate (mma) //! problem space. In general, fragment data is associated with a matrix context (matrix_a, matrix_b or accumulator), a block size (BlockM/N/K), //! a datatype (e.g. single-precision float, etc.) and an in-memory 2D layout (e.g. row_major or col_major). These fragment properties are used //! to define how data is handled and stored locally, and to drive API implementations for loading / storing, mma and transforms. Fragment abstractions are //! designed to promote a simple wavefront programming model, which can accelerate development time. Internal thread-level details are handled by rocWMMA //! which frees the user to focus on wavefront block-wise decomposition. Written purely in device code, the programmer can use this object in their own //! device kernels. //! //! @tparam MatrixT fragment context //! @tparam FragM/N/K fragment dimensions //! @tparam DataT datatype //! @tparam DataLayoutT in-memory layout as col_major or row_major //! @tparam Scheduler wave-wise scheduler //! @note Fragments are stored in packed registers, however vector elements have no guaranteed order or locality. template class __align__(4) fragment { public: //! Input / output traits specific to AMDGCN architecture using IOTraits = typename IOConfig:: IOTraits; struct Traits { private: //! The packed type for element data using PackedElementT = typename PackTraits::PackedT; //! The unpacked type for element data using UnpackedElementT = typename PackTraits::UnpackedT; //! WaveCount sizing factor for cooperative fragments static constexpr uint32_t WaveCount = scheduler_traits::WaveCount; //! Assert the fragment occupies at least one packed register static_assert(IOTraits::PackedVRegCount >= 1, "Fragments must occupy at least one packed register"); //! Assert the fragment is equally splittable among the wave count static_assert(WaveCount >= 1, "WaveCount must be at least 1 for a valid fragment"); static_assert(IOTraits::PackedVRegCount >= WaveCount, "Packed register count must be equal to or greater than WaveCount"); static_assert(IOTraits::PackedVRegCount % WaveCount == 0, "Packed register count must be divisible by WaveCount"); public: constexpr static uint32_t Size = IOTraits::UnpackedSize / WaveCount; //! Unpacked data access view using AccessT = VecT; //! Packed data storage view using StorageT = VecT; static_assert(IOTraits::UnpackedSize % IOTraits::PackedSize == 0, "Unable to pack fragment elements"); }; ROCWMMA_DEVICE fragment() = default; ROCWMMA_DEVICE fragment(const fragment& other); ROCWMMA_DEVICE fragment& operator=(const fragment& other); //! @param index Element index //! @returns Mutable unpacked element accessor at given index ROCWMMA_DEVICE inline DataT& operator[](uint32_t index); //! @param index Element index //! @returns Immutable unpacked element accessor at given index ROCWMMA_DEVICE inline DataT const& operator[](uint32_t index) const; //! @returns Mutable packed storage vector accessor ROCWMMA_DEVICE inline typename Traits::StorageT& operator*(); //! @returns Immutable packed storage vector accessor ROCWMMA_DEVICE inline typename Traits::StorageT const& operator*() const; //! @returns The geometric height of fragment ROCWMMA_DEVICE constexpr static inline uint32_t height(); //! @returns The geometric width of fragment ROCWMMA_DEVICE constexpr static inline uint32_t width(); //! @returns The leading block dimension (non-K) ROCWMMA_DEVICE constexpr static inline uint32_t blockDim(); //! @returns The k dimension ROCWMMA_DEVICE constexpr static inline uint32_t kDim(); //! @returns The size of the unpacked elements vector ROCWMMA_DEVICE constexpr static inline uint32_t size(); //! Internal data storage views. Compatibility with nvcuda::wmma union { typename Traits::StorageT mStorage; // Packed typename Traits::AccessT mAccess; // Unpacked typename Traits::AccessT::Native_vec_ x; // Nuanced access static_assert(sizeof(typename Traits::AccessT) == sizeof(typename Traits::StorageT), "Storage type and access type should be views into the same raw data"); }; // For compatibility constexpr static uint32_t num_elements = Traits::Size; using element_type = DataT; }; //! Fills the entire fragment with the desired value. //! @param frag Fragment of type MatrixT with its associated block sizes, data type and layout //! @param value Fill value of type DataT //! @tparam FragT Opaque fragment type //! @tparam DataT Datatype template ROCWMMA_DEVICE void fill_fragment(FragT& frag, DataT value); //! Loads the entire fragment from the data pointer according to its matrix and data layout contexts. Data pointer may point to either local or global memory. //! @param frag Fragment of type MatrixT with its associated block sizes, data type and layout //! @param data Data pointer to global or local memory //! @param ldm Leading dimension size //! @tparam FragT Opaque fragment type //! @tparam DataT Datatype template ROCWMMA_DEVICE void load_matrix_sync(FragT& frag, const DataT* data, uint32_t ldm); //! Loads the entire fragment from the data pointer according to its matrix layout and data layout contexts. //! Data pointer may point to either local or global memory. This overload provides manual selection of data layout of the incoming memory pointer, //! which will be transformed to conform to the data layout of the fragment. //! @param frag Fragment of type MatrixT with its associated block sizes, data type and layout //! @param data Data pointer to global/local memory //! @param ldm Leading dimension size //! @param layout Data layout //! @tparam FragT Opaque fragment type //! @tparam DataT Datatype template ROCWMMA_DEVICE void load_matrix_sync(FragT& frag, const DataT* data, uint32_t ldm, layout_t layout); //! Stores the entire fragment to the data pointer according to its matrix and data layouts. Data pointer may point to either local or global memory. //! @param frag Fragment of type MatrixT with its associated block sizes, data type and layout //! @param data Data pointer to global/local memory //! @param ldm Leading dimension size //! @tparam FragT Opaque fragment type //! @tparam DataT Datatype template ROCWMMA_DEVICE void store_matrix_sync(DataT* data, FragT const& frag, uint32_t ldm); //! Stores the entire fragment to the data pointer according to its matrix layout and data layout contexts. //! Data pointer may point to either local or global memory. This overload provides manual selection of data layout of the outgoing memory pointer, //! which the data layout of the fragment will be transformed to. //! @param frag Fragment of type MatrixT with its associated block sizes, data type and layout //! @param data Data pointer to global/local memory //! @param ldm Leading dimension size //! @param layout Data layout //! @tparam FragT Opaque fragment type //! @tparam DataT Datatype template ROCWMMA_DEVICE void store_matrix_sync(DataT* data, FragT const& frag, uint32_t ldm, layout_t layout); //! Performs the Multiply-Accumulate operation on the fragments A, B, C and D (D = A * B + C) //! @param d Accumulator output D //! @param a Input fragment A //! @param b Input fragment B //! @param c Input accumulator fragment C //! @tparam FragA Opaque fragment type for matrix A data //! @tparam FragB Opaque fragment type for matrix A data //! @tparam FragAccumIn Opaque fragment type for input accumulation data //! @tparam FragAccumOut Opaque fragment type for output accumulation data //! @note Frag c = d is valid template ROCWMMA_DEVICE void mma_sync(FragAccumOut& d, FragA const& a, FragB const& b, FragAccumIn& c); //! Synchronization point for all wavefronts in a workgroup. Guarantees pending reads / writes to LDS are flushed. ROCWMMA_DEVICE ROCWMMA_INLINE void synchronize_workgroup(); /** @}*/ } // namespace rocwmma #include "rocwmma_impl.hpp" #endif // ROCWMMA_API_HPP