/*
 * Adapted from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/mma.cuh
 * Copyright (c) 2023 by FlashInfer team.
 *
 * Modifications copyright (c) 2024 by SageAttention team.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef SD_CORE_CUDA_SAGEATTENTION_MMA_CUH
#define SD_CORE_CUDA_SAGEATTENTION_MMA_CUH
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <cassert>
#include <type_traits>

namespace sageattention {

    namespace mma {

#if (__CUDACC_VER_MAJOR__ >= 11)
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
#define MMA_F16F16F32_M16N8K16_ENABLED
#define MMA_F16F16F16_M16N8K16_ENABLED
#define MMA_S8S8S32_M16N8K32_ENABLED
#define MMA_S4S4S32_M16N8K64_ENABLED
#endif
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 750))
#define MMA_F16F16F32_M16N8K8_ENABLED
#define MMA_F16F16F16_M16N8K8_ENABLED
#define LDMATRIX_M8N8X2_ENABLED
#define LDMATRIX_M8N8X4_ENABLED
#endif
#endif

#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120400)
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890))
#define MMA_F8F8F32_M16N8K16_ENABLED
#endif
#endif

#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890))
#define MMA_F8F8F16_M16N8K16_ENABLED
#endif
#endif

#if defined(__CUDA_ARCH__)
#define RUNTIME_ASSERT(x) __brkpt()
#else

#define RUNTIME_ASSERT(x) assert(0 && x)
#endif

        enum class MMAMode {
            kInit          = 0U,
            kInplaceUpdate = 1U,
        };

        /*!
         * \brief Wrapper of PTX ldmatrix m8n8.x2 instruction, loads data from shared memory
         *   to fragment
         * \tparam T data type of the fragment
         * \param R pointer to the fragment
         * \param smem_ptr pointer to the shared memory
         */
        template <typename T>
        __device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) {
#ifdef LDMATRIX_M8N8X2_ENABLED
            uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
            asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n"
                         : "=r"(R[0]), "=r"(R[1])
                         : "r"(smem_int_ptr));
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
#endif
        }

        /*!
         * \brief Wrapper of PTX ldmatrix m8n8.x4 instruction, loads data from shared memory
         *   to fragment
         * \tparam T data type of the fragment
         * \param R pointer to the fragment
         * \param smem_ptr pointer to the shared memory
         */
        template <typename T>
        __device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) {
#ifdef LDMATRIX_M8N8X4_ENABLED
            uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
            asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
                         : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3])
                         : "r"(smem_int_ptr));
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
#endif
        }

        /*!
         * \brief Wrapper of PTX ldmatrix m8n8.x4 transposed instruction, loads data from
         *   shared memory to fragment and transposes the fragment
         * \tparam T data type of the fragment
         * \param R pointer to the fragment
         * \param smem_ptr pointer to the shared memory
         */
        template <typename T>
        __device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) {
#ifdef LDMATRIX_M8N8X4_ENABLED
            uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
            asm volatile("ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
                         : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3])
                         : "r"(smem_int_ptr));
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
#endif
        }

        /*!
         * \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix
         *   multiplication, accumulated in f32.
         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
         * \param C pointer to the accumulator
         * \param A pointer to the fragment of matrix A
         * \param B pointer to the fragment of matrix B
         */
        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
        __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_F16F16F32_M16N8K16_ENABLED
            // ! only support half dtype now
            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
                asm volatile(
                    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
                      "f"(C[2]), "f"(C[3]));
            } else if constexpr (mma_mode == MMAMode::kInit) {
                asm volatile(
                    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
                      "f"(0.f), "f"(0.f));
            }
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

        /*!
         * \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix
         *   multiplication, accumulated in f32.
         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
         * \param C pointer to the accumulator
         * \param A pointer to the fragment of matrix A
         * \param B pointer to the fragment of matrix B
         */
        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
        __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_F16F16F32_M16N8K16_ENABLED
            // ! only support half dtype now
            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
                asm volatile(
                    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
                      "f"(C[2]), "f"(C[3]));
                asm volatile(
                    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]),
                      "f"(C[6]), "f"(C[7]));
            } else if constexpr (mma_mode == MMAMode::kInit) {
                asm volatile(
                    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
                      "f"(0.f), "f"(0.f));
                asm volatile(
                    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f),
                      "f"(0.f), "f"(0.f));
            }
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

        /*!
         * \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix
         *   multiplication, accumulated in f16.
         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
         * \param C pointer to the accumulator
         * \param A pointer to the fragment of matrix A
         * \param B pointer to the fragment of matrix B
         */
        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
        __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_F16F16F16_M16N8K16_ENABLED
            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
                asm volatile(
                    "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
                    "{%0,  %1},"
                    "{%2,  %3,  %4,  %5},"
                    "{%6,  %7},"
                    "{%8,  %9};\n"
                    : "=r"(C[0]), "=r"(C[1])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]));
            } else if constexpr (mma_mode == MMAMode::kInit) {
                asm volatile(
                    "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
                    "{%0,  %1},"
                    "{%2,  %3,  %4,  %5},"
                    "{%6,  %7},"
                    "{%8,  %9};\n"
                    : "=r"(C[0]), "=r"(C[1])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0));
            }
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

        /*!
         * \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix
         *   multiplication, accumulated in f16.
         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
         * \param C pointer to the accumulator
         * \param A pointer to the fragment of matrix A
         * \param B pointer to the fragment of matrix B
         */
        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
        __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_F16F16F16_M16N8K16_ENABLED
            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
                asm volatile(
                    "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
                    "{%0,  %1},"
                    "{%2,  %3,  %4,  %5},"
                    "{%6,  %7},"
                    "{%8,  %9};\n"
                    : "=r"(C[0]), "=r"(C[1])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]));
                asm volatile(
                    "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
                    "{%0,  %1},"
                    "{%2,  %3,  %4,  %5},"
                    "{%6,  %7},"
                    "{%8,  %9};\n"
                    : "=r"(C[2]), "=r"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[2]), "r"(C[3]));
            } else if constexpr (mma_mode == MMAMode::kInit) {
                asm volatile(
                    "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
                    "{%0,  %1},"
                    "{%2,  %3,  %4,  %5},"
                    "{%6,  %7},"
                    "{%8,  %9};\n"
                    : "=r"(C[0]), "=r"(C[1])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0));
                asm volatile(
                    "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
                    "{%0,  %1},"
                    "{%2,  %3,  %4,  %5},"
                    "{%6,  %7},"
                    "{%8,  %9};\n"
                    : "=r"(C[2]), "=r"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0));
            }
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

        /*!
         * \brief Wrapper of the mma m16n8k32 instruction for row major and column major int8 matrix
         *   multiplication, accumulated in int32.
         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
         * \param C pointer to the accumulator
         * \param A pointer to the fragment of matrix A
         * \param B pointer to the fragment of matrix B
         */
        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
        __device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_S8S8S32_M16N8K32_ENABLED
            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
                      "r"(C[2]), "r"(C[3]));
            } else if constexpr (mma_mode == MMAMode::kInit) {
                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
                      "r"(0), "r"(0));
            }
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

        /*!
         * \brief Wrapper of the mma m16n16k32 instruction for row major and column major int8 matrix
         *   multiplication, accumulated in int32.
         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
         * \param C pointer to the accumulator
         * \param A pointer to the fragment of matrix A
         * \param B pointer to the fragment of matrix B
         */
        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
        __device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_S8S8S32_M16N8K32_ENABLED
            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
                      "r"(C[2]), "r"(C[3]));
                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]),
                      "r"(C[6]), "r"(C[7]));
            } else if constexpr (mma_mode == MMAMode::kInit) {
                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
                      "r"(0), "r"(0));
                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0),
                      "r"(0), "r"(0));
            }
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

        /*!
         * \brief Wrapper of the mma m16n8k32 instruction for row major and column major int4 matrix
         *   multiplication, accumulated in int32.
         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
         * \param C pointer to the accumulator
         * \param A pointer to the fragment of matrix A
         * \param B pointer to the fragment of matrix B
         */
        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
        __device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_S4S4S32_M16N8K64_ENABLED
            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
                asm volatile(
                    "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
                      "r"(C[2]), "r"(C[3]));
            } else if constexpr (mma_mode == MMAMode::kInit) {
                asm volatile(
                    "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
                      "r"(0), "r"(0));
            }
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

        /*!
         * \brief Wrapper of the mma m16n16k64 instruction for row major and column major int4 matrix
         *   multiplication, accumulated in int32.
         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
         * \param C pointer to the accumulator
         * \param A pointer to the fragment of matrix A
         * \param B pointer to the fragment of matrix B
         */
        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
        __device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_S4S4S32_M16N8K64_ENABLED
            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
                asm volatile(
                    "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
                      "r"(C[2]), "r"(C[3]));
                asm volatile(
                    "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]),
                      "r"(C[6]), "r"(C[7]));
            } else if constexpr (mma_mode == MMAMode::kInit) {
                asm volatile(
                    "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
                      "r"(0), "r"(0));
                asm volatile(
                    "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0),
                      "r"(0), "r"(0));
            }
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

        /*!
         * \brief Wrapper of the mma m16n8k32 instruction for row major and column major fp8 e4m3 matrix
         *   multiplication, accumulated in fp32.
         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
         * \param C pointer to the accumulator
         * \param A pointer to the fragment of matrix A
         * \param B pointer to the fragment of matrix B
         */
        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
        __device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_F8F8F32_M16N8K16_ENABLED
            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
                      "f"(C[2]), "f"(C[3]));
            } else if constexpr (mma_mode == MMAMode::kInit) {
                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
                      "f"(0.f), "f"(0.f));
            }
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

        /*!
         * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix
         *   multiplication, accumulated in fp16.
         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
         * \param C pointer to the accumulator
         * \param A pointer to the fragment of matrix A
         * \param B pointer to the fragment of matrix B
         */
        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
        __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A, uint32_t* B) {
            // uint32_t* C_uint32 = reinterpret_cast<uint32_t*>(C);
#ifdef MMA_F8F8F16_M16N8K16_ENABLED
            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
                    "{%0,  %1},"
                    "{%2,  %3,  %4,  %5},"
                    "{%6,  %7},"
                    "{%8,  %9};\n"
                    : "=r"(C_uint32[0]), "=r"(C_uint32[1])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C_uint32[0]), "r"(C_uint32[1]));

                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
                    "{%0,  %1},"
                    "{%2,  %3,  %4,  %5},"
                    "{%6,  %7},"
                    "{%8,  %9};\n"
                    : "=r"(C_uint32[2]), "=r"(C_uint32[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C_uint32[2]), "r"(C_uint32[3]));
            } else if constexpr (mma_mode == MMAMode::kInit) {
                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
                    "{%0,  %1},"
                    "{%2,  %3,  %4,  %5},"
                    "{%6,  %7},"
                    "{%8,  %9};\n"
                    : "=r"(C_uint32[0]), "=r"(C_uint32[1])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0));

                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
                    "{%0,  %1},"
                    "{%2,  %3,  %4,  %5},"
                    "{%6,  %7},"
                    "{%8,  %9};\n"
                    : "=r"(C_uint32[2]), "=r"(C_uint32[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0));
            }
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

        /*!
         * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix
         *   multiplication, accumulated in fp32.
         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
         * \param C pointer to the accumulator
         * \param A pointer to the fragment of matrix A
         * \param B pointer to the fragment of matrix B
         */
        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
        __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_F8F8F32_M16N8K16_ENABLED
            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
                      "f"(C[2]), "f"(C[3]));

                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]),
                      "f"(C[6]), "f"(C[7]));
            } else if constexpr (mma_mode == MMAMode::kInit) {
                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
                      "f"(0.f), "f"(0.f));

                asm volatile(
                    "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
                    "{%0,  %1,  %2,  %3},"
                    "{%4,  %5,  %6,  %7},"
                    "{%8,  %9},"
                    "{%10, %11, %12, %13};\n"
                    : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f),
                      "f"(0.f), "f"(0.f));
            }
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

        /*!
         * \brief Use mma instructions to compute rowsum.
         */
        __device__ __forceinline__ void rowsum_f16f16f32(float* d, uint32_t* s) {
#ifdef MMA_F16F16F32_M16N8K16_ENABLED
            asm volatile(
                "{\n"
                "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
                "{%0,  _,  %1,  _},"
                "{%2,  %3,  %4,  %5},"
                "{%6,  %7},"
                "{%8,  0.,  %9,  0.};\n"
                "}\n"
                : "=f"(d[0]), "=f"(d[1])
                : "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(1006648320),  // 1006648320 packs two 1.0f in half precision
                  "r"(1006648320), "f"(d[0]), "f"(d[1]));
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

        /*!
         * \brief Use mma instructions to compute rowsum.
         */
        __device__ __forceinline__ void rowsum_f8f8f32(float* d, uint32_t* s) {
#ifdef MMA_F8F8F32_M16N8K16_ENABLED
            asm volatile(
                "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
                "{%0,  _,  %1,  _},"
                "{%2,  %3,  %4,  %5},"
                "{%6,  %7},"
                "{%8,  0.,  %9,  0.};\n"
                : "=f"(d[0]), "=f"(d[1])
                : "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(943208504), "r"(943208504),  // 943208504 packs four 1.0f in e4m3
                  "f"(d[0]), "f"(d[1]));
#else
            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
        }

    }  // namespace mma

}  // namespace sageattention

#endif