/*
 * Adapted from
 * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc7/cpp/tensorrt_llm/kernels/tinygemm2/tinygemm2_kernel.cuh
 * Copyright (c) 2025, The vLLM team.
 * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
 * All rights reserved. SPDX-License-Identifier: Apache-2.0
 *
 * 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.
 */

#include "cuda_bf16.h"
#include <stdint.h>
#include <stdio.h>
#include <vector>

#include "cuda_pipeline.h"
#include <cuda.h>
#include <cuda/barrier>
#include <cuda/std/utility>
#include <cuda_runtime.h>

using barrier = cuda::barrier<cuda::thread_scope_block>;
namespace cde = cuda::device::experimental;
namespace ptx = cuda::ptx;

#define gpuErrChk(ans)                    \
  {                                       \
    gpuAssert((ans), __FILE__, __LINE__); \
  }

inline void gpuAssert(cudaError_t code, char const* file, int line,
                      bool abort = true) {
  if (code != cudaSuccess) {
    fprintf(stderr, "GPUassert: %s %s %d\n", cudaGetErrorString(code), file,
            line);
    if (abort) {
      throw std::runtime_error(cudaGetErrorString(code));
    }
  }
}

#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
__device__ uint64_t gclock64() {
  unsigned long long int rv;
  asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(rv));
  return rv;
}

__device__ void ldmatrix(__nv_bfloat16 rv[2], uint32_t smem_ptr) {
  int dst;
  asm volatile("ldmatrix.sync.aligned.x1.m8n8.shared.b16 {%0}, [%1];\n"
               : "=r"(dst)
               : "r"(smem_ptr));
  int* rvi = reinterpret_cast<int*>(&rv[0]);
  rvi[0] = dst;
}

__device__ void ldmatrix2(__nv_bfloat16 rv[4], uint32_t smem_ptr) {
  int x, y;
  asm volatile("ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];\n"
               : "=r"(x), "=r"(y)
               : "r"(smem_ptr));

  int* rvi = reinterpret_cast<int*>(&rv[0]);
  rvi[0] = x;
  rvi[1] = y;
}

__device__ void ldmatrix4(__nv_bfloat16 rv[8], uint32_t smem_ptr) {
  int x, y, z, w;
  asm volatile(
      "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];"
      : "=r"(x), "=r"(y), "=r"(z), "=r"(w)
      : "r"(smem_ptr));
  int* rvi = reinterpret_cast<int*>(&rv[0]);
  rvi[0] = x;
  rvi[1] = y;
  rvi[2] = z;
  rvi[3] = w;
}

__device__ void HMMA_1688(float d[4], __nv_bfloat16 a[4], __nv_bfloat16 b[2],
                          float c[4]) {
  uint32_t const* A = reinterpret_cast<uint32_t const*>(&a[0]);
  uint32_t const* B = reinterpret_cast<uint32_t const*>(&b[0]);
  float const* C = reinterpret_cast<float const*>(&c[0]);
  float* D = reinterpret_cast<float*>(&d[0]);

  asm volatile(
      "mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 "
      "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
      : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3])
      : "r"(A[0]), "r"(A[1]), "r"(B[0]), "f"(C[0]), "f"(C[1]), "f"(C[2]),
        "f"(C[3]));
}

__device__ void HMMA_16816(float d[4], __nv_bfloat16 a[8], __nv_bfloat16 b[4],
                           float c[4]) {
  uint32_t const* A = reinterpret_cast<uint32_t const*>(&a[0]);
  uint32_t const* B = reinterpret_cast<uint32_t const*>(&b[0]);
  float const* C = reinterpret_cast<float const*>(&c[0]);
  float* D = reinterpret_cast<float*>(&d[0]);

  asm volatile(
      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
      "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
      : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3])
      : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]),
        "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]));
}

__device__ void bar_wait(uint32_t bar_ptr, int phase) {
  asm volatile(
      "{\n"
      ".reg .pred                P1;\n"
      "LAB_WAIT:\n"
      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
      "@P1                       bra.uni DONE;\n"
      "bra.uni                   LAB_WAIT;\n"
      "DONE:\n"
      "}\n" ::"r"(bar_ptr),
      "r"(phase));
}

