#include <tl_templates/cuda/gemm.h>
#include <tl_templates/cuda/copy.h>
#include <tl_templates/cuda/reduce.h>
#include <tl_templates/cuda/ldsm.h>
#include <tl_templates/cuda/threadblock_swizzle.h>
#include <tl_templates/cuda/debug.h>
#ifdef ENABLE_BF16
#include <tl_templates/cuda/cuda_bf16_fallbacks.cuh>
#endif

extern "C" __global__ void hc_head_fuse_tilelang_kernel(const float* fn, const float* hc_base, const float* hc_scale, bfloat16_t* out, const bfloat16_t* residual, int num_tokens);
extern "C" __global__ void __launch_bounds__(128, 1) hc_head_fuse_tilelang_kernel(const float* fn, const float* hc_base, const float* hc_scale, bfloat16_t* out, const bfloat16_t* residual, int num_tokens) {
  extern __shared__ __align__(1024) uchar buf_dyn_shmem[];
  float sqrsum_r[1];
  float mixes_r[4];
  bfloat16_t residual_local_cast[8];
  float x_local[8];
  float fn_local[8];
  float rsqrt_val[1];
  float xl[32];
  float ol[8];
  bfloat16_t xs_local_cast_1[8];
  bfloat16_t out_local_cast_2[8];
  bfloat16_t xs_local_cast_3[8];
  bfloat16_t out_local_cast_4[8];
  bfloat16_t xs_local_cast_5[8];
  bfloat16_t out_local_cast_6[8];
  cudaGridDependencySynchronize();
  sqrsum_r[0] = 0x0p+0f/*0.000000e+00*/;
  float broadcast_var = 0x0p+0f/*0.000000e+00*/;
  *(float4*)(mixes_r + 0) = make_float4(broadcast_var, broadcast_var, broadcast_var, broadcast_var);
  for (int m_c = 0; m_c < 4; ++m_c) {
    for (int i_h = 0; i_h < 4; ++i_h) {
      *(uint4*)(residual_local_cast + 0) = *(uint4*)(residual + ((((((int64_t)((int)blockIdx.x)) * (int64_t)16384) + (((int64_t)m_c) * (int64_t)4096)) + (((int64_t)i_h) * (int64_t)1024)) + (((int64_t)((int)threadIdx.x)) * (int64_t)8)));
      for (int i = 0; i < 2; ++i) {
        float4 __1;
        uint2 v_ = *(uint2*)(residual_local_cast + (i * 4));
        ((float2*)(&__1))[0] = __bfloat1622float2((reinterpret_cast<__nv_bfloat162*>(&v_))[0]);
        ((float2*)(&__1))[1] = __bfloat1622float2((reinterpret_cast<__nv_bfloat162*>(&v_))[1]);
        *(float4*)(x_local + (i * 4)) = __1;
      }
      #pragma unroll
      for (int i_1 = 0; i_1 < 8; ++i_1) {
        sqrsum_r[0] = (sqrsum_r[0] + (x_local[i_1] * x_local[i_1]));
      }
      #pragma unroll
      for (int m_m = 0; m_m < 4; ++m_m) {
        #pragma unroll
        for (int i_2 = 0; i_2 < 2; ++i_2) {
          *(float4*)(fn_local + (i_2 * 4)) = *(float4*)(fn + (((((m_m * 16384) + (m_c * 4096)) + (i_h * 1024)) + (((int)threadIdx.x) * 8)) + (i_2 * 4)));
        }
        #pragma unroll
        for (int i_3 = 0; i_3 < 8; ++i_3) {
          mixes_r[m_m] = (mixes_r[m_m] + (x_local[i_3] * fn_local[i_3]));
        }
      }
    }
  }
  sqrsum_r[0] = tl::AllReduce<tl::SumOp, 128, 1, 0, tl::NamedBarrier<128>>::run(sqrsum_r[0], (&(((float*)buf_dyn_shmem)[0])));
  __syncthreads();
  for (int __finred_0 = 0; __finred_0 < 4; ++__finred_0) {
    mixes_r[__finred_0] = tl::AllReduce<tl::SumOp, 128, 1, 0, tl::NamedBarrier<128>>::run(mixes_r[__finred_0], (&(((float*)buf_dyn_shmem)[0])));
  }
  rsqrt_val[0] = rsqrtf(((sqrsum_r[0] / 0x1p+14f/*1.638400e+04*/) + 0x1.0c6f7a0b5ed8dp-20f/*1.000000e-06*/));
  __syncthreads();
  if ((((int)threadIdx.x) >> 2) == 0) {
    ((float*)buf_dyn_shmem)[(((int)threadIdx.x) & 3)] = ((0x1p+0f/*1.000000e+00*/ / (0x1p+0f/*1.000000e+00*/ + expf((0x0p+0f/*0.000000e+00*/ - (((mixes_r[(((int)threadIdx.x) & 3)] * rsqrt_val[0]) * hc_scale[0]) + hc_base[(((int)threadIdx.x) & 3)]))))) + 0x1.0c6f7a0b5ed8dp-20f/*1.000000e-06*/);
  }
  __syncthreads();
  #pragma unroll
  for (int i_4 = 0; i_4 < 4; ++i_4) {
    tl::cp_async_gs<16>((&(((bfloat16_t*)buf_dyn_shmem)[(((i_4 * 1024) + (((int)threadIdx.x) * 8)) + 256)])), (&(residual[(((((int64_t)((int)blockIdx.x)) * (int64_t)16384) + (((int64_t)i_4) * (int64_t)4096)) + (((int64_t)((int)threadIdx.x)) * (int64_t)8))])));
  }
  tl::cp_async_commit();
  #pragma unroll
  for (int i_5 = 0; i_5 < 4; ++i_5) {
    tl::cp_async_gs<16>((&(((bfloat16_t*)buf_dyn_shmem)[(((i_5 * 1024) + (((int)threadIdx.x) * 8)) + 4352)])), (&(residual[((((((int64_t)((int)blockIdx.x)) * (int64_t)16384) + (((int64_t)i_5) * (int64_t)4096)) + (((int64_t)((int)threadIdx.x)) * (int64_t)8)) + (int64_t)1024)])));
  }
  tl::cp_async_commit();
  for (int i0_h = 0; i0_h < 2; ++i0_h) {
    tl::cp_async_wait<1>();
    __syncthreads();
    #pragma unroll
    for (int i_6 = 0; i_6 < 4; ++i_6) {
      *(uint4*)(xs_local_cast_1 + 0) = *(uint4*)(((bfloat16_t*)buf_dyn_shmem) + ((((i0_h * 4096) + (i_6 * 1024)) + (((int)threadIdx.x) * 8)) + 256));
      for (int vec = 0; vec < 2; ++vec) {
        float4 __2;
        uint2 v__1 = *(uint2*)(xs_local_cast_1 + (vec * 4));
        ((float2*)(&__2))[0] = __bfloat1622float2((reinterpret_cast<__nv_bfloat162*>(&v__1))[0]);
        ((float2*)(&__2))[1] = __bfloat1622float2((reinterpret_cast<__nv_bfloat162*>(&v__1))[1]);
        *(float4*)(xl + ((i_6 * 8) + (vec * 4))) = __2;
      }
    }
    __syncthreads();
    #pragma unroll
    for (int i_7 = 0; i_7 < 4; ++i_7) {
      tl::cp_async_gs<16>((&(((bfloat16_t*)buf_dyn_shmem)[((((i0_h * 4096) + (i_7 * 1024)) + (((int)threadIdx.x) * 8)) + 256)])), (&(residual[(((((((int64_t)((int)blockIdx.x)) * (int64_t)16384) + (((int64_t)i_7) * (int64_t)4096)) + (((int64_t)i0_h) * (int64_t)1024)) + (((int64_t)((int)threadIdx.x)) * (int64_t)8)) + (int64_t)2048)])));
    }
    tl::cp_async_commit();
    #pragma unroll
    for (int i_8 = 0; i_8 < 2; ++i_8) {
      float broadcast_var_1 = 0x0p+0f/*0.000000e+00*/;
      *(float4*)(ol + (i_8 * 4)) = make_float4(broadcast_var_1, broadcast_var_1, broadcast_var_1, broadcast_var_1);
    }
    __syncthreads();
    for (int i_hc = 0; i_hc < 4; ++i_hc) {
      float pre = ((float*)buf_dyn_shmem)[i_hc];
      #pragma unroll
      for (int i_9 = 0; i_9 < 8; ++i_9) {
        ol[i_9] = (ol[i_9] + (pre * xl[((i_hc * 8) + i_9)]));
      }
    }
    for (int i_10 = 0; i_10 < 2; ++i_10) {
      uint2 __3;
      float4 v__2 = *(float4*)(ol + (i_10 * 4));
      (reinterpret_cast<__nv_bfloat162*>(&__3))[0] = __float22bfloat162_rn(((float2*)(&v__2))[0]);
      (reinterpret_cast<__nv_bfloat162*>(&__3))[1] = __float22bfloat162_rn(((float2*)(&v__2))[1]);
      *(uint2*)(out_local_cast_2 + (i_10 * 4)) = __3;
    }
    *(uint4*)(out + (((((int64_t)((int)blockIdx.x)) * (int64_t)4096) + (((int64_t)i0_h) * (int64_t)1024)) + (((int64_t)((int)threadIdx.x)) * (int64_t)8))) = *(uint4*)(out_local_cast_2 + 0);
  }
  tl::cp_async_wait<1>();
  __syncthreads();
  #pragma unroll
  for (int i_11 = 0; i_11 < 4; ++i_11) {
    *(uint4*)(xs_local_cast_3 + 0) = *(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((i_11 * 1024) + (((int)threadIdx.x) * 8)) + 256));
    for (int vec_1 = 0; vec_1 < 2; ++vec_1) {
      float4 __4;
      uint2 v__3 = *(uint2*)(xs_local_cast_3 + (vec_1 * 4));
      ((float2*)(&__4))[0] = __bfloat1622float2((reinterpret_cast<__nv_bfloat162*>(&v__3))[0]);
      ((float2*)(&__4))[1] = __bfloat1622float2((reinterpret_cast<__nv_bfloat162*>(&v__3))[1]);
      *(float4*)(xl + ((i_11 * 8) + (vec_1 * 4))) = __4;
    }
  }
  #pragma unroll
  for (int i_12 = 0; i_12 < 2; ++i_12) {
    float broadcast_var_2 = 0x0p+0f/*0.000000e+00*/;
    *(float4*)(ol + (i_12 * 4)) = make_float4(broadcast_var_2, broadcast_var_2, broadcast_var_2, broadcast_var_2);
  }
  for (int i_hc_1 = 0; i_hc_1 < 4; ++i_hc_1) {
    float pre_1 = ((float*)buf_dyn_shmem)[i_hc_1];
    #pragma unroll
    for (int i_13 = 0; i_13 < 8; ++i_13) {
      ol[i_13] = (ol[i_13] + (pre_1 * xl[((i_hc_1 * 8) + i_13)]));
    }
  }
  for (int i_14 = 0; i_14 < 2; ++i_14) {
    uint2 __5;
    float4 v__4 = *(float4*)(ol + (i_14 * 4));
    (reinterpret_cast<__nv_bfloat162*>(&__5))[0] = __float22bfloat162_rn(((float2*)(&v__4))[0]);
    (reinterpret_cast<__nv_bfloat162*>(&__5))[1] = __float22bfloat162_rn(((float2*)(&v__4))[1]);
    *(uint2*)(out_local_cast_4 + (i_14 * 4)) = __5;
  }
  *(uint4*)(out + (((((int64_t)((int)blockIdx.x)) * (int64_t)4096) + (((int64_t)((int)threadIdx.x)) * (int64_t)8)) + (int64_t)2048)) = *(uint4*)(out_local_cast_4 + 0);
  tl::cp_async_wait<0>();
  __syncthreads();
  #pragma unroll
  for (int i_15 = 0; i_15 < 4; ++i_15) {
    *(uint4*)(xs_local_cast_5 + 0) = *(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((i_15 * 1024) + (((int)threadIdx.x) * 8)) + 4352));
    for (int vec_2 = 0; vec_2 < 2; ++vec_2) {
      float4 __6;
      uint2 v__5 = *(uint2*)(xs_local_cast_5 + (vec_2 * 4));
      ((float2*)(&__6))[0] = __bfloat1622float2((reinterpret_cast<__nv_bfloat162*>(&v__5))[0]);
      ((float2*)(&__6))[1] = __bfloat1622float2((reinterpret_cast<__nv_bfloat162*>(&v__5))[1]);
      *(float4*)(xl + ((i_15 * 8) + (vec_2 * 4))) = __6;
    }
  }
  #pragma unroll
  for (int i_16 = 0; i_16 < 2; ++i_16) {
    float broadcast_var_3 = 0x0p+0f/*0.000000e+00*/;
    *(float4*)(ol + (i_16 * 4)) = make_float4(broadcast_var_3, broadcast_var_3, broadcast_var_3, broadcast_var_3);
  }
  for (int i_hc_2 = 0; i_hc_2 < 4; ++i_hc_2) {
    float pre_2 = ((float*)buf_dyn_shmem)[i_hc_2];
    #pragma unroll
    for (int i_17 = 0; i_17 < 8; ++i_17) {
      ol[i_17] = (ol[i_17] + (pre_2 * xl[((i_hc_2 * 8) + i_17)]));
    }
  }
  for (int i_18 = 0; i_18 < 2; ++i_18) {
    uint2 __7;
    float4 v__6 = *(float4*)(ol + (i_18 * 4));
    (reinterpret_cast<__nv_bfloat162*>(&__7))[0] = __float22bfloat162_rn(((float2*)(&v__6))[0]);
    (reinterpret_cast<__nv_bfloat162*>(&__7))[1] = __float22bfloat162_rn(((float2*)(&v__6))[1]);
    *(uint2*)(out_local_cast_6 + (i_18 * 4)) = __7;
  }
  *(uint4*)(out + (((((int64_t)((int)blockIdx.x)) * (int64_t)4096) + (((int64_t)((int)threadIdx.x)) * (int64_t)8)) + (int64_t)3072)) = *(uint4*)(out_local_cast_6 + 0);
  cudaTriggerProgrammaticLaunchCompletion();
}

