/*
 * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 *
 * 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.
 */

#pragma once
#include "cute/atom/mma_atom.hpp"
#include <cuda_runtime.h>
#include <cute/arch/mma.hpp>
#include <cute/config.hpp>
#include <cute/int_tuple.hpp>
#include <cute/layout.hpp>
#include <cutlass/arch/mma.h>
#include <cutlass/cutlass.h>

#define CUTLASS_HOST_TRACE(x)                                                                                          \
    {                                                                                                                  \
        std::cout << __FILE__ << ":" << __LINE__ << "  " << x << std::endl;                                            \
    }

// Config
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890))
#define CUTE_ARCH_MMA_F32_SM89_ENABLED
#endif

namespace cute
{

// MMA 16x8x32 TN
struct SM89_16x8x32_F32F8F8F32_TN
{
    using DRegisters = float[4];
    using ARegisters = uint32_t[4];
    using BRegisters = uint32_t[2];
    using CRegisters = float[4];

    CUTE_HOST_DEVICE static void fma(float& d0, float& d1, float& d2, float& d3, uint32_t const& a0, uint32_t const& a1,
        uint32_t const& a2, uint32_t const& a3, uint32_t const& b0, uint32_t const& b1, float const& c0,
        float const& c1, float const& c2, float const& c3)
    {
#if defined(CUTE_ARCH_MMA_F32_SM89_ENABLED)
        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"(d0), "=f"(d1), "=f"(d2), "=f"(d3)
            : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "f"(c0), "f"(c1), "f"(c2), "f"(c3));
#else
        CUTE_INVALID_CONTROL_PATH(
            "Attempting to use SM89_16x8x32_F32F8F8F32_TN without "
            "CUTE_ARCH_MMA_F32_SM89_ENABLED");
#endif
    }
};

template <>
struct MMA_Traits<SM89_16x8x32_F32F8F8F32_TN>
{
    using ValTypeD = float;
    using ValTypeA = float_e4m3_t;
    using ValTypeB = float_e4m3_t;
    using ValTypeC = float;

    using Shape_MNK = Shape<_16, _8, _32>;
    using ThrID = Layout<_32>;
    using ALayout = Layout<Shape<Shape<_4, _8>, Shape<_4, _2, _2>>, Stride<Stride<_64, _1>, Stride<_16, _8, _256>>>;
    using BLayout = Layout<Shape<Shape<_4, _8>, Shape<_4, _2>>, Stride<Stride<_32, _1>, Stride<_8, _128>>>;
    using CLayout = SM80_16x8_Row;
};

} // namespace cute

using namespace cute;
using namespace cutlass;
using namespace cutlass::gemm;

namespace ada_blockwise_gemm
{

template <typename Element, typename Arch>
struct DefaultGemm_TensorOp_MMA;

template <>
struct DefaultGemm_TensorOp_MMA<cute::bfloat16_t, cutlass::arch::Sm80>
{
    using ArchTag = cutlass::arch::Sm80;
    using MMA_Atom_Arch = cute::MMA_Atom<cute::SM80_16x8x16_F32BF16BF16F32_TN>;
    using ThreadLayoutMNK = cute::Layout<cute::Shape<cute::_2, cute::_2, cute::_1>>;
    using ValLayoutMNK = cute::Tile<cute::_32, cute::_32, cute::_16>;
    using TiledMma = cute::TiledMMA<MMA_Atom_Arch, ThreadLayoutMNK, ValLayoutMNK>;
};

template <>
struct DefaultGemm_TensorOp_MMA<cute::float_e4m3_t, cutlass::arch::Sm89>
{
    using ArchTag = cutlass::arch::Sm89;
    using MMA_Atom_Arch = cute::MMA_Atom<cute::SM89_16x8x32_F32F8F8F32_TN>;
    using ThreadLayoutMNK = cute::Layout<cute::Shape<cute::_2, cute::_2, cute::_1>>;
    using ValLayoutMNK = cute::Tile<cute::_32, cute::_32, cute::_32>;
    using TiledMma = cute::TiledMMA<MMA_Atom_Arch, ThreadLayoutMNK, ValLayoutMNK>;
};

template <typename ElementType, typename OutElementType, typename AccumElementType, typename BlockScaleElementType,
    int Stages_, int TileM_, int TileN_, int TileK_>
struct AdaBlockwiseGemmTraits
{
    using ElementInput = ElementType;
    using ElementOutput = OutElementType;
    using ElementAccumulator = float;
    using ElementBlockScale = float;

