[GitHub] [incubator-mxnet] sxjscience commented on a change in pull request #18168: [WIP]separate GPU kernel for broadcast_axis

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

sxjscience commented on a change in pull request #18168:
URL: https://github.com/apache/incubator-mxnet/pull/18168#discussion_r416729701



##########
File path: src/operator/tensor/broadcast_reduce_op.h
##########
@@ -1077,6 +1077,58 @@ struct broadcast_kernel {
   }
 };
 
+template<typename OP>
+struct broadcast_kernel_gpu {
+  template<typename IType, typename OType>
+  MSHADOW_XINLINE static void Map(int32_t i,
+                                  IType *input,
+                                  OType *output,
+                                  const int32_t *input_shape,
+                                  const int32_t *output_shape,
+                                  const int32_t *in_stride,
+                                  const int32_t *out_stride,
+                                  const OpReqType req,
+                                  const int32_t ndim) {
+    int32_t idx = i;
+    int32_t in_idx = i;
+    for (int32_t iter = ndim - 1; iter >= 0; --iter) {
+      int32_t dim_idx = idx % output_shape[iter];
+      if (input_shape[iter] != 1) {
+        in_idx += dim_idx * (in_stride[iter] - out_stride[iter]);
+      } else {
+        in_idx -= dim_idx * out_stride[iter];
+      }
+      idx /= output_shape[iter];
+    }
+    KERNEL_ASSIGN(output[i], req, OP::Map(input[in_idx]));
+  }
+};
+
+template<int req>
+struct compute_offset {
+  MSHADOW_XINLINE static void Map(int32_t i,
+                                  int32_t *in_stride,
+                                  int32_t *out_stride,

Review comment:
       If we are accessing a large tensor, will it happen that the in_stride, out_stride can exceed the int32 limit?




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