/*
Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved.

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

#include <hip_test_common.hh>
#include <hip/hip_fp8.h>

#include <type_traits>
#include <vector>
#include <bitset>

/*
 * This catch test is meant for FP8 OCP conversions
 * tests only supported on gfx1200 and gfx1201 archs
 */

static_assert(sizeof(unsigned int) == sizeof(float));

std::string arch_type() {
  hipDeviceProp_t prop;
  int device;
  HIP_CHECK(hipGetDevice(&device));
  HIP_CHECK(hipGetDeviceProperties(&prop, device));
  std::string gfxName(prop.gcnArchName);
  return gfxName;
}

#define ARCH_TYPE_GFX1200(name)                                                                    \
  (name.find("gfx1200") != std::string::npos) || (name.find("gfx1201") != std::string::npos)

#define FP8_OCP_SKIP_TEST                                                                          \
  std::string gfxName = arch_type();                                                               \
  if (!(ARCH_TYPE_GFX1200(gfxName))) {                                                             \
    HipTest::HIP_SKIP_TEST("This test can only be run on GFX1200");                                \
    return;                                                                                        \
  }

#define __FP8_DEVICE__ __device__ static inline

template <typename T> __FP8_DEVICE__ void e4m3_ocp_device(T* val) {
#if (defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
  __hip_fp8_e4m3 tmp(*val);
  *val = tmp;
#else
  *val = 0;
#endif
}

template <typename T> __FP8_DEVICE__ void e5m2_ocp_device(T* val) {
#if (defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
  __hip_fp8_e5m2 tmp(*val);
  *val = tmp;
#else
  *val = 0;
#endif
}

template <typename T, bool is_e4m3_ocp> __global__ void cvt_float_fp8_float(T* in, size_t len) {
  int i = threadIdx.x;
  if (i < len) {
    T val = in[i];
    if constexpr (is_e4m3_ocp) {
      e4m3_ocp_device<T>(&val);
      in[i] = val;
    } else {
      e5m2_ocp_device<T>(&val);
      in[i] = val;
    }
  }
}

template <typename T, bool is_e4m3_ocp>
std::vector<T> cpu_cvt_float_fp8_float(const std::vector<T>& nums) {
  std::vector<T> ret;
  ret.reserve(nums.size());
  for (const auto& num : nums) {
    T out = 0.0;
    if constexpr (is_e4m3_ocp) {
      __hip_fp8_e4m3 tmp(num);
      out = tmp;
    } else {
      __hip_fp8_e5m2 tmp(num);
      out = tmp;
    }
    ret.push_back(out);
  }
  return ret;
}

// This test only makes sense on gfx94x where device side convert will use the builtins to convert
// floats to fp8
TEMPLATE_TEST_CASE("Unit_fp8_ocp_compare_host_device", "", float, double) {
  FP8_OCP_SKIP_TEST

  std::vector<TestType> numbers = {0.0f, 1.0f, 1.1f, 2.0f,  2.1f,  3.0f,  3.2f,
                                   3.3f, 4.0f, 4.5f, 10.0f, 11.0f, 12.2f, 14.1f};
  TestType* d_numbers;
  HIP_CHECK(hipMalloc(&d_numbers, sizeof(TestType) * numbers.size()));
  HIP_CHECK(hipMemcpy(d_numbers, numbers.data(), sizeof(TestType) * numbers.size(),
                      hipMemcpyHostToDevice));

  std::vector<TestType> result(numbers.size(), 0.0f);
  std::vector<TestType> cpu_result;

  SECTION("e4m3_ocp") {
    cpu_result = cpu_cvt_float_fp8_float<TestType, true>(numbers);
    auto kernel = cvt_float_fp8_float<TestType, true>;
    kernel<<<1, numbers.size()>>>(d_numbers, numbers.size());
    HIP_CHECK(hipMemcpy(result.data(), d_numbers, sizeof(TestType) * numbers.size(),
                        hipMemcpyDeviceToHost));
  }

  SECTION("e5m2_ocp") {
    cpu_result = cpu_cvt_float_fp8_float<TestType, false>(numbers);
    auto kernel = cvt_float_fp8_float<TestType, false>;
    kernel<<<1, numbers.size()>>>(d_numbers, numbers.size());
    HIP_CHECK(hipMemcpy(result.data(), d_numbers, sizeof(TestType) * numbers.size(),
                        hipMemcpyDeviceToHost));
  }

  REQUIRE(cpu_result.size() == result.size());
  for (size_t i = 0; i < result.size(); i++) {
    INFO("Checking: " << numbers[i] << " cpu convert: " << cpu_result[i]
                      << " - gpu_result: " << result[i]);
    CHECK(cpu_result[i] == result[i]);
  }

  HIP_CHECK(hipFree(d_numbers));
}

__FP8_DEVICE__ void e4m3_fp8x2_ocp_device(float2* val) {
#if (defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
  __hip_fp8x2_e4m3 tmp(*val);
  *val = tmp;
#else
  *val = float2(0.0, 0.0);
#endif
}

