chore: import upstream snapshot with attribution

test/cpp/cinn/common/broadcast_tree_test.cc

@@ -0,0 +1,118 @@
+// 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.
+
+#include "paddle/cinn/common/broadcast_tree.h"
+
+#include "gtest/gtest.h"
+
+namespace cinn::common {
+using namespace symbol;  // NOLINT
+
+namespace {
+
+DimExpr MakeBroadcastDimExpr(const DimExpr& expr1, const DimExpr& expr2) {
+  List<DimExpr> operands{expr1, expr2};
+  return Broadcast<DimExpr>{operands};
+}
+
+bool DimExprNonBroadcast(const DimExpr& dim_expr) {
+  if (dim_expr.Has<Broadcast<DimExpr>>()) {
+    return false;
+  } else {
+    return true;
+  }
+}
+
+void CheckLeafNonBroadcast(const BroadcastLeaf& leaf) {
+  for (const auto& operands : *leaf) {
+    for (const auto& operand : operands) {
+      ASSERT_TRUE(DimExprNonBroadcast(operand));
+    }
+  }
+}
+
+void CheckInnerBranchNonBroadcast(
+    const BroadcastBranch<BroadcastTree>& branch) {
+  const auto& [_, lhs_eq_rhs_tree, lhs_eq_one_tree, rhs_eq_one_tree] =
+      branch.tuple();
+  ASSERT_TRUE(lhs_eq_rhs_tree.Has<BroadcastLeaf>());
+  ASSERT_TRUE(lhs_eq_one_tree.Has<BroadcastLeaf>());
+  ASSERT_TRUE(rhs_eq_one_tree.Has<BroadcastLeaf>());
+  CheckLeafNonBroadcast(lhs_eq_rhs_tree.Get<BroadcastLeaf>());
+  CheckLeafNonBroadcast(lhs_eq_one_tree.Get<BroadcastLeaf>());
+  CheckLeafNonBroadcast(rhs_eq_one_tree.Get<BroadcastLeaf>());
+}
+
+}  // namespace
+
+TEST(BroadcastTree, Naive) {
+  DimExpr expr1("S1");
+  DimExpr expr2("S2");
+  DimExpr expr3("S3");
+  DimExpr expr4("S4");
+  std::vector<DimExpr> tensor_shape{expr1,
+                                    expr2,
+                                    MakeBroadcastDimExpr(expr1, expr2),
+                                    MakeBroadcastDimExpr(expr3, expr4)};
+  BroadcastLeaf leaf = adt::List<std::vector<DimExpr>>{tensor_shape};
+  int num_of_leaves = 0;
+  BroadcastTree tree = ConstructBroadcastTree(leaf, &num_of_leaves);
+  ASSERT_TRUE(tree.Has<BroadcastBranch<BroadcastTree>>());
+  const auto& branch = tree.Get<BroadcastBranch<BroadcastTree>>();
+  const auto& [cstr_broadcastable,
+               lhs_eq_rhs_tree,
+               lhs_eq_one_tree,
+               rhs_eq_one_tree] = branch.tuple();
+  ASSERT_EQ(cstr_broadcastable->lhs, DimExpr("S1"));
+  ASSERT_EQ(cstr_broadcastable->rhs, DimExpr("S2"));
+  ASSERT_TRUE(lhs_eq_rhs_tree.Has<BroadcastBranch<BroadcastTree>>());
+  ASSERT_TRUE(lhs_eq_one_tree.Has<BroadcastBranch<BroadcastTree>>());
+  ASSERT_TRUE(rhs_eq_one_tree.Has<BroadcastBranch<BroadcastTree>>());
+  CheckInnerBranchNonBroadcast(
+      lhs_eq_rhs_tree.Get<BroadcastBranch<BroadcastTree>>());
+  CheckInnerBranchNonBroadcast(
+      lhs_eq_one_tree.Get<BroadcastBranch<BroadcastTree>>());
+  CheckInnerBranchNonBroadcast(
+      rhs_eq_one_tree.Get<BroadcastBranch<BroadcastTree>>());
+}
+
+TEST(BroadcastTree, SimplifyConstantBroadcast) {
+  DimExpr expr1("S1");
+  DimExpr expr2("S2");
+  DimExpr expr3("S3");
+  DimExpr expr4(4);
+  std::vector<DimExpr> tensor_shape{expr1,
+                                    expr2,
+                                    MakeBroadcastDimExpr(expr1, expr2),
+                                    MakeBroadcastDimExpr(expr3, expr4)};
+  BroadcastLeaf leaf = adt::List<std::vector<DimExpr>>{tensor_shape};
+  int num_of_leaves = 0;
+  BroadcastTree tree = ConstructBroadcastTree(leaf, &num_of_leaves);
+  ASSERT_TRUE(tree.Has<BroadcastBranch<BroadcastTree>>());
+  const auto& branch = tree.Get<BroadcastBranch<BroadcastTree>>();
+  const auto& [cstr_broadcastable,
+               lhs_eq_rhs_tree,
+               lhs_eq_one_tree,
+               rhs_eq_one_tree] = branch.tuple();
+  ASSERT_EQ(cstr_broadcastable->lhs, DimExpr("S1"));
+  ASSERT_EQ(cstr_broadcastable->rhs, DimExpr("S2"));
+  ASSERT_TRUE(lhs_eq_rhs_tree.Has<BroadcastLeaf>());
+  ASSERT_TRUE(lhs_eq_one_tree.Has<BroadcastLeaf>());
+  ASSERT_TRUE(rhs_eq_one_tree.Has<BroadcastLeaf>());
+  CheckLeafNonBroadcast(lhs_eq_rhs_tree.Get<BroadcastLeaf>());
+  CheckLeafNonBroadcast(lhs_eq_one_tree.Get<BroadcastLeaf>());
+  CheckLeafNonBroadcast(rhs_eq_one_tree.Get<BroadcastLeaf>());
+}
+
+}  // namespace cinn::common