Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into multithread-sparse-adam

test=develop

CMakeLists.txt

@@ -55,6 +55,7 @@ option(WITH_DOUBLE      "Compile PaddlePaddle with double precision"    OFF)
 option(WITH_RDMA        "Compile PaddlePaddle with RDMA support"        OFF)
 option(WITH_TIMER       "Compile PaddlePaddle with stats timer"         OFF)
 option(WITH_PROFILER    "Compile PaddlePaddle with GPU profiler and gperftools"        OFF)
+option(WITH_JEMALLOC    "Compile PaddlePaddle with jemalloc"            OFF)
 option(WITH_DOC         "Compile PaddlePaddle with documentation"       OFF)
 option(WITH_COVERAGE    "Compile PaddlePaddle with code coverage"       OFF)
 option(COVERALLS_UPLOAD "Package code coverage data to coveralls"       OFF)
@@ -261,6 +262,12 @@ if (WITH_PROFILER)
     add_definitions(-DWITH_GPERFTOOLS)
 endif()
 
+if (WITH_JEMALLOC)
+    find_package(JeMalloc REQUIRED)
+    include_directories(${JEMALLOC_INCLUDE_DIR})
+    add_definitions(-DWITH_JEMALLOC)
+endif()
+
 include(generic)            # simplify cmake module
 include(package)            # set paddle packages
 include(ccache)             # set ccache for compilation
@@ -290,7 +297,7 @@ if(WITH_PSLIB)
     list(APPEND EXTERNAL_LIBS pslib_brpc)
     list(APPEND EXTERNAL_LIBS libmct)
 endif(WITH_PSLIB)
-    
+
 if(WITH_AMD_GPU)
     find_package(HIP)
     include(hip)

cmake/FindJeMalloc.cmake

+# - Find JeMalloc library
+# Find the native JeMalloc includes and library
+#
+# JEMALLOC_INCLUDE_DIR - where to find jemalloc.h, etc.
+# JEMALLOC_LIBRARIES - List of libraries when using jemalloc.
+# JEMALLOC_FOUND - True if jemalloc found.
+
+find_path(JEMALLOC_INCLUDE_DIR
+  NAMES jemalloc/jemalloc.h
+  HINTS ${JEMALLOC_ROOT_DIR}/include)
+
+find_library(JEMALLOC_LIBRARIES
+  NAMES jemalloc
+  HINTS ${JEMALLOC_ROOT_DIR}/lib)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(jemalloc DEFAULT_MSG JEMALLOC_LIBRARIES JEMALLOC_INCLUDE_DIR)
+
+mark_as_advanced(
+  JEMALLOC_LIBRARIES
+  JEMALLOC_INCLUDE_DIR)
+
+if (JEMALLOC_FOUND)
+  add_library(jemalloc::jemalloc UNKNOWN IMPORTED)
+  set_target_properties(jemalloc::jemalloc PROPERTIES
+    IMPORTED_LOCATION ${JEMALLOC_LIBRARIES}
+    INTERFACE_INCLUDE_DIRECTORIES "${JEMALLOC_INCLUDE_DIR}")
+endif()

cmake/configure.cmake

@@ -134,6 +134,7 @@ if(WITH_GPU)
             message(WARNING "Anakin needs CUDNN >= 7.0 to compile. Force WITH_ANAKIN=OFF")
             set(WITH_ANAKIN OFF CACHE STRING "Anakin is valid only when CUDNN >= 7.0." FORCE)
         endif()
+        add_definitions(-DWITH_ANAKIN)
     endif()
     if(WITH_ANAKIN)
         # NOTICE(minqiyang): the end slash is important because $CUDNN_INCLUDE_DIR

cmake/cuda.cmake

@@ -5,6 +5,8 @@ endif()
 set(paddle_known_gpu_archs "30 35 50 52 60 61 70")
 set(paddle_known_gpu_archs7 "30 35 50 52")
 set(paddle_known_gpu_archs8 "30 35 50 52 60 61")
+set(paddle_known_gpu_archs9 "30 35 50 52 60 61 70")
+set(paddle_known_gpu_archs10 "30 35 50 52 60 61 70 75")
 
 ######################################################################################
 # A function for automatic detection of GPUs installed  (if autodetection is enabled)
@@ -59,7 +61,7 @@ endfunction()
 #   select_nvcc_arch_flags(out_variable)
 function(select_nvcc_arch_flags out_variable)
   # List of arch names
-  set(archs_names "Kepler" "Maxwell" "Pascal" "All" "Manual")
+  set(archs_names "Kepler" "Maxwell" "Pascal" "Volta" "Turing" "All" "Manual")
   set(archs_name_default "All")
   if(NOT CMAKE_CROSSCOMPILING)
     list(APPEND archs_names "Auto")
@@ -93,6 +95,8 @@ function(select_nvcc_arch_flags out_variable)
     set(cuda_arch_bin "60 61")
   elseif(${CUDA_ARCH_NAME} STREQUAL "Volta")
     set(cuda_arch_bin "70")
+  elseif(${CUDA_ARCH_NAME} STREQUAL "Turing")
+    set(cuda_arch_bin "75")
   elseif(${CUDA_ARCH_NAME} STREQUAL "All")
     set(cuda_arch_bin ${paddle_known_gpu_archs})
   elseif(${CUDA_ARCH_NAME} STREQUAL "Auto")
@@ -153,6 +157,16 @@ elseif (${CUDA_VERSION} LESS 9.0) # CUDA 8.x
   # warning for now.
   list(APPEND CUDA_NVCC_FLAGS "-Wno-deprecated-gpu-targets")
   add_definitions("-DPADDLE_CUDA_BINVER=\"80\"")
+elseif (${CUDA_VERSION} LESS 10.0) # CUDA 9.x
+  set(paddle_known_gpu_archs ${paddle_known_gpu_archs9})
+  list(APPEND CUDA_NVCC_FLAGS "-D_MWAITXINTRIN_H_INCLUDED")
+  list(APPEND CUDA_NVCC_FLAGS "-D__STRICT_ANSI__")
+  add_definitions("-DPADDLE_CUDA_BINVER=\"90\"")
+elseif (${CUDA_VERSION} LESS 11.0) # CUDA 10.x
+  set(paddle_known_gpu_archs ${paddle_known_gpu_archs10})
+  list(APPEND CUDA_NVCC_FLAGS "-D_MWAITXINTRIN_H_INCLUDED")
+  list(APPEND CUDA_NVCC_FLAGS "-D__STRICT_ANSI__")
+  add_definitions("-DPADDLE_CUDA_BINVER=\"100\"")
 endif()
 
 include_directories(${CUDA_INCLUDE_DIRS})

cmake/external/boost.cmake

@@ -23,11 +23,8 @@ set(BOOST_PROJECT       "extern_boost")
 # checked that the devtools package of CentOS 6 installs boost 1.41.0.
 # So we use 1.41.0 here.
 set(BOOST_VER           "1.41.0")
-if((NOT DEFINED BOOST_TAR) OR (NOT DEFINED BOOST_URL))
-    message(STATUS "use pre defined download url")
-    set(BOOST_TAR "boost_1_41_0" CACHE STRING "" FORCE)
-    set(BOOST_URL "http://paddlepaddledeps.cdn.bcebos.com/${BOOST_TAR}.tar.gz" CACHE STRING "" FORCE)
-endif()
+set(BOOST_TAR "boost_1_41_0" CACHE STRING "" FORCE)
+set(BOOST_URL "http://paddlepaddledeps.cdn.bcebos.com/${BOOST_TAR}.tar.gz" CACHE STRING "" FORCE)
 
 MESSAGE(STATUS "BOOST_TAR: ${BOOST_TAR}, BOOST_URL: ${BOOST_URL}")
 

cmake/external/gflags.cmake

@@ -63,6 +63,15 @@ ADD_DEPENDENCIES(gflags extern_gflags)
 
 LIST(APPEND external_project_dependencies gflags)
 
+# On Windows (including MinGW), the Shlwapi library is used by gflags if available.
+if (WIN32)
+  include(CheckIncludeFileCXX)
+  check_include_file_cxx("shlwapi.h" HAVE_SHLWAPI)
+  if (HAVE_SHLWAPI)
+    set_property(GLOBAL PROPERTY OS_DEPENDENCY_MODULES shlwapi.lib)
+  endif(HAVE_SHLWAPI)
+endif (WIN32)
+
 IF(WITH_C_API)
   INSTALL(DIRECTORY ${GFLAGS_INCLUDE_DIR} DESTINATION third_party/gflags)
   IF(ANDROID)

cmake/external/mkldnn.cmake

