[GitHub] [incubator-mxnet] reminisce commented on a change in pull request #17208: gather_nd: check bound and wrap negative indices

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

reminisce commented on a change in pull request #17208: gather_nd: check bound and wrap negative indices
URL: https://github.com/apache/incubator-mxnet/pull/17208#discussion_r363435410
 
 

 ##########
 File path: src/operator/tensor/indexing_op.cu
 ##########
 @@ -437,6 +436,98 @@ inline void SparseEmbeddingOpBackwardRspImpl<gpu>(const bool deterministic,
   });
 }
 
+struct gather_nd {
+  template<typename DType, typename IType>
+  MSHADOW_XINLINE static void Map(index_t i, OpReqType req, index_t N, index_t M, index_t K,
+                                  const mshadow::Shape<10> strides,
+                                  const mshadow::Shape<10> mshape,
+                                  DType* out, const DType* data,
+                                  const IType* indices) {
+    index_t offset = 0;
+    for (index_t j = 0; j < M; ++j) {
+      offset += strides[j] * ((static_cast<index_t>(indices[j*N + i]) + mshape[j])%mshape[j]);
+    }
+    for (index_t j = 0; j < K; ++j) {
+      KERNEL_ASSIGN(out[i*K + j], req, data[offset+j]);
+    }
+  }
+};
+
+/*! \brief If there are out-of-bound indices, out will be assigned to 1.
+ */
+struct is_valid_check_gather_nd {
+  template<typename DType>
+  MSHADOW_XINLINE static void Map(int i, char* out, const DType* data,
+                                  const index_t N, const mshadow::Shape<10> mshape) {
+    index_t n = N - 1;
+    while (n >= 0) {
+      if (data[i*N + n] < -mshape[i] || data[i*N + n] > mshape[i] - 1) *out = 1;
+      n--;
+    }
+  }
+};
+
+/*
+ * \brief returns true if all indices is out of bounds for axis x with size y
+ * \param s the stream
+ * \param data_ptr the indices on the stream
+ * \param data_size the number of indices to examine
+ * \param M the number of axises to exmaine
+ * \param mshape the array that stores shape for each dimension
+ * \param is_valid_ptr the temparary workspace
+ */
+template<typename DType>
+bool GatherNDCheckBound(mshadow::Stream<gpu> *s, const DType* data_ptr, index_t N,
+                        index_t M, const mshadow::Shape<10> mshape, char* is_valid_ptr) {
+  using namespace mxnet_op;
+  int32_t is_valid = 0;
+  Kernel<set_zero, gpu>::Launch(s, 1, is_valid_ptr);
+  Kernel<is_valid_check_gather_nd, gpu>::Launch(s, M, is_valid_ptr, data_ptr, N, mshape);
+  CUDA_CALL(cudaMemcpyAsync(&is_valid, is_valid_ptr, sizeof(char),
+                            cudaMemcpyDeviceToHost, mshadow::Stream<gpu>::GetStream(s)));
+  CUDA_CALL(cudaStreamSynchronize(mshadow::Stream<gpu>::GetStream(s)));
+  return is_valid == 0;
+}
+
+template<typename gpu>
+void GatherNDForward(const nnvm::NodeAttrs& attrs,
 
 Review comment:
   Same here. `GatherNDForwardGPU`.

----------------------------------------------------------------
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