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

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

#include <assert.h>
#include <stdio.h>
#include <algorithm>
#include <stdlib.h>
#include <iostream>

#define WIDTH 8
#define HEIGHT 8
#define NUM (WIDTH * HEIGHT)

#define THREADS_PER_BLOCK_X 8
#define THREADS_PER_BLOCK_Y 8
#define THREADS_PER_BLOCK_Z 1

unsigned int firstbit_u32(unsigned int a) {
  if (a == 0) {
    return 32;
  }
  unsigned int pos = 0;
  while ((int)a > 0) {
    a <<= 1;
    pos++;
  }
  return pos;
}

unsigned int firstbit_u64(unsigned long long int a) {
  if (a == 0) {
    return 64;
  }
  unsigned int pos = 0;
  while ((long long int)a > 0) {
    a <<= 1;
    pos++;
  }
  return pos;
}

// Check implicit conversion will not cause ambiguity.
__device__ void test_ambiguity() {
  short s = 1;
  unsigned short us = 2;
  float f = 3;
  int i = 4;
  unsigned int ui = 5;
  __clz(f);
  __clz(s);
  __clz(us);
  __clzll(f);
  __clzll(i);
  __clzll(ui);
}

__global__ void clz_HIP_kernel(unsigned int* a, unsigned int* b, unsigned int* c,
                               unsigned long long int* d, int width, int height) {
  int x = blockDim.x * blockIdx.x + threadIdx.x;
  int y = blockDim.y * blockIdx.y + threadIdx.y;

  int i = y * width + x;
  if (i < (width * height)) {
    a[i] = __clz(b[i]);
    c[i] = __clzll(d[i]);
  }
}

HIP_TEST_CASE(Unit_clz) {
  using namespace std;

  unsigned int* hostA;
  unsigned int* hostB;
  unsigned int* hostC;
  unsigned long long int* hostD;

  unsigned int* deviceA;
  unsigned int* deviceB;
  unsigned int* deviceC;
  unsigned long long int* deviceD;

  hipDeviceProp_t devProp;
  HIP_CHECK(hipGetDeviceProperties(&devProp, 0));
  INFO("System minor : " << devProp.minor);
  INFO("System major : " << devProp.major);
  INFO("agent prop name : " << devProp.name);

  INFO("hip Device prop succeeded");

  unsigned int i;
  int errors;

  hostA = (unsigned int*)malloc(NUM * sizeof(unsigned int));
  hostB = (unsigned int*)malloc(NUM * sizeof(unsigned int));
  hostC = (unsigned int*)malloc(NUM * sizeof(unsigned int));
  hostD = (unsigned long long int*)malloc(NUM * sizeof(unsigned long long int));

  // initialize the input data
  for (i = 0; i < NUM; i++) {
    hostB[i] = 419430 * i;
    hostD[i] = i;
  }

  HIP_CHECK(hipMalloc((void**)&deviceA, NUM * sizeof(unsigned int)));
  HIP_CHECK(hipMalloc((void**)&deviceB, NUM * sizeof(unsigned int)));
  HIP_CHECK(hipMalloc((void**)&deviceC, NUM * sizeof(unsigned int)));
  HIP_CHECK(hipMalloc((void**)&deviceD, NUM * sizeof(unsigned long long int)));

  HIP_CHECK(hipMemcpy(deviceB, hostB, NUM * sizeof(unsigned int), hipMemcpyHostToDevice));
  HIP_CHECK(hipMemcpy(deviceD, hostD, NUM * sizeof(unsigned long long int), hipMemcpyHostToDevice));

  hipLaunchKernelGGL(clz_HIP_kernel,
                     dim3(WIDTH / THREADS_PER_BLOCK_X, HEIGHT / THREADS_PER_BLOCK_Y),
                     dim3(THREADS_PER_BLOCK_X, THREADS_PER_BLOCK_Y), 0, 0, deviceA, deviceB,
                     deviceC, deviceD, WIDTH, HEIGHT);
  HIP_CHECK(hipGetLastError());


  HIP_CHECK(hipMemcpy(hostA, deviceA, NUM * sizeof(unsigned int), hipMemcpyDeviceToHost));
  HIP_CHECK(hipMemcpy(hostC, deviceC, NUM * sizeof(unsigned int), hipMemcpyDeviceToHost));

  // verify the results
  errors = 0;
  for (i = 0; i < NUM; i++) {
    if (hostA[i] != firstbit_u32(hostB[i])) {
      INFO("Match Failed: " << hostA[i] << " - " << firstbit_u32(hostB[i]));
      errors++;
    }
  }

  for (i = 0; i < NUM; i++) {
    if (hostC[i] != firstbit_u64(hostD[i])) {
      INFO("Match Failed: " << hostC[i] << " - " << firstbit_u32(hostD[i]));
      errors++;
    }
  }


  HIP_CHECK(hipFree(deviceA));
  HIP_CHECK(hipFree(deviceB));
  HIP_CHECK(hipFree(deviceC));
  HIP_CHECK(hipFree(deviceD));

  free(hostA);
  free(hostB);
  free(hostC);
  free(hostD);

  REQUIRE(errors == 0);
}