Merge branch 'develop' into fix-fp16-dtype-checking-for-argmin-op

cmake/external/protobuf.cmake

@@ -248,8 +248,13 @@ function(build_protobuf TARGET_NAME BUILD_FOR_HOST)
     # Change the tag to support building with vs2019
     set(PROTOBUF_TAG 01a05a53f40ca2ac5f0af10c6cc0810bee39b792)
   else()
-    set(PROTOBUF_REPOSITORY ${GIT_URL}/protocolbuffers/protobuf.git)
-    set(PROTOBUF_TAG v21.12)
+    if(WITH_PSLIB)
+      set(PROTOBUF_REPOSITORY "https://github.com/google/protobuf.git")
+      set(PROTOBUF_TAG "9f75c5aa851cd877fb0d93ccc31b8567a6706546")
+    else()
+      set(PROTOBUF_REPOSITORY ${GIT_URL}/protocolbuffers/protobuf.git)
+      set(PROTOBUF_TAG v21.12)
+    endif()
     if(WITH_GPU)
       if(${CMAKE_CUDA_COMPILER_VERSION} LESS 12.0
          AND ${CMAKE_CXX_COMPILER_VERSION} VERSION_GREATER 12.0)

cmake/external/pslib.cmake

@@ -23,9 +23,15 @@ if((NOT DEFINED PSLIB_VER) OR (NOT DEFINED PSLIB_URL))
   set(PSLIB_NAME
       "pslib"
       CACHE STRING "" FORCE)
-  set(PSLIB_URL
-      "https://pslib.bj.bcebos.com/pslib.tar.gz"
-      CACHE STRING "" FORCE)
+  if(NOT WITH_HETERPS)
+    set(PSLIB_URL
+        "https://pslib.bj.bcebos.com/pslib.tar.gz"
+        CACHE STRING "" FORCE)
+  else()
+    set(PSLIB_URL
+        "https://pslib.bj.bcebos.com/pslib_3631b2/pslib.tar.gz"
+        CACHE STRING "" FORCE)
+  endif()
 endif()
 message(STATUS "PSLIB_NAME: ${PSLIB_NAME}, PSLIB_URL: ${PSLIB_URL}")
 set(PSLIB_PREFIX_DIR "${THIRD_PARTY_PATH}/pslib")