@@ -55,7 +55,7 @@ ExternalProject_Add(
     ${MKLDNN_PROJECT}
     ${EXTERNAL_PROJECT_LOG_ARGS}
     DEPENDS             ${MKLDNN_DEPENDS}
-    GIT_REPOSITORY      "https://github.com/01org/mkl-dnn.git"
+    GIT_REPOSITORY      "https://github.com/intel/mkl-dnn.git"
     GIT_TAG             "830a10059a018cd2634d94195140cf2d8790a75a"
     PREFIX              ${MKLDNN_SOURCES_DIR}
     UPDATE_COMMAND      ""

cmake/external/mklml.cmake

@@ -16,6 +16,12 @@ IF(NOT ${WITH_MKLML})
   return()
 ENDIF(NOT ${WITH_MKLML})
 
+IF(APPLE)
+    MESSAGE(WARNING "Mac is not supported with MKLML in Paddle yet. Force WITH_MKLML=OFF.")
+    SET(WITH_MKLML OFF CACHE STRING "Disable MKLML package in MacOS" FORCE)
+    return()
+ENDIF()
+
 INCLUDE(ExternalProject)
 SET(MKLML_DST_DIR       "mklml")
 SET(MKLML_INSTALL_ROOT  "${THIRD_PARTY_PATH}/install")
@@ -23,32 +29,24 @@ SET(MKLML_INSTALL_DIR   ${MKLML_INSTALL_ROOT}/${MKLML_DST_DIR})
 SET(MKLML_ROOT          ${MKLML_INSTALL_DIR})
 SET(MKLML_INC_DIR       ${MKLML_ROOT}/include)
 SET(MKLML_LIB_DIR       ${MKLML_ROOT}/lib)
-if(WIN32)
+SET(CMAKE_INSTALL_RPATH "${CMAKE_INSTALL_RPATH}" "${MKLML_ROOT}/lib")
+
+SET(TIME_VERSION "2019.0.1.20181227")
+IF(WIN32)
+    SET(MKLML_VER "mklml_win_${TIME_VERSION}" CACHE STRING "" FORCE)
+    SET(MKLML_URL "https://paddlepaddledeps.cdn.bcebos.com/${MKLML_VER}.zip" CACHE STRING "" FORCE)
     SET(MKLML_LIB                 ${MKLML_LIB_DIR}/mklml.lib)
     SET(MKLML_IOMP_LIB            ${MKLML_LIB_DIR}/libiomp5md.lib)
     SET(MKLML_SHARED_LIB          ${MKLML_LIB_DIR}/mklml.dll)
     SET(MKLML_SHARED_IOMP_LIB     ${MKLML_LIB_DIR}/libiomp5md.dll)
-else()
+ELSE()  
+    SET(MKLML_VER "mklml_lnx_${TIME_VERSION}" CACHE STRING "" FORCE)
+    SET(MKLML_URL "http://paddlepaddledeps.cdn.bcebos.com/${MKLML_VER}.tgz" CACHE STRING "" FORCE)
     SET(MKLML_LIB                 ${MKLML_LIB_DIR}/libmklml_intel.so)
     SET(MKLML_IOMP_LIB            ${MKLML_LIB_DIR}/libiomp5.so)
     SET(MKLML_SHARED_LIB          ${MKLML_LIB_DIR}/libmklml_intel.so)
     SET(MKLML_SHARED_IOMP_LIB     ${MKLML_LIB_DIR}/libiomp5.so)
-endif()
-SET(CMAKE_INSTALL_RPATH "${CMAKE_INSTALL_RPATH}" "${MKLML_ROOT}/lib")
-
-IF((NOT DEFINED MKLML_VER) OR (NOT DEFINED MKLML_URL))
-    MESSAGE(STATUS "use pre defined download url")
-    if(WIN32)
-        SET(MKLML_VER "mklml_win_2019.0.1.20180928" CACHE STRING "" FORCE)
-        SET(MKLML_URL "https://paddlepaddledeps.cdn.bcebos.com/${MKLML_VER}.zip" CACHE STRING "" FORCE)
-    elseif(APPLE)
-        SET(MKLML_VER "mklml_mac_2019.0.1.20180928" CACHE STRING "" FORCE)
-        SET(MKLML_URL "http://paddlepaddledeps.cdn.bcebos.com/${MKLML_VER}.tgz" CACHE STRING "" FORCE)
-    else()
-        SET(MKLML_VER "mklml_lnx_2019.0.1.20180928" CACHE STRING "" FORCE)
-        SET(MKLML_URL "http://paddlepaddledeps.cdn.bcebos.com/${MKLML_VER}.tgz" CACHE STRING "" FORCE)
-    ENDIF()
-endif()
+ENDIF()
 
 SET(MKLML_PROJECT       "extern_mklml")
 MESSAGE(STATUS "MKLML_VER: ${MKLML_VER}, MKLML_URL: ${MKLML_URL}")

cmake/external/ngraph.cmake

@@ -37,14 +37,18 @@ INCLUDE(GNUInstallDirs)
 INCLUDE(ExternalProject)
 
 SET(NGRAPH_PROJECT         "extern_ngraph")
-SET(NGRAPH_GIT_TAG         "v0.10.1")
+SET(NGRAPH_GIT_TAG         "20bd8bbc79ae3a81c57313846a2be7313e5d1dab")
 SET(NGRAPH_SOURCES_DIR     ${THIRD_PARTY_PATH}/ngraph)
 SET(NGRAPH_INSTALL_DIR     ${THIRD_PARTY_PATH}/install/ngraph)
 SET(NGRAPH_INC_DIR         ${NGRAPH_INSTALL_DIR}/include)
 SET(NGRAPH_LIB_DIR         ${NGRAPH_INSTALL_DIR}/${CMAKE_INSTALL_LIBDIR})
 SET(NGRAPH_SHARED_LIB_NAME libngraph.so)
 SET(NGRAPH_CPU_LIB_NAME    libcpu_backend.so)
-SET(NGRAPH_TBB_LIB_NAME    libtbb.so.2)
+if(CMAKE_BUILD_TYPE STREQUAL "Debug")
+    SET(NGRAPH_TBB_LIB_NAME    libtbb_debug.so.2)
+else()
+    SET(NGRAPH_TBB_LIB_NAME    libtbb.so.2)
+endif()
 SET(NGRAPH_GIT_REPO        "https://github.com/NervanaSystems/ngraph.git")
 SET(NGRAPH_SHARED_LIB      ${NGRAPH_LIB_DIR}/${NGRAPH_SHARED_LIB_NAME})
 SET(NGRAPH_CPU_LIB         ${NGRAPH_LIB_DIR}/${NGRAPH_CPU_LIB_NAME})
@@ -66,16 +70,7 @@ ExternalProject_Add(
     CMAKE_ARGS          -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
     CMAKE_ARGS          -DMKLDNN_INCLUDE_DIR=${MKLDNN_INC_DIR}
     CMAKE_ARGS          -DMKLDNN_LIB_DIR=${MKLDNN_INSTALL_DIR}/lib
-)
-
-# Workaround for nGraph expecting mklml to be in mkldnn install directory.
-ExternalProject_Add_Step(
-    ${NGRAPH_PROJECT}
-    PrepareMKL
-    COMMAND ${CMAKE_COMMAND} -E create_symlink ${MKLML_LIB} ${MKLDNN_INSTALL_DIR}/lib/libmklml_intel.so
-    COMMAND ${CMAKE_COMMAND} -E create_symlink ${MKLML_IOMP_LIB} ${MKLDNN_INSTALL_DIR}/lib/libiomp5.so
-    DEPENDEES download
-    DEPENDERS configure
+    CMAKE_ARGS          -DMKLML_LIB_DIR=${MKLML_INSTALL_DIR}/lib
 )
 
 add_dependencies(ngraph ${NGRAPH_PROJECT})

cmake/generic.cmake

@@ -115,6 +115,10 @@ function(common_link TARGET_NAME)
   if (WITH_PROFILER)
     target_link_libraries(${TARGET_NAME} gperftools::profiler)
   endif()
+
+  if (WITH_JEMALLOC)
+    target_link_libraries(${TARGET_NAME} jemalloc::jemalloc)
+  endif()
 endfunction()
 
 
@@ -228,7 +232,7 @@ function(merge_static_libs TARGET_NAME)
       # Get the file names of the libraries to be merged
       set(libfiles ${libfiles} $<TARGET_FILE:${lib}>)
     endforeach()
-    # msvc will put libarary in directory of "/Release/xxxlib" by default 
+    # msvc will put libarary in directory of "/Release/xxxlib" by default
     #       COMMAND cmake -E remove "${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_BUILD_TYPE}/${TARGET_NAME}.lib"
     add_custom_command(TARGET ${TARGET_NAME} POST_BUILD
       COMMAND cmake -E make_directory "${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_BUILD_TYPE}"
@@ -355,6 +359,8 @@ function(cc_binary TARGET_NAME)
     add_dependencies(${TARGET_NAME} ${cc_binary_DEPS})
     common_link(${TARGET_NAME})
   endif()
+  get_property(os_dependency_modules GLOBAL PROPERTY OS_DEPENDENCY_MODULES)
+  target_link_libraries(${TARGET_NAME} ${os_dependency_modules})
 endfunction(cc_binary)
 
 function(cc_test TARGET_NAME)
@@ -363,18 +369,15 @@ function(cc_test TARGET_NAME)
     set(oneValueArgs "")
     set(multiValueArgs SRCS DEPS ARGS)
     cmake_parse_arguments(cc_test "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
+    add_executable(${TARGET_NAME} ${cc_test_SRCS})
     if(WIN32)
-      list(APPEND win32_deps shlwapi)
       if("${cc_test_DEPS};" MATCHES "python;")
         list(REMOVE_ITEM cc_test_DEPS python)
-        list(APPEND win32_deps ${PYTHON_LIBRARIES})
+        target_link_libraries(${TARGET_NAME} ${PYTHON_LIBRARIES})
       endif()
     endif(WIN32)
-    add_executable(${TARGET_NAME} ${cc_test_SRCS})
-    target_link_libraries(${TARGET_NAME} ${cc_test_DEPS} paddle_gtest_main lod_tensor memory gtest gflags glog)
-    if(WIN32)
-      target_link_libraries(${TARGET_NAME} ${win32_deps})
-    endif(WIN32)
+    get_property(os_dependency_modules GLOBAL PROPERTY OS_DEPENDENCY_MODULES)
+    target_link_libraries(${TARGET_NAME} ${cc_test_DEPS} ${os_dependency_modules} paddle_gtest_main lod_tensor memory gtest gflags glog)
     add_dependencies(${TARGET_NAME} ${cc_test_DEPS} paddle_gtest_main lod_tensor memory gtest gflags glog)
     common_link(${TARGET_NAME})
     add_test(NAME ${TARGET_NAME}
@@ -447,7 +450,8 @@ function(nv_test TARGET_NAME)
     set(multiValueArgs SRCS DEPS)
     cmake_parse_arguments(nv_test "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
     cuda_add_executable(${TARGET_NAME} ${nv_test_SRCS})
-    target_link_libraries(${TARGET_NAME} ${nv_test_DEPS} paddle_gtest_main lod_tensor memory gtest gflags glog)
+    get_property(os_dependency_modules GLOBAL PROPERTY OS_DEPENDENCY_MODULES)
+    target_link_libraries(${TARGET_NAME} ${nv_test_DEPS} paddle_gtest_main lod_tensor memory gtest gflags glog ${os_dependency_modules})
     add_dependencies(${TARGET_NAME} ${nv_test_DEPS} paddle_gtest_main lod_tensor memory gtest gflags glog)
     common_link(${TARGET_NAME})
     add_test(${TARGET_NAME} ${TARGET_NAME})
@@ -534,7 +538,8 @@ function(hip_test TARGET_NAME)
     endif()
     add_executable(${TARGET_NAME} ${_cmake_options} ${_generated_files} ${_sources})
     set_target_properties(${TARGET_NAME} PROPERTIES LINKER_LANGUAGE HIP)
-    target_link_libraries(${TARGET_NAME} ${hip_test_DEPS} paddle_gtest_main memory gtest gflags)
+    get_property(os_dependency_modules GLOBAL PROPERTY OS_DEPENDENCY_MODULES)
+    target_link_libraries(${TARGET_NAME} ${hip_test_DEPS} paddle_gtest_main memory gtest gflags ${os_dependency_modules})
     add_dependencies(${TARGET_NAME} ${hip_test_DEPS} paddle_gtest_main memory gtest gflags)
     common_link(${TARGET_NAME})
     add_test(${TARGET_NAME} ${TARGET_NAME})

paddle/fluid/API.spec

@@ -88,6 +88,7 @@ paddle.fluid.layers.pool3d ArgSpec(args=['input', 'pool_size', 'pool_type', 'poo
 paddle.fluid.layers.adaptive_pool2d ArgSpec(args=['input', 'pool_size', 'pool_type', 'require_index', 'name'], varargs=None, keywords=None, defaults=('max', False, None))
 paddle.fluid.layers.adaptive_pool3d ArgSpec(args=['input', 'pool_size', 'pool_type', 'require_index', 'name'], varargs=None, keywords=None, defaults=('max', False, None))
 paddle.fluid.layers.batch_norm ArgSpec(args=['input', 'act', 'is_test', 'momentum', 'epsilon', 'param_attr', 'bias_attr', 'data_layout', 'in_place', 'name', 'moving_mean_name', 'moving_variance_name', 'do_model_average_for_mean_and_var', 'fuse_with_relu', 'use_global_stats'], varargs=None, keywords=None, defaults=(None, False, 0.9, 1e-05, None, None, 'NCHW', False, None, None, None, False, False, False))
+paddle.fluid.layers.data_norm ArgSpec(args=['input', 'act', 'epsilon', 'param_attr', 'data_layout', 'in_place', 'use_mkldnn', 'name', 'moving_mean_name', 'moving_variance_name', 'do_model_average_for_mean_and_var'], varargs=None, keywords=None, defaults=(None, 1e-05, None, 'NCHW', False, False, None, None, None, False))
 paddle.fluid.layers.beam_search_decode ArgSpec(args=['ids', 'scores', 'beam_size', 'end_id', 'name'], varargs=None, keywords=None, defaults=(None,))
 paddle.fluid.layers.conv2d_transpose ArgSpec(args=['input', 'num_filters', 'output_size', 'filter_size', 'padding', 'stride', 'dilation', 'groups', 'param_attr', 'bias_attr', 'use_cudnn', 'act', 'name'], varargs=None, keywords=None, defaults=(None, None, 0, 1, 1, None, None, None, True, None, None))
 paddle.fluid.layers.conv3d_transpose ArgSpec(args=['input', 'num_filters', 'output_size', 'filter_size', 'padding', 'stride', 'dilation', 'groups', 'param_attr', 'bias_attr', 'use_cudnn', 'act', 'name'], varargs=None, keywords=None, defaults=(None, None, 0, 1, 1, None, None, None, True, None, None))
@@ -210,6 +211,7 @@ paddle.fluid.layers.get_tensor_from_selected_rows ArgSpec(args=['x', 'name'], va
 paddle.fluid.layers.lstm ArgSpec(args=['input', 'init_h', 'init_c', 'max_len', 'hidden_size', 'num_layers', 'dropout_prob', 'is_bidirec', 'is_test', 'name', 'default_initializer', 'seed'], varargs=None, keywords=None, defaults=(0.0, False, False, None, None, -1))
 paddle.fluid.layers.py_func ArgSpec(args=['func', 'x', 'out', 'backward_func', 'skip_vars_in_backward_input'], varargs=None, keywords=None, defaults=(None, None))
 paddle.fluid.layers.psroi_pool ArgSpec(args=['input', 'rois', 'output_channels', 'spatial_scale', 'pooled_height', 'pooled_width', 'name'], varargs=None, keywords=None, defaults=(None,))
+paddle.fluid.layers.teacher_student_sigmoid_loss ArgSpec(args=['input', 'label', 'soft_max_up_bound', 'soft_max_lower_bound'], varargs=None, keywords=None, defaults=(15.0, -15.0))
 paddle.fluid.layers.huber_loss ArgSpec(args=['input', 'label', 'delta'], varargs=None, keywords=None, defaults=None)
 paddle.fluid.layers.data ArgSpec(args=['name', 'shape', 'append_batch_size', 'dtype', 'lod_level', 'type', 'stop_gradient'], varargs=None, keywords=None, defaults=(True, 'float32', 0, VarType.LOD_TENSOR, True))
 paddle.fluid.layers.open_files ArgSpec(args=['filenames', 'shapes', 'lod_levels', 'dtypes', 'thread_num', 'buffer_size', 'pass_num', 'is_test'], varargs=None, keywords=None, defaults=(None, None, 1, None))
@@ -405,28 +407,50 @@ paddle.fluid.nets.glu ArgSpec(args=['input', 'dim'], varargs=None, keywords=None
 paddle.fluid.nets.scaled_dot_product_attention ArgSpec(args=['queries', 'keys', 'values', 'num_heads', 'dropout_rate'], varargs=None, keywords=None, defaults=(1, 0.0))
 paddle.fluid.nets.img_conv_group ArgSpec(args=['input', 'conv_num_filter', 'pool_size', 'conv_padding', 'conv_filter_size', 'conv_act', 'param_attr', 'conv_with_batchnorm', 'conv_batchnorm_drop_rate', 'pool_stride', 'pool_type', 'use_cudnn'], varargs=None, keywords=None, defaults=(1, 3, None, None, False, 0.0, 1, 'max', True))
 paddle.fluid.optimizer.SGDOptimizer.__init__ ArgSpec(args=['self', 'learning_rate', 'regularization', 'name'], varargs=None, keywords=None, defaults=(None, None))
+paddle.fluid.optimizer.SGDOptimizer.apply_gradients ArgSpec(args=['self', 'params_grads'], varargs=None, keywords=None, defaults=None)
+paddle.fluid.optimizer.SGDOptimizer.backward ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set', 'callbacks'], varargs=None, keywords=None, defaults=(None, None, None, None))
 paddle.fluid.optimizer.SGDOptimizer.minimize ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.optimizer.MomentumOptimizer.__init__ ArgSpec(args=['self', 'learning_rate', 'momentum', 'use_nesterov', 'regularization', 'name'], varargs=None, keywords=None, defaults=(False, None, None))
+paddle.fluid.optimizer.MomentumOptimizer.apply_gradients ArgSpec(args=['self', 'params_grads'], varargs=None, keywords=None, defaults=None)
+paddle.fluid.optimizer.MomentumOptimizer.backward ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set', 'callbacks'], varargs=None, keywords=None, defaults=(None, None, None, None))
 paddle.fluid.optimizer.MomentumOptimizer.minimize ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.optimizer.AdagradOptimizer.__init__ ArgSpec(args=['self', 'learning_rate', 'epsilon', 'regularization', 'name'], varargs=None, keywords=None, defaults=(1e-06, None, None))
+paddle.fluid.optimizer.AdagradOptimizer.apply_gradients ArgSpec(args=['self', 'params_grads'], varargs=None, keywords=None, defaults=None)
+paddle.fluid.optimizer.AdagradOptimizer.backward ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set', 'callbacks'], varargs=None, keywords=None, defaults=(None, None, None, None))
 paddle.fluid.optimizer.AdagradOptimizer.minimize ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.optimizer.AdamOptimizer.__init__ ArgSpec(args=['self', 'learning_rate', 'beta1', 'beta2', 'epsilon', 'regularization', 'name', 'lazy_mode', 'min_row_size_to_use_multithread'], varargs=None, keywords=None, defaults=(0.001, 0.9, 0.999, 1e-08, None, None, False, 0))
+paddle.fluid.optimizer.AdamOptimizer.apply_gradients ArgSpec(args=['self', 'params_grads'], varargs=None, keywords=None, defaults=None)
+paddle.fluid.optimizer.AdamOptimizer.backward ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set', 'callbacks'], varargs=None, keywords=None, defaults=(None, None, None, None))
 paddle.fluid.optimizer.AdamOptimizer.minimize ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.optimizer.AdamaxOptimizer.__init__ ArgSpec(args=['self', 'learning_rate', 'beta1', 'beta2', 'epsilon', 'regularization', 'name'], varargs=None, keywords=None, defaults=(0.001, 0.9, 0.999, 1e-08, None, None))
+paddle.fluid.optimizer.AdamaxOptimizer.apply_gradients ArgSpec(args=['self', 'params_grads'], varargs=None, keywords=None, defaults=None)
+paddle.fluid.optimizer.AdamaxOptimizer.backward ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set', 'callbacks'], varargs=None, keywords=None, defaults=(None, None, None, None))
 paddle.fluid.optimizer.AdamaxOptimizer.minimize ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.optimizer.DecayedAdagradOptimizer.__init__ ArgSpec(args=['self', 'learning_rate', 'decay', 'epsilon', 'regularization', 'name'], varargs=None, keywords=None, defaults=(0.95, 1e-06, None, None))
+paddle.fluid.optimizer.DecayedAdagradOptimizer.apply_gradients ArgSpec(args=['self', 'params_grads'], varargs=None, keywords=None, defaults=None)
+paddle.fluid.optimizer.DecayedAdagradOptimizer.backward ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set', 'callbacks'], varargs=None, keywords=None, defaults=(None, None, None, None))
 paddle.fluid.optimizer.DecayedAdagradOptimizer.minimize ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.optimizer.FtrlOptimizer.__init__ ArgSpec(args=['self', 'learning_rate', 'l1', 'l2', 'lr_power', 'regularization', 'name'], varargs=None, keywords=None, defaults=(0.0, 0.0, -0.5, None, None))
+paddle.fluid.optimizer.FtrlOptimizer.apply_gradients ArgSpec(args=['self', 'params_grads'], varargs=None, keywords=None, defaults=None)
+paddle.fluid.optimizer.FtrlOptimizer.backward ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set', 'callbacks'], varargs=None, keywords=None, defaults=(None, None, None, None))
 paddle.fluid.optimizer.FtrlOptimizer.minimize ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.optimizer.RMSPropOptimizer.__init__ ArgSpec(args=['self', 'learning_rate', 'rho', 'epsilon', 'momentum', 'centered', 'regularization', 'name'], varargs=None, keywords=None, defaults=(0.95, 1e-06, 0.0, False, None, None))
+paddle.fluid.optimizer.RMSPropOptimizer.apply_gradients ArgSpec(args=['self', 'params_grads'], varargs=None, keywords=None, defaults=None)
+paddle.fluid.optimizer.RMSPropOptimizer.backward ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set', 'callbacks'], varargs=None, keywords=None, defaults=(None, None, None, None))
 paddle.fluid.optimizer.RMSPropOptimizer.minimize ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.optimizer.AdadeltaOptimizer.__init__ ArgSpec(args=['self', 'learning_rate', 'epsilon', 'rho', 'regularization', 'name'], varargs=None, keywords=None, defaults=(1e-06, 0.95, None, None))
+paddle.fluid.optimizer.AdadeltaOptimizer.apply_gradients ArgSpec(args=['self', 'params_grads'], varargs=None, keywords=None, defaults=None)
+paddle.fluid.optimizer.AdadeltaOptimizer.backward ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set', 'callbacks'], varargs=None, keywords=None, defaults=(None, None, None, None))
 paddle.fluid.optimizer.AdadeltaOptimizer.minimize ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.optimizer.ModelAverage.__init__ ArgSpec(args=['self', 'average_window_rate', 'min_average_window', 'max_average_window', 'regularization', 'name'], varargs=None, keywords=None, defaults=(10000, 10000, None, None))
 paddle.fluid.optimizer.ModelAverage.apply ArgSpec(args=[], varargs='args', keywords='kwds', defaults=None)
+paddle.fluid.optimizer.ModelAverage.apply_gradients ArgSpec(args=['self', 'params_grads'], varargs=None, keywords=None, defaults=None)
+paddle.fluid.optimizer.ModelAverage.backward ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set', 'callbacks'], varargs=None, keywords=None, defaults=(None, None, None, None))
 paddle.fluid.optimizer.ModelAverage.minimize ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.optimizer.ModelAverage.restore ArgSpec(args=['self', 'executor'], varargs=None, keywords=None, defaults=None)
 paddle.fluid.optimizer.LarsMomentumOptimizer.__init__ ArgSpec(args=['self', 'learning_rate', 'momentum', 'lars_coeff', 'lars_weight_decay', 'regularization', 'name'], varargs=None, keywords=None, defaults=(0.001, 0.0005, None, None))
+paddle.fluid.optimizer.LarsMomentumOptimizer.apply_gradients ArgSpec(args=['self', 'params_grads'], varargs=None, keywords=None, defaults=None)
+paddle.fluid.optimizer.LarsMomentumOptimizer.backward ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set', 'callbacks'], varargs=None, keywords=None, defaults=(None, None, None, None))
 paddle.fluid.optimizer.LarsMomentumOptimizer.minimize ArgSpec(args=['self', 'loss', 'startup_program', 'parameter_list', 'no_grad_set'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.backward.append_backward ArgSpec(args=['loss', 'parameter_list', 'no_grad_set', 'callbacks'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.regularizer.L1DecayRegularizer.__init__ ArgSpec(args=['self', 'regularization_coeff'], varargs=None, keywords=None, defaults=(0.0,))

paddle/fluid/framework/details/CMakeLists.txt

@@ -94,4 +94,4 @@ cc_library(build_strategy SRCS build_strategy.cc DEPS
         graph_viz_pass multi_devices_graph_pass
         multi_devices_graph_print_pass multi_devices_graph_check_pass
         fuse_elewise_add_act_pass multi_batch_merge_pass
-        memory_optimize_pass)
+        memory_optimize_pass lock_free_optimize_pass)

paddle/fluid/framework/details/build_strategy.cc

@@ -18,7 +18,7 @@ limitations under the License. */
 #include <memory>
 
 #include "paddle/fluid/framework/details/memory_reuse_types.h"
-#include "paddle/fluid/framework/details/multi_devices_graph_check_pass.h"
+#include "paddle/fluid/framework/details/multi_devices_graph_pass.h"
 #include "paddle/fluid/framework/details/multi_devices_graph_print_pass.h"
 #include "paddle/fluid/framework/details/reduce_op_handle.h"
 #include "paddle/fluid/framework/details/sequential_execution_pass.h"
@@ -86,10 +86,8 @@ class ParallelExecutorPassBuilder : public ir::PassBuilder {
     if (strategy.memory_optimize_) {
       auto analysis_var_pass = AppendPass("analysis_var_pass");
     }
-    // Convert graph to run on multi-devices.
-    auto multi_devices_pass = AppendPass("multi_devices_pass");
-    multi_devices_pass->SetNotOwned<const BuildStrategy>("strategy",
-                                                         &strategy_);
+
+    AppendMultiDevPass(strategy);
 
     // Add a graph print pass to record a graph with device info.
     if (!strategy_.debug_graphviz_path_.empty()) {
@@ -115,6 +113,25 @@ class ParallelExecutorPassBuilder : public ir::PassBuilder {
     }
   }
 
+  // Convert graph to run on multi-devices.
+  void AppendMultiDevPass(const BuildStrategy &strategy) {
+    ir::Pass *multi_devices_pass;
+    if (strategy_.is_distribution_) {
+      multi_devices_pass = AppendPass("dist_multi_devices_pass").get();
+    } else {
+      if (strategy.reduce_ == BuildStrategy::ReduceStrategy::kAllReduce) {
+        multi_devices_pass =
+            AppendPass("allreduce_mode_multi_devices_pass").get();
+      } else if (strategy.reduce_ == BuildStrategy::ReduceStrategy::kReduce) {
+        multi_devices_pass = AppendPass("reduce_mode_multi_devices_pass").get();
+      } else {
+        PADDLE_THROW("Unknown reduce strategy.");
+      }
+    }
+    multi_devices_pass->SetNotOwned<const BuildStrategy>("strategy",
+                                                         &strategy_);
+  }
+
  private:
   BuildStrategy strategy_;
 };
@@ -131,6 +148,10 @@ std::shared_ptr<ir::PassBuilder> BuildStrategy::CreatePassesFromStrategy(
   return pass_builder_;
 }
 
+bool BuildStrategy::IsMultiDevPass(const std::string &pass_name) const {
+  return framework::details::MultiDevSSAGraphBuilder().count(pass_name) > 0;
+}
+
 std::unique_ptr<ir::Graph> BuildStrategy::Apply(
     const ProgramDesc &main_program, const std::vector<platform::Place> &places,
     const std::string &loss_var_name, const std::vector<Scope *> &local_scopes,
@@ -145,22 +166,23 @@ std::unique_ptr<ir::Graph> BuildStrategy::Apply(
 
   std::unique_ptr<ir::Graph> graph(new ir::Graph(main_program));
   for (std::shared_ptr<ir::Pass> &pass : pass_builder_->AllPasses()) {
-    if (pass->Type() == "multi_devices_pass") {
-      pass->Erase("places");
-      pass->SetNotOwned<const std::vector<platform::Place>>("places", &places);
-      pass->Erase("loss_var_name");
-      pass->SetNotOwned<const std::string>("loss_var_name", &loss_var_name);
-      pass->Erase("local_scopes");
-      pass->SetNotOwned<const std::vector<Scope *>>("local_scopes",
+    if (IsMultiDevPass(pass->Type())) {
+      pass->Erase(kPlaces);
+      pass->SetNotOwned<const std::vector<platform::Place>>(kPlaces, &places);
+      pass->Erase(kLossVarName);
+      pass->SetNotOwned<const std::string>(kLossVarName, &loss_var_name);
+      pass->Erase(kLocalScopes);
+      pass->SetNotOwned<const std::vector<Scope *>>(kLocalScopes,
                                                     &local_scopes);
-      pass->Erase("nranks");
-      pass->Set<size_t>("nranks", new size_t(nranks));
+      pass->Erase(kNRanks);
+      pass->Set<size_t>(kNRanks, new size_t(nranks));
 
 #if defined(PADDLE_WITH_CUDA) && !defined(_WIN32)
       platform::NCCLContextMap *nctx = use_cuda ? nccl_ctxs : nullptr;
       pass->Erase("nccl_ctxs");
       pass->SetNotOwned<platform::NCCLContextMap>("nccl_ctxs", nctx);
 #endif
+
     } else if (pass->Type() == "analysis_var_pass") {
       const std::vector<OpDesc *> *all_op_descs =
           new std::vector<OpDesc *>(main_program.Block(0).AllOps());
@@ -201,10 +223,13 @@ std::unique_ptr<ir::Graph> BuildStrategy::Apply(
 USE_PASS(fuse_elewise_add_act_pass);
 USE_PASS(graph_viz_pass);
 USE_PASS(multi_batch_merge_pass);
-USE_PASS(multi_devices_pass);
+USE_PASS(reduce_mode_multi_devices_pass);
+USE_PASS(allreduce_mode_multi_devices_pass);
+USE_PASS(dist_multi_devices_pass);
 USE_PASS(multi_devices_check_pass);
 USE_PASS(multi_devices_print_pass);
 USE_PASS(analysis_var_pass);
 USE_PASS(sequential_execution_pass);
 USE_PASS(all_reduce_deps_pass);
 USE_PASS(modify_op_lock_and_record_event_pass);
+USE_PASS(lock_free_optimize_pass);

paddle/fluid/framework/details/build_strategy.h

@@ -74,8 +74,6 @@ struct BuildStrategy {
 
   bool fuse_elewise_add_act_ops_{false};
 
-  bool enable_data_balance_{false};
-
   bool memory_optimize_{false};
 
   bool memory_early_delete_{false};
@@ -84,6 +82,10 @@ struct BuildStrategy {
 
   bool fuse_broadcast_op_{false};
 
+  // FIXME(zcd): is_distribution_ is a temporary field, because in pserver mode,
+  // num_trainers is 1, so the current fields of build_strategy doesn't tell if
+  // it's distributed model.
+  bool is_distribution_{false};
   int num_trainers_{1};
   int trainer_id_{0};
   std::vector<std::string> trainers_endpoints_;
@@ -104,6 +106,8 @@ struct BuildStrategy {
 
   bool IsFinalized() const { return is_finalized_; }
 
+  bool IsMultiDevPass(const std::string &pass_name) const;
+
   // Apply the passes built by the pass_builder_. The passes will be
   // applied to the Program and output an ir::Graph.
   std::unique_ptr<ir::Graph> Apply(const ProgramDesc &main_program,

paddle/fluid/framework/details/multi_devices_graph_check_pass.cc

@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/framework/details/multi_devices_graph_check_pass.h"
 #include <string>
+#include "paddle/fluid/framework/details/multi_devices_helper.h"
 #include "paddle/fluid/framework/ir/graph.h"
 #include "paddle/fluid/framework/ir/graph_helper.h"
 
@@ -21,68 +21,78 @@ namespace paddle {
 namespace framework {
 namespace details {
 
-bool SSAGraghBuilderWithChecker::IsValidGraph(const ir::Graph *graph) const {
-  std::unordered_map<OpHandleBase *, size_t> pending_ops;
-  std::unordered_set<VarHandleBase *> pending_vars;
-  std::unordered_set<VarHandleBase *> ready_vars;
-  std::unordered_set<OpHandleBase *> ready_ops;
+class SSAGraghBuilderWithChecker : public ir::Pass {
+ protected:
+  std::unique_ptr<ir::Graph> ApplyImpl(
+      std::unique_ptr<ir::Graph> graph) const override {
+    PADDLE_ENFORCE(IsValidGraph(graph.get()));
+    return graph;
+  }
 
-  auto insert_pending_var = [&](VarHandleBase *var) {
-    pending_vars.insert(var);
-    if (var->GeneratedOp() == nullptr) {
-      ready_vars.emplace(var);
-    }
-  };
+  bool IsValidGraph(const ir::Graph *graph) const {
+    std::unordered_map<OpHandleBase *, size_t> pending_ops;
+    std::unordered_set<VarHandleBase *> pending_vars;
+    std::unordered_set<VarHandleBase *> ready_vars;
+    std::unordered_set<OpHandleBase *> ready_ops;
 
-  for (auto &var_map : graph->Get<GraphVars>(kGraphVars)) {
-    for (auto &name_pair : var_map) {
-      for (auto &version_pair : name_pair.second) {
-        insert_pending_var(version_pair);
+    auto insert_pending_var = [&](VarHandleBase *var) {
+      pending_vars.insert(var);
+      if (var->GeneratedOp() == nullptr) {
+        ready_vars.emplace(var);
       }
-    }
-  }
+    };
 
-  for (auto &var : graph->Get<GraphDepVars>(kGraphDepVars)) {
-    insert_pending_var(var);
-  }
+    for (auto &var_map : graph->Get<GraphVars>(kGraphVars)) {
+      for (auto &name_pair : var_map) {
+        for (auto &version_pair : name_pair.second) {
+          insert_pending_var(version_pair);
+        }
+      }
+    }
 
-  for (OpHandleBase *op : ir::FilterByNodeWrapper<OpHandleBase>(*graph)) {
-    if (op->Inputs().empty()) {
-      ready_ops.insert(op);
-    } else {
-      pending_ops.insert({op, op->NoDupInputSize()});
+    for (auto &var : graph->Get<GraphDepVars>(kGraphDepVars)) {
+      insert_pending_var(var);
     }
-  }
 
-  auto run_all_ops = [&](std::unordered_set<OpHandleBase *> &set) {
-    for (auto *op : set) {
-      for (auto out : op->Outputs()) {
-        ready_vars.emplace(out);
+    for (OpHandleBase *op : ir::FilterByNodeWrapper<OpHandleBase>(*graph)) {
+      if (op->Inputs().empty()) {
+        ready_ops.insert(op);
+      } else {
+        pending_ops.insert({op, op->NoDupInputSize()});
       }
     }
-    set.clear();
-  };
 
-  while (!pending_vars.empty()) {
-    run_all_ops(ready_ops);
+    auto run_all_ops = [&](std::unordered_set<OpHandleBase *> &set) {
+      for (auto *op : set) {
+        for (auto out : op->Outputs()) {
+          ready_vars.emplace(out);
+        }
+      }
+      set.clear();
+    };
 
-    if (ready_vars.empty()) {
-      return false;
-    }
+    while (!pending_vars.empty()) {
+      run_all_ops(ready_ops);
 
-    for (auto ready_var : ready_vars) {
-      pending_vars.erase(ready_var);
-      for (auto *op : ready_var->PendingOps()) {
-        auto &deps = --pending_ops[op];
-        if (deps == 0) {
-          ready_ops.insert(op);
+      if (ready_vars.empty()) {
+        return false;
+      }
+
+      for (auto ready_var : ready_vars) {
+        pending_vars.erase(ready_var);
+        for (auto *op : ready_var->PendingOps()) {
+          auto &deps = --pending_ops[op];
+          if (deps == 0) {
+            ready_ops.insert(op);
+          }
         }
       }
+      ready_vars.clear();
     }
-    ready_vars.clear();
+    return true;
   }
-  return true;
-}
+};
+
 }  // namespace details
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/details/multi_devices_graph_pass.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include <string>
 #include <utility>
 #include <vector>
@@ -30,78 +31,70 @@ namespace framework {
 class Scope;
 namespace details {
 
-class MultiDevSSAGraphBuilder : public ir::Pass {
+constexpr char kLossVarName[] = "loss_var_name";
+constexpr char kPlaces[] = "places";
+constexpr char kLocalScopes[] = "local_scopes";
+constexpr char kStrategy[] = "strategy";
+constexpr char kNRanks[] = "nranks";
+
+class MultiDevSSAGraphBuilderBase : public ir::Pass {
  protected:
   std::unique_ptr<ir::Graph> ApplyImpl(
       std::unique_ptr<ir::Graph> graph) const override;
 
- private:
-  void CreateOpHandleIOs(ir::Graph *result, ir::Node *node,
-                         size_t device_id) const;
-  void Init() const;
+  virtual void Init() const;
 
-#if defined(PADDLE_WITH_CUDA) && !defined(_WIN32)
-  mutable platform::NCCLContextMap *nccl_ctxs_;
-#endif
+  virtual std::vector<ir::Node *> SortOperations(const ir::Graph &graph) const;
 
-  int GetVarDeviceID(
-      const std::string &varname,
-      const std::unordered_map<std::string, int> &sharded_var_device) const;
+  virtual void InsertCollectiveOp(ir::Graph *result, const std::string &p_name,
+                                  const std::string &g_name) const = 0;
 
-  bool IsScaleLossOp(ir::Node *node) const;
+  virtual bool DealWithSpecialOp(ir::Graph *result, ir::Node *node) const = 0;
+
+  virtual void InsertPostprocessOps(ir::Graph *result) const = 0;
 
-  int CreateRPCOp(
-      ir::Graph *result, ir::Node *node,
-      std::unordered_map<std::string, int> *sharded_var_device) const;
-  int CreateDistTrainOp(
-      ir::Graph *result, ir::Node *node,
-      std::unordered_map<std::string, int> *sharded_var_device) const;
+  bool UseGPU() const;
+
+  bool NeedCollectiveOps() const;
+
+  bool IsScaleLossOp(ir::Node *node) const;
 
   void CreateComputationalOps(ir::Graph *result, ir::Node *node,
                               size_t num_places) const;
 
   void CreateScaleLossGradOp(ir::Graph *result,
                              const std::string &loss_grad_name,
-                             ir::Node *out_var_node,
+                             ir::Node *out_var_node, size_t loss_scale,
                              proto::VarType::Type dtype) const;
 
   VarHandle *CreateReduceOp(ir::Graph *result, const std::string &og,
                             int dst_dev_id) const;
+
   void CreateComputationalOp(ir::Graph *result, ir::Node *node,
                              int dev_id) const;
 
-  int GetOpDeviceID(
-      ir::Node *node,
-      const std::unordered_map<std::string, int> &sharded_var_device) const;
-
-  void InsertAllReduceOp(ir::Graph *result, const std::string &og) const;
+  bool IsSparseGradient(const std::string &og) const;
 
-  void InsertDataBalanceOp(ir::Graph *result,
-                           const std::vector<std::string> &datas) const;
+  void CreateAllReduceOp(ir::Graph *result, const std::string &og) const;
 
   void CreateBroadcastOp(ir::Graph *result, const std::string &p_name,
                          size_t src_dev_id) const;
 
+  void InsertScaleLossGradOp(ir::Graph *result, const ir::Node *node) const;
+
   void CreateFusedBroadcastOp(
       ir::Graph *result,
       const std::vector<std::unordered_set<std::string>> &bcast_varnames) const;
 
-  bool IsSparseGradient(const std::string &og) const;
-
-  size_t GetAppropriateDeviceID(
-      const std::vector<std::string> &var_names) const;
-
   void SetCommunicationContext(OpHandleBase *op_handle,
                                const platform::Place &p) const;
 
-  std::vector<ir::Node *> SortForReduceMode(
-      const std::vector<ir::Node *> &) const;
+  void CreateOpHandleIOs(ir::Graph *result, ir::Node *node,
+                         size_t device_id) const;
 
-  int GetOpDeviceID(
-      ir::Node *node,
-      const std::unordered_map<std::string, int> &shared_var_device,
-      std::unordered_map<std::string, std::vector<ir::Node *>> *delay_ops)
-      const;
+#if defined(PADDLE_WITH_CUDA) && !defined(_WIN32)
+  mutable platform::NCCLContextMap *nccl_ctxs_;
+#endif
 
   mutable std::string loss_var_name_;
   mutable std::vector<platform::Place> places_;
@@ -109,8 +102,83 @@ class MultiDevSSAGraphBuilder : public ir::Pass {
 
   mutable BuildStrategy strategy_;
   mutable std::unordered_map<std::string, VarDesc *> all_vars_;
+};
+
+class AllReduceSSAGraphBuilder : public MultiDevSSAGraphBuilderBase {
+ protected:
+  virtual void InsertCollectiveOp(ir::Graph *result, const std::string &p_name,
+                                  const std::string &g_name) const;
+
+  virtual bool DealWithSpecialOp(ir::Graph *result, ir::Node *node) const {
+    return false;
+  }
+
+  virtual void InsertPostprocessOps(ir::Graph *result) const {}
+};
+
+class BalanceVarSSAGraphBuilder : public MultiDevSSAGraphBuilderBase {
+ protected:
+  int GetVarDeviceID(const std::string &varname) const;
+
+  int GetOpDeviceID(ir::Node *node) const;
+
+  size_t GetAppropriateDeviceID(
+      const std::vector<std::string> &var_names) const;
+
+  virtual void ResetState() const;
+
+  mutable std::unordered_map<std::string, int> sharded_var_device_;
   mutable std::vector<int64_t> balance_vars_;
 };
+
+class ReduceSSAGraphBuilder : public BalanceVarSSAGraphBuilder {
+ protected:
+  virtual void Init() const;
+
+  virtual void InsertCollectiveOp(ir::Graph *result, const std::string &p_name,
+                                  const std::string &g_name) const;
+
+  virtual bool DealWithSpecialOp(ir::Graph *result, ir::Node *node) const;
+
+  virtual void InsertPostprocessOps(ir::Graph *result) const;
+
+  virtual std::vector<ir::Node *> SortOperations(const ir::Graph &graph) const;
+
+  virtual void ResetState() const;
+
+  int GetOpDeviceID(ir::Node *node,
+                    std::unordered_map<std::string, std::vector<ir::Node *>>
+                        *delay_ops) const;
+
+  std::vector<ir::Node *> SortForReduceMode(
+      const std::vector<ir::Node *> &topo_ops) const;
+
+  mutable std::vector<std::unordered_set<std::string>> bcast_var_name_set_;
+};
+
+class DistSSAGraphBuilder : public BalanceVarSSAGraphBuilder {
+ protected:
+  virtual void Init() const;
+
+  virtual bool DealWithSpecialOp(ir::Graph *result, ir::Node *node) const;
+
+  virtual void InsertPostprocessOps(ir::Graph *result) const;
+
+  virtual void InsertCollectiveOp(ir::Graph *result, const std::string &p_name,
+                                  const std::string &g_name) const;
+
+  virtual void ResetState() const;
+
+  int CreateRPCOp(ir::Graph *result, ir::Node *node) const;
+
+  int CreateDistTrainOp(ir::Graph *result, ir::Node *node) const;
+
+  mutable std::vector<std::unordered_set<std::string>> bcast_var_name_set_;
+  mutable bool need_broadcast_var_{false};
+};
+
+std::unordered_set<std::string> &MultiDevSSAGraphBuilder();
+
 }  // namespace details
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/ir/CMakeLists.txt

@@ -31,6 +31,7 @@ cc_library(fuse_pass_base SRCS fuse_pass_base.cc DEPS pass)
 
 pass_library(graph_to_program_pass base)
 pass_library(graph_viz_pass base)
+pass_library(lock_free_optimize_pass base)
 pass_library(fc_fuse_pass inference)
 pass_library(attention_lstm_fuse_pass inference)
 pass_library(infer_clean_graph_pass inference)
@@ -41,11 +42,23 @@ pass_library(seq_concat_fc_fuse_pass inference)
 pass_library(multi_batch_merge_pass base)
 pass_library(conv_bn_fuse_pass inference)
 pass_library(seqconv_eltadd_relu_fuse_pass inference)
+pass_library(seqpool_concat_fuse_pass inference)
 pass_library(is_test_pass base)
 pass_library(conv_elementwise_add_act_fuse_pass inference)
 pass_library(conv_elementwise_add2_act_fuse_pass inference)
 pass_library(conv_elementwise_add_fuse_pass inference)
 pass_library(conv_affine_channel_fuse_pass inference)
+pass_library(transpose_flatten_concat_fuse_pass inference)
+
+# There may be many transpose-flatten structures in a model, and the output of
+# these structures will be used as inputs to the concat Op. This pattern will
+# be detected by our pass. The index here represents the number of structures in the
+# pattern. We use index 3 ~ 6, because these quantities of structures are
+# common in the models.
+foreach (index RANGE 3 6)
+   file(APPEND ${pass_file} "USE_PASS(transpose_flatten${index}_concat_fuse_pass);\n")
+endforeach()
+
 if(WITH_MKLDNN)
     pass_library(mkldnn_placement_pass base)
     pass_library(depthwise_conv_mkldnn_pass base)
@@ -67,6 +80,7 @@ cc_test(graph_helper_test SRCS graph_helper_test.cc DEPS graph graph_helper op_r
 cc_test(graph_to_program_pass_test SRCS graph_to_program_pass_test.cc DEPS graph_to_program_pass)
 cc_test(test_graph_pattern_detector SRCS graph_pattern_detector_tester.cc DEPS graph_pattern_detector)
 cc_test(test_fc_fuse_pass SRCS fc_fuse_pass_tester.cc DEPS fc_fuse_pass framework_proto)
+cc_test(test_seqpool_concat_fuse_pass SRCS seqpool_concat_fuse_pass_tester.cc DEPS seqpool_concat_fuse_pass framework_proto)
 cc_test(test_is_test_pass SRCS is_test_pass_tester.cc DEPS is_test_pass)
 if (WITH_MKLDNN)
     cc_test(test_depthwise_conv_mkldnn_pass SRCS depthwise_conv_mkldnn_pass_tester.cc DEPS depthwise_conv_mkldnn_pass)

paddle/fluid/framework/ir/graph.h

@@ -109,7 +109,6 @@ class Graph {
     attr_dels_[attr_name] = []() {};
   }
 
-  template <typename AttrType>
   void Erase(const std::string &attr_name) {
     PADDLE_ENFORCE(attrs_.count(attr_name) != 0, "%s not set in the graph",
                    attr_name);

paddle/fluid/framework/ir/graph_pattern_detector.cc

@@ -1306,6 +1306,69 @@ PDNode *patterns::ConvAffineChannel::operator()(
   return ac_out_var;
 }
 
+// a -> transpose_op(1) -> transpose_out_a -> flatten_op(1) -> flatten_out_a
+// b -> transpose_op(2) -> transpose_out_b -> flatten_op(2) -> flatten_out_b
+// ...
+// z -> transpose_op(n) -> transpose_out_z -> flatten_op(n) -> flatten_out_z
+// flatten_out_a -> concat_op  flatten_out_b -> concat_op ... flatten_out_z ->
+// concat_op
+PDNode *patterns::TransposeFlattenConcat::operator()(
+    std::vector<PDNode *> conv_in, int times) {
+  // The times represents the repeat times of the
+  // {trans, trans_out, flatten, flatten_out}
+  const int kNumFields = 4;
+  const int kTransOutOffset = 1;
+  const int kFlattenOffset = 2;
+  const int kFlattenOutOffset = 3;
+
+  std::vector<PDNode *> nodes;
+
+  for (int i = 0; i < times; i++) {
+    nodes.push_back(
+        pattern->NewNode(GetNodeName("transpose" + std::to_string(i)))
+            ->assert_is_op("transpose2"));
+    nodes.push_back(
+        pattern->NewNode(GetNodeName("transpose_out" + std::to_string(i)))
+            ->assert_is_op_output("transpose2")
+            ->assert_is_op_input("flatten2", "X")
+            ->AsIntermediate());
+    nodes.push_back(pattern->NewNode(GetNodeName("flatten" + std::to_string(i)))
+                        ->assert_is_op("flatten2"));
+
+    nodes.push_back(
+        pattern->NewNode(GetNodeName("flatten_out" + std::to_string(i)))
+            ->assert_is_op_output("flatten2")
+            ->assert_is_op_nth_input("concat", "X", i)
+            ->AsIntermediate());
+  }
+
+  auto concat_op = pattern->NewNode(GetNodeName("concat"))
+                       ->assert_is_op("concat")
+                       ->assert_op_has_n_inputs("concat", times);
+  auto concat_out = pattern->NewNode(GetNodeName("concat_out"))
+                        ->assert_is_op_output("concat")
+                        ->AsOutput();
+
+  std::vector<PDNode *> flatten_outs;
+  for (int i = 0; i < times; i++) {
+    conv_in[i]->AsInput();
+    // trans
+    nodes[i * kNumFields]->LinksFrom({conv_in[i]});
+    // trans_out
+    nodes[i * kNumFields + kTransOutOffset]->LinksFrom({nodes[i * kNumFields]});
+    // flatten
+    nodes[i * kNumFields + kFlattenOffset]->LinksFrom(
+        {nodes[i * kNumFields + kTransOutOffset]});
+    // flatten_out
+    nodes[i * kNumFields + kFlattenOutOffset]->LinksFrom(
+        {nodes[i * kNumFields + kFlattenOffset]});
+    flatten_outs.push_back(nodes[i * kNumFields + kFlattenOutOffset]);
+  }
+
+  concat_op->LinksFrom(flatten_outs).LinksTo({concat_out});
+  return concat_out;
+}
+
 }  // namespace ir
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/ir/graph_pattern_detector.h

@@ -766,6 +766,21 @@ struct ConvAffineChannel : public PatternBase {
   PATTERN_DECL_NODE(ac_out);  // Out
 };
 
+struct TransposeFlattenConcat : public PatternBase {
+  TransposeFlattenConcat(PDPattern* pattern, const std::string& name_scope)
+      : PatternBase(pattern, name_scope, "transpose_flatten_concat") {}
+
+  PDNode* operator()(std::vector<PDNode*> conv_inputs, int times);
+
+  std::string GetNodeName(const std::string& op_type) {
+    return PDNodeName(name_scope_, repr_, id_, op_type);
+  }
+
+  PDNode* GetPDNode(const std::string& op_type) {
+    return pattern->RetrieveNode(GetNodeName(op_type));
+  }
+};
+
 }  // namespace patterns
 
 // Link two ir::Nodes from each other.

paddle/fluid/framework/ir/lock_free_optimize_pass.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/framework/ir/lock_free_optimize_pass.h"
+
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "paddle/fluid/framework/ir/node.h"
+#include "paddle/fluid/framework/op_proto_maker.h"
+#include "paddle/fluid/framework/operator.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace framework {
+namespace ir {
+
+const char kSumGradOpName[] = "sum";
+// TODO(minqiyang): only support sgd at current time, please add
+// other optimizers later.
+const char kOptimizerType[] = "sgd";
+
+std::unique_ptr<ir::Graph> LockFreeOptimizePass::ApplyImpl(
+    std::unique_ptr<ir::Graph> graph) const {
+  PADDLE_ENFORCE(graph.get());
+
+  // We could collect all weights' name from SGD, where
+  // W1 <- SGD(W0, Grad0)
+  std::unordered_set<std::string> weight_var_set;
+  for (auto* node : graph->Nodes()) {
+    if (IsOpNamed(node, kOptimizerType)) {
+      auto& param_out_vars = node->Op()->Output("ParamOut");
+      PADDLE_ENFORCE(param_out_vars.size() == 1u);
+      weight_var_set.insert(param_out_vars[0]);
+    }
+  }
+
+  // find all grad's merge op via weight name, where
+  // Grad0 <- SUM(Grad1, Grad2, Grad3 ...)
+  std::unordered_set<ir::Node*> grad_sum_op_set;
+  for (ir::Node* node : graph->Nodes()) {
+    if (IsOpNamed(node, kSumGradOpName)) {
+      for (ir::Node* output : node->outputs) {
+        // strip the last grad suffix @GRAD
+        std::string var_name = output->Name();
+        const std::string suffix(kGradVarSuffix);
+        if (var_name != suffix && var_name.size() > suffix.size() &&
+            var_name.substr(var_name.size() - suffix.size()) == suffix) {
+          // if so then strip them off
+          var_name = var_name.substr(0, var_name.size() - suffix.size());
+          if (weight_var_set.find(var_name) != weight_var_set.end()) {
+            grad_sum_op_set.insert(node);
+            break;
+          }
+        }
+      }
+    }
+  }
+
+  // get the forward op and backward op pairs, where
+  // out <- forward(X, W)
+  // Grad1 <- backward(out, X')
+  // Grad0 <- SUM(Grad1, Grad2, Grad3 ...)
+  // W0 <- SGD(W1, Grad0)
+  for (ir::Node* node : grad_sum_op_set) {
+    for (ir::Node* merged_grad_var : node->outputs) {
+      // find the optimizers connected with sum op
+      if (IsVarNameEndsWith(merged_grad_var, kGradVarSuffix) &&
+          merged_grad_var->outputs.size() == 1u) {
+        ir::Node* opt_node = merged_grad_var->outputs[0];
+        VLOG(3) << "Found opt node " << opt_node->Name();
+
+        // find the backward op connected with sum op
+        for (ir::Node* unmerged_grad_var : node->inputs) {
+          if (IsVarNameContains(unmerged_grad_var, kGradVarSuffix) &&
+              unmerged_grad_var->inputs.size() == 1u) {
+            ir::Node* backward_op = unmerged_grad_var->inputs[0];
+
+            VLOG(3) << "Found backward_op " << backward_op->Name();
+
+            // find the forward op related to the backward op
+            ir::Node* forward_op =
+                FindForwardOpViaBackwardOp(graph.get(), backward_op);
+
+            VLOG(3) << "Found forward_op " << forward_op->Name();
+
+            PADDLE_ENFORCE(forward_op);
+
+            Node* new_optimizer_node = CreateNewSGDNode(
+                graph.get(), forward_op, backward_op, node, opt_node);
+
+            PADDLE_ENFORCE(new_optimizer_node);
+          }
+        }
+      }
+    }
+  }
+
+  // Remove the sum_op and its' outputs and connected Optimizers
+  for (Node* sum_op : grad_sum_op_set) {
+    for (Node* sum_op_output : sum_op->outputs) {
+      for (Node* optimize_op : sum_op_output->outputs) {
+        if (optimize_op->NodeType() == Node::Type::kOperation &&
+            optimize_op->Name() == kOptimizerType) {
+          VLOG(3) << "remove optimize_op: " << optimize_op->Name() << "_"
+                  << optimize_op->id();
+          graph->RemoveNode(optimize_op);
+        }
+      }
+      VLOG(3) << "remove sum_op_output: " << sum_op_output->Name() << "_"
+              << sum_op_output->id();
+      graph->RemoveNode(sum_op_output);
+    }
+    VLOG(3) << "remove sum_op: " << sum_op->Name() << "_" << sum_op->id();
+    graph->RemoveNode(sum_op);
+  }
+
+  for (auto* node : graph->Nodes()) {
+    for (Node* output_node : node->outputs) {
+      if (output_node->Name() == "sgd") {
+        VLOG(3) << "Node link to SGD: " << node->Name() << "_" << node->id()
+                << " --> " << output_node->Name() << "_" << output_node->id();
+        for (Node* input_node : node->inputs) {
+          VLOG(3) << "SGD Input link: " << input_node->Name() << "_"
+                  << input_node->id() << " --> " << node->Name() << "_"
+                  << node->id();
+        }
+      }
+    }
+  }
+
+  return graph;
+}
+
+ir::Node* LockFreeOptimizePass::CreateNewSGDNode(
+    ir::Graph* graph, ir::Node* forward_node, ir::Node* backward_node,
+    ir::Node* grad_sum_node, ir::Node* optimize_node) const {
+  PADDLE_ENFORCE(graph);
+  PADDLE_ENFORCE(forward_node);
+  PADDLE_ENFORCE(backward_node);
+  PADDLE_ENFORCE(grad_sum_node);
+  PADDLE_ENFORCE(optimize_node);
+
+  // find the grad var node between the grad sum node and backward_node
+  std::vector<ir::Node*> grad_vars =
+      FindConnectedNode(backward_node, grad_sum_node);
+  ir::Node* grad_node = nullptr;
+  for (ir::Node* node : grad_vars) {
+    if (!ir::IsControlDepVar(*node)) {
+      grad_node = node;
+    }
+  }
+  PADDLE_ENFORCE(grad_node);
+
+  // create a new SGD node
+  OpDesc* old_desc = optimize_node->Op();
+  // keep with the same block between new optimizer and the old one
+  OpDesc new_desc(*old_desc, old_desc->Block());
+  new_desc.SetInput("Param", old_desc->Input("Param"));
+  new_desc.SetInput("LearningRate", old_desc->Input("LearningRate"));
+  new_desc.SetInput("Grad", std::vector<std::string>({grad_node->Name()}));
+  new_desc.SetOutput("ParamOut", old_desc->Output("ParamOut"));
+
+  std::vector<std::string> op_role_vars = boost::get<std::vector<std::string>>(
+      new_desc.GetAttr(framework::OpProtoAndCheckerMaker::OpRoleVarAttrName()));
+  // replace the second op role var, because the grad name was
+  // changed in new optimizer
+  op_role_vars.pop_back();
+  op_role_vars.push_back(grad_node->Name());
+  new_desc.SetAttr(framework::OpProtoAndCheckerMaker::OpRoleVarAttrName(),
+                   op_role_vars);
+  new_desc.SetType(kOptimizerType);
+
+  // set backward op's op role var, this will be used to
+  // set device_id in multi_device_pass
+  backward_node->Op()->SetAttr(
+      framework::OpProtoAndCheckerMaker::OpRoleVarAttrName(), op_role_vars);
+  // backward_node->Op()->SetAttr(
+  // framework::OpProtoAndCheckerMaker::OpRoleVarAttrName(), {});
+
+  // keep with the same output nodes between new optimizer and the
+  // old one
+  Node* sgd_node = graph->CreateOpNode(&new_desc);
+
+  // change all outputs of the optimize_node to the new one
+  ReplaceAllDownstreamNode(optimize_node, sgd_node);
+
+  // find connected node between forward node and optimize node
+  // and replace the optimize node to new sgd node
+  std::vector<ir::Node*> forward_opt_connected_nodes =
+      FindConnectedNode(forward_node, optimize_node);
+  for (ir::Node* node : forward_opt_connected_nodes) {
+    ReplaceUpstreamNode(node, optimize_node, sgd_node);
+  }
+
+  // find connected node between backward node and optimize node
+  // and replace the optimize node to new sgd node
+  std::vector<ir::Node*> backward_opt_connected_nodes =
+      FindConnectedNode(backward_node, optimize_node);
+  for (ir::Node* node : backward_opt_connected_nodes) {
+    ReplaceUpstreamNode(node, optimize_node, sgd_node);
+  }
+
+  // SGD must have only one param and LR in
+  PADDLE_ENFORCE(old_desc->Input("LearningRate").size() == 1u);
+  PADDLE_ENFORCE(old_desc->Input("Param").size() == 1u);
+
+  // LR and weight nodes should be copied
+  for (Node* upstream_node : optimize_node->inputs) {
+    if (upstream_node->Name() == old_desc->Input("LearningRate")[0] ||
+        upstream_node->Name() == old_desc->Input("Param")[0]) {
+      ReplaceUpstreamNode(upstream_node, optimize_node, sgd_node);
+    }
+  }
+
+  VLOG(3) << "Create new opt node" << sgd_node->Name() << "_" << sgd_node->id();
+
+  return sgd_node;
+}
+
+std::vector<ir::Node*> LockFreeOptimizePass::FindConnectedNode(
+    ir::Node* upstream_node, ir::Node* downstream_node) const {
+  std::vector<ir::Node*> result;
+  for (ir::Node* out_node : upstream_node->outputs) {
+    for (ir::Node* in_node : downstream_node->inputs) {
+      if (in_node == out_node) {
+        result.push_back(in_node);
+      }
+    }
+  }
+
+  return result;
+}
+
+void LockFreeOptimizePass::ReplaceUpstreamNode(
+    ir::Node* upstream_node, ir::Node* old_optimizer_node,
+    ir::Node* new_optimizer_node) const {
+  PADDLE_ENFORCE(upstream_node);
+  PADDLE_ENFORCE(old_optimizer_node);
+  PADDLE_ENFORCE(new_optimizer_node);
+
+  // Remove the old_optimizer_node from upstream_node's outputs vector
+  auto& output_node_vec = upstream_node->outputs;
+  for (auto output_node_iter = output_node_vec.begin();
+       output_node_iter != output_node_vec.end();) {
+    if (*output_node_iter == old_optimizer_node) {
+      output_node_vec.erase(output_node_iter);
+      break;
+    } else {
+      ++output_node_iter;
+    }
+  }
+
+  // Add the new_optimizer_node to upstream_node's outputs vector
+  output_node_vec.emplace_back(new_optimizer_node);
+  new_optimizer_node->inputs.emplace_back(upstream_node);
+}
+
+void LockFreeOptimizePass::ReplaceAllDownstreamNode(
+    ir::Node* old_optimizer_node, ir::Node* new_optimizer_node) const {
+  PADDLE_ENFORCE(old_optimizer_node);
+  PADDLE_ENFORCE(new_optimizer_node);
+
+  for (ir::Node* downstream_node : old_optimizer_node->outputs) {
+    // Remove the old_optimizer_node from downstream_node's inputs vector
+    auto& input_node_vec = downstream_node->inputs;
+    for (auto input_node_iter = input_node_vec.begin();
+         input_node_iter != input_node_vec.end();) {
+      if (*input_node_iter == old_optimizer_node) {
+        input_node_vec.erase(input_node_iter);
+        break;
+      } else {
+        ++input_node_iter;
+      }
+    }
+
+    // Add the new_optimizer_node to downstream_node's inputs vector
+    input_node_vec.emplace_back(new_optimizer_node);
+    new_optimizer_node->outputs.emplace_back(downstream_node);
+  }
+}
+
+ir::Node* LockFreeOptimizePass::FindForwardOpViaBackwardOp(
+    ir::Graph* graph, ir::Node* backward_node) const {
+  PADDLE_ENFORCE(graph);
+  PADDLE_ENFORCE(backward_node);
+
+  // strip the suffix _grad of backward_node's name
+  std::string forward_op_name = backward_node->Name();
+  const std::string suffix("_grad");
+  if (forward_op_name != suffix && forward_op_name.size() > suffix.size() &&
+      forward_op_name.substr(forward_op_name.size() - suffix.size()) ==
+          suffix) {
+    // if so then strip them off
+    forward_op_name =
+        forward_op_name.substr(0, forward_op_name.size() - suffix.size());
+  } else {
+    LOG(WARNING) << "Illegal backward node's name " << backward_node->Name()
+                 << " id " << backward_node->id();
+
+    return nullptr;
+  }
+
+  for (ir::Node* node : graph->Nodes()) {
+    if (node->Name() == forward_op_name) {
+      if (node->outputs.size() == 0u) {
+        // if forward_node has no output, then it has NO grad op
+        continue;
+      }
+
+      // check whether all inputs of the backward_op that ends_with @GRAD
+      // comes from the output of forward_op is the input of the backward_op
+      bool is_related_forward_node = true;
+      for (ir::Node* backward_input : backward_node->inputs) {
+        if (IsVarNameEndsWith(backward_input, kGradVarSuffix)) {
+          bool meets_correct_output = false;
+          for (ir::Node* forward_output : node->outputs) {
+            if (forward_output->Name() + kGradVarSuffix ==
+                backward_input->Name()) {
+              meets_correct_output = true;
+              break;
+            }
+          }
+
+          if (!meets_correct_output) {
+            is_related_forward_node = false;
+            break;
+          }
+        }
+      }
+
+      if (is_related_forward_node) {
+        return node;
+      }
+    }
+  }
+
+  return nullptr;
+}
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle
+
+REGISTER_PASS(lock_free_optimize_pass,
+              paddle::framework::ir::LockFreeOptimizePass);

paddle/fluid/framework/ir/lock_free_optimize_pass.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.
+
+#ifndef PADDLE_FLUID_FRAMEWORK_IR_LOCK_FREE_OPTIMIZE_PASS_H_
+#define PADDLE_FLUID_FRAMEWORK_IR_LOCK_FREE_OPTIMIZE_PASS_H_
+
+#include <string>
+#include <vector>
+
+#include <boost/algorithm/string/predicate.hpp>
+
+#include "paddle/fluid/framework/ir/graph.h"
+#include "paddle/fluid/framework/ir/pass.h"
+
+namespace paddle {
+namespace framework {
+namespace ir {
+
+class Node;
+
+/*
+* Remove the sum op of all gradients of the backward op.
+* And remove the dependecies of the optimizer related to the
+* same backward op.
+*
+* Before this pass:
+*
+* forward_op1 forward_op2
+*     |            |
+*  grad_op1    grad_op2
+*        \      /
+*          \  /
+*         sum_op
+*           |
+*         sgd_op
+*
+* After this pass:
+* forward_op1 forward_op2
+*     |            |
+*  grad_op1    grad_op2
+*     |            |
+*  sgd_op1      sgd_op2
+*
+* sgd_op1 and sgd_op2 will update the same weight which holds the same
+* memory, so we could benefits from the acceleration
+*/
+class LockFreeOptimizePass : public Pass {
+ public:
+  virtual ~LockFreeOptimizePass() {}
+
+ protected:
+  std::unique_ptr<ir::Graph> ApplyImpl(std::unique_ptr<ir::Graph> graph) const;
+
+ private:
+  // Create a new sgd node via current optimizer node
+  ir::Node* CreateNewSGDNode(ir::Graph* graph, ir::Node* forward_node,
+                             ir::Node* backward_node, ir::Node* grad_sum_node,
+                             ir::Node* optimize_node) const;
+
+  // Replace the input weight's optimizers
+  void ReplaceUpstreamNode(ir::Node* upstream_node,
+                           ir::Node* old_optimizer_node,
+                           ir::Node* new_optimizer_node) const;
+
+  // Replace the output weight's optimizers
+  void ReplaceAllDownstreamNode(ir::Node* old_optimizer_node,
+                                ir::Node* new_optimizer_node) const;
+
+  // Find all weight variables in graph
+  bool FindAllWeightVars(ir::Graph* graph) const;
+
+  // Find the forward_op node via the backward_op node
+  ir::Node* FindForwardOpViaBackwardOp(ir::Graph* graph,
+                                       ir::Node* backward_node) const;
+
+  std::vector<ir::Node*> FindConnectedNode(ir::Node* upstream_node,
+                                           ir::Node* downstream_node) const;
+
+  inline bool IsOpNamed(ir::Node* node, const std::string& name) const {
+    PADDLE_ENFORCE(node);
+
+    return node->NodeType() == Node::Type::kOperation && node->Name() == name;
+  }
+
+  inline bool IsVarNamed(ir::Node* node, const std::string& name) const {
+    PADDLE_ENFORCE(node);
+
+    return node->NodeType() == Node::Type::kVariable && node->Name() == name;
+  }
+
+  inline bool IsVarNameEndsWith(ir::Node* node, const std::string& name) const {
+    PADDLE_ENFORCE(node);
+
+    return node->NodeType() == Node::Type::kVariable &&
+           boost::algorithm::ends_with(node->Name(), name);
+  }
+
+  inline bool IsVarNameContains(ir::Node* node, const std::string& name) const {
+    PADDLE_ENFORCE(node);
+
+    return node->NodeType() == Node::Type::kVariable &&
+           node->Name().find(name) != std::string::npos;
+  }
+
+  inline bool IsControlDepFrom(ir::Node* ctrl_dep_node, ir::Node* node) const {
+    PADDLE_ENFORCE(ctrl_dep_node);
+    PADDLE_ENFORCE(node);
+
+    return IsControlDepVar(*ctrl_dep_node) &&
+           ctrl_dep_node->inputs.size() >= 1u &&
+           ctrl_dep_node->inputs[0] == node;
+  }
+};
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle
+
+#endif  // PADDLE_FLUID_FRAMEWORK_IR_LOCK_FREE_OPTIMIZE_PASS_H_

paddle/fluid/framework/ir/seqpool_concat_fuse_pass.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/framework/ir/seqpool_concat_fuse_pass.h"
+#include <string>
+#include <vector>
+#include "paddle/fluid/framework/lod_tensor.h"
+
+#define MAX_CONCAT_INPUTS 200
+
+namespace paddle {
+namespace framework {
+namespace ir {
+
+PDNode* BuildSeqPoolConcatPattern(PDPattern* pattern,
+                                  const std::string& name_scope,
+                                  int num_inputs) {
+  auto is_concat_op_with_inputs = [](Node* x, int num) -> bool {
+    return x && x->IsOp() && x->Op()->Type() == "concat" &&
+           x->Op()->Input("X").size() == static_cast<size_t>(num);
+  };
+
+  auto is_nth_input_var_of_concat = [=](Node* x, int idx) -> bool {
+    return x && x->IsVar() && VarLinksToOp(x, "concat") &&
+           x->outputs.size() == 1 && IsNthInput(x, x->outputs[0], "X", idx) &&
+           is_concat_op_with_inputs(x->outputs[0], num_inputs);
+  };
+
+  auto is_seqpool_op_with_pootype_of_nth_input_of_concat = [=](
+      Node* x, const std::string& type, int idx) -> bool {
+    bool this_is_seqpool_op =
+        x && x->IsOp() && x->Op()->Type() == "sequence_pool" &&
+        x->Op()->HasAttr("pooltype") &&
+        boost::get<std::string>(x->Op()->GetAttr("pooltype")) == type &&
+        x->outputs.size() == 2;  // seqpool should only have 2 outputs
+    bool satisfied_all = this_is_seqpool_op;
+    if (this_is_seqpool_op) {
+      // Only one output of seqpool_op is nth_input_var of concat,
+      // the other one should be unused empty var.
+      if (is_nth_input_var_of_concat(x->outputs[0], idx)) {
+        satisfied_all = satisfied_all && x->outputs[1]->IsVar() &&
+                        x->outputs[1]->outputs.empty();
+      } else {
+        satisfied_all =
+            satisfied_all && is_nth_input_var_of_concat(x->outputs[1], idx) &&
+            x->outputs[0]->IsVar() && x->outputs[0]->outputs.size() == 0;
+      }
+    }
+    return satisfied_all;
+  };
+
+  auto* concat_op = pattern->NewNode(
+      [=](Node* x) { return is_concat_op_with_inputs(x, num_inputs); },
+      name_scope + "/concat_op");
+  concat_op->assert_op_attr<int>("axis", 1);
+
+  auto* concat_out_var = pattern->NewNode(
+      [=](Node* x) {
+        return x && x->IsVar() && VarLinksFromOp(x, "concat") &&
+               x->inputs.size() == 1 &&
+               is_concat_op_with_inputs(x->inputs[0], num_inputs);
+      },
+      name_scope + "/concat_out_var");
+  concat_out_var->assert_is_only_output_of_op("concat");
+
+  std::vector<PDNode*> seqpool_ops_input_var(num_inputs);
+  std::vector<PDNode*> seqpool_ops_output_var(num_inputs);
+  std::vector<PDNode*> seqpool_ops_output_unused_var(num_inputs);
+  std::vector<PDNode*> seqpool_ops(num_inputs);
+
+  for (int i = 0; i < num_inputs; ++i) {
+    seqpool_ops_output_var[i] = pattern->NewNode(
+        [=](Node* x) {
+          return x && x->IsVar() && is_nth_input_var_of_concat(x, i) &&
+                 x->inputs.size() == 1 &&
+                 is_seqpool_op_with_pootype_of_nth_input_of_concat(x->inputs[0],
+                                                                   "SUM", i);
+        },
+        name_scope + "/sequence_pool_out_" + std::to_string(i));
+
+    seqpool_ops_output_unused_var[i] = pattern->NewNode(
+        [=](Node* x) {
+          return x && x->IsVar() && x->inputs.size() == 1 &&
+                 x->outputs.size() == 0 &&
+                 is_seqpool_op_with_pootype_of_nth_input_of_concat(x->inputs[0],
+                                                                   "SUM", i);
+        },
+        name_scope + "/sequence_pool_unused_out_" + std::to_string(i));
+
+    seqpool_ops[i] = pattern->NewNode(
+        [=](Node* x) {
+          return x && x->IsOp() &&
+                 is_seqpool_op_with_pootype_of_nth_input_of_concat(x, "SUM", i);
+        },
+        name_scope + "/sequence_pool_op_" + std::to_string(i));
+
+    seqpool_ops_input_var[i] = pattern->NewNode(
+        [=](Node* x) {
+          bool basic = x && x->IsVar() && x->outputs.size() >= 1;
+          bool next_is_fine = false;
+          for (auto* o : x->outputs) {
+            if (is_seqpool_op_with_pootype_of_nth_input_of_concat(o, "SUM",
+                                                                  i)) {
+              next_is_fine = true;
+              break;
+            }
+          }
+          return basic && next_is_fine;
+        },
+        name_scope + "/sequence_pool_in_" + std::to_string(i));
+
+    // Links
+    seqpool_ops[i]
+        ->LinksFrom({seqpool_ops_input_var[i]})
+        .LinksTo({seqpool_ops_output_var[i], seqpool_ops_output_unused_var[i]});
+  }
+  concat_op->LinksFrom(seqpool_ops_output_var).LinksTo({concat_out_var});
+  return concat_out_var;
+}
+
+int BuildFusion(Graph* graph, const std::string& name_scope, int num_inputs) {
+  GraphPatternDetector gpd;
+  auto* pattern = gpd.mutable_pattern();
+  BuildSeqPoolConcatPattern(pattern, name_scope, num_inputs);
+
+  auto retrieve_node = [](const std::string& name,
+                          const GraphPatternDetector::subgraph_t& subgraph,
+                          const PDPattern& pat) -> Node* {
+    PADDLE_ENFORCE(subgraph.count(pat.RetrieveNode(name)),
+                   "pattern has no Node called %s", name.c_str());
+    Node* p = subgraph.at(pat.RetrieveNode(name));
+    PADDLE_ENFORCE_NOT_NULL(p, "subgraph has no node %s", name.c_str());
+    return p;
+  };
+
+  int fusion_count{0};
+  auto handler = [&](const GraphPatternDetector::subgraph_t& subgraph,
+                     Graph* g) {
+    VLOG(4) << "handle SeqPool Concat fuse";
+    std::vector<std::string> input_names(num_inputs);
+    std::vector<Node*> input_vars(num_inputs);
+    auto& fused_pattern = gpd.pattern();
+    for (int i = 0; i < num_inputs; ++i) {
+      input_vars[i] =
+          retrieve_node(name_scope + "/sequence_pool_in_" + std::to_string(i),
+                        subgraph, fused_pattern);
+      input_names[i] = input_vars[i]->Name();
+    }
+    auto* concat_op =
+        retrieve_node(name_scope + "/concat_op", subgraph, fused_pattern);
+    auto* concat_out_var =
+        retrieve_node(name_scope + "/concat_out_var", subgraph, fused_pattern);
+    auto* seqpool_op0 = retrieve_node(name_scope + "/sequence_pool_op_0",
+                                      subgraph, fused_pattern);
+
+    // Create New OpDesc
+    OpDesc op_desc;
+    op_desc.SetType("fusion_seqpool_concat");
+    op_desc.SetInput("X", input_names);
+    op_desc.SetAttr("pooltype", seqpool_op0->Op()->GetAttr("pooltype"));
+    op_desc.SetAttr("axis", concat_op->Op()->GetAttr("axis"));
+    op_desc.SetOutput("Out", {concat_out_var->Name()});
+    auto* op = graph->CreateOpNode(&op_desc);
+    for (size_t i = 0; i < input_vars.size(); ++i) {
+      IR_NODE_LINK_TO(input_vars[i], op);
+    }
+    IR_NODE_LINK_TO(op, concat_out_var);
+
+    std::unordered_set<const Node*> marked_nodes;
+    for (auto& item : subgraph) {
+      marked_nodes.insert(item.second);
+    }
+    for (size_t i = 0; i < input_vars.size(); ++i) {
+      marked_nodes.erase(input_vars[i]);
+    }
+    marked_nodes.erase(concat_out_var);
+    GraphSafeRemoveNodes(graph, marked_nodes);
+    ++fusion_count;
+  };
+
+  gpd(graph, handler);
+  return fusion_count;
+}
+
+std::unique_ptr<ir::Graph> SeqPoolConcatFusePass::ApplyImpl(
+    std::unique_ptr<ir::Graph> graph) const {
+  FusePassBase::Init(name_scope_, graph.get());
+  int fusion_count = 0;
+  for (int i = MAX_CONCAT_INPUTS; i > 0; --i) {
+    fusion_count +=
+        BuildFusion(graph.get(), name_scope_ + "/" + std::to_string(i), i);
+  }
+  AddStatis(fusion_count);
+
+  return graph;
+}
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle
+
+REGISTER_PASS(seqpool_concat_fuse_pass,
+              paddle::framework::ir::SeqPoolConcatFusePass);

paddle/fluid/framework/ir/seqpool_concat_fuse_pass.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 <string>
+#include "paddle/fluid/framework/ir/fuse_pass_base.h"
+#include "paddle/fluid/framework/ir/graph.h"
+#include "paddle/fluid/framework/ir/graph_pattern_detector.h"
+
+namespace paddle {
+namespace framework {
+namespace ir {
+
+/**
+ * Fuse SequencePool(with sum pooltype yet) and Concat;
+ *
+ * Before fuse:
+ *    |         |             |
+ * seq_pool, seq_pool, ... seq_pool
+ *    \         |      ...   /
+ *            concat
+ *              |
+ * After fuse:
+ *    \      |       /
+ *   FusionSeqPoolConcat
+ *           |
+ */
+class SeqPoolConcatFusePass : public FusePassBase {
+ public:
+  virtual ~SeqPoolConcatFusePass() {}
+
+ protected:
+  std::unique_ptr<ir::Graph> ApplyImpl(std::unique_ptr<ir::Graph> graph) const;
+
+  const std::string name_scope_{"seqpool_concat_fuse"};
+};
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle

paddle/fluid/framework/ir/seqpool_concat_fuse_pass_tester.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/framework/ir/seqpool_concat_fuse_pass.h"
+#include <gtest/gtest.h>
+#include "paddle/fluid/framework/op_proto_maker.h"
+
+namespace paddle {
+namespace framework {
+namespace ir {
+
+void SetOp(ProgramDesc* prog, const std::string& type,
+           const std::vector<std::string>& inputs,
+           const std::vector<std::string>& outputs) {
+  auto* op = prog->MutableBlock(0)->AppendOp();
+  op->SetType(type);
+  if (type == "sequence_pool") {
+    op->SetInput("X", {inputs[0]});
+    std::string pooltype = "SUM";
+    op->SetAttr("pooltype", pooltype);
+    op->SetOutput("MaxIndex", {outputs[0]});
+    op->SetOutput("Out", {outputs[1]});
+  } else if (type == "concat") {
+    op->SetInput("X", inputs);
+    op->SetAttr("axis", 1);
+    op->SetOutput("Out", {outputs[0]});
+  } else {
+    op->SetInput("X", inputs);
+    op->SetOutput("Out", outputs);
+  }
+  op->SetAttr(OpProtoAndCheckerMaker::OpRoleAttrName(),
+              static_cast<int>(OpRole::kForward));
+}
+
+int CountOpType(const ir::Graph* graph,
+                const std::string& op_type = "fusion_seqpool_concat") {
+  int count = 0;
+  for (auto* node : graph->Nodes()) {
+    if (node->IsOp() && node->Op()->Type() == op_type) {
+      ++count;
+    }
+  }
+  return count;
+}
+
+std::unique_ptr<ir::Graph> GetNumNodesOfBeforeAfter(
+    std::unique_ptr<ir::Graph> graph, int* before, int* after,
+    const std::string& pass_type = "seqpool_concat_fuse_pass") {
+  auto pass = PassRegistry::Instance().Get(pass_type);
+  *before = graph->Nodes().size();
+  graph = pass->Apply(std::move(graph));
+  *after = graph->Nodes().size();
+  return graph;
+}
+
+/*
+ * Before fuse:
+ *    a         b         c
+ *    |         |         |
+ *   op1       op2       op3
+ *   / \       / \       / \
+ *  d  e      f   g     h   i
+ *      \         |        /
+ *            concat
+ *              |
+ *              j
+ * Type of op1, op2 and op3 are sequence_pool, with "SUM" pooltype attr
+ *
+ * After fuse:
+ *    a         b         c
+ *    \         |        /
+ *    fusion_seqpool_concat
+ *              |
+ *              j
+ */
+TEST(SeqPoolConcatFusePass, basic) {
+  ProgramDesc prog;
+  for (auto& v : std::vector<std::string>(
+           {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"})) {
+    auto* var = prog.MutableBlock(0)->Var(v);
+    var->SetType(proto::VarType::LOD_TENSOR);
+  }
+
+  SetOp(&prog, "sequence_pool", std::vector<std::string>({"a"}),
+        std::vector<std::string>({"d", "e"}));
+  SetOp(&prog, "sequence_pool", std::vector<std::string>({"b"}),
+        std::vector<std::string>({"f", "g"}));
+  SetOp(&prog, "sequence_pool", std::vector<std::string>({"c"}),
+        std::vector<std::string>({"h", "i"}));
+  SetOp(&prog, "concat", std::vector<std::string>({"e", "g", "i"}),
+        std::vector<std::string>({"j"}));
+
+  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+  int before, after;
+  graph = GetNumNodesOfBeforeAfter(std::move(graph), &before, &after);
+  // Remove 10 Nodes: op1, op2, op3, d, e, f, g, h, i, concat_op
+  // Add 1 Node: fusion_seqpool_concat
+  EXPECT_EQ(after, before - 9);
+  EXPECT_EQ(CountOpType(graph.get()), 1);
+}
+
+/*
+ * Before fuse:
+ *    a            b
+ *    |           /  \
+ *   op1        op2  op3
+ *   / \        / \    \
+ *  c  d       e   f    g
+ *      \         /
+ *        concat
+ *          |
+ *          h
+ * Type of op1 and op2 are sequence_pool, with "SUM" pooltype attr
+ *
+ * After fuse:
+ *   a                         b
+ *    \                     /     \
+ *    fusion_seqpool_concat       op3
+ *              |                  |
+ *              h                  g
+ */
+TEST(SeqPoolConcatFusePass, advanced) {
+  ProgramDesc prog;
+  for (auto& v :
+       std::vector<std::string>({"a", "b", "c", "d", "e", "f", "g", "h"})) {
+    auto* var = prog.MutableBlock(0)->Var(v);
+    var->SetType(proto::VarType::LOD_TENSOR);
+  }
+
+  SetOp(&prog, "sequence_pool", std::vector<std::string>({"a"}),
+        std::vector<std::string>({"c", "d"}));
+  SetOp(&prog, "sequence_pool", std::vector<std::string>({"b"}),
+        std::vector<std::string>({"e", "f"}));
+  SetOp(&prog, "op3", std::vector<std::string>({"b"}),
+        std::vector<std::string>({"g"}));
+  SetOp(&prog, "concat", std::vector<std::string>({"d", "f"}),
+        std::vector<std::string>({"h"}));
+
+  std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+  int before, after;
+  graph = GetNumNodesOfBeforeAfter(std::move(graph), &before, &after);
+  // Remove 7 Nodes: op1, op2, c, d, e, f concat_op
+  // Add 1 Node: fusion_seqpool_concat
+  EXPECT_EQ(after, before - 6);
+  EXPECT_EQ(CountOpType(graph.get()), 1);
+}
+
+ProgramDesc BuildProgramDesc(int num_inputs_of_concat) {
+  ProgramDesc prog;
+  auto new_var = [&](const std::string& name) {
+    auto* var = prog.MutableBlock(0)->Var(name);
+    var->SetType(proto::VarType::LOD_TENSOR);
+  };
+  std::vector<std::string> concat_inputs;
+  for (int i = 0; i < num_inputs_of_concat; ++i) {
+    std::string prefix = "seqpool_op_" + i;
+    new_var(prefix + "in");
+    new_var(prefix + "out");
+    new_var(prefix + "out_unused");
+    SetOp(&prog, "sequence_pool", std::vector<std::string>({prefix + "in"}),
+          std::vector<std::string>({prefix + "out", prefix + "out_unused"}));
+    concat_inputs.push_back(prefix + "out");
+  }
+  SetOp(&prog, "concat", concat_inputs,
+        std::vector<std::string>({"concat_out"}));
+  return prog;
+}
+
+// test more inputs of concat
+TEST(SeqPoolConcatFusePass, more_inputs) {
+  for (int num : {1, 2, 10}) {
+    ProgramDesc prog = BuildProgramDesc(num);
+    std::unique_ptr<ir::Graph> graph(new ir::Graph(prog));
+    int before, after;
+    graph = GetNumNodesOfBeforeAfter(std::move(graph), &before, &after);
+    // Remove Nodes: n * (seqpool_op, out, out_unused), and concat_op
+    // Add Node: fusion_seqpool_concat op
+    EXPECT_EQ(after, before - num * 3);
+    EXPECT_EQ(CountOpType(graph.get()), 1);
+  }
+}
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle
+
+USE_PASS(seqpool_concat_fuse_pass);

paddle/fluid/framework/ir/transpose_flatten_concat_fuse_pass.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 <string>
+#include <vector>
+
+#include "paddle/fluid/framework/ir/graph_viz_pass.h"
+#include "paddle/fluid/framework/ir/node.h"
+#include "paddle/fluid/framework/ir/transpose_flatten_concat_fuse_pass.h"
+
+namespace paddle {
+namespace framework {
+namespace ir {
+
+template <int times>
+std::unique_ptr<ir::Graph> TransposeFlattenConcatFusePass<times>::ApplyImpl(
+    std::unique_ptr<ir::Graph> graph) const {
+  const std::string pattern_name =
+      "transpose_flatten" + std::to_string(times) + "_concat_fuse";
+  FusePassBase::Init(pattern_name, graph.get());
+
+  GraphPatternDetector gpd;
+  std::vector<PDNode *> input_nodes;
+  for (int i = 0; i < times; i++) {
+    input_nodes.push_back(gpd.mutable_pattern()
+                              ->NewNode("x" + std::to_string(i))
+                              ->assert_is_op_input("transpose2", "X")
+                              ->AsInput());
+  }
+
+  patterns::TransposeFlattenConcat pattern(gpd.mutable_pattern(), pattern_name);
+  pattern(input_nodes, times);
+
+  auto handler = [&](const GraphPatternDetector::subgraph_t &subgraph,
+                     Graph *g) {
+    const int kNumFields = 5;
+    const int kTransOffset = 1;
+    const int kTransOutOffset = 2;
+    const int kFlattenOffset = 3;
+    const int kFlattenOutOffset = 4;
+    std::vector<Node *> nodes;
+
+    for (int i = 0; i < times; i++) {
+      PADDLE_ENFORCE(
+          subgraph.at(pattern.GetPDNode("transpose" + std::to_string(i))));
+      PADDLE_ENFORCE(
+          subgraph.at(pattern.GetPDNode("transpose_out" + std::to_string(i))));
+      PADDLE_ENFORCE(
+          subgraph.at(pattern.GetPDNode("flatten" + std::to_string(i))));
+      PADDLE_ENFORCE(
+          subgraph.at(pattern.GetPDNode("flatten_out" + std::to_string(i))));
+      PADDLE_ENFORCE(subgraph.at(input_nodes[i]));
+
+      nodes.push_back(subgraph.at(input_nodes[i]));
+      nodes.push_back(
+          subgraph.at(pattern.GetPDNode("transpose" + std::to_string(i))));
+      nodes.push_back(
+          subgraph.at(pattern.GetPDNode("transpose_out" + std::to_string(i))));
+      nodes.push_back(
+          subgraph.at(pattern.GetPDNode("flatten" + std::to_string(i))));
+      nodes.push_back(
+          subgraph.at(pattern.GetPDNode("flatten_out" + std::to_string(i))));
+    }
+
+    Node *concat_op = subgraph.at(pattern.GetPDNode("concat"));
+    Node *concat_out = subgraph.at(pattern.GetPDNode("concat_out"));
+    std::vector<std::string> input_names;
+    std::vector<int> trans_axis = boost::get<std::vector<int>>(
+        nodes[kTransOffset]->Op()->GetAttr("axis"));
+    int flatten_axis =
+        boost::get<int>(nodes[kFlattenOffset]->Op()->GetAttr("axis"));
+    int concat_axis = boost::get<int>(concat_op->Op()->GetAttr("axis"));
+    std::string output_name = concat_out->Name();
+
+    for (int i = 0; i < times; i++) {
+      input_names.push_back(nodes[i * kNumFields]->Name());
+    }
+
+    framework::OpDesc new_op_desc;
+    new_op_desc.SetType("fusion_transpose_flatten_concat");
+    new_op_desc.SetInput("X", input_names);
+    new_op_desc.SetAttr("trans_axis", trans_axis);
+    new_op_desc.SetAttr("flatten_axis", flatten_axis);
+    new_op_desc.SetAttr("concat_axis", concat_axis);
+    new_op_desc.SetOutput("Out", {output_name});
+    new_op_desc.Flush();
+
+    // Create a new node for the fused op.
+    auto *new_conv_op = graph->CreateOpNode(&new_op_desc);
+
+    std::unordered_set<const Node *> delete_nodes;
+
+    for (int i = 0; i < times; i++) {
+      nodes[i * kNumFields]->outputs.push_back(new_conv_op);
+      new_conv_op->inputs.push_back(nodes[i * kNumFields]);
+      delete_nodes.insert(nodes[i * kNumFields + kTransOffset]);
+      delete_nodes.insert(nodes[i * kNumFields + kTransOutOffset]);
+      delete_nodes.insert(nodes[i * kNumFields + kFlattenOffset]);
+      delete_nodes.insert(nodes[i * kNumFields + kFlattenOutOffset]);
+    }
+    delete_nodes.insert(concat_op);
+
+    new_conv_op->outputs.push_back(concat_out);
+    concat_out->inputs.push_back(new_conv_op);
+
+    // Delete the unneeded nodes.
+    GraphSafeRemoveNodes(graph.get(), delete_nodes);
+  };
+
+  gpd(graph.get(), handler);
+  return graph;
+}
+
+template class TransposeFlattenConcatFusePass<1>;
+template class TransposeFlattenConcatFusePass<3>;
+template class TransposeFlattenConcatFusePass<4>;
+template class TransposeFlattenConcatFusePass<5>;
+template class TransposeFlattenConcatFusePass<6>;
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle
+
+REGISTER_PASS(transpose_flatten_concat_fuse_pass,
+              paddle::framework::ir::TransposeFlattenConcatFusePass<1>);
+
+REGISTER_PASS(transpose_flatten3_concat_fuse_pass,
+              paddle::framework::ir::TransposeFlattenConcatFusePass<3>);
+
+REGISTER_PASS(transpose_flatten4_concat_fuse_pass,
+              paddle::framework::ir::TransposeFlattenConcatFusePass<4>);
+
+REGISTER_PASS(transpose_flatten5_concat_fuse_pass,
+              paddle::framework::ir::TransposeFlattenConcatFusePass<5>);
+
+REGISTER_PASS(transpose_flatten6_concat_fuse_pass,
+              paddle::framework::ir::TransposeFlattenConcatFusePass<6>);

paddle/fluid/framework/ir/transpose_flatten_concat_fuse_pass.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/ir/fuse_pass_base.h"
+#include "paddle/fluid/framework/ir/graph_pattern_detector.h"
+
+namespace paddle {
+namespace framework {
+namespace ir {
+
+// There may be many transpose-flatten structures in a model, and the output of
+// these structures will be used as inputs to the concat Op. This pattern will
+// be detected by our pass. The times here represents the repeat times of this
+// structure.
+template <int times>
+class TransposeFlattenConcatFusePass : public FusePassBase {
+ public:
+  virtual ~TransposeFlattenConcatFusePass() {}
+
+ protected:
+  std::unique_ptr<ir::Graph> ApplyImpl(std::unique_ptr<ir::Graph> graph) const;
+};
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle

paddle/fluid/framework/naive_executor.cc

@@ -40,14 +40,14 @@ void NaiveExecutor::Prepare(Scope *scope, const ProgramDesc &program_desc,
 
 void NaiveExecutor::Run() {
 #ifndef PADDLE_ON_INFERENCE
-  LOG_FIRST_N(WARNING, 15) << "The NaiveExecutor can not work properly if the "
-                              "cmake flag ON_INFER is not set.";
-  LOG_FIRST_N(WARNING, 15) << "Unlike the training phase, all the scopes and "
-                              "variables will be reused to save the allocation "
-                              "overhead.";
-  LOG_FIRST_N(WARNING, 15) << "Please re-compile the inference library by "
-                              "setting the cmake flag ON_INFER=ON if you are "
-                              "running Paddle Inference";
+  LOG_FIRST_N(WARNING, 5) << "The NaiveExecutor can not work properly if the "
+                             "cmake flag ON_INFER is not set.";
+  LOG_FIRST_N(WARNING, 5) << "Unlike the training phase, all the scopes and "
+                             "variables will be reused to save the allocation "
+                             "overhead.";
+  LOG_FIRST_N(WARNING, 5) << "Please re-compile the inference library by "
+                             "setting the cmake flag ON_INFER=ON if you are "
+                             "running Paddle Inference";
 #endif  // PADDLE_ON_INFERENCE
   for (auto &op : ops_) {
     VLOG(3) << std::this_thread::get_id() << " run " << op->Type()

paddle/fluid/framework/ngraph_bridge.cc

@@ -32,8 +32,11 @@ std::map<std::string,
                                 std::string, std::shared_ptr<ngraph::Node>>>)>>
     NgraphBridge::NG_NODE_MAP = {
         {"fill_constant", paddle::operators::ngraphs::BuildFillConstantNode},
+        {"mean", paddle::operators::ngraphs::BuildMeanNode},
+        {"mean_grad", paddle::operators::ngraphs::BuildMeanGradNode},
         {"mul", paddle::operators::ngraphs::BuildMulNode},
         {"mul_grad", paddle::operators::ngraphs::BuildMulGradNode},
+        {"scale", paddle::operators::ngraphs::BuildScaleNode},
         {"relu", paddle::operators::ngraphs::BuildUnaryNode<ngraph::op::Relu>},
         {"tanh", paddle::operators::ngraphs::BuildUnaryNode<ngraph::op::Tanh>},
         {"top_k", paddle::operators::ngraphs::BuildTopKNode}};

paddle/fluid/framework/ngraph_operator.cc

@@ -539,7 +539,7 @@ void NgraphEngine::Run(const Scope& scope, const platform::Place& place) const {
     }
   }
 
-  backend_->call(ngraph_function_, t_out, t_in);
+  backend_->call(backend_->compile(ngraph_function_), t_out, t_in);
 }  // NgraphEngine::RunImpl
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/operator.h

@@ -393,7 +393,7 @@ class ExecutionContext {
     PADDLE_ENFORCE(
         dynamic_cast<platform::TemporaryAllocation*>(allocation_ptr) != nullptr,
         "The AllocationPtr must be TemporaryAllocation.");
-    PADDLE_ENFORCE_EQ(allocation_ptr->size(),
+    PADDLE_ENFORCE_GE(allocation_ptr->size(),
                       framework::product(dim) * sizeof(T));
 
     paddle::framework::Tensor temp_tensor(

paddle/fluid/framework/parallel_executor.cc

@@ -193,15 +193,14 @@ ParallelExecutor::ParallelExecutor(
     const std::unordered_set<std::string> &bcast_vars,
     const ProgramDesc &main_program, const std::string &loss_var_name,
     Scope *scope, const std::vector<Scope *> &local_scopes,
-    const ExecutionStrategy &exec_strategy, const BuildStrategy &build_strategy,
-    size_t num_trainers, size_t trainer_id)
+    const ExecutionStrategy &exec_strategy, const BuildStrategy &build_strategy)
     : member_(new ParallelExecutorPrivate(places)) {
   member_->global_scope_ = scope;
   member_->use_cuda_ = exec_strategy.use_cuda_;
   member_->build_strategy_ = build_strategy;
   member_->use_all_reduce_ =
       build_strategy.reduce_ == BuildStrategy::ReduceStrategy::kAllReduce;
-  member_->nranks_ = num_trainers * places.size();
+  member_->nranks_ = build_strategy.num_trainers_ * places.size();
 
   if (!member_->use_all_reduce_) {
     PADDLE_ENFORCE(places.size() > 1,
@@ -253,7 +252,8 @@ ParallelExecutor::ParallelExecutor(
     }
 
     member_->nccl_ctxs_.reset(new platform::NCCLContextMap(
-        member_->places_, nccl_id, num_trainers, trainer_id));
+        member_->places_, nccl_id, build_strategy.num_trainers_,
+        build_strategy.trainer_id_));
 #else
     PADDLE_THROW("Not compiled with CUDA");
 #endif

paddle/fluid/framework/parallel_executor.h

@@ -50,8 +50,7 @@ class ParallelExecutor {
                             const std::string &loss_var_name, Scope *scope,
                             const std::vector<Scope *> &local_scopes,
                             const ExecutionStrategy &exec_strategy,
-                            const BuildStrategy &build_strategy,
-                            size_t num_trainers = 1, size_t trainer_id = 0);
+                            const BuildStrategy &build_strategy);
 
   ~ParallelExecutor();
 

paddle/fluid/framework/python_headers.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
+
+// workaround for Python 2 issue: https://bugs.python.org/issue17120
+#pragma push_macro("_XOPEN_SOURCE")
+#pragma push_macro("_POSIX_C_SOURCE")
+#undef _XOPEN_SOURCE
+#undef _POSIX_C_SOURCE
+
+#include "pybind11/pybind11.h"
+
+#pragma pop_macro("_XOPEN_SOURCE")
+#pragma pop_macro("_POSIX_C_SOURCE")

paddle/fluid/framework/scope.cc

@@ -87,11 +87,12 @@ Variable* Scope::Var(const std::string& name) {
 }
 
 Variable* Scope::Var(std::string* name) {
-  auto new_name = string::Sprintf("%p.%d", this, vars_.size());
+  SCOPE_VARS_WRITER_LOCK
+  auto new_name = std::to_string(reinterpret_cast<uintptr_t>(this)) + "." +
+                  std::to_string(vars_.size());
   if (name != nullptr) {
     *name = new_name;
   }
-  SCOPE_VARS_WRITER_LOCK
   return VarInternal(new_name);
 }
 

paddle/fluid/framework/var_type_traits.cc

@@ -105,13 +105,15 @@ struct VarIdToTypeIndexMapHolder {
 
 }  // namespace detail
 
-const std::type_index &ToTypeIndex(int var_id) {
+const std::type_index &VarTraitIdToTypeIndex(int var_id) {
   return detail::VarIdToTypeIndexMapHolder::ToTypeIndex(var_id);
 }
 
-const char *ToTypeName(int var_id) { return ToTypeIndex(var_id).name(); }
+const char *ToTypeName(int var_id) {
+  return VarTraitIdToTypeIndex(var_id).name();
+}
 
-int ToTypeId(const std::type_index &type) {
+int TypeIndexToVarTraitId(const std::type_index &type) {
   return detail::VarIdToTypeIndexMapHolder::ToTypeId(type);
 }
 

paddle/fluid/framework/var_type_traits.h

@@ -66,8 +66,8 @@ namespace paddle {
 namespace framework {
 
 const char *ToTypeName(int var_id);
-const std::type_index &ToTypeIndex(int var_id);
-int ToTypeId(const std::type_index &type);
+const std::type_index &VarTraitIdToTypeIndex(int var_id);
+int TypeIndexToVarTraitId(const std::type_index &type);
 
 namespace detail {
 

paddle/fluid/framework/var_type_traits_test.cc

@@ -45,10 +45,11 @@ struct TypeIndexChecker {
     constexpr auto kId = VarTypeTrait<Type>::kId;
     std::type_index actual_type(typeid(Type));
     EXPECT_EQ(std::string(ToTypeName(kId)), std::string(actual_type.name()));
-    EXPECT_EQ(ToTypeIndex(kId), actual_type);
-    EXPECT_EQ(ToTypeId(actual_type), kId);
-    EXPECT_EQ(ToTypeIndex(ToTypeId(actual_type)), actual_type);
-    EXPECT_EQ(ToTypeId(ToTypeIndex(kId)), kId);
+    EXPECT_EQ(VarTraitIdToTypeIndex(kId), actual_type);
+    EXPECT_EQ(TypeIndexToVarTraitId(actual_type), kId);
+    EXPECT_EQ(VarTraitIdToTypeIndex(TypeIndexToVarTraitId(actual_type)),
+              actual_type);
+    EXPECT_EQ(TypeIndexToVarTraitId(VarTraitIdToTypeIndex(kId)), kId);
 
     EXPECT_TRUE(var_id_set->count(kId) == 0);              // NOLINT
     EXPECT_TRUE(type_index_set->count(actual_type) == 0);  // NOLINT

paddle/fluid/imperative/layer.cc

@@ -27,6 +27,8 @@
 namespace paddle {
 namespace imperative {
 
+std::map<int, py::object> py_funcs_;
+
 using framework::Variable;
 
 void AddTo(Variable* src, Variable* dst) {
@@ -42,7 +44,7 @@ void AddTo(Variable* src, Variable* dst) {
                  src_tensor->numel());
   float* dst_data = dst_tensor->mutable_data<float>(platform::CPUPlace());
   const float* src_data = src_tensor->data<float>();
-  for (size_t i = 0; i < src_tensor->numel(); ++i) {
+  for (int64_t i = 0; i < src_tensor->numel(); ++i) {
     dst_data[i] += src_data[i];
   }
 }
@@ -55,6 +57,7 @@ class Autograd {
     if (var->stop_gradient_) {
       return;
     }
+    VLOG(3) << "start autograd";
 
     std::deque<OpBase*> ready;
     ready.push_back(var->pre_op_);
@@ -114,57 +117,63 @@ class Autograd {
   }
 };
 
-framework::LoDTensor& VarBase::Grad() {
+framework::LoDTensor& VarBase::GradValue() {
   VLOG(3) << "get var grad " << var_desc_->Name();
-  return *grads_->GetMutable<framework::LoDTensor>();
+  return *(grads_->var_->GetMutable<framework::LoDTensor>());
 }
 
 std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad() {
-  if (!grad_op_desc_) {
+  if (!grad_op_desc_ && backward_id_ <= 0) {
     LOG(WARNING) << "op with no grad: " << op_desc_->Type();
     return {};
   }
-  VLOG(3) << "op grad " << grad_op_desc_->Type();
 
-  std::vector<std::unique_ptr<framework::Variable>> tmp_vars;
   std::map<std::string, std::vector<framework::Variable*>> grad_outputs;
-  for (auto it : grad_output_vars_) {
-    auto& outputs = grad_outputs[it.first];
-    for (size_t i = 0; i < it.second.size(); ++i) {
-      // Allocate a new variable
-      Variable* tmp_var = new framework::Variable();
-      tmp_var->GetMutable<framework::LoDTensor>();
-
-      tmp_vars.emplace_back(tmp_var);
-      outputs.push_back(tmp_var);
+  if (backward_id_ > 0) {
+    VLOG(3) << "py_layer_grad";
+    grad_outputs["Out@GRAD"] =
+        PyLayer::ApplyGrad(backward_id_, grad_input_vars_["X@GRAD"]);
+  } else {
+    VLOG(3) << "op grad " << grad_op_desc_->Type();
+    for (auto it : grad_output_vars_) {
+      auto& outputs = grad_outputs[it.first];
+      for (size_t i = 0; i < it.second.size(); ++i) {
+        // Allocate a new variable
+        Variable* tmp_var = new framework::Variable();
+        tmp_var->GetMutable<framework::LoDTensor>();
+        outputs.push_back(tmp_var);
+      }
     }
-  }
 
-  framework::RuntimeContext ctx(grad_input_vars_, grad_outputs);
+    framework::RuntimeContext ctx(grad_input_vars_, grad_outputs);
 
-  // No need to do compile time infer shape here.
-  // grad_op_desc_->InferShape(*block_);
-  grad_op_desc_->InferVarType(block_);
+    // No need to do compile time infer shape here.
+    // grad_op_desc_->InferShape(*block_);
+    grad_op_desc_->InferVarType(block_);
 
-  std::unique_ptr<framework::OperatorBase> opbase =
-      framework::OpRegistry::CreateOp(*grad_op_desc_);
-  framework::OperatorWithKernel* op_kernel =
-      dynamic_cast<framework::OperatorWithKernel*>(opbase.get());
-  PADDLE_ENFORCE_NOT_NULL(op_kernel, "only support op with kernel");
+    std::unique_ptr<framework::OperatorBase> opbase =
+        framework::OpRegistry::CreateOp(*grad_op_desc_);
+    framework::OperatorWithKernel* op_kernel =
+        dynamic_cast<framework::OperatorWithKernel*>(opbase.get());
+    PADDLE_ENFORCE_NOT_NULL(op_kernel, "only support op with kernel");
 
-  framework::Scope scope;
-  platform::CPUPlace place;
-  PreparedOp p = PreparedOp::Prepare(ctx, *op_kernel, place);
-  p.op.RuntimeInferShape(scope, place, ctx);
-  p.func(framework::ExecutionContext(p.op, scope, *p.dev_ctx, p.ctx));
+    framework::Scope scope;
+    platform::CPUPlace place;
+    PreparedOp p = PreparedOp::Prepare(ctx, *op_kernel, place);
+    p.op.RuntimeInferShape(scope, place, ctx);
+    p.func(framework::ExecutionContext(p.op, scope, *p.dev_ctx, p.ctx));
+  }
 
   for (auto it : grad_output_vars_) {
     auto& outputs = grad_outputs[it.first];
     auto& origin_outputs = it.second;
+    PADDLE_ENFORCE_EQ(outputs.size(), origin_outputs.size());
 
     for (size_t i = 0; i < outputs.size(); ++i) {
+      framework::Variable* grad = outputs[i];
       framework::Variable* orig_grad = origin_outputs[i];
-      AddTo(outputs[i], orig_grad);
+      AddTo(grad, orig_grad);
+      delete grad;
     }
   }
   return input_vars_;
@@ -173,7 +182,8 @@ std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad() {
 void VarBase::RunBackward() {
   if (!pre_op_) return;
 
-  auto grads_t = grads_->GetMutable<framework::LoDTensor>();
+  VLOG(3) << "start backward";
+  auto grads_t = grads_->var_->GetMutable<framework::LoDTensor>();
   float* data = grads_t->mutable_data<float>(platform::CPUPlace());
   std::fill(data, data + grads_t->numel(), 1.0);
 
@@ -183,5 +193,65 @@ void VarBase::RunBackward() {
   Autograd().RunBackward(this);
 }
 
+void PyLayer::RegisterFunc(int func_id, const py::object& py_func) {
+  py_funcs_[func_id] = py_func;
+}
+
+int PyLayer::NumFuncs() { return py_funcs_.size(); }
+
+std::vector<VarBase*> PyLayer::Apply(int func_id,
+                                     const std::vector<VarBase*>& inputs) {
+  std::vector<framework::Variable*> invars;
+  for (const VarBase* in : inputs) {
+    invars.push_back(in->var_);
+  }
+  PADDLE_ENFORCE(py_funcs_.find(func_id) != py_funcs_.end());
+  std::vector<Variable*> outvars = CallPythonFunc(py_funcs_[func_id], invars);
+  std::vector<VarBase*> ret;
+  for (Variable* v : outvars) {
+    ret.push_back(new VarBase(v, new VarBase(true)));
+  }
+  return ret;
+}
+
+std::vector<Variable*> PyLayer::ApplyGrad(
+    int func_id, const std::vector<framework::Variable*>& inputs) {
+  PADDLE_ENFORCE(py_funcs_.find(func_id) != py_funcs_.end());
+  return CallPythonFunc(py_funcs_[func_id], inputs);
+}
+
+std::vector<framework::Variable*> PyLayer::CallPythonFunc(
+    const py::object& callable, const std::vector<framework::Variable*>& ins) {
+  py::gil_scoped_acquire guard;
+  py::tuple in_args(ins.size());
+  for (size_t i = 0; i < ins.size(); ++i) {
+    const framework::LoDTensor& t = ins[i]->Get<framework::LoDTensor>();
+    in_args[i] = t.IsInitialized() ? py::cast(t) : py::cast(nullptr);
+  }
+  VLOG(3) << "pyfunc in " << py::len(in_args);
+
+  // TODO(panyx0718): Who owns the returned LoDTensor.
+  auto ret = callable(in_args);
+  auto ret_tuple = py::cast<py::tuple>(ret);
+  size_t ret_num = py::len(ret_tuple);
+  std::vector<framework::Variable*> outs;
+  VLOG(3) << "pyfunc out " << ret_num;
+  for (size_t i = 0; i < ret_num; ++i) {
+    try {
+      auto* py_out_tensor = py::cast<framework::LoDTensor*>(ret_tuple[i]);
+      PADDLE_ENFORCE_NOT_NULL(py_out_tensor,
+                              "Output tensor %d should not be nullptr", i);
+      auto* var = new framework::Variable();
+      auto* tensor = var->GetMutable<framework::LoDTensor>();
+      tensor->ShareDataWith(*py_out_tensor);
+      tensor->set_lod(py_out_tensor->lod());
+      outs.push_back(var);
+    } catch (py::cast_error&) {
+      PADDLE_THROW("The %d-th output must be LoDTensor", i);
+    }
+  }
+  return outs;
+}
+
 }  // namespace imperative
 }  // namespace paddle

paddle/fluid/imperative/layer.h

@@ -14,17 +14,26 @@
 
 #pragma once
 
-#include <map>
-#include <string>
-#include <vector>
+// clang-format off
+#include "paddle/fluid/framework/python_headers.h"
+// clang-format on
+
+#include <map>     // NOLINT
+#include <string>  // NOLINT
+#include <vector>  // NOLINT
+
 #include "paddle/fluid/framework/op_desc.h"
 #include "paddle/fluid/framework/operator.h"
 #include "paddle/fluid/framework/var_desc.h"
 #include "paddle/fluid/platform/enforce.h"
 
+#include "paddle/fluid/imperative/type_defs.h"
+
 namespace paddle {
 namespace imperative {
 
+namespace py = ::pybind11;
+
 class PreparedOp {
  public:
   PreparedOp(const framework::OperatorBase& op,
@@ -77,31 +86,50 @@ class PreparedOp {
   framework::OperatorWithKernel::OpKernelFunc func;
   platform::DeviceContext* dev_ctx;
 };
+
 class OpBase;
 
+/* The wrapper for Variable which holds a Variable and a VarBase of its
+ * gradient. This object should be managed totally by Python intepreter.
+ *
+ * Nearly all interface should be implemented in C++.
+ */
 class VarBase {
  public:
-  VarBase()
+  VarBase() : VarBase(new framework::Variable(), new VarBase(true)) {}
+
+  // Owns `var` and `grad`
+  VarBase(framework::Variable* var, VarBase* grad)
       : pre_op_(nullptr),
+        pre_op_out_name_(),
         pre_op_out_idx_(-1),
         var_desc_(nullptr),
-        var_(new framework::Variable()),
-        grads_(new framework::Variable()),
+        var_(var),
+        grads_(grad),
         stop_gradient_(false) {}
 
   explicit VarBase(bool stop_gradient)
       : pre_op_(nullptr),
+        pre_op_out_name_(),
         pre_op_out_idx_(-1),
         var_desc_(nullptr),
         var_(new framework::Variable()),
-        grads_(new framework::Variable()),
+        grads_(stop_gradient ? nullptr : new VarBase(true)),
         stop_gradient_(stop_gradient) {}
 
-  virtual ~VarBase() {}
+  virtual ~VarBase() {
+    if (var_) {
+      delete var_;
+    }
+
+    if (grads_) {
+      delete grads_;
+    }
+  }
 
   void RunBackward();
 
-  framework::LoDTensor& Grad();
+  framework::LoDTensor& GradValue();
 
   inline std::string GradName() const {
     PADDLE_ENFORCE(
@@ -115,15 +143,23 @@ class VarBase {
   int pre_op_out_idx_;
 
   framework::VarDesc* var_desc_;
+
   framework::Variable* var_;
-  framework::Variable* grads_;
+  VarBase* grads_;
 
   bool stop_gradient_;
 };
 
+/* The wrapper for OpDesc which holds a OpDesc and a OpDesc of its
+ * gradient. This object should be managed totally by Python intepreter.
+ */
 class OpBase {
  public:
-  OpBase() : op_desc_(nullptr), grad_op_desc_(nullptr) {}
+  OpBase()
+      : op_desc_(nullptr),
+        forward_id_(-1),
+        grad_op_desc_(nullptr),
+        backward_id_(-1) {}
 
   virtual ~OpBase() {
     if (grad_op_desc_) delete grad_op_desc_;
@@ -131,16 +167,22 @@ class OpBase {
 
   std::map<std::string, std::vector<VarBase*>> ApplyGrad();
 
+  // One of `op_desc_` or `forward_id_` is set, not both.
+  // For pure python PyLayer, use `forward_id_`, otherwise, use op_desc_.
   framework::OpDesc* op_desc_;
+  int forward_id_;
+  // When has backward, one of `grad_op_desc_` or `backward_id_` is set,
+  // not both.
   framework::OpDesc* grad_op_desc_;
+  int backward_id_;
 
-  std::map<std::string, std::vector<VarBase*>> input_vars_;
-  std::map<std::string, std::vector<VarBase*>> output_vars_;
-  std::map<std::string, std::vector<OpBase*>> pre_ops_;
+  VarBasePtrMap input_vars_;
+  VarBasePtrMap output_vars_;
+  OpBasePtrMap pre_ops_;
   std::map<std::string, std::vector<int>> pre_ops_out_idx_;
 
-  std::map<std::string, std::vector<framework::Variable*>> grad_input_vars_;
-  std::map<std::string, std::vector<framework::Variable*>> grad_output_vars_;
+  framework::VariableValueMap grad_input_vars_;
+  framework::VariableValueMap grad_output_vars_;
   framework::BlockDesc* block_;
 };
 
@@ -152,8 +194,25 @@ class Layer {
     std::vector<VarBase> vars;
     return vars;
   }
+};
+
+class PyLayer {
+ public:
+  virtual ~PyLayer() {}
+
+  static void RegisterFunc(int func_id, const py::object& py_func);
+
+  static int NumFuncs();
+
+  static std::vector<VarBase*> Apply(int func_id,
+                                     const std::vector<VarBase*>& inputs);
+
+  static std::vector<framework::Variable*> ApplyGrad(
+      int func_id, const std::vector<framework::Variable*>& inputs);
 
-  virtual void Backward() { LOG(ERROR) << "To support customize"; }
+ private:
+  static std::vector<framework::Variable*> CallPythonFunc(
+      const py::object& callable, const std::vector<framework::Variable*>& ins);
 };
 
 }  // namespace imperative

paddle/fluid/imperative/tracer.cc

@@ -15,5 +15,199 @@
 #include "paddle/fluid/imperative/tracer.h"
 
 namespace paddle {
-namespace imperative {}  // namespace imperative
+namespace imperative {
+
+void CreateGradOp(const framework::OpDesc& op_desc,
+                  const std::unordered_set<std::string>& no_grad_set,
+                  const std::vector<framework::BlockDesc*>& grad_sub_block,
+                  framework::OpDesc** grad_op_desc,
+                  std::unordered_map<std::string, std::string>* grad_to_var) {
+  std::vector<std::unique_ptr<framework::OpDesc>> grad_op_descs =
+      framework::OpInfoMap::Instance()
+          .Get(op_desc.Type())
+          .GradOpMaker()(op_desc, no_grad_set, grad_to_var, grad_sub_block);
+  PADDLE_ENFORCE(grad_op_descs.size() == 1, "Only support 1 grad op now.");
+  // TODO(panyx0718): Leak?
+  *grad_op_desc = grad_op_descs[0].release();
+}
+
+void InitVar(framework::Variable* var, framework::Variable* grad_var) {
+  auto& var_t = var->Get<framework::LoDTensor>();
+  float* data =
+      grad_var->GetMutable<framework::LoDTensor>()->mutable_data<float>(
+          var_t.dims(), platform::CPUPlace());
+  std::fill(data, data + var_t.numel(), 0.0);
+}
+
+void Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
+                   const VarBasePtrMap& outputs, framework::BlockDesc* block,
+                   const bool stop_gradient) {
+  std::map<std::string, VarBase*> vars;
+
+  framework::OpDesc* op_desc = op->op_desc_;
+  VLOG(3) << "tracer tracing " << op_desc->Type();
+  op_desc->InferShape(*block);
+  op_desc->InferVarType(block);
+  std::unique_ptr<framework::OperatorBase> op_base =
+      framework::OpRegistry::CreateOp(*op_desc);
+
+  framework::VariableValueMap invars_map;
+  framework::VariableValueMap outvars_map;
+
+  op->input_vars_ = inputs;
+  for (auto it : op->input_vars_) {
+    auto& invars = invars_map[it.first];
+    for (VarBase* inp : it.second) {
+      PADDLE_ENFORCE_NOT_NULL(inp->var_, "op %s input %s nullptr",
+                              op->op_desc_->Type(), inp->var_desc_->Name());
+
+      invars.push_back(inp->var_);
+      vars[inp->var_desc_->Name()] = inp;
+      if (inp->pre_op_) {
+        op->pre_ops_[it.first].push_back(inp->pre_op_);
+        op->pre_ops_out_idx_[it.first].push_back(inp->pre_op_out_idx_);
+      } else {
+        op->pre_ops_[it.first].push_back(nullptr);
+      }
+      VLOG(3) << "input vname " << inp->var_desc_->Name() << " "
+              << inp->var_->IsInitialized();
+    }
+  }
+
+  op->output_vars_ = outputs;
+  for (auto it : op->output_vars_) {
+    auto& outvars = outvars_map[it.first];
+    const std::vector<VarBase*>& outputs = it.second;
+    for (size_t i = 0; i < outputs.size(); ++i) {
+      VarBase* out = outputs[i];
+      outvars.push_back(out->var_);
+      vars[out->var_desc_->Name()] = out;
+
+      framework::VarDesc* var_desc = block->FindVar(out->var_desc_->Name());
+      if (var_desc->GetType() == framework::proto::VarType::LOD_TENSOR) {
+        out->var_->GetMutable<framework::LoDTensor>();
+      } else {
+        LOG(ERROR) << "tracer doesn't support yet";
+      }
+      out->stop_gradient_ = stop_gradient;
+      out->pre_op_ = op;
+      out->pre_op_out_name_ = it.first;
+      out->pre_op_out_idx_ = i;
+
+      VLOG(3) << "output vname " << out->var_desc_->Name() << " "
+              << out->var_->IsInitialized();
+    }
+  }
+
+  VLOG(3) << "tracer running " << op_desc->Type();
+  framework::RuntimeContext ctx(invars_map, outvars_map);
+
+  // TODO(panyx0718): Cache p.
+  framework::OperatorWithKernel* op_kernel =
+      dynamic_cast<framework::OperatorWithKernel*>(op_base.get());
+  PADDLE_ENFORCE_NOT_NULL(op_kernel, "only support op with kernel");
+
+  framework::Scope scope;
+  platform::CPUPlace place;
+  PreparedOp p = PreparedOp::Prepare(ctx, *op_kernel, place);
+  p.op.RuntimeInferShape(scope, place, ctx);
+  p.func(framework::ExecutionContext(p.op, scope, *p.dev_ctx, p.ctx));
+
+  if (!stop_gradient) {
+    framework::OpDesc* grad_op_desc;
+    // TODO(panyx): Is this leaked?
+    std::unique_ptr<std::unordered_map<std::string, std::string>> grad_to_var(
+        new std::unordered_map<std::string, std::string>());
+    CreateGradOp(*op_desc, {}, {block}, &grad_op_desc, grad_to_var.get());
+    op->grad_op_desc_ = grad_op_desc;
+
+    for (auto it : grad_op_desc->Inputs()) {
+      auto& grad_in_vars = op->grad_input_vars_[it.first];
+      for (const std::string& grad_invar : it.second) {
+        block->FindRecursiveOrCreateVar(grad_invar);
+        auto var_it = grad_to_var->find(grad_invar);
+        if (var_it == grad_to_var->end()) {
+          auto fwd_var_it = vars.find(grad_invar);
+          PADDLE_ENFORCE(fwd_var_it != vars.end());
+          // Forward inputs or outputs.
+          grad_in_vars.push_back(fwd_var_it->second->var_);
+        } else {
+          VarBase* var = vars[var_it->second];
+          if (!var->grads_->var_->IsInitialized()) {
+            InitVar(var->var_, var->grads_->var_);
+          }
+          // Douts.
+          grad_in_vars.push_back(var->grads_->var_);
+        }
+      }
+    }
+
+    for (auto it : grad_op_desc->Outputs()) {
+      auto& grad_out_vars = op->grad_output_vars_[it.first];
+      for (const std::string& grad_outvar : it.second) {
+        block->FindRecursiveOrCreateVar(grad_outvar);
+        auto var_it = grad_to_var->find(grad_outvar);
+        PADDLE_ENFORCE(var_it != grad_to_var->end());
+        VarBase* var = vars[var_it->second];
+        if (!var->grads_->var_->IsInitialized()) {
+          InitVar(var->var_, var->grads_->var_);
+        }
+        grad_out_vars.push_back(var->grads_->var_);
+      }
+    }
+  }
+
+  op->block_ = block;
+}
+
+std::vector<VarBase*> Tracer::PyTrace(OpBase* op,
+                                      const std::vector<VarBase*>& inputs,
+                                      bool stop_gradient) {
+  VLOG(3) << "py_trace";
+  op->input_vars_["X"] = inputs;
+  op->output_vars_["Out"] = PyLayer::Apply(op->forward_id_, inputs);
+  for (VarBase* inp : inputs) {
+    if (inp->pre_op_) {
+      op->pre_ops_["X"].push_back(inp->pre_op_);
+      op->pre_ops_out_idx_["X"].push_back(inp->pre_op_out_idx_);
+    } else {
+      op->pre_ops_["X"].push_back(nullptr);
+    }
+  }
+
+  auto& outputs = op->output_vars_["Out"];
+  for (size_t i = 0; i < outputs.size(); ++i) {
+    VarBase* out = outputs[i];
+    out->stop_gradient_ = stop_gradient;
+    out->pre_op_ = op;
+    out->pre_op_out_name_ = "Out";
+    out->pre_op_out_idx_ = i;
+  }
+  if (!stop_gradient) {
+    auto& grad_input_vars = op->grad_input_vars_["X@GRAD"];
+    auto& grad_output_vars = op->grad_output_vars_["Out@GRAD"];
+
+    for (const VarBase* inp : inputs) {
+      grad_input_vars.push_back(inp->var_);
+    }
+    for (VarBase* out : outputs) {
+      grad_input_vars.push_back(out->var_);
+    }
+    for (VarBase* out : outputs) {
+      grad_input_vars.push_back(out->grads_->var_);
+      if (!grad_input_vars.back()->IsInitialized()) {
+        InitVar(out->var_, grad_input_vars.back());
+      }
+    }
+    for (const VarBase* inp : inputs) {
+      grad_output_vars.push_back(inp->grads_->var_);
+      if (!grad_output_vars.back()->IsInitialized()) {
+        InitVar(inp->var_, grad_output_vars.back());
+      }
+    }
+  }
+  return outputs;
+}
+
+}  // namespace imperative
 }  // namespace paddle

paddle/fluid/imperative/tracer.h

@@ -30,23 +30,9 @@ void CreateGradOp(const framework::OpDesc& op_desc,
                   const std::unordered_set<std::string>& no_grad_set,
                   const std::vector<framework::BlockDesc*>& grad_sub_block,
                   framework::OpDesc** grad_op_desc,
-                  std::unordered_map<std::string, std::string>* grad_to_var) {
-  std::vector<std::unique_ptr<framework::OpDesc>> grad_op_descs =
-      framework::OpInfoMap::Instance()
-          .Get(op_desc.Type())
-          .GradOpMaker()(op_desc, no_grad_set, grad_to_var, grad_sub_block);
-  PADDLE_ENFORCE(grad_op_descs.size() == 1, "Only support 1 grad op now.");
-  // TODO(panyx0718): Leak?
-  *grad_op_desc = grad_op_descs[0].release();
-}
+                  std::unordered_map<std::string, std::string>* grad_to_var);
 
-void InitVar(framework::Variable* var, framework::Variable* grad_var) {
-  auto& var_t = var->Get<framework::LoDTensor>();
-  float* data =
-      grad_var->GetMutable<framework::LoDTensor>()->mutable_data<float>(
-          var_t.dims(), platform::CPUPlace());
-  std::fill(data, data + var_t.numel(), 0.0);
-}
+void InitVar(framework::Variable* var, framework::Variable* grad_var);
 
 class Tracer {
  public:
@@ -57,120 +43,10 @@ class Tracer {
   void Trace(OpBase* op,
              const std::map<std::string, std::vector<VarBase*>>& inputs,
              const std::map<std::string, std::vector<VarBase*>>& outputs,
-             framework::BlockDesc* block, const bool stop_gradient = false) {
-    std::map<std::string, VarBase*> vars;
+             framework::BlockDesc* block, const bool stop_gradient = false);
 
-    framework::OpDesc* op_desc = op->op_desc_;
-    VLOG(3) << "tracer tracing " << op_desc->Type();
-    op_desc->InferShape(*block);
-    op_desc->InferVarType(block);
-    std::unique_ptr<framework::OperatorBase> op_base =
-        framework::OpRegistry::CreateOp(*op_desc);
-
-    framework::VariableValueMap invars_map;
-    framework::VariableValueMap outvars_map;
-
-    op->input_vars_ = inputs;
-    for (auto it : op->input_vars_) {
-      auto& invars = invars_map[it.first];
-      for (VarBase* inp : it.second) {
-        PADDLE_ENFORCE_NOT_NULL(inp->var_, "op %s input %s nullptr",
-                                op->op_desc_->Type(), inp->var_desc_->Name());
-
-        invars.push_back(inp->var_);
-        vars[inp->var_desc_->Name()] = inp;
-        if (inp->pre_op_) {
-          op->pre_ops_[it.first].push_back(inp->pre_op_);
-          op->pre_ops_out_idx_[it.first].push_back(inp->pre_op_out_idx_);
-        } else {
-          op->pre_ops_[it.first].push_back(nullptr);
-        }
-        VLOG(3) << "input vname " << inp->var_desc_->Name() << " "
-                << inp->var_->IsInitialized();
-      }
-    }
-
-    op->output_vars_ = outputs;
-    for (auto it : op->output_vars_) {
-      auto& outvars = outvars_map[it.first];
-      const std::vector<VarBase*>& outputs = it.second;
-      for (size_t i = 0; i < outputs.size(); ++i) {
-        VarBase* out = outputs[i];
-        outvars.push_back(out->var_);
-        vars[out->var_desc_->Name()] = out;
-
-        framework::VarDesc* var_desc = block->FindVar(out->var_desc_->Name());
-        if (var_desc->GetType() == framework::proto::VarType::LOD_TENSOR) {
-          out->var_->GetMutable<framework::LoDTensor>();
-        } else {
-          LOG(ERROR) << "tracer doesn't support yet";
-        }
-        out->stop_gradient_ = stop_gradient;
-        out->pre_op_ = op;
-        out->pre_op_out_name_ = it.first;
-        out->pre_op_out_idx_ = i;
-
-        VLOG(3) << "output vname " << out->var_desc_->Name() << " "
-                << out->var_->IsInitialized();
-      }
-    }
-
-    VLOG(3) << "tracer running " << op_desc->Type();
-    framework::RuntimeContext ctx(invars_map, outvars_map);
-
-    // TODO(panyx0718): Cache p.
-    framework::OperatorWithKernel* op_kernel =
-        dynamic_cast<framework::OperatorWithKernel*>(op_base.get());
-    PADDLE_ENFORCE_NOT_NULL(op_kernel, "only support op with kernel");
-
-    framework::Scope scope;
-    platform::CPUPlace place;
-    PreparedOp p = PreparedOp::Prepare(ctx, *op_kernel, place);
-    p.op.RuntimeInferShape(scope, place, ctx);
-    p.func(framework::ExecutionContext(p.op, scope, *p.dev_ctx, p.ctx));
-
-    if (!stop_gradient) {
-      framework::OpDesc* grad_op_desc;
-      auto grad_to_var = new std::unordered_map<std::string, std::string>();
-      CreateGradOp(*op_desc, {}, {block}, &grad_op_desc, grad_to_var);
-      op->grad_op_desc_ = grad_op_desc;
-
-      for (auto it : grad_op_desc->Inputs()) {
-        auto& grad_in_vars = op->grad_input_vars_[it.first];
-        for (const std::string& grad_invar : it.second) {
-          block->FindRecursiveOrCreateVar(grad_invar);
-          auto var_it = grad_to_var->find(grad_invar);
-          if (var_it == grad_to_var->end()) {
-            auto fwd_var_it = vars.find(grad_invar);
-            PADDLE_ENFORCE(fwd_var_it != vars.end());
-            grad_in_vars.push_back(fwd_var_it->second->var_);
-          } else {
-            VarBase* var = vars[var_it->second];
-            if (!var->grads_->IsInitialized()) {
-              InitVar(var->var_, var->grads_);
-            }
-            grad_in_vars.push_back(var->grads_);
-          }
-        }
-      }
-
-      for (auto it : grad_op_desc->Outputs()) {
-        auto& grad_out_vars = op->grad_output_vars_[it.first];
-        for (const std::string& grad_outvar : it.second) {
-          block->FindRecursiveOrCreateVar(grad_outvar);
-          auto var_it = grad_to_var->find(grad_outvar);
-          PADDLE_ENFORCE(var_it != grad_to_var->end());
-          VarBase* var = vars[var_it->second];
-          if (!var->grads_->IsInitialized()) {
-            InitVar(var->var_, var->grads_);
-          }
-          grad_out_vars.push_back(var->grads_);
-        }
-      }
-    }
-
-    op->block_ = block;
-  }
+  std::vector<VarBase*> PyTrace(OpBase* op, const std::vector<VarBase*>& inputs,
+                                bool stop_gradient = false);
 
  private:
   framework::BlockDesc* root_block_;

paddle/fluid/imperative/type_defs.h

+/* Copyright (c) 2016 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 <map>
+#include <string>
+#include <vector>
+
+namespace paddle {
+namespace imperative {
+
+class VarBase;
+class OpBase;
+
+typedef std::map<std::string, std::vector<VarBase*>> VarBasePtrMap;
+typedef std::map<std::string, std::vector<OpBase*>> OpBasePtrMap;
+
+}  // namespace imperative
+}  // namespace paddle

paddle/fluid/inference/analysis/analyzer_tester.cc

@@ -80,8 +80,8 @@ void TestWord2vecPrediction(const std::string& model_path) {
        i++) {
     LOG(INFO) << "data: " << static_cast<float*>(outputs.front().data.data())[i]
               << " result: " << result[i];
-    PADDLE_ENFORCE(static_cast<float*>(outputs.front().data.data())[i],
-                   result[i]);
+    EXPECT_NEAR(static_cast<float*>(outputs.front().data.data())[i], result[i],
+                1e-3);
   }
 }
 

paddle/fluid/inference/analysis/argument.h

@@ -123,8 +123,6 @@ struct Argument {
   DECL_ARGUMENT_FIELD(use_gpu, UseGPU, bool);
   DECL_ARGUMENT_FIELD(gpu_device_id, GPUDeviceId, int);
   DECL_ARGUMENT_FIELD(use_tensorrt, UseTensorRT, bool);
-  DECL_ARGUMENT_FIELD(tensorrt_node_teller, TensorRtNodeTeller,
-                      std::function<bool(const framework::ir::Node*)>);
   DECL_ARGUMENT_FIELD(tensorrt_max_batch_size, TensorRtMaxBatchSize, int);
   DECL_ARGUMENT_FIELD(tensorrt_workspace_size, TensorRtWorkspaceSize, int);
   DECL_ARGUMENT_FIELD(tensorrt_min_subgraph_size, TensorRtMinSubgraphSize, int);

paddle/fluid/inference/analysis/ir_pass_manager.cc

@@ -49,13 +49,6 @@ void IRPassManager::CreatePasses(Argument *argument,
   for (const std::string &pass_name : passes) {
     auto pass = framework::ir::PassRegistry::Instance().Get(pass_name);
 
-    // Set some pass attributes.
-    if (pass_name == "ir_analysis_pass") {
-      pass->Set("tensorrt_node_teller",
-                new SubgraphDetector::NodeInsideSubgraphTeller(
-                    argument->tensorrt_node_teller()));
-    }
-
     if (pass_name == "graph_viz_pass") {
       std::string dot_file_path = std::to_string(pass_num) + "_ir_" +
                                   (pre_pass.empty() ? "origin" : pre_pass) +
@@ -70,9 +63,6 @@ void IRPassManager::CreatePasses(Argument *argument,
     }
 
     if (pass_name == "tensorrt_subgraph_pass") {
-      PADDLE_ENFORCE(argument->tensorrt_node_teller_valid());
-      pass->SetNotOwned("tensorrt_node_teller",
-                        argument->tensorrt_node_teller_ptr());
       pass->Set("workspace_size", new int(argument->tensorrt_workspace_size()));
       pass->Set("max_batch_size", new int(argument->tensorrt_max_batch_size()));
       pass->Set("min_subgraph_size",

paddle/fluid/inference/analysis/ir_passes/CMakeLists.txt

 cc_library(subgraph_detector SRCS subgraph_detector.cc DEPS proto_desc)
-cc_library(tensorrt_subgraph_pass SRCS tensorrt_subgraph_pass.cc DEPS subgraph_detector)
-set(analysis_deps ${analysis_deps}
-        subgraph_detector tensorrt_subgraph_pass
-        CACHE INTERNAL "")
 
-set(pass_file ${PADDLE_BINARY_DIR}/paddle/fluid/inference/api/paddle_inference_pass.h)
-file(APPEND ${pass_file} "USE_PASS(tensorrt_subgraph_pass);\n")
-set(INFER_IR_PASSES ${INFER_IR_PASSES} tensorrt_subgraph_pass CACHE INTERNAL "")
+if (TENSORRT_FOUND)
+  cc_library(tensorrt_subgraph_pass SRCS tensorrt_subgraph_pass.cc DEPS subgraph_detector tensorrt_op_teller)
+
+  set(analysis_deps ${analysis_deps}
+          subgraph_detector tensorrt_subgraph_pass
+          CACHE INTERNAL "")
+
+  set(pass_file ${PADDLE_BINARY_DIR}/paddle/fluid/inference/api/paddle_inference_pass.h)
+  file(APPEND ${pass_file} "USE_PASS(tensorrt_subgraph_pass);\n")
+  set(INFER_IR_PASSES ${INFER_IR_PASSES} tensorrt_subgraph_pass CACHE INTERNAL "")
+endif()

paddle/fluid/inference/analysis/ir_passes/tensorrt_subgraph_pass.cc

@@ -20,6 +20,7 @@
 #include "paddle/fluid/inference/analysis/helper.h"
 #include "paddle/fluid/inference/analysis/ir_passes/subgraph_detector.h"
 #include "paddle/fluid/inference/analysis/ir_passes/tensorrt_subgraph_pass.h"
+#include "paddle/fluid/inference/tensorrt/op_teller.h"
 
 namespace paddle {
 namespace inference {
@@ -35,8 +36,10 @@ std::unique_ptr<framework::ir::Graph> analysis::TensorRtSubgraphPass::ApplyImpl(
     std::unique_ptr<framework::ir::Graph> graph) const {
   framework::ir::FusePassBase::Init("tensorrt_subgraph_pass", graph.get());
 
-  auto teller =
-      Get<SubgraphDetector::NodeInsideSubgraphTeller>("tensorrt_node_teller");
+  auto teller = [](const framework::ir::Node *node) {
+    if (!node->IsOp() || !node->Op()) return false;
+    return tensorrt::OpTeller::Global().Tell(node->Op()->Type(), *node->Op());
+  };
 
   SubGraphFuser fuser(graph.get(), teller,
                       Get<int>("min_subgraph_size") /*min subgraph size*/);
@@ -232,7 +235,6 @@ std::vector<std::string> ExtractParameters(
 
 REGISTER_PASS(tensorrt_subgraph_pass,
               paddle::inference::analysis::TensorRtSubgraphPass)
-    .RequirePassAttr("tensorrt_node_teller")
     .RequirePassAttr("max_batch_size")
     .RequirePassAttr("workspace_size")
     .RequirePassAttr("min_subgraph_size");

paddle/fluid/inference/analysis/passes/CMakeLists.txt

@@ -7,4 +7,5 @@ set(analysis_deps ${analysis_deps}
         ir_graph_build_pass
         ir_analysis_pass
         analysis_passes
+        subgraph_detector
         CACHE INTERNAL "")

paddle/fluid/inference/analysis/passes/ir_analysis_compose_pass.cc

@@ -27,9 +27,6 @@ namespace analysis {
 
 void IrAnalysisComposePass::RunImpl(Argument *argument) {
   ARGUMENT_CHECK_FIELD(argument, ir_analysis_passes);
-  if (argument->use_tensorrt_valid() && argument->use_tensorrt()) {
-    InitTensorRTAttrs(argument);
-  }
   ApplyIrPasses(argument);
   CollectFusionStatis(argument);
 }
@@ -38,26 +35,6 @@ std::string IrAnalysisComposePass::repr() const {
   return "ir-analysis-compose-pass";
 }
 
-void IrAnalysisComposePass::InitTensorRTAttrs(Argument *argument) {
-  if (argument->use_tensorrt_valid() && argument->use_tensorrt()) {
-    LOG(INFO) << "Initing TensorRT pass";
-    argument->SetTensorRtNodeTeller([](const framework::ir::Node *node) {
-      std::unordered_set<std::string> teller_set(
-          {"mul", "conv2d", "pool2d", "relu", "softmax", "sigmoid",
-           "depthwise_conv2d", "batch_norm", "concat", "tanh", "pad",
-           "elementwise_add", "elementwise_mul", "dropout", "split", "prelu",
-           "conv2d_transpose", "leaky_relu"});
-      if (!node->IsOp()) return false;
-
-      if (teller_set.count(node->Op()->Type())) {
-        return true;
-      } else {
-        return false;
-      }
-    });
-  }
-}
-
 void IrAnalysisComposePass::ApplyIrPasses(Argument *argument) {
   std::vector<std::string> passes({
       "ir_graph_build_pass", "ir_analysis_pass",

paddle/fluid/inference/analysis/passes/ir_analysis_compose_pass.h

@@ -33,8 +33,6 @@ class IrAnalysisComposePass : public AnalysisPass {
   std::string repr() const override;
 
  private:
-  void InitTensorRTAttrs(Argument* argument);
-
   void ApplyIrPasses(Argument* argument);
 
   void CollectFusionStatis(Argument* argument);

paddle/fluid/inference/api/analysis_config.cc

@@ -14,86 +14,101 @@
 
 #include "paddle/fluid/framework/lod_tensor.h"
 #include "paddle/fluid/framework/scope.h"
+#include "paddle/fluid/inference/api/paddle_analysis_config.h"
 #include "paddle/fluid/inference/api/paddle_inference_api.h"
+#include "paddle/fluid/inference/api/paddle_pass_builder.h"
 #include "paddle/fluid/platform/enforce.h"
-#include "paddle_pass_builder.h"  // NOLINT
+#include "paddle/fluid/platform/gpu_info.h"
 
 namespace paddle {
 
 PassStrategy *contrib::AnalysisConfig::pass_builder() const {
-  PADDLE_ENFORCE(
-      pass_builder_.get(),
-      "Should call constructor first, that will init the pass_builder_.");
+  if (!pass_builder_.get()) {
+    if (use_gpu_) {
+      LOG(INFO) << "Create GPU IR passes";
+      pass_builder_.reset(new GpuPassStrategy);
+    } else {
+      LOG(INFO) << "Create CPU IR passes";
+      pass_builder_.reset(new CpuPassStrategy);
+    }
+  } else if (pass_builder_->use_gpu() ^ use_gpu()) {
+    LOG(WARNING) << "The use_gpu flag is not compatible between Config and "
+                    "PassBuilder, the flags are "
+                 << use_gpu() << " " << pass_builder_->use_gpu();
+    LOG(WARNING) << "Please make them compatible, still use the existing "
+                    "PassBuilder.";
+  }
+
   return pass_builder_.get();
 }
 
-contrib::AnalysisConfig::AnalysisConfig(bool use_gpu) {
-  this->use_gpu = use_gpu;
-  if (use_gpu) {
-    pass_builder_.reset(new GpuPassStrategy);
-  } else {
-    pass_builder_.reset(new CpuPassStrategy);
-  }
+contrib::AnalysisConfig::AnalysisConfig(const std::string &model_dir) {
+  model_dir_ = model_dir;
+}
+contrib::AnalysisConfig::AnalysisConfig(const std::string &prog_file,
+                                        const std::string &params_file) {
+  prog_file_ = prog_file;
+  params_file_ = params_file;
+}
+void contrib::AnalysisConfig::SetModel(const std::string &prog_file_path,
+                                       const std::string &params_file_path) {
+  prog_file_ = prog_file_path;
+  params_file_ = params_file_path;
+}
+void contrib::AnalysisConfig::EnableUseGpu(uint64_t memory_pool_init_size_mb,
+                                           int device_id) {
+#ifdef PADDLE_WITH_CUDA
+  use_gpu_ = true;
+  memory_pool_init_size_mb_ = memory_pool_init_size_mb;
+  device_id_ = device_id;
+#else
+  LOG(ERROR) << "Please compile with gpu to EnableGpu";
+  use_gpu_ = false;
+#endif
 }
+void contrib::AnalysisConfig::DisableGpu() { use_gpu_ = false; }
 
 contrib::AnalysisConfig::AnalysisConfig(const contrib::AnalysisConfig &other) {
-  // fields from Config
-  model_dir = other.model_dir;
-  // fields from NativeConfig
-  use_gpu = other.use_gpu;
-  device = other.device;
-  fraction_of_gpu_memory = other.fraction_of_gpu_memory;
-  prog_file = other.prog_file;
-  param_file = other.param_file;
-  specify_input_name = other.specify_input_name;
-  cpu_math_library_num_threads_ = other.cpu_math_library_num_threads_;
-  // fields from this.
-  enable_ir_optim = other.enable_ir_optim;
-  // For mkldnn
-  use_mkldnn_ = other.use_mkldnn_;
-  mkldnn_enabled_op_types_ = other.mkldnn_enabled_op_types_;
-
-  use_feed_fetch_ops = other.use_feed_fetch_ops;
-  use_tensorrt_ = other.use_tensorrt_;
-  tensorrt_max_batchsize_ = other.tensorrt_max_batchsize_;
-  tensorrt_workspace_size_ = other.tensorrt_workspace_size_;
-  tensorrt_min_subgraph_size_ = other.tensorrt_min_subgraph_size_;
-  model_from_memory_ = other.model_from_memory_;
-
-  if (use_gpu) {
+#define CP_MEMBER(member__) member__ = other.member__;
+
+  // Model related.
+  CP_MEMBER(model_dir_);
+  CP_MEMBER(prog_file_);
+  CP_MEMBER(params_file_);
+  CP_MEMBER(model_from_memory_);  // the memory model reuses prog_file_ and
+                                  // params_file_ fields.
+  // Gpu releated.
+  CP_MEMBER(use_gpu_);
+  CP_MEMBER(device_id_);
+  CP_MEMBER(memory_pool_init_size_mb_);
+  // TensorRT releated.
+  CP_MEMBER(use_tensorrt_);
+  CP_MEMBER(tensorrt_workspace_size_);
+  CP_MEMBER(tensorrt_max_batchsize_);
+  CP_MEMBER(tensorrt_min_subgraph_size_);
+  // MKLDNN releated.
+  CP_MEMBER(use_mkldnn_);
+  CP_MEMBER(mkldnn_enabled_op_types_);
+
+  // Ir related.
+  CP_MEMBER(enable_ir_optim_);
+  CP_MEMBER(use_feed_fetch_ops_);
+  CP_MEMBER(ir_debug_);
+  CP_MEMBER(specify_input_name_);
+
+  CP_MEMBER(cpu_math_library_num_threads_);
+
+  CP_MEMBER(serialized_info_cache_);
+
+  if (use_gpu_) {
     pass_builder_.reset(new GpuPassStrategy(
         *static_cast<GpuPassStrategy *>(other.pass_builder())));
   } else {
     pass_builder_.reset(new CpuPassStrategy(
         *static_cast<CpuPassStrategy *>(other.pass_builder())));
   }
-}
 
-contrib::AnalysisConfig::AnalysisConfig(contrib::AnalysisConfig &&other) {
-  // fields from Config
-  model_dir = other.model_dir;
-  // fields from NativeConfig
-  use_gpu = other.use_gpu;
-  device = other.device;
-  fraction_of_gpu_memory = other.fraction_of_gpu_memory;
-  prog_file = other.prog_file;
-  param_file = other.param_file;
-  specify_input_name = other.specify_input_name;
-  cpu_math_library_num_threads_ = other.cpu_math_library_num_threads_;
-  // fields from this.
-  enable_ir_optim = other.enable_ir_optim;
-  // For mkldnn
-  use_mkldnn_ = other.use_mkldnn_;
-  mkldnn_enabled_op_types_ = other.mkldnn_enabled_op_types_;
-
-  use_feed_fetch_ops = other.use_feed_fetch_ops;
-  use_tensorrt_ = other.use_tensorrt_;
-  tensorrt_max_batchsize_ = other.tensorrt_max_batchsize_;
-  tensorrt_workspace_size_ = other.tensorrt_workspace_size_;
-  tensorrt_min_subgraph_size_ = other.tensorrt_min_subgraph_size_;
-  model_from_memory_ = other.model_from_memory_;
-
-  pass_builder_ = std::move(other.pass_builder_);
+#undef CP_MEMBER
 }
 
 void contrib::AnalysisConfig::EnableMKLDNN() {
@@ -112,17 +127,91 @@ void contrib::AnalysisConfig::EnableTensorRtEngine(int workspace_size,
   use_tensorrt_ = true;
   tensorrt_workspace_size_ = workspace_size;
   tensorrt_max_batchsize_ = max_batch_size;
-  tensorrt_min_subgraph_size_ = min_subgraph_size;
-  // Append after the conv+affine_channel fuse pass.
-  pass_builder()->InsertPass(3, "tensorrt_subgraph_pass");
+  Update();
+}
+
+void contrib::AnalysisConfig::Update() {
+  auto info = SerializeInfoCache();
+  if (info == serialized_info_cache_) return;
+
+  if (use_gpu_) {
+    pass_builder_.reset(new GpuPassStrategy);
+  } else {
+    pass_builder_.reset(new CpuPassStrategy);
+  }
+
+  if (use_tensorrt_) {
+    if (!use_gpu_) {
+      LOG(ERROR)
+          << "TensorRT engine is not available when EnableGpu() not actived.";
+    } else {
+      // Append after the infer_clean pass.
+      pass_builder()->InsertPass(1, "tensorrt_subgraph_pass");
+    }
+  }
+
+  if (use_mkldnn_) {
+    if (!enable_ir_optim_) {
+      LOG(ERROR)
+          << "EnableMKLDNN() only works when IR optimization is enabled.";
+    }
+#ifdef PADDLE_WITH_MKLDNN
+    pass_builder()->EnableMKLDNN();
+    use_mkldnn_ = true;
+#else
+    LOG(ERROR) << "Please compile with MKLDNN first to use MKLDNN";
+    use_mkldnn_ = false;
+#endif
+  }
+
+  if (ir_debug_) {
+    pass_builder()->TurnOnDebug();
+  }
+}
+
+std::string contrib::AnalysisConfig::SerializeInfoCache() {
+  std::stringstream ss;
+  ss << use_gpu_;
+  ss << memory_pool_init_size_mb_;
+
+  ss << use_tensorrt_;
+  ss << tensorrt_workspace_size_;
+  ss << tensorrt_max_batchsize_;
+
+  ss << use_mkldnn_;
+  ss << enable_ir_optim_;
+  ss << use_feed_fetch_ops_;
+  ss << ir_debug_;
+
+  return ss.str();
+}
+
+void contrib::AnalysisConfig::SetCpuMathLibraryNumThreads(
+    int cpu_math_library_num_threads) {
+  cpu_math_library_num_threads_ = cpu_math_library_num_threads;
+}
+
+float contrib::AnalysisConfig::fraction_of_gpu_memory_for_pool() const {
+#ifdef PADDLE_WITH_CUDA
+  // Get the GPU memory details and calculate the fraction of memory for the
+  // GPU memory pool.
+  size_t gpu_used, gpu_available;
+  platform::GpuMemoryUsage(&gpu_used, &gpu_available);
+  double total_gpu_memory = (gpu_used + gpu_available) / 1024. / 1024.;
+  float fraction_of_gpu_memory =
+      static_cast<double>(memory_pool_init_size_mb()) / total_gpu_memory;
+  return fraction_of_gpu_memory;
+#else
+  return 0.;
+#endif
 }
 
 void contrib::AnalysisConfig::SetModelBuffer(const char *prog_buffer,
                                              size_t prog_buffer_size,
                                              const char *param_buffer,
                                              size_t param_buffer_size) {
-  prog_file = std::string(prog_buffer, prog_buffer + prog_buffer_size);
-  param_file = std::string(param_buffer, param_buffer + param_buffer_size);
+  prog_file_ = std::string(prog_buffer, prog_buffer + prog_buffer_size);
+  params_file_ = std::string(param_buffer, param_buffer + param_buffer_size);
   model_from_memory_ = true;
 }
 

paddle/fluid/inference/api/analysis_predictor.cc

@@ -33,6 +33,7 @@
 #include "paddle/fluid/inference/utils/singleton.h"
 #include "paddle/fluid/memory/memcpy.h"
 #include "paddle/fluid/platform/cpu_helper.h"
+#include "paddle/fluid/platform/gpu_info.h"
 #include "paddle/fluid/platform/profiler.h"
 
 DECLARE_bool(profile);
@@ -59,8 +60,8 @@ bool AnalysisPredictor::Init(
   if (FLAGS_profile) {
     LOG(WARNING) << "Profiler is actived, might affect the performance";
     LOG(INFO) << "You can turn off by set gflags '-profile false'";
-    auto tracking_device = config_.use_gpu ? platform::ProfilerState::kAll
-                                           : platform::ProfilerState::kCPU;
+    auto tracking_device = config_.use_gpu() ? platform::ProfilerState::kAll
+                                             : platform::ProfilerState::kCPU;
     platform::EnableProfiler(tracking_device);
   }
 
@@ -112,7 +113,7 @@ bool AnalysisPredictor::PrepareProgram(
     // Optimize the program, and load parameters and modify them in the
     // scope_.
     // This will change the scope_ address.
-    if (config_.enable_ir_optim) {
+    if (config_.ir_optim()) {
       status_ir_optim_enabled_ = true;
       OptimizeInferenceProgram();
     } else {
@@ -140,9 +141,9 @@ bool AnalysisPredictor::PrepareProgram(
   return true;
 }
 bool AnalysisPredictor::CreateExecutor() {
-  if (config_.use_gpu) {
+  if (config_.use_gpu_) {
     status_use_gpu_ = true;
-    place_ = paddle::platform::CUDAPlace(config_.device);
+    place_ = paddle::platform::CUDAPlace(config_.device_id_);
   } else {
     place_ = paddle::platform::CPUPlace();
   }
@@ -151,7 +152,7 @@ bool AnalysisPredictor::CreateExecutor() {
 }
 bool AnalysisPredictor::PrepareExecutor() {
   executor_->Prepare(sub_scope_, *inference_program_, 0,
-                     config_.use_feed_fetch_ops);
+                     config_.use_feed_fetch_ops_);
 
   PADDLE_ENFORCE_NOT_NULL(sub_scope_);
 
@@ -250,7 +251,7 @@ bool AnalysisPredictor::SetFeed(const std::vector<PaddleTensor> &inputs,
     }
     input.set_lod(lod);
     int idx = -1;
-    if (config_.specify_input_name) {
+    if (config_.specify_input_name_) {
       auto name = inputs[i].name;
       if (feed_names_.find(name) == feed_names_.end()) {
         LOG(ERROR) << "feed names from program do not have name: [" << name
@@ -314,22 +315,22 @@ bool AnalysisPredictor::GetFetch(std::vector<PaddleTensor> *outputs,
 void AnalysisPredictor::OptimizeInferenceProgram() {
   status_program_optimized_ = true;
 
-  argument_.SetUseGPU(config_.use_gpu);
-  argument_.SetGPUDeviceId(config_.device);
+  argument_.SetUseGPU(config_.use_gpu());
+  argument_.SetGPUDeviceId(config_.gpu_device_id());
   argument_.SetModelFromMemory(config_.model_from_memory_);
   // Analyze inference_program
-  if (!config_.model_dir.empty()) {
-    argument_.SetModelDir(config_.model_dir);
+  if (!config_.model_dir().empty()) {
+    argument_.SetModelDir(config_.model_dir());
   } else {
     PADDLE_ENFORCE(
-        !config_.param_file.empty(),
+        !config_.params_file().empty(),
         "Either model_dir or (param_file, prog_file) should be set.");
-    PADDLE_ENFORCE(!config_.prog_file.empty());
-    argument_.SetModelProgramPath(config_.prog_file);
-    argument_.SetModelParamsPath(config_.param_file);
+    PADDLE_ENFORCE(!config_.prog_file().empty());
+    argument_.SetModelProgramPath(config_.prog_file());
+    argument_.SetModelParamsPath(config_.params_file());
   }
 
-  if (config_.use_gpu && config_.use_tensorrt_) {
+  if (config_.use_gpu() && config_.tensorrt_engine_enabled()) {
     argument_.SetUseTensorRT(true);
     argument_.SetTensorRtWorkspaceSize(config_.tensorrt_workspace_size_);
     argument_.SetTensorRtMaxBatchSize(config_.tensorrt_max_batchsize_);
@@ -341,7 +342,7 @@ void AnalysisPredictor::OptimizeInferenceProgram() {
   }
 
   auto passes = config_.pass_builder()->AllPasses();
-  if (!config_.enable_ir_optim) passes.clear();
+  if (!config_.ir_optim()) passes.clear();
   argument_.SetIrAnalysisPasses(passes);
   argument_.SetScopeNotOwned(const_cast<framework::Scope *>(scope_.get()));
   Analyzer().Run(&argument_);
@@ -358,18 +359,26 @@ template <>
 std::unique_ptr<PaddlePredictor> CreatePaddlePredictor<
     AnalysisConfig, PaddleEngineKind::kAnalysis>(const AnalysisConfig &config) {
   VLOG(3) << "create AnalysisConfig";
-  if (config.use_gpu) {
+  if (config.use_gpu()) {
     // 1. GPU memeroy
-    PADDLE_ENFORCE_GT(
-        config.fraction_of_gpu_memory, 0.f,
-        "fraction_of_gpu_memory in the config should be set to range (0., 1.]");
-    PADDLE_ENFORCE_GE(config.device, 0, "Invalid device id %d", config.device);
+    PADDLE_ENFORCE_GT(config.memory_pool_init_size_mb(), 0.f);
+    PADDLE_ENFORCE_GE(config.gpu_device_id(), 0, "Invalid device id %d",
+                      config.gpu_device_id());
     std::vector<std::string> flags;
-    if (config.fraction_of_gpu_memory >= 0.0f ||
-        config.fraction_of_gpu_memory <= 0.95f) {
+
+    float fraction_of_gpu_memory = config.fraction_of_gpu_memory_for_pool();
+    if (fraction_of_gpu_memory > 0.95f) {
+      LOG(ERROR)
+          << "Allocate too much memory for the GPU memory pool, assigned "
+          << config.memory_pool_init_size_mb() << " MB";
+      LOG(ERROR)
+          << "Try to shink the value by setting AnalysisConfig::EnableGpu(...)";
+    }
+
+    if (fraction_of_gpu_memory >= 0.0f || fraction_of_gpu_memory <= 0.95f) {
       flags.push_back("dummpy");
       std::string flag = "--fraction_of_gpu_memory_to_use=" +
-                         std::to_string(config.fraction_of_gpu_memory);
+                         std::to_string(fraction_of_gpu_memory);
       flags.push_back(flag);
       VLOG(3) << "set flag: " << flag;
       framework::InitGflags(flags);
@@ -443,22 +452,22 @@ bool AnalysisPredictor::ZeroCopyRun() {
 bool AnalysisPredictor::LoadProgramDesc() {
   // Initialize the inference program
   std::string filename;
-  if (!config_.model_dir.empty()) {
-    filename = config_.model_dir + "/__model__";
-  } else if (!config_.prog_file.empty() && !config_.param_file.empty()) {
+  if (!config_.model_dir().empty()) {
+    filename = config_.model_dir() + "/__model__";
+  } else if (!config_.prog_file().empty() && !config_.params_file().empty()) {
     // All parameters are saved in a single file.
     // The file names should be consistent with that used
     // in Python API `fluid.io.save_inference_model`.
-    filename = config_.prog_file;
+    filename = config_.prog_file();
   } else {
-    if (config_.model_dir.empty() && config_.prog_file.empty()) {
+    if (config_.model_dir().empty() && config_.prog_file().empty()) {
       LOG(ERROR)
           << "Either model_dir or (prog_file, param_file) should be set.";
       return false;
     }
     LOG(ERROR) << string::Sprintf(
-        "not valid model path '%s' or program path '%s'.", config_.model_dir,
-        config_.param_file);
+        "not valid model path '%s' or program path '%s'.", config_.model_dir(),
+        config_.params_file());
     return false;
   }
 
@@ -478,7 +487,7 @@ bool AnalysisPredictor::LoadProgramDesc() {
 
     proto.ParseFromString(pb_content);
   } else {
-    proto.ParseFromString(config_.prog_file);
+    proto.ParseFromString(config_.prog_file());
   }
   inference_program_.reset(new framework::ProgramDesc(proto));
   return true;
@@ -508,27 +517,27 @@ bool AnalysisPredictor::LoadParameters() {
       new_var->SetLoDLevel(var->GetLoDLevel());
       new_var->SetPersistable(true);
 
-      if (!config_.param_file.empty()) {
+      if (!config_.params_file().empty()) {
         params.push_back(new_var->Name());
       } else {
         // append_op
         framework::OpDesc *op = load_block->AppendOp();
         op->SetType("load");
         op->SetOutput("Out", {new_var->Name()});
-        op->SetAttr("file_path", {config_.model_dir + "/" + new_var->Name()});
+        op->SetAttr("file_path", {config_.model_dir() + "/" + new_var->Name()});
         op->CheckAttrs();
       }
     }
   }
 
-  if (!config_.param_file.empty()) {
+  if (!config_.params_file().empty()) {
     // sort paramlist to have consistent ordering
     std::sort(params.begin(), params.end());
     // append just the load_combine op
     framework::OpDesc *op = load_block->AppendOp();
     op->SetType("load_combine");
     op->SetOutput("Out", params);
-    op->SetAttr("file_path", {config_.param_file});
+    op->SetAttr("file_path", {config_.params_file()});
     op->CheckAttrs();
   }
 

paddle/fluid/inference/api/analysis_predictor.h

@@ -35,8 +35,11 @@ using framework::proto::ProgramDesc;
 using framework::NaiveExecutor;
 using contrib::AnalysisConfig;
 
-/* This predictor is based on the original native predictor with IR and Analysis
- * support. It will optimize IR and Parameters in the runtime.
+/** \brief This predictor is based on the original native predictor with IR and
+ * Analysis support.
+ *
+ * It will optimize IR and Parameters in the runtime.
+ *
  * TODO(Superjomn) Replace the Navive predictor?
  */
 class AnalysisPredictor : public PaddlePredictor {

paddle/fluid/inference/api/analysis_predictor_tester.cc

@@ -25,9 +25,9 @@ namespace paddle {
 using contrib::AnalysisConfig;
 
 TEST(AnalysisPredictor, analysis_off) {
-  AnalysisConfig config(false);
-  config.model_dir = FLAGS_dirname;
-  config.enable_ir_optim = false;
+  AnalysisConfig config;
+  config.SetModel(FLAGS_dirname);
+  config.SwitchIrOptim(false);
 
   auto _predictor = CreatePaddlePredictor<AnalysisConfig>(config);
   auto* predictor = static_cast<AnalysisPredictor*>(_predictor.get());
@@ -55,14 +55,14 @@ TEST(AnalysisPredictor, analysis_off) {
 }
 
 TEST(AnalysisPredictor, analysis_on) {
+  AnalysisConfig config;
+  config.SetModel(FLAGS_dirname);
+  config.SwitchIrOptim(true);
 #ifdef PADDLE_WITH_CUDA
-  AnalysisConfig config(true);
-  config.fraction_of_gpu_memory = 0.15;
+  config.EnableUseGpu(100, 0);
 #else
-  AnalysisConfig config;
+  config.DisableGpu();
 #endif
-  config.model_dir = FLAGS_dirname;
-  config.enable_ir_optim = true;
 
   auto _predictor = CreatePaddlePredictor<AnalysisConfig>(config);
   auto* predictor = static_cast<AnalysisPredictor*>(_predictor.get());
@@ -89,7 +89,8 @@ TEST(AnalysisPredictor, analysis_on) {
   }
 
   // compare with NativePredictor
-  auto naive_predictor = CreatePaddlePredictor<NativeConfig>(config);
+  auto naive_predictor =
+      CreatePaddlePredictor<NativeConfig>(config.ToNativeConfig());
   std::vector<PaddleTensor> naive_outputs;
   ASSERT_TRUE(naive_predictor->Run(inputs, &naive_outputs));
   ASSERT_EQ(naive_outputs.size(), 1UL);
@@ -98,9 +99,8 @@ TEST(AnalysisPredictor, analysis_on) {
 
 TEST(AnalysisPredictor, ZeroCopy) {
   AnalysisConfig config;
-  config.model_dir = FLAGS_dirname;
-  config.use_feed_fetch_ops = false;
-
+  config.SetModel(FLAGS_dirname);
+  config.SwitchUseFeedFetchOps(false);
   auto predictor = CreatePaddlePredictor<AnalysisConfig>(config);
 
   auto w0 = predictor->GetInputTensor("firstw");
@@ -137,9 +137,9 @@ TEST(AnalysisPredictor, ZeroCopy) {
 
 TEST(AnalysisPredictor, Clone) {
   AnalysisConfig config;
-  config.model_dir = FLAGS_dirname;
-  config.use_feed_fetch_ops = true;
-  config.enable_ir_optim = true;
+  config.SetModel(FLAGS_dirname);
+  config.SwitchUseFeedFetchOps(true);
+  config.SwitchIrOptim(true);
 
   std::vector<std::unique_ptr<PaddlePredictor>> predictors;
   predictors.emplace_back(CreatePaddlePredictor(config));

paddle/fluid/inference/api/api_anakin_engine.h

@@ -19,8 +19,6 @@ limitations under the License. */
 
 #pragma once
 
-#define WITH_ANAKIN
-
 #include <vector>
 
 #include "framework/core/net/net.h"

paddle/fluid/inference/api/api_impl.cc

@@ -288,7 +288,7 @@ std::unique_ptr<PaddlePredictor> CreatePaddlePredictor<
   VLOG(3) << "create NativePaddlePredictor";
   if (config.use_gpu) {
     // 1. GPU memeroy
-    PADDLE_ENFORCE_GT(
+    PADDLE_ENFORCE_GE(
         config.fraction_of_gpu_memory, 0.f,
         "fraction_of_gpu_memory in the config should be set to range (0., 1.]");
     PADDLE_ENFORCE_GE(config.device, 0, "Invalid device id %d", config.device);

paddle/fluid/inference/api/api_impl.h

@@ -19,7 +19,6 @@ limitations under the License. */
 #include <memory>
 #include <string>
 #include <vector>
-
 #include "paddle/fluid/framework/ddim.h"
 #include "paddle/fluid/framework/lod_tensor.h"
 #include "paddle/fluid/framework/lod_tensor_array.h"

paddle/fluid/inference/api/api_impl_tester.cc

@@ -295,7 +295,8 @@ TEST(inference_api_native, image_classification_gpu) {
 #endif
 
 TEST(PassBuilder, Delete) {
-  contrib::AnalysisConfig config(false);
+  contrib::AnalysisConfig config;
+  config.DisableGpu();
   config.pass_builder()->DeletePass("attention_lstm_fuse_pass");
   const auto& passes = config.pass_builder()->AllPasses();
   auto it = std::find(passes.begin(), passes.end(), "attention_lstm_fuse_pass");

paddle/fluid/inference/api/demo_ci/CMakeLists.txt

@@ -92,10 +92,10 @@ if(WITH_MKL)
   if(NOT WIN32)
     set(MATH_LIB ${PADDLE_LIB}/third_party/install/mklml/lib/libmklml_intel${CMAKE_SHARED_LIBRARY_SUFFIX}
                  ${PADDLE_LIB}/third_party/install/mklml/lib/libiomp5${CMAKE_SHARED_LIBRARY_SUFFIX})
-  else(WIN32)
+  else()
     set(MATH_LIB ${PADDLE_LIB}/third_party/install/mklml/lib/libmklml${CMAKE_SHARED_LIBRARY_SUFFIX}
             ${PADDLE_LIB}/third_party/install/mklml/lib/libiomp5md${CMAKE_SHARED_LIBRARY_SUFFIX})
-  endif(WIN32)
+  endif()
   set(MKLDNN_PATH "${PADDLE_LIB}/third_party/install/mkldnn")
   if(EXISTS ${MKLDNN_PATH})
     include_directories("${MKLDNN_PATH}/include")
@@ -128,8 +128,8 @@ else()
       ${CMAKE_STATIC_LIBRARY_PREFIX}glog  ${CMAKE_STATIC_LIBRARY_PREFIX}gflags  ${CMAKE_STATIC_LIBRARY_PREFIX}protobuf
       ${CMAKE_STATIC_LIBRARY_PREFIX}snappy ${CMAKE_STATIC_LIBRARY_PREFIX}z ${CMAKE_STATIC_LIBRARY_PREFIX}xxhash
       snappystream ${EXTERNAL_LIB})
-  # NOTE(dzhwinter) shlwapi is deprecated.
-  set(DEPS ${DEPS} libcmt shlwapi)
+  get_property(os_dependency_modules GLOBAL PROPERTY OS_DEPENDENCY_MODULES)
+  set(DEPS ${DEPS} libcmt ${os_dependency_modules})
 endif(NOT WIN32)
 
 if(WITH_GPU)

paddle/fluid/inference/api/demo_ci/run.sh

@@ -116,6 +116,10 @@ D
       --modeldir=$DATA_DIR/mobilenet/model \
       --data=$DATA_DIR/mobilenet/data.txt \
       --refer=$DATA_DIR/mobilenet/result.txt 
+    if [ $? -ne 0 ]; then
+      echo "trt demo trt_mobilenet_demo runs fail."
+      exit 1
+    fi
   fi
 done
 set +x

paddle/fluid/inference/api/demo_ci/trt_mobilenet_demo.cc

@@ -36,12 +36,11 @@ namespace demo {
  */
 void Main() {
   std::unique_ptr<PaddlePredictor> predictor;
-  paddle::contrib::AnalysisConfig config(true);
-  config.param_file = FLAGS_modeldir + "/__params__";
-  config.prog_file = FLAGS_modeldir + "/__model__";
-  config.device = 0;
+  paddle::contrib::AnalysisConfig config;
+  config.EnableUseGpu(100, 0);
+  config.SetModel(FLAGS_modeldir + "/__model__",
+                  FLAGS_modeldir + "/__params__");
   config.EnableTensorRtEngine();
-  config.fraction_of_gpu_memory = 0.1;  // set by yourself
   predictor = CreatePaddlePredictor(config);
 
   VLOG(3) << "begin to process data";

paddle/fluid/inference/api/demo_ci/vis_demo.cc

@@ -40,15 +40,14 @@ using contrib::AnalysisConfig;
  */
 void Main(bool use_gpu) {
   std::unique_ptr<PaddlePredictor> predictor, analysis_predictor;
-  AnalysisConfig config(use_gpu);
-  config.param_file = FLAGS_modeldir + "/__params__";
-  config.prog_file = FLAGS_modeldir + "/__model__";
-  config.device = 0;
-  if (FLAGS_use_gpu) {
-    config.fraction_of_gpu_memory = 0.1;  // set by yourself
+  AnalysisConfig config;
+  if (use_gpu) {
+    config.EnableUseGpu(100, 0);
   }
+  config.SetModel(FLAGS_modeldir + "/__model__",
+                  FLAGS_modeldir + "/__params__");
 
-  predictor = CreatePaddlePredictor<NativeConfig>(config);
+  predictor = CreatePaddlePredictor<NativeConfig>(config.ToNativeConfig());
   analysis_predictor = CreatePaddlePredictor(config);
 
   // Just a single batch of data.

paddle/fluid/inference/api/helper.h

@@ -204,11 +204,14 @@ static std::string DescribeTensor(const PaddleTensor &tensor) {
     os << to_string(l) << "; ";
   }
   os << "\n";
-  os << " - data: ";
+  os << " - memory length: " << tensor.data.length();
+  os << "\n";
 
+  os << " - data: ";
   int dim = VecReduceToInt(tensor.shape);
+  float *pdata = static_cast<float *>(tensor.data.data());
   for (int i = 0; i < dim; i++) {
-    os << static_cast<float *>(tensor.data.data())[i] << " ";
+    os << pdata[i] << " ";
   }
   os << '\n';
   return os.str();
@@ -224,10 +227,12 @@ static std::string DescribeZeroCopyTensor(const ZeroCopyTensor &tensor) {
     os << to_string(l) << "; ";
   }
   os << "\n";
-  os << " - data: ";
   PaddlePlace place;
   int size;
   const auto *data = tensor.data<float>(&place, &size);
+  os << " - numel: " << size;
+  os << "\n";
+  os << " - data: ";
   for (int i = 0; i < size; i++) {
     os << data[i] << " ";
   }

paddle/fluid/inference/api/paddle_inference_api.h

@@ -26,9 +26,8 @@ limitations under the License. */
 #include <string>
 #include <vector>
 
-#include "paddle_api.h"  // NOLINT
-#ifndef WITH_ANAKIN
 #include "paddle_analysis_config.h"  // NOLINT
-#else
+#include "paddle_api.h"              // NOLINT
+#ifdef WITH_ANAKIN
 #include "paddle_anakin_config.h"  // NOLINT
 #endif

paddle/fluid/inference/api/paddle_pass_builder.h

@@ -18,30 +18,39 @@
 #include <string>
 #include <vector>
 
+/*! \file */
+
+/*! \namespace paddle */
 namespace paddle {
-/*
- * This is a pass builder based on string. It is part of inference API.
+
+/** This is a pass builder based on string. It is part of inference API.
  */
 class PaddlePassBuilder {
  public:
   explicit PaddlePassBuilder(const std::vector<std::string> &passes)
       : passes_(passes) {}
 
+  /** Append a pass to the end of the passes. */
   void AppendPass(const std::string &pass_type);
 
+  /** Insert a pass to a specific position.
+   * @param idx the position to insert.
+   * @param pass_type the pass key.
+   */
   void InsertPass(size_t idx, const std::string &pass_type);
 
-  // Delete the `idx`-th pass.
+  /** Delete the `idx`-th pass. */
   void DeletePass(size_t idx);
 
-  // Delete all the passes that has type `pass_type`.
+  /** Delete all the passes that has type `pass_type`. */
   void DeletePass(const std::string &pass_type);
 
-  // Visualize the computation graph after each pass by generating a DOT
-  // language file, one can draw them with the Graphviz toolkit.
+  /** Visualize the computation graph after each pass by generating a DOT
+   * language file, one can draw them with the Graphviz toolkit.
+   */
   void TurnOnDebug();
 
-  // Human-readible information.
+  /** Human-readible information. */
   std::string DebugString();
 
   const std::vector<std::string> &AllPasses() const { return passes_; }
@@ -50,23 +59,27 @@ class PaddlePassBuilder {
   std::vector<std::string> passes_;
 };
 
-/*
- * Pass strategy to help control the IR passes.
+/**Pass strategy to help control the IR passes.
  */
 class PassStrategy : public PaddlePassBuilder {
  public:
   explicit PassStrategy(const std::vector<std::string> &passes)
       : PaddlePassBuilder(passes) {}
 
-  // The MKLDNN control exists in both CPU and GPU mode, because there can be
-  // still some CPU kernels running in CPU mode.
+  /** The MKLDNN control exists in both CPU and GPU mode, because there can be
+   * still some CPU kernels running in CPU mode.
+   */
   virtual void EnableMKLDNN() = 0;
 
+  bool use_gpu() const { return use_gpu_; }
+
   virtual ~PassStrategy() = default;
+
+ protected:
+  bool use_gpu_{false};
 };
 
-/*
- * The CPU passes controller, it is used in AnalysisPredictor with CPU mode.
+/** The CPU passes controller, it is used in AnalysisPredictor with CPU mode.
  */
 class CpuPassStrategy : public PassStrategy {
  public:
@@ -76,6 +89,7 @@ class CpuPassStrategy : public PassStrategy {
     passes_.assign({
         "infer_clean_graph_pass",         //
         "attention_lstm_fuse_pass",       //
+        "seqpool_concat_fuse_pass",       //
         "seqconv_eltadd_relu_fuse_pass",  //
         // "embedding_fc_lstm_fuse_pass", //
         "fc_lstm_fuse_pass",             //
@@ -88,6 +102,7 @@ class CpuPassStrategy : public PassStrategy {
         "conv_eltwiseadd_bn_fuse_pass",  //
         "is_test_pass",                  //
     });
+    use_gpu_ = false;
   }
 
   virtual ~CpuPassStrategy() = default;
@@ -111,8 +126,7 @@ class CpuPassStrategy : public PassStrategy {
   CpuPassStrategy(const CpuPassStrategy &other) : PassStrategy(other.passes_) {}
 };
 
-/*
- * The GPU passes strategy, it is used in
+/** The GPU passes strategy, it is used in AnalysisPredictor with GPU mode.
  */
 class GpuPassStrategy : public PassStrategy {
  public:
@@ -126,10 +140,18 @@ class GpuPassStrategy : public PassStrategy {
         "conv_elementwise_add2_act_fuse_pass",       //
         "conv_elementwise_add_fuse_pass",            //
     });
+
+    for (int i = 6; i >= 3; i--) {
+      passes_.push_back("transpose_flatten" + std::to_string(i) +
+                        "_concat_fuse_pass");
+    }
+    use_gpu_ = true;
   }
 
   GpuPassStrategy(const GpuPassStrategy &other)
-      : PassStrategy(other.AllPasses()) {}
+      : PassStrategy(other.AllPasses()) {
+    use_gpu_ = true;
+  }
 
   void EnableMKLDNN() override;
 

paddle/fluid/inference/tensorrt/CMakeLists.txt

 nv_library(tensorrt_engine SRCS engine.cc DEPS ${GLOB_OPERATOR_DEPS} framework_proto device_context)
+nv_library(tensorrt_op_teller SRCS op_teller.cc DEPS framework_proto)
 nv_test(test_tensorrt SRCS test_tensorrt.cc DEPS dynload_cuda device_context dynamic_loader)
 nv_test(test_tensorrt_engine SRCS test_engine.cc DEPS dynload_cuda tensorrt_engine)
 add_subdirectory(plugin)

paddle/fluid/inference/tests/api/analyzer_dam_tester.cc

@@ -165,12 +165,9 @@ void PrepareInputs(std::vector<PaddleTensor> *input_slots, DataRecord *data,
 }
 
 void SetConfig(contrib::AnalysisConfig *cfg) {
-  cfg->prog_file = FLAGS_infer_model + "/__model__";
-  cfg->param_file = FLAGS_infer_model + "/param";
-  cfg->use_gpu = false;
-  cfg->device = 0;
-  cfg->specify_input_name = true;
-  cfg->enable_ir_optim = true;
+  cfg->SetModel(FLAGS_infer_model + "/__model__", FLAGS_infer_model + "/param");
+  cfg->SwitchSpecifyInputNames();
+  cfg->SwitchIrOptim(true);
 }
 
 void SetInput(std::vector<std::vector<PaddleTensor>> *inputs) {

paddle/fluid/inference/tests/api/analyzer_lac_tester.cc

@@ -105,11 +105,10 @@ void GetOneBatch(std::vector<PaddleTensor> *input_slots, DataRecord *data,
 }
 
 void SetConfig(AnalysisConfig *cfg) {
-  cfg->model_dir = FLAGS_infer_model;
-  cfg->use_gpu = false;
-  cfg->device = 0;
-  cfg->specify_input_name = true;
-  cfg->enable_ir_optim = true;
+  cfg->SetModel(FLAGS_infer_model);
+  cfg->DisableGpu();
+  cfg->SwitchSpecifyInputNames();
+  cfg->SwitchIrOptim();
 }
 
 void SetInput(std::vector<std::vector<PaddleTensor>> *inputs) {

paddle/fluid/inference/utils/CMakeLists.txt

@@ -2,6 +2,3 @@ cc_library(benchmark SRCS benchmark.cc DEPS enforce)
 cc_test(test_benchmark SRCS benchmark_tester.cc DEPS benchmark)
 cc_binary(visualizer SRCS visualizer.cc DEPS analysis
     paddle_pass_builder ir_pass_manager pass graph_viz_pass analysis_passes)
-if(WIN32)
-  target_link_libraries(visualizer shlwapi)
-endif(WIN32)