/*
 * Copyright (C) Advanced Micro Devices, Inc. All rights reserved.
 * Copyright (C) 2024-2026, The vLLM 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.
 */
#pragma once
#include "aiter_hip_common.h"
#include "dispatch_utils.h"
#include <ATen/hip/HIPContext.h>
#include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>
#include <torch/all.h>
#include <torch/extension.h>
#include <torch/torch.h>

#include <hip/hip_bf16.h>
#include <hip/hip_fp16.h>
typedef __hip_bfloat16 nv_bfloat16;

namespace aiter {
template <typename T, typename Operation>
inline __device__ T performOperation(T a, T b);

template <typename Operation>
torch::Tensor aten_compute(torch::Tensor& input, torch::Tensor& other);

struct AddOp
{
    template <typename T>
    inline __device__ static T apply(T a, T b)
    {
        return a + b;
    }

    static torch::Tensor compute(torch::Tensor& input, torch::Tensor& other)
    {
        return torch::add(input, other);
    }
};

struct SubOp
{
    template <typename T>
    inline __device__ static T apply(T a, T b)
    {
        return a - b;
    }

    static torch::Tensor compute(torch::Tensor& input, torch::Tensor& other)
    {
        return torch::sub(input, other);
    }
};

struct MulOp
{
    template <typename T>
    inline __device__ static T apply(T a, T b)
    {
        return a * b;
    }

    static torch::Tensor compute(torch::Tensor& input, torch::Tensor& other)
    {
        return torch::mul(input, other);
    }
};

struct DivOp
{
    template <typename T>
    inline __device__ static T apply(T a, T b)
    {
        // assert(b == static_cast<T>(0));
        return a / b;
    }

