/*
 * Adapted from SGLang's sgl-kernel implementation, which was adapted from
 * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
 * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
 *
 * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
 *
 * 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 <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/core/ScalarType.h>

#include "libtorch_stable/torch_utils.h"

#include <cuda_bf16.h>
#include <cuda_runtime.h>

#include <stdexcept>

namespace {

inline int getSMVersion() {
  auto* props = get_device_prop();
  return props->major * 10 + props->minor;
}

}  // namespace

static constexpr int DEFAULT_NUM_EXPERTS = 256;
static constexpr int KIMI_K2_NUM_EXPERTS = 384;
static constexpr int DEFAULT_HIDDEN_DIM = 7168;
static constexpr int GLM_5_HIDDEN_DIM = 6144;

template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
                                 cudaStream_t stream);

template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a,
                                T const* mat_b, cudaStream_t stream);

template <int kBegin, int kEnd, int kNumExperts, int kHiddenDim>
struct LoopUnroller {
  static void unroll_float_output(int num_tokens, float* output,
                                  __nv_bfloat16 const* input,
                                  __nv_bfloat16 const* weights,
                                  cudaStream_t stream) {
    if (num_tokens == kBegin) {
      invokeRouterGemmFloatOutput<__nv_bfloat16, kBegin, kNumExperts,
                                  kHiddenDim>(output, input, weights, stream);
    } else {
      LoopUnroller<kBegin + 1, kEnd, kNumExperts,
                   kHiddenDim>::unroll_float_output(num_tokens, output, input,
                                                    weights, stream);
    }
  }

  static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output,
                                 __nv_bfloat16 const* input,
                                 __nv_bfloat16 const* weights,
                                 cudaStream_t stream) {
    if (num_tokens == kBegin) {
      invokeRouterGemmBf16Output<__nv_bfloat16, kBegin, kNumExperts,
                                 kHiddenDim>(output, input, weights, stream);
    } else {
      LoopUnroller<kBegin + 1, kEnd, kNumExperts,
                   kHiddenDim>::unroll_bf16_output(num_tokens, output, input,
                                                   weights, stream);
    }
  }
};

template <int kEnd, int kNumExperts, int kHiddenDim>
struct LoopUnroller<kEnd, kEnd, kNumExperts, kHiddenDim> {
  static void unroll_float_output(int num_tokens, float* output,
                                  __nv_bfloat16 const* input,
                                  __nv_bfloat16 const* weights,
                                  cudaStream_t stream) {
    if (num_tokens == kEnd) {
      invokeRouterGemmFloatOutput<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(
          output, input, weights, stream);
    } else {
      throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16");
    }
  }

  static void unroll_bf16_output(int num_tokens, __nv_bfloat16* output,
                                 __nv_bfloat16 const* input,
                                 __nv_bfloat16 const* weights,
                                 cudaStream_t stream) {
    if (num_tokens == kEnd) {
      invokeRouterGemmBf16Output<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(
          output, input, weights, stream);
    } else {
      throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16");
    }
  }
};

void dsv3_router_gemm(
    torch::stable::Tensor& output,       // [num_tokens, num_experts]
    torch::stable::Tensor const& mat_a,  // [num_tokens, hidden_dim]
    torch::stable::Tensor const& mat_b   // [num_experts, hidden_dim]
) {
  STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);

  const int num_tokens = mat_a.size(0);
  const int num_experts = mat_b.size(0);
  const int hidden_dim = mat_a.size(1);

  STD_TORCH_CHECK(mat_a.size(1) == mat_b.size(1),
                  "mat_a and mat_b must have the same hidden_dim");
  STD_TORCH_CHECK(
      hidden_dim == DEFAULT_HIDDEN_DIM || hidden_dim == GLM_5_HIDDEN_DIM,
      "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
      " or hidden_dim=", GLM_5_HIDDEN_DIM, ", but got hidden_dim=", hidden_dim);
  STD_TORCH_CHECK(
      num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS,
      "Expected num_experts=", DEFAULT_NUM_EXPERTS,
      " or num_experts=", KIMI_K2_NUM_EXPERTS,
      ", but got num_experts=", num_experts);
  // KIMI_K2_NUM_EXPERTS is only instantiated for the default hidden_dim.
  STD_TORCH_CHECK(
      hidden_dim == DEFAULT_HIDDEN_DIM || num_experts == DEFAULT_NUM_EXPERTS,
      "hidden_dim=", GLM_5_HIDDEN_DIM,
      " only supports num_experts=", DEFAULT_NUM_EXPERTS,
      ", but got num_experts=", num_experts);
  STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
                  "currently num_tokens must be less than or equal to 16 for "
                  "router_gemm");
  STD_TORCH_CHECK(
      mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16,
      "mat_a must be bf16");
  STD_TORCH_CHECK(
      mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16,
      "mat_b must be bf16");
  STD_TORCH_CHECK(
      output.scalar_type() == torch::headeronly::ScalarType::Float ||
          output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
      "output must be float32 or bf16");

  const int sm = getSMVersion();
  STD_TORCH_CHECK(sm >= 90, "required CUDA ARCH >= SM_90");

  const cudaStream_t stream = get_current_cuda_stream(mat_a.get_device_index());

  __nv_bfloat16 const* a_ptr =
      reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
  __nv_bfloat16 const* b_ptr =
      reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());

  if (output.scalar_type() == torch::headeronly::ScalarType::Float) {
    float* out_ptr = reinterpret_cast<float*>(output.mutable_data_ptr());
    if (hidden_dim == DEFAULT_HIDDEN_DIM) {
      if (num_experts == DEFAULT_NUM_EXPERTS) {
        LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
                     DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens,
                                                              out_ptr, a_ptr,
                                                              b_ptr, stream);
      } else {
        LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS,
                     DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens,
                                                              out_ptr, a_ptr,
                                                              b_ptr, stream);
      }
    } else {  // GLM_5_HIDDEN_DIM
      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
                   GLM_5_HIDDEN_DIM>::unroll_float_output(num_tokens, out_ptr,
                                                          a_ptr, b_ptr, stream);
    }
  } else if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
    __nv_bfloat16* out_ptr =
        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr());
    if (hidden_dim == DEFAULT_HIDDEN_DIM) {
      if (num_experts == DEFAULT_NUM_EXPERTS) {
        LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
                     DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens,
                                                             out_ptr, a_ptr,
                                                             b_ptr, stream);
      } else {
        LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS,
                     DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens,
                                                             out_ptr, a_ptr,
                                                             b_ptr, stream);
      }
    } else {  // GLM_5_HIDDEN_DIM
      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
                   GLM_5_HIDDEN_DIM>::unroll_bf16_output(num_tokens, out_ptr,
                                                         a_ptr, b_ptr, stream);
    }
  }
}

STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
  m.impl("dsv3_router_gemm", TORCH_BOX(&dsv3_router_gemm));
}
