/*!
 * \file tl/cuda/target_utils.cc
 * \brief CUDA target attribute helpers.
 */

#include "cuda/target_utils.h"

#include <tvm/ffi/reflection/registry.h>

#include <string>

#include "dlpack/dlpack.h"
#include "support/check.h"

namespace tvm {
namespace tl {
namespace {

int GetCudaArchInt(Target target) {
  auto s = target->GetAttr<ffi::String>("arch");
  ICHECK(s.has_value());
  const std::string arch_str = s.value();
  ICHECK(arch_str.size() >= 3);
  ICHECK_EQ(arch_str.compare(0, 3, "sm_"), 0)
      << "arch string must start with sm_";
  return std::stoi(arch_str.substr(3));
}

} // namespace

bool TargetIsCuda(Target target) {
  return target->GetTargetDeviceType() == kDLCUDA;
}

bool TargetIsCuTeDSL(Target target) {
  for (const auto &key : target->keys) {
    if (key == "cutedsl")
      return true;
  }
  return false;
}

bool TargetIsVolta(Target target) {
  if (!TargetIsCuda(target))
    return false;
  int arch = GetCudaArchInt(target);
  return arch >= 70 && arch < 75;
}

bool TargetIsTuring(Target target) {
  if (!TargetIsCuda(target))
    return false;
  int arch = GetCudaArchInt(target);
  return arch >= 75 && arch < 80;
}

bool TargetIsAmpere(Target target) {
  if (!TargetIsCuda(target))
    return false;
  int arch = GetCudaArchInt(target);
  return arch >= 80 && arch < 90;
}

bool TargetIsHopper(Target target) {
  if (!TargetIsCuda(target))
    return false;
  int arch = GetCudaArchInt(target);
  return arch >= 90 && arch < 100;
}

bool TargetIsSm100(Target target) {
  if (!TargetIsCuda(target))
    return false;
  int arch = GetCudaArchInt(target);
  return arch >= 100 && arch <= 110;
}

bool TargetIsSM120(Target target) {
  if (!TargetIsCuda(target))
    return false;
  int arch = GetCudaArchInt(target);
  return arch >= 120 && arch < 130;
}

bool TargetCudaHasAsyncCopy(Target target) {
  if (!TargetIsCuda(target))
    return false;
  int arch = GetCudaArchInt(target);
  return arch >= 80;
}

int TargetCudaGetWarpSize(Target target) {
  (void)target;
  return 32;
}

bool TargetHasLdmatrix(Target target) {
  if (!TargetIsCuda(target))
    return false;
  int arch = GetCudaArchInt(target);
  return arch >= 75;
}

bool TargetHasStmatrix(Target target, bool is_m16n8) {
  if (!TargetIsCuda(target))
    return false;
  int arch = GetCudaArchInt(target);
  return is_m16n8 ? arch >= 100 : arch >= 90;
}

bool TargetHasTmem(Target target) {
  if (!TargetIsCuda(target))
    return false;
  return TargetIsSm100(target);
}

bool TargetHasBulkCopy(Target target) {
  if (!TargetIsCuda(target))
    return false;
  int arch = GetCudaArchInt(target);
  return arch >= 90;
}

bool TargetSupportVectorize256(Target target) {
  if (!TargetIsCuda(target))
    return false;
  int arch = GetCudaArchInt(target);
  return arch >= 100;
}

bool TargetHasSMVersionGE(Target target, int version) {
  if (!TargetIsCuda(target))
    return false;
  int arch = GetCudaArchInt(target);
  return arch >= version;
}

bool IsCudaVectorizableFP8(DataType dtype) {
  // NOTE: E8M0 is a special type of FP8 which is not handled here.
  // We only handle FP8 types which can be represented with
  // __nv_fp8_interpretation_t here.
  return dtype.is_float8_e4m3() || dtype.is_float8_e4m3fn() ||
         dtype.is_float8_e5m2();
}

bool IsCudaVectorizableCast(DataType from_ty, DataType target_ty) {
  // float16 -> float32
  if (from_ty.is_float16() && target_ty.is_float() && target_ty.bits() == 32)
    return true;

  // float32 -> float16
  if (from_ty.is_float() && from_ty.bits() == 32 && target_ty.is_float16())
    return true;

  // bfloat16 -> float32
  if (from_ty.is_bfloat16() && target_ty.is_float() && target_ty.bits() == 32)
    return true;

  // float32 -> bfloat16
  if (from_ty.is_float() && from_ty.bits() == 32 && target_ty.is_bfloat16())
    return true;

  // float32 -> float8 (E4M3/E5M2)
  if (from_ty.is_float() && from_ty.bits() == 32 &&
      IsCudaVectorizableFP8(target_ty))
    return true;

  // float8 (E4M3/E5M2) -> float32
  if (IsCudaVectorizableFP8(from_ty) && target_ty.is_float() &&
      target_ty.bits() == 32)
    return true;

  // float8 (E4M3/E5M2) -> float16
  if (IsCudaVectorizableFP8(from_ty) && target_ty.is_float16())
    return true;

  // float8 (E4M3/E5M2) -> bfloat16
  if (IsCudaVectorizableFP8(from_ty) && target_ty.is_bfloat16())
    return true;

  // float16 -> float8 (E4M3/E5M2)
  if (from_ty.is_float16() && IsCudaVectorizableFP8(target_ty))
    return true;

  // bfloat16 -> float8 (E4M3/E5M2)
  if (from_ty.is_bfloat16() && IsCudaVectorizableFP8(target_ty))
    return true;

  // Not implemented for now

  // float64(double) -> float8 (E4M3/E5M2)
  // if (from_ty.is_float() && from_ty.bits() == 64 &&
  //     IsCudaVectorizableFP8(target_ty))
  //   return true;

  // float8 (E4M3/E5M2) -> float64(double)
  // if (IsCudaVectorizableFP8(from_ty) && target_ty.is_float() &&
  //     target_ty.bits() == 64)
  //   return true;

  // float8 (E8M0) -> bfloat16
  if (from_ty.is_float8_e8m0fnu() && target_ty.is_bfloat16())
    return true;

  // bfloat16 -> float8 (E8M0)
  if (from_ty.is_bfloat16() && target_ty.is_float8_e8m0fnu())
    return true;

  // float32 -> float8 (E8M0)
  if (from_ty.is_float() && from_ty.bits() == 32 &&
      target_ty.is_float8_e8m0fnu())
    return true;

  // float64(double) -> float8 (E8M0)
  if (from_ty.is_float() && from_ty.bits() == 64 &&
      target_ty.is_float8_e8m0fnu())
    return true;

  // float4_e2m1fn -> float16
  if (from_ty.is_float4_e2m1fn() && target_ty.is_float16())
    return true;

  // float16 -> float4_e2m1fn
  if (from_ty.is_float16() && target_ty.is_float4_e2m1fn())
    return true;

  // float4_e2m1fn -> float32
  if (from_ty.is_float4_e2m1fn() && target_ty.is_float() &&
      target_ty.bits() == 32)
    return true;

  // float32 -> float4_e2m1fn
  if (from_ty.is_float() && from_ty.bits() == 32 &&
      target_ty.is_float4_e2m1fn())
    return true;

  // float4_e2m1fn -> float64(double)
  if (from_ty.is_float4_e2m1fn() && target_ty.is_float() &&
      target_ty.bits() == 64)
    return true;

  // float64(double) -> float4_e2m1fn
  if (from_ty.is_float() && from_ty.bits() == 64 &&
      target_ty.is_float4_e2m1fn())
    return true;

  // float4_e2m1fn -> bfloat16
  if (from_ty.is_float4_e2m1fn() && target_ty.is_bfloat16())
    return true;

  // bfloat16 -> float4_e2m1fn
  if (from_ty.is_bfloat16() && target_ty.is_float4_e2m1fn())
    return true;

  return false;
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("tl.TargetIsCuda",
           [](Target target) { return TargetIsCuda(target); })
      .def("tl.TargetIsVolta",
           [](Target target) { return TargetIsVolta(target); })
      .def("tl.TargetIsTuring",
           [](Target target) { return TargetIsTuring(target); })
      .def("tl.TargetIsAmpere",
           [](Target target) { return TargetIsAmpere(target); })
      .def("tl.TargetIsHopper",
           [](Target target) { return TargetIsHopper(target); })
      .def("tl.TargetIsSM120",
           [](Target target) { return TargetIsSM120(target); })
      .def("tl.TargetCudaGetWarpSize",
           [](Target target) { return TargetCudaGetWarpSize(target); })
      .def("tl.TargetHasLdmatrix",
           [](Target target) { return TargetHasLdmatrix(target); })
      .def_packed(
          "tl.TargetHasStmatrix",
          [](ffi::PackedArgs args, ffi::Any *ret) {
            ICHECK(args.size() == 1 || args.size() == 2)
                << "TargetHasStmatrix expects target and optional is_m16n8";
            Target target = args[0].cast<Target>();
            bool is_m16n8 = args.size() == 2 ? args[1].cast<bool>() : false;
            *ret = TargetHasStmatrix(target, is_m16n8);
          })
      .def("tl.TargetHasBulkCopy",
           [](Target target) { return TargetHasBulkCopy(target); });
}

} // namespace tl
} // namespace tvm