    static torch::Tensor compute(torch::Tensor& input, torch::Tensor& other)
    {
        return torch::div(input, other);
    }
};

template <typename T, typename Operation, bool order_flag>
inline __device__ T performOperation(T a, T b)
{
    if constexpr(std::is_same_v<Operation, AddOp>)
    {
        return Operation::apply(a, b);
    }
    else if constexpr(std::is_same_v<Operation, SubOp>)
    {
        if constexpr(!order_flag)
        {
            return Operation::apply(b, a);
        }
        else
        {
            return Operation::apply(a, b);
        }
    }
    else if constexpr(std::is_same_v<Operation, MulOp>)
    {
        return Operation::apply(a, b);
    }
    else if constexpr(std::is_same_v<Operation, DivOp>)
    {
        if constexpr(!order_flag)
        {
            return Operation::apply(b, a);
        }
        else
        {
            return Operation::apply(a, b);
        }
    }
    else
    {
        static_assert(false, "Unsupported operation");
    }
}
template <typename Operation>
torch::Tensor aten_compute(torch::Tensor& input, torch::Tensor& other)
{
    if constexpr(std::is_same_v<Operation, AddOp>)
    {
        return Operation::compute(input, other);
    }
    else if constexpr(std::is_same_v<Operation, SubOp>)
    {
        return Operation::compute(input, other);
    }
    else if constexpr(std::is_same_v<Operation, MulOp>)
    {
        return Operation::compute(input, other);
    }
    else if constexpr(std::is_same_v<Operation, DivOp>)
    {
        return Operation::compute(input, other);
    }
    else
    {
        static_assert(false, "Unsupported operation");
    }
}

template <class _T,
          int _WG,
          int BIG_TILE_SIZE_N,
          int BIG_TILE_SIZE_K,
          int M_SWIZZLE,
          typename Operation,
          bool order_flag,
          class _T0,
          class _T1>
__global__ void operator_tn_big_tile_kernel(const void* __restrict a,
                                            const void* __restrict b,
                                            void* __restrict c,
                                            const int N,
                                            const int K,
                                            int stride0,
                                            int stride2,
                                            bool types_match)
{
    // pad LDS row by dword
    constexpr uint32_t LDS_PAD         = 4 / sizeof(_T);
    constexpr uint32_t element_size    = sizeof(_T); // in bytes
    constexpr uint32_t elements_in_16B = 16 / element_size;

    union BLOCK_16B
    {
        _T e[elements_in_16B];
        __uint128_t ow;
    };

    // Round up processing to next full tile
    const uint32_t n_tiles        = (N + BIG_TILE_SIZE_N - 1) / BIG_TILE_SIZE_N;
    const uint32_t k_tiles        = (K + BIG_TILE_SIZE_K - 1) / BIG_TILE_SIZE_K;
    const uint32_t nk_tiles       = n_tiles * k_tiles;
    const uint32_t m_tiles        = gridDim.x / nk_tiles;
    const uint32_t m_tile_swizzle = blockIdx.x / nk_tiles / M_SWIZZLE * M_SWIZZLE;
    /// do m_swizzle when there are enough m_tiles
    const bool swizzle_m = m_tile_swizzle + M_SWIZZLE <= m_tiles;
    const uint32_t current_m =
        swizzle_m ? m_tile_swizzle + blockIdx.x % M_SWIZZLE : blockIdx.x / nk_tiles;

    const uint64_t stride_k      = N;
    const uint64_t out_stride_nk = N * K;

    const uint32_t current_nk =
        swizzle_m ? blockIdx.x / M_SWIZZLE % nk_tiles : blockIdx.x % nk_tiles;
    const uint32_t ti = current_nk / k_tiles;
    const uint32_t tj = current_nk % k_tiles;

    __shared__ _T0 sa[BIG_TILE_SIZE_N][BIG_TILE_SIZE_K + LDS_PAD];

    const uint32_t current_n_size = (ti == (n_tiles - 1) && (N % BIG_TILE_SIZE_N) != 0)
                                        ? (N % BIG_TILE_SIZE_N)
                                        : BIG_TILE_SIZE_N;
    const uint32_t current_k_size = (tj == (k_tiles - 1) && (K % BIG_TILE_SIZE_K) != 0)
                                        ? (K % BIG_TILE_SIZE_K)
                                        : BIG_TILE_SIZE_K;
    // use 128bit load&store whenever possible
    if(current_n_size % elements_in_16B == 0 && current_k_size % 8 == 0)
    {
        // Copy full tile with large loads
        constexpr uint32_t row_bytes       = BIG_TILE_SIZE_K;
        constexpr uint32_t ld_per_row      = row_bytes / elements_in_16B;
        constexpr uint32_t rows_per_wg     = _WG / ld_per_row;
        constexpr uint32_t vmem_per_thread = BIG_TILE_SIZE_N / rows_per_wg;
        // Make sure WG isn't too large
        static_assert(vmem_per_thread >= 1);

        const _T0* pat =
            (const _T0*)a + tj * row_bytes + ti * BIG_TILE_SIZE_N * stride2 + current_m * stride0;
#pragma unroll
        for(uint32_t t = 0; t < vmem_per_thread; t++)
        {
            uint32_t col    = threadIdx.x % ld_per_row;
            uint32_t row    = threadIdx.x / ld_per_row + t * rows_per_wg;
            uint64_t offset = (col * elements_in_16B < current_k_size && row < current_n_size)
                                  ? row * stride2 + col * elements_in_16B
                                  : 0;
            const _T0* pfa  = (const _T0*)(pat + offset);
            // BLOCK_16B d;
            // d.ow = *pfa;
#pragma unroll
            for(uint32_t i = 0; i < elements_in_16B; i++)
            {
                sa[row][col * elements_in_16B + i] = pfa[i];
            }
        }
        __syncthreads();
        // Copy full tile with large loads
        constexpr uint32_t row_bytes_wr    = BIG_TILE_SIZE_N;
        constexpr uint32_t vmem_per_row_wr = row_bytes_wr / elements_in_16B;
        constexpr uint32_t rows_per_wg_wr  = _WG / vmem_per_row_wr;
        constexpr uint32_t wr_per_row      = BIG_TILE_SIZE_K / rows_per_wg_wr;
        // Make sure WG isn't too large
        static_assert(wr_per_row >= 1);

        const _T1* pb = (const _T1*)b + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                        current_m * out_stride_nk;
        _T* pc = (_T*)c + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                 current_m * out_stride_nk;
#pragma unroll
        for(uint32_t t = 0; t < vmem_per_thread; t++)
        {
            uint32_t col = threadIdx.x % vmem_per_row_wr;
            uint32_t row = threadIdx.x / vmem_per_row_wr + t * rows_per_wg_wr;
            if(col * elements_in_16B < current_n_size && row < current_k_size)
            {
                uint64_t offset = row * stride_k + col * elements_in_16B;
                BLOCK_16B d;
                if(types_match)
                {
                    const __uint128_t* pfb = (const __uint128_t*)(pb + offset);
                    d.ow                   = *pfb;
// Transpose tile on read from LDS
#pragma unroll
                    for(uint32_t i = 0; i < elements_in_16B; i++)
                    {
                        d.e[i] = performOperation<_T, Operation, order_flag>(
                            static_cast<_T>(sa[col * elements_in_16B + i][row]), d.e[i]);
                    }
                    __uint128_t* pfc = (__uint128_t*)(pc + offset);
                    *pfc             = d.ow;
                }
                else
                {
                    const _T1* pfb = (const _T1*)(pb + offset);
#pragma unroll
                    for(uint32_t i = 0; i < elements_in_16B; i++)
                    {
                        float a = static_cast<float>(sa[col * elements_in_16B + i][row]);
                        float b = static_cast<float>(pfb[i]);
                        float c = performOperation<float, Operation, order_flag>(a, b);
                        d.e[i]  = static_cast<_T>(c);
                    }
                    __uint128_t* pfc = (__uint128_t*)(pc + offset);
                    *pfc             = d.ow;
                }
            }
        }
    }
    else
    {
        // Copy partial tiles with element accesses
        constexpr uint32_t row_bytes       = BIG_TILE_SIZE_K;
        constexpr uint32_t ld_per_row      = BIG_TILE_SIZE_K;
        constexpr uint32_t rows_per_wg     = _WG / ld_per_row;
        constexpr uint32_t vmem_per_thread = BIG_TILE_SIZE_N / rows_per_wg;
        // Make sure WG isn't too large
        static_assert(vmem_per_thread >= 1);

        const _T0* pat =
            (const _T0*)a + ti * BIG_TILE_SIZE_N * stride2 + tj * row_bytes + current_m * stride0;
#pragma unroll
        for(uint32_t t = 0; t < vmem_per_thread; t++)
        {
            uint32_t col = threadIdx.x % ld_per_row;
            uint32_t row = threadIdx.x / ld_per_row + t * rows_per_wg;
            uint64_t offset =
                (col < current_k_size && row < current_n_size) ? row * stride2 + col : 0;
            const _T0* pfa = (const _T0*)(pat + offset);
            sa[row][col]   = *pfa;
        }
        __syncthreads();

        // Copy full tile with large loads
        constexpr uint32_t row_bytes_wr    = BIG_TILE_SIZE_N;
        constexpr uint32_t vmem_per_row_wr = BIG_TILE_SIZE_N;
        constexpr uint32_t rows_per_wg_wr  = _WG / vmem_per_row_wr;
        constexpr uint32_t wr_per_row      = BIG_TILE_SIZE_K / rows_per_wg_wr;
        const _T1* pb = (const _T1*)b + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                        current_m * out_stride_nk;
        _T* pc = (_T*)c + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                 current_m * out_stride_nk;
#pragma unroll
        for(uint32_t t = 0; t < wr_per_row; t++)
        {
            uint32_t col = threadIdx.x % vmem_per_row_wr;
            uint32_t row = threadIdx.x / vmem_per_row_wr + t * rows_per_wg_wr;
            if(col < current_n_size && row < current_k_size)
            {
                uint64_t offset = row * stride_k + col;
                const _T1* pfb  = (const _T1*)(pb + offset);
                _T* pfc         = (_T*)(pc + offset);
                if(types_match)
                {
                    *pfc = performOperation<_T, Operation, order_flag>(
                        static_cast<_T>(sa[col][row]), static_cast<_T>(*pfb));
                }
                else
                {
                    float a = static_cast<float>(sa[col][row]);
                    float b = static_cast<float>(*pfb);
                    float c = performOperation<float, Operation, order_flag>(a, b);
                    *pfc    = static_cast<_T>(c);
                }
            }
        }
    }
}

template <class _T,
          int _WG,
          int BIG_TILE_SIZE_N,
          int BIG_TILE_SIZE_K,
          int M_SWIZZLE,
          typename Operation,
          bool order_flag,
          class _T0,
          class _T1>
__global__ void operator_bcast_big_tile_kernel(const void* __restrict a,
                                               const void* __restrict b,
                                               void* __restrict c,
                                               const int N,
                                               const int K,
                                               bool types_match)
{
    constexpr uint32_t element_size    = sizeof(_T); // in bytes
    constexpr uint32_t elements_in_16B = 16 / element_size;

    union BLOCK_16B
    {
        _T e[elements_in_16B];
        __uint128_t ow;
    };

    // Round up processing to next full tile
    const uint32_t n_tiles        = (N + BIG_TILE_SIZE_N - 1) / BIG_TILE_SIZE_N;
    const uint32_t k_tiles        = (K + BIG_TILE_SIZE_K - 1) / BIG_TILE_SIZE_K;
    const uint32_t nk_tiles       = n_tiles * k_tiles;
    const uint32_t m_tiles        = gridDim.x / nk_tiles;
    const uint32_t m_tile_swizzle = blockIdx.x / nk_tiles / M_SWIZZLE * M_SWIZZLE;
    /// do m_swizzle when there are enough m_tiles
    const bool swizzle_m = m_tile_swizzle + M_SWIZZLE <= m_tiles;
    const uint32_t current_m =
        swizzle_m ? m_tile_swizzle + blockIdx.x % M_SWIZZLE : blockIdx.x / nk_tiles;

    const uint64_t stride_k      = N;
    const uint64_t out_stride_nk = N * K;

    const uint32_t current_nk =
        swizzle_m ? blockIdx.x / M_SWIZZLE % nk_tiles : blockIdx.x % nk_tiles;
    const uint32_t ti = current_nk / k_tiles;
    const uint32_t tj = current_nk % k_tiles;

    const uint32_t current_n_size = (ti == (n_tiles - 1) && (N % BIG_TILE_SIZE_N) != 0)
                                        ? (N % BIG_TILE_SIZE_N)
                                        : BIG_TILE_SIZE_N;
    const uint32_t current_k_size = (tj == (k_tiles - 1) && (K % BIG_TILE_SIZE_K) != 0)
                                        ? (K % BIG_TILE_SIZE_K)
                                        : BIG_TILE_SIZE_K;

    // use 128bit load&store whenever possible
    if(current_n_size % 8 == 0 && current_k_size % elements_in_16B == 0)
    {
        // Copy full tile with large loads
        constexpr uint32_t row_bytes_wr    = BIG_TILE_SIZE_N;
        constexpr uint32_t vmem_per_row_wr = row_bytes_wr / elements_in_16B;
        constexpr uint32_t rows_per_wg_wr  = _WG / vmem_per_row_wr;
        constexpr uint32_t wr_per_row      = BIG_TILE_SIZE_K / rows_per_wg_wr;
        // Make sure WG isn't too large
        static_assert(wr_per_row >= 1);

        const _T0* pa = (const _T0*)a + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                        current_m * out_stride_nk;
        const _T1* pb = (const _T1*)b + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr;
        const _T* pc  = (const _T*)c + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                       current_m * out_stride_nk;
#pragma unroll
        for(uint32_t t = 0; t < wr_per_row; t++)
        {
            uint32_t col = threadIdx.x % vmem_per_row_wr;
            uint32_t row = threadIdx.x / vmem_per_row_wr + t * rows_per_wg_wr;
            if(col * elements_in_16B < current_n_size && row < current_k_size)
            {
                BLOCK_16B d, f;
                uint64_t offset = row * stride_k + col * elements_in_16B;
                if(types_match)
                {
                    const __uint128_t* pfa = (const __uint128_t*)(pa + offset);
                    const __uint128_t* pfb = (const __uint128_t*)(pb + offset);
                    f.ow                   = *pfa;
                    d.ow                   = *pfb;
#pragma unroll
                    for(uint32_t i = 0; i < elements_in_16B; i++)
                    {
                        d.e[i] = performOperation<_T, Operation, order_flag>(
                            static_cast<_T>(f.e[i]), static_cast<_T>(d.e[i]));
                    }
                    __uint128_t* pfc = (__uint128_t*)(pc + offset);
                    *pfc             = d.ow;
                }
                else
                {
                    const _T0* pfa = (const _T0*)(pa + offset);
                    const _T1* pfb = (const _T1*)(pb + offset);
#pragma unroll
                    for(uint32_t i = 0; i < elements_in_16B; i++)
                    {
                        float a = static_cast<float>(pfa[i]);
                        float b = static_cast<float>(pfb[i]);
                        float c = performOperation<float, Operation, order_flag>(a, b);
                        d.e[i]  = static_cast<_T>(c);
                    }
                    __uint128_t* pfc = (__uint128_t*)(pc + offset);
                    *pfc             = d.ow;
                }
            }
        }
    }
    else
    {
        // Copy full tile with large loads
        constexpr uint32_t row_bytes_wr    = BIG_TILE_SIZE_N;
        constexpr uint32_t vmem_per_row_wr = BIG_TILE_SIZE_N;
        constexpr uint32_t rows_per_wg_wr  = _WG / vmem_per_row_wr;
        constexpr uint32_t wr_per_row      = BIG_TILE_SIZE_K / rows_per_wg_wr;
        const _T0* pa = (const _T0*)a + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                        current_m * out_stride_nk;
        const _T1* pb = (const _T1*)b + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr;
        const _T* pc  = (const _T*)c + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                       current_m * out_stride_nk;
#pragma unroll
        for(uint32_t t = 0; t < wr_per_row; t++)
        {
            uint32_t col = threadIdx.x % vmem_per_row_wr;
            uint32_t row = threadIdx.x / vmem_per_row_wr + t * rows_per_wg_wr;
            if(col < current_n_size && row < current_k_size)
            {
                uint64_t offset = row * stride_k + col;
                const _T0* pfa  = (const _T0*)(pa + offset);
                const _T1* pfb  = (const _T1*)(pb + offset);
                _T* pfc         = (_T*)(pc + offset);
                if(types_match)
                {
                    *pfc = performOperation<_T, Operation, order_flag>(static_cast<_T>(*pfa),
                                                                       static_cast<_T>(*pfb));
                }
                else
                {
                    float a = static_cast<float>(*pfa);
                    float b = static_cast<float>(*pfb);
                    float c = performOperation<float, Operation, order_flag>(a, b);
                    *pfc    = static_cast<_T>(c);
                }
            }
        }
    }
}

template <class _T,
          int _WG,
          int BIG_TILE_SIZE_N,
          int BIG_TILE_SIZE_K,
          int M_SWIZZLE,
          typename Operation,
          bool order_flag,
          class _T0,
          class _T1>
__global__ void operator_bcast1_big_tile_kernel(const void* __restrict a,
                                                const void* __restrict b,
                                                void* __restrict c,
                                                const int N,
                                                const int K,
                                                bool types_match)
{
    // pad LDS row by dword
    constexpr uint32_t element_size    = sizeof(_T); // in bytes
    constexpr uint32_t elements_in_16B = 16 / element_size;

    union BLOCK_16B
    {
        _T e[elements_in_16B];
        __uint128_t ow;
    };

    // Round up processing to next full tile
    const uint32_t n_tiles        = (N + BIG_TILE_SIZE_N - 1) / BIG_TILE_SIZE_N;
    const uint32_t k_tiles        = (K + BIG_TILE_SIZE_K - 1) / BIG_TILE_SIZE_K;
    const uint32_t nk_tiles       = n_tiles * k_tiles;
    const uint32_t m_tiles        = gridDim.x / nk_tiles;
    const uint32_t m_tile_swizzle = blockIdx.x / nk_tiles / M_SWIZZLE * M_SWIZZLE;
    /// do m_swizzle when there are enough m_tiles
    const bool swizzle_m = m_tile_swizzle + M_SWIZZLE <= m_tiles;
    const uint32_t current_m =
        swizzle_m ? m_tile_swizzle + blockIdx.x % M_SWIZZLE : blockIdx.x / nk_tiles;

    const uint64_t stride_k      = N;
    const uint64_t out_stride_nk = N * K;

    const uint32_t current_nk =
        swizzle_m ? blockIdx.x / M_SWIZZLE % nk_tiles : blockIdx.x % nk_tiles;
    const uint32_t ti = current_nk / k_tiles;
    const uint32_t tj = current_nk % k_tiles;

    const uint32_t current_n_size = (ti == (n_tiles - 1) && (N % BIG_TILE_SIZE_N) != 0)
                                        ? (N % BIG_TILE_SIZE_N)
                                        : BIG_TILE_SIZE_N;
    const uint32_t current_k_size = (tj == (k_tiles - 1) && (K % BIG_TILE_SIZE_K) != 0)
                                        ? (K % BIG_TILE_SIZE_K)
                                        : BIG_TILE_SIZE_K;

    // use 128bit load&store whenever possible
    if(current_n_size % 8 == 0 && current_k_size % elements_in_16B == 0)
    {
        // Copy full tile with large loads
        constexpr uint32_t row_bytes_wr    = BIG_TILE_SIZE_N;
        constexpr uint32_t vmem_per_row_wr = row_bytes_wr / elements_in_16B;
        constexpr uint32_t rows_per_wg_wr  = _WG / vmem_per_row_wr;
        constexpr uint32_t wr_per_row      = BIG_TILE_SIZE_K / rows_per_wg_wr;
        // Make sure WG isn't too large
        static_assert(wr_per_row >= 1);

        const _T0* pa = (const _T0*)a + ti * row_bytes_wr + current_m * stride_k;
        const _T1* pb = (const _T1*)b + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                        current_m * out_stride_nk;
        const _T* pc = (const _T*)c + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                       current_m * out_stride_nk;
#pragma unroll
        for(uint32_t t = 0; t < wr_per_row; t++)
        {
            uint32_t col = threadIdx.x % vmem_per_row_wr;
            uint32_t row = threadIdx.x / vmem_per_row_wr + t * rows_per_wg_wr;
            if(col * elements_in_16B < current_n_size && row < current_k_size)
            {
                uint64_t offset_a = col * elements_in_16B;
                uint64_t offset   = row * stride_k + col * elements_in_16B;
                BLOCK_16B d, f;
                if(types_match)
                {
                    const __uint128_t* pfa = (const __uint128_t*)(pa + offset_a);
                    const __uint128_t* pfb = (const __uint128_t*)(pb + offset);
                    f.ow                   = *pfa;
                    d.ow                   = *pfb;
#pragma unroll
                    for(uint32_t i = 0; i < elements_in_16B; i++)
                    {
                        d.e[i] = performOperation<_T, Operation, order_flag>(
                            static_cast<_T>(f.e[i]), static_cast<_T>(d.e[i]));
                    }
                    __uint128_t* pfc = (__uint128_t*)(pc + offset);
                    *pfc             = d.ow;
                }
                else
                {
                    const _T0* pfa = (const _T0*)(pa + offset_a);
                    const _T1* pfb = (const _T1*)(pb + offset);
#pragma unroll
                    for(uint32_t i = 0; i < elements_in_16B; i++)
                    {
                        float a = static_cast<float>(pfa[i]);
                        float b = static_cast<float>(pfb[i]);
                        float c = performOperation<float, Operation, order_flag>(a, b);
                        d.e[i]  = static_cast<_T>(c);
                    }
                    __uint128_t* pfc = (__uint128_t*)(pc + offset);
                    *pfc             = d.ow;
                }
            }
        }
    }
    else
    {
        // Copy full tile with large loads
        constexpr uint32_t row_bytes_wr    = BIG_TILE_SIZE_N;
        constexpr uint32_t vmem_per_row_wr = BIG_TILE_SIZE_N;
        constexpr uint32_t rows_per_wg_wr  = _WG / vmem_per_row_wr;
        constexpr uint32_t wr_per_row      = BIG_TILE_SIZE_K / rows_per_wg_wr;
        const _T0* pa = (const _T0*)a + ti * row_bytes_wr + current_m * stride_k;
        const _T1* pb = (const _T1*)b + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                        current_m * out_stride_nk;
        const _T* pc = (const _T*)c + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                       current_m * out_stride_nk;
#pragma unroll
        for(uint32_t t = 0; t < wr_per_row; t++)
        {
            uint32_t col = threadIdx.x % vmem_per_row_wr;
            uint32_t row = threadIdx.x / vmem_per_row_wr + t * rows_per_wg_wr;
            if(col < current_n_size && row < current_k_size)
            {
                uint64_t offset_a = col;
                uint64_t offset   = row * stride_k + col;
                const _T0* pfa    = (const _T0*)(pa + offset_a);
                const _T1* pfb    = (const _T1*)(pb + offset);
                _T* pfc           = (_T*)(pc + offset);
                if(types_match)
                {
                    *pfc = performOperation<_T, Operation, order_flag>(static_cast<_T>(*pfa),
                                                                       static_cast<_T>(*pfb));
                }
                else
                {
                    float a = static_cast<float>(*pfa);
                    float b = static_cast<float>(*pfb);
                    float c = performOperation<float, Operation, order_flag>(a, b);
                    *pfc    = static_cast<_T>(c);
                }
            }
        }
    }
}

template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_bcast_tile_kernel(const void* __restrict a,
                                           const void* __restrict b,
                                           void* __restrict c,
                                           const int M,
                                           const int N,
                                           const int K,
                                           bool types_match)
{
    constexpr uint32_t element_size    = sizeof(_T); // in bytes
    constexpr uint32_t elements_in_16B = 16 / element_size;
    uint64_t idx                       = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
    uint32_t n_tiles                   = N / _rows;
    uint32_t k_tiles                   = K / elements_in_16B;
    if(idx < (uint64_t)M * n_tiles * k_tiles)
    {
        uint32_t ti        = idx / (k_tiles * n_tiles);
        uint64_t idx_block = idx % (k_tiles * n_tiles);
        uint32_t tj        = (idx_block / k_tiles) % n_tiles;
        uint32_t tk        = idx_block % k_tiles;
        for(int row = 0; row < _rows; row++)
        {
            uint64_t offset_b = (uint64_t)(tj + row * n_tiles) * K + tk * elements_in_16B;
            uint64_t offset_ac =
                (uint64_t)(tj + row * n_tiles) * K + tk * elements_in_16B + (uint64_t)ti * N * K;
            const _T0* pa = reinterpret_cast<const _T0*>(a) + offset_ac;
            const _T1* pb = reinterpret_cast<const _T1*>(b) + offset_b;
            _T* pc        = reinterpret_cast<_T*>(c) + offset_ac;
            for(int col = 0; col < elements_in_16B; col++)
            {
                const _T0* pfa = pa + col;
                const _T1* pfb = pb + col;
                _T* pfc        = pc + col;
                if(types_match)
                {
                    *pfc = performOperation<_T, Operation, order_flag>(static_cast<_T>(*pfa),
                                                                       static_cast<_T>(*pfb));
                }
                else
                {
                    float t0 = static_cast<float>(*pfa);
                    float t1 = static_cast<float>(*pfb);
                    float t2 = performOperation<float, Operation, order_flag>(t0, t1);
                    *pfc     = static_cast<_T>(t2);
                }
            }
        }
    }
}

// (m, n, k), (k,)
template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_bcastK_unroll_vectorize_naive(const void* __restrict a,
                                                       const void* __restrict b,
                                                       void* __restrict c,
                                                       const int m,
                                                       const int n,
                                                       const int k,
                                                       bool types_match)
{
    constexpr uint32_t element_size = sizeof(_T);
    constexpr uint32_t vec_size     = 16 / element_size;
    uint64_t idx     = (uint64_t)(blockIdx.x * blockDim.x * _rows + threadIdx.x) * vec_size;
    const _T0* a_ptr = reinterpret_cast<const _T0*>(a);
    const _T1* b_ptr = reinterpret_cast<const _T1*>(b);
    _T* c_ptr        = reinterpret_cast<_T*>(c);
    for(uint64_t index = idx; index + (uint64_t)((_rows - 1) * blockDim.x * vec_size) < (m * n * k);
        index += gridDim.x * blockDim.x * _rows * vec_size)
    {
#pragma unroll
        for(int unroll_index = 0; unroll_index < _rows; ++unroll_index)
        {
            uint64_t block_offset = (uint64_t)blockDim.x * vec_size * unroll_index;
            uint64_t other_start  = (index + block_offset) % k;
#pragma unroll
            for(int vec_index = 0; vec_index < vec_size; ++vec_index)
            {
                uint64_t input_offset = index + block_offset + vec_index;
                uint64_t other_offset = other_start + vec_index;
                if(types_match)
                {
                    _T t0                   = static_cast<_T>(a_ptr[input_offset]);
                    _T t1                   = static_cast<_T>(b_ptr[other_offset]);
                    *(c_ptr + input_offset) = performOperation<_T, Operation, order_flag>(t0, t1);
                }
                else
                {
                    float t0 = static_cast<float>(a_ptr[input_offset]);
                    float t1 = static_cast<float>(b_ptr[other_offset]);
                    float t2 = performOperation<float, Operation, order_flag>(t0, t1);
                    *(c_ptr + input_offset) = static_cast<_T>(t2);
                }
            }
        }
    }
}

// (m, n, k), (1)
template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_bcast_scalar_unroll_vectorize_naive(const void* __restrict a,
                                                             const void* __restrict b,
                                                             void* __restrict c,
                                                             const int n)
{
    constexpr uint32_t element_size = sizeof(_T);
    constexpr uint32_t vec_size     = 16 / element_size;
    uint64_t idx     = (uint64_t)(blockIdx.x * blockDim.x * _rows + threadIdx.x) * vec_size;
    const _T0* a_ptr = reinterpret_cast<const _T0*>(a);
    const _T1 b_val  = *reinterpret_cast<const _T1*>(b);
    _T* c_ptr        = reinterpret_cast<_T*>(c);
    _T b_T           = static_cast<_T>(b_val);
    for(uint64_t index = idx; index + (uint64_t)((_rows - 1) * blockDim.x * vec_size) < n;
        index += gridDim.x * blockDim.x * _rows * vec_size)
    {
#pragma unroll
        for(int unroll_index = 0; unroll_index < _rows; ++unroll_index)
        {
            uint64_t block_offset = (uint64_t)blockDim.x * vec_size * unroll_index;
#pragma unroll
            for(int vec_index = 0; vec_index < vec_size; ++vec_index)
            {
                uint64_t input_offset   = index + block_offset + vec_index;
                _T t0                   = static_cast<_T>(a_ptr[input_offset]);
                *(c_ptr + input_offset) = performOperation<_T, Operation, order_flag>(t0, b_T);
            }
        }
    }
}

// (m, n, k), (m, 1, k)
template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_bcastM1K_unroll_kernel(const void* __restrict a,
                                                const void* __restrict b,
                                                void* __restrict c,
                                                const int m,
                                                const int n,
                                                const int k,
                                                bool types_match)
{
    constexpr uint32_t element_size = sizeof(_T);
    constexpr uint32_t vec_size     = 16 / element_size;
    uint64_t idx     = (uint64_t)(blockIdx.x * blockDim.x * _rows + threadIdx.x) * vec_size;
    const _T0* a_ptr = reinterpret_cast<const _T0*>(a);
    const _T1* b_ptr = reinterpret_cast<const _T1*>(b);
    _T* c_ptr        = reinterpret_cast<_T*>(c);
    for(uint64_t index = idx; index + (uint64_t)((_rows - 1) * blockDim.x * vec_size) < (m * n * k);
        index += gridDim.x * blockDim.x * _rows * vec_size)
    {
#pragma unroll
        for(int unroll_index = 0; unroll_index < _rows; ++unroll_index)
        {
            uint64_t block_offset  = (uint64_t)blockDim.x * vec_size * unroll_index;
            uint64_t other_m_index = (index + block_offset) / (n * k);
            uint64_t other_k_index = (index + block_offset) % k;
#pragma unroll
            for(int vec_index = 0; vec_index < vec_size; ++vec_index)
            {
                uint64_t input_offset = index + block_offset + vec_index;
                uint64_t other_offset = other_m_index * k + other_k_index + vec_index;
                if(types_match)
                {
                    _T t0                   = static_cast<_T>(a_ptr[input_offset]);
                    _T t1                   = static_cast<_T>(b_ptr[other_offset]);
                    *(c_ptr + input_offset) = performOperation<_T, Operation, order_flag>(t0, t1);
                }
                else
                {
                    float t0 = static_cast<float>(a_ptr[input_offset]);
                    float t1 = static_cast<float>(b_ptr[other_offset]);
                    float t2 = performOperation<float, Operation, order_flag>(t0, t1);
                    *(c_ptr + input_offset) = static_cast<_T>(t2);
                }
            }
        }
    }
}

// (m, n, k), (m, n, 1)
template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_bcastMN1_unroll_vec_naive(const void* __restrict a,
                                                   const void* __restrict b,
                                                   void* __restrict c,
                                                   const int forward_dim,
                                                   const int bcast_dim,
                                                   bool types_match)
{
    constexpr uint32_t element_size = sizeof(_T);
    constexpr uint32_t vec_size     = 16 / element_size;
    uint64_t idx     = (uint64_t)(blockIdx.x * blockDim.x * _rows + threadIdx.x) * vec_size;
    const _T0* a_ptr = reinterpret_cast<const _T0*>(a);
    const _T1* b_ptr = reinterpret_cast<const _T1*>(b);
    _T* c_ptr        = reinterpret_cast<_T*>(c);
    for(uint64_t index = idx;
        index + (uint64_t)((_rows - 1) * blockDim.x * vec_size) < (forward_dim * bcast_dim);
        index += gridDim.x * blockDim.x * _rows * vec_size)
    {
#pragma unroll
        for(int unroll_index = 0; unroll_index < _rows; ++unroll_index)
        {
            uint64_t block_offset = (uint64_t)blockDim.x * vec_size * unroll_index;
            uint64_t other_offset = (index + block_offset) / bcast_dim;
#pragma unroll
            for(int vec_index = 0; vec_index < vec_size; ++vec_index)
            {
                uint64_t input_offset = index + block_offset + vec_index;
                if(types_match)
                {
                    _T t0                   = static_cast<_T>(a_ptr[input_offset]);
                    _T t1                   = static_cast<_T>(b_ptr[other_offset]);
                    *(c_ptr + input_offset) = performOperation<_T, Operation, order_flag>(t0, t1);
                }
                else
                {
                    float t0 = static_cast<float>(a_ptr[input_offset]);
                    float t1 = static_cast<float>(b_ptr[other_offset]);
                    float t2 = performOperation<float, Operation, order_flag>(t0, t1);
                    *(c_ptr + input_offset) = static_cast<_T>(t2);
                }
            }
        }
    }
}

template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_bcastMN1_naive(const void* __restrict a,
                                        const void* __restrict b,
                                        void* __restrict c,
                                        const int forward_dim,
                                        const int bcast_dim,
                                        bool types_match)
{
    uint64_t idx     = blockIdx.x * blockDim.x + threadIdx.x;
    const _T0* a_ptr = reinterpret_cast<const _T0*>(a);
    const _T1* b_ptr = reinterpret_cast<const _T1*>(b);
    _T* c_ptr        = reinterpret_cast<_T*>(c);
    for(uint64_t index = idx; index < forward_dim * bcast_dim; index += blockDim.x * gridDim.x)
    {
        if(types_match)
        {
            _T t0            = static_cast<_T>(a_ptr[index]);
            _T t1            = static_cast<_T>(b_ptr[index / bcast_dim]);
            *(c_ptr + index) = performOperation<_T, Operation, order_flag>(t0, t1);
        }
        else
        {
            float t0         = static_cast<float>(a_ptr[index]);
            float t1         = static_cast<float>(b_ptr[index / bcast_dim]);
            float t2         = performOperation<float, Operation, order_flag>(t0, t1);
            *(c_ptr + index) = static_cast<_T>(t2);
        }
    }
}

template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_bcast1N1_unroll_vec_naive(const void* __restrict a,
                                                   const void* __restrict b,
                                                   void* __restrict c,
                                                   const int m,
                                                   const int n,
                                                   const int k,
                                                   bool types_match)
{
    constexpr uint32_t element_size = sizeof(_T);
    constexpr uint32_t vec_size     = 16 / element_size;
    uint64_t idx     = (uint64_t)(blockIdx.x * blockDim.x * _rows + threadIdx.x) * vec_size;
    const _T0* a_ptr = reinterpret_cast<const _T0*>(a);
    const _T1* b_ptr = reinterpret_cast<const _T1*>(b);
    _T* c_ptr        = reinterpret_cast<_T*>(c);
    for(uint64_t index = idx; index + (uint64_t)((_rows - 1) * blockDim.x * vec_size) < (m * n * k);
        index += gridDim.x * blockDim.x * _rows * vec_size)
    {
#pragma unroll
        for(int unroll_index = 0; unroll_index < _rows; ++unroll_index)
        {
            uint64_t block_offset = (uint64_t)blockDim.x * vec_size * unroll_index;
            uint64_t other_offset = (index + block_offset) % (n * k) / k;
#pragma unroll
            for(int vec_index = 0; vec_index < vec_size; ++vec_index)
            {
                uint64_t input_offset = index + block_offset + vec_index;
                if(types_match)
                {
                    _T t0                   = static_cast<_T>(a_ptr[input_offset]);
                    _T t1                   = static_cast<_T>(b_ptr[other_offset]);
                    *(c_ptr + input_offset) = performOperation<_T, Operation, order_flag>(t0, t1);
                }
                else
                {
                    float t0 = static_cast<float>(a_ptr[input_offset]);
                    float t1 = static_cast<float>(b_ptr[other_offset]);
                    float t2 = performOperation<float, Operation, order_flag>(t0, t1);
                    *(c_ptr + input_offset) = static_cast<_T>(t2);
                }
            }
        }
    }
}

template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_bcast1N1_naive(const void* __restrict a,
                                        const void* __restrict b,
                                        void* __restrict c,
                                        const int m,
                                        const int n,
                                        const int k,
                                        bool types_match)
{
    uint64_t idx     = blockIdx.x * blockDim.x + threadIdx.x;
    const _T0* a_ptr = reinterpret_cast<const _T0*>(a);
    const _T1* b_ptr = reinterpret_cast<const _T1*>(b);
    _T* c_ptr        = reinterpret_cast<_T*>(c);
    for(uint64_t index = idx; index < m * n * k; index += blockDim.x * gridDim.x)
    {
        if(types_match)
        {
            _T t0            = static_cast<_T>(a_ptr[index]);
            _T t1            = static_cast<_T>(b_ptr[index % (n * k) / k]);
            *(c_ptr + index) = performOperation<_T, Operation, order_flag>(t0, t1);
        }
        else
        {
            float t0         = static_cast<float>(a_ptr[index]);
            float t1         = static_cast<float>(b_ptr[index % (n * k) / k]);
            float t2         = performOperation<float, Operation, order_flag>(t0, t1);
            *(c_ptr + index) = static_cast<_T>(t2);
        }
    }
}

template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_bcastN11_unroll_vec_naive(const void* __restrict a,
                                                   const void* __restrict b,
                                                   void* __restrict c,
                                                   const int m,
                                                   const int n,
                                                   const int k,
                                                   bool types_match)
{
    constexpr uint32_t element_size = sizeof(_T);
    constexpr uint32_t vec_size     = 16 / element_size;
    uint64_t idx     = (uint64_t)(blockIdx.x * blockDim.x * _rows + threadIdx.x) * vec_size;
    const _T0* a_ptr = reinterpret_cast<const _T0*>(a);
    const _T1* b_ptr = reinterpret_cast<const _T1*>(b);
    _T* c_ptr        = reinterpret_cast<_T*>(c);
    for(uint64_t index = idx; index + (uint64_t)((_rows - 1) * blockDim.x * vec_size) < (m * n * k);
        index += gridDim.x * blockDim.x * _rows * vec_size)
    {
#pragma unroll
        for(int unroll_index = 0; unroll_index < _rows; ++unroll_index)
        {
            uint64_t block_offset = (uint64_t)blockDim.x * vec_size * unroll_index;
            // uint64_t other_offset = (index + block_offset) / (n * k);
#pragma unroll
            for(int vec_index = 0; vec_index < vec_size; ++vec_index)
            {
                uint64_t input_offset = index + block_offset + vec_index;
                uint64_t other_offset = input_offset / (n * k);
                if(types_match)
                {
                    _T t0                   = static_cast<_T>(a_ptr[input_offset]);
                    _T t1                   = static_cast<_T>(b_ptr[other_offset]);
                    *(c_ptr + input_offset) = performOperation<_T, Operation, order_flag>(t0, t1);
                }
                else
                {
                    float t0 = static_cast<float>(a_ptr[input_offset]);
                    float t1 = static_cast<float>(b_ptr[other_offset]);
                    float t2 = performOperation<float, Operation, order_flag>(t0, t1);
                    *(c_ptr + input_offset) = static_cast<_T>(t2);
                }
            }
        }
    }
}

template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_bcastN11_unroll_vec_pad(const void* __restrict a,
                                                 const void* __restrict b,
                                                 void* __restrict c,
                                                 const int m,
                                                 const int n,
                                                 const int k,
                                                 const int padded_size,
                                                 bool types_match)
{
    constexpr uint32_t element_size = sizeof(_T);
    constexpr uint32_t vec_size     = 16 / element_size;
    uint64_t idx     = (uint64_t)(blockIdx.x * blockDim.x * _rows + threadIdx.x) * vec_size;
    const _T0* a_ptr = reinterpret_cast<const _T0*>(a);
    const _T1* b_ptr = reinterpret_cast<const _T1*>(b);
    _T* c_ptr        = reinterpret_cast<_T*>(c);
    for(uint64_t index = idx; index + (uint64_t)((_rows - 1) * blockDim.x * vec_size) < padded_size;
        index += gridDim.x * blockDim.x * _rows * vec_size)
    {
#pragma unroll
        for(int unroll_index = 0; unroll_index < _rows; ++unroll_index)
        {
            uint64_t block_offset = (uint64_t)blockDim.x * vec_size * unroll_index;
            // uint64_t other_offset = (index + block_offset) / (n * k);
            if(index + block_offset < (m * n * k))
            {
#pragma unroll
                for(int vec_index = 0; vec_index < vec_size; ++vec_index)
                {
                    uint64_t input_offset = index + block_offset + vec_index;
                    uint64_t other_offset = input_offset / (n * k);
                    if(types_match)
                    {
                        _T t0 = static_cast<_T>(a_ptr[input_offset]);
                        _T t1 = static_cast<_T>(b_ptr[other_offset]);
                        *(c_ptr + input_offset) =
                            performOperation<_T, Operation, order_flag>(t0, t1);
                    }
                    else
                    {
                        float t0 = static_cast<float>(a_ptr[input_offset]);
                        float t1 = static_cast<float>(b_ptr[other_offset]);
                        float t2 = performOperation<float, Operation, order_flag>(t0, t1);
                        *(c_ptr + input_offset) = static_cast<_T>(t2);
                    }
                }
            }
        }
    }
}

template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_bcastN11_naive(const void* __restrict a,
                                        const void* __restrict b,
                                        void* __restrict c,
                                        const int m,
                                        const int n,
                                        const int k,
                                        bool types_match)
{
    uint64_t idx     = blockIdx.x * blockDim.x + threadIdx.x;
    const _T0* a_ptr = reinterpret_cast<const _T0*>(a);
    const _T1* b_ptr = reinterpret_cast<const _T1*>(b);
    _T* c_ptr        = reinterpret_cast<_T*>(c);
    for(uint64_t index = idx; index < m * n * k; index += blockDim.x * gridDim.x)
    {
        if(types_match)
        {
            _T t0            = static_cast<_T>(a_ptr[index]);
            _T t1            = static_cast<_T>(b_ptr[index / (n * k)]);
            *(c_ptr + index) = performOperation<_T, Operation, order_flag>(t0, t1);
        }
        else
        {
            float t0         = static_cast<float>(a_ptr[index]);
            float t1         = static_cast<float>(b_ptr[index / (n * k)]);
            float t2         = performOperation<float, Operation, order_flag>(t0, t1);
            *(c_ptr + index) = static_cast<_T>(t2);
        }
    }
}

template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_contiguous_kernel_naive(const void* __restrict a,
                                                 const void* __restrict b,
                                                 void* __restrict c,
                                                 const int n,
                                                 bool types_match)
{
    constexpr uint32_t element_size = sizeof(_T);
    constexpr uint32_t vec_size     = 16 / element_size;
    uint64_t idx     = (uint64_t)(blockIdx.x * blockDim.x * _rows + threadIdx.x) * vec_size;
    const _T0* a_ptr = reinterpret_cast<const _T0*>(a);
    const _T1* b_ptr = reinterpret_cast<const _T1*>(b);
    _T* c_ptr        = reinterpret_cast<_T*>(c);
    for(uint64_t index = idx; index + (uint64_t)(_rows - 1) * blockDim.x * vec_size < n;
        index += gridDim.x * blockDim.x * _rows * vec_size)
    {
#pragma unroll
        for(int unroll_index = 0; unroll_index < _rows; ++unroll_index)
        {
            uint64_t block_offset = (uint64_t)blockDim.x * vec_size * unroll_index;
#pragma unroll
            for(int vec_index = 0; vec_index < vec_size; ++vec_index)
            {
                uint64_t offset = index + block_offset + vec_index;
                if(types_match)
                {
                    _T t0             = static_cast<_T>(a_ptr[offset]);
                    _T t1             = static_cast<_T>(b_ptr[offset]);
                    *(c_ptr + offset) = performOperation<_T, Operation, order_flag>(t0, t1);
                }
                else
                {
                    float t0          = static_cast<float>(a_ptr[offset]);
                    float t1          = static_cast<float>(b_ptr[offset]);
                    float t2          = performOperation<float, Operation, order_flag>(t0, t1);
                    *(c_ptr + offset) = static_cast<_T>(t2);
                }
            }
        }
    }
}

template <class _T, int _rows, typename Operation, bool order_flag, class _T0, class _T1>
__global__ void operator_contiguous_kernel(const void* __restrict a,
                                           const void* __restrict b,
                                           void* __restrict c,
                                           const int M,
                                           const int N,
                                           const int K,
                                           bool types_match)
{
    constexpr uint32_t element_size    = sizeof(_T); // in bytes
    constexpr uint32_t elements_in_16B = 16 / element_size;
    uint64_t idx                       = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
    uint32_t n_tiles                   = N / _rows;
    uint32_t k_tiles                   = K / elements_in_16B;
    if(idx < (uint64_t)M * n_tiles * k_tiles)
    {
        uint32_t ti        = idx / (k_tiles * n_tiles);
        uint64_t idx_block = idx % (k_tiles * n_tiles);
        uint32_t tj        = (idx_block / k_tiles) % n_tiles;
        uint32_t tk        = idx_block % k_tiles;
        for(int row = 0; row < _rows; row++)
        {
            uint64_t offset =
                (uint64_t)(tj + row * n_tiles) * K + tk * elements_in_16B + (uint64_t)ti * N * K;
            const _T0* pa = reinterpret_cast<const _T0*>(a) + offset;
            const _T1* pb = reinterpret_cast<const _T1*>(b) + offset;
            _T* pc        = reinterpret_cast<_T*>(c) + offset;
            for(int col = 0; col < elements_in_16B; col++)
            {
                const _T0* pfa = pa + col;
                const _T1* pfb = pb + col;
                _T* pfc        = pc + col;
                if(types_match)
                {
                    *pfc = performOperation<_T, Operation, order_flag>(static_cast<_T>(*pfa),
                                                                       static_cast<_T>(*pfb));
                }
                else
                {
                    float t0 = static_cast<float>(*pfa);
                    float t1 = static_cast<float>(*pfb);
                    float t2 = performOperation<float, Operation, order_flag>(t0, t1);
                    *pfc     = static_cast<_T>(t2);
                }
            }
        }
    }
}

template <class _T, typename Operation, class _T0, class _T1>
__global__ void operator_element_kernel(const void* __restrict a,
                                        const void* __restrict b,
                                        void* __restrict c,
                                        const int size,
                                        bool types_match)
{
    constexpr uint32_t element_size    = sizeof(_T); // in bytes
    constexpr uint32_t elements_in_16B = 16 / element_size;
    uint64_t idx                       = ((uint64_t)blockIdx.x * blockDim.x + threadIdx.x);
    if(idx * elements_in_16B < size)
    {
        int offset    = idx * elements_in_16B;
        const _T0* pa = reinterpret_cast<const _T0*>(a) + offset;
        const _T1* pb = reinterpret_cast<const _T1*>(b) + offset;
        _T* pc        = reinterpret_cast<_T*>(c) + offset;
#pragma unroll
        for(uint32_t v = 0; v < elements_in_16B; v++)
        {
            if(types_match)
            {
                pc[v] = performOperation<_T, Operation, true>(static_cast<_T>(pa[v]),
                                                              static_cast<_T>(pb[v]));
            }
            else
            {
                float t0 = static_cast<float>(pa[v]);
                float t1 = static_cast<float>(pb[v]);
                float t2 = performOperation<float, Operation, true>(t0, t1);
                pc[v]    = static_cast<_T>(t2);
            }
        }
    }
}

template <class _T,
          int _WG,
          int BIG_TILE_SIZE_N,
          int BIG_TILE_SIZE_K,
          int M_SWIZZLE,
          typename Operation,
          bool order_flag,
          class _T0,
          class _T1>
__global__ void operator_contiguous_big_tile_kernel(const void* __restrict a,
                                                    const void* __restrict b,
                                                    void* __restrict c,
                                                    const int N,
                                                    const int K,
                                                    bool types_match)
{
    constexpr uint32_t element_size    = sizeof(_T); // in bytes
    constexpr uint32_t elements_in_16B = 16 / element_size;

    union BLOCK_16B
    {
        _T e[elements_in_16B];
        __uint128_t ow;
    };

    // Round up processing to next full tile
    const uint32_t n_tiles        = (N + BIG_TILE_SIZE_N - 1) / BIG_TILE_SIZE_N;
    const uint32_t k_tiles        = (K + BIG_TILE_SIZE_K - 1) / BIG_TILE_SIZE_K;
    const uint32_t nk_tiles       = n_tiles * k_tiles;
    const uint32_t m_tiles        = gridDim.x / nk_tiles;
    const uint32_t m_tile_swizzle = blockIdx.x / nk_tiles / M_SWIZZLE * M_SWIZZLE;
    /// do m_swizzle when there are enough m_tiles
    const bool swizzle_m = m_tile_swizzle + M_SWIZZLE <= m_tiles;
    const uint32_t current_m =
        swizzle_m ? m_tile_swizzle + blockIdx.x % M_SWIZZLE : blockIdx.x / nk_tiles;

    const uint64_t stride_k      = N;
    const uint64_t out_stride_nk = N * K;

    const uint32_t current_nk =
        swizzle_m ? blockIdx.x / M_SWIZZLE % nk_tiles : blockIdx.x % nk_tiles;
    const uint32_t ti = current_nk / k_tiles;
    const uint32_t tj = current_nk % k_tiles;

    const uint32_t current_n_size = (ti == (n_tiles - 1) && (N % BIG_TILE_SIZE_N) != 0)
                                        ? (N % BIG_TILE_SIZE_N)
                                        : BIG_TILE_SIZE_N;
    const uint32_t current_k_size = (tj == (k_tiles - 1) && (K % BIG_TILE_SIZE_K) != 0)
                                        ? (K % BIG_TILE_SIZE_K)
                                        : BIG_TILE_SIZE_K;

    // use 128bit load&store whenever possible
    if(current_n_size % 8 == 0 && current_k_size % elements_in_16B == 0)
    {
        // Copy full tile with large loads
        constexpr uint32_t row_bytes_wr    = BIG_TILE_SIZE_N;
        constexpr uint32_t vmem_per_row_wr = row_bytes_wr / elements_in_16B;
        constexpr uint32_t rows_per_wg_wr  = _WG / vmem_per_row_wr;
        constexpr uint32_t wr_per_row      = BIG_TILE_SIZE_K / rows_per_wg_wr;
        // Make sure WG isn't too large
        static_assert(wr_per_row >= 1);

        const _T0* pa = (const _T0*)a + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                        current_m * out_stride_nk;
        const _T1* pb = (const _T1*)b + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                        current_m * out_stride_nk;
        const _T* pc = (const _T*)c + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                       current_m * out_stride_nk;
#pragma unroll
        for(uint32_t t = 0; t < wr_per_row; t++)
        {
            uint32_t col = threadIdx.x % vmem_per_row_wr;
            uint32_t row = threadIdx.x / vmem_per_row_wr + t * rows_per_wg_wr;
            if(col * elements_in_16B < current_n_size && row < current_k_size)
            {
                BLOCK_16B d, f;
                uint64_t offset = row * stride_k + col * elements_in_16B;
                if(types_match)
                {
                    const __uint128_t* pfa = (const __uint128_t*)(pa + offset);
                    const __uint128_t* pfb = (const __uint128_t*)(pb + offset);
                    f.ow                   = *pfa;
                    d.ow                   = *pfb;
#pragma unroll
                    for(uint32_t i = 0; i < elements_in_16B; i++)
                    {
                        d.e[i] = performOperation<_T, Operation, order_flag>(
                            static_cast<_T>(f.e[i]), static_cast<_T>(d.e[i]));
                    }
                    __uint128_t* pfc = (__uint128_t*)(pc + offset);
                    *pfc             = d.ow;
                }
                else
                {
                    const _T0* pfa = (const _T0*)(pa + offset);
                    const _T1* pfb = (const _T1*)(pb + offset);
#pragma unroll
                    for(uint32_t i = 0; i < elements_in_16B; i++)
                    {
                        float a = static_cast<float>(pfa[i]);
                        float b = static_cast<float>(pfb[i]);
                        float c = performOperation<float, Operation, order_flag>(a, b);
                        d.e[i]  = static_cast<_T>(c);
                    }
                    __uint128_t* pfc = (__uint128_t*)(pc + offset);
                    *pfc             = d.ow;
                }
            }
        }
    }
    else
    {
        // Copy full tile with large loads
        constexpr uint32_t row_bytes_wr    = BIG_TILE_SIZE_N;
        constexpr uint32_t vmem_per_row_wr = BIG_TILE_SIZE_N;
        constexpr uint32_t rows_per_wg_wr  = _WG / vmem_per_row_wr;
        constexpr uint32_t wr_per_row      = BIG_TILE_SIZE_K / rows_per_wg_wr;
        const _T0* pa = (const _T0*)a + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                        current_m * out_stride_nk;
        const _T1* pb = (const _T1*)b + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                        current_m * out_stride_nk;
        const _T* pc = (const _T*)c + tj * BIG_TILE_SIZE_K * stride_k + ti * row_bytes_wr +
                       current_m * out_stride_nk;
#pragma unroll
        for(uint32_t t = 0; t < wr_per_row; t++)
        {
            uint32_t col = threadIdx.x % vmem_per_row_wr;
            uint32_t row = threadIdx.x / vmem_per_row_wr + t * rows_per_wg_wr;
            if(col < current_n_size && row < current_k_size)
            {
                uint64_t offset = row * stride_k + col;
                const _T0* pfa  = (const _T0*)(pa + offset);
                const _T1* pfb  = (const _T1*)(pb + offset);
                _T* pfc         = (_T*)(pc + offset);
                if(types_match)
                {
                    *pfc = performOperation<_T, Operation, order_flag>(static_cast<_T>(*pfa),
                                                                       static_cast<_T>(*pfb));
                }
                else
                {
                    float a = static_cast<float>(*pfa);
                    float b = static_cast<float>(*pfb);
                    float c = performOperation<float, Operation, order_flag>(a, b);
                    *pfc    = static_cast<_T>(c);
                }
            }
        }
    }
}
} // namespace aiter

