Caffe-plugin Sync cpu data/diff to gpu (apache#3017)

* [caffe-plugin] Sync cpu data/diff to gpu.

* [caffe-plugin] fix bug

* [caffe-plugin] Do sync only for gpu code.

plugin/caffe/caffe_loss-inl.h

@@ -102,6 +102,10 @@ class CaffeLoss : public Operator {
     caffeOp_->Forward(bot_, top_);
 
 #if defined(__CUDACC__)
+    // Sync cpu data to gpu data
+    for (uint32_t i = 0; i < top_.size(); ++i)
+      top_[i]->gpu_data();
+
     CHECK_EQ(cudaStreamSynchronize(NULL), cudaSuccess);
 #endif  // __CUDACC__
   }
@@ -151,6 +155,10 @@ class CaffeLoss : public Operator {
     caffeOp_->Backward(top_, flags_, bot_);
 
 #if defined(__CUDACC__)
+    // Sync cpu diff to gpu diff
+    for (uint32_t i = 0; i < bot_.size(); ++i)
+      bot_[i]->gpu_diff();
+
     CHECK_EQ(cudaStreamSynchronize(NULL), cudaSuccess);
 #endif  // __CUDACC__
   }

plugin/caffe/caffe_op-inl.h

@@ -79,9 +79,9 @@ class CaffeOp : public Operator {
     using std::vector;
     using namespace mshadow;
     using namespace mshadow::expr;
-    for (index_t i = 0; i < req.size(); ++i)
+    for (uint32_t i = 0; i < req.size(); ++i)
       CHECK_EQ(req[i], kWriteTo);
-    index_t expected_num_data = param_.num_weight + param_.num_data;
+    int expected_num_data = param_.num_weight + param_.num_data;
     CHECK_EQ(in_data.size(), expected_num_data);
     CHECK_EQ(out_data.size(), param_.num_out);
 
@@ -116,6 +116,10 @@ class CaffeOp : public Operator {
     caffeOp_->Forward(bot_, top_);
 
 #if defined(__CUDACC__)
+    // Sync cpu data to gpu data
+    for (uint32_t i = 0; i < top_.size(); ++i)
+      top_[i]->gpu_data();
+
     CHECK_EQ(cudaStreamSynchronize(NULL), cudaSuccess);
 #endif  // __CUDACC__
   }
@@ -140,10 +144,10 @@ class CaffeOp : public Operator {
     using namespace mshadow;
     using namespace mshadow::expr;
     CHECK_EQ(out_grad.size(), param_.num_out);
-    for (index_t i = 0; i < param_.num_data; ++i)
+    for (int i = 0; i < param_.num_data; ++i)
       CHECK(req[i] != kAddTo) << "caffe doesn't accm diff on bottom data";
 
-    index_t expected_num_data = param_.num_weight + param_.num_data;
+    int expected_num_data = param_.num_weight + param_.num_data;
     CHECK(in_data.size() == expected_num_data && in_grad.size() == expected_num_data);
     CHECK_EQ(req.size(), expected_num_data);
 
@@ -173,16 +177,20 @@ class CaffeOp : public Operator {
     }
 
     // Handle OpReqType of weights
-    for (index_t i = param_.num_data; i < expected_num_data; ++i)
+    for (int i = param_.num_data; i < expected_num_data; ++i)
       HandleOpReq(s, req[i], in_grad[i]);
 
     // Set BP flag
-    for (index_t i = 0; i < param_.num_data; ++i)
+    for (int i = 0; i < param_.num_data; ++i)
       flags_[i] = req[i] != kNullOp;
 
     caffeOp_->Backward(top_, flags_, bot_);
 
 #if defined(__CUDACC__)
+    // Sync cpu diff to gpu diff
+    for (uint32_t i = 0; i < bot_.size(); ++i)
+      bot_[i]->gpu_diff();
+
     CHECK_EQ(cudaStreamSynchronize(NULL), cudaSuccess);
 #endif  // __CUDACC__
   }