// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

#pragma once

#include "ck_tile/core.hpp"
#include "ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma_impl.hpp"

namespace ck_tile {

// Number of groups of consecutive elements to fill in a ABKLane
enum class WGAttrNumAccessEnum
{
    Single  = 1,
    Double  = 2,
    Quad    = 4,
    Invalid = -1
};

template <WGAttrNumAccessEnum AttrNumAccess>
struct get_wgattr_num_access
{
    private:
    static constexpr index_t getAccesses()
    {
        if constexpr(AttrNumAccess == WGAttrNumAccessEnum::Single)
        {
            return 1;
        }
        else if constexpr(AttrNumAccess == WGAttrNumAccessEnum::Double)
        {
            return 2;
        }
        else if constexpr(AttrNumAccess == WGAttrNumAccessEnum::Quad)
        {
            return 4;
        }
        else
        {
            static_assert(false, "unsupported AttrNumAccess");
            return 0;
        }
    }

    public:
    static constexpr auto value = getAccesses();
};

template <typename WarpGemmAttributeMfmaImpl_,
          WGAttrNumAccessEnum AttrNumAccessA_ = WGAttrNumAccessEnum::Single,
          WGAttrNumAccessEnum AttrNumAccessB_ = AttrNumAccessA_>
struct WarpGemmAttributeMfma
{
    using Impl                            = remove_cvref_t<WarpGemmAttributeMfmaImpl_>;
    static constexpr auto AttrNumAccessA  = AttrNumAccessA_;
    static constexpr auto AttrNumAccessAV = get_wgattr_num_access<AttrNumAccessA>::value;
    static constexpr auto AttrNumAccessB  = AttrNumAccessB_;
    static constexpr auto AttrNumAccessBV = get_wgattr_num_access<AttrNumAccessB>::value;

    static constexpr bool UsePackNumAccess = AttrNumAccessA != AttrNumAccessB;

    using ADataType = typename Impl::ADataType;
    using BDataType = typename Impl::BDataType;
    using CDataType = typename Impl::CDataType;

    using AVecType = typename Impl::AVecType;
    using BVecType = typename Impl::BVecType;
    using CVecType = typename Impl::CVecType;

    static constexpr index_t kM          = Impl::kM;
    static constexpr index_t kN          = Impl::kN;
    static constexpr index_t kK          = Impl::kK;
    static constexpr index_t kKPerThread = Impl::kABKPerLane;
    static constexpr index_t kCMLane     = Impl::kCMLane;

    CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return 1; }

    static_assert(Impl::kAMBlock == 1 && Impl::kBNBlock == 1,
                  "Multi-block WarpGemmAttributeMfmaImpl is not supported");

    template <index_t kMNLane, index_t AttrNumAccessV_>
    static constexpr auto get_warp_dstr_encoding()
    {
        static_assert(kKPerThread % AttrNumAccessV_ == 0,
                      "kKPerThread must be divisible by NumAccess");
        if constexpr(AttrNumAccessV_ == 1)
        {
            return tile_distribution_encoding<
                sequence<>,
                tuple<sequence<kMNLane>, sequence<Impl::kABKLane, Impl::kABKPerLane>>,
                tuple<sequence<2, 1>>,
                tuple<sequence<0, 0>>,
                sequence<2>,
                sequence<1>>{};
        }
        else
        {
            // AttrNumAccess splits the kABKPerLane
            // We can split them but still have them contiguous (packed) or have them interleaved.
            // The reason to split the dimension but still have it packed is to match load transpose
            // encoding when A and B use different AttrNumAccess (they have different types in LDS)
            // Example
            // A: 16bit, B: 8bit
            // Load transpose B: lane0 -> K=0..7 (only 1 instruction)
            // Load transpose A: lane0 -> K=0..3 first instruction, K=4..7 second instruction
            // In this way the data in register are consistent between A and B
            if constexpr(UsePackNumAccess)
            {
                return tile_distribution_encoding<
                    sequence<>,
                    tuple<sequence<kMNLane>,
                          sequence<Impl::kABKLane,
                                   AttrNumAccessV_,
                                   Impl::kABKPerLane / AttrNumAccessV_>>,
                    tuple<sequence<2, 1>>,
                    tuple<sequence<0, 0>>,
                    sequence<2, 2>,
                    sequence<1, 2>>{};
            }
            else
            {
                return tile_distribution_encoding<
                    sequence<>,
                    tuple<sequence<kMNLane>,
                          sequence<AttrNumAccessV_,
                                   Impl::kABKLane,
                                   Impl::kABKPerLane / AttrNumAccessV_>>,
                    tuple<sequence<2, 1>>,
                    tuple<sequence<1, 0>>,
                    sequence<2, 2>,
                    sequence<0, 2>>{};
            }
        }
    }
    using AWarpDstrEncoding = decltype(get_warp_dstr_encoding<Impl::kAMLane, AttrNumAccessAV>());
    using BWarpDstrEncoding = decltype(get_warp_dstr_encoding<Impl::kBNLane, AttrNumAccessBV>());