__FP8_DEVICE__ void e5m2_fp8x2_ocp_device(float2* val) {
#if (defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
  __hip_fp8x2_e5m2 tmp(*val);
  *val = tmp;
#else
  *val = float2(0.0, 0.0);
#endif
}

template <bool is_e4m3_ocp> __global__ void cvt_float2_fp8x2_float2(float2* in, size_t size) {
  int i = threadIdx.x;
  if (i < size) {
    float2 val = in[i];
    if constexpr (is_e4m3_ocp) {
      e4m3_fp8x2_ocp_device(&val);
      in[i] = val;
    } else {
      e5m2_fp8x2_ocp_device(&val);
      in[i] = val;
    }
  }
}

template <bool is_e4m3_ocp>
std::vector<float2> cpu_cvt_float2_fp8x2_float2(const std::vector<float2>& nums) {
  std::vector<float2> ret;
  ret.reserve(nums.size());
  for (const auto& num : nums) {
    float2 out = {0.0f, 0.0f};
    if constexpr (is_e4m3_ocp) {
      __hip_fp8x2_e4m3 tmp(num);
      out = tmp;
    } else {
      __hip_fp8x2_e5m2 tmp(num);
      out = tmp;
    }
    ret.push_back(out);
  }
  return ret;
}

TEST_CASE("Unit_fp8x2_ocp_compare_host_device") {
  FP8_OCP_SKIP_TEST
  std::vector<float> numbers_input = {0.0f, 1.0f, 1.1f, 2.0f,  2.1f,  3.0f,  3.2f,
                                      3.3f, 4.0f, 4.5f, 10.0f, 11.0f, 12.2f, 14.1f};

  std::vector<float2> numbers;
  numbers.reserve(numbers_input.size());
  for (size_t i = 0, end = numbers_input.size() - 1; i < numbers_input.size(); i++, end--) {
    float2 ret(numbers_input[i], numbers_input[end]);
    numbers.push_back(ret);
  }

  float2* d_numbers;
  HIP_CHECK(hipMalloc(&d_numbers, sizeof(float2) * numbers.size()));
  HIP_CHECK(
      hipMemcpy(d_numbers, numbers.data(), sizeof(float2) * numbers.size(), hipMemcpyHostToDevice));

  std::vector<float2> result(numbers.size(), float2{0.0f, 0.0f});
  std::vector<float2> cpu_result;

  SECTION("e4m3_ocp") {
    cpu_result = cpu_cvt_float2_fp8x2_float2<true>(numbers);
    auto kernel = cvt_float2_fp8x2_float2<true>;
    kernel<<<1, numbers.size()>>>(d_numbers, numbers.size());
    HIP_CHECK(hipMemcpy(result.data(), d_numbers, sizeof(float2) * numbers.size(),
                        hipMemcpyDeviceToHost));
  }

  SECTION("e5m2_ocp") {
    cpu_result = cpu_cvt_float2_fp8x2_float2<false>(numbers);
    auto kernel = cvt_float2_fp8x2_float2<false>;
    kernel<<<1, numbers.size()>>>(d_numbers, numbers.size());
    HIP_CHECK(hipMemcpy(result.data(), d_numbers, sizeof(float2) * numbers.size(),
                        hipMemcpyDeviceToHost));
  }

  REQUIRE(cpu_result.size() == result.size());
  for (size_t i = 0; i < result.size(); i++) {
    CHECK(cpu_result[i] == result[i]);
  }

  HIP_CHECK(hipFree(d_numbers));
}

TEST_CASE("Unit_fp8x2_ocp_split_compare") {
  FP8_OCP_SKIP_TEST
  std::vector<float> numbers_input = {0.0f, 1.0f, 1.1f, 2.0f,  2.1f,  3.0f,  3.2f,
                                      3.3f, 4.0f, 4.5f, 10.0f, 11.0f, 12.2f, 14.1f};

  std::vector<float2> numbers;
  numbers.reserve(numbers_input.size());
  for (size_t i = 0, end = numbers_input.size() - 1; i < numbers_input.size(); i++, end--) {
    float2 ret(numbers_input[i], numbers_input[end]);
    numbers.push_back(ret);
  }

  float2* d_numbers;
  HIP_CHECK(hipMalloc(&d_numbers, sizeof(float2) * numbers.size()));
  HIP_CHECK(
      hipMemcpy(d_numbers, numbers.data(), sizeof(float2) * numbers.size(), hipMemcpyHostToDevice));

  std::vector<float2> result(numbers.size(), float2{0.0f, 0.0f});
  std::vector<float2> cpu_result;
  cpu_result.reserve(result.size());

  SECTION("e4m3_ocp") {
    for (const auto& num : numbers) {
      __hip_fp8_e4m3 t_a(num.x);
      __hip_fp8_e4m3 t_b(num.y);
      float a = t_a, b = t_b;
      cpu_result.push_back(float2(a, b));
    }
    auto kernel = cvt_float2_fp8x2_float2<true>;
    kernel<<<1, numbers.size()>>>(d_numbers, numbers.size());
    HIP_CHECK(hipMemcpy(result.data(), d_numbers, sizeof(float2) * numbers.size(),
                        hipMemcpyDeviceToHost));
  }

  SECTION("e5m2_ocp") {
    for (const auto& num : numbers) {
      __hip_fp8_e5m2 t_a(num.x);
      __hip_fp8_e5m2 t_b(num.y);
      float a = t_a, b = t_b;
      cpu_result.push_back(float2(a, b));
    }
    auto kernel = cvt_float2_fp8x2_float2<false>;
    kernel<<<1, numbers.size()>>>(d_numbers, numbers.size());
    HIP_CHECK(hipMemcpy(result.data(), d_numbers, sizeof(float2) * numbers.size(),
                        hipMemcpyDeviceToHost));
  }

  HIP_CHECK(hipDeviceSynchronize());
  REQUIRE(cpu_result.size() == result.size());
  for (size_t i = 0; i < result.size(); i++) {
    INFO("cpu x: " << cpu_result[i].x << " y: " << cpu_result[i].y << " gpu x: " << result[i].x
                   << " y: " << result[i].y);
    CHECK(cpu_result[i] == result[i]);
  }

  HIP_CHECK(hipFree(d_numbers));
}

