PaddlePaddle
diff --git a/‎paddle/cinn/api/op_topo_pattern.h
Copy file name to clipboard
+77Lines changed: 77 additions & 0 deletions b/‎paddle/cinn/api/op_topo_pattern.h
Copy file name to clipboard
+77Lines changed: 77 additions & 0 deletions
diff --git a/‎paddle/cinn/ast_gen_ius/ast_gen.cc
Copy file name to clipboard
+1-22Lines changed: 1 addition & 22 deletions b/‎paddle/cinn/ast_gen_ius/ast_gen.cc
Copy file name to clipboard
+1-22Lines changed: 1 addition & 22 deletions
diff --git a/‎paddle/cinn/backends/codegen_cuda_util.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/cinn/backends/codegen_cuda_util.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/cinn/frontend/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/cinn/frontend/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/cinn/frontend/group_cluster/CMakeLists.txt
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎paddle/cinn/frontend/group_cluster/CMakeLists.txt
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
diff --git a/‎paddle/cinn/frontend/group_cluster/cluster_policy/CMakeLists.txt
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/cinn/frontend/group_cluster/cluster_policy/CMakeLists.txt
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/cinn/frontend/group_cluster/cluster_policy/general_topo_policy.cc
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions b/‎paddle/cinn/frontend/group_cluster/cluster_policy/general_topo_policy.cc
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions
diff --git a/‎paddle/cinn/frontend/group_cluster/cluster_policy/general_topo_policy.h
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions b/‎paddle/cinn/frontend/group_cluster/cluster_policy/general_topo_policy.h
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions
diff --git a/‎paddle/cinn/frontend/group_cluster/cluster_policy/policy_manager.cc
Copy file name to clipboard
+28Lines changed: 28 additions & 0 deletions b/‎paddle/cinn/frontend/group_cluster/cluster_policy/policy_manager.cc
Copy file name to clipboard
+28Lines changed: 28 additions & 0 deletions
diff --git a/‎paddle/cinn/frontend/group_cluster/cluster_policy/policy_manager.h
Copy file name to clipboard
+39Lines changed: 39 additions & 0 deletions b/‎paddle/cinn/frontend/group_cluster/cluster_policy/policy_manager.h
Copy file name to clipboard
+39Lines changed: 39 additions & 0 deletions
@@ -0,0 +1,77 @@
+// Copyright (c) 2024 PaddlePaddle 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 <list>
+#include <variant>
+#include <vector>
+
+namespace cinn::api {
+
+template <typename T>
+struct ErrorPattern {};
+
+// ElementWise/Broadcast/Injective Ops without reduction ancestors.
+template <typename T>
+struct InjectiveSourcePattern {};
+
+// Reduce op
+template <typename T>
+struct SingleReductionOpPattern {};
+
+// ElementWise/Broadcast ops which have shardable dimentions and reduction
+// ancestors.
+template <typename T>
+struct PartialShardablePattern {};
+
+// Reduce base pattern
+template <typename T>
+struct ReductionPattern {
+  using Nothing = std::monostate;
+  std::variant<Nothing, InjectiveSourcePattern<T>, PartialShardablePattern<T>>
+      input;
+  SingleReductionOpPattern<T> reduce_op_pattern;
+
+  bool HasFusedInput() const {
+    return !std::holds_alternative<Nothing>(this->input);
+  }
+};
+
+// Stmt := IS | R | PS
+// ops in StmtPattern will be lowered into a inlined cuda code.
+template <typename T>
+using StmtPattern = std::variant<InjectiveSourcePattern<T>,
+                                 ReductionPattern<T>,
+                                 PartialShardablePattern<T>>;
+
+// Stmts := [Stmt]
+template <typename T>
+using StmtPatternVec = std::vector<StmtPattern<T>>;
+// fuse rules:
+//  1. IS * IS -> IS
+//  2. PS * PS -> PS
+//  3. IS * PS -> PS
+//  4. IS * R -> R
+//  5. PS * R -> R
+// lifting rules:
+//  1. R -> Stmts
+//  2. PS -> Stmts
+//  3. Stmts * Stmts -> Stmts
+// OpTopoPattern := Error | Stmts
+
+template <typename T>
+using OpTopoPattern = std::variant<ErrorPattern<T>, StmtPatternVec<T>>;
+
+}  // namespace cinn::api
@@ -100,13 +100,6 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
     const std::vector<ir::Var>& reduce_axis = tensor->reduce_axis;
     VLOG(4) << "ast gen: tensor init_body is " << init_body;
     for (int i = 0; i < shape.size(); ++i) {
-      bool is_keep_dim = axis[i]->is_keepdim;
-      if (FLAGS_group_schedule_tiling_first && is_keep_dim) {
-        // if tiling first, we need to replace the reduce axis with 0, but don't
-        // deal with the non-reduce axis
-        optim::ReplaceVarWithExpr(&init_body, axis[i], Expr(0));
-        continue;
-      }
       if (!FLAGS_group_schedule_tiling_first &&
           FLAGS_cinn_new_group_scheduler && shape[i] == Expr(1)) {
         optim::ReplaceVarWithExpr(&init_body, axis[i], Expr(0));
@@ -144,13 +137,6 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
     // for same axis so we re-create objects
     std::vector<Var> reduce_axis_vars = cinn::common::GenDefaultAxis(axis_len);
     for (int i = 0; i < shape.size(); ++i) {
-      bool is_keep_dim = axis[i]->is_keepdim;
-      if (FLAGS_group_schedule_tiling_first && is_keep_dim) {
-        // if tiling first, we need to replace the reduce axis with 0, but don't
-        // deal with the non-reduce axis
-        optim::ReplaceVarWithExpr(&reduce_body, axis[i], Expr(0));
-        continue;
-      }
       if (!FLAGS_group_schedule_tiling_first &&
           FLAGS_cinn_new_group_scheduler && shape[i] == Expr(1)) {
         optim::ReplaceVarWithExpr(&reduce_body, axis[i], Expr(0));
@@ -185,10 +171,7 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
       std::vector<ir::Var> non_reduce_axis_vars = [&]() {
         std::vector<ir::Var> res;
         for (int i = 0; i < shape.size(); ++i) {
-          bool is_keep_dim = axis[i]->is_keepdim;
-          if (!is_keep_dim) {
-            res.push_back(axis[i]);
-          }
+          res.push_back(axis[i]);
         }
         return res;
       }();
@@ -240,10 +223,6 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
     // Put the two parts together
     ir::Expr body = ir::Block::Make({init_body, reduce_body});
     for (int i = static_cast<int>(axis_len) - 1; i >= 0; --i) {
-      bool is_keep_dim = axis[i]->is_keepdim;
-      if (FLAGS_group_schedule_tiling_first && is_keep_dim) {
-        continue;
-      }
       if ((!FLAGS_group_schedule_tiling_first || !FLAGS_cinn_bucket_compile) &&
           shape[i] == Expr(1)) {
         continue;
 
@@ -78,6 +78,7 @@ detail::CollectBucketStrategyHostFunctionVisitor::GenDeviceKernelName(
 
 void detail::CollectBucketStrategyHostFunctionVisitor::ProcessLoweredFunc(
     ir::Expr func, ir::Expr predicate) {
+  VLOG(4) << "Process Lowered Func" << func;
   ir::_LoweredFunc_ *func_node = func.as_lowered_func();
   CHECK(func_node);
   if (!func_node->cuda_axis_info.valid()) {
 
@@ -62,6 +62,7 @@ add_subdirectory(paddle)
 add_subdirectory(decomposer)
 add_subdirectory(op_mappers)
 add_subdirectory(pass)
+add_subdirectory(group_cluster)
 
 cinn_cc_test(test_op_mapper_registry SRCS op_mapper_registry_test.cc DEPS
              cinncore)
@@ -0,0 +1,6 @@
+gather_srcs(group_cluster_src SRCS common_utils.cc pattern_node.cc
+            pattern_graph.cc)
+
+add_subdirectory(cluster_policy)
+
+cc_library(group_cluster SRCS ${group_cluster_src})
@@ -0,0 +1,3 @@
+gather_srcs(group_cluster_src SRCS general_topo_policy.cc policy_manager.cc)
+
+add_subdirectory(shardable_axes_policy)
@@ -0,0 +1,25 @@
+// Copyright (c) 2024 PaddlePaddle 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/frontend/group_cluster/cluster_policy/general_topo_policy.h"
+
+namespace cinn::frontend::group_cluster::policy {
+
+bool GeneralTopoPolicy::CanFuse(const PatternNodePtr upstream,
+                                const PatternNodePtr downstream) {
+  // TODO(wuzhanfei) topo policy (if lead to loop)
+  return false;
+}
+
+}  // namespace cinn::frontend::group_cluster::policy
@@ -0,0 +1,25 @@
+// Copyright (c) 2024 PaddlePaddle 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/frontend/group_cluster/cluster_policy/policy_manager.h"
+
+namespace cinn::frontend::group_cluster::policy {
+
+class GeneralTopoPolicy final : virtual public Policy {
+ public:
+  bool CanFuse(const PatternNodePtr upstream, const PatternNodePtr downstream);
+};
+
+}  // namespace cinn::frontend::group_cluster::policy
@@ -0,0 +1,28 @@
+// Copyright (c) 2024 PaddlePaddle 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/frontend/group_cluster/cluster_policy/policy_manager.h"
+#include "paddle/common/enforce.h"
+
+namespace cinn::frontend::group_cluster::policy {
+
+bool PolicyManager::CanFuse(const PatternNodePtr upstream,
+                            const PatternNodePtr downstream) {
+  for (const auto& policy : policies_) {
+    if (!policy->CanFuse(upstream, downstream)) return false;
+  }
+  return true;
+}
+
+}  // namespace cinn::frontend::group_cluster::policy
@@ -0,0 +1,39 @@
+// Copyright (c) 2024 PaddlePaddle 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/frontend/group_cluster/pattern_node.h"
+
+namespace cinn::frontend::group_cluster::policy {
+
+class Policy {
+ public:
+  virtual bool CanFuse(const PatternNodePtr upstream,
+                       const PatternNodePtr downstream) = 0;
+};
+
+using PolicyPtr = std::shared_ptr<Policy>;
+
+class PolicyManager {
+ public:
+  explicit PolicyManager(const std::vector<PolicyPtr>& policies)
+      : policies_(policies) {}
+  bool CanFuse(const PatternNodePtr upstream, const PatternNodePtr downstream);
+
+ private:
+  std::vector<PolicyPtr> policies_;
+};
+
+}  // namespace cinn::frontend::group_cluster::policy
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+gather_srcs(group_cluster_src SRCS general_topo_policy.cc policy_manager.cc)`
	`2`	`+`
	`3`	`+add_subdirectory(shardable_axes_policy)`