Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into...

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into add_FLAGS_use_deterministic_algo

benchmark/fluid/machine_translation.py

@@ -21,7 +21,7 @@ import argparse
 import time
 import distutils.util
 
-import paddle.v2 as paddle
+import paddle
 import paddle.fluid as fluid
 import paddle.fluid.core as core
 import paddle.fluid.framework as framework

benchmark/fluid/mnist.py

@@ -20,7 +20,7 @@ import numpy as np
 import argparse
 import time
 
-import paddle.v2 as paddle
+import paddle
 import paddle.fluid as fluid
 import paddle.fluid.profiler as profiler
 

benchmark/fluid/resnet.py

@@ -23,7 +23,7 @@ import time
 
 import cProfile, pstats, StringIO
 
-import paddle.v2 as paddle
+import paddle
 import paddle.fluid as fluid
 import paddle.fluid.core as core
 import paddle.fluid.profiler as profiler

benchmark/fluid/stacked_dynamic_lstm.py

@@ -23,10 +23,10 @@ import random
 import time
 
 import numpy
-import paddle.v2 as paddle
-import paddle.v2.dataset.imdb as imdb
+import paddle
+import paddle.dataset.imdb as imdb
 import paddle.fluid as fluid
-from paddle.v2 import batch
+import paddle.batch as batch
 import paddle.fluid.profiler as profiler
 
 

benchmark/fluid/vgg.py

@@ -17,7 +17,7 @@ from __future__ import print_function
 import sys
 import time
 import numpy as np
-import paddle.v2 as paddle
+import paddle
 import paddle.fluid as fluid
 import paddle.fluid.core as core
 import argparse

doc/fluid/api/data/data_reader.rst

@@ -56,11 +56,11 @@ DataFeeder
 Reader
 ======
 
-..  automodule:: paddle.v2.reader
+..  automodule:: paddle.reader
     :members:
     :noindex:
 
-..  automodule:: paddle.v2.reader.creator
+..  automodule:: paddle.reader.creator
     :members:
     :noindex:
 

doc/fluid/api/layers.rst

@@ -479,6 +479,13 @@ label_smooth
 ..  autofunction:: paddle.fluid.layers.label_smooth
     :noindex:
 
+roi_pool
+---------
+
+..  autofunction:: paddle.fluid.layers.roi_pool
+    :noindex:
+
+    
 ops
 ===
 
@@ -820,3 +827,5 @@ topk
 
 ..  autofunction:: paddle.fluid.layers.topk
     :noindex:
+
+

doc/fluid/design/data_type/float16.md

@@ -3,7 +3,7 @@
 ## Why float16
 Half precision (float16) is a binary floating-point format that occupies 16 bits in memory. float16 is half the size of traditional 32-bit single precision format (float) and has lower precision and smaller range. 
 
-When high precision computation is not required, using float16 data type could potentially 
+When high precision computation is not required (which is usually the case at least in the deep learning inference stage), using float16 data type could potentially 
 
 - reduce storage space, memory bandwidth, and power usages; 
 - increase the chance of data fitting into a smaller cache of lower latency; 