__inline__ std::vector<int64_t> broadcastShapes(const torch::Tensor& tensor1,
                                                const torch::Tensor& tensor2)
{
    auto shape1 = tensor1.sizes().vec();
    auto shape2 = tensor2.sizes().vec();

    int64_t max_dim = std::max(shape1.size(), shape2.size());
    shape1.insert(shape1.begin(), max_dim - shape1.size(), 1);
    shape2.insert(shape2.begin(), max_dim - shape2.size(), 1);

    std::vector<int64_t> result_shape(max_dim, 1);
    for(int64_t i = 0; i < max_dim; ++i)
    {
        if(shape1[i] == 1)
        {
            result_shape[i] = shape2[i];
        }
        else if(shape2[i] == 1)
        {
            result_shape[i] = shape1[i];
        }
        else if(shape1[i] == shape2[i])
        {
            result_shape[i] = shape1[i];
        }
        else
        {
            throw std::invalid_argument("Incompatible shapes for binary operator.");
        }
    }

    return result_shape;
}

template <int pattern, typename Operation, class _T0, class _T1>
struct BinaryOperationPattern;

// PATTERN_TRANSPOSE
template <typename Operation, class _T0, class _T1>
struct BinaryOperationPattern<1, Operation, _T0, _T1>
{
    static void
    apply(torch::Tensor& input, torch::Tensor& other, torch::Tensor& output, bool order_flag)
    {
        int dim     = input.dim();
        auto shape  = output.sizes().vec();
        void* buf_a = reinterpret_cast<void*>(input.data_ptr());
        void* buf_b = reinterpret_cast<void*>(other.data_ptr());
        void* buf_c = reinterpret_cast<void*>(output.data_ptr());

        int num_elements = output.numel();
        int rem_dim_size = num_elements / (shape[dim - 2] * shape[dim - 1]);
        int M            = dim == 2 ? 1 : rem_dim_size;
        int N            = shape[dim - 2];
        int K            = shape[dim - 1];

        auto tensor_not_conti              = input.is_contiguous() ? other : input;
        int stride0                        = tensor_not_conti.stride(0);
        int stride2                        = tensor_not_conti.stride(2);
        constexpr uint32_t BIG_TILE_SIZE_N = 64;
        constexpr uint32_t BIG_TILE_SIZE_K = 64;
        constexpr uint32_t M_SWIZZLE       = 8;
        const int grid_x                   = M * ((N + BIG_TILE_SIZE_N - 1) / BIG_TILE_SIZE_N) *
                           ((K + BIG_TILE_SIZE_K - 1) / BIG_TILE_SIZE_K);
        const dim3 grid_dim(grid_x, 1, 1);
        const dim3 block_dim(256, 1, 1);
        const hipStream_t stream = at::hip::getCurrentHIPStream();
        bool types_match         = typeid(_T0) == typeid(_T1);

        if(order_flag)
        {
            VLLM_DISPATCH_FLOATING_TYPES(output.scalar_type(), "operator_tn_big_tile_kernel", [&] {
                aiter::operator_tn_big_tile_kernel<scalar_t,
                                                   256,
                                                   BIG_TILE_SIZE_N,
                                                   BIG_TILE_SIZE_K,
                                                   M_SWIZZLE,
                                                   Operation,
                                                   true,
                                                   _T0,
                                                   _T1><<<grid_dim, block_dim, 0, stream>>>(
                    buf_a, buf_b, buf_c, K, N, stride0, stride2, types_match);
            });
        }
        else
        {
            VLLM_DISPATCH_FLOATING_TYPES(output.scalar_type(), "operator_tn_big_tile_kernel", [&] {
                aiter::operator_tn_big_tile_kernel<scalar_t,
                                                   256,
                                                   BIG_TILE_SIZE_N,
                                                   BIG_TILE_SIZE_K,
                                                   M_SWIZZLE,
                                                   Operation,
                                                   false,
                                                   _T1,
                                                   _T0><<<grid_dim, block_dim, 0, stream>>>(
                    buf_b, buf_a, buf_c, K, N, stride0, stride2, types_match);
            });
        }
    }
};