@@ -67,3 +73,4 @@ ExternalProject_Add(
 add_library(pslib SHARED IMPORTED GLOBAL)
 set_property(TARGET pslib PROPERTY IMPORTED_LOCATION ${PSLIB_LIB})
 add_dependencies(pslib ${PSLIB_PROJECT})
+target_link_libraries(pslib INTERFACE ${PSLIB_LIB_DIR}/libjvm.so)

cmake/external/xpu.cmake

@@ -8,7 +8,7 @@ set(XPU_API_LIB_NAME "libxpuapi.so")
 set(XPU_RT_LIB_NAME "libxpurt.so")
 
 set(XPU_BASE_DATE "20230220")
-set(XPU_XCCL_BASE_VERSION "1.0.9")
+set(XPU_XCCL_BASE_VERSION "1.0.10")
 
 if(NOT DEFINED XPU_BASE_URL)
   set(XPU_BASE_URL_WITHOUT_DATE

paddle/fluid/distributed/CMakeLists.txt

@@ -7,7 +7,7 @@ if(WITH_PYTHON)
   file(MAKE_DIRECTORY
        ${PADDLE_BINARY_DIR}/python/paddle/distributed/fleet/proto)
   set(PSLIB_PROTO_DSTPATH
-      "${PADDLE_SOURCE_DIR}/python/paddle/incubate/fleet/parameter_server/pslib/"
+      "${PADDLE_SOURCE_DIR}/python/paddle/incubate/distributed/fleet/parameter_server/pslib/"
   )
   if(NOT WIN32)
     add_custom_command(

paddle/fluid/distributed/collective/process_group_bkcl.cc

@@ -367,6 +367,34 @@ std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::Reduce(
       use_calc_stream);
 }
 
+std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::ReduceScatter(
+    phi::DenseTensor* out_tensor,
+    const phi::DenseTensor& in_tensor,
+    const ReduceScatterOptions& opts,
+    bool sync_op,
+    bool use_calc_stream) {
+  return Collective(
+      out_tensor,
+      in_tensor,
+      [&](phi::DenseTensor* output,
+          const phi::DenseTensor& input,
+          BKCLContext_t comm,
+          const XPUStream& stream) {
+        return bkcl_reduce_scatter(
+            comm,
+            input.data(),
+            output->data(),
+            output->numel(),
+            platform::ToBKCLDataType(
+                framework::TransToProtoVarType(input.type())),
+            ToBKCLRedType(opts.reduce_op),
+            stream);
+      },
+      CommType::REDUCE_SCATTER,
+      sync_op,
+      use_calc_stream);
+}
+
 std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::Barrier(
     const BarrierOptions& opts) {
   PADDLE_ENFORCE_GE(opts.device_id,

paddle/fluid/distributed/collective/process_group_bkcl.h

@@ -115,6 +115,13 @@ class ProcessGroupBKCL : public ProcessGroupWithStream {
                                              bool sync_op,
                                              bool use_calc_stream) override;
 
+  std::shared_ptr<ProcessGroup::Task> ReduceScatter(
+      phi::DenseTensor* out_tensor,
+      const phi::DenseTensor& in_tensor,
+      const ReduceScatterOptions& opts,
+      bool sync_op,
+      bool use_calc_stream) override;
+
   std::shared_ptr<ProcessGroup::Task> Recv(phi::DenseTensor* tensor,
                                            int src_rank,
                                            int64_t offset,

paddle/fluid/framework/ir/conv_bn_fuse_pass.cc

@@ -18,6 +18,7 @@
 
 #include "paddle/fluid/framework/convert_utils.h"
 #include "paddle/fluid/framework/eigen.h"
+#include "paddle/fluid/framework/ir/mkldnn/mkldnn_pass_util.h"
 #include "paddle/fluid/framework/lod_tensor.h"
 #include "paddle/fluid/framework/op_version_registry.h"
 #include "paddle/fluid/framework/tensor.h"
@@ -413,7 +414,7 @@ void ConvBNFusePass::ApplyImpl(ir::Graph* graph) const {
     if (is_mkldnn) {
       if (conv->Op()->Type() == "conv2d" ||
           conv->Op()->Type() == "depthwise_conv2d") {
-        conv->Op()->SetType("fused_conv2d");
+        ConvertToFusedOp(conv->Op());
       }
       if (mkldnn_with_bias) {
         // reuse existing conv bias node

paddle/fluid/framework/ir/mkldnn/activation_onednn_fuse_pass.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+#include "paddle/fluid/framework/ir/mkldnn/mkldnn_pass_util.h"
 #include "paddle/fluid/framework/op_desc.h"
 
 namespace paddle {

paddle/fluid/framework/ir/mkldnn/conv_activation_mkldnn_fuse_pass.cc

@@ -13,7 +13,6 @@
 // limitations under the License.
 
 #include "paddle/fluid/framework/ir/mkldnn/conv_activation_mkldnn_fuse_pass.h"
-
 #include "paddle/fluid/framework/ir/mkldnn/activation_onednn_fuse_pass.h"
 #include "paddle/fluid/framework/op_version_registry.h"
 #include "paddle/utils/string/pretty_log.h"
@@ -66,7 +65,7 @@ void ConvActivationMkldnnFusePass::FuseConvAct(Graph* graph,
     OpDesc* conv_op = conv->Op();
 
     if (conv_op->Type() == "conv2d") {
-      conv_op->SetType("fused_conv2d");
+      ConvertToFusedOp(conv_op);
     }
 
     SetActivationAttrs(conv_op, activation->Op(), act_type);
@@ -138,7 +137,7 @@ void ConvActivationMkldnnFusePass::FuseConvConcatAct(
     for (auto node : concat_inputs) {
       OpDesc* conv_op = node->inputs[0]->Op();
       if (conv_op->Type() == "conv2d") {
-        conv_op->SetType("fused_conv2d");
+        ConvertToFusedOp(conv_op);
       }
 
       SetActivationAttrs(conv_op, activation_op->Op(), act_type);

paddle/fluid/framework/ir/mkldnn/conv_elementwise_add_mkldnn_fuse_pass.cc

@@ -15,6 +15,7 @@
 #include "paddle/fluid/framework/ir/mkldnn/conv_elementwise_add_mkldnn_fuse_pass.h"
 
 #include "paddle/fluid/framework/ir/graph_traits.h"
+#include "paddle/fluid/framework/ir/mkldnn/mkldnn_pass_util.h"
 #include "paddle/fluid/framework/op_version_registry.h"
 #include "paddle/utils/string/pretty_log.h"
 
@@ -166,7 +167,7 @@ GraphWithStats ResidualConnectionMKLDNNFusePass::FuseConv(
     }
 
     if (conv_op->Op()->Type() == "conv2d") {
-      conv_op->Op()->SetType("fused_conv2d");
+      ConvertToFusedOp(conv_op->Op());
     }
 
     conv_op->Op()->SetInput("ResidualData", {residual_data->Name()});
@@ -259,7 +260,7 @@ GraphWithStats ResidualConnectionMKLDNNFusePass::FuseProjectionConv(
     if (HasFusedActivation(residual_conv_op)) return;
 
     if (residual_conv_op->Op()->Type() == "conv2d") {
-      residual_conv_op->Op()->SetType("fused_conv2d");
+      ConvertToFusedOp(residual_conv_op->Op());
     }
     residual_conv_op->Op()->SetInput("ResidualData", {projection_node->Name()});
     residual_conv_op->Op()->SetOutput("Output", {elementwise_out->Name()});

paddle/fluid/framework/ir/mkldnn/cpu_quantize_pass.cc

@@ -453,7 +453,7 @@ void CPUQuantizePass::QuantizeConv(Graph* graph,
     VLOG(4) << "Quantize conv2d op";
     GET_IR_NODE_FROM_SUBGRAPH(conv_op, conv_op, conv_pattern);
     if (conv_op->Op()->Type() == "conv2d") {
-      conv_op->Op()->SetType("fused_conv2d");
+      ConvertToFusedOp(conv_op->Op());
     }
 
     // skip if should not be quantized

paddle/fluid/framework/ir/mkldnn/cpu_quantize_placement_pass.cc

@@ -15,25 +15,14 @@ limitations under the License. */
 #include "paddle/fluid/framework/ir/mkldnn/cpu_quantize_placement_pass.h"
 
 #include <unordered_set>
+#include "paddle/fluid/framework/ir/mkldnn/mkldnn_pass_util.h"
 
 namespace paddle {
 namespace framework {
 namespace ir {
 
 class Graph;
 
-void ReplaceWithFusedOp(Node* op) {
-  const std::string matmul_type = op->Op()->Type();
-  if (matmul_type == "matmul" || matmul_type == "matmul_v2") {
-    op->Op()->SetType("fused_matmul");
-    if (matmul_type == "matmul") {
-      op->Op()->SetAttr("trans_x", op->Op()->GetAttr("transpose_X"));
-      op->Op()->SetAttr("trans_y", op->Op()->GetAttr("transpose_Y"));
-      op->Op()->SetAttr("matmul_alpha", op->Op()->GetAttr("alpha"));
-    }
-  }
-}
-
 void CPUQuantizePlacementPass::ApplyImpl(ir::Graph* graph) const {
   VLOG(3) << "Marks operators which are to be quantized.";
   std::unordered_set<std::string> supported_op_types =
@@ -97,7 +86,7 @@ void CPUQuantizePlacementPass::ApplyImpl(ir::Graph* graph) const {
       return;
     }
 
-    ReplaceWithFusedOp(op);
+    ConvertToFusedOp(op->Op());
     op->Op()->SetAttr("mkldnn_data_type", std::string("int8"));
   };
   gpd(graph, handler);

paddle/fluid/framework/ir/mkldnn/cpu_quantize_squash_pass.cc

@@ -18,6 +18,7 @@
 #include <string>
 #include <vector>
 
+#include "paddle/fluid/framework/ir/mkldnn/mkldnn_pass_util.h"
 #include "paddle/fluid/platform/mkldnn_helper.h"
 #include "paddle/phi/core/enforce.h"
 #include "paddle/utils/string/pretty_log.h"
@@ -355,7 +356,7 @@ void CPUQuantizeSquashPass::OpDequantSquash(Graph* graph) const {
 
       if (output_name.empty()) return;
       if (any_op->Op()->Type() == "conv2d") {
-        any_op->Op()->SetType("fused_conv2d");
+        ConvertToFusedOp(any_op->Op());
       }
       any_op->Op()->SetAttr("force_fp32_output", true);
       any_op->Op()->SetOutput(output_name,

paddle/fluid/framework/ir/mkldnn/int8_scale_calculation_mkldnn_pass.cc

@@ -14,6 +14,7 @@
 
 #include "paddle/fluid/framework/ir/mkldnn/int8_scale_calculation_mkldnn_pass.h"
 
+#include "paddle/fluid/framework/ir/mkldnn/mkldnn_pass_util.h"
 #include "paddle/fluid/framework/op_version_registry.h"
 #include "paddle/fluid/platform/mkldnn_helper.h"
 #include "paddle/phi/core/enforce.h"
@@ -124,7 +125,7 @@ void Int8ScaleCalculationMkldnnPass::Int8ScaleImpl(
     }
     GET_IR_NODE_FROM_SUBGRAPH(conv_op, conv_op, conv_pattern);
     if (conv_op->Op()->Type() == "conv2d") {
-      conv_op->Op()->SetType("fused_conv2d");
+      ConvertToFusedOp(conv_op->Op());
     }
 
     if (!platform::HasOpINT8DataType(conv_op->Op()) ||

paddle/fluid/framework/ir/mkldnn/matmul_activation_mkldnn_fuse_pass.cc

@@ -15,6 +15,7 @@
 #include "paddle/fluid/framework/ir/mkldnn/matmul_activation_mkldnn_fuse_pass.h"
 
 #include "paddle/fluid/framework/ir/mkldnn/activation_onednn_fuse_pass.h"
+#include "paddle/fluid/framework/ir/mkldnn/mkldnn_pass_util.h"
 #include "paddle/fluid/framework/op_version_registry.h"
 #include "paddle/utils/string/pretty_log.h"
 
@@ -63,13 +64,7 @@ void MatmulActivationMkldnnFusePass::FuseMatmulAct(
 
     OpDesc* matmul_op = matmul->Op();
 
-    matmul_op->SetType("fused_matmul");
-    if (matmul_type == "matmul") {
-      matmul_op->SetAttr("trans_x", matmul_op->GetAttr("transpose_X"));
-      matmul_op->SetAttr("trans_y", matmul_op->GetAttr("transpose_Y"));
-      matmul_op->SetAttr("matmul_alpha", matmul_op->GetAttr("alpha"));
-    }
-
+    ConvertToFusedOp(matmul_op);
     SetActivationAttrs(matmul_op, activation->Op(), act_type);
     matmul_op->SetOutput("Out", {activation_out->Name()});
 

paddle/fluid/framework/ir/mkldnn/matmul_elementwise_add_mkldnn_fuse_pass.cc

@@ -15,6 +15,7 @@
 #include "paddle/fluid/framework/ir/mkldnn/matmul_elementwise_add_mkldnn_fuse_pass.h"
 
 #include "paddle/fluid/framework/ir/graph_traits.h"
+#include "paddle/fluid/framework/ir/mkldnn/mkldnn_pass_util.h"
 #include "paddle/fluid/framework/op_version_registry.h"
 #include "paddle/utils/string/pretty_log.h"
 
@@ -65,12 +66,7 @@ void MatmulElementwiseAddMKLDNNFusePass::FuseMatmulElementwiseAdd(
       return;
     }
 
-    matmul->Op()->SetType("fused_matmul");
-    if (matmul_type == "matmul") {
-      matmul->Op()->SetAttr("trans_x", matmul->Op()->GetAttr("transpose_X"));
-      matmul->Op()->SetAttr("trans_y", matmul->Op()->GetAttr("transpose_Y"));
-      matmul->Op()->SetAttr("matmul_alpha", matmul->Op()->GetAttr("alpha"));
-    }
+    ConvertToFusedOp(matmul->Op());
     matmul->Op()->SetInput("ResidualData", {elementwise_addend->Name()});
     matmul->Op()->SetOutput("Out", {elementwise_add_out->Name()});
 

paddle/fluid/framework/ir/mkldnn/matmul_transpose_reshape_mkldnn_fuse_pass.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "paddle/fluid/framework/ir/mkldnn/matmul_transpose_reshape_mkldnn_fuse_pass.h"
+#include "paddle/fluid/framework/ir/mkldnn/mkldnn_pass_util.h"
 #include "paddle/fluid/framework/op_version_registry.h"
 #include "paddle/phi/core/enforce.h"
 #include "paddle/utils/string/pretty_log.h"
@@ -84,12 +85,7 @@ void MatmulTransposeReshapeMKLDNNPass::Fuse(
     }
 
     OpDesc *matmul_desc = matmul_op->Op();
-    matmul_desc->SetType("fused_matmul");
-    if (matmul_type == "matmul") {
-      matmul_desc->SetAttr("trans_x", matmul_desc->GetAttr("transpose_X"));
-      matmul_desc->SetAttr("trans_y", matmul_desc->GetAttr("transpose_Y"));
-      matmul_desc->SetAttr("matmul_alpha", matmul_desc->GetAttr("alpha"));
-    }
+    ConvertToFusedOp(matmul_desc);
     matmul_desc->SetOutput("Out", {reshape_out->Name()});
     matmul_desc->SetAttr("fused_reshape_Out", reshape_shape);
     matmul_desc->SetAttr("fused_transpose_Out", transpose_axis);

paddle/fluid/framework/ir/mkldnn/mkldnn_pass_util.h

@@ -155,6 +155,26 @@ static void GetInfoFromTheFirstOp(ir::Graph* graph,
   }
 }
 
+inline void ConvertToFusedOp(OpDesc* op) {
+  const std::map<std::string, std::string> fused_ops = {
+      {"conv2d", "fused_conv2d"},
+      {"depthwise_conv2d", "fused_conv2d"},
+      {"matmul", "fused_matmul"},
+      {"matmul_v2", "fused_matmul"}};
+
+  if (op->Type() == "matmul") {
+    op->SetAttr("trans_x", op->GetAttr("transpose_X"));
+    op->SetAttr("trans_y", op->GetAttr("transpose_Y"));
+    op->SetAttr("matmul_alpha", op->GetAttr("alpha"));
+  }
+
+  auto it = fused_ops.find(op->Type());
+  if (it != fused_ops.end()) {
+    op->SetType(it->second);
+    VLOG(3) << "Converted " << it->first << " to " << it->second;
+  }
+}
+
 }  // namespace ir
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/ir/mkldnn/operator_scale_onednn_fuse_pass.cc

@@ -14,6 +14,7 @@
 
 #include "paddle/fluid/framework/ir/mkldnn/operator_scale_onednn_fuse_pass.h"
 
+#include "paddle/fluid/framework/ir/mkldnn/mkldnn_pass_util.h"
 #include "paddle/fluid/framework/op_version_registry.h"
 #include "paddle/phi/backends/onednn/onednn_reuse.h"
 #include "paddle/utils/string/pretty_log.h"
@@ -86,17 +87,8 @@ void FuseOperatorScaleOneDNNPass::FuseScale(Graph *graph,
       scale = *(scale_tensor->data<float>());
     }
 
-    if (op_type == "matmul") {
-      operator_op->Op()->SetType("fused_matmul");
-      operator_op->Op()->SetAttr("trans_x",
-                                 operator_op->Op()->GetAttr("transpose_X"));
-      operator_op->Op()->SetAttr("trans_y",
-                                 operator_op->Op()->GetAttr("transpose_Y"));
-      operator_op->Op()->SetAttr("matmul_alpha",
-                                 operator_op->Op()->GetAttr("alpha"));
-    }
-    if (op_type == "matmul_v2") {
-      operator_op->Op()->SetType("fused_matmul");
+    if (op_type == "matmul" || op_type == "matmul_v2") {
+      ConvertToFusedOp(operator_op->Op());
     }
 
     operator_op->Op()->SetAttr("fused_output_scale", scale);