__device__ bool bar_try_wait(uint32_t bar_ptr, int phase) {
  uint32_t success;
  #ifdef INTERNAL
  asm volatile(".pragma \"set knob DontInsertYield\";\n" : : : "memory");
  #endif
  asm volatile(
      "{\n\t"
      ".reg .pred P1; \n\t"
      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t"
      "selp.b32 %0, 1, 0, P1; \n\t"
      "}"
      : "=r"(success)
      : "r"(bar_ptr), "r"(phase));
  return success;
}

__device__ uint32_t elect_one_sync() {
  uint32_t pred = 0;
  uint32_t laneid = 0;
  asm volatile(
      "{\n"
      ".reg .b32 %%rx;\n"
      ".reg .pred %%px;\n"
      "     elect.sync %%rx|%%px, %2;\n"
      "@%%px mov.s32 %1, 1;\n"
      "     mov.s32 %0, %%rx;\n"
      "}\n"
      : "+r"(laneid), "+r"(pred)
      : "r"(0xFFFFFFFF));
  return pred;
}
#endif

struct Profile {
  uint64_t start;
  uint64_t weight_load_start;
  uint64_t act_load_start;
  uint64_t compute_start;
  uint64_t complete;
};

template <int WARP_TILE_M, int TILE_M, int TILE_N, int TILE_K, int STAGES,
          int STAGE_UNROLL, bool PROFILE>