    using index_t = uint32_t;
    static_assert(TileM_ % 16 == 0);
    static_assert(TileN_ % 32 == 0);
    static_assert(TileK_ % 32 == 0);
    static constexpr int Stages = Stages_;
    static constexpr int kTileM = TileM_;
    static constexpr int kTileN = TileN_;
    static constexpr int kTileK = TileK_;

    using TileShape = Shape<Int<kTileM>, Int<kTileN>, Int<kTileK>>;
    static constexpr int kWarpsCount = 4;
    static constexpr int kThreadCount = kWarpsCount * 32;

    static constexpr int ScaleGranularityM = 1;
    static constexpr int ScaleGranularityN = 128;
    static constexpr int ScaleGranularityK = 128;

    static constexpr int ScaleMsPerTile = (kTileM + ScaleGranularityM - 1) / ScaleGranularityM;
    static constexpr int ScaleNsPerTile = (kTileN + ScaleGranularityN - 1) / ScaleGranularityN;
    static constexpr int ScaleKsPerTile = (kTileK + ScaleGranularityK - 1) / ScaleGranularityK;

    using ScaleGranularity = Shape<Int<ScaleGranularityM>, Int<ScaleGranularityN>, Int<ScaleGranularityK>>;
    using ScalePerTileShape = Shape<Int<ScaleMsPerTile>, Int<ScaleNsPerTile>, Int<ScaleKsPerTile>>;

    // ====== mma ======
    static constexpr int kMmaPermM = 32;
    static constexpr int kMmaPermN = 32;
    static constexpr int kMmaPermK = 32;
    constexpr static int NUM_GROUP_M = kTileM / kMmaPermM;
    constexpr static int NUM_GROUP_N = kTileN / kMmaPermN;
    constexpr static int NUM_GROUP_K = kTileK / kMmaPermK;
    using MMA_Atom = MMA_Atom<SM89_16x8x32_F32F8F8F32_TN>;
    using AtomLayoutMNK = Layout<Shape<_2, _2, _1>>;
    using PermutationMNK = Tile<Int<kMmaPermM>, Int<kMmaPermN>, Int<kMmaPermK>>;
    using TiledMma = TiledMMA<MMA_Atom, AtomLayoutMNK, PermutationMNK>;

    // ====== load gmem -> smem ======
    using GmemTiledCopyLoad = decltype(make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, ElementInput>{},
        Layout<Shape<_16, _8>, Stride<_8, _1>>{}, Layout<Shape<_1, _16>>{}));

    using GmemTiledCopyA = GmemTiledCopyLoad;
    using GmemTiledCopyB = GmemTiledCopyLoad;

    // ====== load smem -> rf ======
    using SmemAtomLayoutLoad = decltype(composition(Swizzle<3, 4, 3>{}, Layout<Shape<_16, _128>, Stride<_128, _1>>{}));
    using SmemLayoutA = decltype(tile_to_shape(SmemAtomLayoutLoad{}, Shape<Int<kTileM>, Int<kTileK>, Int<Stages>>{}));
    using SmemLayoutB = decltype(tile_to_shape(SmemAtomLayoutLoad{}, Shape<Int<kTileN>, Int<kTileK>, Int<Stages>>{}));

    using SmemCopyAtomLoad = Copy_Atom<SM75_U32x4_LDSM_N, ElementInput>;
    using SmemCopyAtomA = SmemCopyAtomLoad;
    using SmemCopyAtomB = SmemCopyAtomLoad;

    // ====== store rf -> smem ======
    using SmemAtomLayoutStore = decltype(composition(
        Swizzle<3, 3, 3>{}, Layout<Shape<_8, Shape<_8, _8>>, Stride<_8, Stride<_1, _64>>>{})); //  8x64

    using SmemLayoutO = decltype(tile_to_shape(SmemAtomLayoutStore{}, Shape<Int<kTileM>, Int<kTileN>>{}));

    using SmemCopyAtomR2S = Copy_Atom<AutoVectorizingCopy, ElementOutput>;

