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 cast_storage-inl.h
   * \brief CPU implementation of cast_storage operator for dense and sparse tensors
   */
  #ifndef MXNET_OPERATOR_TENSOR_CAST_STORAGE_INL_H_
  #define MXNET_OPERATOR_TENSOR_CAST_STORAGE_INL_H_
  
  #include <dmlc/timer.h>
  #include <mxnet/ndarray.h>
  #include <vector>
  #include <algorithm>
  #include "../mxnet_op.h"
  #include "../operator_common.h"
  #include "../../src/operator/tensor/init_op.h"
  #ifdef __CUDACC__
  #include "./cast_storage-inl.cuh"
  #endif  // __CUDACC__
  #if MXNET_USE_MKLDNN == 1
  #include "../nn/mkldnn/mkldnn_base-inl.h"
  #endif
  
  
  namespace mxnet {
  namespace op {
  
  /*!
   * \brief CPU Kernel for marking row_idx of a RSP tensor per row.
   */
  struct MarkRspRowIdx {
    // i represents the row index of the tensor data
    template<typename DType, typename RType>
    MSHADOW_CINLINE static void Map(int i,
                                    RType* row_idx,
                                    const DType* data,
                                    const nnvm::dim_t row_length) {
      using nnvm::dim_t;
      dim_t j = 0;
      dim_t offset = i * row_length;
      for (; j < row_length; ++j) {
        if (data[offset+j] != 0) {
          break;
        }
      }
      if (row_length == j) {
        row_idx[i] = 0;  // mark as zero for zero row
      } else {
        row_idx[i] = 1;  // mark as one for non-zero row
      }
    }
  };
  
  /*!
   * \brief CPU implementation of casting a dns tensor to rsp type.
   */
  inline void CastStorageDnsRspImpl(const OpContext& ctx,
                                    const cpu& cpu_dev,
                                    const TBlob& dns,
                                    NDArray* rsp) {
    using namespace rowsparse;
    using namespace mshadow;
    using nnvm::dim_t;
    CHECK(rsp != nullptr);
    CHECK_EQ(rsp->storage_type(), kRowSparseStorage);
    CHECK_EQ(dns.shape_, rsp->shape());
    mshadow::Stream<cpu>* s = ctx.get_stream<cpu>();
    MSHADOW_TYPE_SWITCH(dns.type_flag_, DType, {  // data type
      MSHADOW_IDX_TYPE_SWITCH(rsp->aux_type(kIdx), RType, {  // row idx type
        const dim_t num_rows = dns.shape_[0];
        const dim_t row_length = dns.shape_.ProdShape(1, dns.shape_.ndim());
        rsp->CheckAndAllocAuxData(kIdx, Shape1(num_rows));
        TBlob row_idx_blob = rsp->aux_data(kIdx);
        RType* row_idx = row_idx_blob.dptr<RType>();
        dim_t num_threads = num_rows;
        mxnet_op::Kernel<MarkRspRowIdx, cpu>::Launch(s, num_threads,
            row_idx, dns.dptr<DType>(), row_length);
        dim_t nnr = 0;
        nnr = common::ParallelAccumulate(row_idx, num_rows, nnr);
        rsp->set_aux_shape(kIdx, Shape1(nnr));
        if (0 == nnr) return;
        auto storage_shape = dns.shape_;
        storage_shape[0] = nnr;
        rsp->CheckAndAllocData(storage_shape);
        auto dns_data = dns.get_with_shape<cpu, 2, DType>(Shape2(num_rows, row_length), s);
        auto rsp_data = rsp->data().get_with_shape<cpu, 2, DType>(Shape2(nnr, row_length), s);
        dim_t idx = 0;
        for (dim_t i = 0; i < num_rows; ++i) {
          if (row_idx[i] > 0) {
            row_idx[idx] = i;
            Copy(rsp_data[idx], dns_data[i], s);
            ++idx;
          }
        }
      });
    });
  }
  
  // TODO(haibin) Use memcopy instead will be much faster than assigning each individual element
  struct CastStorageRspDnsKernel {
    template<typename DType, typename IType>
    MSHADOW_XINLINE static void Map(int i,
                                    const nnvm::dim_t row_length,
                                    const IType* idx,
                                    const DType *data,
                                    DType* dns) {
      using nnvm::dim_t;
      IType rid = idx[i];
      dim_t dns_offset = rid * row_length;
      dim_t rsp_offset = i * row_length;
      for (dim_t col = 0; col < row_length; col++) {
        dns[dns_offset + col] = data[rsp_offset + col];
      }
    }
  };
  
  /*!
   * \brief This function assumes that the memory for dns has been allocated already
   * since the shape is known at binding stage.
   */
  template<typename xpu>
  void CastStorageRspDnsImpl(const OpContext& ctx,
                             const NDArray& rsp,
                             TBlob* dns) {
    mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
    CHECK_EQ(rsp.storage_type(), kRowSparseStorage);
    using nnvm::dim_t;
    MSHADOW_TYPE_SWITCH(dns->type_flag_, DType, {
      MSHADOW_IDX_TYPE_SWITCH(rsp.aux_type(rowsparse::kIdx), IType, {
        // assign zeros
        mxnet_op::Kernel<mxnet_op::set_zero, xpu>::Launch(s, dns->Size(), dns->dptr<DType>());
        if (rsp.storage_initialized()) {
          // copy over row by row
          auto in_idx = rsp.aux_data(rowsparse::kIdx).FlatTo1D<xpu, IType>(s).dptr_;
          auto in_data = rsp.data().dptr<DType>();
          auto out_data = dns->dptr<DType>();
          auto shape = rsp.shape();
          const dim_t num_rows = rsp.aux_shape(rowsparse::kIdx).Size();
          const dim_t row_length = shape.ProdShape(1, shape.ndim());
          const dim_t num_threads = num_rows;
          mxnet_op::Kernel<CastStorageRspDnsKernel, xpu>::Launch(s, num_threads,
              row_length, in_idx, in_data, out_data);
        }
      });
    });
  }
  
  /*!
   * \brief CPU kernel for initializing the indptr in a csr matrix.
   */
  struct FillCsrIndPtr {
    /*!
     * \brief
     * \param i         the i-th row of the dns tensor
     * \param indptr    the indptr of the csr tensor
     * \param dns       the dns tensor
     * \param num_rows  number of rows of the dns tensor
     * \param num_cols  number of columns of the dns tensor
     */
    template<typename DType, typename IType>
    MSHADOW_CINLINE static void Map(int i,
                                    IType* indptr,
                                    const DType* dns,
                                    const nnvm::dim_t num_rows,
                                    const nnvm::dim_t num_cols) {
      using nnvm::dim_t;
      indptr[i+1] = 0;
      const dim_t offset = i * num_cols;
      for (dim_t j = 0; j < num_cols; ++j) {
        if (dns[offset+j] != 0) {
          ++indptr[i+1];
        }
      }
    }
  };
  
  /*!
   * \brief CPU kernel for initializing the col_idx and value array of the csr matrix.
   */
  struct FillCsrColIdxAndVals {
    /*!
     * \brief
     * \param i         the i-th row of the dns tensor
     * \param val       value array of the csr tensor
     * \param col_idx   column idx array of the csr tensor
     * \param indptr    indptr array of the csr tensor
     * \param dns       dns tensor
     * \param num_rows  number of rows of the dns tensor
     * \param num_cols  number of columns of the dns tensor
     */
    template<typename DType, typename IType, typename CType>
    MSHADOW_CINLINE static void Map(int i,
                                    DType* val,
                                    CType* col_idx,
                                    const IType* indptr,
                                    const DType* dns,
                                    const nnvm::dim_t num_rows,
                                    const nnvm::dim_t num_cols) {
      using nnvm::dim_t;
      const dim_t offset = i * num_cols;
      IType k = indptr[i];
      for (dim_t j = 0; j < num_cols; ++j) {
        if (dns[offset+j] != 0) {
          val[k] = dns[offset+j];
          col_idx[k] = j;
          ++k;
        }
      }
    }
  };
  
  /*!
   * \brief CPU implementation of casting a dns matrix to csr type.
   */
  inline void CastStorageDnsCsrImpl(const OpContext& ctx,
                                    const cpu& cpu_dev,
                                    const TBlob& dns,
                                    NDArray* csr) {
    CHECK(csr != nullptr);
    CHECK_EQ(csr->storage_type(), kCSRStorage);
    CHECK_EQ(dns.shape_.ndim(), 2);
    CHECK_EQ(dns.shape_, csr->shape());
    using mshadow::Shape1;
    using nnvm::dim_t;
    mshadow::Stream<cpu>* s = ctx.get_stream<cpu>();
    MSHADOW_TYPE_SWITCH(dns.type_flag_, DType, {  // data type
      MSHADOW_IDX_TYPE_SWITCH(csr->aux_type(csr::kIndPtr), IType, {  // indptr type
        MSHADOW_IDX_TYPE_SWITCH(csr->aux_type(csr::kIdx), CType, {  // col idx type
          const dim_t num_rows = dns.shape_[0];
          const dim_t num_cols = dns.shape_[1];
          csr->CheckAndAllocAuxData(csr::kIndPtr, mshadow::Shape1(num_rows+1));
          IType* indptr = csr->aux_data(csr::kIndPtr).dptr<IType>();
          DType* dns_data = dns.dptr<DType>();
          dim_t num_threads = num_rows;
          mxnet_op::Kernel<FillCsrIndPtr, cpu>::Launch(s, num_threads,
              indptr, dns_data, num_rows, num_cols);
          // single thread to accumulate indptr
          // indptr[num_rows] indicates the number of non-zero elements
          indptr[0] = 0;
          for (dim_t i = 0; i < num_rows; ++i) {
            indptr[i+1] += indptr[i];
          }
          // allocate column idx array and value array
          csr->CheckAndAllocAuxData(csr::kIdx, Shape1(static_cast<index_t>(indptr[num_rows])));
          csr->CheckAndAllocData(Shape1(static_cast<index_t>(indptr[num_rows])));
          // fill col_idx and value arrays of the csr
          mxnet_op::Kernel<FillCsrColIdxAndVals, cpu>::Launch(s, num_threads,
              csr->data().dptr<DType>(), csr->aux_data(csr::kIdx).dptr<CType>(),
              indptr, dns_data, num_rows, num_cols);
        });
      });
    });
  }
  
  /*!
   * \brief This is the kernel for copying csr.data to its corresponding dns matrix.
   */
  struct CopyCsrDataToDns {
    /*!
     * \brief
     * \param i         the i-th row of the dns tensor
     * \param dns_data  data blob of the dns tensor
     * \param col_idx   column idx array of the csr tensor
     * \param indptr    indptr array of the csr tensor
     * \param csr_data  data blob of the csr tensor
     * \param num_cols  number of columns of the dns tensor
     */
    template<typename DType, typename IType, typename CType>
    MSHADOW_XINLINE static void Map(int i,
                                    DType* dns_data,
                                    const CType* col_idx,
                                    const IType* indptr,
                                    const DType* csr_data,
                                    const nnvm::dim_t num_cols) {
      const nnvm::dim_t offset = i * num_cols;
      for (IType j = indptr[i]; j < indptr[i+1]; ++j) {
        dns_data[offset+col_idx[j]] = csr_data[j];
      }
    }
  };
  
  /*!
   * \brief Casts a csr matrix to dns format.
   */
  template<typename xpu>
  void CastStorageCsrDnsImpl(const OpContext& ctx,
                             const NDArray& csr,
                             TBlob* dns) {
    CHECK(dns != nullptr);
    CHECK_EQ(csr.storage_type(), kCSRStorage);
    CHECK_EQ(dns->shape_.ndim(), 2);
    CHECK_EQ(dns->shape_, csr.shape());
    using nnvm::dim_t;
    mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
    MSHADOW_TYPE_SWITCH(dns->type_flag_, DType, {  // data type
      MSHADOW_IDX_TYPE_SWITCH(csr.aux_type(csr::kIndPtr), IType, {  // indptr type
        MSHADOW_IDX_TYPE_SWITCH(csr.aux_type(csr::kIdx), CType, {  // col idx type
          const dim_t num_rows = dns->shape_[0];
          const dim_t num_cols = dns->shape_[1];
          DType* dns_data = dns->dptr<DType>();
          dim_t num_threads = dns->shape_.Size();
          mxnet_op::Kernel<mxnet_op::set_zero, xpu>::Launch(s, num_threads, dns_data);
          if (!csr.storage_initialized()) return;
          const IType* indptr = csr.aux_data(csr::kIndPtr).dptr<IType>();
          const CType* col_idx = csr.aux_data(csr::kIdx).dptr<CType>();
          const DType* csr_data = csr.data().dptr<DType>();
          num_threads = num_rows;
          mxnet_op::Kernel<CopyCsrDataToDns, xpu>::Launch(s, num_threads,
              dns_data, col_idx, indptr, csr_data, num_cols);
        });
      });
    });
  }
  
  /*!
   * \brief Casts a csr matrix to another csr.
   */
  template <typename xpu>
  void CastStorageCsrCsrImpl(const OpContext& ctx, const NDArray& csr,
                             NDArray* output) {
    mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
    if (!csr.storage_initialized()) {
      FillZerosCsrImpl(s, *output);
      return;
    }
    std::vector<TShape> aux_shapes({csr.aux_shape(csr::kIndPtr), csr.aux_shape(csr::kIdx)});
    output->CheckAndAlloc(aux_shapes);
    const TBlob& val = output->data();
    const TBlob& indptr = output->aux_data(csr::kIndPtr);
    const TBlob& idx = output->aux_data(csr::kIdx);
    mxnet_op::copy(s, val, csr.data());
    mxnet_op::copy(s, indptr, csr.aux_data(csr::kIndPtr));
    mxnet_op::copy(s, idx, csr.aux_data(csr::kIdx));
  }
  
  /*!
   * \brief Casts a rsp matrix to another rsp.
   */
  template <typename xpu>
  void CastStorageRspRspImpl(const OpContext& ctx, const NDArray& rsp,
                             NDArray* output) {
    CHECK_EQ(rsp.storage_type(), output->storage_type())
        << "Copying with different storage type";
    mshadow::Stream<xpu>* s = ctx.get_stream<xpu>();
    if (!rsp.storage_initialized()) {
      FillZerosRspImpl(s, *output);
      return;
    }
    auto aux_shape = rsp.aux_shape(rowsparse::kIdx);
    output->CheckAndAlloc({aux_shape});
    const TBlob& val = output->data();
    const TBlob& idx = output->aux_data(rowsparse::kIdx);
    const TBlob& from_val = rsp.data();
    const TBlob& from_idx = rsp.aux_data(rowsparse::kIdx);
    mxnet_op::copy(s, val, from_val);
    mxnet_op::copy(s, idx, from_idx);
  }
  
  template<typename xpu>
  void CastStorageComputeImpl(const OpContext& ctx,
                              const NDArray& input,
                              const NDArray& output) {
    const auto src_stype = input.storage_type();
    const auto dst_stype = output.storage_type();
    if (src_stype == kRowSparseStorage && dst_stype == kDefaultStorage) {
      TBlob ret = output.data();
      CastStorageRspDnsImpl<xpu>(ctx, input, &ret);
    } else if (src_stype == kDefaultStorage && dst_stype == kRowSparseStorage) {
      NDArray ret = output;  // get rid of the const qualifer
      CastStorageDnsRspImpl(ctx, xpu(), input.data(), &ret);
    } else if (src_stype == kDefaultStorage && dst_stype == kCSRStorage) {
      NDArray ret = output;  // get rid of the const qualifer
      CastStorageDnsCsrImpl(ctx, xpu(), input.data(), &ret);
    } else if (src_stype == kCSRStorage && dst_stype == kDefaultStorage) {
      TBlob ret = output.data();
      CastStorageCsrDnsImpl<xpu>(ctx, input, &ret);
    } else if (src_stype == kCSRStorage && dst_stype == kCSRStorage) {
      NDArray ret = output;
      CastStorageCsrCsrImpl<xpu>(ctx, input, &ret);
    } else if (src_stype == kRowSparseStorage && dst_stype == kRowSparseStorage) {
      NDArray ret = output;
      CastStorageRspRspImpl<xpu>(ctx, input, &ret);
  #if MXNET_USE_MKLDNN == 1
    } else if (src_stype == kDefaultStorage && dst_stype == kDefaultStorage) {
      CHECK_EQ(output.ctx().dev_type, input.ctx().dev_type);
      // If one of them uses the MKLDNN layout.
      if (input.IsMKLDNNData() || output.IsMKLDNNData()) {
        NDArray tmp_input = input;
        // If the input data is MKLDNN and is a view, we need to reorder the input
        // data first.
        if (input.IsMKLDNNData() && input.IsView())
          tmp_input = input.Reorder2Default();
        const mkldnn::memory *in_mem = tmp_input.GetMKLDNNData();
        const_cast<NDArray &>(output).CopyFrom(*in_mem);
        MKLDNNStream::Get()->Submit();
      } else {
        mxnet_op::copy(ctx.get_stream<xpu>(), output.data(), input.data());
      }
  #endif
    } else {
      LOG(FATAL) << "Not implemented from " << src_stype << " to " << dst_stype;
    }
  }
  
  struct CastStorageParam : public dmlc::Parameter<CastStorageParam> {
    int stype;
    DMLC_DECLARE_PARAMETER(CastStorageParam) {
      DMLC_DECLARE_FIELD(stype)
      .add_enum("default", kDefaultStorage)
      .add_enum("row_sparse", kRowSparseStorage)
      .add_enum("csr", kCSRStorage)
      .describe("Output storage type.");
    }
  };
  
  inline bool CastStorageInferStorageType(const nnvm::NodeAttrs& attrs,
                                          const int dev_mask,
                                          DispatchMode* dispatch_mode,
                                          std::vector<int> *in_attrs,
                                          std::vector<int> *out_attrs) {
    CHECK_EQ(in_attrs->size(), 1U);
    CHECK_EQ(out_attrs->size(), 1U);
    CHECK_NE(in_attrs->at(0), kUndefinedStorage)
      << "src ndarray's storage type must be specified";
    const CastStorageParam& param = nnvm::get<CastStorageParam>(attrs.parsed);
    CHECK_NE(param.stype, kUndefinedStorage)
      << "dst ndarray's storage type must be specified";
    const auto& in_stype = in_attrs->at(0);
    const auto& param_stype = static_cast<NDArrayStorageType>(param.stype);
    bool dispatched = false;
    // dns -> dns, dns -> rsp, dns -> csr
    if (!dispatched && in_stype == kDefaultStorage && param_stype == kDefaultStorage) {
      // dns -> dns
      DispatchMode mode = DispatchMode::kFCompute;
  #if MXNET_USE_MKLDNN == 1
      // If we use MKLDNN and the arrays are in CPU memory, the array may store
      // MKLDNN layout, we should convert its layout explicitly.
      if (dev_mask == kCPU)
        mode = DispatchMode::kFComputeEx;
  #endif
      dispatched = storage_type_assign(out_attrs, kDefaultStorage, dispatch_mode, mode);
    }
    if (!dispatched && in_stype == kDefaultStorage &&
      (param_stype == kRowSparseStorage || param_stype == kCSRStorage)) {
      // dns -> rsp, dns -> csr
      dispatched = storage_type_assign(out_attrs, param_stype,
                                       dispatch_mode, DispatchMode::kFComputeEx);
    }
    if (!dispatched && in_stype == kRowSparseStorage &&
        (param_stype == kRowSparseStorage || param_stype == kDefaultStorage)) {
      // rsp -> rsp, rsp -> dns
      dispatched = storage_type_assign(out_attrs, param_stype,
                                       dispatch_mode, DispatchMode::kFComputeEx);
    }
    if (!dispatched && in_stype == kCSRStorage &&
        (param_stype == kCSRStorage || param_stype == kDefaultStorage)) {
      // csr -> csr, csr -> dns
      dispatched = storage_type_assign(out_attrs, param_stype,
                                       dispatch_mode, DispatchMode::kFComputeEx);
    }
    return dispatched;
  }
  
  template<typename xpu>
  void CastStorageComputeEx(const nnvm::NodeAttrs& attrs,
                            const OpContext& ctx,
                            const std::vector<NDArray>& inputs,
                            const std::vector<OpReqType>& req,
                            const std::vector<NDArray>& outputs) {
    CHECK_EQ(inputs.size(), 1);
    CHECK_EQ(outputs.size(), 1);
    if (req[0] == kNullOp) return;
    CHECK_EQ(req[0], kWriteTo) << "CastStorageComputeEx expects req[0] == kWriteTo";
    CastStorageComputeImpl<xpu>(ctx, inputs[0], outputs[0]);
  }
  
  }  // namespace op
  }  // namespace mxnet
  
  #endif  // MXNET_OPERATOR_TENSOR_CAST_STORAGE_INL_H_
  