Merge branch 'develop' of https://github.com/PaddlePaddle/CINN into develop

Author: MayYouBeProsperous

cinn/auto_schedule/auto_tuner.cc

@@ -25,6 +25,7 @@
 #include "cinn/auto_schedule/measure/simple_builder.h"
 #include "cinn/auto_schedule/measure/simple_runner.h"
 #include "cinn/auto_schedule/task/task_creator.h"
+#include "cinn/auto_schedule/task/task_registry.h"
 #include "cinn/auto_schedule/task/tune_task.h"
 #include "cinn/auto_schedule/task_scheduler/task_scheduler.h"
 #include "cinn/common/context.h"
@@ -49,11 +50,21 @@ void AutoTuner::Initialize(const Config& config, hlir::framework::GraphCompiler*
   const auto& dtype_dict = graph_->GetAttrs<absl::flat_hash_map<std::string, common::Type>>("inferdtype");
   const auto& shape_dict = graph_->GetAttrs<absl::flat_hash_map<std::string, hlir::framework::shape_t>>("infershape");
 
-  op_lowerer_ = std::make_unique<hlir::framework::OpLowerer>(dtype_dict, shape_dict, target_);
+  op_lowerer_                        = std::make_unique<hlir::framework::OpLowerer>(dtype_dict, shape_dict, target_);
+  InitialTaskRegistry* task_registry = InitialTaskRegistry::Global();
   for (TuneTask& task : tasks_) {
     task.SetOpLowerer(op_lowerer_.get());
     task.TaskGraphToUnoptLoweredFunc();
     task.SerializeToString(shape_dict, dtype_dict);
+
+    // Register the initial ModuleExpr corresponding to the task
+    std::vector<ir::Expr> exprs(task.lowered_funcs.size());
+    std::transform(
+        task.lowered_funcs.begin(), task.lowered_funcs.end(), exprs.begin(), [&](const ir::LoweredFunc& func) {
+          return func->body;
+        });
+    task_registry->Regist(task.serialized_key, ir::ModuleExpr(exprs));
+
     VLOG(3) << "Add a task with serialized_key:\n" << task.serialized_key;
   }
 

cinn/auto_schedule/task/CMakeLists.txt

@@ -9,3 +9,4 @@ gather_srcs(cinnapi_src SRCS task_creator.cc task_optimizer.cc)
 
 cc_test(test_task_creator SRCS task_creator_test.cc DEPS cinncore)
 cc_test(test_tune_task SRCS tune_task_test.cc DEPS cinncore)
+cc_test(test_task_registry SRCS task_registry_test.cc DEPS cinncore)

cinn/auto_schedule/task/task_registry.h

@@ -0,0 +1,79 @@
+// Copyright (c) 2022 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 <gflags/gflags.h>
+
+#include <mutex>
+#include <string>
+
+#include "cinn/ir/ir_schedule.h"
+#include "cinn/optim/ir_copy.h"
+#include "cinn/utils/registry.h"
+
+namespace cinn {
+
+namespace auto_schedule {
+
+struct InitialTaskInfo {
+  std::string task_key;
+  ir::ModuleExpr module_expr;
+
+  InitialTaskInfo(const std::string& task_key, const ir::ModuleExpr& module_expr)
+      : task_key(task_key), module_expr(module_expr) {}
+};
+
+// Global task registry, used to save the initial ModuleExpr of each task.
+class InitialTaskRegistry : public Registry<InitialTaskInfo> {
+ public:
+  static InitialTaskRegistry* Global() {
+    static InitialTaskRegistry x;
+    return &x;
+  }
+
+  // Get the initial ModuleExpr of a task.
+  inline const InitialTaskInfo* Get(const std::string& task_key) {
+    const InitialTaskInfo* task_info = Registry<InitialTaskInfo>::Find(task_key);
+    CHECK(task_info) << "InitialTaskInfo [" << task_key << "] is not registered";
+    return task_info;
+  }
+
+  // Check if the task info with task_key exists;
+  inline const bool Has(const std::string& task_key) { return nullptr != Registry<InitialTaskInfo>::Find(task_key); }
+
+  // Regist the initial ModuleExpr of a task into the map
+  inline void Regist(const std::string& task_key, const ir::ModuleExpr& module_expr) {
+    std::lock_guard<std::mutex> guard(registering_mutex);
+    if (fmap_.count(task_key) == 0) {
+      InitialTaskInfo* task_info = new InitialTaskInfo(task_key, optim::IRCopy(module_expr));
+      __REGISTER__(task_key, task_info);
+    }
+  }
+
+ private:
+  InitialTaskRegistry() = default;
+  CINN_DISALLOW_COPY_AND_ASSIGN(InitialTaskRegistry);
+
+  // Regist the initial ModuleExpr of a task.
+  inline InitialTaskInfo* __REGISTER__(const std::string& task_key, InitialTaskInfo* task_info) {
+    fmap_[task_key] = task_info;
+    const_list_.push_back(task_info);
+    entry_list_.push_back(task_info);
+    return task_info;
+  }
+};
+
+}  // namespace auto_schedule
+}  // namespace cinn

cinn/auto_schedule/task/task_registry_test.cc

@@ -0,0 +1,107 @@
+// Copyright (c) 2022 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/auto_schedule/task/task_registry.h"
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include <cstdlib>
+
+#include "cinn/auto_schedule/task/task_creator.h"
+#include "cinn/auto_schedule/task/tune_task.h"
+#include "cinn/frontend/net_builder.h"
+#include "cinn/hlir/framework/graph.h"
+#include "cinn/hlir/framework/graph_compiler.h"
+#include "cinn/hlir/framework/op_lowering.h"
+#include "cinn/utils/string.h"
+#include "cinn/utils/type_defs.h"
+
+DECLARE_bool(auto_schedule_use_cost_model);
+DECLARE_bool(cinn_ir_schedule);
+
+namespace cinn {
+namespace auto_schedule {
+
+std::vector<TuneTask> CreateTasks(hlir::framework::Graph* graph, const common::Target& target) {
+  // create tasks
+  TaskCreator task_creator;
+  std::vector<TuneTask> tasks = task_creator.CreateTuneTaskOpLevel(graph);
+
+  const auto& dtype_dict = graph->GetAttrs<absl::flat_hash_map<std::string, common::Type>>("inferdtype");
+  const auto& shape_dict = graph->GetAttrs<absl::flat_hash_map<std::string, hlir::framework::shape_t>>("infershape");
+
+  std::unique_ptr<hlir::framework::OpLowerer> op_lowerer =
+      std::make_unique<hlir::framework::OpLowerer>(dtype_dict, shape_dict, target);
+  for (TuneTask& task : tasks) {
+    task.SetOpLowerer(op_lowerer.get());
+    task.TaskGraphToUnoptLoweredFunc();
+    task.SerializeToString(shape_dict, dtype_dict);
+    VLOG(3) << "Add a task with serialized_key:\n" << task.serialized_key;
+  }
+
+  return tasks;
+}
+
+std::shared_ptr<hlir::framework::Graph> CreateAddProgram(const common::Target& target) {
+  frontend::NetBuilder builder("test");
+
+  auto a = builder.CreateInput(Float(32), {1, 64, 112, 112}, "A");
+  auto b = builder.CreateInput(Float(32), {64}, "B");
+  auto c = builder.Add(a, b, 1);
+
+  return std::make_shared<hlir::framework::Graph>(builder.Build(), target);
+}
+
+TEST(TestTaskRegistry, basic) {
+  FLAGS_auto_schedule_use_cost_model = true;
+  FLAGS_cinn_ir_schedule             = true;
+
+#ifdef CINN_WITH_CUDA
+  Target target = common::DefaultNVGPUTarget();
+#else
+  Target target = common::DefaultHostTarget();
+#endif
+  std::shared_ptr<hlir::framework::Graph> graph = CreateAddProgram(target);
+  std::vector<TuneTask> tasks                   = CreateTasks(graph.get(), target);
+
+  InitialTaskRegistry* task_registry = InitialTaskRegistry::Global();
+
+  std::vector<ir::ModuleExpr> module_exprs;
+  for (const TuneTask& task : tasks) {
+    module_exprs.emplace_back(task.GetLoweredFuncBodyExprs());
+    task_registry->Regist(task.serialized_key, module_exprs.back());
+  }
+
+  for (int i = 0; i < tasks.size(); ++i) {
+    std::string key = tasks[i].serialized_key;
+    VLOG(3) << "serialized_key = " << key;
+    ir::ModuleExpr new_expr = task_registry->Get(key)->module_expr;
+
+    ASSERT_EQ(new_expr.GetExprs().size(), module_exprs[i].GetExprs().size());
+    for (int j = 0; j < new_expr.GetExprs().size(); ++j) {
+      VLOG(3) << "expr " << j << " of task " << key << " : " << new_expr.GetExprs().at(j);
+      ASSERT_EQ(utils::GetStreamCnt(new_expr.GetExprs().at(j)), utils::GetStreamCnt(module_exprs[i].GetExprs().at(j)));
+    }
+  }
+
+  bool flag = task_registry->Has(tasks[0].serialized_key);
+  ASSERT_EQ(flag, true);
+
+  flag = task_registry->Has("not_exist");
+  ASSERT_EQ(flag, false);
+}
+
+}  // namespace auto_schedule
+}  // namespace cinn

cinn/frontend/computation_test.cc

@@ -48,7 +48,7 @@ Program CreateTestProgram() {
   auto i = builder.Max(e, h);
   auto j = builder.Min(e, h);
   auto k = builder.Multiply(i, j);
-  auto l = builder.ConstScalar<bool>(1, "condition");
+  auto l = builder.Constant<bool>(1, "condition");
   auto m = builder.BroadcastTo(l, {B, M, N}, {0});
   auto n = builder.Select(m, j, k);
   auto o = builder.ReduceSum(n, {0, 1, 2});

cinn/frontend/decomposer/batch_norm.cc

@@ -54,7 +54,7 @@ struct BatchNormHelper {
 
   template <typename T>
   Variable GetTensorFromScalar(T value, std::string name, const std::vector<int>& shape) {
-    // return builder->BroadcastTo(builder->ConstScalar<T>(value, common::UniqName(name)), shape, {0});
+    // return builder->BroadcastTo(builder->Constant<T>(value, common::UniqName(name)), shape, {0});
     return builder->FillConstant<T>(shape, value, common::UniqName(name));
   }
 

cinn/frontend/net_builder.cc

@@ -220,6 +220,10 @@ std::vector<Variable> NetBuilder::Split(const Variable& operand, const std::vect
 }
 
 Variable NetBuilder::Concat(const std::vector<Variable>& input_vars, int axis) {
+  CHECK(!input_vars.empty()) << "The inputs of concat op should not be empty! Please check.";
+  if (input_vars.size() == 1UL) {
+    return Identity(input_vars.front());
+  }
   return CustomInstr("concat", input_vars, {{"axis", axis}}).front();
 }
 
@@ -446,6 +450,24 @@ Variable NetBuilder::Sort(const Variable& operand, const int& axis, const bool&
   return instr.GetOutput(0);
 }
 
+Variable NetBuilder::Argmax(const Variable& x, const int& axis, const bool& keep_dim) {
+  Instruction instr("argmax", {x});
+  instr.SetAttr("axis", axis);
+  instr.SetAttr("keep_dim", keep_dim);
+  InferShape(instr);
+  AppendInstruction(instr);
+  return instr.GetOutput(0);
+}
+
+Variable NetBuilder::Argmin(const Variable& x, const int& axis, const bool& keep_dim) {
+  Instruction instr("argmin", {x});
+  instr.SetAttr("axis", axis);
+  instr.SetAttr("keep_dim", keep_dim);
+  InferShape(instr);
+  AppendInstruction(instr);
+  return instr.GetOutput(0);
+}
+
 Variable NetBuilder::Conv2d(const Variable& a,
                             const Variable& b,
                             const std::vector<int>& strides,
@@ -502,6 +524,10 @@ Variable NetBuilder::Pool2d(const Variable& a,
       .front();
 }
 
+Variable NetBuilder::Repeat(const Variable& x, int repeats, int axis) {
+  return CustomInstr("repeat", {x}, {{"repeats", repeats}, {"axis", axis}}).front();
+}
+
 std::vector<Variable> NetBuilder::BatchNorm(const Variable& a,
                                             const Variable& scale,
                                             const Variable& bias,
@@ -557,6 +583,14 @@ Variable NetBuilder::Arange(const float start, const float stop, const float ste
   return CustomInstr("arange", {}, {{"start", start}, {"stop", stop}, {"step", step}, {"dtype", dtype}}).front();
 }
 
+Variable NetBuilder::Flip(const Variable& operand, const std::vector<int>& axes) {
+  Instruction instr("flip", {operand});
+  instr.SetAttr("axes", axes);
+  InferShape(instr);
+  AppendInstruction(instr);
+  return instr.GetOutput(0);
+}
+
 // conv2d grad, output(grad_x, grad_w)
 std::vector<Variable> NetBuilder::Conv2dGrad(const Variable& dy,
                                              const Variable& x,