// Copyright 2019 Bytedance Inc. or its affiliates. All Rights Reserved.
// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// 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 "tensor_util.h"

namespace byteps {
namespace mxnet {

// Define all types for TensorUtil.
const DataType TensorUtil::GetDType(NDArray* tensor) {
  switch (tensor->dtype()) {
    case mshadow::kFloat32:
      return DataType::BYTEPS_FLOAT32;
    case mshadow::kFloat64:
      return DataType::BYTEPS_FLOAT64;
    case mshadow::kFloat16:
      return DataType::BYTEPS_FLOAT16;
    case mshadow::kUint8:
      return DataType::BYTEPS_UINT8;
    case mshadow::kInt32:
      return DataType::BYTEPS_INT32;
    case mshadow::kInt8:
      return DataType::BYTEPS_INT8;
    case mshadow::kInt64:
      return DataType::BYTEPS_INT64;
    default:
      throw std::logic_error("GetDType: Type " +
                             std::to_string(tensor->dtype()) +
                             " is not supported.");
  }
}

// Return shape of tensor (similar to TShape)
const TensorShape TensorUtil::GetShape(NDArray* tensor) {
  TensorShape shape;
  TShape mx_shape = tensor->shape();
  for (unsigned idx = 0; idx < mx_shape.ndim(); idx++) {
    shape.AddDim(mx_shape[idx]);
  }
  return shape;
}

// Return data of tensor
const void* TensorUtil::GetData(NDArray* tensor) {
  // The following returns an error:
  // return tensor->data().dptr<void>();
  switch (tensor->dtype()) {
    case mshadow::kFloat32:
      return static_cast<void*>(tensor->data().dptr<float>());
    case mshadow::kFloat64:
      return static_cast<void*>(tensor->data().dptr<double>());
    case mshadow::kFloat16:
      return static_cast<void*>(tensor->data().dptr<mshadow::half::half_t>());
    case mshadow::kUint8:
      return static_cast<void*>(tensor->data().dptr<uint8_t>());
    case mshadow::kInt32:
      return static_cast<void*>(tensor->data().dptr<int32_t>());
    case mshadow::kInt8:
      return static_cast<void*>(tensor->data().dptr<int8_t>());
    case mshadow::kInt64:
      return static_cast<void*>(tensor->data().dptr<int64_t>());
    default:
      throw std::logic_error("Type " + std::to_string(tensor->dtype()) +
                             " is not supported in BytePS.");
  }
}

// Return size of tensor in bytes
int64_t TensorUtil::GetSize(NDArray* tensor) {
  int64_t element_size = 0;
  switch (tensor->dtype()) {
    case mshadow::kFloat32:
      element_size = kFloat32Size;
      break;
    case mshadow::kFloat64:
      element_size = kFloat64Size;
      break;
    case mshadow::kFloat16:
      element_size = kFloat16Size;
      break;
    case mshadow::kUint8:
      element_size = kUInt8Size;
      break;
    case mshadow::kInt32:
      element_size = kInt32Size;
      break;
    case mshadow::kInt8:
      element_size = kInt8Size;
      break;
    case mshadow::kInt64:
      element_size = kInt64Size;
      break;
    default:
      throw std::logic_error("Type " + std::to_string(tensor->dtype()) +
                             " is not supported in BytePS.");
  }
  return (int64_t)(tensor->shape().Size()) * element_size;
}

// If Tensor on GPU, return device id
// Otherwise return CPU_DEVICE_ID (-1)
int TensorUtil::GetDevice(NDArray* tensor) {
  int dev_mask = tensor->ctx().dev_mask();
  if (dev_mask == gpu::kDevMask) return tensor->ctx().real_dev_id();
  return CPU_DEVICE_ID;
}

// Returns pointer to newly created NDArray
// If dev_id equal to CPU_DEVICE_ID, construct Tensor on CPU
// Otherwise construct on GPU
NDArray* TensorUtil::New(int device, int dtype) {
  NDArrayHandle array;
  int dev_type = gpu::kDevMask;
  if (device == CPU_DEVICE_ID) {
    dev_type = cpu::kDevMask;
    device = 0;
  }
  MXNDArrayCreateEx(nullptr, 0, dev_type, device, true, dtype, &array);
  return static_cast<NDArray*>(array);
}

void TensorUtil::Free(NDArray* tensor) { delete tensor; }

// Copy from tensor to output
void TensorUtil::Copy(NDArray* output, NDArray* tensor) {
  if (tensor->shape() != output->shape())
    output->ReshapeAndAlloc(tensor->shape());
  CopyFromTo(*tensor, output, 0);
}

// Elementwise division of tensor by value in-place
void TensorUtil::DivideTensorInPlace(NDArray* tensor, int value) {
  *tensor /= value;
}

#ifdef HAVE_CUDA
void TensorUtil::CopyCPUToCuda(NDArray* cpu, NDArray* cuda) {
  TensorUtil::Copy(cuda, cpu);
}

void TensorUtil::AsyncCopyCudaToCPU(NDArray* cuda, NDArray* cpu) {
  TensorUtil::Copy(cpu, cuda);
}
#endif

}  // namespace mxnet
}  // namespace byteps