    using CWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kCM0PerLane, Impl::kCMLane, Impl::kCM1PerLane>,
              sequence<Impl::kCNLane>>,
        tuple<sequence<1, 2>>,
        tuple<sequence<1, 0>>,
        sequence<1, 1>,
        sequence<0, 2>>;

    // c_vec += a_vec * b_vec
    template <bool post_nop_ = false>
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const BVecType& b_vec,
                                   bool_constant<post_nop_> = {}) const
    {
        Impl{}(c_vec, a_vec, b_vec, bool_constant<post_nop_>{});
    }

    // c_vec += a_vec * b_vec
    template <index_t opselA, index_t opselB, bool post_nop_ = false>
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const int32_t& a_scale,
                                   const BVecType& b_vec,
                                   const int32_t& b_scale,
                                   bool_constant<post_nop_> = {}) const
    {
        Impl{}.template operator()<opselA, opselB>(
            c_vec, a_vec, a_scale, b_vec, b_scale, bool_constant<post_nop_>{});
    }

    // c_vec = a_vec * b_vec
    CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const
    {
        return Impl{}(a_vec, b_vec);
    }

    // c_vec = a_vec * b_vec
    template <index_t opselA, index_t opselB>
    CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec,
                                       const int32_t& a_scale,
                                       const BVecType& b_vec,
                                       const int32_t& b_scale) const
    {
        return Impl{}.template operator()<opselA, opselB>(a_vec, a_scale, b_vec, b_scale);
    }
};

template <typename WarpGemmAttributeMfmaImpl_,
          index_t kKIter,
          WGAttrNumAccessEnum AttrNumAccessA_ = WGAttrNumAccessEnum::Single,
          WGAttrNumAccessEnum AttrNumAccessB_ = AttrNumAccessA_>
struct WarpGemmAttributeMfmaIterateK
{
    static_assert(kKIter > 0, "wrong!");

    using Impl                            = remove_cvref_t<WarpGemmAttributeMfmaImpl_>;
    static constexpr auto AttrNumAccessA  = AttrNumAccessA_;
    static constexpr auto AttrNumAccessAV = get_wgattr_num_access<AttrNumAccessA>::value;
    static constexpr auto AttrNumAccessB  = AttrNumAccessB_;
    static constexpr auto AttrNumAccessBV = get_wgattr_num_access<AttrNumAccessB>::value;

    static constexpr bool UsePackNumAccess = AttrNumAccessA != AttrNumAccessB;

    using ADataType = typename Impl::ADataType;
    using BDataType = typename Impl::BDataType;
    using CDataType = typename Impl::CDataType;

    using AVecType =
        ext_vector_t<ADataType, vector_traits<typename Impl::AVecType>::vector_size * kKIter>;
    using BVecType =
        ext_vector_t<BDataType, vector_traits<typename Impl::BVecType>::vector_size * kKIter>;
    using CVecType = typename Impl::CVecType;

    static constexpr index_t kM          = Impl::kM;
    static constexpr index_t kN          = Impl::kN;
    static constexpr index_t kK          = Impl::kK * kKIter;
    static constexpr index_t kKPerThread = Impl::kABKPerLane * kKIter;
    static constexpr index_t kCMLane     = Impl::kCMLane;

    CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return kKIter; }

    static_assert(Impl::kAMBlock == 1 || Impl::kBNBlock == 1,
                  "Multi-block on both M & N directions is not supported");

    template <index_t kMNLane, index_t kMNBlock, index_t kNMBlock, index_t AttrNumAccessV_>
    CK_TILE_DEVICE static constexpr auto get_warp_dstr_encoding()
    {
        if constexpr(kMNBlock == 1 && kNMBlock == 1)
        {
            static_assert(kKPerThread % AttrNumAccessV_ == 0,
                          "kKPerThread must be divisible by NumAccess");
            if constexpr(AttrNumAccessV_ == 1)
            {
                return tile_distribution_encoding<
                    sequence<>,
                    tuple<sequence<kMNLane>, sequence<Impl::kABKLane, Impl::kABKPerLane * kKIter>>,
                    tuple<sequence<2, 1>>,
                    tuple<sequence<0, 0>>,
                    sequence<2>,
                    sequence<1>>{};
            }
            else
            {
                if constexpr(UsePackNumAccess)
                {
                    return tile_distribution_encoding<
                        sequence<>,
                        tuple<sequence<kMNLane>,
                              sequence<Impl::kABKLane,
                                       AttrNumAccessV_,
                                       Impl::kABKPerLane * kKIter / AttrNumAccessV_>>,
                        tuple<sequence<2, 1>>,
                        tuple<sequence<0, 0>>,
                        sequence<2, 2>,
                        sequence<1, 2>>{};
                }
                else
                {
                    return tile_distribution_encoding<
                        sequence<>,
                        tuple<sequence<kMNLane>,
                              sequence<AttrNumAccessV_,
                                       Impl::kABKLane,
                                       Impl::kABKPerLane * kKIter / AttrNumAccessV_>>,
                        tuple<sequence<2, 1>>,
                        tuple<sequence<1, 0>>,
                        sequence<2, 2>,
                        sequence<0, 2>>{};
                }
            }
        }
        else if constexpr(kMNBlock == 1 && 1 < kNMBlock)
        {
            static_assert(AttrNumAccessV_ == 1,
                          "Multiple access is not supported when using multi-block");
            // each M/N blocks share the same data
            return tile_distribution_encoding<
                sequence<kNMBlock>,
                tuple<sequence<kMNLane>, sequence<Impl::kABKLane, Impl::kABKPerLane * kKIter>>,
                tuple<sequence<0, 2, 1>>,
                tuple<sequence<0, 0, 0>>,
                sequence<2>,
                sequence<1>>{};
        }
        else if constexpr(1 < kMNBlock && kNMBlock == 1)
        {
            static_assert(AttrNumAccessV_ == 1,
                          "Multiple access is not supported when using multi-block");
            // single block to multi-block thread mapping
            return tile_distribution_encoding<
                sequence<>,
                tuple<sequence<kMNBlock, kMNLane>,
                      sequence<Impl::kABKLane, Impl::kABKPerLane * kKIter>>,
                tuple<sequence<1, 2, 1>>,
                tuple<sequence<0, 0, 1>>,
                sequence<2>,
                sequence<1>>{};
        }
    }

    CK_TILE_DEVICE static constexpr auto get_cwarp_dstr_encoding()
    {
        if constexpr(Impl::kAMBlock == 1 && Impl::kBNBlock == 1)
        {
            return tile_distribution_encoding<
                sequence<>,
                tuple<sequence<Impl::kCM0PerLane, Impl::kCMLane, Impl::kCM1PerLane>,
                      sequence<Impl::kCNLane>>,
                tuple<sequence<1, 2>>,
                tuple<sequence<1, 0>>,
                sequence<1, 1>,
                sequence<0, 2>>{};
        }
        else if constexpr(Impl::kAMBlock == 1 && 1 < Impl::kBNBlock)
        {
            return tile_distribution_encoding<
                sequence<>,
                tuple<sequence<Impl::kCM0PerLane, Impl::kCMLane, Impl::kCM1PerLane>,
                      sequence<Impl::kBNBlock * Impl::kCNLane>>,
                tuple<sequence<1, 2>>,
                tuple<sequence<1, 0>>,
                sequence<1, 1>,
                sequence<0, 2>>{};
        }
        else if constexpr(1 < Impl::kAMBlock && Impl::kBNBlock == 1)
        {
            return tile_distribution_encoding<
                sequence<>,
                tuple<
                    sequence<Impl::kCM0PerLane, Impl::kAMBlock * Impl::kCMLane, Impl::kCM1PerLane>,
                    sequence<Impl::kCNLane>>,
                tuple<sequence<1, 2>>,
                tuple<sequence<1, 0>>,
                sequence<1, 1>,
                sequence<0, 2>>{};
        }
    }

    using AWarpDstrEncoding = decltype(get_warp_dstr_encoding<Impl::kAMLane,
                                                              Impl::kAMBlock,
                                                              Impl::kBNBlock,
                                                              AttrNumAccessAV>());
    using BWarpDstrEncoding = decltype(get_warp_dstr_encoding<Impl::kBNLane,
                                                              Impl::kBNBlock,
                                                              Impl::kAMBlock,
                                                              AttrNumAccessBV>());
    using CWarpDstrEncoding = decltype(get_cwarp_dstr_encoding());

    // c_vec += a_vec * b_vec
    template <bool post_nop_ = false>
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const BVecType& b_vec,
                                   bool_constant<post_nop_> = {}) const
    {
        static_for<0, kKIter, 1>{}([&](auto iKIter) { operator()(c_vec, a_vec, b_vec, iKIter); });
    }

    template <index_t iKIter, bool post_nop_ = false>
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const BVecType& b_vec,
                                   number<iKIter>,
                                   bool_constant<post_nop_> = {}) const
    {
        using buf_a = thread_buffer<typename Impl::AVecType, kKIter>;
        using buf_b = thread_buffer<typename Impl::BVecType, kKIter>;

        static_assert(iKIter < kKIter);

        Impl{}(c_vec,
               reinterpret_cast<const buf_a&>(a_vec)
                   .template get_as<typename Impl::AVecType>()[iKIter],
               reinterpret_cast<const buf_b&>(b_vec)
                   .template get_as<typename Impl::BVecType>()[iKIter],
               bool_constant<post_nop_>{});
    }

    // c_vec = a_vec * b_vec
    CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const
    {
        constexpr auto I0 = number<0>{};
        using buf_a       = thread_buffer<typename Impl::AVecType, kKIter>;
        using buf_b       = thread_buffer<typename Impl::BVecType, kKIter>;

        // c = a * b
        auto c_vec = Impl{}(
            reinterpret_cast<const buf_a&>(a_vec).template get_as<typename Impl::AVecType>()[I0],
            reinterpret_cast<const buf_b&>(b_vec).template get_as<typename Impl::BVecType>()[I0]);

        // c += a * b
        static_for<1, kKIter, 1>{}([&](auto iKIter) { operator()(c_vec, a_vec, b_vec, iKIter); });

        return c_vec;
    }
};