// PATTERN_BROADCAST_0
template <typename Operation, class _T0, class _T1>
struct BinaryOperationPattern<2, Operation, _T0, _T1>
{
    static void
    apply(torch::Tensor& input, torch::Tensor& other, torch::Tensor& output, bool order_flag)
    {
        int dim    = input.dim();
        auto shape = output.sizes().vec();

        void* buf_a      = reinterpret_cast<void*>(input.data_ptr());
        void* buf_b      = reinterpret_cast<void*>(other.data_ptr());
        void* buf_c      = reinterpret_cast<void*>(output.data_ptr());
        int num_elements = output.numel();
        int rem_dim_size = num_elements / (shape[dim - 2] * shape[dim - 1]);
        int M            = dim != 3 ? 1 : rem_dim_size;
        int N            = shape[dim - 2];
        int K            = shape[dim - 1];
        if(dim == 4)
        {
            N = shape[0] * shape[1] * shape[2];
            K = shape[3];
        }
        const hipStream_t stream = at::hip::getCurrentHIPStream();
        bool types_match         = typeid(_T0) == typeid(_T1);

        const uint32_t rows = 8;
        int vec             = 16 / output.element_size();

        hipDevice_t dev;
        hipDeviceProp_t dev_prop;
        hipGetDevice(&dev);
        hipGetDeviceProperties(&dev_prop, dev);
        uint32_t num_cu = dev_prop.multiProcessorCount;

        bool bcast_k_dim = (input.dim() == 1 && input.size(0) == shape[2]) ||
                           (other.dim() == 1 && other.size(0) == shape[2]);
        bool bcast_scalar =
            (input.dim() == 1 && input.size(0) == 1) || (other.dim() == 1 && other.size(0) == 1);
        bool vec_unroll_able = num_elements % (rows * vec * 256) == 0 && shape[2] % vec == 0;

        //  (m,n,k), (k,)
        if(bcast_k_dim && vec_unroll_able && output.dim() == 3)
        {
            int grid_x = (num_elements / (rows * vec) + 256 - 1) / 256;
            int occupancy;
            auto kernel_ptr =
                aiter::operator_bcastK_unroll_vectorize_naive<_T0, rows, Operation, true, _T0, _T1>;
            hipOccupancyMaxActiveBlocksPerMultiprocessor(
                &occupancy, reinterpret_cast<const void*>(kernel_ptr), 256, 0);
            grid_x = grid_x < num_cu * occupancy ? grid_x : num_cu * occupancy;
            const dim3 block_dim(256, 1, 1);
            const dim3 grid_dim(grid_x, 1, 1);
            if(order_flag)
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcastK_unroll_vectorize_naive", [&] {
                        aiter::operator_bcastK_unroll_vectorize_naive<scalar_t,
                                                                      rows,
                                                                      Operation,
                                                                      true,
                                                                      _T0,
                                                                      _T1>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_a, buf_b, buf_c, M, N, K, types_match);
                    });
            }
            else
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcastK_unroll_vectorize_naive", [&] {
                        aiter::operator_bcastK_unroll_vectorize_naive<scalar_t,
                                                                      rows,
                                                                      Operation,
                                                                      false,
                                                                      _T1,
                                                                      _T0>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_b, buf_a, buf_c, shape[0], shape[1], shape[2], types_match);
                    });
            }
        }
        // (m, n, k), (1)
        else if(bcast_scalar && num_elements % (rows * vec * 256) == 0)
        {
            int grid_x = (num_elements / (rows * vec) + 256 - 1) / 256;
            int occupancy;
            auto kernel_ptr = aiter::
                operator_bcast_scalar_unroll_vectorize_naive<_T0, rows, Operation, true, _T0, _T1>;
            hipOccupancyMaxActiveBlocksPerMultiprocessor(
                &occupancy, reinterpret_cast<const void*>(kernel_ptr), 256, 0);
            grid_x = grid_x < num_cu * occupancy ? grid_x : num_cu * occupancy;
            const dim3 block_dim(256, 1, 1);
            const dim3 grid_dim(grid_x, 1, 1);
            if(order_flag)
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcast_scalar_unroll_vectorize_naive", [&] {
                        aiter::operator_bcast_scalar_unroll_vectorize_naive<scalar_t,
                                                                            rows,
                                                                            Operation,
                                                                            true,
                                                                            _T0,
                                                                            _T1>
                            <<<grid_dim, block_dim, 0, stream>>>(buf_a, buf_b, buf_c, num_elements);
                    });
            }
            else
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcast_scalar_unroll_vectorize_naive", [&] {
                        aiter::operator_bcast_scalar_unroll_vectorize_naive<scalar_t,
                                                                            rows,
                                                                            Operation,
                                                                            false,
                                                                            _T1,
                                                                            _T0>
                            <<<grid_dim, block_dim, 0, stream>>>(buf_b, buf_a, buf_c, num_elements);
                    });
            }
        }
        else if(N % rows == 0 && K % vec == 0)
        {
            constexpr uint32_t wg = 64;
            int grid_x            = (num_elements / (rows * vec) + wg - 1) / wg;
            const dim3 grid_dim(grid_x, 1, 1);
            const dim3 block_dim(wg, 1, 1);

            if(order_flag)
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcast_tile_kernel", [&] {
                        aiter::operator_bcast_tile_kernel<scalar_t, rows, Operation, true, _T0, _T1>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_a, buf_b, buf_c, M, N, K, types_match);
                    });
            }
            else
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcast_tile_kernel", [&] {
                        aiter::
                            operator_bcast_tile_kernel<scalar_t, rows, Operation, false, _T1, _T0>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_b, buf_a, buf_c, M, N, K, types_match);
                    });
            }
        }
        else
        {
            constexpr uint32_t BIG_TILE_SIZE_N = 64;
            constexpr uint32_t BIG_TILE_SIZE_K = 64;
            constexpr uint32_t M_SWIZZLE       = 8;
            const int grid_x                   = M * ((N + BIG_TILE_SIZE_N - 1) / BIG_TILE_SIZE_N) *
                               ((K + BIG_TILE_SIZE_K - 1) / BIG_TILE_SIZE_K);
            const dim3 grid_dim(grid_x, 1, 1);
            const dim3 block_dim(256, 1, 1);

            if(order_flag)
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcast_big_tile_kernel", [&] {
                        aiter::operator_bcast_big_tile_kernel<scalar_t,
                                                              256,
                                                              BIG_TILE_SIZE_N,
                                                              BIG_TILE_SIZE_K,
                                                              M_SWIZZLE,
                                                              Operation,
                                                              true,
                                                              _T0,
                                                              _T1>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_a, buf_b, buf_c, K, N, types_match);
                    });
            }
            else
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcast_big_tile_kernel", [&] {
                        aiter::operator_bcast_big_tile_kernel<scalar_t,
                                                              256,
                                                              BIG_TILE_SIZE_N,
                                                              BIG_TILE_SIZE_K,
                                                              M_SWIZZLE,
                                                              Operation,
                                                              false,
                                                              _T1,
                                                              _T0>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_b, buf_a, buf_c, K, N, types_match);
                    });
            }
        }
    }
};

