text: /*
   * 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(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(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(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(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 continugous 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 != NULL) << "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 = NULL;
        void* d_temp_storage = NULL;
        size_t temp_storage_bytes = 0;
        cub::DeviceScan::InclusiveSum(d_temp_storage,
                                      temp_storage_bytes,
                                      row_flg,
                                      row_flg,
                                      num_rows,
                                      mshadow::Stream<gpu>::GetStream(s));
        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,
                                      mshadow::Stream<gpu>::GetStream(s));
        // Get total number of output non-zero rows from GPU and allocate out data and row_idx
        dim_t nnr_out = 0;
        CUDA_CALL(cudaMemcpy(&nnr_out, &row_flg[num_rows-1], sizeof(dim_t),
                             cudaMemcpyDeviceToHost));
        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
  