template <typename WarpGemmAttributeMfmaImpl_,
          WGAttrNumAccessEnum AttrNumAccess_ = WGAttrNumAccessEnum::Single>
struct WarpGemmAttributeMfmaTransposedCDistribution
{
    using Impl                           = remove_cvref_t<WarpGemmAttributeMfmaImpl_>;
    static constexpr auto AttrNumAccess  = AttrNumAccess_;
    static constexpr auto AttrNumAccessV = static_cast<index_t>(AttrNumAccess);

    using ADataType = typename Impl::BDataType;
    using BDataType = typename Impl::ADataType;
    using CDataType = typename Impl::CDataType;

    using AVecType = typename Impl::BVecType;
    using BVecType = typename Impl::AVecType;
    using CVecType = typename Impl::CVecType;

    static constexpr index_t kM          = Impl::kN;
    static constexpr index_t kN          = Impl::kM;
    static constexpr index_t kK          = Impl::kK;
    static constexpr index_t kKPerThread = Impl::kABKPerLane;
    static constexpr index_t kCMLane     = Impl::kCMLane;

    CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return 1; }

    static_assert(Impl::kAMBlock == 1 && Impl::kBNBlock == 1,
                  "Multi-block WarpGemmAttributeMfmaImpl is not supported");

    using AWarpDstrEncoding =
        typename WarpGemmAttributeMfma<Impl, AttrNumAccess>::BWarpDstrEncoding;
    using BWarpDstrEncoding =
        typename WarpGemmAttributeMfma<Impl, AttrNumAccess>::AWarpDstrEncoding;

    using CWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kCNLane>,
              sequence<Impl::kCM0PerLane, Impl::kCMLane, Impl::kCM1PerLane>>,
        tuple<sequence<2, 1>>,
        tuple<sequence<1, 0>>,
        sequence<2, 2>,
        sequence<0, 2>>;

    // c_vec += a_vec * b_vec
    template <bool post_nop_ = false>
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const BVecType& b_vec,
                                   bool_constant<post_nop_> = {}) const
    {
        // swap A and B
        Impl{}(c_vec, b_vec, a_vec, bool_constant<post_nop_>{});
    }

    template <index_t opselA, index_t opselB, bool post_nop_ = false>
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const int32_t& a_scale,
                                   const BVecType& b_vec,
                                   const int32_t& b_scale,
                                   bool_constant<post_nop_> = {}) const
    {
        // swap A and B
        Impl{}.template operator()<opselB, opselA>(
            c_vec, b_vec, b_scale, a_vec, a_scale, bool_constant<post_nop_>{});
    }

    // c_vec = a_vec * b_vec
    CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const
    {
        // swap A and B
        return Impl{}(b_vec, a_vec);
    }
};