// PATTERN_BROADCAST_1
template <typename Operation, class _T0, class _T1>
struct BinaryOperationPattern<3, Operation, _T0, _T1>
{
    static void
    apply(torch::Tensor& input, torch::Tensor& other, torch::Tensor& output, bool order_flag)
    {
        int dim     = input.dim();
        auto shape  = output.sizes().vec();
        void* buf_a = reinterpret_cast<void*>(input.data_ptr());
        void* buf_b = reinterpret_cast<void*>(other.data_ptr());
        void* buf_c = reinterpret_cast<void*>(output.data_ptr());

        int num_elements = output.numel();
        int rem_dim_size = num_elements / (shape[dim - 2] * shape[dim - 1]);
        int M            = dim == 2 ? 1 : rem_dim_size;
        int N            = shape[dim - 2];
        int K            = shape[dim - 1];

        // (m, n, p, q), (m, 1, p, q)
        if(dim == 4 && input.size(1) != other.size(1) && (input.size(1) == 1 || other.size(1)))
        {
            M = shape[0];
            N = shape[1];
            K = shape[2] * shape[3];
        }

        constexpr uint32_t BIG_TILE_SIZE_N = 64;
        constexpr uint32_t BIG_TILE_SIZE_K = 64;
        constexpr uint32_t M_SWIZZLE       = 8;
        int grid_x                         = M * ((N + BIG_TILE_SIZE_N - 1) / BIG_TILE_SIZE_N) *
                     ((K + BIG_TILE_SIZE_K - 1) / BIG_TILE_SIZE_K);
        const dim3 grid_dim(grid_x, 1, 1);
        const dim3 block_dim(256, 1, 1);
        const hipStream_t stream = at::hip::getCurrentHIPStream();
        bool types_match         = typeid(_T0) == typeid(_T1);
        constexpr int rows       = 8;
        int vec_size             = 16 / output.element_size();
        if(K % vec_size == 0 && num_elements % (256 * rows * vec_size) == 0)
        {
            hipDevice_t dev;
            hipDeviceProp_t dev_prop;
            hipGetDevice(&dev);
            hipGetDeviceProperties(&dev_prop, dev);
            uint32_t num_cu = dev_prop.multiProcessorCount;
            grid_x          = (num_elements / (rows * vec_size) + 256 - 1) / 256;
            int occupancy;
            auto kernel_ptr =
                aiter::operator_bcastM1K_unroll_kernel<_T0, rows, Operation, true, _T0, _T1>;
            hipOccupancyMaxActiveBlocksPerMultiprocessor(
                &occupancy, reinterpret_cast<const void*>(kernel_ptr), 256, 0);
            grid_x = grid_x < num_cu * occupancy ? grid_x : num_cu * occupancy;

            const dim3 grid_dim(grid_x, 1, 1);
            const dim3 block_dim(256, 1, 1);
            if(order_flag)
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcastM1K_unroll_kernel", [&] {
                        aiter::operator_bcastM1K_unroll_kernel<scalar_t,
                                                               rows,
                                                               Operation,
                                                               true,
                                                               _T0,
                                                               _T1>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_a, buf_b, buf_c, M, N, K, types_match);
                    });
            }
            else
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcastM1K_unroll_kernel", [&] {
                        aiter::operator_bcastM1K_unroll_kernel<scalar_t,
                                                               rows,
                                                               Operation,
                                                               false,
                                                               _T0,
                                                               _T1>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_b, buf_a, buf_c, M, N, K, types_match);
                    });
            }
        }
        else
        {
            if(order_flag)
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcast1_big_tile_kernel", [&] {
                        aiter::operator_bcast1_big_tile_kernel<scalar_t,
                                                               256,
                                                               BIG_TILE_SIZE_N,
                                                               BIG_TILE_SIZE_K,
                                                               M_SWIZZLE,
                                                               Operation,
                                                               true,
                                                               _T0,
                                                               _T1>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_a, buf_b, buf_c, K, N, types_match);
                    });
            }
            else
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcast1_big_tile_kernel", [&] {
                        aiter::operator_bcast1_big_tile_kernel<scalar_t,
                                                               256,
                                                               BIG_TILE_SIZE_N,
                                                               BIG_TILE_SIZE_K,
                                                               M_SWIZZLE,
                                                               Operation,
                                                               false,
                                                               _T1,
                                                               _T0>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_b, buf_a, buf_c, K, N, types_match);
                    });
            }
        }
    }
};