__FP8_DEVICE__ void e4m3_fp8x4_ocp_device(float4* val) {
#if (defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
  __hip_fp8x4_e4m3 tmp(*val);
  *val = tmp;
#else
  *val = float4(0.0, 0.0, 0.0, 0.0);
#endif
}

__FP8_DEVICE__ void e5m2_fp8x4_ocp_device(float4* val) {
#if (defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
  __hip_fp8x4_e5m2 tmp(*val);
  *val = tmp;
#else
  *val = float4(0.0, 0.0, 0.0, 0.0);
#endif
}
template <bool is_e4m3_ocp> __global__ void cvt_float4_fp8x4_float4(float4* in, size_t size) {
  int i = threadIdx.x;
  if (i < size) {
    float4 val = in[i];
    if constexpr (is_e4m3_ocp) {
      e4m3_fp8x4_ocp_device(&val);
      in[i] = val;
    } else {
      e5m2_fp8x4_ocp_device(&val);
      in[i] = val;
    }
  }
}

TEST_CASE("Unit_fp8x4_ocp_split_compare") {
  FP8_OCP_SKIP_TEST

  std::vector<float> numbers_input = {0.0f, 1.0f, 1.1f, 2.0f,  2.1f,  3.0f,  3.2f,
                                      3.3f, 4.0f, 4.5f, 10.0f, 11.0f, 12.2f, 14.1f};
  std::vector<float> numbers_input2 = {1.3f, 1.6f, 1.8f, 2.5f,  2.9f,  3.8f,  3.9f,
                                       5.5f, 7.1f, 8.5f, 11.2f, 13.5f, 16.1f, 19.4f};

  std::vector<float4> numbers;
  numbers.reserve(numbers_input.size());
  for (size_t i = 0, end = numbers_input.size() - 1; i < numbers_input.size(); i++, end--) {
    float4 ret(numbers_input[i], numbers_input[end], numbers_input2[i], numbers_input2[end]);
    numbers.push_back(ret);
  }

  float4* d_numbers;
  HIP_CHECK(hipMalloc(&d_numbers, sizeof(float4) * numbers.size()));
  HIP_CHECK(
      hipMemcpy(d_numbers, numbers.data(), sizeof(float4) * numbers.size(), hipMemcpyHostToDevice));

  std::vector<float4> result(numbers.size(), float4{0.0f, 0.0f, 0.0f, 0.0f});
  std::vector<float4> cpu_result;
  cpu_result.reserve(result.size());

  SECTION("e4m3_ocp") {
    for (const auto& num : numbers) {
      __hip_fp8_e4m3 t_a(num.x);
      __hip_fp8_e4m3 t_b(num.y);
      __hip_fp8_e4m3 t_c(num.z);
      __hip_fp8_e4m3 t_d(num.w);
      float a = t_a, b = t_b, c = t_c, d = t_d;
      cpu_result.push_back(float4(a, b, c, d));
    }
    auto kernel = cvt_float4_fp8x4_float4<true>;
    kernel<<<1, numbers.size()>>>(d_numbers, numbers.size());
    HIP_CHECK(hipMemcpy(result.data(), d_numbers, sizeof(float4) * numbers.size(),
                        hipMemcpyDeviceToHost));
  }

  SECTION("e5m2_ocp") {
    for (const auto& num : numbers) {
      __hip_fp8_e5m2 t_a(num.x);
      __hip_fp8_e5m2 t_b(num.y);
      __hip_fp8_e5m2 t_c(num.z);
      __hip_fp8_e5m2 t_d(num.w);
      float a = t_a, b = t_b, c = t_c, d = t_d;
      cpu_result.push_back(float4(a, b, c, d));
    }
    auto kernel = cvt_float4_fp8x4_float4<false>;
    kernel<<<1, numbers.size()>>>(d_numbers, numbers.size());
    HIP_CHECK(hipMemcpy(result.data(), d_numbers, sizeof(float4) * numbers.size(),
                        hipMemcpyDeviceToHost));
  }

  HIP_CHECK(hipDeviceSynchronize());
  REQUIRE(cpu_result.size() == result.size());
  for (size_t i = 0; i < result.size(); i++) {
    INFO("original: x: " << numbers[i].x << " y: " << numbers[i].y << " z: " << numbers[i].z
                         << " w: " << numbers[i].w);
    INFO("cpu x: " << cpu_result[i].x << " y: " << cpu_result[i].y << " z: " << cpu_result[i].z
                   << " w: " << cpu_result[i].w);
    INFO("gpu x: " << result[i].x << " y: " << result[i].y << " z: " << result[i].z
                   << " w: " << result[i].w);
    CHECK(cpu_result[i] == result[i]);
  }

  HIP_CHECK(hipFree(d_numbers));
}

