PaddlePaddle
diff --git a/‎paddle/cinn/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/cinn/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/cinn/ast_gen_ius/CMakeLists.txt
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎paddle/cinn/ast_gen_ius/CMakeLists.txt
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
diff --git a/‎paddle/cinn/ast_gen_ius/ast_gen.cc
Copy file name to clipboard
+52Lines changed: 52 additions & 0 deletions b/‎paddle/cinn/ast_gen_ius/ast_gen.cc
Copy file name to clipboard
+52Lines changed: 52 additions & 0 deletions
diff --git a/‎paddle/cinn/ast_gen_ius/ast_gen.h
Copy file name to clipboard
+30Lines changed: 30 additions & 0 deletions b/‎paddle/cinn/ast_gen_ius/ast_gen.h
Copy file name to clipboard
+30Lines changed: 30 additions & 0 deletions
diff --git a/‎paddle/cinn/ast_gen_ius/ast_gen_test.cc
Copy file name to clipboard
+44Lines changed: 44 additions & 0 deletions b/‎paddle/cinn/ast_gen_ius/ast_gen_test.cc
Copy file name to clipboard
+44Lines changed: 44 additions & 0 deletions
diff --git a/‎paddle/cinn/ast_gen_ius/tensor_group.cc
Copy file name to clipboard
+198Lines changed: 198 additions & 0 deletions b/‎paddle/cinn/ast_gen_ius/tensor_group.cc
Copy file name to clipboard
+198Lines changed: 198 additions & 0 deletions
@@ -3,6 +3,7 @@ if(WITH_TESTING)
 endif()
 
 add_subdirectory(api)
+add_subdirectory(ast_gen_ius)
 add_subdirectory(auto_schedule)
 add_subdirectory(common)
 add_subdirectory(utils)
 
