【PaddlePaddle Hackathon No.75】add sort op

cinn/frontend/net_builder.cc

@@ -103,6 +103,24 @@ Variable NetBuilder::Squeeze(const Variable& operand, const std::vector<int>& ax
   return instr.GetOutput(0);
 }
 
+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

@@ -102,6 +102,16 @@ class NetBuilder : public BaseBuilder {
    */
   Variable Squeeze(const Variable& operand, const std::vector<int>& axes);
 
+  /**
+   * Sort Variable x along the given axis and return index.
+   */
+  Variable ArgSort(const Variable& operand, const int& axis, const bool& is_ascend = true);
+
+  /**
+   * Sort Variable x along the given axis.
+   */
+  Variable Sort(const Variable& operand, const int& axis, const bool& is_ascend = true);
+
   /**
    * The convolution2D layer calculates the output based on the input, filter
    * and strides, paddings, dilations, groups parameters.

cinn/frontend/net_builder_test.cc

@@ -453,5 +453,115 @@ 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;
+    }
+  }
+}
+
 }  // namespace frontend
 }  // namespace cinn

cinn/hlir/op/contrib/CMakeLists.txt

@@ -2,10 +2,12 @@ core_gather_headers()
 
 gather_srcs(cinnapi_src SRCS
         cast.cc
-        squeeze.cc
         clip.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)