PaddlePaddle
diff --git a/‎paddle/fluid/ir_adaptor/translator/op_compat_gen.py
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/fluid/ir_adaptor/translator/op_compat_gen.py
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/fluid/operators/mkldnn/reshape_mkldnn_op.cc
Copy file name to clipboard
-42Lines changed: 0 additions & 42 deletions b/‎paddle/fluid/operators/mkldnn/reshape_mkldnn_op.cc
Copy file name to clipboard
-42Lines changed: 0 additions & 42 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
Copy file name to clipboard
+12-16Lines changed: 12 additions & 16 deletions b/‎paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
Copy file name to clipboard
+12-16Lines changed: 12 additions & 16 deletions
diff --git a/‎paddle/phi/kernels/onednn/flatten_kernel.cc
Copy file name to clipboard
+85Lines changed: 85 additions & 0 deletions b/‎paddle/phi/kernels/onednn/flatten_kernel.cc
Copy file name to clipboard
+85Lines changed: 85 additions & 0 deletions
diff --git a/‎paddle/phi/kernels/onednn/flatten_kernel_grad.cc
Copy file name to clipboard
+61Lines changed: 61 additions & 0 deletions b/‎paddle/phi/kernels/onednn/flatten_kernel_grad.cc
Copy file name to clipboard
+61Lines changed: 61 additions & 0 deletions
diff --git a/‎paddle/phi/kernels/onednn/reshape_grad_kernel.cc
Copy file name to clipboard
+56Lines changed: 56 additions & 0 deletions b/‎paddle/phi/kernels/onednn/reshape_grad_kernel.cc
Copy file name to clipboard
+56Lines changed: 56 additions & 0 deletions
@@ -84,7 +84,7 @@ def insert_new_mutable_attributes(
                     if k == 'tensor_name' or k == 'tensors_name':
                         op_mutable_attribute_infos[op_name][
                             attribute_name
-                        ].append(v)
+                        ].insert(0, v)
 
         _, legacy_name = insert_new_mappings(op_compat_item["op"])
         legacy_backward_op_names = []
 
@@ -20,10 +20,8 @@ limitations under the License. */
 namespace {
 enum class ReshapeKernelOpName {
   reshape,
-  reshape2,
   squeeze,
   flatten,
-  flatten2,
 };
 }  // anonymous namespace
 
