// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "core/providers/cuda/math/clip_impl.h"
#include "core/providers/cuda/cu_inc/common.cuh"

namespace onnxruntime {
namespace cuda {
template <typename T>
__global__ void _Clip(const T* input, T* output, const T* min, const T* max, T min_default, T max_default, size_t N) {
  auto min_val = (min) ? *min : min_default;
  auto max_val = (max) ? *max : max_default;
  CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);

  // output = Min(max, Max(input, min)). Note that min might be larger than max, so we need to compute in two steps.
  auto value = input[id];
  value = (value < min_val) ? min_val : value;
  output[id] = (value > max_val) ? max_val : value;
}

template <typename T>
void ClipImpl(cudaStream_t stream, const T* input_data, T* output_data, const T* min, const T* max, T min_default, T max_default, size_t count) {
  typedef typename ToCudaType<T>::MappedType CudaT;

  int blocksPerGrid = (int)(ceil(static_cast<float>(count) / GridDim::maxThreadsPerBlock));
  union ConstAliasUnion {
    const T* t;
    const CudaT* cudaT;
    ConstAliasUnion(const T* _t) { t = _t; }
  };
  union AliasUnion {
    T* t;
    CudaT* cudaT;
    AliasUnion(T* _t) { t = _t; }
  };
  _Clip<CudaT><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(((union ConstAliasUnion)input_data).cudaT,
                                                                          ((union AliasUnion)output_data).cudaT,
                                                                          ((union ConstAliasUnion)min).cudaT,
                                                                          ((union ConstAliasUnion)max).cudaT,
                                                                          *((union AliasUnion) & min_default).cudaT,
                                                                          *((union AliasUnion) & max_default).cudaT,
                                                                          count);
}

template void ClipImpl<float>(cudaStream_t stream, const float* input_data, float* output_data, const float* min, const float* max, float min_default, float max_default, size_t count);
template void ClipImpl<double>(cudaStream_t stream, const double* input_data, double* output_data, const double* min, const double* max, double min_default, double max_default, size_t count);
template void ClipImpl<MLFloat16>(cudaStream_t stream, const MLFloat16* input_data, MLFloat16* output_data, const MLFloat16* min, const MLFloat16* max, MLFloat16 min_default, MLFloat16 max_default, size_t count);
template void ClipImpl<int8_t>(cudaStream_t stream, const int8_t* input_data, int8_t* output_data, const int8_t* min, const int8_t* max, int8_t min_default, int8_t max_default, size_t count);
template void ClipImpl<uint8_t>(cudaStream_t stream, const uint8_t* input_data, uint8_t* output_data, const uint8_t* min, const uint8_t* max, uint8_t min_default, uint8_t max_default, size_t count);
template void ClipImpl<int64_t>(cudaStream_t stream, const int64_t* input_data, int64_t* output_data, const int64_t* min, const int64_t* max, int64_t min_default, int64_t max_default, size_t count);
template void ClipImpl<uint64_t>(cudaStream_t stream, const uint64_t* input_data, uint64_t* output_data, const uint64_t* min, const uint64_t* max, uint64_t min_default, uint64_t max_default, size_t count);

}  // namespace cuda
}  // namespace onnxruntime