    // ====== store smem -> gmem ======
    using SmemCopyAtomS2R = Copy_Atom<UniversalCopy<uint128_t>, ElementOutput>;
    using GmemCopyAtomR2G = SmemCopyAtomS2R;

    using TiledCopyS2G = decltype(make_tiled_copy(
        SmemCopyAtomS2R{}, Layout<Shape<_16, _8>, Stride<_8, _1>>{}, Layout<Shape<_1, _8>>{})); // 16x64

    // ====== load scale gmem -> smem ======
    using GmemCopyAtomScale = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
    using GmemLayoutTVSFA = Layout<Shape<Shape<Int<ScaleMsPerTile>, Int<kThreadCount / ScaleMsPerTile>>, Shape<_1, _1>>,
        Stride<Stride<_1, _0>, Stride<_1, _1>>>;
    using GmemTileShapeSFA = Shape<Int<ScaleMsPerTile>, Int<ScaleKsPerTile>>;
    using GmemTiledCopySFA = decltype(make_tiled_copy_impl(GmemCopyAtomScale{}, GmemLayoutTVSFA{}, GmemTileShapeSFA{}));

    using GmemLayoutTVSFB = Layout<Shape<Shape<_32, _4>, Shape<_1, _1>>, Stride<Stride<_0, _0>, Stride<_1, _1>>>;
    using GmemTileShapeSFB = Shape<Int<ScaleNsPerTile>, Int<ScaleKsPerTile>>;
    using GmemTiledCopySFB = decltype(make_tiled_copy_impl(GmemCopyAtomScale{}, GmemLayoutTVSFB{}, GmemTileShapeSFB{}));

    // ====== load scale smem -> rf ======
    using SmemCopyAtomScale = Copy_Atom<UniversalCopy<ElementBlockScale>, ElementBlockScale>;
    using SmemLayoutTVSFA
        = Layout<Shape<Shape<_4, _8, _2, _2>, Shape<_2>>, Stride<Stride<_0, _1, _16, _0>, Stride<_8, _0>>>;
    using SmemTileShapeSFA = Shape<Int<kMmaPermM>, _1>;
    using SmemTiledCopySFA = decltype(make_tiled_copy_impl(SmemCopyAtomScale{}, SmemLayoutTVSFA{}, SmemTileShapeSFA{}));

    using SmemLayoutSFA = decltype(tile_to_shape(make_layout(SmemTileShapeSFA{}),
        make_shape(
            shape<0>(ScalePerTileShape{}), shape<2>(ScalePerTileShape{}), Int<Stages>{}))); // BLK_M, BLK_K, Stages

    using SmemLayoutTVSFB
        = Layout<Shape<Shape<_4, _8, _2, _2>, Shape<_1>>, Stride<Stride<_0, _0, _0, _0>, Stride<_0, _0>>>;
    using SmemTileShapeSFB = Shape<_1, _1>;
    using SmemTiledCopySFB = decltype(make_tiled_copy_impl(SmemCopyAtomScale{}, SmemLayoutTVSFB{}, SmemTileShapeSFB{}));

    using SmemLayoutSFB = decltype(tile_to_shape(make_layout(SmemTileShapeSFB{}),
        make_shape(
            shape<1>(ScalePerTileShape{}), shape<2>(ScalePerTileShape{}), Int<Stages>{}))); // BLK_N, BLK_K, Stages

    // we need at least 2 stages..
    static_assert(Stages >= 2);

    struct SharedStorageLoad : aligned_struct<128>
    {
        array_aligned<ElementInput, cosize_v<SmemLayoutA>> smem_a;
        array_aligned<ElementInput, cosize_v<SmemLayoutB>> smem_b;
        array_aligned<float, cosize_v<SmemLayoutSFA>> smem_sfa;
        array_aligned<float, cosize_v<SmemLayoutSFB>> smem_sfb;
    };

    struct SharedStorageStore : aligned_struct<128>
    {
        array_aligned<ElementOutput, cosize_v<SmemLayoutO>> smem_o;
    };

    union SharedStorage
    {
        SharedStorageLoad load;
        SharedStorageStore store;
    };

    static constexpr int kSmemSize = static_cast<int>(sizeof(SharedStorage));
};

} // namespace ada_blockwise_gemm