@@ -105,9 +103,6 @@ class ReshapeMKLDNNKernel : public framework::OpKernel<T> {
         InferShapeSqueezeOp(ctx, x_dims, out_dims);
         break;
       case ReshapeKernelOpName::flatten:
-      case ReshapeKernelOpName::flatten2:
-        InferShapeFlattenOp(ctx, x_dims, out_dims);
-        break;
       default:
         PADDLE_THROW(paddle::platform::errors::OutOfRange(
             "Reshape kernel doesn not support that operator name"));
@@ -317,18 +312,12 @@ class ReshapeGradMKLDNNKernel : public ReshapeMKLDNNKernel<T, op_name> {
       case ReshapeKernelOpName::reshape:
         InferShapeReshapeSqueezeGradOp(ctx, x_dims);
         break;
-      case ReshapeKernelOpName::reshape2:
-        InferShapeReshape2Flatten2GradOp(ctx, x_dims);
-        break;
       case ReshapeKernelOpName::squeeze:
         InferShapeReshapeSqueezeGradOp(ctx, x_dims);
         break;
       case ReshapeKernelOpName::flatten:
         InferShapeFlattenGradOp(ctx, x_dims);
         break;
-      case ReshapeKernelOpName::flatten2:
-        InferShapeReshape2Flatten2GradOp(ctx, x_dims);
-        break;
       default:
         PADDLE_THROW(paddle::platform::errors::OutOfRange(
             "Reshape grad kernel doesn not support that operator name"));
@@ -342,13 +331,6 @@ class ReshapeGradMKLDNNKernel : public ReshapeMKLDNNKernel<T, op_name> {
     dx_dims = dx->dims();
   }
 
-  void InferShapeReshape2Flatten2GradOp(
-      const framework::ExecutionContext& ctx,
-      framework::DDim& dx_dims) const {  // NOLINT
-    auto xshape_dims = ctx.Input<phi::DenseTensor>("XShape")->dims();
-    dx_dims = common::slice_ddim(xshape_dims, 1, xshape_dims.size());
-  }
-
   void InferShapeFlattenGradOp(const framework::ExecutionContext& ctx,
                                framework::DDim& dx_dims) const {  // NOLINT
     dx_dims = ctx.Input<phi::DenseTensor>("X")->dims();
@@ -390,14 +372,6 @@ REGISTER_OP_KERNEL(
     ops::ReshapeGradMKLDNNKernel<paddle::platform::bfloat16,
                                  ReshapeKernelOpName::reshape>);
 
-REGISTER_OP_KERNEL(
-    reshape2_grad,
-    MKLDNN,
-    phi::CPUPlace,
-    ops::ReshapeGradMKLDNNKernel<float, ReshapeKernelOpName::reshape2>,
-    ops::ReshapeGradMKLDNNKernel<paddle::platform::bfloat16,
-                                 ReshapeKernelOpName::reshape2>);
-
 REGISTER_OP_KERNEL(
     flatten,
     MKLDNN,
@@ -413,19 +387,3 @@ REGISTER_OP_KERNEL(
     ops::ReshapeGradMKLDNNKernel<float, ReshapeKernelOpName::flatten>,
     ops::ReshapeGradMKLDNNKernel<paddle::platform::bfloat16,
                                  ReshapeKernelOpName::flatten>);
-
-REGISTER_OP_KERNEL(
-    flatten2,
-    MKLDNN,
-    phi::CPUPlace,
-    ops::ReshapeMKLDNNKernel<float, ReshapeKernelOpName::flatten2>,
-    ops::ReshapeMKLDNNKernel<paddle::platform::bfloat16,
-                             ReshapeKernelOpName::flatten2>);
-
-REGISTER_OP_KERNEL(
-    flatten2_grad,
-    MKLDNN,
-    phi::CPUPlace,
-    ops::ReshapeGradMKLDNNKernel<float, ReshapeKernelOpName::flatten2>,
-    ops::ReshapeGradMKLDNNKernel<paddle::platform::bfloat16,
-                                 ReshapeKernelOpName::flatten2>);
@@ -86,13 +86,11 @@
 - op : fc
   extra_args : bool ALL_KERNELS_MUST_COMPUTE_RUNTIME_SHAPE=true, bool use_quantizer=false, str mkldnn_data_type="float32", float scale_in=1.0, float[] scale_weights={1.0f}, float scale_out=1.0, bool force_fp32_output=false
 
-# - op : flatten
-
-# - op : flatten_grad
-
-# - op : flatten2
+- op : flatten
+  extra_args : str mkldnn_data_type="float32"
 
-# - op : flatten2_grad
+- op : flatten_grad
+  extra_args : str mkldnn_data_type="float32"
 
 - op : full
 
@@ -226,13 +224,11 @@
 - op : relu6_grad
   extra_args : float threshold=6.0
 
-# - op : reshape
-
-# - op : reshape_grad
+- op : reshape
+  extra_args : str mkldnn_data_type="float32", bool use_quantizer=false
 
-# - op : reshape_infer
-
-# - op : reshape2_grad
+- op : reshape_grad
+  extra_args : str mkldnn_data_type="float32", bool use_quantizer=false
 
 - op : round
 
@@ -273,11 +269,11 @@
 
 - op : sqrt_grad
 
-# - op : squeeze
-
-# - op : squeeze_grad
+- op : squeeze
+  extra_args : str mkldnn_data_type="float32"
 
-# - op : squeeze_infer
+- op : squeeze_grad
+  extra_args : str mkldnn_data_type="float32"
 
 # - op : stack
 
 
@@ -0,0 +1,85 @@
+// Copyright (c) 2024 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
+//
+//     https://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/flatten_kernel.h"
+
+#include "paddle/phi/backends/onednn/onednn_reuse.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void ExecuteFlatten(const Context& dev_ctx,
+                    const DenseTensor& x,
+                    const DDim& x_dims,
+                    const DDim& out_dims,
+                    DenseTensor* out) {
+  auto x_vec_dims = common::vectorize(x_dims);
+
+  funcs::ReorderOneDNNHandler reorder_handler(
+      x_vec_dims,
+      x.dtype(),
+      funcs::ToOneDNNDataType(x.dtype()),
+      dev_ctx.GetEngine());
+
+  auto reorder_src_memory_p = reorder_handler.AcquireSrcMemory(
+      x.mem_desc(), funcs::to_void_cast(x.data<T>()));
+  out->Resize(x_dims);  // to match x numel, format is changed later
+  // reorder is done into a plain tag to allow usage with blocked formats
+  auto reorder_dst_memory_p = reorder_handler.AcquireDstMemory(
+      out, funcs::GetPlainOneDNNFormat(x_dims.size()), dev_ctx.GetPlace());
+  auto reorder_p = reorder_handler.AcquireReorder(reorder_dst_memory_p,
+                                                  reorder_src_memory_p);
+  auto& astream = OneDNNContext::tls().get_stream();
+  reorder_p->execute(astream, *reorder_src_memory_p, *reorder_dst_memory_p);
+  astream.wait();
+
+  out->Resize(out_dims);
+
+  auto reshape_dims = out_dims.size() != 0 ? common::vectorize(out_dims)
+                                           : std::vector<int64_t>{1};
+  out->set_mem_desc(reorder_dst_memory_p->get_desc().reshape(reshape_dims));
+}
+
+template <typename T, typename Context>
+void FlattenInferKernel(const Context& dev_ctx,
+                        const DenseTensor& x,
+                        int start_axis,
+                        int stop_axis,
+                        DenseTensor* out) {
+  auto x_dims = x.dims();
+  auto out_dims = out->dims();
+  ExecuteFlatten<T, Context>(dev_ctx, x, x_dims, out_dims, out);
+}
+
+template <typename T, typename Context>
+void FlattenKernel(const Context& dev_ctx,
+                   const DenseTensor& x,
+                   int start_axis,
+                   int stop_axis,
+                   DenseTensor* out,
+                   DenseTensor* xshape UNUSED) {
+  FlattenInferKernel<T, Context>(dev_ctx, x, start_axis, stop_axis, out);
+}
+
+}  // namespace phi
+PD_REGISTER_KERNEL(flatten_infer,
+                   OneDNN,
+                   ONEDNN,
+                   phi::FlattenInferKernel,
+                   float,
+                   phi::dtype::bfloat16) {}
+
+PD_REGISTER_KERNEL(
+    flatten, OneDNN, ONEDNN, phi::FlattenKernel, float, phi::dtype::bfloat16) {}
@@ -0,0 +1,61 @@
+// Copyright (c) 2024 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
+//
+//     https://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/flatten_grad_kernel.h"
+
+#include "paddle/phi/backends/onednn/onednn_reuse.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void FlattenGradKernel(const Context& dev_ctx,
+                       const DenseTensor& xshape,
+                       const DenseTensor& out_grad,
+                       DenseTensor* x_grad) {
+  auto out_grad_vec_dims = out_grad.dims().size() != 0
+                               ? common::vectorize(out_grad.dims())
+                               : std::vector<int64_t>{1};
+
+  auto out_grad_type = funcs::ToOneDNNDataType(out_grad.dtype());
+
+  funcs::ReorderOneDNNHandler reorder_handler(
+      out_grad_vec_dims, out_grad.dtype(), out_grad_type, dev_ctx.GetEngine());
+
+  auto reorder_src_memory_p = reorder_handler.AcquireSrcMemory(
+      out_grad.mem_desc(), funcs::to_void_cast(out_grad.data<T>()));
+  auto reorder_dst_memory_p = reorder_handler.AcquireDstMemory(
+      x_grad,
+      funcs::GetPlainOneDNNFormat(out_grad_vec_dims.size()),
+      dev_ctx.GetPlace());
+  auto reorder_p = reorder_handler.AcquireReorder(reorder_dst_memory_p,
+                                                  reorder_src_memory_p);
+
+  auto& astream = OneDNNContext::tls().get_stream();
+  reorder_p->execute(astream, *reorder_src_memory_p, *reorder_dst_memory_p);
+  astream.wait();
+
+  auto x_grad_dims = slice_ddim(xshape.dims(), 1, xshape.dims().size());
+  x_grad->Resize(x_grad_dims);
+  reorder_dst_memory_p->get_desc().reshape(common::vectorize(x_grad_dims));
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(flatten_grad,
+                   OneDNN,
+                   ONEDNN,
+                   phi::FlattenGradKernel,
+                   float,
+                   phi::dtype::bfloat16) {}
@@ -0,0 +1,56 @@
+/* Copyright (c) 2024 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
+    https://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/backends/onednn/onednn_reuse.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void ReshapeGradKernel(const Context& dev_ctx,
+                       const DenseTensor& out_grad,
+                       DenseTensor* x_grad) {
+  auto out_grad_vec_dims = out_grad.dims().size() != 0
+                               ? common::vectorize(out_grad.dims())
+                               : std::vector<int64_t>{1};
+
+  auto out_grad_type = funcs::ToOneDNNDataType(out_grad.dtype());
+
+  funcs::ReorderOneDNNHandler reorder_handler(
+      out_grad_vec_dims, out_grad.dtype(), out_grad_type, dev_ctx.GetEngine());
+
+  auto reorder_src_memory_p = reorder_handler.AcquireSrcMemory(
+      out_grad.mem_desc(), funcs::to_void_cast(out_grad.data<T>()));
+  auto reorder_dst_memory_p = reorder_handler.AcquireDstMemory(
+      x_grad,
+      funcs::GetPlainOneDNNFormat(out_grad_vec_dims.size()),
+      dev_ctx.GetPlace());
+  auto reorder_p = reorder_handler.AcquireReorder(reorder_dst_memory_p,
+                                                  reorder_src_memory_p);
+
+  auto& astream = OneDNNContext::tls().get_stream();
+  reorder_p->execute(astream, *reorder_src_memory_p, *reorder_dst_memory_p);
+  astream.wait();
+
+  auto grad_shape = x_grad->dims().size() == 0
+                        ? std::vector<int64_t>{1}
+                        : phi::vectorize<int64_t>(x_grad->dims());
+  reorder_dst_memory_p->get_desc().reshape(grad_shape);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(reshape_grad,
+                   OneDNN,
+                   ONEDNN,
+                   phi::ReshapeGradKernel,
+                   float,
+                   phi::dtype::bfloat16) {}