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

/**
Testcase Scenarios :
1) Test to verify hipExtStreamGetCUMask api returning default CU Mask or global CU Mask.
2) Test to verify hipExtStreamGetCUMask api returns custom mask set.
3) Negative tests for hipExtStreamGetCUMask api.
*/

#include <hip_test_common.hh>

#include <cmath>
#include <vector>


/**
 * Scenario to verify hipExtStreamGetCUMask api returning default CU Mask or global CU Mask.
 * Scenario to verify hipExtStreamGetCUMask api returns custom mask set.
 */
HIP_TEST_CASE(Unit_hipExtStreamGetCUMask_verifyDefaultAndCustomMask) {
  constexpr unsigned maxCUPerValue = 32;
  hipDeviceProp_t props;
  std::stringstream ss;
  std::string gCUMask;
  std::string globalCUMask("");
  std::vector<uint32_t> defaultCUMask;

  int nGpu = 0;
  HIP_CHECK(hipGetDeviceCount(&nGpu));
  if (nGpu < 1) {
    HIP_SKIP_TEST(HipTest::SkipReason::kNoGpuDevice);
  }

  HIP_CHECK(hipSetDevice(0));
  HIP_CHECK(hipGetDeviceProperties(&props, 0));
  INFO("info: running on bus " << "0x" << props.pciBusID << " " << props.name << " with "
                               << props.multiProcessorCount << " CUs");

  const unsigned int maxNum =
      static_cast<unsigned int>(std::ceil((props.multiProcessorCount * 1.0f) / maxCUPerValue));
  std::vector<uint32_t> cuMask(maxNum);

  // Get global CU Mask if exists
  gCUMask = TestContext::getEnvVar("ROC_GLOBAL_CU_MASK");
  if (!gCUMask.empty()) {
    globalCUMask.assign(gCUMask);

    for_each(globalCUMask.begin(), globalCUMask.end(), [](char& c) { c = ::tolower(c); });
  }

  // Create default CU Mask
  uint32_t temp = 0;
  uint32_t bit_index = 0;
  for (uint32_t i = 0; i < (uint32_t)props.multiProcessorCount; i++) {
    temp |= 1UL << bit_index;
    if (bit_index >= 32) {
      defaultCUMask.push_back(temp);
      temp = 0;
      bit_index = 0;
      temp |= 1UL << bit_index;
    }
    bit_index += 1;
  }
  if (bit_index != 0) {
    defaultCUMask.push_back(temp);
  }

  SECTION("Verify with default CU Mask or global CU Mask") {
    // make a default CU mask bit-array where all CUs are active
    // this default mask is expected to be returned when there is no
    // custom or global CU mask defined

    HIP_CHECK(hipExtStreamGetCUMask(0, cuMask.size(), &cuMask[0]));

    ss << std::hex;
    for (int i = cuMask.size() - 1; i >= 0; i--) {
      ss << cuMask[i];
    }

    // remove extra 0 from ss if any present
    size_t found = ss.str().find_first_not_of("0");
    if (found != std::string::npos) {
      ss.str(ss.str().substr(found, ss.str().length()));
    }

    INFO("info: CU mask for the default stream is: 0x" << ss.str().c_str());
    if (globalCUMask.size() > 0) {
      if (ss.str().compare(globalCUMask) != 0) {
        INFO("Expected CU mask:" << globalCUMask.c_str() << ", api returned:" << ss.str().c_str());
        REQUIRE(false);
      }
    } else {
      for (int i = 0; i < min(cuMask.size(), defaultCUMask.size()); i++) {
        if (cuMask[i] != defaultCUMask[i]) {
          INFO("Expected CU mask " << defaultCUMask[i] << ", api returned:" << cuMask[i]);
          REQUIRE(false);
        }
      }
    }
  }

  SECTION("Verify with custom mask set") {
    hipDeviceProp_t props;
    std::vector<uint32_t> customMask(defaultCUMask);
    hipStream_t stream;
    int deviceId;

    HIP_CHECK(hipGetDevice(&deviceId));
    HIP_CHECK(hipGetDeviceProperties(&props, deviceId));

    if (props.major >= 10) {
      // For gfx >= 10, one work group processor encompasses 2 CUs &
      // hence the CUs need to be enabled in pair
      customMask[0] = 0xc;
    } else {
      customMask[0] = 0xe;
    }

    HIP_CHECK(hipExtStreamCreateWithCUMask(&stream, customMask.size(), customMask.data()));
    ss.str("");
    for (int i = customMask.size() - 1; i >= 0; i--) {
      ss << customMask[i];
    }
    INFO("info: setting a custom CU mask 0x" << ss.str());

    HIP_CHECK(hipExtStreamGetCUMask(stream, cuMask.size(), &cuMask[0]));
    ss.str("");
    for (int i = cuMask.size() - 1; i >= 0; i--) {
      ss << cuMask[i];
    }

    size_t found = ss.str().find_first_not_of("0");
    if (found != std::string::npos) {
      ss.str(ss.str().substr(found, ss.str().length()));
    }

    INFO("info: reading back CU mask 0x" << ss.str() << " for stream " << stream);

    if (!gCUMask.empty()) {
      for (size_t i = 0; i < customMask.size(); i++) {
        if (customMask[i] != cuMask[i]) {
          INFO("Error! expected CU mask:" << customMask[i] << ", Received CU mask:" << cuMask[i]);
          REQUIRE(false);
        }
      }
    }

    HIP_CHECK(hipStreamDestroy(stream));
  }
}

/**
 * Negative tests for hipExtStreamGetCUMask.
 */
HIP_TEST_CASE(Unit_hipExtStreamGetCUMask_Negative) {
  hipError_t ret;
  constexpr int maxNum = 6;
  std::vector<uint32_t> cuMask(maxNum);

  SECTION("cuMask is nullptr") {
    ret = hipExtStreamGetCUMask(0, cuMask.size(), nullptr);
    REQUIRE(ret == hipErrorInvalidValue);
  }

  SECTION("cuMaskSize is 0") {
    ret = hipExtStreamGetCUMask(0, 0, &cuMask[0]);
    REQUIRE(ret == hipErrorInvalidValue);
  }
}