template <typename WarpGemmAttributeMfmaImpl_, index_t SFactor_ = 2>
struct WarpGemmAttributeMfmaTransposedCDistribution_SwizzleB
{
    using Impl = remove_cvref_t<WarpGemmAttributeMfmaImpl_>;

    using ADataType = typename Impl::BDataType;
    using BDataType = typename Impl::ADataType;
    using CDataType = typename Impl::CDataType;

    using AVecType = typename Impl::BVecType;
    using BVecType = typename Impl::AVecType;
    using CVecType = typename Impl::CVecType;

    static constexpr index_t kM          = Impl::kN;
    static constexpr index_t kN          = Impl::kM;
    static constexpr index_t kK          = Impl::kK;
    static constexpr index_t kKPerThread = Impl::kABKPerLane;
    static constexpr index_t SFactor     = SFactor_; // group how many CM1 together

    CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return 1; }

    static_assert(Impl::kAMBlock == 1 && Impl::kBNBlock == 1,
                  "Multi-block WarpGemmAttributeMfmaImpl is not supported");

    using AWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kBNLane>, sequence<Impl::kABKLane, Impl::kABKPerLane>>,
        tuple<sequence<2, 1>>,
        tuple<sequence<0, 0>>,
        sequence<2>,
        sequence<1>>;
#if 0
    using BWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kAMLane / (Impl::kABKPerLane * Impl::kABKLane * 2),
                       Impl::kABKLane,
                       2,
                       Impl::kABKPerLane>,
              sequence<Impl::kABKLane, Impl::kABKPerLane>>,
        tuple<sequence<2, 1, 1, 1, 1>>,
        tuple<sequence<0, 0, 2, 1, 3>>,
        sequence<2>,
        sequence<1>>;

    using CWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kCNLane>,
              sequence<Impl::kCM0PerLane / 2, Impl::kCMLane, Impl::kCM1PerLane * 2>>,
        tuple<sequence<2, 1>>,
        tuple<sequence<1, 0>>,
        sequence<2, 2>,
        sequence<0, 2>>;
#else
    // TODO: more test not only 32x32
    using BWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kAMLane / (Impl::kCMLane * SFactor * Impl::kCM1PerLane),
                       Impl::kCMLane,
                       SFactor,
                       Impl::kCM1PerLane>,
              sequence<Impl::kABKLane, Impl::kABKPerLane>>,
        tuple<sequence<2, 1, 1, 1, 1>>,
        tuple<sequence<0, 0, 2, 1, 3>>,
        sequence<2>,
        sequence<1>>;

    using CWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kCNLane>,
              sequence<Impl::kCM0PerLane / SFactor, Impl::kCMLane, Impl::kCM1PerLane * SFactor>>,
        tuple<sequence<2, 1>>,
        tuple<sequence<1, 0>>,
        sequence<2, 2>,
        sequence<0, 2>>;
#endif
    template <bool post_nop_ = false>
    // c_vec += a_vec * b_vec
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const BVecType& b_vec,
                                   bool_constant<post_nop_> = {}) const
    {
        // swap A and B
        Impl{}(c_vec, b_vec, a_vec, bool_constant<post_nop_>{});
    }

    template <index_t opselA, index_t opselB, bool post_nop_ = false>
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const int32_t& a_scale,
                                   const BVecType& b_vec,
                                   const int32_t& b_scale,
                                   bool_constant<post_nop_> = {}) const
    {
        // swap A and B
        Impl{}.template operator()<opselB, opselA>(
            c_vec, b_vec, b_scale, a_vec, a_scale, bool_constant<post_nop_>{});
    }

    // c_vec = a_vec * b_vec
    CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const
    {
        // swap A and B
        return Impl{}(b_vec, a_vec);
    }
};

template <typename WarpGemmAttributeMfmaImpl_,
          index_t kKIter,
          WGAttrNumAccessEnum AttrNumAccess_ = WGAttrNumAccessEnum::Single>
struct WarpGemmAttributeMfmaIterateKAndTransposedCDistribution
{
    using Impl                          = remove_cvref_t<WarpGemmAttributeMfmaImpl_>;
    static constexpr auto AttrNumAccess = AttrNumAccess_;

    // swap A and B
    using ADataType = typename Impl::BDataType;
    using BDataType = typename Impl::ADataType;
    using CDataType = typename Impl::CDataType;

    using AVecType =
        ext_vector_t<ADataType, vector_traits<typename Impl::AVecType>::vector_size * kKIter>;
    using BVecType =
        ext_vector_t<BDataType, vector_traits<typename Impl::BVecType>::vector_size * kKIter>;
    using CVecType = typename Impl::CVecType;

