【PaddlePaddle Hackathon No.75】add sort op

cinn/frontend/net_builder.cc

@@ -383,6 +383,24 @@ Variable NetBuilder::Conv(const Variable& lhs,
       .front();
 }
 
+Variable NetBuilder::ArgSort(const Variable& operand, const int& axis, const bool& is_ascend) {
+  Instruction instr("argsort", {operand});
+  instr.SetAttr("axis", axis);
+  instr.SetAttr("is_ascend", is_ascend);
+  InferShape(instr);
+  AppendInstruction(instr);
+  return instr.GetOutput(0);
+}
+
+Variable NetBuilder::Sort(const Variable& operand, const int& axis, const bool& is_ascend) {
+  Instruction instr("sort", {operand});
+  instr.SetAttr("axis", axis);
+  instr.SetAttr("is_ascend", is_ascend);
+  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

@@ -792,6 +792,26 @@ class NetBuilder {
                                       const float epsilon            = 1e-5f,
                                       const std::string& data_layout = "NCHW");
 
+  /**
+   * @brief Sort Variable x along the given axis. The original Variable x will not be changed.
+   * @param operand The variable that will be sorted.
+   * @param axis Specify the axis to operate on the input. Default: 0.
+   * @param is_ascend Sort mode.
+   * Defalut “NCHW”.
+   * @return `Sorted variable index`.
+   */
+  Variable ArgSort(const Variable& operand, const int& axis, const bool& is_ascend = true);
+
+  /**
+   * @brief Sort Variable x along the given axis. The original Variable x will not be changed.
+   * @param operand The variable that will be sorted.
+   * @param axis Specify the axis to operate on the input. Default: 0.
+   * @param is_ascend Sort mode.
+   * Defalut “NCHW”.
+   * @return `Sorted variable`.
+   */
+  Variable Sort(const Variable& operand, const int& axis, const bool& is_ascend = true);
+
  private:
   CINN_DISALLOW_COPY_AND_ASSIGN(NetBuilder);
 };

cinn/frontend/net_builder_test.cc

@@ -453,6 +453,116 @@ TEST(net_build, program_execute_squeeze_case3) {
   }
 }
 
+TEST(net_build, program_execute_argsort) {
+  const int B = 4;
+  const int H = 7;
+
+  NetBuilder builder("net_builder");
+  Placeholder input = builder.CreateInput(Float(32), {B, H}, "In");
+  Variable output   = builder.ArgSort(input, 0, true);
+  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()));
+  SetRandData<float>(input_tensor, target);
+  auto* 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_tensor->type(), Int(32));
+  EXPECT_EQ(output_shape.size(), 2UL);
+  EXPECT_EQ(output_shape[0], B);
+  EXPECT_EQ(output_shape[1], H);
+
+  int* output_data = output_tensor->mutable_data<int>(target);
+  VLOG(6) << "Visualize output_data";
+  for (int h = 0; h < H; ++h) {
+    std::vector<float> sorted_data;
+    std::vector<float> out_sorted_data(H);
+    for (int b = 0; b < B; ++b) {
+      int index = h + H * b;
+      sorted_data.push_back(input_data[index]);
+      out_sorted_data[output_data[index]] = input_data[index];
+    }
+    std::sort(sorted_data.begin(), sorted_data.begin() + B);
+
+    for (int b = 0; b < B; ++b) {
+      std::string line;
+      int index       = h + H * b;
+      float true_data = sorted_data[b];
+      float out_data  = out_sorted_data[b];
+      line += (std::to_string(out_data) + ", ");
+      EXPECT_EQ(true_data, out_data);
+      VLOG(6) << line;
+    }
+  }
+}
+
+TEST(net_build, program_execute_sort) {
+  const int B = 4;
+  const int H = 7;
+
+  NetBuilder builder("net_builder");
+  Placeholder input = builder.CreateInput(Float(32), {B, H}, "In");
+  Variable output   = builder.Sort(input, 0, true);
+  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()));
+  SetRandData<float>(input_tensor, target);
+  auto* 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_tensor->type(), Float(32));
+  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 h = 0; h < H; ++h) {
+    std::vector<float> sorted_data;
+    for (int b = 0; b < B; ++b) {
+      int index = h + H * b;
+      sorted_data.push_back(input_data[index]);
+    }
+    std::sort(sorted_data.begin(), sorted_data.begin() + B);
+
+    for (int b = 0; b < B; ++b) {
+      std::string line;
+      int index       = h + H * b;
+      float true_data = sorted_data[b];
+      float out_data  = output_data[index];
+      line += (std::to_string(out_data) + ", ");
+      EXPECT_EQ(true_data, out_data);
+      VLOG(6) << line;
+    }
+  }
+}
+
 TEST(net_build, program_execute_arange_float) {
   const float start       = 1.5F;
   const float stop        = 31.5F;

cinn/hlir/op/contrib/CMakeLists.txt

@@ -5,9 +5,12 @@ gather_srcs(cinnapi_src SRCS
         squeeze.cc
         clip.cc
         arange.cc
+        sort.cc
+        squeeze.cc
         )
 
 cc_test(test_cast SRCS cast_test.cc DEPS cinncore)
 cc_test(test_squeeze SRCS squeeze_test.cc DEPS cinncore)
 cc_test(test_clip SRCS clip_test.cc DEPS cinncore)
+cc_test(test_sort SRCS sort_test.cc DEPS cinncore)
 cc_test(test_arange SRCS arange_test.cc DEPS cinncore)