__FP8_DEVICE__ bool e4m3_bool_ocp_device(float val) {
  bool x = false;
  float y = val;
#if (defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
  __hip_fp8_e4m3 tmp(y);
  x = tmp;
#else
  x = (y == 0);
#endif
  return x;
}

__FP8_DEVICE__ bool e5m2_bool_ocp_device(float val) {
  bool x = false;
  float y = val;
#if (defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
  __hip_fp8_e5m2 tmp(y);
  x = tmp;
#else
  x = (y == 0);
#endif
  return x;
}
template <bool is_e4m3_ocp> __global__ void fp8_2_bool(float* f, bool* ret, size_t size) {
  int i = threadIdx.x;
  bool r = false;
  if (i < size) {
    if constexpr (is_e4m3_ocp) {
      r = e4m3_bool_ocp_device(f[i]);
    } else {
      r = e5m2_bool_ocp_device(f[i]);
    }
    ret[i] = r;
  }
}

TEST_CASE("Unit_fp8_ocp_bool_device") {
  FP8_OCP_SKIP_TEST
  // clang-format off
  std::vector<float> fvals{-10.0f, -1.0f, -0.0f,  0.0f, 1.0f, 10.0f};
  std::vector<bool> tvals   {true,  true, false, false, true,  true};
  // clang-format on

  bool result[] = {false, false, false,
                   false, false, false};  // cant use std::vector coz data() = delete

  SECTION("e4m3_ocp-gpu") {
    float* d_in{nullptr};
    bool* d_res{nullptr};
    HIP_CHECK(hipMalloc(&d_in, sizeof(float) * tvals.size()));
    HIP_CHECK(hipMalloc(&d_res, sizeof(bool) * tvals.size()));

    auto kernel = fp8_2_bool<true>;
    HIP_CHECK(hipMemcpy(d_in, fvals.data(), sizeof(float) * fvals.size(), hipMemcpyHostToDevice));
    kernel<<<1, tvals.size()>>>(d_in, d_res, tvals.size());

    HIP_CHECK(hipMemcpy(result, d_res, sizeof(bool) * tvals.size(), hipMemcpyDeviceToHost));

    HIP_CHECK(hipFree(d_in));
    HIP_CHECK(hipFree(d_res));
  }

  SECTION("e5m2_ocp-gpu") {
    float* d_in{nullptr};
    bool* d_res{nullptr};
    HIP_CHECK(hipMalloc(&d_in, sizeof(float) * tvals.size()));
    HIP_CHECK(hipMalloc(&d_res, sizeof(bool) * tvals.size()));

    HIP_CHECK(hipMemcpy(d_in, fvals.data(), sizeof(float) * fvals.size(), hipMemcpyHostToDevice));
    auto kernel = fp8_2_bool<false>;
    kernel<<<1, tvals.size()>>>(d_in, d_res, tvals.size());

    HIP_CHECK(hipMemcpy(result, d_res, sizeof(bool) * tvals.size(), hipMemcpyDeviceToHost));

    HIP_CHECK(hipFree(d_in));
    HIP_CHECK(hipFree(d_res));
  }

  for (size_t i = 0; i < tvals.size(); i++) {
    INFO("Check for: " << fvals[i] << " expected: " << tvals[i] << " result: " << result[i]);
    REQUIRE(result[i] == tvals[i]);
  }
}

std::vector<__hip_fp8_storage_t> get_all_fp8_nums(bool is_e4m3_ocp) {
  std::vector<__hip_fp8_storage_t> ret;
  constexpr unsigned short max_fp8_num = 0b1111'1111;
  ret.reserve(max_fp8_num + 1);

  for (unsigned short i = 0; i <= max_fp8_num; i++) {
    if (is_e4m3_ocp) {
      if ((i & 0x7f) != 0x7f) {  // 0xff and 0x7f are nan
        ret.push_back(static_cast<__hip_fp8_storage_t>(i));
      }
    } else {
      if ((i & 0x7f) < 0x7c) {  // 0x7c 0x7d 0x7e are nan and 0x7f is inf
        ret.push_back(static_cast<__hip_fp8_storage_t>(i));
      }
    }
  }
  return ret;
}

__FP8_DEVICE__ __hip_fp8_storage_t e4m3_ocp_fp8_device(float val) {
  __hip_fp8_storage_t x = 0;
  float y = val;
#if (defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
  __hip_fp8_e4m3 tmp(y);
  x = tmp.__x;
#else
  x = (y == 0) ? 0x0 : 0x7f;
#endif
  return x;
}