// PATTERN_BROADCAST_2
template <typename Operation, class _T0, class _T1>
struct BinaryOperationPattern<5, Operation, _T0, _T1>
{
    static void
    apply(torch::Tensor& input, torch::Tensor& other, torch::Tensor& output, bool order_flag)
    {
        int dim     = input.dim();
        auto shape  = output.sizes().vec();
        void* buf_a = reinterpret_cast<void*>(input.data_ptr());
        void* buf_b = reinterpret_cast<void*>(other.data_ptr());
        void* buf_c = reinterpret_cast<void*>(output.data_ptr());

        int bcast_dim = order_flag ? input.numel() / other.numel() : other.numel() / input.numel();
        int forward_dim = order_flag ? other.numel() : input.numel();

        int num_elements         = output.numel();
        int vec_size             = 16 / output.element_size();
        constexpr uint32_t row   = 8;
        const hipStream_t stream = at::hip::getCurrentHIPStream();
        bool types_match         = typeid(_T0) == typeid(_T1);

        // optimize kernel
        if(bcast_dim % vec_size == 0 && forward_dim % row == 0 &&
           num_elements % (256 * vec_size * row) == 0)
        {
            hipDevice_t dev;
            hipDeviceProp_t dev_prop;
            hipGetDevice(&dev);
            hipGetDeviceProperties(&dev_prop, dev);
            uint32_t num_cu = dev_prop.multiProcessorCount;

            constexpr uint32_t wg = 256;
            int grid_x            = (num_elements / (row * vec_size) + wg - 1) / wg;
            int occupancy;
            auto kernel_ptr =
                aiter::operator_bcastMN1_unroll_vec_naive<_T0, row, Operation, true, _T0, _T1>;
            hipOccupancyMaxActiveBlocksPerMultiprocessor(
                &occupancy, reinterpret_cast<const void*>(kernel_ptr), wg, 0);
            grid_x = grid_x < num_cu * occupancy ? grid_x : num_cu * occupancy;

            const dim3 grid_dim(grid_x, 1, 1);
            const dim3 block_dim(wg, 1, 1);
            if(order_flag)
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcastMN1_unroll_vec_naive", [&] {
                        aiter::operator_bcastMN1_unroll_vec_naive<scalar_t,
                                                                  row,
                                                                  Operation,
                                                                  true,
                                                                  _T0,
                                                                  _T1>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_a, buf_b, buf_c, forward_dim, bcast_dim, types_match);
                    });
            }
            else
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcastMN1_unroll_vec_naive", [&] {
                        aiter::operator_bcastMN1_unroll_vec_naive<scalar_t,
                                                                  row,
                                                                  Operation,
                                                                  false,
                                                                  _T0,
                                                                  _T1>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_b, buf_a, buf_c, forward_dim, bcast_dim, types_match);
                    });
            }
        }
        // fallback
        else
        {
            const dim3 block_dim(256, 1, 1);
            const dim3 grid_dim((num_elements + 256 - 1) / 256, 1, 1);
            if(order_flag)
            {
                VLLM_DISPATCH_FLOATING_TYPES(output.scalar_type(), "operator_bcastMN1_naive", [&] {
                    aiter::operator_bcastMN1_naive<scalar_t, row, Operation, true, _T0, _T1>
                        <<<grid_dim, block_dim, 0, stream>>>(
                            buf_a, buf_b, buf_c, forward_dim, bcast_dim, types_match);
                });
            }
            else
            {
                VLLM_DISPATCH_FLOATING_TYPES(output.scalar_type(), "operator_bcastMN1_naive", [&] {
                    aiter::operator_bcastMN1_naive<scalar_t, row, Operation, false, _T0, _T1>
                        <<<grid_dim, block_dim, 0, stream>>>(
                            buf_b, buf_a, buf_c, forward_dim, bcast_dim, types_match);
                });
            }
        }
    }
};

