move share_buffer kernel to phi (PaddlePaddle#48858)

* move share_buffer kernel to phi * fix ut * add source file * fix window links
lxsbupt · Dec 9, 2022 · 47c2821 · 47c2821
1 parent 5c52e38
commit 47c2821
Show file tree

Hide file tree

Showing 14 changed files with 144 additions and 148 deletions.
diff --git a/paddle/fluid/eager/auto_code_generator/generator/codegen_utils.py b/paddle/fluid/eager/auto_code_generator/generator/codegen_utils.py
@@ -81,6 +81,7 @@
     'str': 'std::string',
     'str[]': 'std::vector<std::string>',
     'float[]': 'std::vector<float>',
+    'bool[]': 'std::vector<bool>',
     'Place': 'paddle::Place',
     'DataLayout': 'phi::DataLayout',
     'DataType': 'paddle::experimental::DataType',

diff --git a/paddle/fluid/framework/operator.cc b/paddle/fluid/framework/operator.cc
@@ -3258,6 +3258,10 @@ void OperatorWithKernel::BuildPhiKernelContext(
             phi_kernel_context->EmplaceBackAttr(
                 PADDLE_GET_CONST(std::vector<int>, attr_iter->second));
             break;
+          case phi::AttributeType::BOOLS:
+            phi_kernel_context->EmplaceBackAttr(
+                PADDLE_GET_CONST(std::vector<bool>, attr_iter->second));
+            break;
           case phi::AttributeType::DATA_TYPE: {
             auto data_type = framework::TransToPhiDataType(
                 static_cast<framework::proto::VarType::Type>(

diff --git a/paddle/fluid/operators/CMakeLists.txt b/paddle/fluid/operators/CMakeLists.txt
@@ -207,7 +207,7 @@ elseif(WITH_ROCM)
 else()
     cc_test(test_leaky_relu_grad_grad_functor SRCS test_leaky_relu_grad_grad_functor.cc DEPS tensor device_context eigen3)
 endif()
-cc_test(share_buffer_op_cpp_test SRCS share_buffer_op_test.cc DEPS lod_tensor device_context share_buffer_op)
+cc_test(share_buffer_op_cpp_test SRCS share_buffer_op_test.cc DEPS lod_tensor device_context generated_op)
 
 cc_library(tensor_formatter SRCS tensor_formatter.cc DEPS ${OP_HEADER_DEPS})
 if (WITH_PYTHON)

diff --git a/paddle/fluid/operators/share_buffer_op.cc b/paddle/fluid/operators/share_buffer_op.cc
diff --git a/paddle/fluid/operators/share_buffer_op.h b/paddle/fluid/operators/share_buffer_op.h
diff --git a/paddle/fluid/operators/share_buffer_op_test.cc b/paddle/fluid/operators/share_buffer_op_test.cc
@@ -18,8 +18,15 @@
 #include "paddle/fluid/framework/operator.h"
 #include "paddle/fluid/framework/scope.h"
 #include "paddle/fluid/platform/place.h"
+#include "paddle/phi/core/kernel_registry.h"
 
-USE_OP(share_buffer);
+USE_OP_ITSELF(share_buffer);
+
+PD_DECLARE_KERNEL(share_buffer, CPU, ALL_LAYOUT);
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_DECLARE_KERNEL(share_buffer, GPU, ALL_LAYOUT);
+#endif
 
 namespace paddle {
 namespace framework {

diff --git a/paddle/phi/api/yaml/generator/api_base.py b/paddle/phi/api/yaml/generator/api_base.py
@@ -166,6 +166,7 @@ def parse_input_and_attr(self, api_name, args_config, optional_vars=[]):
             'float[]': 'const std::vector<float>&',
             'double': 'double',
             'bool': 'bool',
+            'bool[]': 'const std::vector<bool>&',
             'str': 'const std::string&',
             'str[]': 'const std::vector<std::string>&',
             'Place': 'const Place&',

diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -1019,6 +1019,13 @@
   extra :
     attrs : [bool use_mkldnn = false, str mkldnn_data_type = "float32"]
 
+- op : share_buffer
+  inputs :
+    x : X
+  outputs :
+    out : Out
+    xout : XOut
+
 - op : shuffle_channel
   backward : shuffle_channel_grad
   extra :

diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -952,3 +952,11 @@
   kernel :
     func : unfold
   backward : unfold_grad
+
+- op: share_buffer
+  args : (Tensor[] x, bool[] share_dims_and_dtype={})
+  output : Tensor[](out){x.size()}, Tensor[](xout){x.size()}
+  infer_meta :
+    func : ShareBufferInferMeta
+  kernel :
+    func : share_buffer
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -2668,6 +2668,30 @@ void UnchangedMultiInferMeta(const std::vector<const MetaTensor*>& x,
   }
 }
 
+void ShareBufferInferMeta(const std::vector<const MetaTensor*>& xs,
+                          const std::vector<bool>& share_dims_and_dtype,
+                          std::vector<MetaTensor*> outs,
+                          std::vector<MetaTensor*> xouts) {
+  if (share_dims_and_dtype.empty()) {
+    return;
+  }
+  PADDLE_ENFORCE_EQ(xs.size(),
+                    share_dims_and_dtype.size(),
+                    phi::errors::PermissionDenied(
+                        "The input(X) and attribute share_dims_and_dtype "
+                        "should have the same size, but got size of input(X) "
+                        "is %d and size of share_dims_and_dtype is %d.",
+                        xs.size(),
+                        share_dims_and_dtype.size()));
+
+  for (size_t i = 0; i < xs.size(); ++i) {
+    if (share_dims_and_dtype[i]) {
+      outs[i]->set_dims(xs[i]->dims());
+      outs[i]->set_dtype(xs[i]->dtype());
+    }
+  }
+}
+
 void UpdateLossScalingInferMeta(const std::vector<const MetaTensor*>& xs,
                                 const MetaTensor& found_infinite,
                                 const MetaTensor& prev_loss_scaling,

diff --git a/paddle/phi/infermeta/multiary.h b/paddle/phi/infermeta/multiary.h
@@ -479,6 +479,11 @@ void StackInferMeta(const std::vector<const MetaTensor*>& x,
 void UnchangedMultiInferMeta(const std::vector<const MetaTensor*>& x,
                              std::vector<MetaTensor*> out);
 
+void ShareBufferInferMeta(const std::vector<const MetaTensor*>& x,
+                          const std::vector<bool>& share_dims_and_dtype,
+                          std::vector<MetaTensor*> out,
+                          std::vector<MetaTensor*> xout);
+
 void UpdateLossScalingInferMeta(const std::vector<const MetaTensor*>& xs,
                                 const MetaTensor& found_infinite,
                                 const MetaTensor& prev_loss_scaling,

diff --git a/paddle/phi/kernels/shard_index_kernel.h b/paddle/phi/kernels/shard_index_kernel.h
@@ -14,6 +14,8 @@
 
 #pragma once
 
+#include <vector>
+
 #include "paddle/phi/core/dense_tensor.h"
 
 namespace phi {

diff --git a/paddle/phi/kernels/share_buffer_kernel.cc b/paddle/phi/kernels/share_buffer_kernel.cc
@@ -0,0 +1,67 @@
+// Copyright (c) 2022 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/phi/kernels/share_buffer_kernel.h"
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/visit_type.h"
+
+namespace phi {
+
+template <typename Context>
+void ShareBufferKernel(const Context &dev_ctx,
+                       const std::vector<const DenseTensor *> &x,
+                       const std::vector<bool> &share_dims_and_dtype,
+                       std::vector<DenseTensor *> out,
+                       std::vector<DenseTensor *> xout) {
+  PADDLE_ENFORCE_EQ(
+      x.size(),
+      out.size(),
+      phi::errors::PermissionDenied(
+          "The input(X) and Output(out) should have the same size, but got "
+          "size of Input(X) is %d and size of Output(out) is %d.",
+          x.size(),
+          out.size()));
+  for (size_t i = 0; i < x.size(); ++i) {
+    if (x[i] == nullptr || out[i] == nullptr) {
+      continue;
+    }
+    out[i]->ShareBufferWith(*x[i]);
+    VLOG(10) << "Share tensor buffer ";
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_GENERAL_KERNEL(share_buffer,
+                           CPU,
+                           ALL_LAYOUT,
+                           phi::ShareBufferKernel<phi::CPUContext>,
+                           ALL_DTYPE) {}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_GENERAL_KERNEL(share_buffer,
+                           GPU,
+                           ALL_LAYOUT,
+                           phi::ShareBufferKernel<phi::GPUContext>,
+                           ALL_DTYPE) {}
+#endif
+
+#ifdef PADDLE_WITH_XPU
+PD_REGISTER_GENERAL_KERNEL(share_buffer,
+                           XPU,
+                           ALL_LAYOUT,
+                           phi::ShareBufferKernel<phi::XPUContext>,
+                           ALL_DTYPE) {}
+#endif
diff --git a/paddle/fluid/operators/share_buffer_op.cu → paddle/phi/kernels/share_buffer_kernel.h b/paddle/fluid/operators/share_buffer_op.cu → paddle/phi/kernels/share_buffer_kernel.h
@@ -1,4 +1,4 @@
-// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2022 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.
@@ -12,7 +12,19 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/operators/share_buffer_op.h"
+#pragma once
 
-namespace ops = paddle::operators;
-REGISTER_OP_CUDA_KERNEL(share_buffer, ops::ShareBufferOpKernel<float>);
+#include <vector>
+
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+template <typename Context>
+void ShareBufferKernel(const Context &dev_ctx,
+                       const std::vector<const DenseTensor *> &x,
+                       const std::vector<bool> &share_dims_and_dtype,
+                       std::vector<DenseTensor *> out,
+                       std::vector<DenseTensor *> xout);
+
+}  // namespace phi