/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

/*!
 * \file ndarray_function.cu
 * \brief GPU Implementation of ndarray function.
 */

// this will be invoked by nvcc and compile GPU version
#include <cub/cub.cuh>
#include <dmlc/logging.h>
#include "../operator/tensor/elemwise_binary_op-inl.h"
#include "../operator/tensor/elemwise_sum.h"
#include "../operator/tensor/indexing_op.h"
#include "../operator/tensor/init_op.h"
#include "../operator/tensor/util/tensor_util-inl.h"
#include "../operator/tensor/util/tensor_util-inl.cuh"
#include "../common/cuda/utils.h"
#include "./ndarray_function.h"
#include "./ndarray_function-inl.h"
#include "./ndarray_function-inl.cuh"

namespace mxnet {
namespace ndarray {
template <>
void Copy<cpu, gpu>(const TBlob& from,
                    TBlob* to,
                    Context from_ctx,
                    Context to_ctx,
                    RunContext ctx) {
  CHECK_EQ(to->type_flag_, from.type_flag_)
      << "Source and target must have the same data type when copying across devices.";
  MSHADOW_TYPE_SWITCH_EXT_WITH_BOOL(to->type_flag_, DType, {
    mshadow::Copy(to->FlatTo1D<gpu, DType>(), from.FlatTo1D<cpu, DType>(), ctx.get_stream<gpu>());
  });
}

template <>
void Copy<gpu, cpu>(const TBlob& from,
                    TBlob* to,
                    Context from_ctx,
                    Context to_ctx,
                    RunContext ctx) {
  CHECK_EQ(to->type_flag_, from.type_flag_)
      << "Source and target must have the same data type when copying across devices.";
  MSHADOW_TYPE_SWITCH_EXT_WITH_BOOL(to->type_flag_, DType, {
    mshadow::Copy(to->FlatTo1D<cpu, DType>(), from.FlatTo1D<gpu, DType>(), ctx.get_stream<gpu>());
  });
}

template <>
void Copy<gpu, gpu>(const TBlob& from,
                    TBlob* to,
                    Context from_ctx,
                    Context to_ctx,
                    RunContext ctx) {
  if (from_ctx.dev_id == to_ctx.dev_id) {
    mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
    MSHADOW_TYPE_SWITCH_WITH_BOOL(to->type_flag_, DType, {
      if (to->type_flag_ == from.type_flag_) {
        mshadow::Copy(to->FlatTo1D<gpu, DType>(s), from.FlatTo1D<gpu, DType>(s), s);
      } else {
        MSHADOW_TYPE_SWITCH_WITH_BOOL(from.type_flag_, SrcDType, {
          to->FlatTo1D<gpu, DType>(s) =
              mshadow::expr::tcast<DType>(from.FlatTo1D<gpu, SrcDType>(s));
        })
      }
    })
  } else {
    CHECK(from.CheckContiguous() && to->CheckContiguous())
        << "copy across only support contiguous memory";
    CHECK_EQ(to->type_flag_, from.type_flag_)
        << "Source and target must have the same data type when copying across devices.";
    mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
    CHECK(s != nullptr) << "need stream in GPU context";
    cudaMemcpyPeerAsync(to->dptr_,
                        to_ctx.dev_id,
                        from.dptr_,
                        from_ctx.dev_id,
                        from.shape_.Size() * mshadow::mshadow_sizeof(to->type_flag_),
                        s->stream_);
  }
}

/*!
 * \brief GPU impl of elemwise sum for rowsparse tensors.
 */
void ElementwiseSumRspImpl(mshadow::Stream<gpu>* s,
                           const Resource& rsc,
                           const std::vector<NDArray>& nds,
                           NDArray* out) {
  using namespace mxnet::op;
  using namespace rowsparse;
  using nnvm::dim_t;
  CHECK_EQ(out->storage_type(), kRowSparseStorage)
      << "Expected rowsparse storage_type (" << out->storage_type() << " given)";
  int init = 0;
  for (const auto& nd : nds) {
    if (nd.storage_initialized()) {
      init++;
      break;
    }
  }
  if (init == 0) {
    FillZerosRspImpl(s, *out);
    return;
  }
  const dim_t num_rows   = out->shape()[0];
  const dim_t row_length = out->shape().ProdShape(1, out->shape().ndim());
  MSHADOW_TYPE_SWITCH(out->dtype(), DType, {               // data type
    MSHADOW_IDX_TYPE_SWITCH(out->aux_type(kIdx), IType, {  // row_idx type
      // Allocate temporary storage for row_flg array and cub's prefix sum operation
      IType* row_flg            = nullptr;
      void* d_temp_storage      = nullptr;
      size_t temp_storage_bytes = 0;
      cudaStream_t stream       = mshadow::Stream<gpu>::GetStream(s);
      cub::DeviceScan::InclusiveSum(
          d_temp_storage, temp_storage_bytes, row_flg, row_flg, num_rows, stream);
      mshadow::Tensor<gpu, 1, char> workspace = rsc.get_space_typed<gpu, 1, char>(
          mshadow::Shape1(num_rows * sizeof(IType) + temp_storage_bytes), s);
      row_flg        = reinterpret_cast<IType*>(workspace.dptr_);
      d_temp_storage = workspace.dptr_ + num_rows * sizeof(IType);
      // Mark row_flg array with 0 for zero rows and 1 for non-zero rows
      dim_t num_threads = num_rows;
      mxnet_op::Kernel<mxnet_op::set_zero, gpu>::Launch(s, num_threads, row_flg);
      for (const auto& nd : nds) {
        if (nd.storage_initialized()) {
          const IType* nd_row_idx = nd.aux_data(kIdx).dptr<IType>();
          const dim_t nd_nnr      = nd.storage_shape()[0];
          num_threads             = nd_nnr;
          mxnet_op::Kernel<MarkRspRowFlgKernel, gpu>::Launch(
              s, num_threads, row_flg, nd_row_idx, nd_nnr);
        }
      }
      // Compute inclusive prefix sum over row_flg
      cub::DeviceScan::InclusiveSum(
          d_temp_storage, temp_storage_bytes, row_flg, row_flg, num_rows, stream);
      // Get total number of output non-zero rows from GPU and allocate out data and row_idx
      dim_t nnr_out = 0;
      CUDA_CALL(cudaMemcpyAsync(
          &nnr_out, &row_flg[num_rows - 1], sizeof(dim_t), cudaMemcpyDeviceToHost, stream));
      CUDA_CALL(cudaStreamSynchronize(stream));
      out->CheckAndAlloc({mshadow::Shape1(nnr_out)});
      IType* out_row_idx = out->aux_data(kIdx).dptr<IType>();
      DType* out_data    = out->data().dptr<DType>();
      // Fill row_idx array of output using row_flg
      num_threads = num_rows;
      mxnet_op::Kernel<FillRspRowIdxKernel, gpu>::Launch(
          s, num_threads, out_row_idx, row_flg, num_rows);
      // Perform elementwise addition, writing to output data
      num_threads = nnr_out * row_length;
      mxnet_op::Kernel<mxnet_op::set_zero, gpu>::Launch(s, num_threads, out_data);
      for (const auto& nd : nds) {
        if (nd.storage_initialized()) {
          const IType* nd_row_idx = nd.aux_data(kIdx).dptr<IType>();
          const DType* nd_data    = nd.data().dptr<DType>();
          const dim_t nd_nnr      = nd.storage_shape()[0];
          num_threads             = nd_nnr * row_length;
          mxnet_op::Kernel<ElementWiseRspAdditionKernel, gpu>::Launch(
              s, num_threads, out_data, row_flg, nd_row_idx, nd_data, nd_nnr, row_length);
        }
      }
    });
  });
}

void ElementwiseSumDnsCsrDnsImpl(mshadow::Stream<gpu>* s,
                                 const Resource& rsc,
                                 const std::vector<NDArray>& nds,
                                 NDArray* out) {
  using namespace mxnet::op;
  using namespace mxnet::op::mxnet_op;
  const TBlob& out_data = out->data();
  MSHADOW_TYPE_SWITCH(out->dtype(), DType, {  // data type
    Kernel<Sum, gpu>::Launch(s,
                             out_data.Size(),
                             out_data.dptr<DType>(),
                             kWriteTo,
                             nds[0].data().dptr<DType>(),
                             nds[2].data().dptr<DType>());
    const TBlob& csr_data      = nds[1].data();
    const TBlob& csr_indices   = nds[1].aux_data(csr::kIdx);
    const TBlob& csr_indptr    = nds[1].aux_data(csr::kIndPtr);
    const nnvm::dim_t num_rows = nds[1].shape()[0];
    const nnvm::dim_t num_cols = nds[1].shape()[1];
    MSHADOW_IDX_TYPE_SWITCH(csr_indices.type_flag_, IType, {   // indices type
      MSHADOW_IDX_TYPE_SWITCH(csr_indptr.type_flag_, CType, {  // indptr type
        if (nds[1].storage_initialized()) {
          Kernel<ElemwiseDnsCsrDnsWarpKernel<kWriteTo, mshadow_op::plus>, gpu>::Launch(
              s,
              kWarpSize * num_rows,
              out_data.dptr<DType>(),
              out_data.dptr<DType>(),
              csr_data.dptr<DType>(),
              csr_indices.dptr<IType>(),
              csr_indptr.dptr<CType>(),
              num_rows,
              num_cols);
        }
      });
    });
  });
}

void ElementwiseSumContainsDnsImpl(mshadow::Stream<gpu>* s,
                                   const Resource& rsc,
                                   const std::vector<NDArray>& nds,
                                   NDArray* out) {
  using namespace mxnet::op;
  using namespace mxnet::op::mxnet_op;
  const TBlob& out_data = out->data();
  MSHADOW_TYPE_SWITCH(out->dtype(), DType, {  // data type
    for (size_t i = 0; i < nds.size(); ++i) {
      const NDArray& nd          = nds[i];
      const nnvm::dim_t num_rows = nd.shape()[0];
      const nnvm::dim_t num_cols = nd.shape()[1];
      const TBlob& nd_data       = nd.data();

      if (i == 0) {
        if (nd.storage_type() == kDefaultStorage) {
          Kernel<op_with_req<mshadow_op::identity, kWriteTo>, gpu>::Launch(
              s, out_data.Size(), out_data.dptr<DType>(), nd_data.dptr<DType>());
          continue;
        } else {
          Kernel<set_zero, gpu>::Launch(s, out_data.Size(), out_data.dptr<DType>());
        }
      }

      switch (nd.storage_type()) {
        case kDefaultStorage: {
          Kernel<op_with_req<mshadow_op::plus, kWriteTo>, gpu>::Launch(s,
                                                                       out_data.Size(),
                                                                       out_data.dptr<DType>(),
                                                                       out_data.dptr<DType>(),
                                                                       nd_data.dptr<DType>());
          break;
        }
        case kCSRStorage: {
          const TBlob& nd_indices = nd.aux_data(csr::kIdx);
          const TBlob& nd_indptr  = nd.aux_data(csr::kIndPtr);
          MSHADOW_IDX_TYPE_SWITCH(nd_indices.type_flag_, IType, {   // indices type
            MSHADOW_IDX_TYPE_SWITCH(nd_indptr.type_flag_, CType, {  // indptr type
              if (nd.storage_initialized()) {
                Kernel<ElemwiseDnsCsrDnsWarpKernel<kWriteTo, mshadow_op::plus>, gpu>::Launch(
                    s,
                    kWarpSize * num_rows,
                    out_data.dptr<DType>(),
                    out_data.dptr<DType>(),
                    nd_data.dptr<DType>(),
                    nd_indices.dptr<IType>(),
                    nd_indptr.dptr<CType>(),
                    num_rows,
                    num_cols);
              }
            });
          });
          break;
        }
        case kRowSparseStorage: {
          const TBlob& nd_indices = nd.aux_data(rowsparse::kIdx);
          MSHADOW_IDX_TYPE_SWITCH(nd_indices.type_flag_, IType, {  // indices type
            if (nd.storage_initialized()) {
              const nnvm::dim_t nz_rows = nd_indices.Size();
              Kernel<ElemwiseDnsRspDnsKernel<kWriteTo, mshadow_op::plus>, gpu>::Launch(
                  s,
                  nz_rows * num_cols,
                  out_data.dptr<DType>(),
                  out_data.dptr<DType>(),
                  nd_data.dptr<DType>(),
                  nd_indices.dptr<IType>(),
                  num_rows,
                  nz_rows,
                  num_cols);
            }
          });
          break;
        }
        default:
          LOG(FATAL) << "unknown storage type " << nd.storage_type() << "encountered...";
      }
    }
  });
}

/*!
 * \brief Parallel gpu impl of elemwise sum for sparse tensors.
 * Currently only support row sparse sum.
 */
template <>
void ElementwiseSum<gpu>(mshadow::Stream<gpu>* s,
                         const Resource& rsc,
                         const std::vector<NDArray>& nds,
                         NDArray* out) {
  if (nds.empty())
    return;
  if (common::ContainsOnlyStorage(nds, kRowSparseStorage)) {
    ElementwiseSumRspImpl(s, rsc, nds, out);
  } else if (nds.size() == 3U && nds[0].storage_type() == kDefaultStorage &&
             nds[1].storage_type() == kCSRStorage && nds[2].storage_type() == kDefaultStorage &&
             out->storage_type() == kDefaultStorage) {
    ElementwiseSumDnsCsrDnsImpl(s, rsc, nds, out);
  } else if (nds.size() > 4U && common::ContainsStorageType(nds, kDefaultStorage) &&
             out->storage_type() == kDefaultStorage) {
    ElementwiseSumContainsDnsImpl(s, rsc, nds, out);
  } else {
    LOG(FATAL) << "ElementwiseSum<gpu> has not been implemented for storage_type = << "
               << nds[0].storage_type();
  }
}

template <>
void Eval<gpu>(mshadow::Stream<gpu>* s, const real_t val, const NDArray& dst) {
  NDArray temp                   = dst;
  const NDArrayStorageType stype = temp.storage_type();
  if (stype == kRowSparseStorage) {
    SetValueRspImpl(s, val, &temp);
  } else {
    LOG(FATAL) << "Not implemented for storage type" << stype;
  }
}

}  // namespace ndarray
}  // namespace mxnet