__global__ __launch_bounds__(384, 1) void gpt_oss_router_gemm_kernel(
    __nv_bfloat16* output, __nv_bfloat16* weights, __nv_bfloat16* activations,
    __nv_bfloat16* bias, int M, int N, int K,
    const __grid_constant__ CUtensorMap weight_map,
    const __grid_constant__ CUtensorMap activation_map,
    Profile* profile = nullptr) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)

  if (PROFILE && threadIdx.x == 0 && blockIdx.y == 0)
    profile[blockIdx.x].start = gclock64();

  extern __shared__ __align__(128) char smem[];

  __nv_bfloat16* sh_weights = (__nv_bfloat16*)&smem[0];
  __nv_bfloat16* sh_activations =
      (__nv_bfloat16*)&smem[STAGES * STAGE_UNROLL * TILE_M * TILE_K *
                            sizeof(__nv_bfloat16)];

  #pragma nv_diag_suppress static_var_with_dynamic_init
  __shared__ barrier bar_wt_ready[STAGES];
  __shared__ barrier bar_act_ready[STAGES];
  __shared__ barrier bar_data_consumed[STAGES];

  __shared__ float4 reduction_buffer[128];

  __shared__ nv_bfloat16 sh_bias[TILE_M];

  if (threadIdx.x == 0) {
    for (int i = 0; i < STAGES; i++) {
      init(&bar_wt_ready[i], 1);
      init(&bar_act_ready[i], 1);
      init(&bar_data_consumed[i], 32);
    }
    ptx::fence_proxy_async(ptx::space_shared);
    asm volatile("prefetch.tensormap [%0];"
                 :
                 : "l"(reinterpret_cast<uint64_t>(&weight_map))
                 : "memory");
    asm volatile("prefetch.tensormap [%0];"
                 :
                 : "l"(reinterpret_cast<uint64_t>(&activation_map))
                 : "memory");
  }
  __syncthreads();

  int warp_id = threadIdx.x / 32;
  int lane_id = threadIdx.x % 32;

  int phase = 0;

  int mib = blockIdx.x * TILE_M;
  int ni = blockIdx.y * TILE_N;

  float accum[4];
  for (int i = 0; i < 4; i++) accum[i] = 0.f;

  int const K_LOOPS_DMA =
      (K + 4 * TILE_K * STAGE_UNROLL - 1) / (4 * (TILE_K * STAGE_UNROLL));
  int const K_LOOPS_COMPUTE = K_LOOPS_DMA;

  // Data loading thread
  if (warp_id >= 4 && elect_one_sync()) {
    int stage = warp_id % 4;

    bool weight_warp = warp_id < 8;
    if (!weight_warp) {
      cudaGridDependencySynchronize();
      cudaTriggerProgrammaticLaunchCompletion();
    }

    for (int ki = 0; ki < K_LOOPS_DMA; ki++) {
      int k = (ki * 4 + (warp_id % 4)) * TILE_K * STAGE_UNROLL;

      uint64_t desc_ptr_wt = reinterpret_cast<uint64_t>(&weight_map);
      uint64_t desc_ptr_act = reinterpret_cast<uint64_t>(&activation_map);

      uint32_t bar_ptr_wt = __cvta_generic_to_shared(&bar_wt_ready[stage]);
      uint32_t bar_ptr_act = __cvta_generic_to_shared(&bar_act_ready[stage]);
      int bytes_wt = TILE_M * TILE_K * sizeof(__nv_bfloat16);
      int bytes_act = TILE_N * TILE_K * sizeof(__nv_bfloat16);

      bar_wait(__cvta_generic_to_shared(&bar_data_consumed[stage]), phase ^ 1);

      if (weight_warp)
        asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;"
                     :
                     : "r"(bar_ptr_wt), "r"(STAGE_UNROLL * bytes_wt));
      if (!weight_warp)
        asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;"
                     :
                     : "r"(bar_ptr_act), "r"(STAGE_UNROLL * bytes_act));

      if (PROFILE && blockIdx.y == 0 && ki == 0 && weight_warp)
        profile[blockIdx.x].weight_load_start = gclock64();
      if (PROFILE && blockIdx.y == 0 && ki == 0 && !weight_warp)
        profile[blockIdx.x].act_load_start = gclock64();

      for (int i = 0; i < STAGE_UNROLL; i++) {
        uint32_t smem_ptr_wt = __cvta_generic_to_shared(
            &sh_weights[(stage * STAGE_UNROLL + i) * TILE_M * TILE_K]);
        uint32_t crd0 = k + i * TILE_K;
        uint32_t crd1 = mib;
        if (weight_warp)
          asm volatile(
              "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_"
              "tx::bytes [%0], [%1, {%3,%4}], "
              "[%2];"
              :
              : "r"(smem_ptr_wt), "l"(desc_ptr_wt), "r"(bar_ptr_wt), "r"(crd0),
                "r"(crd1)
              : "memory");

        uint32_t smem_ptr_act = __cvta_generic_to_shared(
            &sh_activations[(stage * STAGE_UNROLL + i) * TILE_N * TILE_K]);
        crd0 = k + i * TILE_K;
        crd1 = ni;
        if (!weight_warp)
          asm volatile(
              "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_"
              "tx::bytes [%0], [%1, {%3,%4}], "
              "[%2];"
              :
              : "r"(smem_ptr_act), "l"(desc_ptr_act), "r"(bar_ptr_act),
                "r"(crd0), "r"(crd1)
              : "memory");
      }

      stage += 4;
      if (stage >= STAGES) {
        stage = warp_id % 4;
        phase ^= 1;
      }
    }
    // Wait for pending loads to be consumed before exiting, to avoid race
    for (int i = 0; i < (STAGES / 4) - 1; i++) {
      bar_wait(__cvta_generic_to_shared(&bar_data_consumed[stage]), phase ^ 1);
      stage += 4;
      if (stage >= STAGES) {
        stage = warp_id % 4;
        phase ^= 1;
      }
    }
  }
  // Compute threads
  else if (warp_id < 4) {
    // Sneak the bias load into the compute warps since they're just waiting for
    // stuff anyway
    if (threadIdx.x < TILE_M) sh_bias[threadIdx.x] = bias[mib + threadIdx.x];

    int stage = warp_id;

    int phase = 0;
    int lane_id_div8 = lane_id / 8;
    int lane_id_mod8 = lane_id % 8;

    int lane_row_offset_wt = (lane_id_div8 % 2) ? 8 : 0;
    int lane_col_offset_wt = (lane_id_div8 / 2) ? 1 : 0;

    int row_wt = lane_id_mod8 + lane_row_offset_wt;
    int row_act = lane_id_mod8;

    int row_offset_wt = (reinterpret_cast<uintptr_t>(sh_weights) / 128) % 8;
    int row_offset_act = row_offset_wt;

    uint32_t bar_ptr_wt = __cvta_generic_to_shared(&bar_wt_ready[stage]);
    uint32_t bar_ptr_act = __cvta_generic_to_shared(&bar_act_ready[stage]);

    bool weight_ready = bar_try_wait(bar_ptr_wt, phase);
    bool act_ready = bar_try_wait(bar_ptr_act, phase);

  #pragma unroll 2
    for (int ki = 0; ki < K_LOOPS_COMPUTE; ki++) {
      int next_stage = stage + 4;
      int next_phase = phase;
      if (next_stage >= STAGES) {
        next_stage = warp_id;
        next_phase ^= 1;
      }

      while (!weight_ready || !act_ready) {
        weight_ready = bar_try_wait(bar_ptr_wt, phase);
        act_ready = bar_try_wait(bar_ptr_act, phase);
      }

      if (PROFILE && blockIdx.y == 0 && threadIdx.x == 0 && ki == 0)
        profile[blockIdx.x].compute_start = gclock64();

      if (ki + 1 < K_LOOPS_COMPUTE) {
        weight_ready = bar_try_wait(
            __cvta_generic_to_shared(&bar_wt_ready[next_stage]), next_phase);
        act_ready = bar_try_wait(
            __cvta_generic_to_shared(&bar_act_ready[next_stage]), next_phase);
      }

  #pragma unroll
      for (int su = 0; su < STAGE_UNROLL; su++) {
        __nv_bfloat16* ptr_weights =
            &sh_weights[(stage * STAGE_UNROLL + su) * TILE_M * TILE_K];
        __nv_bfloat16* ptr_act =
            &sh_activations[(stage * STAGE_UNROLL + su) * TILE_N * TILE_K];

  #pragma unroll
        for (int kii = 0; kii < TILE_K / 16; kii++) {
          __nv_bfloat16 a[8];
          __nv_bfloat16 b[4];

          int col = 2 * kii + lane_col_offset_wt;
          int col_sw = ((row_wt + row_offset_wt) % 8) ^ col;

          ldmatrix4(a, __cvta_generic_to_shared(
                           &ptr_weights[row_wt * TILE_K + col_sw * 8]));

          col = 2 * kii + lane_id_div8;
          col_sw = ((row_act + row_offset_act) % 8) ^ col;

          ldmatrix2(b, __cvta_generic_to_shared(
                           &ptr_act[row_act * TILE_K + 8 * col_sw]));

          HMMA_16816(accum, a, b, accum);
        }
      }

      uint32_t bar_c = __cvta_generic_to_shared(&bar_data_consumed[stage]);
      asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" : : "r"(bar_c));

      stage = next_stage;
      phase = next_phase;
    }

    float4 accum4;
    accum4.x = accum[0];
    accum4.y = accum[1];
    accum4.z = accum[2];
    accum4.w = accum[3];
    reduction_buffer[threadIdx.x] = accum4;

    __syncthreads();

    if (warp_id == 0) {
      int mi = mib + warp_id * WARP_TILE_M;
      int tm = mi + lane_id / 4;
      int tn = ni + 2 * (lane_id % 4);

      float4 accum1 = reduction_buffer[32 + threadIdx.x];
      float4 accum2 = reduction_buffer[64 + threadIdx.x];
      float4 accum3 = reduction_buffer[96 + threadIdx.x];

      accum[0] = accum[0] + accum1.x + accum2.x + accum3.x;
      accum[1] = accum[1] + accum1.y + accum2.y + accum3.y;
      accum[2] = accum[2] + accum1.z + accum2.z + accum3.z;
      accum[3] = accum[3] + accum1.w + accum2.w + accum3.w;

      float bias_lo = __bfloat162float(sh_bias[tm - mib]);
      float bias_hi = __bfloat162float(sh_bias[tm + 8 - mib]);

      if (tn < N && tm < M)
        output[tn * M + tm] = __float2bfloat16(accum[0] + bias_lo);
      if (tn + 1 < N && tm < M)
        output[(tn + 1) * M + tm] = __float2bfloat16(accum[1] + bias_lo);
      if (tn < N && tm + 8 < M)
        output[tn * M + tm + 8] = __float2bfloat16(accum[2] + bias_hi);
      if (tn + 1 < N && tm + 8 < M)
        output[(tn + 1) * M + tm + 8] = __float2bfloat16(accum[3] + bias_hi);

      if (PROFILE && blockIdx.y == 0 && threadIdx.x == 0)
        profile[blockIdx.x].complete = gclock64();
    }
  }
#endif  // end if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
}