    static constexpr index_t kM          = Impl::kN;
    static constexpr index_t kN          = Impl::kM;
    static constexpr index_t kK          = Impl::kK * kKIter;
    static constexpr index_t kKPerThread = Impl::kABKPerLane * kKIter;

    CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return kKIter; }

    static_assert(Impl::kAMBlock == 1 || Impl::kBNBlock == 1,
                  "Multi-block on both M & N directions is not supported");

    CK_TILE_DEVICE static constexpr auto get_cwarp_dstr_encoding()
    {
        if constexpr(Impl::kAMBlock == 1 && Impl::kBNBlock == 1)
        {
            return tile_distribution_encoding<
                sequence<>,
                tuple<sequence<Impl::kCNLane>,
                      sequence<Impl::kCM0PerLane, Impl::kCMLane, Impl::kCM1PerLane>>,
                tuple<sequence<2, 1>>,
                tuple<sequence<1, 0>>,
                sequence<2, 2>,
                sequence<0, 2>>{};
        }
        else if constexpr(Impl::kAMBlock == 1 && 1 < Impl::kBNBlock)
        {
            return tile_distribution_encoding<
                sequence<>,
                tuple<sequence<Impl::kBNBlock * Impl::kCNLane>,
                      sequence<Impl::kCM0PerLane, Impl::kCMLane, Impl::kCM1PerLane>>,
                tuple<sequence<2, 1>>,
                tuple<sequence<1, 0>>,
                sequence<2, 2>,
                sequence<0, 2>>{};
        }
        else if constexpr(1 < Impl::kAMBlock && Impl::kBNBlock == 1)
        {
            return tile_distribution_encoding<
                sequence<>,
                tuple<
                    sequence<Impl::kCNLane>,
                    sequence<Impl::kCM0PerLane, Impl::kAMBlock * Impl::kCMLane, Impl::kCM1PerLane>>,
                tuple<sequence<2, 1>>,
                tuple<sequence<1, 0>>,
                sequence<2, 2>,
                sequence<0, 2>>{};
        }
    }

    using AWarpDstrEncoding =
        typename WarpGemmAttributeMfmaIterateK<Impl, kKIter, AttrNumAccess>::BWarpDstrEncoding;
    using BWarpDstrEncoding =
        typename WarpGemmAttributeMfmaIterateK<Impl, kKIter, AttrNumAccess>::AWarpDstrEncoding;
    using CWarpDstrEncoding = decltype(get_cwarp_dstr_encoding());

    // c_vec += a_vec * b_vec
    template <bool post_nop_ = false>
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const BVecType& b_vec,
                                   bool_constant<post_nop_> = {}) const
    {
        static_for<0, kKIter, 1>{}([&](auto iKIter) { operator()(c_vec, a_vec, b_vec, iKIter); });
    }

    template <index_t iKIter, bool post_nop_ = false>
    // c_vec += a_vec * b_vec
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const BVecType& b_vec,
                                   number<iKIter>,
                                   bool_constant<post_nop_> = {}) const
    {
        using buf_a = thread_buffer<typename Impl::AVecType, kKIter>;
        using buf_b = thread_buffer<typename Impl::BVecType, kKIter>;

        static_assert(iKIter < kKIter);
        // swap A and B, value and type
        Impl{}(c_vec,
               reinterpret_cast<const buf_b&>(b_vec)
                   .template get_as<typename Impl::BVecType>()[iKIter],
               reinterpret_cast<const buf_a&>(a_vec)
                   .template get_as<typename Impl::AVecType>()[iKIter],
               bool_constant<post_nop_>{});
    }

    // c_vec = a_vec * b_vec
    CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const
    {
        constexpr auto I0 = number<0>{};
        using buf_a       = thread_buffer<typename Impl::AVecType, kKIter>;
        using buf_b       = thread_buffer<typename Impl::BVecType, kKIter>;

        // swap A and B, value and type
        auto c_vec = Impl{}(
            reinterpret_cast<const buf_b&>(b_vec).template get_as<typename Impl::BVecType>()[I0],
            reinterpret_cast<const buf_a&>(a_vec).template get_as<typename Impl::AVecType>()[I0]);

        static_for<1, kKIter, 1>{}([&](auto iKIter) { operator()(c_vec, a_vec, b_vec, iKIter); });

        return c_vec;
    }
};

template <typename WarpGemmAttributeMfmaImpl_, index_t kKIter, index_t SFactor_ = 2>
struct WarpGemmAttributeMfmaIterateKAndTransposedCDistribution_SwizzleB
{
    using Impl = remove_cvref_t<WarpGemmAttributeMfmaImpl_>;

