text: /*
   * Copyright (c) 2020, NVIDIA CORPORATION.
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  #include <diagnose.hpp>
  #include <limits>
  #include <utils.cuh>
  
  namespace HugeCTR {
  
  namespace diagnose {
  
  __device__ float atomicMin(float* address, float val) {
    float old = val;
    do {
      val = old;
      old = atomicExch(address, val);
    } while (old < val);
    return old;
  }
  
  __device__ float atomicMax(float* address, float val) {
    float old = val;
    do {
      val = old;
      old = atomicExch(address, val);
    } while (old > val);
    return old;
  }
  
  template <typename T>
  __global__ void histogram_kernel(const T* arr, size_t len, float* range) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < len; i += blockDim.x * gridDim.x) {
      float val = TypeConvertFunc<float, T>::convert(arr[i]);
      if (val <= 0) {
        atomicMin(range + 0, val);
        atomicMax(range + 1, val);
      }
      if (val >= 0) {
        atomicMin(range + 2, val);
        atomicMax(range + 3, val);
      }
    }
  }
  
  template <typename T>
  __global__ void verify_kernel(const T* arr, size_t len, int* flag);
  
  template <>
  __global__ void verify_kernel<float>(const float* arr, size_t len, int* flag) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < len; i += blockDim.x * gridDim.x) {
      if (isnan(arr[i])) atomicAdd(flag, 1);
    }
  }
  
  template <>
  __global__ void verify_kernel(const __half* arr, size_t len, int* flag) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < len; i += blockDim.x * gridDim.x) {
      if (__hisnan(arr[i])) {
        atomicAdd(flag, 1);
      }
    }
  }
  
  template <typename T>
  __global__ void sample_kernel(const T* arr, int len, float* arr_sample, int stride,
                                int max_sample_len) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < len; i += blockDim.x * gridDim.x) {
      if (i % stride == 0) {
        int j = i / stride;
        if (j < max_sample_len) {
          arr_sample[j] = TypeConvertFunc<float, T>::convert(arr[i]);
        }
      }
    }
  }
  
  template <typename T>
  void verify_and_histogram(const char* category, const Tensor2<T>& tensor,
                            const cudaStream_t& stream) {
    float h_array[4]{0.0f, -std::numeric_limits<float>::infinity(),
                     std::numeric_limits<float>::infinity(), 0.0f};
    int h_flag;
    float* d_array;
    int* d_flag;
    CK_CUDA_THROW_(cudaMalloc(&d_array, sizeof(h_array)));
    CK_CUDA_THROW_(cudaMalloc(&d_flag, sizeof(int)));
    CK_CUDA_THROW_(
        cudaMemcpyAsync(d_array, h_array, sizeof(h_array), cudaMemcpyHostToDevice, stream));
    CK_CUDA_THROW_(cudaMemsetAsync(d_flag, 0, sizeof(int), stream));
    histogram_kernel<<<160, 1024, 0, stream>>>(tensor.get_ptr(), tensor.get_num_elements(), d_array);
    verify_kernel<<<160, 1024, 0, stream>>>(tensor.get_ptr(), tensor.get_num_elements(), d_flag);
    CK_CUDA_THROW_(
        cudaMemcpyAsync(h_array, d_array, sizeof(h_array), cudaMemcpyDeviceToHost, stream));
    CK_CUDA_THROW_(cudaMemcpyAsync(&h_flag, d_flag, sizeof(int), cudaMemcpyDeviceToHost, stream));
    CK_CUDA_THROW_(cudaStreamSynchronize(stream));
  
    std::stringstream ss;
    ss << "Diagnose for (" << category << "), Histogram [" << h_array[0] << ", " << h_array[1] << "]"
       << ", [" << h_array[2] << ", " << h_array[3] << "]" << std::endl;
    MESSAGE_(ss.str());
  
    if (h_flag != 0) {
      CK_THROW_(Error_t::DataCheckError, std::string("Nan assert for ") + category + " failed(" +
                                             std::to_string(h_flag) + ").");
    }
    CK_CUDA_THROW_(cudaFree(d_array));
    CK_CUDA_THROW_(cudaFree(d_flag));
  }
  
  template <typename T>
  void sample_and_print(const char* category, const Tensor2<T>& tensor, size_t sample_count,
                        const cudaStream_t& stream) {
    if (sample_count == 0) return;
  
    std::unique_ptr<float[]> h_array(new float[sample_count]);
  
    float* d_array;
    CK_CUDA_THROW_(cudaMalloc(&d_array, sample_count * sizeof(float)));
    CK_CUDA_THROW_(cudaMemsetAsync(d_array, 0, sample_count * sizeof(float), stream));
    sample_kernel<<<160, 1024, 0, stream>>>(tensor.get_ptr(), tensor.get_num_elements(), d_array,
                                            tensor.get_num_elements() / sample_count, sample_count);
    CK_CUDA_THROW_(cudaMemcpyAsync(h_array.get(), d_array, sample_count * sizeof(float),
                                   cudaMemcpyDeviceToHost, stream));
    CK_CUDA_THROW_(cudaStreamSynchronize(stream));
  
    std::stringstream ss;
    ss << "Diagnose for (" << category << "), Sampling [";
    for (size_t i = 0; i < min(sample_count, tensor.get_num_elements()); i++) {
      if (i != 0) ss << ",";
      ss << h_array[i];
    }
    ss << "]" << std::endl;
    MESSAGE_(ss.str());
  
    CK_CUDA_THROW_(cudaFree(d_array));
  }
  
  template <typename T>
  void sample_and_print(const char* category, const Tensor2<T>& tensor, int begin, int end,
                        const cudaStream_t& stream) {
    if (begin >= 0 && end <= static_cast<int>(tensor.get_num_elements()) && end > begin) {
    } else if (end < 0 && begin >= -static_cast<int>(tensor.get_num_elements()) && end > begin) {
      begin += tensor.get_num_elements();
      end += tensor.get_num_elements();
    } else {
      return;
    }
  
    std::unique_ptr<T[]> h_array(new T[end - begin]);
    CK_CUDA_THROW_(cudaMemcpyAsync(h_array.get(), tensor.get_ptr() + begin,
                                   (begin - end) * sizeof(float), cudaMemcpyDeviceToHost, stream));
    CK_CUDA_THROW_(cudaStreamSynchronize(stream));
  
    std::stringstream ss;
    ss << "Diagnose for (" << category << "), Sampling [";
    for (size_t i = 0; i < end - begin; i++) {
      if (i != 0) ss << ",";
      ss << h_array[i];
    }
    ss << "]" << std::endl;
    MESSAGE_(ss.str());
  }
  
  template void verify_and_histogram<float>(const char* category, const Tensor2<float>& tensor,
                                            const cudaStream_t& stream);
  
  }  // namespace diagnose
  
  }  // namespace HugeCTR