// PATTERN_BROADCAST_3
template <typename Operation, class _T0, class _T1>
struct BinaryOperationPattern<6, Operation, _T0, _T1>
{
    static void
    apply(torch::Tensor& input, torch::Tensor& other, torch::Tensor& output, bool order_flag)
    {
        int dim     = output.dim();
        auto shape  = output.sizes().vec();
        void* buf_a = reinterpret_cast<void*>(input.data_ptr());
        void* buf_b = reinterpret_cast<void*>(other.data_ptr());
        void* buf_c = reinterpret_cast<void*>(output.data_ptr());

        int m = output.size(0);
        int n = output.size(1);
        int k = output.size(2);

        int num_elements         = output.numel();
        int vec_size             = 16 / output.element_size();
        constexpr uint32_t row   = 8;
        const hipStream_t stream = at::hip::getCurrentHIPStream();
        bool types_match         = typeid(_T0) == typeid(_T1);

        // optimize kernel
        if(k % vec_size == 0 && num_elements % (256 * vec_size * row) == 0)
        {
            hipDevice_t dev;
            hipDeviceProp_t dev_prop;
            hipGetDevice(&dev);
            hipGetDeviceProperties(&dev_prop, dev);
            uint32_t num_cu = dev_prop.multiProcessorCount;

            constexpr uint32_t wg = 256;
            int grid_x            = (num_elements / (row * vec_size) + wg - 1) / wg;
            int occupancy;
            auto kernel_ptr =
                aiter::operator_bcast1N1_unroll_vec_naive<_T0, row, Operation, true, _T0, _T1>;
            hipOccupancyMaxActiveBlocksPerMultiprocessor(
                &occupancy, reinterpret_cast<const void*>(kernel_ptr), wg, 0);
            grid_x = grid_x < num_cu * occupancy ? grid_x : num_cu * occupancy;

            const dim3 grid_dim(grid_x, 1, 1);
            const dim3 block_dim(wg, 1, 1);
            if(order_flag)
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcast1N1_unroll_vec_naive", [&] {
                        aiter::operator_bcast1N1_unroll_vec_naive<scalar_t,
                                                                  row,
                                                                  Operation,
                                                                  true,
                                                                  _T0,
                                                                  _T1>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_a, buf_b, buf_c, m, n, k, types_match);
                    });
            }
            else
            {
                VLLM_DISPATCH_FLOATING_TYPES(
                    output.scalar_type(), "operator_bcast1N1_unroll_vec_naive", [&] {
                        aiter::operator_bcast1N1_unroll_vec_naive<scalar_t,
                                                                  row,
                                                                  Operation,
                                                                  false,
                                                                  _T0,
                                                                  _T1>
                            <<<grid_dim, block_dim, 0, stream>>>(
                                buf_b, buf_a, buf_c, m, n, k, types_match);
                    });
            }
        }
        // fallback
        else
        {
            const dim3 block_dim(256, 1, 1);
            const dim3 grid_dim((num_elements + 256 - 1) / 256, 1, 1);
            if(order_flag)
            {
                VLLM_DISPATCH_FLOATING_TYPES(output.scalar_type(), "operator_bcast1N1_naive", [&] {
                    aiter::operator_bcast1N1_naive<scalar_t, row, Operation, true, _T0, _T1>
                        <<<grid_dim, block_dim, 0, stream>>>(
                            buf_a, buf_b, buf_c, m, n, k, types_match);
                });
            }
            else
            {
                VLLM_DISPATCH_FLOATING_TYPES(output.scalar_type(), "operator_bcast1N1_naive", [&] {
                    aiter::operator_bcast1N1_naive<scalar_t, row, Operation, false, _T0, _T1>
                        <<<grid_dim, block_dim, 0, stream>>>(
                            buf_b, buf_a, buf_c, m, n, k, types_match);
                });
            }
        }
    }
};