    // swap A and B
    using ADataType = typename Impl::BDataType;
    using BDataType = typename Impl::ADataType;
    using CDataType = typename Impl::CDataType;

    using AVecType =
        ext_vector_t<ADataType, vector_traits<typename Impl::AVecType>::vector_size * kKIter>;
    using BVecType =
        ext_vector_t<BDataType, vector_traits<typename Impl::BVecType>::vector_size * kKIter>;
    using CVecType = typename Impl::CVecType;

    static constexpr index_t kM          = Impl::kN;
    static constexpr index_t kN          = Impl::kM;
    static constexpr index_t kK          = Impl::kK * kKIter;
    static constexpr index_t kKPerThread = Impl::kABKPerLane * kKIter;
    static constexpr index_t SFactor     = SFactor_; // group how many CM1 together

    CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return kKIter; }

    static_assert(Impl::kAMBlock == 1 && Impl::kBNBlock == 1,
                  "Multi-block WarpGemmAttributeMfmaImpl is not supported");

    using AWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kBNLane>, sequence<Impl::kABKLane, Impl::kABKPerLane * kKIter>>,
        tuple<sequence<2, 1>>,
        tuple<sequence<0, 0>>,
        sequence<2>,
        sequence<1>>;
#if 0
    using BWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kAMLane / (Impl::kABKPerLane * Impl::kABKLane * 2),
                       Impl::kABKLane,
                       2,
                       Impl::kABKPerLane>,
              sequence<Impl::kABKLane, Impl::kABKPerLane * kKIter>>,
        tuple<sequence<2, 1, 1, 1, 1>>,
        tuple<sequence<0, 0, 2, 1, 3>>,
        sequence<2>,
        sequence<1>>;

    using CWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kCNLane>,
              sequence<Impl::kCM0PerLane / 2, Impl::kCMLane, Impl::kCM1PerLane * 2>>,
        tuple<sequence<2, 1>>,
        tuple<sequence<1, 0>>,
        sequence<2, 2>,
        sequence<0, 2>>;
#else
    // TODO: more test not only 32x32
    using BWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kAMLane / (Impl::kCMLane * SFactor * Impl::kCM1PerLane),
                       Impl::kCMLane,
                       SFactor,
                       Impl::kCM1PerLane>,
              sequence<Impl::kABKLane, Impl::kABKPerLane * kKIter>>,
        tuple<sequence<2, 1, 1, 1, 1>>,
        tuple<sequence<0, 0, 2, 1, 3>>,
        sequence<2>,
        sequence<1>>;

    using CWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kCNLane>,
              sequence<Impl::kCM0PerLane / SFactor, Impl::kCMLane, Impl::kCM1PerLane * SFactor>>,
        tuple<sequence<2, 1>>,
        tuple<sequence<1, 0>>,
        sequence<2, 2>,
        sequence<0, 2>>;
#endif
    // c_vec += a_vec * b_vec
    template <bool post_nop_ = false>
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const BVecType& b_vec,
                                   bool_constant<post_nop_> = {}) const
    {
        // swap A and B, value and type
        static_for<0, kKIter, 1>{}([&](auto iKIter) { operator()(c_vec, a_vec, b_vec, iKIter); });
    }

    template <index_t iKIter, bool post_nop_ = false>
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const BVecType& b_vec,
                                   number<iKIter>,
                                   bool_constant<post_nop_> = {}) const
    {
        using buf_a = thread_buffer<typename Impl::AVecType, kKIter>;
        using buf_b = thread_buffer<typename Impl::BVecType, kKIter>;

        static_assert(iKIter < kKIter);
        // swap A and B, value and type
        Impl{}(c_vec,
               reinterpret_cast<const buf_b&>(b_vec)
                   .template get_as<typename Impl::BVecType>()[iKIter],
               reinterpret_cast<const buf_a&>(a_vec)
                   .template get_as<typename Impl::AVecType>()[iKIter],
               bool_constant<post_nop_>{});
    }

    // c_vec = a_vec * b_vec
    CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const
    {
        using buf_a       = thread_buffer<typename Impl::AVecType, kKIter>;
        using buf_b       = thread_buffer<typename Impl::BVecType, kKIter>;
        constexpr auto I0 = number<0>{};

        // swap A and B, value and type
        auto c_vec = Impl{}(
            reinterpret_cast<const buf_b&>(b_vec).template get_as<typename Impl::BVecType>()[I0],
            reinterpret_cast<const buf_a&>(a_vec).template get_as<typename Impl::AVecType>()[I0]);

        static_for<1, kKIter, 1>{}([&](auto iKIter) { operator()(c_vec, a_vec, b_vec, iKIter); });

        return c_vec;
    }
};

