larroy commented on a change in pull request #8915: NVLink communication 
pattern updated 
URL: https://github.com/apache/incubator-mxnet/pull/8915#discussion_r167594127
 
 

 ##########
 File path: src/kvstore/comm.h
 ##########
 @@ -644,38 +803,44 @@ class CommDevice : public Comm {
         CopyFromTo(src, out, priority);
       } else {
         CHECK_EQ(out->storage_type(), kRowSparseStorage)
-                 << "BroadcastRowSparse expects row_sparse dst NDArray";
+            << "BroadcastRowSparse expects row_sparse dst NDArray";
 
         const bool is_diff_ctx = out->ctx() != src.ctx();
-        NDArray out_gpu = is_diff_ctx? NDArray(kRowSparseStorage, out->shape(),
-            src.ctx(), true, out->dtype(), out->aux_types()) : *out;
+        NDArray out_gpu =
+            is_diff_ctx ? NDArray(kRowSparseStorage, out->shape(), src.ctx(),
+                                  true, out->dtype(), out->aux_types())
+                        : *out;
 
         CHECK_EQ(row_id.ctx(), src.ctx())
-                << "row_id and src are expected to be on the same context";
-
-        Engine::Get()->PushAsync([=](RunContext rctx, 
Engine::CallbackOnComplete on_complete) {
-            NDArray temp = out_gpu;
-            const TBlob& indices = row_id.data();
-            switch (temp.ctx().dev_mask()) {
-              case cpu::kDevMask: {
-                
mxnet::common::SparseRetainOpForwardRspWrapper<cpu>(rctx.get_stream<cpu>(),
-                    src, indices, kWriteTo, &temp);
-                break;
-              }
+            << "row_id and src are expected to be on the same context";
+
+        Engine::Get()->PushAsync(
+            [=](RunContext rctx, Engine::CallbackOnComplete on_complete) {
+              NDArray temp = out_gpu;
+              const TBlob& indices = row_id.data();
+              switch (temp.ctx().dev_mask()) {
+                case cpu::kDevMask: {
+                  mxnet::common::SparseRetainOpForwardRspWrapper<cpu>(
+                      rctx.get_stream<cpu>(), src, indices, kWriteTo, &temp);
+                  break;
+                }
 #if MXNET_USE_CUDA
-              case gpu::kDevMask: {
-                
mxnet::common::SparseRetainOpForwardRspWrapper<gpu>(rctx.get_stream<gpu>(),
-                    src, indices, kWriteTo, &temp);
-                // wait for GPU operations to complete
-                rctx.get_stream<gpu>()->Wait();
-                break;
-              }
+                case gpu::kDevMask: {
 
 Review comment:
   Wrong indentation?

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


With regards,
Apache Git Services

Reply via email to