__FP8_DEVICE__ __hip_fp8_storage_t e5m2_ocp_fp8_device(float val) {
  __hip_fp8_storage_t x = 0;
  float y = val;
#if (defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
  __hip_fp8_e5m2 tmp(y);
  x = tmp.__x;
#else
  x = (y == 0) ? 0x0 : 0x7f;
#endif
  return x;
}

template <bool is_e4m3_ocp>
__global__ void Type_to_fp8(float* f, __hip_fp8_storage_t* res, size_t size) {
  auto i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < size) {
    if constexpr (is_e4m3_ocp) {
      res[i] = e4m3_ocp_fp8_device(f[i]);
    } else {
      res[i] = e5m2_ocp_fp8_device(f[i]);
    }
  }
}

TEST_CASE("Unit_all_fp8_ocp_cvt") {
  FP8_OCP_SKIP_TEST

  bool is_e4m3_ocp = GENERATE(true, false);
  std::vector<float> f_vals;
  std::vector<__hip_fp8_storage_t> all_vals;

  SECTION("all representable number") {
    all_vals = get_all_fp8_nums(is_e4m3_ocp);
    f_vals.reserve(all_vals.size());

    for (const auto& fp8 : all_vals) {
      float f = 0.0f;
      if (is_e4m3_ocp) {
        __hip_fp8_e4m3 tmp;
        tmp.__x = fp8;
        f = tmp;
      } else {
        __hip_fp8_e5m2 tmp;
        tmp.__x = fp8;
        f = tmp;
      }
      f_vals.push_back(f);
    }
  }

  SECTION("Range stepped numbers") {
    constexpr float lhs = -200.0f;
    constexpr float rhs = 200.0f;
    constexpr float step = 0.1234f;

    f_vals.reserve(4000);
    all_vals.reserve(4000);

    for (float fval = lhs; fval <= rhs; fval += step) {
      if (is_e4m3_ocp) {
        __hip_fp8_e4m3 tmp = fval;
        all_vals.push_back(tmp.__x);
      } else {
        __hip_fp8_e5m2 tmp = fval;
        all_vals.push_back(tmp.__x);
      }
      f_vals.push_back(fval);
    }
  }

  float* d_f_vals;
  __hip_fp8_storage_t* d_res;

  HIP_CHECK(hipMalloc(&d_f_vals, sizeof(float) * f_vals.size()));
  HIP_CHECK(hipMalloc(&d_res, sizeof(__hip_fp8_storage_t) * f_vals.size()));

  HIP_CHECK(
      hipMemcpy(d_f_vals, f_vals.data(), sizeof(float) * f_vals.size(), hipMemcpyHostToDevice));

  auto fp8_kernel = is_e4m3_ocp ? Type_to_fp8<true> : Type_to_fp8<false>;
  fp8_kernel<<<(f_vals.size() / 256) + 1, 256>>>(d_f_vals, d_res, f_vals.size());

  std::vector<__hip_fp8_storage_t> final_res(f_vals.size(), static_cast<__hip_fp8_storage_t>(0));

  HIP_CHECK(hipMemcpy(final_res.data(), d_res, sizeof(__hip_fp8_storage_t) * final_res.size(),
                      hipMemcpyDeviceToHost));

  for (size_t i = 0; i < final_res.size(); i++) {
    INFO("Checking: " << f_vals[i] << " for: " << (is_e4m3_ocp ? "e4m3_ocp" : "e5m2_ocp")
                      << " original: " << (int)all_vals[i]
                      << " convert back: " << (int)final_res[i]);
    float gpu_cvt_res = 0.0f, cpu_cvt_res = 0.0f;
    if (is_e4m3_ocp) {
      __hip_fp8_e4m3 gtmp;
      gtmp.__x = final_res[i];
      gpu_cvt_res = gtmp;
      __hip_fp8_e4m3 ctmp;
      ctmp.__x = all_vals[i];
      cpu_cvt_res = ctmp;
    } else {
      __hip_fp8_e5m2 gtmp;
      gtmp.__x = final_res[i];
      gpu_cvt_res = gtmp;
      __hip_fp8_e5m2 ctmp;
      ctmp.__x = all_vals[i];
      cpu_cvt_res = ctmp;
    }

    INFO("cpu cvt val: " << cpu_cvt_res << " gpu cvt val: " << gpu_cvt_res);
    REQUIRE(cpu_cvt_res == gpu_cvt_res);
  }

  HIP_CHECK(hipFree(d_f_vals));
  HIP_CHECK(hipFree(d_res));
}

template <typename T> __FP8_DEVICE__ void e4m3_ocp_fp8_cvt(T val, float* cvt1, float* cvt2) {
  T y = val;
#if (defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
  __hip_fp8_e4m3 tmp(y);
  *cvt1 = tmp;

  __hip_fp8_e4m3 tmp1;
  tmp1.__x = std::is_same<T, float>::value
                 ? __hip_cvt_float_to_fp8(val, __HIP_SATFINITE, __HIP_E4M3)
                 : __hip_cvt_double_to_fp8(val, __HIP_SATFINITE, __HIP_E4M3);
  ;
  *cvt2 = tmp1;
#else
  *cvt1 = (y == 0) ? 0 : y;
  *cvt2 = (y == 0) ? 0 : y;
#endif
}

