/*
 * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 *
 * SPDX-License-Identifier: MIT
 */

#include <hip_test_common.hh>


// Renamed from Unit_hip_library_load_co (which collided with the more
// comprehensive case of the same name in loadlib_co.cc). This case is
// specifically a positive smoke test for hipKernelSetAttribute(
// MAX_DYNAMIC_SHARED_SIZE_BYTES) followed by a kernel launch via
// hipKernelGetFunction.
HIP_TEST_CASE(Unit_hipKernelSetAttribute_Positive_LaunchAfterSet) {
  constexpr size_t size = 32;
  constexpr size_t size_overwrite = 8;
  std::vector<float> input1, input2;
  input1.resize(size);
  input2.resize(size);
  for (size_t i = 0; i < size; i++) {
    input1[i] = (i + 1) * 2;
    input2[i] = i;
  }

  float *d_in1, *d_in2, *d_out;
  HIP_CHECK(hipMalloc(&d_in1, sizeof(float) * size));
  HIP_CHECK(hipMalloc(&d_in2, sizeof(float) * size));
  HIP_CHECK(hipMalloc(&d_out, sizeof(float) * size));

  HIP_CHECK(hipMemset(d_out, 0, sizeof(float) * size));
  HIP_CHECK(hipMemcpy(d_in1, input1.data(), sizeof(float) * size, hipMemcpyHostToDevice));
  HIP_CHECK(hipMemcpy(d_in2, input2.data(), sizeof(float) * size, hipMemcpyHostToDevice));

  hipStream_t stream;
  HIP_CHECK(hipStreamCreate(&stream));
  std::string lib_co = "library_code_load.code";

  hipLibrary_t library;
  hipFunction_t function;
  hipKernel_t kernel;

  HIP_CHECK(
      hipLibraryLoadFromFile(&library, lib_co.data(), nullptr, nullptr, 0, nullptr, nullptr, 0));
  HIP_CHECK(hipLibraryGetKernel(&kernel, library, "add_kernel"));
  HIP_CHECK(hipKernelGetFunction(&function, kernel));
  HIP_CHECK(hipKernelSetAttribute(HIP_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES ,sizeof(float) * size_overwrite, kernel, 0));

  void* args[] = {&d_out, &d_in1, &d_in2};

  HIP_CHECK(hipLaunchKernel(function, 1, size, args, 0, stream));
  HIP_CHECK(hipStreamSynchronize(stream));
  HIP_CHECK(hipLibraryUnload(library));

  std::vector<float> out(size, 0);
  HIP_CHECK(hipMemcpy(out.data(), d_out, sizeof(float) * size, hipMemcpyDeviceToHost));
  for (size_t i = 0; i < size; i++) {
    float tmp = input1[i] + input2[i];
    INFO("Index: " << i << " cpu res: " << tmp << " gpu res: " << out[i]);
    REQUIRE(out[i] == tmp);
  }

  HIP_CHECK(hipStreamDestroy(stream));
  HIP_CHECK(hipFree(d_in1));
  HIP_CHECK(hipFree(d_in2));
  HIP_CHECK(hipFree(d_out));
}

HIP_TEST_CASE(Unit_hipKernelSetAttribute_Negative_Parameters) {
  hipStream_t stream;
  HIP_CHECK(hipStreamCreate(&stream));
  std::string lib_co = "library_code_load.code";

  hipLibrary_t library;
  hipKernel_t kernel;

  HIP_CHECK(
      hipLibraryLoadFromFile(&library, lib_co.data(), nullptr, nullptr, 0, nullptr, nullptr, 0));
  HIP_CHECK(hipLibraryGetKernel(&kernel, library, "add_kernel"));

  int device_id = 0;

  SECTION("kernel == nullptr") {
    HIP_CHECK_ERROR(
        hipKernelSetAttribute(HIP_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, 0, nullptr,
                              device_id),
        hipErrorInvalidValue);
  }

  SECTION("invalid attribute") {
    HIP_CHECK_ERROR(
        hipKernelSetAttribute(static_cast<hipFunction_attribute>(-1), 0, kernel, device_id),
        hipErrorInvalidValue);
  }

  SECTION("read-only attribute MAX_THREADS_PER_BLOCK") {
    HIP_CHECK_ERROR(
        hipKernelSetAttribute(HIP_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK, 0, kernel, device_id),
        hipErrorInvalidValue);
  }

  SECTION("read-only attribute NUM_REGS") {
    HIP_CHECK_ERROR(
        hipKernelSetAttribute(HIP_FUNC_ATTRIBUTE_NUM_REGS, 0, kernel, device_id),
        hipErrorInvalidValue);
  }

  SECTION("read-only attribute SHARED_SIZE_BYTES") {
    HIP_CHECK_ERROR(
        hipKernelSetAttribute(HIP_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, 0, kernel, device_id),
        hipErrorInvalidValue);
  }

  SECTION("read-only attribute BINARY_VERSION") {
    HIP_CHECK_ERROR(
        hipKernelSetAttribute(HIP_FUNC_ATTRIBUTE_BINARY_VERSION, 0, kernel, device_id),
        hipErrorInvalidValue);
  }

  SECTION("MAX_DYNAMIC_SHARED_SIZE_BYTES negative value") {
    HIP_CHECK_ERROR(
        hipKernelSetAttribute(HIP_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, -1, kernel,
                              device_id),
        hipErrorInvalidValue);
  }

  SECTION("invalid device id") {
    int device_count = 0;
    HIP_CHECK(hipGetDeviceCount(&device_count));
    HIP_CHECK_ERROR(
        hipKernelSetAttribute(HIP_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, 0, kernel,
                              device_count),
        hipErrorInvalidDevice);
  }

  HIP_CHECK(hipLibraryUnload(library));
  HIP_CHECK(hipStreamDestroy(stream));
}

HIP_TEST_CASE(Unit_hipKernelGetFunction_Negative_Parameters) {
  hipStream_t stream;
  HIP_CHECK(hipStreamCreate(&stream));
  std::string lib_co = "library_code_load.code";

  hipLibrary_t library;
  hipKernel_t kernel;

  HIP_CHECK(
      hipLibraryLoadFromFile(&library, lib_co.data(), nullptr, nullptr, 0, nullptr, nullptr, 0));
  HIP_CHECK(hipLibraryGetKernel(&kernel, library, "add_kernel"));

  SECTION("pFunc == nullptr") {
    HIP_CHECK_ERROR(hipKernelGetFunction(nullptr, kernel), hipErrorInvalidValue);
  }

  SECTION("kernel == nullptr") {
    hipFunction_t function;
    HIP_CHECK_ERROR(hipKernelGetFunction(&function, nullptr), hipErrorInvalidValue);
  }

  HIP_CHECK(hipLibraryUnload(library));
  HIP_CHECK(hipStreamDestroy(stream));
}