template <typename WarpGemmAttributeMfmaImpl_, index_t kKIter, index_t SFactor_ = 2>
struct WarpGemmAttributeMfmaIterateK_SwizzleA
{
    using Impl = remove_cvref_t<WarpGemmAttributeMfmaImpl_>;

    using ADataType = typename Impl::ADataType;
    using BDataType = typename Impl::BDataType;
    using CDataType = typename Impl::CDataType;

    using AVecType =
        ext_vector_t<ADataType, vector_traits<typename Impl::AVecType>::vector_size * kKIter>;
    using BVecType =
        ext_vector_t<BDataType, vector_traits<typename Impl::BVecType>::vector_size * kKIter>;
    using CVecType = typename Impl::CVecType;

    static constexpr index_t kM          = Impl::kM;
    static constexpr index_t kN          = Impl::kN;
    static constexpr index_t kK          = Impl::kK * kKIter;
    static constexpr index_t kKPerThread = Impl::kABKPerLane * kKIter;
    static constexpr index_t SFactor     = SFactor_; // group how many CM1 together

    CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return kKIter; }

    static_assert(Impl::kAMBlock == 1 && Impl::kBNBlock == 1,
                  "Multi-block WarpGemmAttributeMfmaImpl is not supported");

    using AWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kAMLane / (Impl::kCMLane * SFactor * Impl::kCM1PerLane),
                       Impl::kCMLane,
                       SFactor,
                       Impl::kCM1PerLane>,
              sequence<Impl::kABKLane, Impl::kABKPerLane * kKIter>>,
        tuple<sequence<2, 1, 1, 1, 1>>,
        tuple<sequence<0, 0, 2, 1, 3>>,
        sequence<2>,
        sequence<1>>;

    using BWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kBNLane>, sequence<Impl::kABKLane, Impl::kABKPerLane * kKIter>>,
        tuple<sequence<2, 1>>,
        tuple<sequence<0, 0>>,
        sequence<2>,
        sequence<1>>;

    using CWarpDstrEncoding = tile_distribution_encoding<
        sequence<>,
        tuple<sequence<Impl::kCM0PerLane / SFactor, Impl::kCMLane, Impl::kCM1PerLane * SFactor>,
              sequence<Impl::kCNLane>>,
        tuple<sequence<1, 2>>,
        tuple<sequence<1, 0>>,
        sequence<1, 1>,
        sequence<0, 2>>;

    // c_vec += a_vec * b_vec
    template <bool post_nop_ = false>
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const BVecType& b_vec,
                                   bool_constant<post_nop_> = {}) const
    {
        static_for<0, kKIter, 1>{}([&](auto iKIter) { operator()(c_vec, a_vec, b_vec, iKIter); });
    }

    template <index_t iKIter, bool post_nop_ = false>
    CK_TILE_DEVICE void operator()(CVecType& c_vec,
                                   const AVecType& a_vec,
                                   const BVecType& b_vec,
                                   number<iKIter>,
                                   bool_constant<post_nop_> = {}) const
    {
        using buf_a = thread_buffer<typename Impl::AVecType, kKIter>;
        using buf_b = thread_buffer<typename Impl::BVecType, kKIter>;

        static_assert(iKIter < kKIter);

        Impl{}(c_vec,
               reinterpret_cast<const buf_a&>(a_vec)
                   .template get_as<typename Impl::AVecType>()[iKIter],
               reinterpret_cast<const buf_b&>(b_vec)
                   .template get_as<typename Impl::BVecType>()[iKIter],
               bool_constant<post_nop_>{});
    }

    // c_vec = a_vec * b_vec
    CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const
    {
        constexpr auto I0 = number<0>{};
        using buf_a       = thread_buffer<typename Impl::AVecType, kKIter>;
        using buf_b       = thread_buffer<typename Impl::BVecType, kKIter>;

        auto c_vec = Impl{}(
            reinterpret_cast<const buf_a&>(a_vec).template get_as<typename Impl::AVecType>()[I0],
            reinterpret_cast<const buf_b&>(b_vec).template get_as<typename Impl::BVecType>()[I0]);

        static_for<1, kKIter, 1>{}([&](auto iKIter) { operator()(c_vec, a_vec, b_vec, iKIter); });

        return c_vec;
    }
};

} // namespace ck_tile
