[CINN / PIR] Cinn trivalop fuse

paddle/cinn/api/op_topo_pattern.h

@@ -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
+//
+//     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 <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

paddle/cinn/ast_gen_ius/ast_gen.cc

@@ -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;

paddle/cinn/backends/codegen_cuda_util.cc

@@ -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()) {

paddle/cinn/frontend/CMakeLists.txt

@@ -62,6 +62,7 @@ add_subdirectory(paddle)
 add_subdirectory(decomposer)
 add_subdirectory(op_mappers)
 add_subdirectory(pass)
+add_subdirectory(cluster_ops)
 
 cinn_cc_test(test_op_mapper_registry SRCS op_mapper_registry_test.cc DEPS
              cinncore)

paddle/cinn/frontend/cluster_ops/CMakeLists.txt

@@ -0,0 +1,13 @@
+gather_srcs(
+  cluster_ops_src
+  SRCS
+  common_utils.cc
+  shardable_axes_inferer.cc
+  shardable_axes_provider.cc
+  shardable_axes_utils.cc
+  pattern_utils.cc
+  fusion_helper.cc
+  cluster_policy.cc
+  clustering_engine.cc)
+
+cc_library(cluster_ops SRCS ${cluster_ops_src})

paddle/cinn/frontend/cluster_ops/cluster_ops.h

@@ -0,0 +1,57 @@
+// 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
+//
+//     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 "paddle/cinn/frontend/cluster_ops/clustering_engine.h"
+
+namespace cinn::frontend {
+
+cluster_ops::ClusteringResult ClusterOps(
+    const cinn::dialect::GroupOp& group_op) {
+  const auto& ops = [&] {
+    std::vector<const pir::Operation*> ops;
+    for (const auto& op : group_op.GetOperators()) {
+      ops.push_back(op);
+    }
+    return ops;
+  }();
+
+  VLOG(4) << "Start Cluster Ops!";
+  VLOG(4) << "Input Group with size " << ops.size() << " :\n"
+          << cluster_ops::OpsDebugStr(ops);
+
+  auto shardable_axes_provider = [&] {
+    auto* program = group_op->GetParentProgram();
+    const auto* shape_analysis =
+        &pir::ShapeAnalysisManager::Instance().Get(program);
+    return cluster_ops::MakeDefaultShardableAxesProvider(shape_analysis);
+  }();
+
+  auto cluster_policy = [&] {
+    auto* program = group_op->GetParentProgram();
+    const auto* shape_analysis =
+        &pir::ShapeAnalysisManager::Instance().Get(program);
+    return cluster_ops::MakeLoopAlignableClusteringPolicy(shape_analysis);
+  }();
+
+  cluster_ops::ShardableAxesInferer inferer(shardable_axes_provider);
+  cluster_ops::ClusteringEngine engine(ops, inferer, cluster_policy);
+
+  auto result = engine.ClusterOps();
+  VLOG(4) << result.DebugStr();
+  VLOG(4) << "Finished Cluster Ops!";
+  return result;
+}
+}  // namespace cinn::frontend