[oneDNN] Reshape attr_axes when going to oneDNN kernel (#59641)

Author: LLee233

paddle/phi/kernels/onednn/squeeze_kernel.cc

@@ -59,7 +59,29 @@ void SqueezeInferKernel(const Context& dev_ctx,
                         const IntArray& axes,
                         DenseTensor* out) {
   auto x_dims = x.dims();
+  auto x_dims_tz = x_dims.size();
   std::vector<int32_t> tmp(axes.GetData().begin(), axes.GetData().end());
+
+  // Currently there is only tranformation for tensors, while attr axes still
+  // follows default dtype instead of oneDNN dtype, so here manually change it
+  if ((x_dims_tz >= 3) &&
+      (phi::OneDNNContext::tls().get_cur_paddle_data_layout() ==
+           phi::DataLayout::NDHWC ||
+       phi::OneDNNContext::tls().get_cur_paddle_data_layout() ==
+           phi::DataLayout::NHWC)) {
+    int axes_size = tmp.size();
+    for (int i = 0; i < axes_size; i++) {
+      if (tmp[i] < 0) {
+        tmp[i] += x_dims_tz;
+      }
+      if (tmp[i] >= 1 && tmp[i] < (x_dims_tz - 1)) {
+        tmp[i] += 1;
+      } else if (tmp[i] == (x_dims_tz - 1)) {
+        tmp[i] = 1;
+      }
+    }
+  }
+
   auto out_dims = funcs::GetOutputSqueezeShape(tmp, x_dims, true);
   ExecuteSqueeze<T, Context>(dev_ctx, x, x_dims, out_dims, out);
 }

test/cpp/fluid/mkldnn/CMakeLists.txt

@@ -104,3 +104,18 @@ cc_test(
        scope
        device_context
        enforce)
+
+cc_test(
+  test_mkldnn_squeeze
+  SRCS test_mkldnn_squeeze.cc
+  DEPS fleet_executor
+       conditional_block_op
+       standalone_executor
+       executor
+       op_registry
+       generated_static_op
+       generated_op
+       phi
+       scope
+       device_context
+       enforce)

test/cpp/fluid/mkldnn/test_mkldnn_squeeze.cc

@@ -0,0 +1,106 @@
+/* Copyright (c) 2023 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 <gtest/gtest.h>
+
+#include <fstream>
+
+#include "paddle/fluid/framework/lod_tensor.h"
+#include "paddle/fluid/framework/naive_executor.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+#include "paddle/fluid/framework/scope.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+template <typename DataType>
+void AddVarToScope(const std::string var_name,
+                   paddle::framework::Scope* scope,
+                   const paddle::framework::DDim& dims) {
+  std::random_device seed;
+  std::default_random_engine engine(seed());
+  std::uniform_real_distribution<float> dist(0, 100);
+
+  phi::DenseTensor tmp_tensor;
+  auto* tmp_data =
+      tmp_tensor.mutable_data<DataType>(dims, paddle::platform::CPUPlace());
+  auto* tensor = scope->Var(var_name)->GetMutable<phi::DenseTensor>();
+  tensor->mutable_data<DataType>(dims, paddle::platform::CPUPlace());
+  for (auto i = 0; i < tensor->numel(); ++i) {
+    tmp_data[i] = static_cast<DataType>(dist(engine));
+  }
+  paddle::framework::TensorCopySync(
+      tmp_tensor, paddle::platform::CPUPlace(), tensor);
+}
+void test_squeeze() {
+  framework::Scope scope;
+  paddle::platform::CPUPlace cpu_place;
+  // Prepare Op description
+  framework::OpDesc desc;
+  // We assume it is kNHWC, so that can use this transformation
+  phi::OneDNNContext::tls().set_cur_paddle_data_layout(DataLayout::kNHWC);
+  desc.SetType("squeeze2");
+  desc.SetInput("X", {"squeeze-X"});
+  desc.SetOutput("Out", {"squeeze-Out"});
+  // DataLayout::kNHWC will make it become {2, 3, 2, 1}
+  AddVarToScope<float>("squeeze-X", &scope, {2, 2, 1, 3});
+  AddVarToScope<float>("squeeze-Out", &scope, {2, 3, 2});
+  // transform will make it become -1
+  std::vector<int> axes({-2});
+
+  desc.SetAttr("axes", axes);
+  desc.SetAttr("use_mkldnn", true);
+
+  auto op = paddle::framework::OpRegistry::CreateOp(desc);
+
+  op->Run(scope, cpu_place);
+}
+
+void test_squeeze2() {
+  framework::Scope scope;
+  paddle::platform::CPUPlace cpu_place;
+  // Prepare Op description
+  framework::OpDesc desc;
+  // We assume it is HNWC, so that can use this transformation
+  phi::OneDNNContext::tls().set_cur_paddle_data_layout(DataLayout::kNHWC);
+  desc.SetType("squeeze2");
+  desc.SetInput("X", {"squeeze-X"});
+  desc.SetOutput("Out", {"squeeze-Out"});
+  // DataLayout::kNHWC will make it become {2, 1, 3, 2}
+  AddVarToScope<float>("squeeze-X", &scope, {2, 3, 2, 1});
+  AddVarToScope<float>("squeeze-Out", &scope, {2, 3, 2});
+  // transform will make it become -3(1)
+  std::vector<int> axes({-1});
+
+  desc.SetAttr("axes", axes);
+  desc.SetAttr("use_mkldnn", true);
+
+  auto op = paddle::framework::OpRegistry::CreateOp(desc);
+
+  op->Run(scope, cpu_place);
+}
+
+TEST(SqueezeOpConverter, normal) { test_squeeze(); }
+TEST(SqueezeOpConverter_2, normal) { test_squeeze2(); }
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle
+
+USE_OP_ITSELF(squeeze2);
+PD_DECLARE_KERNEL(squeeze_infer, OneDNN, ONEDNN);