[GitHub] [incubator-mxnet] wkcn commented on a change in pull request #14359: Speedup _contrib_index_copy

[GitBox <gi...@apache.org>]

wkcn commented on a change in pull request #14359: Speedup _contrib_index_copy
URL: https://github.com/apache/incubator-mxnet/pull/14359#discussion_r264194219
 
 

 ##########
 File path: src/operator/contrib/index_copy.cu
 ##########
 @@ -18,14 +18,101 @@
  */
 
 /*!
- * \file index_copy.cc
+ * \file index_copy.cu
  * \brief
  */
 #include "./index_copy-inl.h"
 
 namespace mxnet {
 namespace op {
 
+struct index_copy_fwd_gpu {
+  template<typename DType, typename IType>
+  MSHADOW_XINLINE static void Map(int i,
+                                  const DType* new_tensor,
+                                  const IType* idx,
+                                  DType* out_tensor,
+                                  int dim_size) {
+    int index = static_cast<int>(idx[i / dim_size]);
+    out_tensor[index * dim_size + i % dim_size] = new_tensor[i];
+  }
+};
+
+template<>
+void IndexCopyForward<gpu>(const nnvm::NodeAttrs& attrs,
+                           const OpContext& ctx,
+                           const std::vector<TBlob>& inputs,
+                           const std::vector<OpReqType>& req,
+                           const std::vector<TBlob>& outputs) {
+  using namespace mshadow;
+  using namespace mxnet_op;
+  CHECK_EQ(inputs.size(), 3U);
+  CHECK_EQ(outputs.size(), 1U);
+  CHECK_EQ(req.size(), 1U);
+  mshadow::Stream<gpu> *s = ctx.get_stream<gpu>();
+  const TBlob& out = outputs[0];
+  const TBlob& original_tensor = inputs[0];
+  const TBlob& idx_vector = inputs[1];
+  const TBlob& copied_tensor = inputs[2];
+  int dim_size = inputs[2].Size() / inputs[1].Size();
+  // copy original tensor to output
+  copy(s, out, original_tensor);
+  // index copy
+  MSHADOW_TYPE_SWITCH(out.type_flag_, DType, {
+    MSHADOW_TYPE_SWITCH(idx_vector.type_flag_, IType, {
+      Kernel<index_copy_fwd_gpu, gpu>::Launch(
+        s, copied_tensor.Size(), copied_tensor.dptr<DType>(),
+        idx_vector.dptr<IType>(), out.dptr<DType>(), dim_size);
+    });
+  });
+}
+
+struct index_copy_bwd_gpu {
+  template<typename DType, typename IType>
+  MSHADOW_XINLINE static void Map(int i,
+                                  const DType* out_grad,
+                                  DType* orig_grad,
+                                  DType* new_grad,
+                                  const IType* idx,
+                                  int dim_size,
+                                  int idx_size) {
+    orig_grad[i] = out_grad[i];
+    if (i / dim_size < idx_size) {
 
 Review comment:
   the same question as that in CPU.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services