// PATTERN_BROADCAST_4
// broadcast last 2 dim, (m, n, k) (m, 1, 1)
template <typename Operation, class _T0, class _T1>
struct BinaryOperationPattern<7, Operation, _T0, _T1>
{
    static void
    apply(torch::Tensor& input, torch::Tensor& other, torch::Tensor& output, bool order_flag)
    {
        int dim     = output.dim();
        auto shape  = output.sizes().vec();
        void* buf_a = reinterpret_cast<void*>(input.data_ptr());
        void* buf_b = reinterpret_cast<void*>(other.data_ptr());
        void* buf_c = reinterpret_cast<void*>(output.data_ptr());

        int m = output.size(0);
        int n = output.size(1);
        int k = output.size(2);

        int num_elements         = output.numel();
        int vec_size             = 16 / output.element_size();
        constexpr uint32_t row   = 8;
        const hipStream_t stream = at::hip::getCurrentHIPStream();
        bool types_match         = typeid(_T0) == typeid(_T1);

        // optimize kernel
        // if (k % vec_size == 0 && num_elements % (256 * vec_size * row) == 0)
        if(num_elements % vec_size == 0)
        {
            hipDevice_t dev;
            hipDeviceProp_t dev_prop;
            hipGetDevice(&dev);
            hipGetDeviceProperties(&dev_prop, dev);
            uint32_t num_cu = dev_prop.multiProcessorCount;

            constexpr uint32_t wg = 256;
            int tmp_row; // = row;
            int grid_x = (num_elements / (row * vec_size) + wg - 1) / wg;
            int occupancy;
            bool need_pad   = true;
            int padded_size = num_elements;

            auto ifNeedPad = [=](int tmp_row) {
                return num_elements % (wg * tmp_row * vec_size) != 0;
            };

            auto getPaddedSize = [=](int tmp_row) {
                int elem_num_per_block = wg * vec_size * tmp_row;
                return ((num_elements + elem_num_per_block - 1) / elem_num_per_block) *
                       elem_num_per_block;
            };

#define GET_PATTERN(_row)                                                                    \
    do                                                                                       \
    {                                                                                        \
        tmp_row = _row;                                                                      \
        grid_x  = (num_elements / (tmp_row * vec_size) + wg - 1) / wg;                       \
        auto kernel_ptr =                                                                    \
            aiter::operator_bcast1N1_unroll_vec_naive<_T0, _row, Operation, true, _T0, _T1>; \
        hipOccupancyMaxActiveBlocksPerMultiprocessor(                                        \
            &occupancy, reinterpret_cast<const void*>(kernel_ptr), wg, 0);                   \
        need_pad = ifNeedPad(tmp_row);                                                       \
        if(need_pad)                                                                         \
        {                                                                                    \
            padded_size = getPaddedSize(tmp_row);                                            \
        }                                                                                    \
    } while(0)

            if(grid_x > 512 && num_elements % (8 * vec_size) == 0) // row = 8
            {
                GET_PATTERN(8);
            }
            else if(grid_x > 256 && num_elements % (4 * vec_size) == 0) // row = 4
            {
                GET_PATTERN(4);
            }
            else if(grid_x > 128 && num_elements % (2 * vec_size) == 0) // row = 2
            {
                GET_PATTERN(2);
            }
            else // row = 1
            {
                GET_PATTERN(1);
            }
            grid_x = grid_x < num_cu * occupancy ? grid_x : num_cu * occupancy;

#define BCAST_CASE(case_row, normal_tensor, bcast_tensor, order_flag)                  \
    do                                                                                 \
    {                                                                                  \
    case case_row: {                                                                   \
        if(!need_pad)                                                                  \
        {                                                                              \
            VLLM_DISPATCH_FLOATING_TYPES(                                              \
                output.scalar_type(), "operator_bcastN11_unroll_vec_naive", [&] {      \
                    aiter::operator_bcastN11_unroll_vec_naive<scalar_t,                \
                                                              case_row,                \
                                                              Operation,               \
                                                              order_flag,              \
                                                              _T0,                     \
                                                              _T1>                     \
                        <<<grid_dim, block_dim, 0, stream>>>(                          \
                            normal_tensor, bcast_tensor, buf_c, m, n, k, types_match); \
                });                                                                    \
        }                                                                              \
        else                                                                           \
        {                                                                              \
            VLLM_DISPATCH_FLOATING_TYPES(                                              \
                output.scalar_type(), "operator_bcastN11_unroll_vec_pad", [&] {        \
                    aiter::operator_bcastN11_unroll_vec_pad<scalar_t,                  \
                                                            case_row,                  \
                                                            Operation,                 \
                                                            order_flag,                \
                                                            _T0,                       \
                                                            _T1>                       \
                        <<<grid_dim, block_dim, 0, stream>>>(normal_tensor,            \
                                                             bcast_tensor,             \
                                                             buf_c,                    \
                                                             m,                        \
                                                             n,                        \
                                                             k,                        \
                                                             padded_size,              \
                                                             types_match);             \
                });                                                                    \
        }                                                                              \
        return;                                                                        \
    }                                                                                  \
    } while(0)

            const dim3 grid_dim(grid_x);
            const dim3 block_dim(wg);
            if(order_flag)
            {
                switch(tmp_row)
                {
                    BCAST_CASE(8, buf_a, buf_b, true);
                    BCAST_CASE(4, buf_a, buf_b, true);
                    BCAST_CASE(2, buf_a, buf_b, true);
                    BCAST_CASE(1, buf_a, buf_b, true);
                }
            }
            else
            {
                switch(tmp_row)
                {
                    BCAST_CASE(8, buf_b, buf_a, false);
                    BCAST_CASE(4, buf_b, buf_a, false);
                    BCAST_CASE(2, buf_b, buf_a, false);
                    BCAST_CASE(1, buf_b, buf_a, false);
                }
            }
        }
        // fallback
        else
        {
            const dim3 block_dim(256, 1, 1);
            const dim3 grid_dim((num_elements + 256 - 1) / 256, 1, 1);
            if(order_flag)
            {
                VLLM_DISPATCH_FLOATING_TYPES(output.scalar_type(), "operator_bcastN11_naive", [&] {
                    aiter::operator_bcastN11_naive<scalar_t, row, Operation, true, _T0, _T1>
                        <<<grid_dim, block_dim, 0, stream>>>(
                            buf_a, buf_b, buf_c, m, n, k, types_match);
                });
            }
            else
            {
                VLLM_DISPATCH_FLOATING_TYPES(output.scalar_type(), "operator_bcastN11_naive", [&] {
                    aiter::operator_bcastN11_naive<scalar_t, row, Operation, false, _T0, _T1>
                        <<<grid_dim, block_dim, 0, stream>>>(
                            buf_b, buf_a, buf_c, m, n, k, types_match);
                });
            }
        }
    }
};

// PATTERN_CONTIGUOUS
template <typename Operation, class _T0, class _T1>
struct BinaryOperationPattern<4, Operation, _T0, _T1>
{
    static void
    apply(torch::Tensor& input, torch::Tensor& other, torch::Tensor& output, bool order_flag)
    {
        int dim    = input.dim();
        auto shape = output.sizes().vec();

        const uint32_t rows = 8;
        void* buf_a         = reinterpret_cast<void*>(input.data_ptr());
        void* buf_b         = reinterpret_cast<void*>(other.data_ptr());
        void* buf_c         = reinterpret_cast<void*>(output.data_ptr());
        int num_elements    = output.numel();
        int rem_dim_size    = 1;
        int M, N, K;
        if(dim == 1)
        {
            M = 1;
            N = input.numel() / 128;
            K = 128;
        }
        else
        {
            for(int i = 0; i < dim - 2; ++i)
            {
                rem_dim_size *= shape[i];
            }
            M = dim == 3 ? shape[0] : rem_dim_size;
            N = shape[dim - 2];
            K = shape[dim - 1];
            if(N < rows)
            {
                K = N * K;
                N = M;
                M = 1;
            }
        }

        const hipStream_t stream = at::hip::getCurrentHIPStream();
        bool types_match         = typeid(_T0) == typeid(_T1);
        int vec                  = 16 / output.element_size();
        hipDevice_t dev;
        hipDeviceProp_t dev_prop;
        hipGetDevice(&dev);
        hipGetDeviceProperties(&dev_prop, dev);
        uint32_t num_cu = dev_prop.multiProcessorCount;

        if(num_elements % vec == 0 && num_elements < num_cu * 256 * vec)
        {
            constexpr uint32_t wg = 256;
            const int grid_x      = (num_elements / vec + wg - 1) / wg;
            const dim3 grid_dim(grid_x, 1, 1);
            const dim3 block_dim(wg, 1, 1);
            VLLM_DISPATCH_FLOATING_TYPES(output.scalar_type(), "operator_element_kernel", [&] {
                aiter::operator_element_kernel<scalar_t, Operation, _T0, _T1>
                    <<<grid_dim, block_dim, 0, stream>>>(
                        buf_a, buf_b, buf_c, num_elements, types_match);
            });
        }
        else if(num_elements % (rows * vec * 256) == 0)
        {
            constexpr uint32_t wg = 256;
            int grid_x            = (num_elements / (rows * vec) + wg - 1) / wg;
            int occupancy;
            auto kernel_ptr =
                aiter::operator_contiguous_kernel_naive<_T0, rows, Operation, true, _T0, _T1>;
            hipOccupancyMaxActiveBlocksPerMultiprocessor(
                &occupancy, reinterpret_cast<const void*>(kernel_ptr), wg, 0);
            grid_x = grid_x < num_cu * occupancy ? grid_x : num_cu * occupancy;
            const dim3 grid_dim(grid_x, 1, 1);
            const dim3 block_dim(wg, 1, 1);
            VLLM_DISPATCH_FLOATING_TYPES(
                output.scalar_type(), "operator_contiguous_kernel_naive", [&] {
                    aiter::
                        operator_contiguous_kernel_naive<scalar_t, rows, Operation, true, _T0, _T1>
                        <<<grid_dim, block_dim, 0, stream>>>(
                            buf_a, buf_b, buf_c, num_elements, types_match);
                });
        }
        else if(N % rows == 0 && K % vec == 0)
        {
            constexpr uint32_t wg = 256;
            const int grid_x      = (num_elements / (rows * vec) + wg - 1) / wg;
            const dim3 grid_dim(grid_x, 1, 1);
            const dim3 block_dim(wg, 1, 1);

            VLLM_DISPATCH_FLOATING_TYPES(output.scalar_type(), "operator_contiguous_kernel", [&] {
                aiter::operator_contiguous_kernel<scalar_t, rows, Operation, true, _T0, _T1>
                    <<<grid_dim, block_dim, 0, stream>>>(buf_a, buf_b, buf_c, M, N, K, types_match);
            });
        }
        else
        {
            constexpr uint32_t wg              = 256;
            constexpr uint32_t BIG_TILE_SIZE_N = 64;
            constexpr uint32_t BIG_TILE_SIZE_K = 64;
            constexpr uint32_t M_SWIZZLE       = 8;
            const int grid_x                   = M * ((N + BIG_TILE_SIZE_N - 1) / BIG_TILE_SIZE_N) *
                               ((K + BIG_TILE_SIZE_K - 1) / BIG_TILE_SIZE_K);
            const dim3 grid_dim(grid_x, 1, 1);
            const dim3 block_dim(wg, 1, 1);

            VLLM_DISPATCH_FLOATING_TYPES(
                output.scalar_type(), "operator_contiguous_big_tile_kernel", [&] {
                    aiter::operator_contiguous_big_tile_kernel<scalar_t,
                                                               256,
                                                               BIG_TILE_SIZE_N,
                                                               BIG_TILE_SIZE_K,
                                                               M_SWIZZLE,
                                                               Operation,
                                                               true,
                                                               _T0,
                                                               _T1>
                        <<<grid_dim, block_dim, 0, stream>>>(
                            buf_a, buf_b, buf_c, K, N, types_match);
                });
        }
    }
};

template <int pattern, typename Operation, class _T0, class _T1>
void binary_operation_process(torch::Tensor& input,
                              torch::Tensor& other,
                              torch::Tensor& output,
                              bool order_flag)
{
    BinaryOperationPattern<pattern, Operation, _T0, _T1>::apply(input, other, output, order_flag);
}

void binary_op_dispatch(const std::string& op_type,
                        torch::Tensor& input,
                        torch::Tensor& other,
                        torch::Tensor& output);