template <typename T> __FP8_DEVICE__ void e5m2_ocp_fp8_cvt(T val, float* cvt1, float* cvt2) {
  T y = val;
#if (defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
  __hip_fp8_e5m2 tmp(y);
  *cvt1 = tmp;

  __hip_fp8_e5m2 tmp1;
  tmp1.__x = std::is_same<T, float>::value
                 ? __hip_cvt_float_to_fp8(val, __HIP_SATFINITE, __HIP_E5M2)
                 : __hip_cvt_double_to_fp8(val, __HIP_SATFINITE, __HIP_E5M2);
  ;
  *cvt2 = tmp1;
#else
  *cvt1 = (y == 0) ? 0 : y;
  *cvt2 = (y == 0) ? 0 : y;
#endif
}


template <typename T, bool is_e4m3_ocp>
__global__ void Type_to_fp8_cvt(T* f, float* cvt1, float* cvt2, size_t size) {
  auto i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < size) {
    if constexpr (is_e4m3_ocp) {
      e4m3_ocp_fp8_cvt(f[i], &cvt1[i], &cvt2[i]);
    } else {
      e5m2_ocp_fp8_cvt(f[i], &cvt1[i], &cvt2[i]);
    }
  }
}

TEMPLATE_TEST_CASE("Unit_fp8_ocp_correctness_device", "", float, double) {
  FP8_OCP_SKIP_TEST

  SECTION("e4m3_ocp") {
    /* These are basically all the fp8 - e4m3_ocp type numbers.
     * They can be generated by iterating over 0'0000'000 and converting them to fp32 number
     * skipping the nan/inf */
    std::vector<TestType> e4m3_ocp_nums = {
        0,          0.00195312, 0.00390625,  0.00585938,  0.0078125,   0.00976562, 0.0117188,
        0.0136719,  0.015625,   0.0175781,   0.0195312,   0.0214844,   0.0234375,  0.0253906,
        0.0273438,  0.0292969,  0.03125,     0.0351562,   0.0390625,   0.0429688,  0.046875,
        0.0507812,  0.0546875,  0.0585938,   0.0625,      0.0703125,   0.078125,   0.0859375,
        0.09375,    0.101562,   0.109375,    0.117188,    0.125,       0.140625,   0.15625,
        0.171875,   0.1875,     0.203125,    0.21875,     0.234375,    0.25,       0.28125,
        0.3125,     0.34375,    0.375,       0.40625,     0.4375,      0.46875,    0.5,
        0.5625,     0.625,      0.6875,      0.75,        0.8125,      0.875,      0.9375,
        1,          1.125,      1.25,        1.375,       1.5,         1.625,      1.75,
        1.875,      2,          2.25,        2.5,         2.75,        3,          3.25,
        3.5,        3.75,       4,           4.5,         5,           5.5,        6,
        6.5,        7,          7.5,         8,           9,           10,         11,
        12,         13,         14,          15,          16,          18,         20,
        22,         24,         26,          28,          30,          32,         36,
        40,         44,         48,          52,          56,          60,         64,
        72,         80,         88,          96,          104,         112,        120,
        128,        144,        160,         176,         192,         208,        224,
        240,        256,        288,         320,         352,         384,        416,
        448,        -0,         -0.00195312, -0.00390625, -0.00585938, -0.0078125, -0.00976562,
        -0.0117188, -0.0136719, -0.015625,   -0.0175781,  -0.0195312,  -0.0214844, -0.0234375,
        -0.0253906, -0.0273438, -0.0292969,  -0.03125,    -0.0351562,  -0.0390625, -0.0429688,
        -0.046875,  -0.0507812, -0.0546875,  -0.0585938,  -0.0625,     -0.0703125, -0.078125,
        -0.0859375, -0.09375,   -0.101562,   -0.109375,   -0.117188,   -0.125,     -0.140625,
        -0.15625,   -0.171875,  -0.1875,     -0.203125,   -0.21875,    -0.234375,  -0.25,
        -0.28125,   -0.3125,    -0.34375,    -0.375,      -0.40625,    -0.4375,    -0.46875,
        -0.5,       -0.5625,    -0.625,      -0.6875,     -0.75,       -0.8125,    -0.875,
        -0.9375,    -1,         -1.125,      -1.25,       -1.375,      -1.5,       -1.625,
        -1.75,      -1.875,     -2,          -2.25,       -2.5,        -2.75,      -3,
        -3.25,      -3.5,       -3.75,       -4,          -4.5,        -5,         -5.5,
        -6,         -6.5,       -7,          -7.5,        -8,          -9,         -10,
        -11,        -12,        -13,         -14,         -15,         -16,        -18,
        -20,        -22,        -24,         -26,         -28,         -30,        -32,
        -36,        -40,        -44,         -48,         -52,         -56,        -60,
        -64,        -72,        -80,         -88,         -96,         -104,       -112,
        -120,       -128,       -144,        -160,        -176,        -192,       -208,
        -224,       -240,       -256,        -288,        -320,        -352,       -384,
        -416,       -448};
    size_t totalnums = e4m3_ocp_nums.size();
    TestType* fnums;
    HIP_CHECK(hipMalloc((void**)&fnums, totalnums * sizeof(TestType)));
    float* cvt1_dev;
    HIP_CHECK(hipMalloc((void**)&cvt1_dev, totalnums * sizeof(TestType)));
    float* cvt2_dev;
    HIP_CHECK(hipMalloc((void**)&cvt2_dev, totalnums * sizeof(TestType)));

    HIP_CHECK(hipMemcpy(fnums, e4m3_ocp_nums.data(), totalnums * sizeof(TestType),
                        hipMemcpyHostToDevice));

    auto fp8_kernel = Type_to_fp8_cvt<TestType, true>;
    fp8_kernel<<<totalnums / 256 + 1, 256>>>(fnums, cvt1_dev, cvt2_dev, totalnums);

    float* cvt1_host = (float*)malloc(sizeof(float) * totalnums);
    float* cvt2_host = (float*)malloc(sizeof(float) * totalnums);

    HIP_CHECK(hipMemcpy(cvt1_host, cvt1_dev, totalnums * sizeof(float), hipMemcpyDeviceToHost));
    HIP_CHECK(hipMemcpy(cvt2_host, cvt2_dev, totalnums * sizeof(float), hipMemcpyDeviceToHost));

    HIP_CHECK(hipDeviceSynchronize());

    for (size_t idx = 0; idx < totalnums; idx++) {
      TestType orig = e4m3_ocp_nums[idx];
      float cvt1 = cvt1_host[idx];
      float cvt2 = cvt2_host[idx];

      INFO("Original: " << std::bitset<32>(*reinterpret_cast<const unsigned int*>(&orig)));
      INFO("Cvt back: " << std::bitset<32>(*reinterpret_cast<const unsigned int*>(&cvt1)));
      REQUIRE(cvt1 == Catch::Approx(orig));
      REQUIRE(cvt2 == cvt1);
    }

    HIP_CHECK(hipFree(fnums));
    HIP_CHECK(hipFree(cvt1_dev));
    HIP_CHECK(hipFree(cvt2_dev));
    free(cvt1_host);
    free(cvt2_host);
  }

  SECTION("e5m2_ocp") {
    /* These are basically all the fp8 - e5m2_ocp type numbers.
     * They can be generated by iterating over 0'00000'00 converting them to fp32 number skipping
     * the nan/inf */
    std::vector<TestType> e5m2_ocp_nums = {0,
                                           1.52588e-05,
                                           3.05176e-05,
                                           4.57764e-05,
                                           6.10352e-05,
                                           7.62939e-05,
                                           9.15527e-05,
                                           0.000106812,
                                           0.00012207,
                                           0.000152588,
                                           0.000183105,
                                           0.000213623,
                                           0.000244141,
                                           0.000305176,
                                           0.000366211,
                                           0.000427246,
                                           0.000488281,
                                           0.000610352,
                                           0.000732422,
                                           0.000854492,
                                           0.000976562,
                                           0.0012207,
                                           0.00146484,
                                           0.00170898,
                                           0.00195312,
                                           0.00244141,
                                           0.00292969,
                                           0.00341797,
                                           0.00390625,
                                           0.00488281,
                                           0.00585938,
                                           0.00683594,
                                           0.0078125,
                                           0.00976562,
                                           0.0117188,
                                           0.0136719,
                                           0.015625,
                                           0.0195312,
                                           0.0234375,
                                           0.0273438,
                                           0.03125,
                                           0.0390625,
                                           0.046875,
                                           0.0546875,
                                           0.0625,
                                           0.078125,
                                           0.09375,
                                           0.109375,
                                           0.125,
                                           0.15625,
                                           0.1875,
                                           0.21875,
                                           0.25,
                                           0.3125,
                                           0.375,
                                           0.4375,
                                           0.5,
                                           0.625,
                                           0.75,
                                           0.875,
                                           1,
                                           1.25,
                                           1.5,
                                           1.75,
                                           2,
                                           2.5,
                                           3,
                                           3.5,
                                           4,
                                           5,
                                           6,
                                           7,
                                           8,
                                           10,
                                           12,
                                           14,
                                           16,
                                           20,
                                           24,
                                           28,
                                           32,
                                           40,
                                           48,
                                           56,
                                           64,
                                           80,
                                           96,
                                           112,
                                           128,
                                           160,
                                           192,
                                           224,
                                           256,
                                           320,
                                           384,
                                           448,
                                           512,
                                           640,
                                           768,
                                           896,
                                           1024,
                                           1280,
                                           1536,
                                           1792,
                                           2048,
                                           2560,
                                           3072,
                                           3584,
                                           4096,
                                           5120,
                                           6144,
                                           7168,
                                           8192,
                                           10240,
                                           12288,
                                           14336,
                                           16384,
                                           20480,
                                           24576,
                                           28672,
                                           32768,
                                           40960,
                                           49152,
                                           57344,
                                           -0,
                                           -1.52588e-05,
                                           -3.05176e-05,
                                           -4.57764e-05,
                                           -6.10352e-05,
                                           -7.62939e-05,
                                           -9.15527e-05,
                                           -0.000106812,
                                           -0.00012207,
                                           -0.000152588,
                                           -0.000183105,
                                           -0.000213623,
                                           -0.000244141,
                                           -0.000305176,
                                           -0.000366211,
                                           -0.000427246,
                                           -0.000488281,
                                           -0.000610352,
                                           -0.000732422,
                                           -0.000854492,
                                           -0.000976562,
                                           -0.0012207,
                                           -0.00146484,
                                           -0.00170898,
                                           -0.00195312,
                                           -0.00244141,
                                           -0.00292969,
                                           -0.00341797,
                                           -0.00390625,
                                           -0.00488281,
                                           -0.00585938,
                                           -0.00683594,
                                           -0.0078125,
                                           -0.00976562,
                                           -0.0117188,
                                           -0.0136719,
                                           -0.015625,
                                           -0.0195312,
                                           -0.0234375,
                                           -0.0273438,
                                           -0.03125,
                                           -0.0390625,
                                           -0.046875,
                                           -0.0546875,
                                           -0.0625,
                                           -0.078125,
                                           -0.09375,
                                           -0.109375,
                                           -0.125,
                                           -0.15625,
                                           -0.1875,
                                           -0.21875,
                                           -0.25,
                                           -0.3125,
                                           -0.375,
                                           -0.4375,
                                           -0.5,
                                           -0.625,
                                           -0.75,
                                           -0.875,
                                           -1,
                                           -1.25,
                                           -1.5,
                                           -1.75,
                                           -2,
                                           -2.5,
                                           -3,
                                           -3.5,
                                           -4,
                                           -5,
                                           -6,
                                           -7,
                                           -8,
                                           -10,
                                           -12,
                                           -14,
                                           -16,
                                           -20,
                                           -24,
                                           -28,
                                           -32,
                                           -40,
                                           -48,
                                           -56,
                                           -64,
                                           -80,
                                           -96,
                                           -112,
                                           -128,
                                           -160,
                                           -192,
                                           -224,
                                           -256,
                                           -320,
                                           -384,
                                           -448,
                                           -512,
                                           -640,
                                           -768,
                                           -896,
                                           -1024,
                                           -1280,
                                           -1536,
                                           -1792,
                                           -2048,
                                           -2560,
                                           -3072,
                                           -3584,
                                           -4096,
                                           -5120,
                                           -6144,
                                           -7168,
                                           -8192,
                                           -10240,
                                           -12288,
                                           -14336,
                                           -16384,
                                           -20480,
                                           -24576,
                                           -28672,
                                           -32768,
                                           -40960,
                                           -49152,
                                           -57344};
    size_t totalnums = e5m2_ocp_nums.size();
    TestType* fnums;
    HIP_CHECK(hipMalloc((void**)&fnums, totalnums * sizeof(TestType)));
    float* cvt1_dev;
    HIP_CHECK(hipMalloc((void**)&cvt1_dev, totalnums * sizeof(TestType)));
    float* cvt2_dev;
    HIP_CHECK(hipMalloc((void**)&cvt2_dev, totalnums * sizeof(TestType)));

    HIP_CHECK(hipMemcpy(fnums, e5m2_ocp_nums.data(), totalnums * sizeof(TestType),
                        hipMemcpyHostToDevice));

    auto fp8_kernel = Type_to_fp8_cvt<TestType, false>;
    fp8_kernel<<<totalnums / 256 + 1, 256>>>(fnums, cvt1_dev, cvt2_dev, totalnums);

    float* cvt1_host = (float*)malloc(sizeof(float) * totalnums);
    float* cvt2_host = (float*)malloc(sizeof(float) * totalnums);

    HIP_CHECK(hipMemcpy(cvt1_host, cvt1_dev, totalnums * sizeof(float), hipMemcpyDeviceToHost));
    HIP_CHECK(hipMemcpy(cvt2_host, cvt2_dev, totalnums * sizeof(float), hipMemcpyDeviceToHost));

    HIP_CHECK(hipDeviceSynchronize());

    for (size_t idx = 0; idx < totalnums; idx++) {
      TestType orig = e5m2_ocp_nums[idx];
      float cvt1 = cvt1_host[idx];
      float cvt2 = cvt2_host[idx];

      INFO("Original: " << std::bitset<32>(*reinterpret_cast<const unsigned int*>(&orig)));
      INFO("Cvt back: " << std::bitset<32>(*reinterpret_cast<const unsigned int*>(&cvt1)));
      REQUIRE(cvt1 == Catch::Approx(orig));
      REQUIRE(cvt2 == cvt1);
    }

    HIP_CHECK(hipFree(fnums));
    HIP_CHECK(hipFree(cvt1_dev));
    HIP_CHECK(hipFree(cvt2_dev));
    free(cvt1_host);
    free(cvt2_host);
  }
}