#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 mhc_fused_tilelang_kernel(const float* comb_mix, const float* post_mix, const bfloat16_t* residual_in, bfloat16_t* residual_out, float* rp_out, const float* weight_t, const bfloat16_t* x_in, float* yp_out, int num_tokens, int split_k);
extern "C" __global__ void __launch_bounds__(256, 1) mhc_fused_tilelang_kernel(const float* comb_mix, const float* post_mix, const bfloat16_t* residual_in, bfloat16_t* residual_out, float* rp_out, const float* weight_t, const bfloat16_t* x_in, float* yp_out, int num_tokens, int split_k) {
  extern __shared__ __align__(1024) uchar buf_dyn_shmem[];
  float acc[2];
  float sqr[1];
  float pm[4];
  float cm[16];
  float new_r[4];
  float v = 0x0p+0f/*0.000000e+00*/;
  float v2 = 0x0p+0f/*0.000000e+00*/;
  float broadcast_var = 0x0p+0f/*0.000000e+00*/;
  *(float2*)(acc + 0) = make_float2(broadcast_var, broadcast_var);
  sqr[0] = 0x0p+0f/*0.000000e+00*/;
  cudaGridDependencySynchronize();
  if (((int)threadIdx.x) < 4) {
    ((float*)buf_dyn_shmem)[((int)threadIdx.x)] = post_mix[((((int64_t)((int)blockIdx.x)) * (int64_t)4) + ((int64_t)((int)threadIdx.x)))];
  }
  if (((int)threadIdx.x) < 16) {
    ((float*)buf_dyn_shmem)[(((int)threadIdx.x) + 4)] = comb_mix[((((int64_t)((int)blockIdx.x)) * (int64_t)16) + ((int64_t)((int)threadIdx.x)))];
  }
  __syncthreads();
  #pragma unroll
  for (int j = 0; j < 4; ++j) {
    pm[j] = ((float*)buf_dyn_shmem)[j];
  }
  #pragma unroll
  for (int j_1 = 0; j_1 < 4; ++j_1) {
    #pragma unroll
    for (int k = 0; k < 4; ++k) {
      cm[((k * 4) + j_1)] = ((float*)buf_dyn_shmem)[(((k * 4) + j_1) + 4)];
    }
  }
  for (int it = 0; it < ((4096 / split_k) >> 8); ++it) {
    #pragma unroll
    for (int j_2 = 0; j_2 < 4; ++j_2) {
      new_r[j_2] = (pm[j_2] * ((float)x_in[((((((int64_t)((int)blockIdx.x)) * (int64_t)4096) + (((int64_t)it) * (int64_t)256)) + (((int64_t)((int)blockIdx.z)) * ((int64_t)4096 / ((int64_t)split_k)))) + ((int64_t)((int)threadIdx.x)))]));
      #pragma unroll
      for (int k_1 = 0; k_1 < 4; ++k_1) {
        new_r[j_2] = (new_r[j_2] + (cm[((k_1 * 4) + j_2)] * ((float)residual_in[(((((((int64_t)((int)blockIdx.x)) * (int64_t)16384) + (((int64_t)k_1) * (int64_t)4096)) + (((int64_t)it) * (int64_t)256)) + (((int64_t)((int)blockIdx.z)) * ((int64_t)4096 / ((int64_t)split_k)))) + ((int64_t)((int)threadIdx.x)))])));
      }
    }
    if (((int)blockIdx.y) == 0) {
      #pragma unroll
      for (int j_3 = 0; j_3 < 4; ++j_3) {
        residual_out[(((((((int64_t)((int)blockIdx.x)) * (int64_t)16384) + (((int64_t)j_3) * (int64_t)4096)) + (((int64_t)it) * (int64_t)256)) + (((int64_t)((int)blockIdx.z)) * ((int64_t)4096 / ((int64_t)split_k)))) + ((int64_t)((int)threadIdx.x)))] = ((bfloat16_t)new_r[j_3]);
        sqr[0] = (sqr[0] + (new_r[j_3] * new_r[j_3]));
      }
    }
    #pragma unroll
    for (int n = 0; n < 2; ++n) {
      #pragma unroll
      for (int j_4 = 0; j_4 < 4; ++j_4) {
        acc[n] = (acc[n] + (weight_t[((((((((int64_t)((int)blockIdx.y)) * (int64_t)32768) + (((int64_t)n) * (int64_t)16384)) + (((int64_t)j_4) * (int64_t)4096)) + (((int64_t)it) * (int64_t)256)) + (((int64_t)((int)blockIdx.z)) * ((int64_t)4096 / ((int64_t)split_k)))) + ((int64_t)((int)threadIdx.x)))] * new_r[j_4]));
      }
    }
  }
  #pragma unroll
  for (int n_1 = 0; n_1 < 2; ++n_1) {
    acc[n_1] = tl::warp_reduce_sum(acc[n_1]);
  }
  if (((int)blockIdx.y) == 0) {
    sqr[0] = tl::warp_reduce_sum(sqr[0]);
  }
  if ((((int)threadIdx.x) % 32) == 0) {
    #pragma unroll
    for (int n_2 = 0; n_2 < 2; ++n_2) {
      ((float*)buf_dyn_shmem)[((((((int)threadIdx.x) >> 5) * 3) + n_2) + 20)] = acc[n_2];
    }
    if (((int)blockIdx.y) == 0) {
      ((float*)buf_dyn_shmem)[(((((int)threadIdx.x) >> 5) * 3) + 22)] = sqr[0];
    }
  }
  __syncthreads();
  if ((((int)threadIdx.x) >> 5) == 0) {
    if ((((int)threadIdx.x) & 31) < 2) {
      #pragma unroll
      for (int w = 0; w < 8; ++w) {
        v = (v + ((float*)buf_dyn_shmem)[(((w * 3) + (((int)threadIdx.x) & 31)) + 20)]);
      }
      yp_out[((((((int64_t)((int)blockIdx.x)) * (int64_t)24) + ((((int64_t)((int)blockIdx.z)) * ((int64_t)num_tokens)) * (int64_t)24)) + (((int64_t)((int)blockIdx.y)) * (int64_t)2)) + (((int64_t)((int)threadIdx.x)) & (int64_t)31))] = v;
    }
    if ((((int)blockIdx.y) == 0) && ((((int)threadIdx.x) % 32) == 0)) {
      #pragma unroll
      for (int w_1 = 0; w_1 < 8; ++w_1) {
        v2 = (v2 + ((float*)buf_dyn_shmem)[((w_1 * 3) + 22)]);
      }
      rp_out[((((int64_t)((int)blockIdx.z)) * ((int64_t)num_tokens)) + ((int64_t)((int)blockIdx.x)))] = v2;
    }
  }
  cudaTriggerProgrammaticLaunchCompletion();
}

