【PaddlePaddle Hackathon 77】Add squeeze op

cinn/frontend/net_builder.cc

@@ -78,6 +78,14 @@ Variable NetBuilder::ReduceSum(const Variable& x, const std::vector<int>& dim, b
   return Reduce(x, ReduceKind::kSum, dim, keep_dim);
 }
 
+Variable NetBuilder::Squeeze(const Variable& operand, const std::vector<int>& axes) {
+  Instruction instr("squeeze", {operand});
+  instr.SetAttr("axes", axes);
+  InferShape(instr);
+  AppendInstruction(instr);
+  return instr.GetOutput(0);
+}
+
 Variable NetBuilder::Conv2d(const Variable& a,
                             const Variable& b,
                             const std::vector<int>& strides,

cinn/frontend/net_builder.h

@@ -82,6 +82,11 @@ class NetBuilder : public BaseBuilder {
    */
   Variable ReduceSum(const Variable& x, const std::vector<int>& dim, bool keep_dim = false);
 
+  /**
+   * Squeeze Variable x along the given axes.
+   */
+  Variable Squeeze(const Variable& operand, const std::vector<int>& axes);
+
   /**
    * The convolution2D layer calculates the output based on the input, filter
    * and strides, paddings, dilations, groups parameters.

cinn/frontend/net_builder_test.cc

@@ -189,5 +189,57 @@ TEST(net_build, program_execute_reverse) {
   runtime_program->Execute();
 }
 
+TEST(net_build, program_execute_pool2d_grad) {
+  const int B = 4;
+  const int C = 1;
+  const int H = 7;
+  const int W = 1;
+
+  NetBuilder builder("net_builder");
+  Placeholder input = builder.CreateInput(Float(32), {B, C, H, W}, "In");
+  Variable output   = builder.Squeeze(input, {1, 3});
+  auto program      = builder.Build();
+
+  Target target = common::DefaultHostTarget();
+
+  auto graph = std::make_shared<hlir::framework::Graph>(program, target);
+  auto scope = BuildScope(target, graph);
+  hlir::framework::GraphCompiler gc(target, scope, graph);
+  auto runtime_program = gc.Build();
+
+  scope->Var<hlir::framework::Tensor>(std::string(input.id()));
+  scope->Var<hlir::framework::Tensor>(std::string(output->id));
+
+  auto input_tensor = scope->GetTensor(std::string(input.id()));
+  float* input_data = input_tensor->mutable_data<float>(target);
+
+  runtime_program->Execute();
+
+  auto output_tensor                   = scope->GetTensor(std::string(output->id));
+  const std::vector<int>& output_shape = output_tensor->shape().data();
+  EXPECT_EQ(output_shape.size(), 2UL);
+  EXPECT_EQ(output_shape[0], B);
+  EXPECT_EQ(output_shape[1], H);
+
+  float* output_data = output_tensor->mutable_data<float>(target);
+  VLOG(6) << "Visualize output_data";
+  for (int b = 0; b < B; ++b) {
+    for (int c = 0; c < C; ++c) {
+      VLOG(6) << "b = " << b << ", c = " << c;
+      for (int h = 0; h < H; ++h) {
+        std::string line;
+        for (int w = 0; w < W; ++w) {
+          int index      = w + W * (h + H * (c + C * b));
+          float in_data  = input_data[index];
+          float out_data = output_data[index];
+          line += (std::to_string(out_data) + ", ");
+          EXPECT_EQ(in_data, out_data);
+        }
+        VLOG(6) << line;
+      }
+    }
+  }
+}
+
 }  // namespace frontend
 }  // namespace cinn

cinn/hlir/op/CMakeLists.txt

@@ -1,3 +1,5 @@
+add_subdirectory(contrib)
+
 core_gather_headers()
 
 gather_srcs(cinnapi_src SRCS

cinn/hlir/op/contrib/CMakeLists.txt

@@ -0,0 +1,7 @@
+core_gather_headers()
+
+gather_srcs(cinnapi_src SRCS
+    squeeze.cc
+    )
+
+cc_test(test_squeeze SRCS squeeze_test.cc DEPS cinncore)

cinn/hlir/op/contrib/squeeze.cc

@@ -0,0 +1,228 @@
+// Copyright (c) 2021 CINN 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 "cinn/hlir/op/contrib/squeeze.h"
+
+#include <gflags/gflags.h>
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "cinn/common/cas.h"
+#include "cinn/common/common.h"
+#include "cinn/common/context.h"
+#include "cinn/common/macros.h"
+#include "cinn/hlir/framework/node.h"
+#include "cinn/hlir/framework/op.h"
+#include "cinn/hlir/framework/op_strategy.h"
+#include "cinn/hlir/pe/elementwise.h"
+#include "cinn/ir/ir.h"
+#include "cinn/ir/ir_base.h"
+#include "cinn/ir/tensor.h"
+#include "cinn/lang/builtin.h"
+#include "cinn/lang/compute.h"
+
+DECLARE_bool(cinn_ir_schedule);
+
+namespace cinn {
+namespace hlir {
+namespace op {
+
+using common::CINNValue;
+using common::CINNValuePack;
+
+ir::Tensor Squeeze(const ir::Tensor &A, const std::vector<int> &axes, const std::string &name) {
+  std::vector<Expr> new_expr_shape;
+  std::vector<Expr> A_expr_shape = A->shape;
+  if (axes.size() != 0) {
+    for (int i = 0; i < A_expr_shape.size(); ++i) {
+      CHECK(A_expr_shape[i].is_constant()) << "Input tensor's shape should be constant value.";
+      if (std::find(axes.begin(), axes.end(), i) != axes.end()) {
+        CHECK_EQ(A_expr_shape[i], Expr(1));
+      } else {
+        new_expr_shape.push_back(A_expr_shape[i]);
+      }
+    }
+  } else {
+    for (auto &i : A_expr_shape) {
+      CHECK(i.is_constant()) << "Input tensor's shape should be constant value.";
+      if (i != Expr(1)) {
+        new_expr_shape.push_back(i);
+      }
+    }
+  }
+  auto res = Compute(
+      new_expr_shape,
+      [=](const std::vector<Expr> &indices) {
+        Expr offset = Expr(0);
+        for (int i = 0; i < indices.size(); i++) {
+          offset = offset * new_expr_shape[i] + indices[i];
+        }
+        std::vector<Expr> indices_a;
+        for (int i = A_expr_shape.size() - 1; i >= 0; i--) {
+          auto temp = offset % A_expr_shape[i];
+          indices_a.insert(indices_a.begin(), common::AutoSimplify(temp));
+          offset = (offset - temp) / A_expr_shape[i];
+        }
+        return lang::Identity(A(indices_a));
+      },
+      name);
+  return res;
+}
+
+std::shared_ptr<framework::OpStrategy> StrategyForSqueeze(const framework::NodeAttr &attrs,
+                                                          const std::vector<ir::Tensor> &inputs,
+                                                          const std::vector<Type> &out_type,
+                                                          const std::vector<std::vector<int>> &output_shapes,
+                                                          const Target &target) {
+  CHECK(attrs.attr_store.count("axes")) << "find no attr of axes";
+  std::vector<int> axes = absl::get<std::vector<int>>(attrs.attr_store.at("axes"));
+
+  framework::CINNCompute squeeze_compute([=](lang::Args args, lang::RetValue *ret) {
+    CHECK(!args.empty()) << "The input arguments of Squeeze compute is empty! Please check.\n";
+    CINNValuePack a = args[0];
+    CHECK_GE(a.size(), 1U) << "at least 1 input tensors for Squeeze compute\n";
+    Expr A = a[0];
+    CHECK(A.as_tensor());
+    CHECK(!output_shapes.empty());
+    auto tensor_A = A.as_tensor_ref();
+    auto stages   = CreateStages({tensor_A});
+    VLOG(3) << "A shape: " << utils::Join(tensor_A->shape, ", ")
+            << ", output_shapes: " << utils::Join(output_shapes[0], ", ");
+    ir::Tensor out = Squeeze(tensor_A, axes, UniqName("Squeeze_out"));
+    std::vector<CINNValue> res;
+    stages->InsertLazily(out);
+    res.push_back(CINNValue(out));
+    CHECK(!out_type.empty()) << "Output type of Squeeze is empty! Please check.\n";
+    res.push_back(CINNValue(stages));
+    *ret = CINNValuePack{res};
+  });
+
+  framework::CINNSchedule squeeze_schedule([=](lang::Args args, lang::RetValue *ret) {
+    CHECK(!args.empty()) << "The input argument of reshape schedule is empty! Please check.\n";
+    CINNValuePack arg_pack = args[0];
+    Expr out               = arg_pack[0];
+    CHECK(out.as_tensor());
+    *ret = arg_pack;
+  });
+
+  auto strategy = std::make_shared<framework::OpStrategy>();
+  strategy->AddImpl(squeeze_compute, squeeze_schedule, "strategy.squeeze.x86", 1);
+  return strategy;
+}
+
+std::vector<std::vector<int>> InferShapeForSqueeze(const std::vector<std::vector<int>> &inputs_shape,
+                                                   const framework::AttrMapType &attrs) {
+  CHECK_EQ(inputs_shape.size(), 1U) << "The input's shape size should be 1! Please check again.";
+  std::vector<int> axes;
+  for (auto &iter : attrs) {
+    if (iter.first == "axes") {
+      axes = absl::get<std::vector<int>>(iter.second);
+      break;
+    }
+  }
+
+  std::vector<int> output_shape;
+  int tensor_size = 1;
+  if (axes.size() != 0) {
+    std::vector<int> temp_shape = inputs_shape[0];
+    for (auto &a : axes) {
+      CHECK(a < temp_shape.size());
+      temp_shape[a] = 0;
+    }
+    for (auto &i : temp_shape) {
+      if (i != 0) {
+        output_shape.push_back(i);
+        tensor_size *= i;
+      }
+    }
+  } else {
+    for (auto &i : inputs_shape[0]) {
+      if (i != 1) {
+        output_shape.push_back(i);
+        tensor_size *= i;
+      }
+    }
+  }
+
+  CHECK(!output_shape.empty()) << "infer_shape for squeeze turns out to be empty. Please check\n";
+  int flag_index = -1;
+  for (int i = 0; i < output_shape.size(); i++) {
+    if (output_shape[i] > 0) {
+      CHECK_EQ(tensor_size % output_shape[i], 0)
+          << "Incompatible input shape and output shape in op reshape: " << tensor_size << ", " << output_shape[i];
+      tensor_size /= output_shape[i];
+    } else if (output_shape[i] == 0) {
+      CHECK_LT(i, inputs_shape[0].size())
+          << "In op reshape, when attribute shape[i] == 0, shape[i] = input_shape[i]. But now the size of input_shape "
+             "<= i, which is incompatible. Please check!";
+      output_shape[i] = inputs_shape[0][i];
+      CHECK_EQ(tensor_size % output_shape[i], 0)
+          << "Incompatible input shape and output shape in op reshape: " << tensor_size << ", " << output_shape[i];
+      tensor_size /= output_shape[i];
+    } else if (output_shape[i] == -1 && flag_index == -1) {
+      flag_index = i;
+    } else if (output_shape[i] == -1) {
+      LOG(FATAL) << "More than one -1 in output_shape of op reshape.";
+    } else {
+      LOG(FATAL) << "Unsupported output_shape " << output_shape[i];
+    }
+  }
+  if (flag_index >= 0) output_shape[flag_index] = tensor_size;
+  std::vector<std::vector<int>> res{output_shape};
+  return res;
+}
+
+std::vector<Type> InferDtypeForSqueeze(const std::vector<Type> &inputs_type, const framework::AttrMapType &attrs) {
+  CHECK(!inputs_type.empty()) << "The input's type size is 0! Please check again.";
+  std::vector<Type> res{inputs_type[0]};
+  return res;
+}
+
+std::vector<std::vector<std::string>> InferLayoutForSqueeze(const std::vector<framework::shape_t> &input_shapes,
+                                                            const std::vector<std::string> &input_layouts,
+                                                            const framework::NodeAttr &attrs,
+                                                            const Target &target) {
+  CHECK_EQ(input_shapes.size(), 1U) << "The input's shape size is not 1! Please check again.";
+  CHECK_EQ(input_layouts.size(), 1U) << "The input's layout size is not 1! Please check again.";
+  std::vector<std::string> new_input_layouts = input_layouts;
+  if (input_shapes[0].size() > 4) {
+    // alter input layout back
+    new_input_layouts[0] = "NCHW";
+    VLOG(3) << "alter input layout from " << input_layouts[0] << " to " << new_input_layouts[0];
+  }
+  return {new_input_layouts, new_input_layouts};
+}
+
+}  // namespace op
+}  // namespace hlir
+}  // namespace cinn
+
+CINN_REGISTER_HELPER(squeeze_ops) {
+  CINN_REGISTER_OP(squeeze)
+      .describe("Squeeze.")
+      .set_num_inputs(1)
+      .set_num_outputs(1)
+      .set_attr<cinn::hlir::framework::StrategyFunction>("CINNStrategy", cinn::hlir::op::StrategyForSqueeze)
+      .set_attr("infershape", MakeOpFunction(cinn::hlir::op::InferShapeForSqueeze))
+      .set_attr("inferdtype", MakeOpFunction(cinn::hlir::op::InferDtypeForSqueeze))
+#ifndef CINN_WITH_CUDA
+      .set_attr("inferlayout", MakeOpFunction(cinn::hlir::op::InferLayoutForSqueeze))
+#endif
+      .set_support_level(4);
+
+  return true;
+}

cinn/hlir/op/contrib/squeeze.h

@@ -0,0 +1,32 @@
+// Copyright (c) 2021 CINN 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 <vector>
+
+#include "cinn/ir/ir.h"
+#include "cinn/ir/ir_base.h"
+#include "cinn/ir/tensor.h"
+
+namespace cinn {
+namespace hlir {
+namespace op {
+
+ir::Tensor Squeeze(const ir::Tensor& A, const std::vector<int>& axis, const std::string& name);
+
+}  // namespace op
+}  // namespace hlir
+}  // namespace cinn