@@ -12,7 +12,7 @@ When high precision computation is not required, using float16 data type could p
 ## Survey of current float16 support
 A brief survey of float16 support on different compilers, hardwares, and libraries can be found below. Interested readers can refer to [link1](https://github.com/PaddlePaddle/Paddle/issues/4853) and [link2](https://github.com/Xreki/Xreki.github.io/blob/master/multi_data_types_in_dl_framework/ppt/float16_and_quantized_type.md) for more info.
 
-The goal of float16 is to serve as a key for the executor to find and run the correct version of compute method specialized for float16 in operator kernel. It should be compatible with various natively supported float16 implementations including `__half` for cuda, `float16_t` for ARM, and `Eigen::half` for Eigen to make writing customized float16 kernels easier. 
+The goal of float16 is to serve as a key for the executor to find and run the correct version of compute method specialized for float16 in operator kernels. It should be compatible with various natively supported float16 implementations including `__half` for cuda, `float16_t` for ARM, and `Eigen::half` for Eigen to make writing customized float16 kernels easier. 
 
 ### Compiler
 - nvcc supports `__half` data type after CUDA 7.5.
@@ -95,11 +95,89 @@ float half_to_float(float16 h);
 ```
 which provides one-to-one conversion between float32 and float16. These twos functions will do different conversion routines based on the current hardware. CUDA/ARM instrinsics will be used when the corresonding hardware is available. If the hardware or compiler level does not support float32 to float16 conversion, software emulation will be performed to do the conversion.
 
-## To do
-After float16 class is available, some of the future items are below:
+## float16 inference
+In Fluid, a neural network is represented as a protobuf message called [ProgramDesc](https://github.com/PaddlePaddle/Paddle/blob/develop/doc/fluid/design/concepts/program.md), whose Python wrapper is a [Program](https://github.com/PaddlePaddle/Paddle/blob/develop/doc/fluid/design/modules/python_api.md#program). The basic structure of a program is some nested [blocks](https://github.com/PaddlePaddle/Paddle/blob/develop/doc/fluid/design/modules/python_api.md#block), where each block consists of some [variable](https://github.com/PaddlePaddle/Paddle/blob/develop/doc/fluid/design/modules/python_api.md#variable) definitions and a sequence of [operators](https://github.com/PaddlePaddle/Paddle/blob/develop/doc/fluid/design/modules/python_api.md#operator). An [executor](https://github.com/PaddlePaddle/Paddle/blob/develop/doc/fluid/design/concepts/executor.md) will run a given program desc by executing the sequence of operators in the entrance block of the program one by one.  
 
-- Update pybind/tensor_py.h to bind c++ float16 with numpy float16. 
+### Operator level requirement
+Each operator has many kernels for different data types, devices, and library types. The operator will select the appropriate kernel to run based on, among other things, the data type of the input variables. By default, every Fluid operator has a float data type kernel that takes float variables as input and generates float output. 
 
-- Modify `GetKernelType()` method in `framework/operator.h` to make it compatible with float16.
+This means that if we provide float input to the first operator in a program, then each opeartor will use float kernel to compute float output and send it as input to the next operator to trigger the float kernel. Overall, the program will run in float mode and give us a final output of float data type.
 
-- Create a type-casting operator that can convert the data type in tensor between float16 and other types.
+The same principle applies if we want a program to run in float16 mode. We provide input variable of float16 data type to the first operator, and then one by one, each operator in the program will run the float16 kernel (provided that each operator in this program has float16 kernels registered) until we finally obtain a float16 output variable.
+
+So the preliminary requirement for float16 inference is to add float16 kernel to operators that are needed in a specific kind of program. For example, float16 inference on an image classification neural network like Vgg or Resnet, typically requires the following operators to have float16 kernels: convolution, pooling, multiplication, addition, batch norm, dropout, relu, and softmax. Please refer to [new_op_en](https://github.com/PaddlePaddle/Paddle/blob/develop/doc/fluid/dev/new_op_en.md) for details of how to add new kernels to an operator.
+
+### Variable level requirement
+Operators including convolution and multiplication (used in fully-connected layers) takes as input not only the variables generated by the preceding operators but also [parameter](https://github.com/PaddlePaddle/Paddle/blob/develop/doc/fluid/design/modules/python_api.md#parameter) variables, which contains the trained weights to apply to the input data. These weights are obtained in the Fluid training process and are by default of float data type.
+
+When these operators are running in float16 mode, the float16 kernel requires those parameter variables to contain weights of Fluid float16 data type. Thus, we need a convenient way to convert the original float weights to float16 weights. 
+
+In Fluid, we use tensor to hold actual data for a variable on the c++ end. [Pybind](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/fluid/pybind/tensor_py.h) is used to bind c++ tensors of certain data type with numpy array of the correponding numpy data type on the Python end. Each common c++ built-in data type has a corresponding numpy data type of the same name. However, since there is no built-in float16 type in c++, we cannot directly bind numpy float16 data type with the Fluid float16 class. Since both Fluid float16 and numpy float16 use uint16 as the internal data storage type, we use c++ built-in type `uint16_t` and the corresponding numpy uint16 data type to bridge the gap via [Pybind](https://github.com/PaddlePaddle/Paddle/blob/develop/paddle/fluid/pybind/tensor_py.h). 
+
+The following code demonstrates how to do the tensor conversion.
+```Python
+# var is the variable of float weights
+# tensor is a numpy array of data copied from the tensor data in var 
+# fp16_var is the variable that will contain float16 weights converted from var  
+tensor = numpy.array(var.get_tensor())
+fp16_tensor = fp16_var.get_tensor()
+
+# After the original tensor data is converted to numpy float16 data type, 
+# view(numpy.uint16) is used so that the internal memory of the numpy array 
+# will be reinterpreted to be of uint16 data type, which is binded to 
+# Fluid float16 class via pybind with the help of uint16_t built-in c++ type
+fp16_tensor.set(tensor.astype(numpy.float16).view(numpy.uint16), GPUPlace)  
+```
+
+### Consistent API requirement
+The basic inference in float16 mode requires users to feed input and obtain output both of float16 data type. However, in this way, the inference APIs are not consistent between float16 mode and float mode, and users may find it confusing and diffcult to use float16 inference since they need to do extra steps to provide float16 input data and convert float16 output data back to float. To have consistent API for different inference modes, we need to transpile the program desc in some way so that we can run float16 inference by feeding and fetching variables of float data type.
+
+This problem can be solved by introducing a type-casting operator which takes an input variable of certain data type, cast it to another specified data type, and put the casted data into the output variable. Insert cast operator where needed can make a program internally run in float16 mode.   
+
+### float16 transpiler
+Put all the above requirements in mind, we designed a float16 inference transpiler that can tranpile a float32 mode inference program desc to a float16 mode one.
+
+Given a float inference program and the corresponding variables of float32 weights in the [scope](https://github.com/PaddlePaddle/Paddle/blob/develop/doc/fluid/design/concepts/scope.md),
+this transpiler mainly does the following modifications:
+
+1. Insert cast operators at the beginning of the program so that the input float data will be converted to float16 data type before feeding to subsequent operators to invoke the float16 kernel. 
+
+2. Insert cast operators at the end of the program so that the output float16 data will be converted back to float data type before users obtain the result.
+
+3. For each parameter variable of float weights, create in the scope a corresponding variable of float16 weights which are converted from the corresponding float weights and add this new float16 variable to the program.
+
+4. Update the operator information in the program so that each relevant operator use the newly created float16 variable instead of its float counterpart.
+
+Below is an example of usage:
+```Python
+# Get the float inference program
+[float_inference_program, feed_target_names,
+ fetch_targets] = fluid.io.load_inference_model(save_dirname, exe)
+
+# Prepare the float input data
+tensor_img = numpy.random.rand(1, 3, 32, 32).astype(numpy.float32)
+
+# Running inference_program in float mode
+float_results = exe.run(float_inference_program,
+                        feed={feed_target_names[0]: tensor_img},
+                        fetch_list=fetch_targets)
+
+# Use float16 transpiler to speedup
+float16_inference_program = float_inference_program.clone()
+t = fluid.InferenceTranspiler()
+t.float16_transpile(float16_inference_program, GPUPlace)
+
+# Running 
+float16_results = exe.run(float16_inference_program,
+                          feed={feed_target_names[0]: tensor_img},
+                          fetch_list=fetch_targets)
+```
+
+As we can see from the example above, users can simply use the `float16_transpile` method provided by the infernece transpiler class on an existing float inference program to run inference in float16 mode.
+
+### Speedup on GPU
+Currently, Fluid inference in float16 mode is only supported on Nvidia GPU device. There is no motivation to support float16 inference on non-ARM CPUs because float16 is not natively supported there and float16 calculation will only be slower than its float counterpart. 
+
+Nvidia started to support its native float16 data type (which has the same internal memory representation as Fluid float16 class) on CUDA 7.5. Moreover, float16 speedups on common computational intensive tasks including GEMM (general matrix-matrix multiplication) and convolution are supported since cublas 7.5 and cuDNN 5.0.
+
+Recently, the introduction of [tensor core](https://devblogs.nvidia.com/programming-tensor-cores-cuda-9/) in volta architecture GPUs and the support of tensor core calculation in CUDA 9.0 and cuDNN 7.0 make float16 truly superior to float in certain deep learning applications. Please refer to this [benchmark report](https://github.com/kexinzhao/Paddle_benchmark/blob/master/float16_benchmark.md) for more details.

doc/v2/api/data/data_reader.rst

@@ -56,11 +56,11 @@ DataFeeder
 Reader
 ======
 
-..  automodule:: paddle.v2.reader
+..  automodule:: paddle.reader
     :members:
     :noindex:
 
-..  automodule:: paddle.v2.reader.creator
+..  automodule:: paddle.reader.creator
     :members:
     :noindex:
 

doc/v2/api/data/dataset.rst

 Dataset
 =======
 
-..  automodule:: paddle.v2.dataset
+..  automodule:: paddle.dataset
     :members:
     :noindex:
 
 mnist
 +++++
 
-..  automodule:: paddle.v2.dataset.mnist
+..  automodule:: paddle.dataset.mnist
     :members:
     :noindex:
 
 cifar
 +++++
 
-..  automodule:: paddle.v2.dataset.cifar
+..  automodule:: paddle.dataset.cifar
     :members:
     :noindex:
 
 conll05
 +++++++
 
-..  automodule:: paddle.v2.dataset.conll05
+..  automodule:: paddle.dataset.conll05
     :members: get_dict,get_embedding,test
     :noindex:
 
 imdb
 ++++
 
-..  automodule:: paddle.v2.dataset.imdb
+..  automodule:: paddle.dataset.imdb
     :members:
     :noindex:
 
 imikolov
 ++++++++
 
-..  automodule:: paddle.v2.dataset.imikolov
+..  automodule:: paddle.dataset.imikolov
     :members:
     :noindex:
 
 movielens
 +++++++++
 
-..  automodule:: paddle.v2.dataset.movielens
+..  automodule:: paddle.dataset.movielens
     :members:
     :noindex:
 
-..  autoclass:: paddle.v2.dataset.movielens.MovieInfo
+..  autoclass:: paddle.dataset.movielens.MovieInfo
     :noindex:
-    
-..  autoclass:: paddle.v2.dataset.movielens.UserInfo
+
+..  autoclass:: paddle.dataset.movielens.UserInfo
     :noindex:
 
 sentiment
 +++++++++
 
-..  automodule:: paddle.v2.dataset.sentiment
+..  automodule:: paddle.dataset.sentiment
     :members:
     :noindex:
 
 uci_housing
 +++++++++++
 
-..  automodule:: paddle.v2.dataset.uci_housing
+..  automodule:: paddle.dataset.uci_housing
     :members:
     :noindex:
 
 wmt14
 +++++
 
-..  automodule:: paddle.v2.dataset.wmt14
+..  automodule:: paddle.dataset.wmt14
     :members:
     :noindex:
 
 wmt16
 +++++
 
-..  automodule:: paddle.v2.dataset.wmt16
+..  automodule:: paddle.dataset.wmt16
     :members:
     :noindex:

doc/v2/howto/cluster/multi_cluster/index_en.rst

 Use different clusters
 ======================
 
-PaddlePaddle supports running jobs on several platforms including:
-- `Kubernetes <http://kubernetes.io>`_ open-source system for automating deployment, scaling, and management of containerized applications from Google.
-- `OpenMPI <https://www.open-mpi.org>`_ Mature high performance parallel computing framework.
-- `Fabric <http://www.fabfile.org>`_ A cluster management tool. Write scripts to submit jobs or manage the cluster.
+The user's cluster environment is not the same. To facilitate everyone's deployment, we provide a variety of cluster deployment methods to facilitate the submission of cluster training tasks, which will be introduced as follows:
 
-We'll introduce cluster job management on these platforms. The examples can be found under `cluster_train_v2 <https://github.com/PaddlePaddle/Paddle/tree/develop/paddle/scripts/cluster_train_v2>`_ .
+`Kubernetes <http://kubernetes.io>`_ is a scheduling framework of Google open source container cluster, supporting a complete cluster solution for large-scale cluster production environment. The following guidelines show PaddlePaddle's support for Kubernetes:
 
-These cluster platforms provide API or environment variables for training processes, when the job is dispatched to different nodes. Like node ID, IP or total number of nodes etc.
+..  toctree::
+  :maxdepth: 1
+
+  k8s_cn.md
+  k8s_distributed_cn.md
+
+`OpenMPI <https://www.open-mpi.org>`_ is a mature high-performance parallel computing framework, which is widely used in the field of HPC. The following guide describes how to use OpenMPI to build PaddlePaddle's cluster training task:
 
 ..  toctree::
   :maxdepth: 1
 
-  fabric_en.md
-  openmpi_en.md
-  k8s_en.md
-  k8s_aws_en.md
+  openmpi_cn.md
+
+`Fabric <http://www.fabfile.org>`_ is a convenient tool for program deployment and management. We provide a way to deploy and manage with Fabric. If you want to know more about it, please read the following guidelines:
+
+..  toctree::
+  :maxdepth: 1
+
+  fabric_cn.md
+
+We also support the deployment of PaddlePaddle on AWS. Learn more about:
+
+..  toctree::
+  :maxdepth: 1
+
+  k8s_aws_cn.md
+
+The examples can be found under `cluster_train_v2 <https://github.com/PaddlePaddle/Paddle/tree/develop/paddle/scripts/cluster_train_v2>`_ .
\ No newline at end of file

paddle/fluid/framework/details/broadcast_op_handle_test.cc

@@ -139,7 +139,7 @@ struct TestBroadcastOpHandle {
       PADDLE_ENFORCE_EQ(out_tensor.lod(), lod, "lod is not equal.");
 
       f::Tensor result_tensor;
-      f::TensorCopy(out_tensor, cpu_place, *(ctxs_[j]), &result_tensor);
+      f::TensorCopySync(out_tensor, cpu_place, &result_tensor);
       float* ct = result_tensor.mutable_data<float>(cpu_place);
 
       for (int64_t i = 0; i < f::product(kDims); ++i) {
@@ -185,7 +185,7 @@ struct TestBroadcastOpHandle {
       }
 
       f::Tensor result_tensor;
-      f::TensorCopy(rt, cpu_place, *(ctxs_[j]), &result_tensor);
+      f::TensorCopySync(rt, cpu_place, &result_tensor);
       float* ct = result_tensor.data<float>();
 
       for (int64_t i = 0; i < f::product(kDims); ++i) {

paddle/fluid/framework/details/fetch_op_handle.cc

@@ -66,8 +66,7 @@ void FetchOpHandle::RunImpl() {
     auto &t = var->Get<framework::LoDTensor>();
     if (platform::is_gpu_place(t.place())) {
 #ifdef PADDLE_WITH_CUDA
-      TensorCopy(t, cpu, *dev_ctxes_[t.place()], &tensors_[i], true);
-      dev_ctxes_.at(t.place())->Wait();
+      TensorCopySync(t, cpu, &tensors_[i]);
 #endif
     } else {
       tensors_[i].ShareDataWith(t);

paddle/fluid/framework/details/multi_devices_graph_builder.cc

@@ -58,23 +58,20 @@ MultiDevSSAGraphBuilder::MultiDevSSAGraphBuilder(
 
 void MultiDevSSAGraphBuilder::CreateOpHandleIOs(SSAGraph *result,
                                                 const OpDesc &op,
-                                                const platform::Place &p,
-                                                const size_t &i) const {
+                                                size_t place_id) const {
+  auto p = places_[place_id];
   auto *op_handle = result->ops_.back().get();
   op_handle->SetDeviceContext(p,
                               platform::DeviceContextPool::Instance().Get(p));
 
-  auto var_names = op.InputArgumentNames();
-
-  for (auto &each_var_name : var_names) {
-    VarHandle *var = CreateOrGetLatestVarHandle(result, each_var_name, p, i);
+  for (auto &each_var_name : op.InputArgumentNames()) {
+    VarHandle *var =
+        CreateOrGetLatestVarHandle(result, each_var_name, p, place_id);
     op_handle->AddInput(var);
   }
 
-  var_names = op.OutputArgumentNames();
-
-  for (auto &each_var_name : var_names) {
-    CreateOpOutput(result, op_handle, each_var_name, p, i);
+  for (auto &each_var_name : op.OutputArgumentNames()) {
+    CreateOpOutput(result, op_handle, each_var_name, p, place_id);
   }
 }
 
@@ -84,17 +81,18 @@ bool MultiDevSSAGraphBuilder::IsDistTrainOp(const OpDesc &op,
     return false;
   }
 
-  auto checker = [&](const std::vector<std::string> opvars,
-                     const std::vector<std::string> sendvars) -> bool {
-    bool is_dist_train_op = false;
+  /**
+   * Check any of opvars contains `.block` and in sendvars
+   */
+  auto checker = [](const std::vector<std::string> &opvars,
+                    const std::vector<std::string> &sendvars) -> bool {
     for (auto &var : opvars) {
       if (var.find(".block") != std::string::npos &&
           std::find(sendvars.begin(), sendvars.end(), var) != sendvars.end()) {
-        is_dist_train_op = true;
-        break;
+        return true;
       }
     }
-    return is_dist_train_op;
+    return false;
   };
 
   if (op.Type() == "split") {
@@ -117,13 +115,7 @@ std::unique_ptr<SSAGraph> MultiDevSSAGraphBuilder::Build(
       places_.size());
 
   // Find "send" op first for split is in front of send.
-  OpDesc *send_op = nullptr;
-  for (auto *op : program.Block(0).AllOps()) {
-    if (op->Type() == "send") {
-      send_op = op;
-      break;
-    }
-  }
+  OpDesc *send_op = GetSendOpDesc(program);
 
   bool is_forwarding = true;
   for (auto *op : program.Block(0).AllOps()) {
@@ -134,6 +126,7 @@ std::unique_ptr<SSAGraph> MultiDevSSAGraphBuilder::Build(
     } else if (IsDistTrainOp(*op, send_op)) {
       CreateComputationalOps(&result, *op, 1);
     } else if (IsScaleLossOp(*op)) {
+      // user can customize loss@grad if skip_scale_loss_
       if (!skip_scale_loss_) {
         CreateScaleLossGradOp(&result);
       }
@@ -142,10 +135,7 @@ std::unique_ptr<SSAGraph> MultiDevSSAGraphBuilder::Build(
       CreateComputationalOps(&result, *op, places_.size());
       if (!is_forwarding) {
         // Currently, we assume that once gradient is generated, it can be
-        // broadcast, and each gradient is only broadcast once. But there are no
-        // other cases, for example, we need to adjust the gradient according to
-        // the input when we get the gradient, which is not considered at
-        // present.
+        // broadcast, and each gradient is only broadcast once.
         for (auto &og : op->OutputArgumentNames()) {
           if (IsParameterGradientOnce(og, &og_has_been_broadcast)) {
             InsertNCCLAllReduceOp(&result, og);
@@ -175,6 +165,16 @@ std::unique_ptr<SSAGraph> MultiDevSSAGraphBuilder::Build(
   return std::unique_ptr<SSAGraph>(graph);
 }
 
+OpDesc *MultiDevSSAGraphBuilder::GetSendOpDesc(
+    const ProgramDesc &program) const {
+  for (auto *op : program.Block(0).AllOps()) {
+    if (op->Type() == "send") {
+      return op;
+    }
+  }
+  return nullptr;
+}
+
 void MultiDevSSAGraphBuilder::InsertNCCLAllReduceOp(
     SSAGraph *result, const std::string &og) const {
 #ifdef PADDLE_WITH_CUDA
@@ -243,7 +243,7 @@ void MultiDevSSAGraphBuilder::CreateComputationalOps(SSAGraph *result,
     auto p = places_[scope_idx];
     auto s = local_scopes_[scope_idx];
     result->ops_.emplace_back(new ComputationOpHandle(op, s, p));
-    CreateOpHandleIOs(result, op, p, scope_idx);
+    CreateOpHandleIOs(result, op, scope_idx);
   }
 }
 
@@ -255,7 +255,7 @@ void MultiDevSSAGraphBuilder::CreateSendOp(SSAGraph *result,
   result->ops_.emplace_back(new SendOpHandle(op, s, p));
   // Create inputs for output on original place and no ssa output
   // is created for send op.
-  CreateOpHandleIOs(result, op, p, 0);
+  CreateOpHandleIOs(result, op, 0);
 }
 
 bool MultiDevSSAGraphBuilder::IsScaleLossOp(const OpDesc &op) const {

paddle/fluid/framework/details/multi_devices_graph_builder.h

@@ -48,7 +48,7 @@ class MultiDevSSAGraphBuilder : public SSAGraphBuilder {
 
  private:
   void CreateOpHandleIOs(SSAGraph *result, const OpDesc &op,
-                         const platform::Place &p, const size_t &i) const;
+                         size_t place_id) const;
 
  private:
   std::string loss_var_name_;
@@ -65,6 +65,9 @@ class MultiDevSSAGraphBuilder : public SSAGraphBuilder {
 
   void CreateSendOp(SSAGraph *result, const OpDesc &op) const;
 
+  /**
+   * Is this operator as the end-point operator before/after send operator.
+   */
   bool IsDistTrainOp(const OpDesc &op, OpDesc *send_op) const;
 
   void CreateComputationalOps(SSAGraph *result, const OpDesc &op,
@@ -77,6 +80,12 @@ class MultiDevSSAGraphBuilder : public SSAGraphBuilder {
       std::unordered_set<std::string> *og_has_been_broadcast) const;
 
   void InsertNCCLAllReduceOp(SSAGraph *result, const std::string &og) const;
+
+  /**
+   * Get send op in the global block of program.
+   * nullptr if not found.
+   */
+  OpDesc *GetSendOpDesc(const ProgramDesc &program) const;
 };
 }  // namespace details
 }  // namespace framework

paddle/fluid/framework/details/reduce_op_handle_test.cc

@@ -194,7 +194,7 @@ struct TestReduceOpHandle {
     }
 
     f::Tensor result_tensor;
-    f::TensorCopy(rt, cpu_place, *(ctxs_[output_scope_idx]), &result_tensor);
+    f::TensorCopySync(rt, cpu_place, &result_tensor);
     float *ct = result_tensor.data<float>();
 
     for (int64_t j = 0; j < f::product(result_tensor.dims()); ++j) {
@@ -239,7 +239,7 @@ struct TestReduceOpHandle {
     auto &rt = out_var->Get<f::LoDTensor>();
 
     f::Tensor result_tensor;
-    f::TensorCopy(rt, cpu_place, *(ctxs_[output_scope_idx]), &result_tensor);
+    f::TensorCopySync(rt, cpu_place, &result_tensor);
     float *ct = result_tensor.data<float>();
 
     for (int64_t j = 0; j < f::product(result_tensor.dims()); ++j) {

paddle/fluid/framework/details/scale_loss_grad_op_handle.cc

@@ -46,6 +46,7 @@ void ScaleLossGradOpHandle::RunImpl() {
               ->stream();
       memory::Copy(boost::get<platform::CUDAPlace>(place_), tmp,
                    platform::CPUPlace(), &coeff_, sizeof(float), stream);
+      VLOG(1) << place_ << "RUN Scale loss grad op";
     });
 #endif
   }

paddle/fluid/framework/details/ssa_graph.h

@@ -25,12 +25,22 @@ namespace paddle {
 namespace framework {
 namespace details {
 
+// A SSA graph used by parallel executor.
 struct SSAGraph {
+  // all variable in each devices.
+  // The outside vector is the device vector. Each element of this vector is a
+  // map from variable name to variables. The variables, who have the same name,
+  // will have a different version. The offset in the
+  // `std::vector<std::unique_ptr<VarHandle>>` is the version of varaibles.
   std::vector<
       std::unordered_map<std::string, std::vector<std::unique_ptr<VarHandle>>>>
       vars_;
+
   // aux variables to represent dependency. Useful to resolve data hazard.
   std::unordered_set<std::unique_ptr<VarHandleBase>> dep_vars_;
+
+  // all operators. NOTE that even we use a vector here, the operators is
+  // unordered.
   std::vector<std::unique_ptr<OpHandleBase>> ops_;
 };
 

paddle/fluid/framework/details/ssa_graph_builder.h

@@ -48,6 +48,8 @@ class SSAGraphBuilder {
                                                const platform::Place &place,
                                                size_t place_offset);
 
+  // Add an output variable (each_var_name, place, place_offset) to op_handle,
+  // which belongs to graph
   static void CreateOpOutput(SSAGraph *graph, OpHandleBase *op_handle,
                              const std::string &each_var_name,
                              const platform::Place &place, size_t place_offset);

paddle/fluid/framework/executor.cc

@@ -226,15 +226,15 @@ static bool has_fetch_operators(
 }
 
 void Executor::Run(const ProgramDesc& program, Scope* scope,
-                   std::map<std::string, const LoDTensor*>& feed_targets,
-                   std::map<std::string, LoDTensor*>& fetch_targets,
+                   std::map<std::string, const LoDTensor*>* feed_targets,
+                   std::map<std::string, LoDTensor*>* fetch_targets,
                    bool create_vars, const std::string& feed_holder_name,
                    const std::string& fetch_holder_name) {
   platform::RecordBlock b(kProgramId);
   bool has_feed_ops =
-      has_feed_operators(program.Block(0), feed_targets, feed_holder_name);
+      has_feed_operators(program.Block(0), *feed_targets, feed_holder_name);
   bool has_fetch_ops =
-      has_fetch_operators(program.Block(0), fetch_targets, fetch_holder_name);
+      has_fetch_operators(program.Block(0), *fetch_targets, fetch_holder_name);
 
   ProgramDesc* copy_program = const_cast<ProgramDesc*>(&program);
   if (!has_feed_ops || !has_fetch_ops) {
@@ -250,7 +250,7 @@ void Executor::Run(const ProgramDesc& program, Scope* scope,
     feed_holder->SetPersistable(true);
 
     int i = 0;
-    for (auto& feed_target : feed_targets) {
+    for (auto& feed_target : (*feed_targets)) {
       std::string var_name = feed_target.first;
       VLOG(3) << "feed target's name: " << var_name;
 
@@ -273,7 +273,7 @@ void Executor::Run(const ProgramDesc& program, Scope* scope,
     fetch_holder->SetPersistable(true);
 
     int i = 0;
-    for (auto& fetch_target : fetch_targets) {
+    for (auto& fetch_target : (*fetch_targets)) {
       std::string var_name = fetch_target.first;
       VLOG(3) << "fetch target's name: " << var_name;
 
@@ -361,16 +361,16 @@ void Executor::RunPreparedContext(ExecutorPrepareContext* ctx, Scope* scope,
 
 void Executor::RunPreparedContext(
     ExecutorPrepareContext* ctx, Scope* scope,
-    std::map<std::string, const LoDTensor*>& feed_targets,
-    std::map<std::string, LoDTensor*>& fetch_targets, bool create_vars,
+    std::map<std::string, const LoDTensor*>* feed_targets,
+    std::map<std::string, LoDTensor*>* fetch_targets, bool create_vars,
     const std::string& feed_holder_name, const std::string& fetch_holder_name) {
   auto& global_block = ctx->prog_.Block(ctx->block_id_);
 
   PADDLE_ENFORCE(
-      has_feed_operators(global_block, feed_targets, feed_holder_name),
+      has_feed_operators(global_block, *feed_targets, feed_holder_name),
       "Program in ExecutorPrepareContext should has feed_ops.");
   PADDLE_ENFORCE(
-      has_fetch_operators(global_block, fetch_targets, fetch_holder_name),
+      has_fetch_operators(global_block, *fetch_targets, fetch_holder_name),
       "Program in the prepared context should has fetch_ops.");
 
   // map the data of feed_targets to feed_holder
@@ -378,8 +378,8 @@ void Executor::RunPreparedContext(
     if (op->Type() == kFeedOpType) {
       std::string feed_target_name = op->Output("Out")[0];
       int idx = boost::get<int>(op->GetAttr("col"));
-      SetFeedVariable(scope, *feed_targets[feed_target_name], feed_holder_name,
-                      idx);
+      SetFeedVariable(scope, *(*feed_targets)[feed_target_name],
+                      feed_holder_name, idx);
     }
   }
 
@@ -390,7 +390,7 @@ void Executor::RunPreparedContext(
     if (op->Type() == kFetchOpType) {
       std::string fetch_target_name = op->Input("X")[0];
       int idx = boost::get<int>(op->GetAttr("col"));
-      *fetch_targets[fetch_target_name] =
+      *(*fetch_targets)[fetch_target_name] =
           GetFetchVariable(*scope, fetch_holder_name, idx);
     }
   }

paddle/fluid/framework/executor.h

@@ -55,8 +55,8 @@ class Executor {
            bool create_local_scope = true, bool create_vars = true);
 
   void Run(const ProgramDesc& program, Scope* scope,
-           std::map<std::string, const LoDTensor*>& feed_targets,
-           std::map<std::string, LoDTensor*>& fetch_targets,
+           std::map<std::string, const LoDTensor*>* feed_targets,
+           std::map<std::string, LoDTensor*>* fetch_targets,
            bool create_vars = true,
            const std::string& feed_holder_name = "feed",
            const std::string& fetch_holder_name = "fetch");
@@ -74,8 +74,8 @@ class Executor {
                           bool create_vars = true);
 
   void RunPreparedContext(ExecutorPrepareContext* ctx, Scope* scope,
-                          std::map<std::string, const LoDTensor*>& feed_targets,
-                          std::map<std::string, LoDTensor*>& fetch_targets,
+                          std::map<std::string, const LoDTensor*>* feed_targets,
+                          std::map<std::string, LoDTensor*>* fetch_targets,
                           bool create_vars = true,
                           const std::string& feed_holder_name = "feed",
                           const std::string& fetch_holder_name = "fetch");

paddle/fluid/framework/init.cc

@@ -15,7 +15,6 @@ limitations under the License. */
 #include <algorithm>
 #include <stdexcept>
 #include <string>
-#include <vector>
 
 #include "paddle/fluid/framework/init.h"
 #include "paddle/fluid/framework/operator.h"
@@ -31,6 +30,7 @@ std::once_flag p2p_init_flag;
 
 void InitGflags(std::vector<std::string> argv) {
   std::call_once(gflags_init_flag, [&]() {
+    argv.insert(argv.begin(), "dummy");
     int argc = argv.size();
     char **arr = new char *[argv.size()];
     std::string line;
@@ -44,20 +44,23 @@ void InitGflags(std::vector<std::string> argv) {
   });
 }
 
-void InitP2P(int count) {
+void InitP2P(std::vector<int> devices) {
 #ifdef PADDLE_WITH_CUDA
   std::call_once(p2p_init_flag, [&]() {
+    int count = devices.size();
     for (int i = 0; i < count; ++i) {
       for (int j = 0; j < count; ++j) {
-        if (i == j) continue;
+        if (devices[i] == devices[j]) continue;
         int can_acess = -1;
-        PADDLE_ENFORCE(cudaDeviceCanAccessPeer(&can_acess, i, j),
-                       "Failed to test P2P access.");
+        PADDLE_ENFORCE(
+            cudaDeviceCanAccessPeer(&can_acess, devices[i], devices[j]),
+            "Failed to test P2P access.");
         if (can_acess != 1) {
-          LOG(WARNING) << "Cannot enable P2P access from " << i << " to " << j;
+          LOG(WARNING) << "Cannot enable P2P access from " << devices[i]
+                       << " to " << devices[j];
         } else {
-          cudaSetDevice(i);
-          cudaDeviceEnablePeerAccess(j, 0);
+          cudaSetDevice(devices[i]);
+          cudaDeviceEnablePeerAccess(devices[j], 0);
         }
       }
     }
@@ -67,11 +70,26 @@ void InitP2P(int count) {
 
 void InitDevices(bool init_p2p) {
   /*Init all available devices by default */
+  std::vector<int> devices;
+#ifdef PADDLE_WITH_CUDA
+  try {
+    int count = platform::GetCUDADeviceCount();
+    for (int i = 0; i < count; ++i) {
+      devices.push_back(i);
+    }
+  } catch (const std::exception &exp) {
+    LOG(WARNING) << "Compiled with WITH_GPU, but no GPU found in runtime.";
+  }
+#else
+  LOG(WARNING)
+      << "'CUDA' is not supported, Please re-compile with WITH_GPU option";
+#endif
+  InitDevices(init_p2p, devices);
+}
 
+void InitDevices(bool init_p2p, const std::vector<int> devices) {
   std::vector<platform::Place> places;
-  places.emplace_back(platform::CPUPlace());
   int count = 0;
-
 #ifdef PADDLE_WITH_CUDA
   try {
     count = platform::GetCUDADeviceCount();
@@ -83,12 +101,17 @@ void InitDevices(bool init_p2p) {
       << "'CUDA' is not supported, Please re-compile with WITH_GPU option";
 #endif
 
-  for (int i = 0; i < count; ++i) {
-    places.emplace_back(platform::CUDAPlace(i));
+  for (size_t i = 0; i < devices.size(); ++i) {
+    if (devices[i] >= count || devices[i] < 0) {
+      LOG(WARNING) << "Invalid devices id.";
+      continue;
+    }
+    places.emplace_back(platform::CUDAPlace(devices[i]));
   }
   if (init_p2p) {
-    InitP2P(count);
+    InitP2P(devices);
   }
+  places.emplace_back(platform::CPUPlace());
   platform::DeviceContextPool::Init(places);
 }
 

paddle/fluid/framework/init.h

@@ -28,5 +28,7 @@ void InitGLOG(const std::string &prog_name);
 
 void InitDevices(bool init_p2p);
 
+void InitDevices(bool init_p2p, const std::vector<int> devices);
+
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/tensor_util.cc

@@ -20,7 +20,7 @@ namespace paddle {
 namespace framework {
 
 void TensorCopy(const Tensor& src, const platform::Place& dst_place,
-                const platform::DeviceContext& ctx, Tensor* dst, bool sync) {
+                const platform::DeviceContext& ctx, Tensor* dst) {
   VLOG(3) << "TensorCopy " << src.dims() << " from " << src.place() << " to "
           << dst_place;
   src.check_memory_size();
@@ -48,9 +48,7 @@ void TensorCopy(const Tensor& src, const platform::Place& dst_place,
     auto ctx_gpu_place = boost::get<platform::CUDAPlace>(ctx_place);
     PADDLE_ENFORCE_EQ(src_gpu_place, ctx_gpu_place);
     auto stream =
-        sync ? nullptr
-             : reinterpret_cast<const platform::CUDADeviceContext&>(ctx)
-                   .stream();
+        reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream();
     memory::Copy(dst_cpu_place, dst_ptr, src_gpu_place, src_ptr, size, stream);
   } else if (platform::is_cpu_place(src_place) &&
              platform::is_gpu_place(dst_place)) {
@@ -61,9 +59,7 @@ void TensorCopy(const Tensor& src, const platform::Place& dst_place,
     auto ctx_gpu_place = boost::get<platform::CUDAPlace>(ctx_place);
     PADDLE_ENFORCE_EQ(dst_gpu_place, ctx_gpu_place);
     auto stream =
-        sync ? nullptr
-             : reinterpret_cast<const platform::CUDADeviceContext&>(ctx)
-                   .stream();
+        reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream();
     memory::Copy(dst_gpu_place, dst_ptr, src_cpu_place, src_ptr, size, stream);
   } else if (platform::is_gpu_place(src_place) &&
              platform::is_gpu_place(dst_place)) {
@@ -72,9 +68,7 @@ void TensorCopy(const Tensor& src, const platform::Place& dst_place,
     auto ctx_place = ctx.GetPlace();
     PADDLE_ENFORCE(platform::is_gpu_place(ctx_place));
     auto stream =
-        sync ? nullptr
-             : reinterpret_cast<const platform::CUDADeviceContext&>(ctx)
-                   .stream();
+        reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream();
     memory::Copy(dst_gpu_place, dst_ptr, src_gpu_place, src_ptr, size, stream);
   }
 #endif
@@ -92,6 +86,41 @@ void TensorCopy(const Tensor& src, const platform::Place& dst_place,
   TensorCopy(src, dst_place, *dev_ctx, dst);
 }
 
+void TensorCopySync(const Tensor& src, const platform::Place& dst_place,
+                    Tensor* dst) {
+  VLOG(3) << "TensorCopySync " << src.dims() << " from " << src.place()
+          << " to " << dst_place;
+  src.check_memory_size();
+  dst->Resize(src.dims());
+  dst->set_layout(src.layout());
+  auto src_place = src.place();
+  auto src_ptr = src.data<void>();
+  auto dst_ptr = dst->mutable_data(dst_place, src.type());
+  auto size = src.numel() * SizeOfType(src.type());
+  if (platform::is_cpu_place(src_place) && platform::is_cpu_place(dst_place)) {
+    memory::Copy(boost::get<platform::CPUPlace>(dst_place), dst_ptr,
+                 boost::get<platform::CPUPlace>(src_place), src_ptr, size);
+  }
+#ifdef PADDLE_WITH_CUDA
+  else if (platform::is_gpu_place(src_place) &&  // NOLINT
+           platform::is_cpu_place(dst_place)) {
+    auto src_gpu_place = boost::get<platform::CUDAPlace>(src_place);
+    auto dst_cpu_place = boost::get<platform::CPUPlace>(dst_place);
+    memory::Copy(dst_cpu_place, dst_ptr, src_gpu_place, src_ptr, size, nullptr);
+  } else if (platform::is_cpu_place(src_place) &&
+             platform::is_gpu_place(dst_place)) {
+    auto src_cpu_place = boost::get<platform::CPUPlace>(src_place);
+    auto dst_gpu_place = boost::get<platform::CUDAPlace>(dst_place);
+    memory::Copy(dst_gpu_place, dst_ptr, src_cpu_place, src_ptr, size, nullptr);
+  } else if (platform::is_gpu_place(src_place) &&
+             platform::is_gpu_place(dst_place)) {
+    auto src_gpu_place = boost::get<platform::CUDAPlace>(src_place);
+    auto dst_gpu_place = boost::get<platform::CUDAPlace>(dst_place);
+    memory::Copy(dst_gpu_place, dst_ptr, src_gpu_place, src_ptr, size, nullptr);
+  }
+#endif
+}
+
 template <typename Predicate, typename DevCtx>
 struct AnyDTypeVisitor {
   Predicate predicate_;

paddle/fluid/framework/tensor_util.h

@@ -24,10 +24,11 @@ namespace paddle {
 namespace framework {
 
 void TensorCopy(const Tensor& src, const platform::Place& dst_place,
-                const platform::DeviceContext& ctx, Tensor* dst,
-                bool sync = false);
+                const platform::DeviceContext& ctx, Tensor* dst);
 void TensorCopy(const Tensor& src, const platform::Place& dst_place,
                 Tensor* dst);
+void TensorCopySync(const Tensor& src, const platform::Place& dst_place,
+                    Tensor* dst);
 
 template <typename T>
 void TensorFromVector(const std::vector<T>& src,

paddle/fluid/inference/engine.h

+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include "paddle/fluid/framework/framework.pb.h"
+
+namespace paddle {
+namespace inference {
+
+/*
+ * EngineBase is the base class of all inference engines. An inference engine
+ * takes a paddle program as input, and outputs the result in fluid Tensor
+ * format. It can be used to optimize performance of computation sub-blocks, for
+ * example, break down the original block into sub-blocks and execute each
+ * sub-blocks in different engines.
+ *
+ * For example:
+ *   When inference, the resnet50 model can put most of the model into subgraph
+ * and run it on a TensorRT engine.
+ *
+ * There are several engines such as TensorRT and other frameworks, so an
+ * EngineBase is put forward to give an unified interface for all the
+ * different engine implemention.
+ */
+class EngineBase {
+ public:
+  using DescType = ::paddle::framework::proto::BlockDesc;
+
+  // Build the model and do some preparation, for example, in TensorRT, run
+  // createInferBuilder, buildCudaEngine.
+  virtual void Build(const DescType& paddle_model) = 0;
+
+  // Execute the engine, that will run the inference network.
+  virtual void Execute(int batch_size) = 0;
+
+  virtual ~EngineBase() {}
+};  // class EngineBase
+
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/io.cc

@@ -16,17 +16,29 @@ limitations under the License. */
 
 #include <algorithm>
 #include <fstream>
+#include <vector>
 #include "paddle/fluid/framework/block_desc.h"
 #include "paddle/fluid/framework/feed_fetch_type.h"
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/pybind/pybind.h"
 
+DEFINE_string(devices, "", "The devices to be used which is joined by comma.");
+DEFINE_bool(init_p2p, false, "Whether to init p2p.");
+
 namespace paddle {
 namespace inference {
 
-// Temporarily add this function for exposing framework::InitDevices() when
-// linking the inference shared library.
-void Init(bool init_p2p) { framework::InitDevices(init_p2p); }
+void Init(const std::vector<std::string> argv) {
+  framework::InitGflags(argv);
+  // init devices
+  std::vector<int> devices;
+  std::string token;
+  std::istringstream tokenStream(FLAGS_devices);
+  while (std::getline(tokenStream, token, ',')) {
+    devices.push_back(std::stoi(token));
+  }
+  framework::InitDevices(FLAGS_init_p2p, devices);
+}
 
 void ReadBinaryFile(const std::string& filename, std::string* contents) {
   std::ifstream fin(filename, std::ios::in | std::ios::binary);

paddle/fluid/inference/io.h

@@ -25,7 +25,7 @@ limitations under the License. */
 namespace paddle {
 namespace inference {
 
-void Init(bool init_p2p);
+void Init(const std::vector<std::string> argv);
 
 void LoadPersistables(framework::Executor* executor, framework::Scope* scope,
                       const framework::ProgramDesc& main_program,

paddle/fluid/inference/tensorrt/CMakeLists.txt

-nv_test(test_tensorrt SRCS test_tensorrt.cc DEPS dynload_cuda device_context dynamic_loader)
+if(WITH_TESTING)
+  nv_test(test_tensorrt SRCS test_tensorrt.cc DEPS dynload_cuda device_context dynamic_loader)
+  nv_test(test_tensorrt_engine SRCS test_engine.cc engine.cc DEPS dynload_cuda)
+endif()

paddle/fluid/inference/tensorrt/engine.cc

+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+
+#include <NvInfer.h>
+#include <cuda.h>
+#include <glog/logging.h>
+#include <string>
+#include "paddle/fluid/inference/tensorrt/helper.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+void TensorRTEngine::Build(const DescType& paddle_model) {
+  PADDLE_ENFORCE(false, "not implemented");
+}
+
+void TensorRTEngine::Execute(int batch_size) {
+  infer_context_->enqueue(batch_size, buffers_.data(), *stream_, nullptr);
+  cudaStreamSynchronize(*stream_);
+}
+
+TensorRTEngine::~TensorRTEngine() {
+  // clean buffer
+  for (auto& buffer : buffers_) {
+    if (buffer != nullptr) {
+      PADDLE_ENFORCE_EQ(0, cudaFree(buffer));
+      buffer = nullptr;
+    }
+  }
+}
+
+void TensorRTEngine::FreezeNetwork() {
+  PADDLE_ENFORCE(infer_builder_ != nullptr,
+                 "Call InitNetwork first to initialize network.");
+  PADDLE_ENFORCE(infer_network_ != nullptr,
+                 "Call InitNetwork first to initialize network.");
+  // build engine.
+  infer_builder_->setMaxBatchSize(max_batch_);
+  infer_builder_->setMaxWorkspaceSize(max_workspace_);
+
+  infer_engine_.reset(infer_builder_->buildCudaEngine(*infer_network_));
+  PADDLE_ENFORCE(infer_engine_ != nullptr, "build cuda engine failed!");
+
+  infer_context_.reset(infer_engine_->createExecutionContext());
+
+  // allocate GPU buffers.
+  buffers_.resize(buffer_sizes_.size(), nullptr);
+  for (auto& item : buffer_sizes_) {
+    if (item.second == 0) {
+      auto slot_offset = infer_engine_->getBindingIndex(item.first.c_str());
+      item.second = kDataTypeSize[static_cast<int>(
+                        infer_engine_->getBindingDataType(slot_offset))] *
+                    AccumDims(infer_engine_->getBindingDimensions(slot_offset));
+    }
+    PADDLE_ENFORCE_EQ(0, cudaMalloc(&buffer(item.first), item.second));
+  }
+}
+
+nvinfer1::ITensor* TensorRTEngine::DeclareInput(const std::string& name,
+                                                nvinfer1::DataType dtype,
+                                                const nvinfer1::Dims& dim) {
+  PADDLE_ENFORCE_EQ(0, buffer_sizes_.count(name), "duplicate input name %s",
+                    name);
+
+  PADDLE_ENFORCE(infer_network_ != nullptr, "should initnetwork first");
+  auto* input = infer_network_->addInput(name.c_str(), dtype, dim);
+  PADDLE_ENFORCE(input, "infer network add input %s failed", name);
+
+  buffer_sizes_[name] = kDataTypeSize[static_cast<int>(dtype)] * AccumDims(dim);
+  return input;
+}
+
+void TensorRTEngine::DeclareOutput(const nvinfer1::ILayer* layer, int offset,
+                                   const std::string& name) {
+  PADDLE_ENFORCE_EQ(0, buffer_sizes_.count(name), "duplicate output name %s",
+                    name);
+
+  auto* output = layer->getOutput(offset);
+  PADDLE_ENFORCE(output != nullptr);
+  output->setName(name.c_str());
+  infer_network_->markOutput(*output);
+  // output buffers' size can only be decided latter, set zero here to mark this
+  // and will reset latter.
+  buffer_sizes_[name] = 0;
+}
+
+void* TensorRTEngine::GetOutputInGPU(const std::string& name) {
+  return buffer(name);
+}
+
+void TensorRTEngine::GetOutputInCPU(const std::string& name, void* dst,
+                                    size_t max_size) {
+  // determine data size
+  auto it = buffer_sizes_.find(name);
+  PADDLE_ENFORCE(it != buffer_sizes_.end());
+  PADDLE_ENFORCE_GT(it->second, 0);
+  PADDLE_ENFORCE_GE(max_size, it->second);
+
+  PADDLE_ENFORCE_EQ(0, cudaMemcpyAsync(dst, buffer(name), it->second,
+                                       cudaMemcpyDeviceToHost, *stream_));
+}
+
+void*& TensorRTEngine::buffer(const std::string& name) {
+  PADDLE_ENFORCE(infer_engine_ != nullptr, "call FreezeNetwork first.");
+  auto it = buffer_sizes_.find(name);
+  PADDLE_ENFORCE(it != buffer_sizes_.end());
+  auto slot_offset = infer_engine_->getBindingIndex(name.c_str());
+  return buffers_[slot_offset];
+}
+
+void TensorRTEngine::SetInputFromCPU(const std::string& name, void* data,
+                                     size_t size) {
+  void* buf = buffer(name);
+  PADDLE_ENFORCE_EQ(
+      0, cudaMemcpyAsync(buf, data, size, cudaMemcpyHostToDevice, *stream_));
+}
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/engine.h

+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include <NvInfer.h>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+#include "paddle/fluid/inference/engine.h"
+#include "paddle/fluid/inference/tensorrt/helper.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+/*
+ * TensorRT Engine.
+ *
+ * There are two alternative ways to use it, one is  to build from a paddle
+ * protobuf model, another way is to manully construct the network.
+ */
+class TensorRTEngine : public EngineBase {
+ public:
+  // Weight is model parameter.
+  class Weight {
+   public:
+    Weight(nvinfer1::DataType dtype, void* value, int num_elem) {
+      w_.type = dtype;
+      w_.values = value;
+      w_.count = num_elem;
+    }
+    const nvinfer1::Weights& get() { return w_; }
+
+   private:
+    nvinfer1::Weights w_;
+  };
+
+  TensorRTEngine(int max_batch, int max_workspace, cudaStream_t* stream,
+                 nvinfer1::ILogger& logger = NaiveLogger::Global())
+      : max_batch_(max_batch),
+        max_workspace_(max_workspace),
+        stream_(stream),
+        logger_(logger) {}
+
+  virtual ~TensorRTEngine();
+
+  // TODO(Superjomn) implement it later when graph segmentation is supported.
+  void Build(const DescType& paddle_model) override;
+
+  void Execute(int batch_size) override;
+
+  // Initialize the inference network, so that TensorRT layers can add to this
+  // network.
+  void InitNetwork() {
+    infer_builder_.reset(createInferBuilder(logger_));
+    infer_network_.reset(infer_builder_->createNetwork());
+  }
+  // After finishing adding ops, freeze this network and creates the executation
+  // environment.
+  void FreezeNetwork();
+
+  // Add an input and set its name, data type and dimention.
+  nvinfer1::ITensor* DeclareInput(const std::string& name,
+                                  nvinfer1::DataType dtype,
+                                  const nvinfer1::Dims& dim);
+  // Set the offset-th output from a layer as the network's output, and set its
+  // name.
+  void DeclareOutput(const nvinfer1::ILayer* layer, int offset,
+                     const std::string& name);
+
+  // GPU memory address for an ITensor with specific name. One can operate on
+  // these memory directly for acceleration, for example, output the converted
+  // data directly to the buffer to save data copy overhead.
+  // NOTE this should be used after calling `FreezeNetwork`.
+  void*& buffer(const std::string& name);
+
+  // Fill an input from CPU memory with name and size.
+  void SetInputFromCPU(const std::string& name, void* data, size_t size);
+  // TODO(Superjomn) is this method necessary given that buffer(xxx) can be
+  // accessed directly. Fill an input from GPU memory with name and size.
+  void SetInputFromGPU(const std::string& name, void* data, size_t size);
+  // Get an output called name, the output of tensorrt is in GPU, so this method
+  // will just return the output's GPU memory address.
+  void* GetOutputInGPU(const std::string& name);
+  // LOW EFFICENCY! Get output to CPU, this will trigger a memory copy from GPU
+  // to CPU.
+  void GetOutputInCPU(const std::string& name, void* dst, size_t max_size);
+
+  nvinfer1::ICudaEngine* engine() { return infer_engine_.get(); }
+  nvinfer1::INetworkDefinition* network() { return infer_network_.get(); }
+
+ private:
+  // the max batch size
+  int max_batch_;
+  // the max memory size the engine uses
+  int max_workspace_;
+  cudaStream_t* stream_;
+  nvinfer1::ILogger& logger_;
+
+  std::vector<void*> buffers_;
+  // max data size for the buffers.
+  std::unordered_map<std::string /*name*/, size_t /*max size*/> buffer_sizes_;
+
+  // TensorRT related internal members
+  template <typename T>
+  struct Destroyer {
+    void operator()(T* x) { x->destroy(); }
+  };
+  template <typename T>
+  using infer_ptr = std::unique_ptr<T, Destroyer<T>>;
+  infer_ptr<nvinfer1::IBuilder> infer_builder_;
+  infer_ptr<nvinfer1::INetworkDefinition> infer_network_;
+  infer_ptr<nvinfer1::ICudaEngine> infer_engine_;
+  infer_ptr<nvinfer1::IExecutionContext> infer_context_;
+};  // class TensorRTEngine
+
+// Add an layer__ into engine__ with args ARGS.
+// For example:
+//   TRT_ENGINE_ADD_LAYER(xxx, FullyConnected, input, dim, weights, bias)
+//
+// Reference
+// https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#charRNN_define_network
+//
+// will add a fully connected layer into the engine.
+// TensorRT has too many layers, so that is not wise to add member functions for
+// them, and an macro like this is more extensible when underlying TensorRT
+// library add new layer supports.
+#define TRT_ENGINE_ADD_LAYER(engine__, layer__, ARGS...) \
+  engine__->network()->add##layer__(ARGS);
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/helper.h

+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License. */
+
+#pragma once
+
+#include <NvInfer.h>
+#include <cuda.h>
+#include <glog/logging.h>
+#include "paddle/fluid/platform/dynload/tensorrt.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+namespace dy = paddle::platform::dynload;
+
+static size_t AccumDims(nvinfer1::Dims dims) {
+  size_t num = dims.nbDims == 0 ? 0 : 1;
+  for (int i = 0; i < dims.nbDims; i++) {
+    PADDLE_ENFORCE_GT(dims.d[i], 0);
+    num *= dims.d[i];
+  }
+  return num;
+}
+
+// TensorRT data type to size
+const int kDataTypeSize[] = {
+    4,  // kFLOAT
+    2,  // kHALF
+    1,  // kINT8
+    4   // kINT32
+};
+
+// The following two API are implemented in TensorRT's header file, cannot load
+// from the dynamic library. So create our own implementation and directly
+// trigger the method from the dynamic library.
+static nvinfer1::IBuilder* createInferBuilder(nvinfer1::ILogger& logger) {
+  return static_cast<nvinfer1::IBuilder*>(
+      dy::createInferBuilder_INTERNAL(&logger, NV_TENSORRT_VERSION));
+}
+static nvinfer1::IRuntime* createInferRuntime(nvinfer1::ILogger& logger) {
+  return static_cast<nvinfer1::IRuntime*>(
+      dy::createInferRuntime_INTERNAL(&logger, NV_TENSORRT_VERSION));
+}
+
+// A logger for create TensorRT infer builder.
+class NaiveLogger : public nvinfer1::ILogger {
+ public:
+  void log(nvinfer1::ILogger::Severity severity, const char* msg) override {
+    switch (severity) {
+      case Severity::kINFO:
+        LOG(INFO) << msg;
+        break;
+      case Severity::kWARNING:
+        LOG(WARNING) << msg;
+        break;
+      case Severity::kINTERNAL_ERROR:
+      case Severity::kERROR:
+        LOG(ERROR) << msg;
+        break;
+      default:
+        break;
+    }
+  }
+
+  static nvinfer1::ILogger& Global() {
+    static nvinfer1::ILogger* x = new NaiveLogger;
+    return *x;
+  }
+
+  virtual ~NaiveLogger() override {}
+};
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/test_engine.cc

+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include <cuda.h>
+#include <cuda_runtime_api.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+class TensorRTEngineTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    ASSERT_EQ(0, cudaStreamCreate(&stream_));
+    engine_ = new TensorRTEngine(1, 1 << 10, &stream_);
+    engine_->InitNetwork();
+  }
+
+  void TearDown() override {
+    delete engine_;
+    cudaStreamDestroy(stream_);
+  }
+
+ protected:
+  TensorRTEngine* engine_;
+  cudaStream_t stream_;
+};
+
+TEST_F(TensorRTEngineTest, add_layer) {
+  const int size = 1;
+
+  float raw_weight[size] = {2.};  // Weight in CPU memory.
+  float raw_bias[size] = {3.};
+
+  LOG(INFO) << "create weights";
+  TensorRTEngine::Weight weight(nvinfer1::DataType::kFLOAT, raw_weight, size);
+  TensorRTEngine::Weight bias(nvinfer1::DataType::kFLOAT, raw_bias, size);
+  auto* x = engine_->DeclareInput("x", nvinfer1::DataType::kFLOAT,
+                                  nvinfer1::DimsCHW{1, 1, 1});
+  auto* fc_layer = TRT_ENGINE_ADD_LAYER(engine_, FullyConnected, *x, size,
+                                        weight.get(), bias.get());
+  PADDLE_ENFORCE(fc_layer != nullptr);
+
+  engine_->DeclareOutput(fc_layer, 0, "y");
+  LOG(INFO) << "freeze network";
+  engine_->FreezeNetwork();
+  ASSERT_EQ(engine_->engine()->getNbBindings(), 2);
+
+  // fill in real data
+  float x_v = 1234;
+  engine_->SetInputFromCPU("x", reinterpret_cast<void*>(&x_v),
+                           1 * sizeof(float));
+  LOG(INFO) << "to execute";
+  engine_->Execute(1);
+
+  LOG(INFO) << "to get output";
+  // void* y_v =
+  float y_cpu;
+  engine_->GetOutputInCPU("y", &y_cpu, sizeof(float));
+
+  LOG(INFO) << "to checkout output";
+  ASSERT_EQ(y_cpu, x_v * 2 + 3);
+}
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/test_tensorrt.cc

 /* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
 
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
 
-   http://www.apache.org/licenses/LICENSE-2.0
+http://www.apache.org/licenses/LICENSE-2.0
 
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License. */
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
 
 #include <glog/logging.h>
 #include <gtest/gtest.h>

paddle/fluid/inference/tests/book/test_inference_image_classification.cc

@@ -62,5 +62,21 @@ TEST(inference, image_classification) {
   LOG(INFO) << output2.dims();
 
   CheckError<float>(output1, output2);
+
+  // float16 inference requires cuda GPUs with >= 5.3 compute capability
+  if (paddle::platform::GetCUDAComputeCapability(0) >= 53) {
+    paddle::framework::LoDTensor output3;
+    std::vector<paddle::framework::LoDTensor*> cpu_fetchs3;
+    cpu_fetchs3.push_back(&output3);
+
+    LOG(INFO) << "--- GPU Runs in float16 mode: ---";
+    std::string fp16_dirname = dirname;
+    fp16_dirname.replace(fp16_dirname.find("book/"),
+                         std::string("book/").size(), "book/float16_");
+    TestInference<paddle::platform::CUDAPlace, false, true>(
+        fp16_dirname, cpu_feeds, cpu_fetchs3, FLAGS_repeat);
+
+    CheckError<float>(output2, output3);
+  }
 #endif
 }

paddle/fluid/inference/tests/test_helper.h

@@ -178,10 +178,10 @@ void TestInference(const std::string& dirname,
     std::unique_ptr<paddle::framework::ExecutorPrepareContext> ctx;
     if (PrepareContext) {
       ctx = executor.Prepare(*inference_program, 0);
-      executor.RunPreparedContext(ctx.get(), scope, feed_targets, fetch_targets,
-                                  CreateVars);
+      executor.RunPreparedContext(ctx.get(), scope, &feed_targets,
+                                  &fetch_targets, CreateVars);
     } else {
-      executor.Run(*inference_program, scope, feed_targets, fetch_targets,
+      executor.Run(*inference_program, scope, &feed_targets, &fetch_targets,
                    CreateVars);
     }
 
@@ -197,10 +197,10 @@ void TestInference(const std::string& dirname,
       if (PrepareContext) {
         // Note: if you change the inference_program, you need to call
         // executor.Prepare() again to get a new ExecutorPrepareContext.
-        executor.RunPreparedContext(ctx.get(), scope, feed_targets,
-                                    fetch_targets, CreateVars);
+        executor.RunPreparedContext(ctx.get(), scope, &feed_targets,
+                                    &fetch_targets, CreateVars);
       } else {
-        executor.Run(*inference_program, scope, feed_targets, fetch_targets,
+        executor.Run(*inference_program, scope, &feed_targets, &fetch_targets,
                      CreateVars);
       }
     }

paddle/fluid/operators/adam_op.h

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #pragma once
 #include <math.h>  // for sqrt in CPU and CUDA
+#include <Eigen/Dense>
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/operators/detail/safe_ref.h"
 #include "paddle/fluid/operators/math/selected_rows_functor.h"
@@ -24,8 +25,14 @@ namespace operators {
 
 namespace scatter = paddle::operators::math::scatter;
 
+struct GPUAdam;
+struct CPUAdam;
+
+template <typename T, typename Flavour>
+struct AdamFunctor;
+
 template <typename T>
-struct AdamFunctor {
+struct AdamFunctor<T, GPUAdam> {
   T beta1_;
   T beta2_;
   T epsilon_;
@@ -71,6 +78,7 @@ struct AdamFunctor {
 
     // Calculation
     lr *= sqrt(1 - beta2_pow) / (1 - beta1_pow);
+
     mom1 = beta1_ * mom1 + (1 - beta1_) * g;
     mom2 = beta2_ * mom2 + (1 - beta2_) * g * g;
     p -= lr * (mom1 / (sqrt(mom2) + epsilon_));
@@ -82,6 +90,71 @@ struct AdamFunctor {
   }
 };
 
+template <typename T>
+struct AdamFunctor<T, CPUAdam> {
+  T beta1_;
+  T beta2_;
+  T epsilon_;
+
+  const T* beta1_pow_;
+  const T* beta2_pow_;
+  const T* moment1_;
+  T* moment1_out_;
+  const T* moment2_;
+  T* moment2_out_;
+  const T* lr_;
+  const T* grad_;
+  const T* param_;
+  T* param_out_;
+
+  AdamFunctor(T beta1, T beta2, T epsilon, const T* beta1_pow,
+              const T* beta2_pow, const T* mom1, T* mom1_out, const T* mom2,
+              T* mom2_out, const T* lr, const T* grad, const T* param,
+              T* param_out)
+      : beta1_(beta1),
+        beta2_(beta2),
+        epsilon_(epsilon),
+        beta1_pow_(beta1_pow),
+        beta2_pow_(beta2_pow),
+        moment1_(mom1),
+        moment1_out_(mom1_out),
+        moment2_(mom2),
+        moment2_out_(mom2_out),
+        lr_(lr),
+        grad_(grad),
+        param_(param),
+        param_out_(param_out) {}
+
+  void operator()(size_t numel) const {
+    Eigen::Map<const Eigen::Array<T, 1, Eigen::Dynamic>> g{
+        grad_, static_cast<Eigen::Index>(numel)};
+    Eigen::Map<const Eigen::Array<T, 1, Eigen::Dynamic>> mom1{
+        moment1_, static_cast<Eigen::Index>(numel)};
+    Eigen::Map<const Eigen::Array<T, 1, Eigen::Dynamic>> mom2{
+        moment2_, static_cast<Eigen::Index>(numel)};
+    Eigen::Map<const Eigen::Array<T, 1, Eigen::Dynamic>> param{
+        param_, static_cast<Eigen::Index>(numel)};
+
+    Eigen::Map<Eigen::Array<T, 1, Eigen::Dynamic>> param_out{
+        param_out_, static_cast<Eigen::Index>(numel)};
+    Eigen::Map<Eigen::Array<T, 1, Eigen::Dynamic>> moment1_out{
+        moment1_out_, static_cast<Eigen::Index>(numel)};
+    Eigen::Map<Eigen::Array<T, 1, Eigen::Dynamic>> moment2_out{
+        moment2_out_, static_cast<Eigen::Index>(numel)};
+
+    T lr = *lr_;
+    T beta1_pow = *beta1_pow_;
+    T beta2_pow = *beta2_pow_;
+
+    // Calculation
+    lr *= sqrt(1 - beta2_pow) / (1 - beta1_pow);
+
+    moment1_out = beta1_ * mom1 + (1 - beta1_) * g;
+    moment2_out = beta2_ * mom2 + (1 - beta2_) * g * g;
+    param_out = param - lr * (moment1_out / (moment2_out.sqrt() + epsilon_));
+  }
+};
+
 template <typename T>
 struct SparseAdamFunctor {
   T beta1_;
@@ -134,6 +207,7 @@ struct SparseAdamFunctor {
       T p = param_[rows_[i] * row_numel_ + j];
 
       lr *= sqrt(1 - beta2_pow) / (1 - beta1_pow);
+
       mom1 = beta1_ * mom1 + (1 - beta1_) * g;
       mom2 = beta2_ * mom2 + (1 - beta2_) * g * g;
       p -= lr * (mom1 / (sqrt(mom2) + epsilon_));
@@ -177,19 +251,34 @@ class AdamOpKernel : public framework::OpKernel<T> {
 
     if (grad_var->IsType<framework::LoDTensor>()) {
       auto& grad = Ref(ctx.Input<LoDTensor>("Grad"), "Must set Grad");
-      AdamFunctor<T> functor(
-          beta1, beta2, epsilon, beta1_pow.template data<T>(),
-          beta2_pow.template data<T>(), mom1.template data<T>(),
-          mom1_out.template mutable_data<T>(ctx.GetPlace()),
-          mom2.template data<T>(),
-          mom2_out.template mutable_data<T>(ctx.GetPlace()),
-          lr.template data<T>(), grad.template data<T>(),
-          param.template data<T>(),
-          param_out.template mutable_data<T>(ctx.GetPlace()));
-      platform::ForRange<DeviceContext> for_range(
-          static_cast<const DeviceContext&>(ctx.device_context()),
-          param.numel());
-      for_range(functor);
+
+      if (platform::is_cpu_place(ctx.GetPlace())) {
+        AdamFunctor<T, CPUAdam> functor(
+            beta1, beta2, epsilon, beta1_pow.template data<T>(),
+            beta2_pow.template data<T>(), mom1.template data<T>(),
+            mom1_out.template mutable_data<T>(ctx.GetPlace()),
+            mom2.template data<T>(),
+            mom2_out.template mutable_data<T>(ctx.GetPlace()),
+            lr.template data<T>(), grad.template data<T>(),
+            param.template data<T>(),
+            param_out.template mutable_data<T>(ctx.GetPlace()));
+        functor(param.numel());
+      } else if (platform::is_gpu_place(ctx.GetPlace())) {
+        AdamFunctor<T, GPUAdam> functor(
+            beta1, beta2, epsilon, beta1_pow.template data<T>(),
+            beta2_pow.template data<T>(), mom1.template data<T>(),
+            mom1_out.template mutable_data<T>(ctx.GetPlace()),
+            mom2.template data<T>(),
+            mom2_out.template mutable_data<T>(ctx.GetPlace()),
+            lr.template data<T>(), grad.template data<T>(),
+            param.template data<T>(),
+            param_out.template mutable_data<T>(ctx.GetPlace()));
+
+        platform::ForRange<DeviceContext> for_range(
+            static_cast<const DeviceContext&>(ctx.device_context()),
+            param.numel());
+        for_range(functor);
+      }
     } else if (grad_var->IsType<framework::SelectedRows>()) {
       auto& grad =
           Ref(ctx.Input<framework::SelectedRows>("Grad"), "Must set Grad");

paddle/fluid/operators/concurrency/channel_util.cc

@@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
 
-#include "channel_util.h"
+#include "paddle/fluid/operators/concurrency/channel_util.h"
 #include "paddle/fluid/framework/var_type.h"
 
 namespace poc = paddle::operators::concurrency;

paddle/fluid/operators/detail/grpc_server.cc

@@ -30,9 +30,13 @@ enum CallStatus { PROCESS = 0, FINISH };
 class RequestBase {
  public:
   explicit RequestBase(GrpcService::AsyncService* service,
-                       ::grpc::ServerCompletionQueue* cq,
+                       ::grpc::ServerCompletionQueue* cq, bool sync_mode,
                        const platform::DeviceContext* dev_ctx)
-      : service_(service), cq_(cq), status_(PROCESS), dev_ctx_(dev_ctx) {
+      : service_(service),
+        cq_(cq),
+        sync_mode_(sync_mode),
+        status_(PROCESS),
+        dev_ctx_(dev_ctx) {
     PADDLE_ENFORCE(cq_);
   }
   virtual ~RequestBase() {}
@@ -49,6 +53,7 @@ class RequestBase {
   ::grpc::ServerContext ctx_;
   GrpcService::AsyncService* service_;
   ::grpc::ServerCompletionQueue* cq_;
+  const bool sync_mode_;
   CallStatus status_;
   const platform::DeviceContext* dev_ctx_;
 };
@@ -56,11 +61,17 @@ class RequestBase {
 class RequestSend final : public RequestBase {
  public:
   explicit RequestSend(GrpcService::AsyncService* service,
-                       ::grpc::ServerCompletionQueue* cq,
+                       ::grpc::ServerCompletionQueue* cq, bool sync_mode,
                        framework::Scope* scope, ReceivedQueue* queue,
                        const platform::DeviceContext* dev_ctx)
-      : RequestBase(service, cq, dev_ctx), queue_(queue), responder_(&ctx_) {
-    request_.reset(new VariableResponse(scope, dev_ctx_));
+      : RequestBase(service, cq, sync_mode, dev_ctx),
+        queue_(queue),
+        responder_(&ctx_) {
+    if (sync_mode_) {
+      request_.reset(new VariableResponse(scope, dev_ctx_, false));
+    } else {
+      request_.reset(new VariableResponse(scope, dev_ctx_, true));
+    }
     int method_id = static_cast<int>(detail::GrpcMethod::kSendVariable);
     service_->RequestAsyncUnary(method_id, &ctx_, request_.get(), &responder_,
                                 cq_, cq_, this);
@@ -87,11 +98,11 @@ class RequestSend final : public RequestBase {
 class RequestGet final : public RequestBase {
  public:
   explicit RequestGet(GrpcService::AsyncService* service,
-                      ::grpc::ServerCompletionQueue* cq,
+                      ::grpc::ServerCompletionQueue* cq, bool sync_mode,
                       framework::Scope* scope,
                       const platform::DeviceContext* dev_ctx,
                       framework::BlockingQueue<MessageWithName>* queue)
-      : RequestBase(service, cq, dev_ctx),
+      : RequestBase(service, cq, sync_mode, dev_ctx),
         responder_(&ctx_),
         scope_(scope),
         queue_(queue) {
@@ -134,19 +145,23 @@ class RequestGet final : public RequestBase {
 class RequestPrefetch final : public RequestBase {
  public:
   explicit RequestPrefetch(GrpcService::AsyncService* service,
-                           ::grpc::ServerCompletionQueue* cq,
+                           ::grpc::ServerCompletionQueue* cq, bool sync_mode,
                            framework::Scope* scope,
                            const platform::DeviceContext* dev_ctx,
                            framework::Executor* executor,
                            framework::ProgramDesc* program,
                            framework::ExecutorPrepareContext* prefetch_ctx)
-      : RequestBase(service, cq, dev_ctx),
+      : RequestBase(service, cq, sync_mode, dev_ctx),
         responder_(&ctx_),
         scope_(scope),
         executor_(executor),
         program_(program),
         prefetch_ctx_(prefetch_ctx) {
-    request_.reset(new VariableResponse(scope, dev_ctx_));
+    if (sync_mode_) {
+      request_.reset(new VariableResponse(scope, dev_ctx_, false));
+    } else {
+      request_.reset(new VariableResponse(scope, dev_ctx_, true));
+    }
     int method_id = static_cast<int>(detail::GrpcMethod::kPrefetchVariable);
     service_->RequestAsyncUnary(method_id, &ctx_, request_.get(), &responder_,
                                 cq_, cq_, this);
@@ -181,7 +196,6 @@ class RequestPrefetch final : public RequestBase {
   framework::Executor* executor_;
   framework::ProgramDesc* program_;
   framework::ExecutorPrepareContext* prefetch_ctx_;
-  int blkid_;
 };
 
 void AsyncGRPCServer::WaitClientGet(int count) {
@@ -254,8 +268,8 @@ void AsyncGRPCServer::TryToRegisterNewSendOne() {
     VLOG(3) << "shutdown, do not TryToRegisterNewSendOne";
     return;
   }
-  RequestSend* send = new RequestSend(&service_, cq_send_.get(), scope_,
-                                      &var_recv_queue_, dev_ctx_);
+  RequestSend* send = new RequestSend(&service_, cq_send_.get(), sync_mode_,
+                                      scope_, &var_recv_queue_, dev_ctx_);
   VLOG(4) << "Create RequestSend status:" << send->Status();
 }
 
@@ -265,8 +279,8 @@ void AsyncGRPCServer::TryToRegisterNewGetOne() {
     VLOG(3) << "shutdown, do not TryToRegisterNewGetOne";
     return;
   }
-  RequestGet* get = new RequestGet(&service_, cq_get_.get(), scope_, dev_ctx_,
-                                   &var_get_queue_);
+  RequestGet* get = new RequestGet(&service_, cq_get_.get(), sync_mode_, scope_,
+                                   dev_ctx_, &var_get_queue_);
   VLOG(4) << "Create RequestGet status:" << get->Status();
 }
 
@@ -277,8 +291,8 @@ void AsyncGRPCServer::TryToRegisterNewPrefetchOne() {
     return;
   }
   RequestPrefetch* prefetch =
-      new RequestPrefetch(&service_, cq_prefetch_.get(), scope_, dev_ctx_,
-                          executor_, program_, prefetch_ctx_);
+      new RequestPrefetch(&service_, cq_prefetch_.get(), sync_mode_, scope_,
+                          dev_ctx_, executor_, program_, prefetch_ctx_);
 
   VLOG(4) << "Create RequestPrefetch status:" << prefetch->Status();
 }
@@ -301,9 +315,11 @@ void AsyncGRPCServer::HandleRequest(::grpc::ServerCompletionQueue* cq,
     VLOG(3) << "HandleRequest for " << cq_name << " while after Next";
 
     PADDLE_ENFORCE(tag);
-    // FIXME(typhoonzero): de-couple the barriers with recv_op
-    if (!is_shut_down_ && cq_name == "cq_get") WaitCond(1);
-    if (!is_shut_down_ && cq_name == "cq_send") WaitCond(0);
+    if (sync_mode_) {
+      // FIXME(typhoonzero): de-couple the barriers with recv_op
+      if (!is_shut_down_ && cq_name == "cq_get") WaitCond(1);
+      if (!is_shut_down_ && cq_name == "cq_send") WaitCond(0);
+    }
 
     RequestBase* base = reinterpret_cast<RequestBase*>(tag);
     // reference:
@@ -320,13 +336,13 @@ void AsyncGRPCServer::HandleRequest(::grpc::ServerCompletionQueue* cq,
 
     switch (base->Status()) {
       case PROCESS: {
-        VLOG(4) << cq_name << " status:" << base->Status();
+        VLOG(4) << cq_name << " PROCESS status:" << base->Status();
         TryToRegisterNewOne();
         base->Process();
         break;
       }
       case FINISH: {
-        VLOG(4) << cq_name << " status:" << base->Status();
+        VLOG(4) << cq_name << " FINISH status:" << base->Status();
         delete base;
         break;
       }

paddle/fluid/operators/detail/grpc_server.h

@@ -44,7 +44,8 @@ class RequestBase;
 
 class AsyncGRPCServer final {
  public:
-  explicit AsyncGRPCServer(const std::string &address) : address_(address) {}
+  explicit AsyncGRPCServer(const std::string &address, bool sync_mode)
+      : address_(address), sync_mode_(sync_mode) {}
 
   void RunSyncUpdate();
 
@@ -95,6 +96,7 @@ class AsyncGRPCServer final {
   std::unique_ptr<::grpc::Server> server_;
 
   std::string address_;
+  const bool sync_mode_;
   framework::Scope *scope_;
   const platform::DeviceContext *dev_ctx_;
 

paddle/fluid/operators/detail/grpc_server_test.cc

@@ -89,7 +89,7 @@ void InitTensorsOnServer(framework::Scope* scope, platform::CPUPlace* place,
 }
 
 void StartServer(const std::string& endpoint) {
-  rpc_service_.reset(new detail::AsyncGRPCServer(endpoint));
+  rpc_service_.reset(new detail::AsyncGRPCServer(endpoint, true));
   framework::ProgramDesc program;
   framework::Scope scope;
   platform::CPUPlace place;

paddle/fluid/operators/detail/variable_response.h

@@ -46,7 +46,9 @@ class VariableResponse {
   }
 
   virtual ~VariableResponse() {
-    if (create_scope_) scope_->DeleteScope(local_scope_);
+    if (create_scope_) {
+      scope_->DeleteScope(local_scope_);
+    }
   }
 
   // return:
@@ -63,6 +65,8 @@ class VariableResponse {
 
   const framework::Scope& GetLocalScope() const { return *local_scope_; }
 
+  framework::Scope* GetMutableLocalScope() const { return local_scope_; }
+
   inline std::string Varname() { return meta_.varname(); }
   inline std::string OutVarname() { return meta_.out_varname(); }
 

paddle/fluid/operators/fetch_op.cc

@@ -57,10 +57,7 @@ class FetchOp : public framework::OperatorBase {
 
     // FIXME(yuyang18): Should we assume the fetch operator always generate
     // CPU outputs?
-    auto &dev_ctx = *pool.Get(src_item.place());
-
-    TensorCopy(src_item, platform::CPUPlace(), dev_ctx, &dst_item);
-    dev_ctx.Wait();
+    TensorCopySync(src_item, platform::CPUPlace(), &dst_item);
     dst_item.set_lod(src_item.lod());
 
     VLOG(3) << "Fetch variable " << fetch_var_name << " to " << out_name;

paddle/fluid/operators/gru_op.h

@@ -34,7 +34,7 @@ inline void ReorderInitState(const DeviceContext& ctx,
                              framework::Tensor* dst, bool indexed_src) {
   math::CopyMatrixRowsFunctor<DeviceContext, T> row_shuffle;
   dst->mutable_data<T>(src.dims(), ctx.GetPlace());
-  row_shuffle(ctx, src, index_lod, *dst, indexed_src);
+  row_shuffle(ctx, src, index_lod, dst, indexed_src);
 }
 
 template <typename DeviceContext, typename T>
@@ -61,7 +61,7 @@ class GRUKernel : public framework::OpKernel<T> {
     bool is_reverse = context.Attr<bool>("is_reverse");
     math::LoDTensor2BatchFunctor<DeviceContext, T> to_batch;
     auto& dev_ctx = context.template device_context<DeviceContext>();
-    to_batch(dev_ctx, *input, *batch_gate, true, is_reverse);
+    to_batch(dev_ctx, *input, batch_gate, true, is_reverse);
 
     if (bias) {
       math::RowwiseAdd<DeviceContext, T> add_bias;
@@ -113,7 +113,7 @@ class GRUKernel : public framework::OpKernel<T> {
 
     math::Batch2LoDTensorFunctor<DeviceContext, T> to_seq;
     batch_hidden->set_lod(batch_gate->lod());
-    to_seq(dev_ctx, *batch_hidden, *hidden);
+    to_seq(dev_ctx, *batch_hidden, hidden);
   }
 
   void Compute(const framework::ExecutionContext& context) const override {
@@ -174,7 +174,7 @@ class GRUGradKernel : public framework::OpKernel<T> {
 
     bool is_reverse = context.Attr<bool>("is_reverse");
     batch_hidden_grad.set_lod(batch_hidden->lod());
-    to_batch(dev_ctx, *hidden_grad, batch_hidden_grad, false, is_reverse);
+    to_batch(dev_ctx, *hidden_grad, &batch_hidden_grad, false, is_reverse);
 
     math::GRUMetaValue<T> gru_value;
     gru_value.gate_weight = const_cast<T*>(weight_data);
@@ -236,7 +236,7 @@ class GRUGradKernel : public framework::OpKernel<T> {
       input_grad->mutable_data<T>(context.GetPlace());
       math::Batch2LoDTensorFunctor<DeviceContext, T> to_seq;
       batch_gate_grad.set_lod(batch_gate->lod());
-      to_seq(dev_ctx, batch_gate_grad, *input_grad);
+      to_seq(dev_ctx, batch_gate_grad, input_grad);
     }
     if (bias_grad) {
       bias_grad->mutable_data<T>(context.GetPlace());

paddle/fluid/operators/iou_similarity_op.h

@@ -41,22 +41,24 @@ struct IOUSimilarityFunctor {
   IOUSimilarityFunctor(const T* x, const T* y, T* z, int cols)
       : x_(x), y_(y), z_(z), cols_(static_cast<size_t>(cols)) {}
 
-  inline HOSTDEVICE void operator()(size_t row_id) const {
+  inline HOSTDEVICE void operator()(size_t tid) const {
+    size_t row_id = tid / cols_;
+    size_t col_id = tid % cols_;
+
     T x_min1 = x_[row_id * 4];
     T y_min1 = x_[row_id * 4 + 1];
     T x_max1 = x_[row_id * 4 + 2];
     T y_max1 = x_[row_id * 4 + 3];
-    for (size_t i = 0; i < cols_; ++i) {
-      T x_min2 = y_[i * 4];
-      T y_min2 = y_[i * 4 + 1];
-      T x_max2 = y_[i * 4 + 2];
-      T y_max2 = y_[i * 4 + 3];
 
-      T sim = IOUSimilarity(x_min1, y_min1, x_max1, y_max1, x_min2, y_min2,
-                            x_max2, y_max2);
+    T x_min2 = y_[col_id * 4];
+    T y_min2 = y_[col_id * 4 + 1];
+    T x_max2 = y_[col_id * 4 + 2];
+    T y_max2 = y_[col_id * 4 + 3];
+
+    T sim = IOUSimilarity(x_min1, y_min1, x_max1, y_max1, x_min2, y_min2,
+                          x_max2, y_max2);
 
-      z_[row_id * cols_ + i] = sim;
-    }
+    z_[row_id * cols_ + col_id] = sim;
   }
   const T* x_;
   const T* y_;
@@ -81,7 +83,7 @@ class IOUSimilarityKernel : public framework::OpKernel<T> {
                                     out->mutable_data<T>(ctx.GetPlace()), y_n);
 
     platform::ForRange<DeviceContext> for_range(
-        static_cast<const DeviceContext&>(ctx.device_context()), x_n);
+        static_cast<const DeviceContext&>(ctx.device_context()), x_n * y_n);
     for_range(functor);
   }
 };  // namespace operators

paddle/fluid/operators/listen_and_serv_op.cc

@@ -27,6 +27,38 @@ void RunServer(std::shared_ptr<detail::AsyncGRPCServer> service) {
   VLOG(4) << "RunServer thread end";
 }
 
+static void split(const std::string &str, char sep,
+                  std::vector<std::string> *pieces) {
+  pieces->clear();
+  if (str.empty()) {
+    return;
+  }
+  size_t pos = 0;
+  size_t next = str.find(sep, pos);
+  while (next != std::string::npos) {
+    pieces->push_back(str.substr(pos, next - pos));
+    pos = next + 1;
+    next = str.find(sep, pos);
+  }
+  if (!str.substr(pos).empty()) {
+    pieces->push_back(str.substr(pos));
+  }
+}
+
+static void AsyncExecuteBlock(framework::Executor *executor,
+                              framework::ExecutorPrepareContext *prepared,
+                              framework::Scope *scope) {
+  std::future<void> future = framework::Async([&executor, &prepared, &scope]() {
+    try {
+      executor->RunPreparedContext(prepared, scope, false, false);
+    } catch (std::exception &e) {
+      LOG(ERROR) << "run sub program error " << e.what();
+    }
+  });
+  // TODO(qiao) maybe we can remove this
+  future.wait();
+}
+
 static void ParallelExecuteBlocks(
     const std::vector<size_t> &parallel_blkids, framework::Executor *executor,
     const std::vector<std::shared_ptr<framework::ExecutorPrepareContext>>
@@ -169,15 +201,82 @@ void ListenAndServOp::RunSyncLoop(framework::Executor *executor,
   }  // while(true)
 }
 
+void ListenAndServOp::RunAsyncLoop(framework::Executor *executor,
+                                   framework::ProgramDesc *program,
+                                   framework::Scope *recv_scope,
+                                   framework::BlockDesc *prefetch_block) const {
+  VLOG(3) << "RunAsyncLoop in";
+  // grad name to block id
+  std::unordered_map<std::string, int32_t> grad_to_block_id;
+  std::unordered_map<int32_t, std::string> id_to_grad;
+
+  auto grad_to_block_id_str =
+      Attr<std::vector<std::string>>("grad_to_block_id");
+  for (auto &grad_and_id : grad_to_block_id_str) {
+    std::vector<std::string> pieces;
+    split(grad_and_id, ':', &pieces);
+    VLOG(3) << "after split, grad = " << pieces[0] << ", id=" << pieces[1];
+    PADDLE_ENFORCE_EQ(pieces.size(), 2);
+    PADDLE_ENFORCE_EQ(grad_to_block_id.count(pieces[0]), 0);
+    int block_id = std::stoi(pieces[1]);
+    grad_to_block_id[pieces[0]] = block_id;
+    id_to_grad[block_id] = pieces[0];
+  }
+  size_t num_blocks = program->Size();
+  PADDLE_ENFORCE_GE(num_blocks, 2,
+                    "server program should have at least 2 blocks");
+
+  std::vector<int> block_list;
+  for (size_t blkid = 1; blkid < num_blocks; ++blkid) {
+    block_list.push_back(blkid);
+  }
+  auto optimize_prepared = executor->Prepare(*program, block_list);
+  std::unordered_map<std::string,
+                     std::shared_ptr<framework::ExecutorPrepareContext>>
+      grad_to_prepared_ctx;
+  for (size_t i = 0; i < block_list.size(); ++i) {
+    grad_to_prepared_ctx[id_to_grad[block_list[i]]] = optimize_prepared[i];
+  }
+
+  VLOG(3) << "RunAsyncLoop into while";
+  bool exit_flag = false;
+  while (!exit_flag) {
+    const detail::ReceivedMessage v = rpc_service_->Get();
+    auto recv_var_name = v.first;
+    if (recv_var_name == LISTEN_TERMINATE_MESSAGE) {
+      LOG(INFO) << "received terminate message and exit";
+      exit_flag = true;
+      break;
+    } else {
+      VLOG(3) << "received grad: " << recv_var_name;
+      auto var = v.second->GetVar();
+      if (var == nullptr) {
+        LOG(ERROR) << "Can not find server side var: " << recv_var_name;
+        PADDLE_THROW("Can not find server side var");
+      }
+      AsyncExecuteBlock(executor, grad_to_prepared_ctx[recv_var_name].get(),
+                        v.second->GetMutableLocalScope());
+    }
+
+    if (exit_flag) {
+      rpc_service_->ShutDown();
+      break;
+    }
+  }  // while(true)
+}
+
 void ListenAndServOp::RunImpl(const framework::Scope &scope,
                               const platform::Place &dev_place) const {
   platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
   auto &dev_ctx = *pool.Get(dev_place);
   framework::Scope &recv_scope = scope.NewScope();
 
+  bool sync_mode = Attr<bool>("sync_mode");
+
   PADDLE_ENFORCE(!rpc_service_);
   std::string endpoint = Attr<std::string>("endpoint");
-  rpc_service_.reset(new detail::AsyncGRPCServer(endpoint));
+
+  rpc_service_.reset(new detail::AsyncGRPCServer(endpoint, sync_mode));
 
   auto *optimize_block = Attr<framework::BlockDesc *>(kOptimizeBlock);
   auto *prefetch_block = Attr<framework::BlockDesc *>(kPrefetchBlock);
@@ -202,7 +301,11 @@ void ListenAndServOp::RunImpl(const framework::Scope &scope,
   sleep(5);
   // Write to a file of server selected port for python use.
   SavePort(rpc_service_);
-  RunSyncLoop(&executor, program, &recv_scope, prefetch_block);
+  if (sync_mode) {
+    RunSyncLoop(&executor, program, &recv_scope, prefetch_block);
+  } else {
+    RunAsyncLoop(&executor, program, &recv_scope, prefetch_block);
+  }
 }
 
 class ListenAndServOpMaker : public framework::OpProtoAndCheckerMaker {
@@ -221,6 +324,12 @@ from send_op and send back variables to recv_op.
                          "IP address to listen on.")
         .SetDefault("127.0.0.1:6164")
         .AddCustomChecker([](const std::string &ip) { return !ip.empty(); });
+    AddAttr<std::vector<std::string>>(
+        "grad_to_block_id",
+        "['param1@GRAD.block0:1', 'param2@GRAD.blockn:2'] "
+        "a map from grad name to it's optimize block id")
+        .SetDefault({});
+    AddAttr<bool>("sync_mode", "if works at sync_mode or not").SetDefault(true);
     AddAttr<framework::BlockDesc *>(kOptimizeBlock,
                                     "BlockID to run on server side.");
     AddAttr<framework::BlockDesc *>(kPrefetchBlock,

paddle/fluid/operators/listen_and_serv_op.h

@@ -46,6 +46,11 @@ class ListenAndServOp : public framework::OperatorBase {
                    framework::Scope* recv_scope,
                    framework::BlockDesc* prefetch_block) const;
 
+  void RunAsyncLoop(framework::Executor* executor,
+                    framework::ProgramDesc* program,
+                    framework::Scope* recv_scope,
+                    framework::BlockDesc* prefetch_block) const;
+
   void Stop() override;
 
   void RunImpl(const framework::Scope& scope,

paddle/fluid/operators/lstm_op.h

@@ -33,7 +33,7 @@ inline void ReorderInitState(const DeviceContext& ctx,
                              framework::Tensor* dst, bool indexed_src) {
   math::CopyMatrixRowsFunctor<DeviceContext, T> row_shuffle;
   dst->mutable_data<T>(src.dims(), ctx.GetPlace());
-  row_shuffle(ctx, src, index_lod, *dst, indexed_src);
+  row_shuffle(ctx, src, index_lod, dst, indexed_src);
 }
 
 template <typename DeviceContext, typename T>
@@ -57,7 +57,7 @@ class LSTMKernel : public framework::OpKernel<T> {
     bool is_reverse = ctx.Attr<bool>("is_reverse");
     math::LoDTensor2BatchFunctor<DeviceContext, T> to_batch;
     auto& device_ctx = ctx.template device_context<DeviceContext>();
-    to_batch(device_ctx, *input, *batch_gate, true, is_reverse);
+    to_batch(device_ctx, *input, batch_gate, true, is_reverse);
 
     auto in_dims = input->dims();
     int frame_size = static_cast<int>(in_dims[1] / 4);
@@ -161,11 +161,11 @@ class LSTMKernel : public framework::OpKernel<T> {
     math::Batch2LoDTensorFunctor<DeviceContext, T> to_seq;
     batch_hidden.set_lod(batch_gate->lod());
     // restore the output hidden in LoDTensor from the batch hidden
-    to_seq(device_ctx, batch_hidden, *hidden_out);
+    to_seq(device_ctx, batch_hidden, hidden_out);
 
     batch_cell.set_lod(batch_gate->lod());
     // restore the output cell state in LoDTensor from the batch cell
-    to_seq(device_ctx, batch_cell, *cell_out);
+    to_seq(device_ctx, batch_cell, cell_out);
   }
 };
 
@@ -257,7 +257,7 @@ class LSTMGradKernel : public framework::OpKernel<T> {
         const framework::DDim& dims, framework::LoDTensor& dst) {
       dst.mutable_data<T>(dims, ctx.GetPlace());
       dst.set_lod(batch_gate->lod());
-      to_batch(ctx, src, dst, false);
+      to_batch(ctx, src, &dst, false);
     };
 
     LoDTensor batch_hidden, batch_hidden_g, batch_cell;
@@ -351,7 +351,7 @@ class LSTMGradKernel : public framework::OpKernel<T> {
     if (in_g) {
       /* backward data */
       in_g->mutable_data<T>(ctx.GetPlace());
-      to_seq(device_ctx, batch_gate_g, *in_g);
+      to_seq(device_ctx, batch_gate_g, in_g);
     }
     if (bias && bias_g) {
       /* backward bias */

paddle/fluid/operators/lstmp_op.h

@@ -40,7 +40,7 @@ inline void ReorderInitState(const DeviceContext& ctx,
                              framework::Tensor* dst, bool indexed_src) {
   math::CopyMatrixRowsFunctor<DeviceContext, T> row_shuffle;
   dst->mutable_data<T>(src.dims(), ctx.GetPlace());
-  row_shuffle(ctx, src, index, *dst, indexed_src);
+  row_shuffle(ctx, src, index, dst, indexed_src);
 }
 
 template <typename DeviceContext, typename T>
@@ -81,7 +81,7 @@ class LSTMPKernel : public framework::OpKernel<T> {
     bool is_reverse = ctx.Attr<bool>("is_reverse");
     math::LoDTensor2BatchFunctor<DeviceContext, T> to_batch;
     auto& device_ctx = ctx.template device_context<DeviceContext>();
-    to_batch(device_ctx, *input, *batch_gate, true, is_reverse);
+    to_batch(device_ctx, *input, batch_gate, true, is_reverse);
 
     auto in_dims = input->dims();
     int frame_size = static_cast<int>(in_dims[1] / 4);
@@ -208,11 +208,11 @@ class LSTMPKernel : public framework::OpKernel<T> {
     math::Batch2LoDTensorFunctor<DeviceContext, T> to_seq;
     batch_proj.set_lod(batch_gate->lod());
     // restore the output hidden in LoDTensor from the batch hidden
-    to_seq(device_ctx, batch_proj, *proj_out);
+    to_seq(device_ctx, batch_proj, proj_out);
 
     batch_cell.set_lod(batch_gate->lod());
     // restore the output cell state in LoDTensor from the batch cell
-    to_seq(device_ctx, batch_cell, *cell_out);
+    to_seq(device_ctx, batch_cell, cell_out);
   }
 };
 
@@ -332,7 +332,7 @@ class LSTMPGradKernel : public framework::OpKernel<T> {
         const framework::DDim& dims, framework::LoDTensor& dst) {
       dst.mutable_data<T>(dims, ctx.GetPlace());
       dst.set_lod(batch_gate->lod());
-      to_batch(ctx, src, dst, false);
+      to_batch(ctx, src, &dst, false);
     };
 
     LoDTensor batch_hidden_g, batch_proj, batch_proj_g, batch_cell;
@@ -471,7 +471,7 @@ class LSTMPGradKernel : public framework::OpKernel<T> {
     if (in_g) {
       /* backward data */
       in_g->mutable_data<T>(ctx.GetPlace());
-      to_seq(device_ctx, batch_gate_g, *in_g);
+      to_seq(device_ctx, batch_gate_g, in_g);
     }
     if (bias && bias_g) {
       /* backward bias */

paddle/fluid/operators/math/concat_test.cc

@@ -17,17 +17,14 @@ limitations under the License. */
 #include <vector>
 #include "paddle/fluid/framework/tensor_util.h"
 
-using namespace paddle::framework;
-using namespace paddle::platform;
-
 template <typename DeviceContext, typename Place>
 void testConcat() {
-  Tensor input_a_cpu;
-  Tensor input_b_cpu;
-  Tensor out_cpu;
-  Tensor input_a;
-  Tensor input_b;
-  Tensor out;
+  paddle::framework::Tensor input_a_cpu;
+  paddle::framework::Tensor input_b_cpu;
+  paddle::framework::Tensor out_cpu;
+  paddle::framework::Tensor input_a;
+  paddle::framework::Tensor input_b;
+  paddle::framework::Tensor out;
 
   DeviceContext* context = new DeviceContext(Place());
   //  DeviceContext context(Place());
@@ -40,18 +37,18 @@ void testConcat() {
    *    output:
    *        out.shape: [5, 3, 4]
    */
-  auto dim_a = make_ddim({2, 3, 4});
-  auto dim_b = make_ddim({3, 3, 4});
-  auto dim_out = make_ddim({5, 3, 4});
+  auto dim_a = paddle::framework::make_ddim({2, 3, 4});
+  auto dim_b = paddle::framework::make_ddim({3, 3, 4});
+  auto dim_out = paddle::framework::make_ddim({5, 3, 4});
 
   input_a.mutable_data<int>(dim_a, Place());
   input_b.mutable_data<int>(dim_b, Place());
   out.mutable_data<int>(dim_out, Place());
 
   if (paddle::platform::is_gpu_place(Place())) {
-    input_a_cpu.mutable_data<int>(dim_a, CPUPlace());
-    input_b_cpu.mutable_data<int>(dim_b, CPUPlace());
-    out_cpu.mutable_data<int>(dim_out, CPUPlace());
+    input_a_cpu.mutable_data<int>(dim_a, paddle::platform::CPUPlace());
+    input_b_cpu.mutable_data<int>(dim_b, paddle::platform::CPUPlace());
+    out_cpu.mutable_data<int>(dim_out, paddle::platform::CPUPlace());
   }
 
   int* a_ptr;
@@ -72,11 +69,11 @@ void testConcat() {
   }
 
   if (paddle::platform::is_gpu_place(Place())) {
-    TensorCopy(input_a_cpu, Place(), *context, &input_a);
-    TensorCopy(input_b_cpu, Place(), *context, &input_b);
+    paddle::framework::TensorCopy(input_a_cpu, Place(), *context, &input_a);
+    paddle::framework::TensorCopy(input_b_cpu, Place(), *context, &input_b);
   }
 
-  std::vector<Tensor> input;
+  std::vector<paddle::framework::Tensor> input;
   input.push_back(input_a);
   input.push_back(input_b);
 
@@ -89,7 +86,8 @@ void testConcat() {
 
   int* out_ptr;
   if (paddle::platform::is_gpu_place(Place())) {
-    TensorCopy(out, CPUPlace(), *context, &out_cpu);
+    paddle::framework::TensorCopy(out, paddle::platform::CPUPlace(), *context,
+                                  &out_cpu);
     out_ptr = out_cpu.data<int>();
   } else {
     out_ptr = out.data<int>();
@@ -115,9 +113,9 @@ void testConcat() {
     *    output:
     *        out.shape: [2, 7, 4]
     */
-  dim_a = make_ddim({2, 3, 4});
-  dim_b = make_ddim({2, 4, 4});
-  dim_out = make_ddim({2, 7, 4});
+  dim_a = paddle::framework::make_ddim({2, 3, 4});
+  dim_b = paddle::framework::make_ddim({2, 4, 4});
+  dim_out = paddle::framework::make_ddim({2, 7, 4});
 
   input_a.Resize(dim_a);
   input_b.Resize(dim_b);
@@ -144,8 +142,8 @@ void testConcat() {
   }
 
   if (paddle::platform::is_gpu_place(Place())) {
-    TensorCopy(input_a_cpu, Place(), *context, &input_a);
-    TensorCopy(input_b_cpu, Place(), *context, &input_b);
+    paddle::framework::TensorCopy(input_a_cpu, Place(), *context, &input_a);
+    paddle::framework::TensorCopy(input_b_cpu, Place(), *context, &input_b);
   }
 
   input.clear();
@@ -159,7 +157,8 @@ void testConcat() {
   PADDLE_ENFORCE_EQ(input_b.dims(), dim_b);
 
   if (paddle::platform::is_gpu_place(Place())) {
-    TensorCopy(out, CPUPlace(), *context, &out_cpu);
+    paddle::framework::TensorCopy(out, paddle::platform::CPUPlace(), *context,
+                                  &out_cpu);
     out_ptr = out_cpu.data<int>();
   } else {
     out_ptr = out.data<int>();
@@ -187,9 +186,9 @@ void testConcat() {
     *    output:
     *        out.shape: [2, 3, 9]
     */
-  dim_a = make_ddim({2, 3, 4});
-  dim_b = make_ddim({2, 3, 5});
-  dim_out = make_ddim({2, 3, 9});
+  dim_a = paddle::framework::make_ddim({2, 3, 4});
+  dim_b = paddle::framework::make_ddim({2, 3, 5});
+  dim_out = paddle::framework::make_ddim({2, 3, 9});
 
   input_a.Resize(dim_a);
   input_b.Resize(dim_b);
@@ -216,8 +215,8 @@ void testConcat() {
   }
 
   if (paddle::platform::is_gpu_place(Place())) {
-    TensorCopy(input_a_cpu, Place(), *context, &input_a);
-    TensorCopy(input_b_cpu, Place(), *context, &input_b);
+    paddle::framework::TensorCopy(input_a_cpu, Place(), *context, &input_a);
+    paddle::framework::TensorCopy(input_b_cpu, Place(), *context, &input_b);
   }
 
   input.clear();
@@ -231,7 +230,8 @@ void testConcat() {
   PADDLE_ENFORCE_EQ(input_b.dims(), dim_b);
 
   if (paddle::platform::is_gpu_place(Place())) {
-    TensorCopy(out, CPUPlace(), *context, &out_cpu);
+    paddle::framework::TensorCopy(out, paddle::platform::CPUPlace(), *context,
+                                  &out_cpu);
     out_ptr = out_cpu.data<int>();
   } else {
     out_ptr = out.data<int>();
@@ -261,9 +261,9 @@ void testConcat() {
     *    output:
     *        out.shape: [2, 6, 4]
     */
-  dim_a = make_ddim({2, 3, 4});
-  dim_b = make_ddim({2, 3, 4});
-  dim_out = make_ddim({2, 6, 4});
+  dim_a = paddle::framework::make_ddim({2, 3, 4});
+  dim_b = paddle::framework::make_ddim({2, 3, 4});
+  dim_out = paddle::framework::make_ddim({2, 6, 4});
 
   input_a.Resize(dim_a);
   input_b.Resize(dim_b);
@@ -290,8 +290,8 @@ void testConcat() {
   }
 
   if (paddle::platform::is_gpu_place(Place())) {
-    TensorCopy(input_a_cpu, Place(), *context, &input_a);
-    TensorCopy(input_b_cpu, Place(), *context, &input_b);
+    paddle::framework::TensorCopy(input_a_cpu, Place(), *context, &input_a);
+    paddle::framework::TensorCopy(input_b_cpu, Place(), *context, &input_b);
   }
 
   input.clear();
@@ -305,7 +305,8 @@ void testConcat() {
   PADDLE_ENFORCE_EQ(input_b.dims(), dim_b);
 
   if (paddle::platform::is_gpu_place(Place())) {
-    TensorCopy(out, CPUPlace(), *context, &out_cpu);
+    paddle::framework::TensorCopy(out, paddle::platform::CPUPlace(), *context,
+                                  &out_cpu);
     out_ptr = out_cpu.data<int>();
   } else {
     out_ptr = out.data<int>();

paddle/fluid/operators/math/context_project.h

@@ -14,6 +14,8 @@ limitations under the License. */
 
 #pragma once
 
+#include <algorithm>
+#include <vector>
 #include "paddle/fluid/framework/lod_tensor.h"
 #include "paddle/fluid/operators/math/im2col.h"
 #include "paddle/fluid/operators/math/math_function.h"

paddle/fluid/operators/math/cross_entropy.cu

@@ -108,7 +108,9 @@ class CrossEntropyFunctor<platform::CUDADeviceContext, T> {
 
     if (softLabel) {
       const T* label_data = labels->data<T>();
-      int block = class_num > 512 ? 512 : pow(2, int(std::log2(class_num)));
+      int block = class_num > 512
+                      ? 512
+                      : pow(2, static_cast<int>(std::log2(class_num)));
 
       SoftCrossEntropyKernel<T><<<
           batch_size, block, block * sizeof(T),

paddle/fluid/operators/math/depthwise_conv.cu

@@ -12,6 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
 
+#include <vector>
 #include "paddle/fluid/operators/math/depthwise_conv.h"
 #include "paddle/fluid/platform/cuda_helper.h"
 

paddle/fluid/operators/math/depthwise_conv.h

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #pragma once
+#include <vector>
 #include "paddle/fluid/framework/tensor.h"
 #include "paddle/fluid/platform/device_context.h"
 #include "paddle/fluid/platform/hostdevice.h"

paddle/fluid/operators/math/detail/activation_functions.h

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #pragma once
 #include <math.h>
+#include <string>
 #include "paddle/fluid/platform/enforce.h"
 #include "paddle/fluid/platform/hostdevice.h"
 

paddle/fluid/operators/math/detail/lstm_gpu_kernel.h

@@ -13,13 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #pragma once
+#include <type_traits>
+
 #include "paddle/fluid/operators/math/detail/activation_functions.h"
 #include "paddle/fluid/operators/math/lstm_compute.h"
 #include "paddle/fluid/platform/cuda_helper.h"
 #include "paddle/fluid/platform/device_context.h"
 
-#include <type_traits>
-
 namespace paddle {
 namespace operators {
 namespace math {

paddle/fluid/operators/math/im2col.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/operators/math/im2col.h"
+#include <vector>
 
 namespace paddle {
 namespace operators {

paddle/fluid/operators/math/im2col.cu

@@ -12,6 +12,8 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
 
+#include <algorithm>
+#include <vector>
 #include "paddle/fluid/operators/math/im2col.h"
 #include "paddle/fluid/platform/cuda_helper.h"
 

paddle/fluid/operators/math/im2col.h

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #pragma once
 
+#include <vector>
 #include "paddle/fluid/framework/tensor.h"
 #include "paddle/fluid/framework/tensor_util.h"
 #include "paddle/fluid/platform/device_context.h"

paddle/fluid/operators/math/im2col_test.cc

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #include "paddle/fluid/operators/math/im2col.h"
 #include <gtest/gtest.h>
+#include <vector>
 
 template <typename DeviceContext, typename Place>
 void testIm2col() {
@@ -62,7 +63,7 @@ void testIm2col() {
   if (paddle::platform::is_cpu_place(*place)) {
     input = input_tmp;
   } else {
-    TensorCopy(input_tmp, *place, *context, &input);
+    TensorCopySync(input_tmp, *place, &input);
   }
   output_cfo.mutable_data<float>(
       {1, filter_size, filter_size, output_height, output_width}, *place);
@@ -87,7 +88,7 @@ void testIm2col() {
   if (paddle::platform::is_cpu_place(*place)) {
     out_cfo_ptr = output_cfo.data<float>();
   } else {
-    TensorCopy(output_cfo, paddle::platform::CPUPlace(), *context, &output_tmp);
+    TensorCopySync(output_cfo, paddle::platform::CPUPlace(), &output_tmp);
     out_cfo_ptr = output_tmp.data<float>();
   }
   for (int i = 0; i < 6; ++i) {
@@ -98,7 +99,7 @@ void testIm2col() {
   if (paddle::platform::is_cpu_place(*place)) {
     out_ocf_ptr = output_ocf.data<float>();
   } else {
-    TensorCopy(output_ocf, paddle::platform::CPUPlace(), *context, &output_tmp);
+    TensorCopySync(output_ocf, paddle::platform::CPUPlace(), &output_tmp);
     out_ocf_ptr = output_tmp.data<float>();
   }
 
@@ -119,7 +120,7 @@ void testIm2col() {
   if (paddle::platform::is_cpu_place(*place)) {
     input = input_tmp;
   } else {
-    TensorCopy(input_tmp, *place, *context, &input);
+    TensorCopySync(input_tmp, *place, &input);
   }
 
   col2im(*context, output_cfo, dilation, stride, padding, &input);
@@ -128,7 +129,7 @@ void testIm2col() {
   if (paddle::platform::is_cpu_place(*place)) {
     in_ptr = input.data<float>();
   } else {
-    TensorCopy(input, paddle::platform::CPUPlace(), *context, &input_tmp);
+    TensorCopySync(input, paddle::platform::CPUPlace(), &input_tmp);
     in_ptr = input_tmp.data<float>();
   }
   for (int i = 0; i < 6; ++i) {
@@ -140,7 +141,7 @@ void testIm2col() {
   if (paddle::platform::is_cpu_place(*place)) {
     input = input_tmp;
   } else {
-    TensorCopy(input_tmp, *place, *context, &input);
+    TensorCopySync(input_tmp, *place, &input);
   }
 
   col2im_ocf(*context, output_ocf, dilation, stride, padding, &input);
@@ -148,7 +149,7 @@ void testIm2col() {
   if (paddle::platform::is_cpu_place(*place)) {
     in_ptr = input.data<float>();
   } else {
-    TensorCopy(input, paddle::platform::CPUPlace(), *context, &input_tmp);
+    TensorCopySync(input, paddle::platform::CPUPlace(), &input_tmp);
     in_ptr = input_tmp.data<float>();
   }
   for (int i = 0; i < 6; ++i) {

paddle/fluid/operators/math/math_function.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/operators/math/math_function.h"
+#include <vector>
 #include "paddle/fluid/framework/data_type.h"
 #include "paddle/fluid/operators/math/math_function_impl.h"
 #include "paddle/fluid/platform/float16.h"
@@ -161,7 +162,8 @@ void batched_gemm<platform::CPUDeviceContext, float16>(
     const platform::CPUDeviceContext& context, const CBLAS_TRANSPOSE transA,
     const CBLAS_TRANSPOSE transB, const int M, const int N, const int K,
     const float16 alpha, const float16* A, const float16* B, const float16 beta,
-    float16* C, const int batchCount, const int strideA, const int strideB) {
+    float16* C, const int batchCount, const int64_t strideA,
+    const int64_t strideB) {
   PADDLE_THROW("float16 batched_gemm not supported on CPU");
 }
 
@@ -172,7 +174,8 @@ void batched_gemm<platform::CPUDeviceContext, float>(
     const platform::CPUDeviceContext& context, const CBLAS_TRANSPOSE transA,
     const CBLAS_TRANSPOSE transB, const int M, const int N, const int K,
     const float alpha, const float* A, const float* B, const float beta,
-    float* C, const int batchCount, const int strideA, const int strideB) {
+    float* C, const int batchCount, const int64_t strideA,
+    const int64_t strideB) {
   int lda = (transA == CblasNoTrans) ? K : M;
   int ldb = (transB == CblasNoTrans) ? N : K;
   int ldc = N;
@@ -194,7 +197,8 @@ void batched_gemm<platform::CPUDeviceContext, double>(
     const platform::CPUDeviceContext& context, const CBLAS_TRANSPOSE transA,
     const CBLAS_TRANSPOSE transB, const int M, const int N, const int K,
     const double alpha, const double* A, const double* B, const double beta,
-    double* C, const int batchCount, const int strideA, const int strideB) {
+    double* C, const int batchCount, const int64_t strideA,
+    const int64_t strideB) {
   int lda = (transA == CblasNoTrans) ? K : M;
   int ldb = (transB == CblasNoTrans) ? N : K;
   int ldc = N;
@@ -220,7 +224,8 @@ void batched_gemm<platform::CPUDeviceContext, float>(
     const platform::CPUDeviceContext& context, const CBLAS_TRANSPOSE transA,
     const CBLAS_TRANSPOSE transB, const int M, const int N, const int K,
     const float alpha, const float* A, const float* B, const float beta,
-    float* C, const int batchCount, const int strideA, const int strideB) {
+    float* C, const int batchCount, const int64_t strideA,
+    const int64_t strideB) {
   for (int k = 0; k < batchCount; ++k) {
     const float* Ak = &A[k * strideA];
     const float* Bk = &B[k * strideB];
@@ -235,7 +240,8 @@ void batched_gemm<platform::CPUDeviceContext, double>(
     const platform::CPUDeviceContext& context, const CBLAS_TRANSPOSE transA,
     const CBLAS_TRANSPOSE transB, const int M, const int N, const int K,
     const double alpha, const double* A, const double* B, const double beta,
-    double* C, const int batchCount, const int strideA, const int strideB) {
+    double* C, const int batchCount, const int64_t strideA,
+    const int64_t strideB) {
   for (int k = 0; k < batchCount; ++k) {
     const double* Ak = &A[k * strideA];
     const double* Bk = &B[k * strideB];

paddle/fluid/operators/math/math_function.cu

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #define EIGEN_USE_GPU
+#include <vector>
 #include "paddle/fluid/framework/data_type.h"
 #include "paddle/fluid/operators/math/math_function.h"
 #include "paddle/fluid/operators/math/math_function_impl.h"
@@ -267,7 +268,8 @@ void batched_gemm<platform::CUDADeviceContext, float16>(
     const platform::CUDADeviceContext& context, const CBLAS_TRANSPOSE transA,
     const CBLAS_TRANSPOSE transB, const int M, const int N, const int K,
     const float16 alpha, const float16* A, const float16* B, const float16 beta,
-    float16* C, const int batchCount, const int strideA, const int strideB) {
+    float16* C, const int batchCount, const int64_t strideA,
+    const int64_t strideB) {
 #if CUDA_VERSION >= 8000
   // Note that cublas follows fortran order, so the order is different from
   // the cblas convention.
@@ -278,7 +280,7 @@ void batched_gemm<platform::CUDADeviceContext, float16>(
       (transA == CblasNoTrans) ? CUBLAS_OP_N : CUBLAS_OP_T;
   cublasOperation_t cuTransB =
       (transB == CblasNoTrans) ? CUBLAS_OP_N : CUBLAS_OP_T;
-  const int strideC = M * N;
+  const int64_t strideC = M * N;
 
   const half h_alpha = static_cast<const half>(alpha);
   const half h_beta = static_cast<const half>(beta);
@@ -303,7 +305,8 @@ void batched_gemm<platform::CUDADeviceContext, float>(
     const platform::CUDADeviceContext& context, const CBLAS_TRANSPOSE transA,
     const CBLAS_TRANSPOSE transB, const int M, const int N, const int K,
     const float alpha, const float* A, const float* B, const float beta,
-    float* C, const int batchCount, const int strideA, const int strideB) {
+    float* C, const int batchCount, const int64_t strideA,
+    const int64_t strideB) {
 #if CUDA_VERSION >= 8000
   // Note that cublas follows fortran order, so the order is different from
   // the cblas convention.
@@ -314,7 +317,7 @@ void batched_gemm<platform::CUDADeviceContext, float>(
       (transA == CblasNoTrans) ? CUBLAS_OP_N : CUBLAS_OP_T;
   cublasOperation_t cuTransB =
       (transB == CblasNoTrans) ? CUBLAS_OP_N : CUBLAS_OP_T;
-  const int strideC = M * N;
+  const int64_t strideC = M * N;
 
   PADDLE_ENFORCE(platform::dynload::cublasSgemmStridedBatched(
       context.cublas_handle(), cuTransB, cuTransA, N, M, K, &alpha, B, ldb,
@@ -329,7 +332,8 @@ void batched_gemm<platform::CUDADeviceContext, double>(
     const platform::CUDADeviceContext& context, const CBLAS_TRANSPOSE transA,
     const CBLAS_TRANSPOSE transB, const int M, const int N, const int K,
     const double alpha, const double* A, const double* B, const double beta,
-    double* C, const int batchCount, const int strideA, const int strideB) {
+    double* C, const int batchCount, const int64_t strideA,
+    const int64_t strideB) {
 #if CUDA_VERSION >= 8000
   // Note that cublas follows fortran order, so the order is different from
   // the cblas convention.
@@ -340,7 +344,7 @@ void batched_gemm<platform::CUDADeviceContext, double>(
       (transA == CblasNoTrans) ? CUBLAS_OP_N : CUBLAS_OP_T;
   cublasOperation_t cuTransB =
       (transB == CblasNoTrans) ? CUBLAS_OP_N : CUBLAS_OP_T;
-  const int strideC = M * N;
+  const int64_t strideC = M * N;
 
   PADDLE_ENFORCE(platform::dynload::cublasDgemmStridedBatched(
       context.cublas_handle(), cuTransB, cuTransA, N, M, K, &alpha, B, ldb,

paddle/fluid/operators/math/math_function.h

@@ -26,7 +26,7 @@ limitations under the License. */
 
 #ifndef LAPACK_FOUND
 extern "C" {
-#include <cblas.h>
+#include <cblas.h>  // NOLINT
 int LAPACKE_sgetrf(int matrix_layout, int m, int n, float* a, int lda,
                    int* ipiv);
 int LAPACKE_dgetrf(int matrix_layout, int m, int n, double* a, int lda,
@@ -39,6 +39,7 @@ int LAPACKE_dgetri(int matrix_layout, int n, double* a, int lda,
 #endif
 
 #include <cmath>
+#include <vector>
 
 #include "paddle/fluid/framework/eigen.h"
 #include "paddle/fluid/framework/tensor.h"
@@ -78,8 +79,8 @@ template <typename DeviceContext, typename T>
 void batched_gemm(const DeviceContext& context, const CBLAS_TRANSPOSE transA,
                   const CBLAS_TRANSPOSE transB, const int M, const int N,
                   const int K, const T alpha, const T* A, const T* B,
-                  const T beta, T* C, const int batchCount, const int strideA,
-                  const int strideB);
+                  const T beta, T* C, const int batchCount,
+                  const int64_t strideA, const int64_t strideB);
 
 template <typename DeviceContext, typename T>
 void gemv(const DeviceContext& context, const bool trans_a, const int M,

paddle/fluid/operators/math/math_function_impl.h

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #pragma once
+#include <vector>
 #include "paddle/fluid/framework/data_type.h"
 #include "paddle/fluid/operators/math/math_function.h"
 

paddle/fluid/operators/math/math_function_test.cu

@@ -40,15 +40,15 @@ TEST(math_function, notrans_mul_trans_fp32) {
   float arr[6] = {0, 1, 2, 3, 4, 5};
   memcpy(input1_ptr, arr, 6 * sizeof(float));
 
-  TensorCopy(input1, gpu_place, context, &input1_gpu);
-  TensorCopy(input1, gpu_place, context, &input2_gpu);
+  TensorCopySync(input1, gpu_place, &input1_gpu);
+  TensorCopySync(input1, gpu_place, &input2_gpu);
 
   out_gpu.mutable_data<float>({2, 2}, gpu_place);
 
   paddle::operators::math::matmul<CUDADeviceContext, float>(
       context, input1_gpu, false, input2_gpu, true, 1, &out_gpu, 0);
 
-  TensorCopy(out_gpu, cpu_place, context, &out);
+  TensorCopySync(out_gpu, cpu_place, &out);
 
   float* out_ptr = out.data<float>();
   context.Wait();
@@ -80,8 +80,8 @@ TEST(math_function, notrans_mul_trans_fp16) {
   float16* input1_ptr = input1.mutable_data<float16>({2, 3}, cpu_place);
   fill_fp16_data(input1_ptr, input1.numel(), {0, 1, 2, 3, 4, 5});
 
-  TensorCopy(input1, gpu_place, context, &input1_gpu);
-  TensorCopy(input1, gpu_place, context, &input2_gpu);
+  TensorCopySync(input1, gpu_place, &input1_gpu);
+  TensorCopySync(input1, gpu_place, &input2_gpu);
 
   out_gpu.mutable_data<float16>({2, 2}, gpu_place);
 
@@ -89,7 +89,7 @@ TEST(math_function, notrans_mul_trans_fp16) {
       context, input1_gpu, false, input2_gpu, true, float16(1), &out_gpu,
       float16(0));
 
-  TensorCopy(out_gpu, cpu_place, context, &out);
+  TensorCopySync(out_gpu, cpu_place, &out);
 
   float16* out_ptr = out.data<float16>();
   context.Wait();
@@ -117,15 +117,15 @@ TEST(math_function, trans_mul_notrans_fp32) {
   float arr[6] = {0, 1, 2, 3, 4, 5};
   memcpy(input1_ptr, arr, 6 * sizeof(float));
 
-  TensorCopy(input1, gpu_place, context, &input1_gpu);
-  TensorCopy(input1, gpu_place, context, &input2_gpu);
+  TensorCopySync(input1, gpu_place, &input1_gpu);
+  TensorCopySync(input1, gpu_place, &input2_gpu);
 
   out_gpu.mutable_data<float>({3, 3}, gpu_place);
 
   paddle::operators::math::matmul<paddle::platform::CUDADeviceContext, float>(
       context, input1_gpu, true, input2_gpu, false, 1, &out_gpu, 0);
 
-  TensorCopy(out_gpu, cpu_place, context, &out);
+  TensorCopySync(out_gpu, cpu_place, &out);
 
   float* out_ptr = out.data<float>();
   context.Wait();
@@ -162,8 +162,8 @@ TEST(math_function, trans_mul_notrans_fp16) {
   float16* input1_ptr = input1.mutable_data<float16>({2, 3}, cpu_place);
   fill_fp16_data(input1_ptr, input1.numel(), {0, 1, 2, 3, 4, 5});
 
-  TensorCopy(input1, gpu_place, context, &input1_gpu);
-  TensorCopy(input1, gpu_place, context, &input2_gpu);
+  TensorCopySync(input1, gpu_place, &input1_gpu);
+  TensorCopySync(input1, gpu_place, &input2_gpu);
 
   out_gpu.mutable_data<float16>({3, 3}, gpu_place);
 
@@ -171,7 +171,7 @@ TEST(math_function, trans_mul_notrans_fp16) {
       context, input1_gpu, true, input2_gpu, false, float16(1), &out_gpu,
       float16(0));
 
-  TensorCopy(out_gpu, cpu_place, context, &out);
+  TensorCopySync(out_gpu, cpu_place, &out);
 
   float16* out_ptr = out.data<float16>();
   context.Wait();
@@ -214,9 +214,9 @@ TEST(math_function, gemm_notrans_cublas_fp32) {
   float arr3[8] = {0, 1, 2, 3, 4, 5, 6, 7};
   memcpy(input3_ptr, arr3, 8 * sizeof(float));
 
-  TensorCopy(input1, gpu_place, context, &input1_gpu);
-  TensorCopy(input2, gpu_place, context, &input2_gpu);
-  TensorCopy(input3, gpu_place, context, &input3_gpu);
+  TensorCopySync(input1, gpu_place, &input1_gpu);
+  TensorCopySync(input2, gpu_place, &input2_gpu);
+  TensorCopySync(input3, gpu_place, &input3_gpu);
   float* a = input1_gpu.data<float>();
   float* b = input2_gpu.data<float>();
   float* c = input3_gpu.mutable_data<float>(gpu_place);
@@ -224,7 +224,7 @@ TEST(math_function, gemm_notrans_cublas_fp32) {
   paddle::operators::math::gemm<paddle::platform::CUDADeviceContext, float>(
       context, false, false, m, n, k, 1, a, 3, b + 1, 4, 1, c + 1, 4);
 
-  TensorCopy(input3_gpu, cpu_place, context, &input3);
+  TensorCopySync(input3_gpu, cpu_place, &input3);
 
   // numpy code:
   // a = np.arange(6).reshape(2, 3)
@@ -274,9 +274,9 @@ TEST(math_function, gemm_notrans_cublas_fp16) {
   float16* input3_ptr = input3.mutable_data<float16>({2, 4}, cpu_place);
   fill_fp16_data(input3_ptr, input3.numel(), {0, 1, 2, 3, 4, 5, 6, 7});
 
-  TensorCopy(input1, gpu_place, context, &input1_gpu);
-  TensorCopy(input2, gpu_place, context, &input2_gpu);
-  TensorCopy(input3, gpu_place, context, &input3_gpu);
+  TensorCopySync(input1, gpu_place, &input1_gpu);
+  TensorCopySync(input2, gpu_place, &input2_gpu);
+  TensorCopySync(input3, gpu_place, &input3_gpu);
   float16* a = input1_gpu.data<float16>();
   float16* b = input2_gpu.data<float16>();
   float16* c = input3_gpu.mutable_data<float16>(gpu_place);
@@ -285,7 +285,7 @@ TEST(math_function, gemm_notrans_cublas_fp16) {
       context, false, false, m, n, k, float16(1), a, 3, b + 1, 4, float16(1),
       c + 1, 4);
 
-  TensorCopy(input3_gpu, cpu_place, context, &input3);
+  TensorCopySync(input3_gpu, cpu_place, &input3);
 
   // numpy code:
   // a = np.arange(6).reshape(2, 3)
@@ -332,9 +332,9 @@ TEST(math_function, gemm_trans_cublas_fp32) {
   float arr3[8] = {0, 1, 2, 3, 4, 5, 6, 7};
   memcpy(input3_ptr, arr3, 8 * sizeof(float));
 
-  TensorCopy(input1, gpu_place, context, &input1_gpu);
-  TensorCopy(input2, gpu_place, context, &input2_gpu);
-  TensorCopy(input3, gpu_place, context, &input3_gpu);
+  TensorCopySync(input1, gpu_place, &input1_gpu);
+  TensorCopySync(input2, gpu_place, &input2_gpu);
+  TensorCopySync(input3, gpu_place, &input3_gpu);
   float* a = input1_gpu.data<float>();
   float* b = input2_gpu.data<float>();
   float* c = input3_gpu.mutable_data<float>(gpu_place);
@@ -342,7 +342,7 @@ TEST(math_function, gemm_trans_cublas_fp32) {
   paddle::operators::math::gemm<paddle::platform::CUDADeviceContext, float>(
       context, false, true, m, n, k, 1, a, 3, b + 3, 3, 1, c + 1, 4);
 
-  TensorCopy(input3_gpu, cpu_place, context, &input3);
+  TensorCopySync(input3_gpu, cpu_place, &input3);
 
   context.Wait();
   EXPECT_EQ(input3_ptr[0], 0);
@@ -386,9 +386,9 @@ TEST(math_function, gemm_trans_cublas_fp16) {
   float16* input3_ptr = input3.mutable_data<float16>({2, 4}, cpu_place);
   fill_fp16_data(input3_ptr, input3.numel(), {0, 1, 2, 3, 4, 5, 6, 7});
 
-  TensorCopy(input1, gpu_place, context, &input1_gpu);
-  TensorCopy(input2, gpu_place, context, &input2_gpu);
-  TensorCopy(input3, gpu_place, context, &input3_gpu);
+  TensorCopySync(input1, gpu_place, &input1_gpu);
+  TensorCopySync(input2, gpu_place, &input2_gpu);
+  TensorCopySync(input3, gpu_place, &input3_gpu);
   float16* a = input1_gpu.data<float16>();
   float16* b = input2_gpu.data<float16>();
   float16* c = input3_gpu.mutable_data<float16>(gpu_place);
@@ -397,7 +397,7 @@ TEST(math_function, gemm_trans_cublas_fp16) {
       context, false, true, m, n, k, float16(1), a, 3, b + 3, 3, float16(1),
       c + 1, 4);
 
-  TensorCopy(input3_gpu, cpu_place, context, &input3);
+  TensorCopySync(input3_gpu, cpu_place, &input3);
 
   context.Wait();
   EXPECT_EQ(static_cast<float>(input3_ptr[0]), 0);
@@ -441,14 +441,14 @@ void GemvTest(int m, int n, bool trans) {
     data_b[i] = static_cast<T>(i);
   }
 
-  TensorCopy(mat_a, gpu_place, context, &g_mat_a);
-  TensorCopy(vec_b, gpu_place, context, &g_vec_b);
+  TensorCopySync(mat_a, gpu_place, &g_mat_a);
+  TensorCopySync(vec_b, gpu_place, &g_vec_b);
 
   paddle::operators::math::gemv<CUDADeviceContext, T>(
       context, trans, static_cast<int>(m), static_cast<int>(n), 1., g_data_a,
       g_data_b, 0., g_data_c);
 
-  TensorCopy(g_vec_c, cpu_place, context, &vec_c);
+  TensorCopySync(g_vec_c, cpu_place, &vec_c);
 
   if (!trans) {
     for (int i = 0; i < m; ++i) {

paddle/fluid/operators/math/matmul.h

@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #pragma once
+#include <algorithm>
+#include <vector>
 #include "paddle/fluid/operators/math/math_function.h"
 
 namespace paddle {

paddle/fluid/operators/math/sampler.cc

@@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
 
-#include "sampler.h"
+#include "paddle/fluid/operators/math/sampler.h"
 
 namespace paddle {
 namespace random {

paddle/fluid/operators/math/sampler.h

@@ -13,9 +13,9 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #pragma once
+#include <cstdint>
 #include <memory>
 #include <random>
-typedef long int64;
 namespace paddle {
 namespace operators {
 namespace math {
@@ -27,25 +27,25 @@ namespace math {
 */
 class Sampler {
  public:
-  explicit Sampler(int64 range) : range_(range) {
+  explicit Sampler(int64_t range) : range_(range) {
     PADDLE_ENFORCE_GT(range, 0);
     std::random_device r;
     seed_ = r();
   }
-  explicit Sampler(int64 range, unsigned int seed)
+  explicit Sampler(int64_t range, unsigned int seed)
       : range_(range), seed_(seed) {
     PADDLE_ENFORCE_GT(range, 0);
   }
   virtual ~Sampler();
   // Sample a single value
-  virtual int64 Sample() const = 0;
+  virtual int64_t Sample() const = 0;
   // The probability that a single call to Sample() returns the given value.
-  virtual float Probability(int64 value) const = 0;
+  virtual float Probability(int64_t value) const = 0;
 
-  int64 range() { return range_; };
+  int64 range() { return range_; }
 
  protected:
-  const int64 range_;
+  const int64_t range_;
   unsigned int seed_;
 };
 
@@ -56,15 +56,15 @@ class Sampler {
  */
 class UniformSampler : public Sampler {
  public:
-  explicit UniformSampler(int64 range);
+  explicit UniformSampler(int64_t range);
 
-  explicit UniformSampler(int64 range, unsigned int seed);
+  explicit UniformSampler(int64_t range, unsigned int seed);
 
   ~UniformSampler() override {}
 
   int64 Sample() const override;
 
-  float Probability(int64 value) const override;
+  float Probability(int64_t value) const override;
 
  private:
   const float inv_range_;
@@ -79,15 +79,15 @@ class UniformSampler : public Sampler {
  */
 class LogUniformSampler : public Sampler {
  public:
-  explicit LogUniformSampler(int64 range);
+  explicit LogUniformSampler(int64_t range);
 
-  explicit LogUniformSampler(int64 range, unsigned int seed);
+  explicit LogUniformSampler(int64_t range, unsigned int seed);
 
   ~LogUniformSampler() override {}
 
   int64 Sample() const override;
 
-  float Probability(int64 value) const override;
+  float Probability(int64_t value) const override;
 
  private:
   const float log_range_;
@@ -95,6 +95,6 @@ class LogUniformSampler : public Sampler {
   std::shared_ptr<std::uniform_real_distribution<>> dist_;
 };
 
-}  // math
+}  // namespace math
 }  // namespace operators
 }  // namespace paddle

paddle/fluid/operators/math/selected_rows_functor.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include <set>
+#include <vector>
 
 #include "paddle/fluid/operators/math/math_function.h"
 #include "paddle/fluid/operators/math/selected_rows_functor.h"

paddle/fluid/operators/math/selected_rows_functor.cu

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include <set>
+#include <vector>
 
 #include "paddle/fluid/operators/math/math_function.h"
 #include "paddle/fluid/operators/math/selected_rows_functor.h"

paddle/fluid/operators/math/selected_rows_functor_test.cc

@@ -13,41 +13,50 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/operators/math/selected_rows_functor.h"
+#include <vector>
 #include "gtest/gtest.h"
 #include "paddle/fluid/operators/math/math_function.h"
 
 TEST(selected_rows_functor, cpu_add) {
-  using namespace paddle::framework;
-  using namespace paddle::platform;
-  using namespace paddle::operators::math;
-
-  CPUPlace cpu_place;
-  CPUDeviceContext ctx(cpu_place);
-  SetConstant<CPUDeviceContext, float> functor;
+  paddle::platform::CPUPlace cpu_place;
+  paddle::platform::CPUDeviceContext ctx(cpu_place);
+  paddle::operators::math::SetConstant<paddle::platform::CPUDeviceContext,
+                                       float>
+      functor;
   int64_t height = 10;
   int64_t row_numel = 10;
 
   std::vector<int64_t> rows1{0, 4, 7};
-  std::unique_ptr<SelectedRows> selected_rows1{new SelectedRows(rows1, height)};
+  std::unique_ptr<paddle::framework::SelectedRows> selected_rows1{
+      new paddle::framework::SelectedRows(rows1, height)};
   auto* in1_value = selected_rows1->mutable_value();
   in1_value->mutable_data<float>(
-      make_ddim({static_cast<int64_t>(rows1.size()), row_numel}), cpu_place);
+      paddle::framework::make_ddim(
+          {static_cast<int64_t>(rows1.size()), row_numel}),
+      cpu_place);
   functor(ctx, in1_value, 1.0);
 
   std::vector<int64_t> rows2{0, 5, 7, 9};
-  std::unique_ptr<SelectedRows> selected_rows2{new SelectedRows(rows2, height)};
+  std::unique_ptr<paddle::framework::SelectedRows> selected_rows2{
+      new paddle::framework::SelectedRows(rows2, height)};
   auto* in2_value = selected_rows2->mutable_value();
   in2_value->mutable_data<float>(
-      make_ddim({static_cast<int64_t>(rows2.size()), row_numel}), cpu_place);
+      paddle::framework::make_ddim(
+          {static_cast<int64_t>(rows2.size()), row_numel}),
+      cpu_place);
   functor(ctx, in2_value, 2.0);
 
-  std::unique_ptr<SelectedRows> output{new SelectedRows()};
+  std::unique_ptr<paddle::framework::SelectedRows> output{
+      new paddle::framework::SelectedRows()};
   auto* out_value = output->mutable_value();
 
   // simplely concat two SelectedRows
-  out_value->mutable_data<float>(make_ddim({7, 10}), cpu_place);
+  out_value->mutable_data<float>(paddle::framework::make_ddim({7, 10}),
+                                 cpu_place);
 
-  SelectedRowsAdd<CPUDeviceContext, float> add_functor;
+  paddle::operators::math::SelectedRowsAdd<paddle::platform::CPUDeviceContext,
+                                           float>
+      add_functor;
   add_functor(ctx, *selected_rows1, *selected_rows2, output.get());
 
   auto out_height = output->height();
@@ -78,14 +87,20 @@ TEST(selected_rows_functor, cpu_add) {
   EXPECT_EQ(out_data[5 * row_numel + 7], 2.0);
   EXPECT_EQ(out_data[6 * row_numel + 9], 2.0);
 
-  std::unique_ptr<Tensor> tensor1{new Tensor()};
-  tensor1->mutable_data<float>(make_ddim({height, row_numel}), cpu_place);
+  std::unique_ptr<paddle::framework::Tensor> tensor1{
+      new paddle::framework::Tensor()};
+  tensor1->mutable_data<float>(
+      paddle::framework::make_ddim({height, row_numel}), cpu_place);
   functor(ctx, tensor1.get(), 3.0);
 
-  std::unique_ptr<Tensor> tensor2{new Tensor()};
-  tensor2->mutable_data<float>(make_ddim({height, row_numel}), cpu_place);
+  std::unique_ptr<paddle::framework::Tensor> tensor2{
+      new paddle::framework::Tensor()};
+  tensor2->mutable_data<float>(
+      paddle::framework::make_ddim({height, row_numel}), cpu_place);
 
-  SelectedRowsAddTensor<CPUDeviceContext, float> add_tensor_functor;
+  paddle::operators::math::SelectedRowsAddTensor<
+      paddle::platform::CPUDeviceContext, float>
+      add_tensor_functor;
   add_tensor_functor(ctx, *output, *tensor1, tensor2.get());
 
   auto* tensor2_data = tensor2->data<float>();
@@ -106,38 +121,46 @@ TEST(selected_rows_functor, cpu_add) {
 }
 
 TEST(selected_rows_functor, cpu_add_to) {
-  using namespace paddle::framework;
-  using namespace paddle::platform;
-  using namespace paddle::operators::math;
-
-  CPUPlace cpu_place;
-  CPUDeviceContext ctx(cpu_place);
-  SetConstant<CPUDeviceContext, float> functor;
+  paddle::platform::CPUPlace cpu_place;
+  paddle::platform::CPUDeviceContext ctx(cpu_place);
+  paddle::operators::math::SetConstant<paddle::platform::CPUDeviceContext,
+                                       float>
+      functor;
   int64_t height = 10;
   int64_t row_numel = 10;
 
   std::vector<int64_t> rows1{0, 4, 7};
-  std::unique_ptr<SelectedRows> selected_rows1{new SelectedRows(rows1, height)};
+  std::unique_ptr<paddle::framework::SelectedRows> selected_rows1{
+      new paddle::framework::SelectedRows(rows1, height)};
   auto* in1_value = selected_rows1->mutable_value();
   in1_value->mutable_data<float>(
-      make_ddim({static_cast<int64_t>(rows1.size()), row_numel}), cpu_place);
+      paddle::framework::make_ddim(
+          {static_cast<int64_t>(rows1.size()), row_numel}),
+      cpu_place);
   functor(ctx, in1_value, 1.0);
 
   std::vector<int64_t> rows2{0, 5, 7, 9};
-  std::unique_ptr<SelectedRows> selected_rows2{new SelectedRows(rows2, height)};
+  std::unique_ptr<paddle::framework::SelectedRows> selected_rows2{
+      new paddle::framework::SelectedRows(rows2, height)};
   auto* in2_value = selected_rows2->mutable_value();
   in2_value->mutable_data<float>(
-      make_ddim({static_cast<int64_t>(rows2.size()), row_numel}), cpu_place);
+      paddle::framework::make_ddim(
+          {static_cast<int64_t>(rows2.size()), row_numel}),
+      cpu_place);
   functor(ctx, in2_value, 2.0);
 
-  std::unique_ptr<SelectedRows> output{new SelectedRows()};
+  std::unique_ptr<paddle::framework::SelectedRows> output{
+      new paddle::framework::SelectedRows()};
   output->set_height(height);
   auto* out_value = output->mutable_value();
 
   // simplely concat two SelectedRows
-  out_value->mutable_data<float>(make_ddim({7, 10}), cpu_place);
+  out_value->mutable_data<float>(paddle::framework::make_ddim({7, 10}),
+                                 cpu_place);
 
-  SelectedRowsAddTo<CPUDeviceContext, float> add_to_functor;
+  paddle::operators::math::SelectedRowsAddTo<paddle::platform::CPUDeviceContext,
+                                             float>
+      add_to_functor;
   add_to_functor(ctx, *selected_rows1, 0, output.get());
   add_to_functor(ctx, *selected_rows2, in1_value->numel(), output.get());
 
@@ -169,11 +192,15 @@ TEST(selected_rows_functor, cpu_add_to) {
   EXPECT_EQ(out_data[5 * row_numel + 7], 2.0);
   EXPECT_EQ(out_data[6 * row_numel + 9], 2.0);
 
-  std::unique_ptr<Tensor> tensor1{new Tensor()};
-  tensor1->mutable_data<float>(make_ddim({height, row_numel}), cpu_place);
+  std::unique_ptr<paddle::framework::Tensor> tensor1{
+      new paddle::framework::Tensor()};
+  tensor1->mutable_data<float>(
+      paddle::framework::make_ddim({height, row_numel}), cpu_place);
   functor(ctx, tensor1.get(), 3.0);
 
-  SelectedRowsAddToTensor<CPUDeviceContext, float> add_to_tensor_functor;
+  paddle::operators::math::SelectedRowsAddToTensor<
+      paddle::platform::CPUDeviceContext, float>
+      add_to_tensor_functor;
   add_to_tensor_functor(ctx, *output, tensor1.get());
 
   auto* tensor1_data = tensor1->data<float>();

paddle/fluid/operators/math/sequence2batch.cc

@@ -23,11 +23,11 @@ class CopyMatrixRowsFunctor<platform::CPUDeviceContext, T> {
  public:
   void operator()(const platform::CPUDeviceContext& context,
                   const framework::Tensor& src,
-                  framework::Vector<size_t> index_lod, framework::Tensor& dst,
+                  framework::Vector<size_t> index_lod, framework::Tensor* dst,
                   bool is_src_index) {
     size_t* index = index_lod.data();
     auto src_dims = src.dims();
-    auto dst_dims = dst.dims();
+    auto dst_dims = dst->dims();
     PADDLE_ENFORCE_EQ(src_dims.size(), 2UL,
                       "The src must be matrix with rank 2.");
     PADDLE_ENFORCE_EQ(dst_dims.size(), 2UL,
@@ -37,7 +37,7 @@ class CopyMatrixRowsFunctor<platform::CPUDeviceContext, T> {
     auto height = dst_dims[0];
     auto width = dst_dims[1];
     auto* src_data = src.data<T>();
-    auto* dst_data = dst.data<T>();
+    auto* dst_data = dst->data<T>();
     for (int i = 0; i < height; ++i) {
       if (is_src_index) {
         memcpy(dst_data + i * width, src_data + index[i] * width,

paddle/fluid/operators/math/sequence2batch.cu

@@ -43,10 +43,10 @@ class CopyMatrixRowsFunctor<platform::CUDADeviceContext, T> {
  public:
   void operator()(const platform::CUDADeviceContext& context,
                   const framework::Tensor& src,
-                  framework::Vector<size_t> index_lod, framework::Tensor& dst,
+                  framework::Vector<size_t> index_lod, framework::Tensor* dst,
                   bool is_src_index) {
     auto src_dims = src.dims();
-    auto dst_dims = dst.dims();
+    auto dst_dims = dst->dims();
     PADDLE_ENFORCE_EQ(src_dims.size(), 2,
                       "The src must be matrix with rank 2.");
     PADDLE_ENFORCE_EQ(dst_dims.size(), 2,
@@ -56,7 +56,7 @@ class CopyMatrixRowsFunctor<platform::CUDADeviceContext, T> {
     auto height = dst_dims[0];
     auto width = dst_dims[1];
     auto* src_data = src.data<T>();
-    auto* dst_data = dst.data<T>();
+    auto* dst_data = dst->data<T>();
 
     dim3 threads(128, 8);
     dim3 grid(8, 1);

paddle/fluid/operators/math/sequence2batch.h

@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #pragma once
+#include <algorithm>
+#include <vector>
 #include "paddle/fluid/framework/eigen.h"
 #include "paddle/fluid/framework/lod_tensor.h"
 #include "paddle/fluid/framework/tensor.h"
@@ -35,7 +37,7 @@ class CopyMatrixRowsFunctor {
   // copy the input src to the indexed rows of output dst.
   // The indexed rows are based on the input index.
   void operator()(const DeviceContext& context, const framework::Tensor& src,
-                  framework::Vector<size_t> index_lod, framework::Tensor& dst,
+                  framework::Vector<size_t> index_lod, framework::Tensor* dst,
                   bool is_src_index);
 };
 
@@ -58,10 +60,10 @@ class LoDTensor2BatchFunctor {
  public:
   void operator()(const DeviceContext& context,
                   const framework::LoDTensor& lod_tensor,
-                  framework::LoDTensor& batch, bool is_cal_batch_lod,
+                  framework::LoDTensor* batch, bool is_cal_batch_lod,
                   bool is_reverse = false) const {
     if (!is_cal_batch_lod) {
-      auto lods = batch.lod();
+      auto lods = batch->lod();
       PADDLE_ENFORCE_GT(lods.size(), 2UL);
       PADDLE_ENFORCE_EQ(lods[1].size(),
                         static_cast<size_t>(lod_tensor.dims()[0]));
@@ -141,7 +143,7 @@ class LoDTensor2BatchFunctor {
     for (size_t i = 0; i < seq_info.size(); ++i) {
       seq_order[i] = seq_info[i].seq_idx;
     }
-    batch.set_lod(batch_lods);
+    batch->set_lod(batch_lods);
 
     CopyMatrixRowsFunctor<DeviceContext, T> to_batch;
     to_batch(context, lod_tensor, batch_lods[1], batch, true);
@@ -153,11 +155,11 @@ class Batch2LoDTensorFunctor {
  public:
   void operator()(const DeviceContext& context,
                   const framework::LoDTensor& batch,
-                  framework::LoDTensor& lod_tensor) const {
+                  framework::LoDTensor* lod_tensor) const {
     auto in_lod = batch.lod();
     PADDLE_ENFORCE_GT(in_lod.size(), 2UL);
     PADDLE_ENFORCE_EQ(in_lod[1].size(),
-                      static_cast<size_t>(lod_tensor.dims()[0]));
+                      static_cast<size_t>(lod_tensor->dims()[0]));
     CopyMatrixRowsFunctor<DeviceContext, T> to_seq;
     to_seq(context, batch, in_lod[1], lod_tensor, false);
   }

paddle/fluid/operators/math/sequence_padding_test.cc

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #include "paddle/fluid/operators/math/sequence_padding.h"
 #include <gtest/gtest.h>
+#include <vector>
 
 template <typename DeviceContext, typename Place, typename T>
 void TestSequencePadding(const paddle::framework::LoD& lod,
@@ -75,7 +76,7 @@ void TestSequencePadding(const paddle::framework::LoD& lod,
 
   delete place;
   delete context;
-};
+}
 
 TEST(Seq2BatchPadding, CPU) {
   paddle::framework::LoD lod1;

paddle/fluid/operators/math/sequence_pooling.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/operators/math/sequence_pooling.h"
+#include <string>
 #include "paddle/fluid/operators/math/math_function.h"
 
 namespace paddle {

paddle/fluid/operators/math/sequence_pooling.cu

@@ -12,6 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
 
+#include <string>
 #include "paddle/fluid/operators/math/math_function.h"
 #include "paddle/fluid/operators/math/sequence_pooling.h"
 #include "paddle/fluid/platform/cuda_helper.h"

paddle/fluid/operators/math/sequence_pooling.h

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #pragma once
+#include <string>
 #include "paddle/fluid/framework/lod_tensor.h"
 #include "paddle/fluid/framework/tensor.h"
 #include "paddle/fluid/platform/device_context.h"

paddle/fluid/operators/math/sequence_scale.cc

@@ -21,15 +21,15 @@ namespace math {
 template <typename T>
 class ScaleLoDTensorFunctor<platform::CPUDeviceContext, T> {
  public:
-  void operator()(const platform::CPUDeviceContext& context,
-                  framework::LoDTensor& seq, const T* scales) {
+  void operator()(const platform::CPUDeviceContext& context, const T* scales,
+                  framework::LoDTensor* seq) {
     const size_t level = 0;
-    auto lod = seq.lod();
+    auto lod = seq->lod();
     const size_t num_seq = lod[level].size() - 1;
-    size_t seq_width = seq.dims()[1];
+    size_t seq_width = seq->dims()[1];
     framework::LoD abs_offset_lod = framework::ToAbsOffset(lod);
 
-    T* seq_data = seq.mutable_data<T>(context.GetPlace());
+    T* seq_data = seq->mutable_data<T>(context.GetPlace());
     for (size_t i = 0; i < num_seq; ++i) {
       for (size_t j = lod[level][i] * seq_width;
            j < lod[level][i + 1] * seq_width; ++j) {

paddle/fluid/operators/math/sequence_scale.cu

@@ -35,14 +35,14 @@ __global__ void SequenceScaleKernel(T* seq, size_t* lod, const T* scales,
 template <typename T>
 class ScaleLoDTensorFunctor<platform::CUDADeviceContext, T> {
  public:
-  void operator()(const platform::CUDADeviceContext& context,
-                  framework::LoDTensor& seq, const T* scales) {
+  void operator()(const platform::CUDADeviceContext& context, const T* scales,
+                  framework::LoDTensor* seq) {
     const size_t level = 0;
-    auto lod = seq.lod();
+    auto lod = seq->lod();
     const size_t num_seq = lod[level].size() - 1;
-    const size_t seq_width = seq.numel() / seq.dims()[0];
+    const size_t seq_width = seq->numel() / seq->dims()[0];
     framework::LoD abs_offset_lod = framework::ToAbsOffset(lod);
-    T* seq_data = seq.mutable_data<T>(context.GetPlace());
+    T* seq_data = seq->mutable_data<T>(context.GetPlace());
 
     SequenceScaleKernel<T, PADDLE_CUDA_NUM_THREADS><<<
         num_seq, PADDLE_CUDA_NUM_THREADS, 0, context.stream()>>>(

paddle/fluid/operators/math/sequence_scale.h

@@ -46,8 +46,8 @@ namespace math {
 template <typename DeviceContext, typename T>
 class ScaleLoDTensorFunctor {
  public:
-  void operator()(const DeviceContext& context, framework::LoDTensor& seq,
-                  const T* scales);
+  void operator()(const DeviceContext& context, const T* scales,
+                  framework::LoDTensor* seq);
 };
 
 }  // namespace math

paddle/fluid/operators/math/vol2col.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/operators/math/vol2col.h"
+#include <vector>
 
 namespace paddle {
 namespace operators {

paddle/fluid/operators/math/vol2col.cu

@@ -12,6 +12,8 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
 
+#include <algorithm>
+#include <vector>
 #include "paddle/fluid/operators/math/vol2col.h"
 #include "paddle/fluid/platform/cuda_helper.h"
 

paddle/fluid/operators/math/vol2col.h

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #pragma once
 
+#include <vector>
 #include "paddle/fluid/framework/tensor.h"
 #include "paddle/fluid/framework/tensor_util.h"
 #include "paddle/fluid/platform/device_context.h"

paddle/fluid/operators/math/vol2col_test.cc

@@ -15,6 +15,7 @@ limitations under the License. */
 #include "paddle/fluid/operators/math/vol2col.h"
 #include <gtest/gtest.h>
 #include <iostream>
+#include <vector>
 
 template <typename DeviceContext, typename Place>
 void testVol2col() {
@@ -71,7 +72,7 @@ void testVol2col() {
   if (paddle::platform::is_cpu_place(*place)) {
     input = input_tmp;
   } else {
-    paddle::framework::TensorCopy(input_tmp, *place, *context, &input);
+    paddle::framework::TensorCopySync(input_tmp, *place, &input);
   }
   output.mutable_data<float>({1, filter_size, filter_size, filter_size,
                               output_depth, output_height, output_width},
@@ -85,7 +86,7 @@ void testVol2col() {
   if (paddle::platform::is_cpu_place(*place)) {
     out_cfo_ptr = output.data<float>();
   } else {
-    TensorCopy(output, paddle::platform::CPUPlace(), *context, &output_tmp);
+    TensorCopySync(output, paddle::platform::CPUPlace(), &output_tmp);
     out_cfo_ptr = output_tmp.data<float>();
   }
 
@@ -99,7 +100,7 @@ void testVol2col() {
   if (paddle::platform::is_cpu_place(*place)) {
     input = input_tmp;
   } else {
-    TensorCopy(input_tmp, *place, *context, &input);
+    TensorCopySync(input_tmp, *place, &input);
   }
 
   paddle::operators::math::Col2VolFunctor<DeviceContext, float> col2vol;
@@ -109,7 +110,7 @@ void testVol2col() {
   if (paddle::platform::is_cpu_place(*place)) {
     in_ptr = input.data<float>();
   } else {
-    TensorCopy(input, paddle::platform::CPUPlace(), *context, &input_tmp);
+    TensorCopySync(input, paddle::platform::CPUPlace(), &input_tmp);
     in_ptr = input_tmp.data<float>();
   }
 

paddle/fluid/operators/nccl/nccl_gpu_common.h

@@ -15,9 +15,9 @@ limitations under the License. */
 #pragma once
 
 #include <algorithm>
-#include <condition_variable>
+#include <condition_variable>  // NOLINT
 #include <memory>
-#include <mutex>
+#include <mutex>  // NOLINT
 #include <string>
 #include <unordered_map>
 #include <vector>

paddle/fluid/operators/nccl_op_test.cu.cc

@@ -228,10 +228,8 @@ TEST_F(NCCLTester, ncclReduceOp) {
   result_tensor->Resize(kDims);
   auto *ct = result_tensor->mutable_data<float>(cpu_place);
 
-  paddle::memory::Copy(
-      cpu_place, ct, p::CUDAPlace(gpu_list_[kRoot]), rt,
-      recv_tensor.numel() * sizeof(float),
-      static_cast<p::CUDADeviceContext *>(dev_ctxs_[kRoot])->stream());
+  paddle::memory::Copy(cpu_place, ct, p::CUDAPlace(gpu_list_[kRoot]), rt,
+                       recv_tensor.numel() * sizeof(float), nullptr);
 
   for (int64_t j = 0; j < f::product(kDims); ++j) {
     ASSERT_NEAR(ct[j], expected_result, 1e-5);

paddle/fluid/operators/reader/CMakeLists.txt

@@ -23,5 +23,7 @@ reader_library(create_recordio_file_reader_op SRCS create_recordio_file_reader_o
 reader_library(create_double_buffer_reader_op SRCS create_double_buffer_reader_op.cc)
 reader_library(create_multi_pass_reader_op SRCS create_multi_pass_reader_op.cc)
 reader_library(create_threaded_reader_op SRCS create_threaded_reader_op.cc)
+
+cc_test(reader_blocking_queue_test SRCS reader_blocking_queue_test.cc)
 # Export local libraries to parent
 set(READER_LIBRARY ${LOCAL_READER_LIBS} PARENT_SCOPE)

paddle/fluid/operators/reader/blocking_queue.h

+//   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <condition_variable>  // NOLINT
+#include <deque>
+
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace operators {
+namespace reader {
+
+template <typename T>
+class BlockingQueue {
+  // BlockingQueue is for buffered reading and is supposed to use only the
+  // reader package. It is true that we could and we should have been using
+  // framework::Channel, but which has currently a deadlock bug. BlockingQueue
+  // is a workaround and a simplified version of framework::Channel as it
+  // doesn't support GPU and it implements on buffered blocking queue.
+ public:
+  explicit BlockingQueue(size_t capacity)
+      : capacity_(capacity), closed_(false) {
+    PADDLE_ENFORCE_GT(
+        capacity_, 0,
+        "The capacity of a reader::BlockingQueue must be greater than 0.");
+  }
+
+  bool Send(const T& elem) {
+    std::unique_lock<std::mutex> lock(mutex_);
+    send_cv_.wait(lock, [&] { return queue_.size() < capacity_ || closed_; });
+    if (closed_) {
+      VLOG(5)
+          << "WARNING: Sending an element to a closed reader::BlokcingQueue.";
+      return false;
+    }
+    PADDLE_ENFORCE_LT(queue_.size(), capacity_);
+    queue_.push_back(elem);
+    receive_cv_.notify_one();
+    return true;
+  }
+
+  bool Send(T&& elem) {
+    std::unique_lock<std::mutex> lock(mutex_);
+    send_cv_.wait(lock, [&] { return queue_.size() < capacity_ || closed_; });
+    if (closed_) {
+      VLOG(5)
+          << "WARNING: Sending an element to a closed reader::BlokcingQueue.";
+      return false;
+    }
+    PADDLE_ENFORCE_LT(queue_.size(), capacity_);
+    queue_.emplace_back(std::move(elem));
+    receive_cv_.notify_one();
+    return true;
+  }
+
+  bool Receive(T* elem) {
+    std::unique_lock<std::mutex> lock(mutex_);
+    receive_cv_.wait(lock, [&] { return !queue_.empty() || closed_; });
+    if (!queue_.empty()) {
+      PADDLE_ENFORCE_NOT_NULL(elem);
+      *elem = queue_.front();
+      queue_.pop_front();
+      send_cv_.notify_one();
+      return true;
+    } else {
+      PADDLE_ENFORCE(closed_);
+      return false;
+    }
+  }
+
+  void Close() {
+    std::lock_guard<std::mutex> lock(mutex_);
+    closed_ = true;
+    send_cv_.notify_all();
+    receive_cv_.notify_all();
+  }
+
+  bool IsClosed() {
+    std::lock_guard<std::mutex> lock(mutex_);
+    return closed_;
+  }
+
+  size_t Cap() {
+    std::lock_guard<std::mutex> lock(mutex_);
+    return capacity_;
+  }
+
+ private:
+  size_t capacity_;
+  bool closed_;
+  std::deque<T> queue_;
+
+  std::mutex mutex_;
+  std::condition_variable receive_cv_;
+  std::condition_variable send_cv_;
+};
+}  // namespace reader
+}  // namespace operators
+}  // namespace paddle

paddle/fluid/operators/reader/create_double_buffer_reader_op.cc

@@ -14,7 +14,7 @@
 
 #include <thread>  // NOLINT
 
-#include "paddle/fluid/framework/channel.h"
+#include "paddle/fluid/operators/reader/blocking_queue.h"
 #include "paddle/fluid/operators/reader/reader_op_registry.h"
 
 namespace paddle {
@@ -23,13 +23,13 @@ namespace reader {
 
 // 'Double buffer' means we shall maintain two batches of input data at the same
 // time. So the kCacheSize shoul be at least 2.
-static constexpr size_t kCacheSize = 2;
+static constexpr size_t kCacheSize = 3;
 // There will be two bacthes out of the channel during training:
 // 1. the one waiting to be sent to the channel
 // 2. the one just be received from the channel, which is also being used by
 // subsequent operators.
 // So the channel size should be kChacheSize - 2
-static constexpr size_t kChannelSize = 0;  // kCacheSize - 2
+static constexpr size_t kChannelSize = 1;  // kCacheSize - 2
 
 class DoubleBufferReader : public framework::DecoratedReader {
  public:
@@ -55,10 +55,8 @@ class DoubleBufferReader : public framework::DecoratedReader {
   ~DoubleBufferReader() { EndPrefetcher(); }
 
  private:
-  bool HasNext() const;
-
   void StartPrefetcher() {
-    channel_ = framework::MakeChannel<size_t>(kChannelSize);
+    channel_ = new reader::BlockingQueue<size_t>(kChannelSize);
     prefetcher_ = std::thread([this] { PrefetchThreadFunc(); });
   }
 
@@ -74,7 +72,7 @@ class DoubleBufferReader : public framework::DecoratedReader {
   void PrefetchThreadFunc();
 
   std::thread prefetcher_;
-  framework::Channel<size_t>* channel_;
+  reader::BlockingQueue<size_t>* channel_;
   platform::Place place_;
   std::vector<std::vector<framework::LoDTensor>> cpu_tensor_cache_;
   std::vector<std::vector<framework::LoDTensor>> gpu_tensor_cache_;
@@ -139,17 +137,16 @@ class CreateDoubleBufferReaderOpMaker : public DecoratedReaderMakerBase {
 };
 
 void DoubleBufferReader::ReadNext(std::vector<framework::LoDTensor>* out) {
-  out->clear();
-  if (HasNext()) {
-    size_t cached_tensor_id;
-    channel_->Receive(&cached_tensor_id);
+  size_t cached_tensor_id;
+  if (channel_->Receive(&cached_tensor_id)) {
     if (platform::is_gpu_place(place_)) {
       *out = gpu_tensor_cache_[cached_tensor_id];
-      ctxs_[cached_tensor_id]->Wait();
     } else {
       // CPU place
       *out = cpu_tensor_cache_[cached_tensor_id];
     }
+  } else {
+    out->clear();
   }
 }
 
@@ -159,12 +156,6 @@ void DoubleBufferReader::ReInit() {
   StartPrefetcher();
 }
 
-bool DoubleBufferReader::HasNext() const {
-  while (!channel_->IsClosed() && !channel_->CanReceive()) {
-  }
-  return channel_->CanReceive();
-}
-
 void DoubleBufferReader::PrefetchThreadFunc() {
   VLOG(5) << "A new prefetch thread starts.";
   size_t cached_tensor_id = 0;
@@ -177,18 +168,14 @@ void DoubleBufferReader::PrefetchThreadFunc() {
     }
     if (platform::is_gpu_place(place_)) {
       auto& gpu_batch = gpu_tensor_cache_[cached_tensor_id];
-      auto* gpu_ctx = ctxs_[cached_tensor_id].get();
       gpu_batch.resize(cpu_batch.size());
       for (size_t i = 0; i < cpu_batch.size(); ++i) {
-        framework::TensorCopy(cpu_batch[i], place_, *gpu_ctx, &gpu_batch[i],
-                              true);
+        // TODO(fengjiayi): Use asynchronous TensorCopy instead
+        framework::TensorCopySync(cpu_batch[i], place_, &gpu_batch[i]);
         gpu_batch[i].set_lod(cpu_batch[i].lod());
       }
     }
-    try {
-      size_t tmp = cached_tensor_id;
-      channel_->Send(&tmp);
-    } catch (paddle::platform::EnforceNotMet e) {
+    if (!channel_->Send(cached_tensor_id)) {
       VLOG(5) << "WARNING: The double buffer channel has been closed. The "
                  "prefetch thread will terminate.";
       break;

paddle/fluid/operators/reader/open_files_op.cc

@@ -14,7 +14,7 @@
 
 #include <thread>  // NOLINT
 
-#include "paddle/fluid/framework/channel.h"
+#include "paddle/fluid/operators/reader/blocking_queue.h"
 #include "paddle/fluid/operators/reader/reader_op_registry.h"
 
 namespace paddle {
@@ -37,7 +37,6 @@ class MultiFileReader : public framework::ReaderBase {
   ~MultiFileReader() { EndScheduler(); }
 
  private:
-  bool HasNext();
   void StartNewScheduler();
   void EndScheduler();
   void ScheduleThreadFunc();
@@ -48,15 +47,14 @@ class MultiFileReader : public framework::ReaderBase {
   std::thread scheduler_;
   std::vector<std::thread> prefetchers_;
   size_t buffer_size_;
-  framework::Channel<size_t>* waiting_file_idx_;
-  framework::Channel<size_t>* available_thread_idx_;
-  framework::Channel<std::vector<framework::LoDTensor>>* buffer_;
+  reader::BlockingQueue<size_t>* waiting_file_idx_;
+  reader::BlockingQueue<size_t>* available_thread_idx_;
+  reader::BlockingQueue<std::vector<framework::LoDTensor>>* buffer_;
 };
 
 void MultiFileReader::ReadNext(std::vector<framework::LoDTensor>* out) {
-  out->clear();
-  if (HasNext()) {
-    buffer_->Receive(out);
+  if (!buffer_->Receive(out)) {
+    out->clear();
   }
 }
 
@@ -65,25 +63,19 @@ void MultiFileReader::ReInit() {
   StartNewScheduler();
 }
 
-bool MultiFileReader::HasNext() {
-  while (!buffer_->IsClosed() && !buffer_->CanReceive()) {
-  }
-  return buffer_->CanReceive();
-}
-
 void MultiFileReader::StartNewScheduler() {
   size_t thread_num = prefetchers_.size();
-  waiting_file_idx_ = framework::MakeChannel<size_t>(file_names_.size());
-  available_thread_idx_ = framework::MakeChannel<size_t>(thread_num);
-  buffer_ =
-      framework::MakeChannel<std::vector<framework::LoDTensor>>(buffer_size_);
+  waiting_file_idx_ = new reader::BlockingQueue<size_t>(file_names_.size());
+  available_thread_idx_ = new reader::BlockingQueue<size_t>(thread_num);
+  buffer_ = new reader::BlockingQueue<std::vector<framework::LoDTensor>>(
+      buffer_size_);
 
   for (size_t i = 0; i < file_names_.size(); ++i) {
-    waiting_file_idx_->Send(&i);
+    waiting_file_idx_->Send(i);
   }
   waiting_file_idx_->Close();
   for (size_t i = 0; i < thread_num; ++i) {
-    available_thread_idx_->Send(&i);
+    available_thread_idx_->Send(i);
   }
 
   scheduler_ = std::thread([this] { ScheduleThreadFunc(); });
@@ -149,7 +141,7 @@ void MultiFileReader::PrefetchThreadFunc(std::string file_name,
       break;
     }
     try {
-      buffer_->Send(&ins);
+      buffer_->Send(std::move(ins));
     } catch (paddle::platform::EnforceNotMet e) {
       VLOG(5) << "WARNING: The buffer channel has been closed. The prefetch "
                  "thread of file '"
@@ -158,9 +150,7 @@ void MultiFileReader::PrefetchThreadFunc(std::string file_name,
     }
   }
 
-  try {
-    available_thread_idx_->Send(&thread_idx);
-  } catch (paddle::platform::EnforceNotMet e) {
+  if (!available_thread_idx_->Send(thread_idx)) {
     VLOG(5) << "WARNING: The available_thread_idx_ channel has been closed. "
                "Fail to send thread_idx.";
   }

paddle/fluid/operators/reader/reader_blocking_queue_test.cc

+//   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <chrono>  // NOLINT
+#include <set>
+#include <thread>  // NOLINT
+#include <vector>
+#include "gtest/gtest.h"
+
+#include "paddle/fluid/operators/reader/blocking_queue.h"
+
+using paddle::operators::reader::BlockingQueue;
+
+TEST(BlockingQueue, CapacityTest) {
+  size_t cap = 10;
+  BlockingQueue<int> q(cap);
+  EXPECT_EQ(q.Cap(), cap);
+}
+
+void FirstInFirstOut(size_t queue_cap, size_t elem_num, size_t send_time_gap,
+                     size_t receive_time_gap) {
+  BlockingQueue<size_t> q(queue_cap);
+  std::thread sender([&]() {
+    for (size_t i = 0; i < elem_num; ++i) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(send_time_gap));
+      EXPECT_TRUE(q.Send(i));
+    }
+    q.Close();
+  });
+  size_t count = 0;
+  while (true) {
+    std::this_thread::sleep_for(std::chrono::milliseconds(receive_time_gap));
+    size_t elem;
+    if (!q.Receive(&elem)) {
+      break;
+    }
+    EXPECT_EQ(elem, count++);
+  }
+  sender.join();
+  EXPECT_EQ(count, elem_num);
+  EXPECT_TRUE(q.IsClosed());
+}
+
+TEST(BlockingQueue, FirstInFirstOutTest) {
+  FirstInFirstOut(2, 5, 2, 50);
+  FirstInFirstOut(2, 5, 50, 2);
+  FirstInFirstOut(10, 3, 50, 2);
+  FirstInFirstOut(10, 3, 2, 50);
+}
+
+TEST(BlockingQueue, SenderBlockingTest) {
+  const size_t queue_cap = 2;
+  BlockingQueue<size_t> q(queue_cap);
+  size_t send_count = 0;
+  std::thread sender([&]() {
+    for (size_t i = 0; i < 5; ++i) {
+      if (!q.Send(i)) {
+        break;
+      }
+      ++send_count;
+    }
+  });
+  std::this_thread::sleep_for(std::chrono::milliseconds(200));
+  q.Close();
+  sender.join();
+  EXPECT_EQ(send_count, queue_cap);
+  std::vector<size_t> res;
+  while (true) {
+    size_t elem;
+    if (!q.Receive(&elem)) {
+      break;
+    }
+    res.push_back(elem);
+  }
+  EXPECT_EQ(res.size(), queue_cap);
+  for (size_t i = 0; i < res.size(); ++i) {
+    EXPECT_EQ(res[i], i);
+  }
+}
+
+TEST(BlockingQueue, ReceiverBlockingTest) {
+  const size_t queue_cap = 5;
+  BlockingQueue<size_t> q(queue_cap);
+  std::vector<size_t> receive_res;
+  std::thread receiver([&]() {
+    size_t elem;
+    while (true) {
+      if (!q.Receive(&elem)) {
+        break;
+      }
+      receive_res.push_back(elem);
+    }
+  });
+  std::vector<size_t> to_send{2, 1, 7};
+  for (auto e : to_send) {
+    q.Send(e);
+  }
+  q.Close();
+  receiver.join();
+  EXPECT_EQ(receive_res.size(), to_send.size());
+  for (size_t i = 0; i < to_send.size(); ++i) {
+    EXPECT_EQ(receive_res[i], to_send[i]);
+  }
+}
+
+void CheckIsUnorderedSame(const std::vector<std::vector<size_t>>& v1,
+                          const std::vector<std::vector<size_t>>& v2) {
+  std::set<size_t> s1;
+  std::set<size_t> s2;
+  for (auto vec : v1) {
+    for (size_t elem : vec) {
+      s1.insert(elem);
+    }
+  }
+  for (auto vec : v2) {
+    for (size_t elem : vec) {
+      s2.insert(elem);
+    }
+  }
+  EXPECT_EQ(s1.size(), s2.size());
+  auto it1 = s1.begin();
+  auto it2 = s2.begin();
+  while (it1 != s1.end()) {
+    EXPECT_EQ(*it1, *it2);
+    ++it1;
+    ++it2;
+  }
+}
+
+void MultiSenderMultiReceiver(const size_t queue_cap,
+                              const std::vector<std::vector<size_t>>& to_send,
+                              size_t receiver_num, size_t send_time_gap,
+                              size_t receive_time_gap) {
+  BlockingQueue<size_t> q(queue_cap);
+  size_t sender_num = to_send.size();
+  std::vector<std::thread> senders;
+  for (size_t s_idx = 0; s_idx < sender_num; ++s_idx) {
+    senders.emplace_back(std::thread([&, s_idx] {
+      for (size_t elem : to_send[s_idx]) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(send_time_gap));
+        EXPECT_TRUE(q.Send(elem));
+      }
+    }));
+  }
+  std::vector<std::thread> receivers;
+  std::mutex mu;
+  std::vector<std::vector<size_t>> res;
+  for (size_t r_idx = 0; r_idx < receiver_num; ++r_idx) {
+    receivers.emplace_back(std::thread([&] {
+      std::vector<size_t> receiver_res;
+      while (true) {
+        std::this_thread::sleep_for(
+            std::chrono::milliseconds(receive_time_gap));
+        size_t elem;
+        if (!q.Receive(&elem)) {
+          break;
+        }
+        receiver_res.push_back(elem);
+      }
+      std::lock_guard<std::mutex> lock(mu);
+      res.push_back(receiver_res);
+    }));
+  }
+  for (auto& t : senders) {
+    t.join();
+  }
+  q.Close();
+  for (auto& t : receivers) {
+    t.join();
+  }
+  CheckIsUnorderedSame(to_send, res);
+}
+
+TEST(BlockingQueue, MultiSenderMultiReaderTest) {
+  std::vector<std::vector<size_t>> to_send_1{{2, 3, 4}, {9}, {0, 7, 15, 6}};
+  MultiSenderMultiReceiver(2, to_send_1, 2, 0, 0);
+  MultiSenderMultiReceiver(10, to_send_1, 2, 0, 0);
+  MultiSenderMultiReceiver(2, to_send_1, 20, 0, 0);
+  MultiSenderMultiReceiver(2, to_send_1, 2, 50, 0);
+  MultiSenderMultiReceiver(2, to_send_1, 2, 0, 50);
+
+  std::vector<std::vector<size_t>> to_send_2{
+      {2, 3, 4}, {}, {0, 7, 15, 6, 9, 32}};
+  MultiSenderMultiReceiver(2, to_send_2, 3, 0, 0);
+  MultiSenderMultiReceiver(20, to_send_2, 3, 0, 0);
+  MultiSenderMultiReceiver(2, to_send_2, 30, 0, 0);
+  MultiSenderMultiReceiver(2, to_send_2, 3, 50, 0);
+  MultiSenderMultiReceiver(2, to_send_2, 3, 0, 50);
+}
+
+struct MyClass {
+  MyClass() : val_(0) {}
+  explicit MyClass(int val) : val_(val) {}
+  MyClass(const MyClass& b) { val_ = b.val_; }
+  MyClass(MyClass&& b) { val_ = b.val_; }
+  void operator=(const MyClass& b) { val_ = b.val_; }
+
+  int val_;
+};
+
+TEST(BlockingQueue, MyClassTest) {
+  BlockingQueue<MyClass> q(2);
+  MyClass a(200);
+  q.Send(std::move(a));
+  MyClass b;
+  q.Receive(&b);
+  EXPECT_EQ(a.val_, b.val_);
+}

paddle/fluid/operators/reader/reader_op_registry.cc

@@ -12,7 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "reader_op_registry.h"
+#include "paddle/fluid/operators/reader/reader_op_registry.h"
+#include <string>
+#include <vector>
 
 namespace paddle {
 namespace operators {

paddle/fluid/operators/reader/reader_op_registry.h

@@ -14,6 +14,8 @@
 
 #pragma once
 
+#include <string>
+#include <vector>
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/framework/reader.h"
 

paddle/fluid/operators/reshape_op.h

@@ -93,8 +93,14 @@ class ReshapeOp : public framework::OperatorWithKernel {
 
     if (unk_dim_idx != -1) {
       output_shape[unk_dim_idx] = -in_size / capacity;
-      PADDLE_ENFORCE_EQ(output_shape[unk_dim_idx] * capacity, -in_size,
-                        "Invalid shape is given.");
+      // in_size < 0 and is un-determinate in compile time, skip the check,
+      // for example, in_dims = [-1, 8, 1, 1], shape = [-1, 3, 8],
+      // capacity = -24, in_size = -8, output_shape[0] = 0
+      // the following check will fail.
+      if (in_size > 0) {
+        PADDLE_ENFORCE_EQ(output_shape[unk_dim_idx] * capacity, -in_size,
+                          "Invalid shape is given.");
+      }
     } else {
       PADDLE_ENFORCE_EQ(capacity, in_size, "Invalid shape is given.");
     }
@@ -124,10 +130,8 @@ class ReshapeKernel : public framework::OpKernel<T> {
       auto *shape_data = shape_tensor->data<int>();
       framework::Tensor cpu_shape_tensor;
       if (platform::is_gpu_place(ctx.GetPlace())) {
-        TensorCopy(*shape_tensor, platform::CPUPlace(), ctx.device_context(),
-                   &cpu_shape_tensor);
+        TensorCopySync(*shape_tensor, platform::CPUPlace(), &cpu_shape_tensor);
         shape_data = cpu_shape_tensor.data<int>();
-        ctx.device_context().Wait();
       }
       auto shape =
           std::vector<int>(shape_data, shape_data + shape_tensor->numel());
@@ -146,9 +150,7 @@ class ReshapeKernel : public framework::OpKernel<T> {
     out->Resize(out_dims);
     if (!inplace) {
       out->mutable_data<T>(ctx.GetPlace());
-      framework::TensorCopy(*in, ctx.GetPlace(), ctx.device_context(), out);
-      ctx.device_context().Wait();
-      // TensorCopy will resize to in_dims.
+      framework::TensorCopySync(*in, ctx.GetPlace(), out);
       out->Resize(out_dims);
     } else {
       out->ShareDataWith(*in);

paddle/fluid/operators/roi_pool_op.cc

@@ -18,8 +18,7 @@ namespace paddle {
 namespace operators {
 
 using Tensor = framework::Tensor;
-
-static constexpr int kROISize = 5;
+using LoDTensor = framework::LoDTensor;
 
 class ROIPoolOp : public framework::OperatorWithKernel {
  public:
@@ -40,11 +39,11 @@ class ROIPoolOp : public framework::OperatorWithKernel {
     PADDLE_ENFORCE(input_dims.size() == 4,
                    "The format of input tensor is NCHW.");
     PADDLE_ENFORCE(rois_dims.size() == 2,
-                   "ROIs should be a 2-D tensor of shape (num_rois, 5)"
-                   "given as [[batch_id, x1, y1, x2, y2], …].");
+                   "ROIs should be a 2-D LoDTensor of shape (num_rois, 4)"
+                   "given as [[x1, y1, x2, y2], …].");
     PADDLE_ENFORCE(rois_dims[1] == kROISize,
-                   "ROIs should be a 2-D tensor of shape (num_rois, 5)"
-                   "given as [[batch_id, x1, y1, x2, y2], …].");
+                   "ROIs should be a 2-D LoDTensor of shape (num_rois, 4)"
+                   "given as [[x1, y1, x2, y2], …].");
 
     int pooled_height = ctx->Attrs().Get<int>("pooled_height");
     int pooled_width = ctx->Attrs().Get<int>("pooled_width");
@@ -109,10 +108,10 @@ class ROIPoolOpMaker : public framework::OpProtoAndCheckerMaker {
              "H is the height of the feature, and "
              "W is the width of the feature.");
     AddInput("ROIs",
-             "(Tensor), "
+             "(LoDTensor), "
              "ROIs (Regions of Interest) to pool over. "
-             "should be a 2-D tensor of shape (num_rois, 5)"
-             "given as [[batch_id, x1, y1, x2, y2], …]. "
+             "should be a 2-D LoDTensor of shape (num_rois, 4)"
+             "given as [[x1, y1, x2, y2], …]. "
              "Where batch_id is the id of the data, "
              "(x1, y1) is the top left coordinates, and "
              "(x2, y2) is the bottom right coordinates.");

paddle/fluid/operators/roi_pool_op.cu

@@ -19,10 +19,10 @@ namespace paddle {
 namespace operators {
 
 using Tensor = framework::Tensor;
+using LoDTensor = framework::LoDTensor;
 
 static constexpr int kNumCUDAThreads = 512;
 static constexpr int kNumMaxinumNumBlocks = 4096;
-static constexpr int kROISize = 5;
 
 static inline int NumBlocks(const int N) {
   return std::min((N + kNumCUDAThreads - 1) / kNumCUDAThreads,
@@ -30,13 +30,11 @@ static inline int NumBlocks(const int N) {
 }
 
 template <typename T>
-__global__ void GPUROIPoolForward(const int nthreads, const T* input_data,
-                                  const int64_t* input_rois,
-                                  const float spatial_scale, const int channels,
-                                  const int height, const int width,
-                                  const int pooled_height,
-                                  const int pooled_width, T* output_data,
-                                  int64_t* argmax_data) {
+__global__ void GPUROIPoolForward(
+    const int nthreads, const T* input_data, const int64_t* input_rois,
+    const float spatial_scale, const int channels, const int height,
+    const int width, const int pooled_height, const int pooled_width,
+    int* roi_batch_id_data, T* output_data, int64_t* argmax_data) {
   int index = blockIdx.x * blockDim.x + threadIdx.x;
   int offset = blockDim.x * gridDim.x;
   for (size_t i = index; i < nthreads; i += offset) {
@@ -46,11 +44,11 @@ __global__ void GPUROIPoolForward(const int nthreads, const T* input_data,
     int n = index / pooled_width / pooled_height / channels;
 
     const int64_t* offset_input_rois = input_rois + n * kROISize;
-    int roi_batch_ind = offset_input_rois[0];
-    int roi_start_w = round(offset_input_rois[1] * spatial_scale);
-    int roi_start_h = round(offset_input_rois[2] * spatial_scale);
-    int roi_end_w = round(offset_input_rois[3] * spatial_scale);
-    int roi_end_h = round(offset_input_rois[4] * spatial_scale);
+    int roi_batch_ind = roi_batch_id_data[n];
+    int roi_start_w = round(offset_input_rois[0] * spatial_scale);
+    int roi_start_h = round(offset_input_rois[1] * spatial_scale);
+    int roi_end_w = round(offset_input_rois[2] * spatial_scale);
+    int roi_end_h = round(offset_input_rois[3] * spatial_scale);
 
     int roi_width = max(roi_end_w - roi_start_w + 1, 1);
     int roi_height = max(roi_end_h - roi_start_h + 1, 1);
@@ -93,7 +91,8 @@ __global__ void GPUROIPoolBackward(
     const int nthreads, const int64_t* input_rois, const T* output_grad,
     const int64_t* argmax_data, const int num_rois, const float spatial_scale,
     const int channels, const int height, const int width,
-    const int pooled_height, const int pooled_width, T* input_grad) {
+    const int pooled_height, const int pooled_width, int* roi_batch_id_data,
+    T* input_grad) {
   int index = blockIdx.x * blockDim.x + threadIdx.x;
   int offset = blockDim.x * gridDim.x;
   for (int i = index; i < nthreads; i += offset) {
@@ -102,8 +101,7 @@ __global__ void GPUROIPoolBackward(
     int c = (index / pooled_width / pooled_height) % channels;
     int n = index / pooled_width / pooled_height / channels;
 
-    const int64_t* offset_input_rois = input_rois + n * kROISize;
-    int roi_batch_ind = offset_input_rois[0];
+    int roi_batch_ind = roi_batch_id_data[n];
     int input_offset = (roi_batch_ind * channels + c) * height * width;
     int output_offset = (n * channels + c) * pooled_height * pooled_width;
     const T* offset_output_grad = output_grad + output_offset;
@@ -124,7 +122,7 @@ class GPUROIPoolOpKernel : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext& ctx) const override {
     auto* in = ctx.Input<Tensor>("X");
-    auto* rois = ctx.Input<Tensor>("ROIs");
+    auto* rois = ctx.Input<LoDTensor>("ROIs");
     auto* out = ctx.Output<Tensor>("Out");
     auto* argmax = ctx.Output<Tensor>("Argmax");
 
@@ -133,23 +131,46 @@ class GPUROIPoolOpKernel : public framework::OpKernel<T> {
     auto spatial_scale = ctx.Attr<float>("spatial_scale");
 
     auto in_dims = in->dims();
+    int batch_size = in_dims[0];
     auto in_stride = framework::stride(in_dims);
     int channels = in_dims[1];
     int height = in_dims[2];
     int width = in_dims[3];
 
-    size_t rois_num = rois->dims()[0];
+    int rois_num = rois->dims()[0];
     if (rois_num == 0) return;
 
     int output_size = out->numel();
     int blocks = NumBlocks(output_size);
     int threads = kNumCUDAThreads;
 
+    framework::Tensor roi_batch_id_list;
+    roi_batch_id_list.Resize({rois_num});
+    int* roi_batch_id_data =
+        roi_batch_id_list.mutable_data<int>(platform::CPUPlace());
+    auto rois_lod = rois->lod().back();
+    int rois_batch_size = rois_lod.size() - 1;
+    PADDLE_ENFORCE_EQ(
+        rois_batch_size, batch_size,
+        "The rois_batch_size and imgs batch_size must be the same.");
+    int rois_num_with_lod = rois_lod[rois_batch_size];
+    PADDLE_ENFORCE_EQ(rois_num, rois_num_with_lod,
+                      "The rois_num from input and lod must be the same.");
+    for (int n = 0; n < rois_batch_size; ++n) {
+      for (size_t i = rois_lod[n]; i < rois_lod[n + 1]; ++i) {
+        roi_batch_id_data[i] = n;
+      }
+    }
+
+    framework::Tensor roi_batch_id_list_gpu;
+    framework::TensorCopy(roi_batch_id_list, ctx.GetPlace(),
+                          ctx.device_context(), &roi_batch_id_list_gpu);
+
     GPUROIPoolForward<
         T><<<blocks, threads, 0, ctx.cuda_device_context().stream()>>>(
         output_size, in->data<T>(), rois->data<int64_t>(), spatial_scale,
         channels, height, width, pooled_height, pooled_width,
-        out->mutable_data<T>(ctx.GetPlace()),
+        roi_batch_id_list_gpu.data<int>(), out->mutable_data<T>(ctx.GetPlace()),
         argmax->mutable_data<int64_t>(ctx.GetPlace()));
   }
 };
@@ -159,7 +180,7 @@ class GPUROIPoolGradOpKernel : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext& ctx) const override {
     auto* in = ctx.Input<Tensor>("X");
-    auto* rois = ctx.Input<Tensor>("ROIs");
+    auto* rois = ctx.Input<LoDTensor>("ROIs");
     auto* argmax = ctx.Input<Tensor>("Argmax");
 
     auto* out_grad = ctx.Input<Tensor>(framework::GradVarName("Out"));
@@ -169,12 +190,27 @@ class GPUROIPoolGradOpKernel : public framework::OpKernel<T> {
     auto pooled_width = ctx.Attr<int>("pooled_width");
     auto spatial_scale = ctx.Attr<float>("spatial_scale");
 
-    size_t rois_num = rois->dims()[0];
+    int rois_num = rois->dims()[0];
     int channels = in->dims()[1];
     int height = in->dims()[2];
     int width = in->dims()[3];
 
     if (x_grad) {
+      framework::Tensor roi_batch_id_list;
+      roi_batch_id_list.Resize({rois_num});
+      int* roi_batch_id_data =
+          roi_batch_id_list.mutable_data<int>(platform::CPUPlace());
+      auto rois_lod = rois->lod().back();
+      int rois_batch_size = rois_lod.size() - 1;
+      for (int n = 0; n < rois_batch_size; ++n) {
+        for (size_t i = rois_lod[n]; i < rois_lod[n + 1]; ++i) {
+          roi_batch_id_data[i] = n;
+        }
+      }
+      framework::Tensor roi_batch_id_list_gpu;
+      framework::TensorCopy(roi_batch_id_list, ctx.GetPlace(),
+                            ctx.device_context(), &roi_batch_id_list_gpu);
+
       x_grad->mutable_data<T>(ctx.GetPlace());
       math::SetConstant<Place, T> set_zero;
       set_zero(ctx.cuda_device_context(), x_grad, static_cast<T>(0));
@@ -189,6 +225,7 @@ class GPUROIPoolGradOpKernel : public framework::OpKernel<T> {
             output_grad_size, rois->data<int64_t>(), out_grad->data<T>(),
             argmax->data<int64_t>(), rois_num, spatial_scale, channels, height,
             width, pooled_height, pooled_width,
+            roi_batch_id_list_gpu.data<int>(),
             x_grad->mutable_data<T>(ctx.GetPlace()));
       }
     }

paddle/fluid/operators/roi_pool_op.h

@@ -21,12 +21,14 @@ limitations under the License. */
 namespace paddle {
 namespace operators {
 
+static constexpr int kROISize = 4;
+
 template <typename DeviceContext, typename T>
 class CPUROIPoolOpKernel : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext& ctx) const override {
     auto* in = ctx.Input<framework::Tensor>("X");
-    auto* rois = ctx.Input<framework::Tensor>("ROIs");
+    auto* rois = ctx.Input<framework::LoDTensor>("ROIs");
     auto* out = ctx.Output<framework::Tensor>("Out");
     auto* argmax = ctx.Output<framework::Tensor>("Argmax");
 
@@ -47,24 +49,36 @@ class CPUROIPoolOpKernel : public framework::OpKernel<T> {
     auto out_stride = framework::stride(out->dims());
 
     const T* input_data = in->data<T>();
-    const int64_t* rois_data = rois->data<int64_t>();
-    T* output_data = out->mutable_data<T>(ctx.GetPlace());
-    int64_t* argmax_data = argmax->mutable_data<int64_t>(ctx.GetPlace());
 
-    for (int n = 0; n < rois_num; ++n) {
-      int roi_batch_id = rois_data[0];
-      PADDLE_ENFORCE_GE(roi_batch_id, 0);
-      PADDLE_ENFORCE_LT(roi_batch_id, batch_size);
-      rois_data += roi_stride[0];
+    framework::Tensor roi_batch_id_list;
+    roi_batch_id_list.Resize({rois_num});
+    int* roi_batch_id_data =
+        roi_batch_id_list.mutable_data<int>(ctx.GetPlace());
+
+    auto rois_lod = rois->lod().back();
+    int rois_batch_size = rois_lod.size() - 1;
+    PADDLE_ENFORCE_EQ(
+        rois_batch_size, batch_size,
+        "The rois_batch_size and imgs batch_size must be the same.");
+    int rois_num_with_lod = rois_lod[rois_batch_size];
+    PADDLE_ENFORCE_EQ(rois_num, rois_num_with_lod,
+                      "The rois_num from input and lod must be the same.");
+    for (int n = 0; n < rois_batch_size; ++n) {
+      for (size_t i = rois_lod[n]; i < rois_lod[n + 1]; ++i) {
+        roi_batch_id_data[i] = n;
+      }
     }
 
-    rois_data = rois->data<int64_t>();
+    T* output_data = out->mutable_data<T>(ctx.GetPlace());
+    int64_t* argmax_data = argmax->mutable_data<int64_t>(ctx.GetPlace());
+
+    const int64_t* rois_data = rois->data<int64_t>();
     for (int n = 0; n < rois_num; ++n) {
-      int roi_batch_id = rois_data[0];
-      int roi_start_w = round(rois_data[1] * spatial_scale);
-      int roi_start_h = round(rois_data[2] * spatial_scale);
-      int roi_end_w = round(rois_data[3] * spatial_scale);
-      int roi_end_h = round(rois_data[4] * spatial_scale);
+      int roi_batch_id = roi_batch_id_data[n];
+      int roi_start_w = round(rois_data[0] * spatial_scale);
+      int roi_start_h = round(rois_data[1] * spatial_scale);
+      int roi_end_w = round(rois_data[2] * spatial_scale);
+      int roi_end_h = round(rois_data[3] * spatial_scale);
 
       // Force malformed ROIs to be 1x1
       int roi_height = std::max(roi_end_h - roi_start_h + 1, 1);
@@ -133,7 +147,7 @@ class CPUROIPoolGradOpKernel : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext& ctx) const override {
     auto* in = ctx.Input<framework::Tensor>("X");
-    auto* rois = ctx.Input<framework::Tensor>("ROIs");
+    auto* rois = ctx.Input<framework::LoDTensor>("ROIs");
     auto* argmax = ctx.Input<framework::Tensor>("Argmax");
     auto* out_grad =
         ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
@@ -143,6 +157,20 @@ class CPUROIPoolGradOpKernel : public framework::OpKernel<T> {
     auto pooled_width = ctx.Attr<int>("pooled_width");
 
     if (in_grad) {
+      int rois_num = rois->dims()[0];
+      framework::Tensor roi_batch_id_list;
+      roi_batch_id_list.Resize({rois_num});
+      int* roi_batch_id_data =
+          roi_batch_id_list.mutable_data<int>(ctx.GetPlace());
+
+      auto rois_lod = rois->lod().back();
+      int rois_batch_size = rois_lod.size() - 1;
+      for (int n = 0; n < rois_batch_size; ++n) {
+        for (size_t i = rois_lod[n]; i < rois_lod[n + 1]; ++i) {
+          roi_batch_id_data[i] = n;
+        }
+      }
+
       const int64_t* rois_data = rois->data<int64_t>();
       const T* out_grad_data = out_grad->data<T>();
       const int64_t* argmax_data = argmax->data<int64_t>();
@@ -156,11 +184,10 @@ class CPUROIPoolGradOpKernel : public framework::OpKernel<T> {
       auto roi_stride = framework::stride(rois->dims());
       auto out_stride = framework::stride(out_grad->dims());
 
-      int rois_num = rois->dims()[0];
       int channels = in->dims()[1];
 
       for (int n = 0; n < rois_num; ++n) {
-        int roi_batch_idx = rois_data[0];
+        int roi_batch_idx = roi_batch_id_data[n];
         T* batch_grad_data = in_grad_data + roi_batch_idx * in_stride[0];
         for (int c = 0; c < channels; ++c) {
           for (int ph = 0; ph < pooled_height; ++ph) {

paddle/fluid/operators/send_op.cc

@@ -41,6 +41,8 @@ class SendOp : public framework::OperatorBase {
     std::vector<std::string> endpoints =
         Attr<std::vector<std::string>>("endpoints");
 
+    bool sync_mode = Attr<bool>("sync_mode");
+
     platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
     auto& ctx = *pool.Get(place);
 
@@ -64,11 +66,13 @@ class SendOp : public framework::OperatorBase {
     }
     PADDLE_ENFORCE(rpc_client->Wait());
 
-    for (auto& ep : endpoints) {
-      VLOG(3) << "batch barrier, ep: " << ep;
-      rpc_client->AsyncSendBatchBarrier(ep);
+    if (sync_mode) {
+      for (auto& ep : endpoints) {
+        VLOG(3) << "batch barrier, ep: " << ep;
+        rpc_client->AsyncSendBatchBarrier(ep);
+      }
+      PADDLE_ENFORCE(rpc_client->Wait());
     }
-    PADDLE_ENFORCE(rpc_client->Wait());
 
     if (outs.size() > 0) {
       for (size_t i = 0; i < outs.size(); i++) {
@@ -112,6 +116,7 @@ This operator will send tensor to recv_op at the parameter server.
                                       "Server endpoints in the order of input "
                                       "variables for mapping")
         .SetDefault({});
+    AddAttr<bool>("sync_mode", "work in sync_mode or not").SetDefault(true);
   }
 };
 

paddle/fluid/operators/send_recv_op_test.cc

@@ -137,6 +137,8 @@ void StartServerNet(bool is_sparse) {
   attrs.insert({"GradList", std::vector<std::string>({"x1"})});
   attrs.insert({"OptimizeBlock", optimize_block});
   attrs.insert({"PrefetchBlock", prefetch_block});
+  attrs.insert({"grad_to_block_id", std::vector<std::string>({""})});
+  attrs.insert({"sync_mode", true});
   listen_and_serv_op =
       f::OpRegistry::CreateOp("listen_and_serv", {{"X", {"x1"}}}, {}, attrs);
   listen_and_serv_op->Run(scope, place);