@@ -0,0 +1,6 @@
+core_gather_headers()
+
+gather_srcs(cinnapi_src SRCS ast_gen.cc tensor_group.cc)
+
+cinn_cc_test(test_ast_gen_ius SRCS ast_gen_test.cc DEPS cinncore)
+cinn_cc_test(test_tensor_group SRCS tensor_group_test.cc DEPS cinncore)
@@ -0,0 +1,52 @@
+// Copyright (c) 2023 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
+//
+//     https://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 "paddle/cinn/ast_gen_ius/ast_gen.h"
+#include "paddle/cinn/ir/ir.h"
+#include "paddle/cinn/ir/ir_base.h"
+#include "paddle/cinn/ir/operation.h"
+#include "paddle/cinn/ir/tensor.h"
+#include "paddle/cinn/ir/utils/ir_printer.h"
+
+namespace cinn {
+namespace ast_gen_ius {
+
+ir::Expr AstGen::Build(const ir::Tensor& tensor) {
+  const std::vector<ir::Var>& axis = tensor->axis();
+  const std::vector<ir::Expr>& shape = tensor->shape;
+  size_t axis_len = axis.size();
+  CHECK_EQ(shape.size(), axis_len)
+      << "Internal Error: Tensor has different shape and axis length in AstGen";
+
+  std::vector<ir::Expr> axis_exprs;
+  for (const auto& a : axis) {
+    axis_exprs.push_back(a);
+  }
+  ir::Expr body = ir::Store::Make(tensor, tensor->body(), axis_exprs);
+
+  for (int i = static_cast<int>(axis_len) - 1; i >= 0; --i) {
+    ir::Var loop_var = axis[i];
+    ir::Expr loop_extent = shape[i];
+    body = ir::For::Make(loop_var,
+                         Expr(0),
+                         loop_extent,
+                         ir::ForType::Serial,
+                         ir::DeviceAPI::Host,
+                         ir::Block::Make({body}));
+  }
+  return body;
+}
+
+}  // namespace ast_gen_ius
+}  // namespace cinn
@@ -0,0 +1,30 @@
+// Copyright (c) 2023 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
+//
+//     https://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 "paddle/cinn/ir/ir.h"
+#include "paddle/cinn/ir/ir_base.h"
+#include "paddle/cinn/ir/tensor.h"
+
+namespace cinn {
+namespace ast_gen_ius {
+
+class AstGen {
+ public:
+  static ir::Expr Build(const ir::Tensor& tensor);
+};
+
+}  // namespace ast_gen_ius
+}  // namespace cinn
@@ -0,0 +1,44 @@
+// Copyright (c) 2023 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
+//
+//     https://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 <vector>
+
+#include "paddle/cinn/ast_gen_ius/ast_gen.h"
+#include "paddle/cinn/ir/ir.h"
+#include "paddle/cinn/ir/ir_base.h"
+#include "paddle/cinn/ir/tensor.h"
+#include "paddle/cinn/lang/builtin.h"
+#include "paddle/cinn/lang/compute.h"
+#include "paddle/cinn/lang/placeholder.h"
+
+namespace cinn {
+namespace ast_gen_ius {
+
+using cinn::ir::Expr;
+using cinn::ir::Tensor;
+
+TEST(AstGen, Build) {
+  std::vector<Expr> shape = {Expr(10), Expr(10), Expr(10), Expr(10)};
+  lang::Placeholder<float> A("A", shape);
+  Tensor B = lang::Compute(
+      shape,
+      [&](const std::vector<Expr>& indice) { return lang::Relu(A(indice), 0); },
+      "relu_test");
+  Expr out = AstGen::Build(B);
+  LOG(INFO) << out;
+}
+
+}  // namespace ast_gen_ius
+}  // namespace cinn
@@ -0,0 +1,198 @@
+// Copyright (c) 2023 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
+//
+//     https://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 "paddle/cinn/ast_gen_ius/tensor_group.h"
+
+#include <unordered_map>
+#include <vector>
+
+#include "paddle/cinn/ir/ir.h"
+#include "paddle/cinn/ir/ir_base.h"
+#include "paddle/cinn/ir/tensor.h"
+#include "paddle/cinn/ir/utils/ir_nodes_collector.h"
+
+namespace cinn {
+namespace ast_gen_ius {
+
+TensorGroup::TensorGroup(const std::vector<ir::Tensor>& tensors) {
+  std::set<ir::Tensor> all_tensors(tensors.begin(), tensors.end());
+
+  for (auto& tensor : tensors) {
+    output_tensor_names_.insert(tensor->name);
+    std::set<ir::Expr> used_tensors = ir::CollectIRNodes(
+        tensor->body(), [](const Expr* x) { return x->as_tensor(); });
+    for (const Expr& x : used_tensors) {
+      const ir::Tensor to_dep = x.as_tensor_ref();
+      all_tensors.insert(to_dep);
+      this->CtrlDepend(tensor, to_dep);
+    }
+  }
+
+  for (const ir::Tensor& t : all_tensors) {
+    name_to_tensor_.insert({t->name, t});
+  }
+}
+
+TensorGroup::~TensorGroup() {}
+
+bool TensorGroup::Contain(const std::string& name) const {
+  return name_to_tensor_.find(name) != name_to_tensor_.end();
+}
+
+void TensorGroup::Insert(const ir::Tensor& tensor) {
+  name_to_tensor_.insert({tensor->name, tensor});
+}
+
+ir::Tensor TensorGroup::Get(const std::string& name) {
+  return name_to_tensor_[name];
+}
+
+std::set<ir::Tensor> TensorGroup::GetAllTensors() {
+  std::set<ir::Tensor> all_tensors;
+  for (const std::pair<std::string, ir::Tensor>& p : name_to_tensor_) {
+    all_tensors.insert(p.second);
+  }
+  return all_tensors;
+}
+
+std::vector<ir::Tensor> TensorGroup::GetGenFuncTopoOrder(
+    const std::vector<ir::Tensor>& func_args) {
+  std::unordered_map<std::string, int> in_degree;
+  for (const auto& dep_pair : ctrl_dep_) {
+    const std::unordered_set<std::string>& dep_tensor_names = dep_pair.second;
+    in_degree[dep_pair.first] = dep_tensor_names.size();
+  }
+
+  std::vector<ir::Tensor> ret;
+  std::vector<std::string> stack;
+  for (const auto& name_tensor : name_to_tensor_) {
+    if (!in_degree.count(name_tensor.first)) {
+      stack.emplace_back(name_tensor.first);
+    }
+  }
+
+  std::set<std::string> input_arg_names;
+  for (const ir::Tensor& arg : func_args) {
+    input_arg_names.insert(arg->name);
+  }
+  for (const std::string& name : output_tensor_names_) {
+    input_arg_names.erase(name);
+  }
+
+  while (!stack.empty()) {
+    const std::string& cur = stack.back();
+    stack.pop_back();
+
+    if (!input_arg_names.count(cur)) {
+      ret.push_back(name_to_tensor_[cur]);
+    }
+
+    for (const auto& dep_pair : ctrl_dep_) {
+      const std::unordered_set<std::string>& dep_tensor_names = dep_pair.second;
+      if (dep_tensor_names.count(cur)) {
+        --in_degree[dep_pair.first];
+        if (in_degree[dep_pair.first] == 0) {
+          stack.emplace_back(dep_pair.first);
+        }
+      }
+    }
+  }
+  return ret;
+}
+
+bool TensorGroup::HasMarkedReduceInit(const std::string& tensor_name) const {
+  return tensor_name_needs_reduce_init_.count(tensor_name);
+}
+
+ir::Tensor TensorGroup::MarkReduceInit(const std::string& tensor_name) {
+  // TODO(zhhsplendid): add check
+  tensor_name_needs_reduce_init_.insert(tensor_name);
+}
+
+void TensorGroup::CtrlDepend(const ir::Tensor& tensor,
+                             const ir::Tensor& to_dep) {
+  ctrl_dep_[tensor->name].insert(to_dep->name);
+  if (!name_to_tensor_.count(tensor->name)) {
+    name_to_tensor_[tensor->name] = tensor;
+  }
+  if (!name_to_tensor_.count(to_dep->name)) {
+    name_to_tensor_[to_dep->name] = to_dep;
+  }
+}
+
+std::set<ir::Tensor> TensorGroup::GetCrtlDepTensors(
+    const std::string& tensor_name) {
+  if (!ctrl_dep_.count(tensor_name)) {
+    return {};
+  }
+  std::set<ir::Tensor> ret;
+  for (const std::string& dep_name : ctrl_dep_[tensor_name]) {
+    ret.insert(name_to_tensor_[dep_name]);
+  }
+  return ret;
+}
+
+std::string TensorGroup::GetShareMemRootName(const std::string& tensor_name) {
+  if (!share_memory_tensor_.count(tensor_name)) {
+    share_memory_tensor_[tensor_name] = tensor_name;
+    return tensor_name;
+  }
+  if (share_memory_tensor_[tensor_name] == tensor_name) {
+    return tensor_name;
+  }
+  share_memory_tensor_[tensor_name] =
+      GetShareMemRootName(share_memory_tensor_[tensor_name]);
+  return share_memory_tensor_[tensor_name];
+}
+
+void TensorGroup::ShareMemoryBuffer(const ir::Tensor& tensor,
+                                    const ir::Tensor& to_share) {
+  share_memory_tensor_[GetShareMemRootName(to_share->name)] =
+      GetShareMemRootName(tensor->name);
+}
+
+absl::flat_hash_map<std::string, ir::Tensor> TensorGroup::AllocateBuffers() {
+  std::unordered_set<std::string> allocated_roots;
+  for (auto& name_tensor : name_to_tensor_) {
+    std::string root_name = GetShareMemRootName(name_tensor.first);
+
+    // Allocate root buffer
+    if (!allocated_roots.count(root_name)) {
+      ir::Tensor root_tensor = name_to_tensor_[root_name];
+      if (!root_tensor->buffer.defined() && !root_tensor->type().is_void()) {
+        root_tensor->WithBuffer();
+        VLOG(6) << "Bind root_tensor " << root_name << " with buffer "
+                << root_tensor->buffer->name;
+      }
+      allocated_roots.insert(root_name);
+    }
+
+    // Share buffer
+    if (root_name != name_tensor.first) {
+      ir::Tensor& root_tensor = name_to_tensor_[root_name];
+      ir::Tensor& tensor = name_tensor.second;
+
+      auto keep_shape = root_tensor->buffer->shape;
+      tensor->Bind(root_tensor->buffer);
+      root_tensor->buffer->shape = keep_shape;
+      tensor->buffer->shape = keep_shape;
+      VLOG(6) << "Share buffer " << root_name << " with " << name_tensor.first;
+    }
+  }
+
+  return name_to_tensor_;
+}
+
+}  // namespace ast_gen_ius
+}  // namespace cinn