// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//
// CuTe tile configurations for GEMM with gather/scatter.
// Each specialization provides TiledMMA, smem layout/copy atoms
// for a specific (ElementType, TileTag) combination.
//
// All configurations use 2×2 warp layout (128 threads) following
// CUTLASS's DefaultGemmConfigurationToCutlass3Types for SM80 FP16/BF16.
// Larger tiles (128×64, 64×128, 128×128) use more MMA iterations
// per thread rather than more warps.
//
// The gmem → smem copy is done manually by the kernel (element-wise
// with gather indices), so no GmemTiledCopy types are defined here.
// The SmemLayoutAtom and SmemCopyAtom are used for smem → register
// copies via LDSM instructions.

#pragma once

#include "cute/atom/copy_atom.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/layout.hpp"
#include "cute/tensor.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/numeric_types.h"
#include "gemm_mma_tiles.h"

namespace warpconvnet {
namespace cute_gemm {

// Primary template — must be specialized per (Element, Tile) pair
template <class ElementInput, class TileTag>
struct CuteTileConfig;

// SmemLayoutAtom for operand A: K-contiguous (row-major) with Swizzle<2,3,3>.
// 8 rows × 32 columns, K is the contiguous (stride-1) dimension.
// From CUTLASS DefaultGemm_TensorOpSm80_OperandA<half_t, RowMajor>.
using SmemLayoutAtomA_FP16 = decltype(cute::composition(
    cute::Swizzle<2, 3, 3>{},
    cute::Layout<cute::Shape<cute::_8, cute::_32>, cute::Stride<cute::_32, cute::_1>>{}));

// SmemLayoutAtom for operand B: N-contiguous (column-major) with Swizzle<3,3,3>.
// 64 rows × 8 columns, N is the contiguous (stride-1) dimension.
// From CUTLASS DefaultGemm_TensorOpSm80_OperandB<half_t, RowMajor> which maps to
// OperandA<half_t, ColumnMajor>. Used with SM75_U16x8_LDSM_T (transposing LDSM)
// that transposes N-contiguous smem data to K-contiguous registers for the MMA.
using SmemLayoutAtomB_FP16 = decltype(cute::composition(
    cute::Swizzle<3, 3, 3>{},
    cute::Layout<cute::Shape<cute::_64, cute::_8>, cute::Stride<cute::_1, cute::_64>>{}));

// ============================================================================
// Macro for FP16 tile configs — reduces boilerplate with cute:: qualification
// ============================================================================

#define DEFINE_CUTE_TILE_CONFIG_FP16(TileTag, M_DIM, N_DIM, K_DIM)              \
  template <>                                                                    \
  struct CuteTileConfig<cutlass::half_t, gemm::TileTag> {                        \
    using ElementInput = cutlass::half_t;                                        \
    using TileShape = cute::Shape<cute::Int<M_DIM>, cute::Int<N_DIM>,            \
                                  cute::Int<K_DIM>>;                             \
    using TiledMma =                                                             \
        cute::TiledMMA<cute::MMA_Atom<cute::SM80_16x8x16_F32F16F16F32_TN>,      \
                       cute::Layout<cute::Shape<cute::_2, cute::_2, cute::_1>>,  \
                       cute::Tile<cute::_32, cute::_32, cute::_16>>;             \
    using SmemLayoutAtomA = SmemLayoutAtomA_FP16;                                \
    using SmemLayoutAtomB = SmemLayoutAtomB_FP16;                                \
    using SmemCopyAtomA = cute::Copy_Atom<cute::SM75_U32x4_LDSM_N, ElementInput>;\
    using SmemCopyAtomB = cute::Copy_Atom<cute::SM75_U16x8_LDSM_T, ElementInput>;\
    using GmemTiledCopyA = void;                                                 \
    using GmemTiledCopyB = void;                                                 \
    static constexpr int NumStages = 2;                                          \
    static constexpr int AlignmentA = 4;                                         \
    static constexpr int AlignmentB = 4;                                         \
    static constexpr bool UseCpAsyncGatherA = false;                             \
  };

#define DEFINE_CUTE_TILE_CONFIG_BF16(TileTag, M_DIM, N_DIM, K_DIM)              \
  template <>                                                                    \
  struct CuteTileConfig<cutlass::bfloat16_t, gemm::TileTag> {                    \
    using ElementInput = cutlass::bfloat16_t;                                    \
    using TileShape = cute::Shape<cute::Int<M_DIM>, cute::Int<N_DIM>,            \
                                  cute::Int<K_DIM>>;                             \
    using TiledMma =                                                             \
        cute::TiledMMA<cute::MMA_Atom<cute::SM80_16x8x16_F32BF16BF16F32_TN>,    \
                       cute::Layout<cute::Shape<cute::_2, cute::_2, cute::_1>>,  \
                       cute::Tile<cute::_32, cute::_32, cute::_16>>;             \
    using SmemLayoutAtomA = SmemLayoutAtomA_FP16;                                \
    using SmemLayoutAtomB = SmemLayoutAtomB_FP16;                                \
    using SmemCopyAtomA = cute::Copy_Atom<cute::SM75_U32x4_LDSM_N, ElementInput>;\
    using SmemCopyAtomB = cute::Copy_Atom<cute::SM75_U16x8_LDSM_T, ElementInput>;\
    using GmemTiledCopyA = void;                                                 \
    using GmemTiledCopyB = void;                                                 \
    static constexpr int NumStages = 2;                                          \
    static constexpr int AlignmentA = 4;                                         \
    static constexpr int AlignmentB = 4;                                         \
    static constexpr bool UseCpAsyncGatherA = false;                             \
  };

// ============================================================================
// All tile specializations
// ============================================================================

// tK=32 tiles
DEFINE_CUTE_TILE_CONFIG_FP16(Tile64x64x32,    64,  64, 32)
DEFINE_CUTE_TILE_CONFIG_FP16(Tile128x64x32,  128,  64, 32)
DEFINE_CUTE_TILE_CONFIG_FP16(Tile64x128x32,   64, 128, 32)
DEFINE_CUTE_TILE_CONFIG_FP16(Tile128x128x32, 128, 128, 32)

DEFINE_CUTE_TILE_CONFIG_BF16(Tile64x64x32,    64,  64, 32)
DEFINE_CUTE_TILE_CONFIG_BF16(Tile128x64x32,  128,  64, 32)
DEFINE_CUTE_TILE_CONFIG_BF16(Tile64x128x32,   64, 128, 32)
DEFINE_CUTE_TILE_CONFIG_BF16(Tile128x128x32, 128, 128, 32)

// tK=64 tiles
DEFINE_CUTE_TILE_CONFIG_FP16(Tile64x64x64,    64,  64, 64)
DEFINE_CUTE_TILE_CONFIG_FP16(Tile128x64x64,  128,  64, 64)
DEFINE_CUTE_TILE_CONFIG_FP16(Tile64x128x64,   64, 128, 64)
DEFINE_CUTE_TILE_CONFIG_FP16(Tile128x128x64, 128, 128, 64)

DEFINE_CUTE_TILE_CONFIG_BF16(Tile64x64x64,    64,  64, 64)
DEFINE_CUTE_TILE_CONFIG_BF16(Tile128x64x64,  128,  64, 64)
DEFINE_CUTE_TILE_CONFIG_BF16(Tile64x128x64,   64, 128, 64)
DEFINE_CUTE_TILE_CONFIG_BF16(Tile128x128x64, 128, 128, 64)

// Asymmetric M/N tiles (tK=32)
DEFINE_CUTE_TILE_CONFIG_FP16(Tile256x64x32,  256,  64, 32)
DEFINE_CUTE_TILE_CONFIG_FP16(Tile64x256x32,   64, 256, 32)

DEFINE_CUTE_TILE_CONFIG_BF16(Tile256x64x32,  256,  64, 32)
DEFINE_CUTE_TILE_CONFIG_BF16(Tile64x256x32,   64, 256, 32)

#undef DEFINE_CUTE_TILE_CONFIG_FP16
#undef DEFINE_CUTE_TILE_CONFIG_BF16

/// Override wrapper: inherits all types from BaseConfig, overrides NumStages and UseCpAsyncGatherA.
template <class BaseConfig, int NumStages_, bool UseCpAsyncGatherA_>
struct CuteTileConfigOverride {
  using ElementInput    = typename BaseConfig::ElementInput;
  using TileShape       = typename BaseConfig::TileShape;
  using TiledMma        = typename BaseConfig::TiledMma;
  using SmemLayoutAtomA = typename BaseConfig::SmemLayoutAtomA;
  using SmemLayoutAtomB = typename BaseConfig::SmemLayoutAtomB;
  using SmemCopyAtomA   = typename BaseConfig::SmemCopyAtomA;
  using SmemCopyAtomB   = typename BaseConfig::SmemCopyAtomB;
  using GmemTiledCopyA  = typename BaseConfig::GmemTiledCopyA;
  using GmemTiledCopyB  = typename BaseConfig::GmemTiledCopyB;
  static constexpr int AlignmentA = BaseConfig::AlignmentA;
  static constexpr int AlignmentB = BaseConfig::AlignmentB;
  static constexpr int NumStages = NumStages_;
  static constexpr bool UseCpAsyncGatherA = UseCpAsyncGatherA_;
};

}  // namespace cute_gemm
}  // namespace warpconvnet
