chore: import upstream snapshot with attribution

test/cpp/pir/cinn/generate_shape_util_test.cc

@@ -0,0 +1,93 @@
+// Copyright (c) 2023 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 "gtest/gtest.h"
+
+#include "paddle/cinn/hlir/dialect/operator/ir/generate_shape_util.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+#include "paddle/pir/include/dialect/shape/utils/dim_expr_builder.h"
+
+#include "test/cpp/pir/tools/test_pir_utils.h"
+
+namespace cinn::dialect {
+using namespace symbol;  // NOLINT
+
+namespace {
+DimExpr CreateExampleDimExpr() {
+  DimExprBuilder dim_expr_builder;
+  DimExpr sym0 = DimExpr("S0");
+  DimExpr sym1 = DimExpr("S1");
+  DimExpr constant = DimExpr(2);
+  DimExpr expr1 = (sym0 - sym1) * constant / sym0;
+  DimExpr expr2 = dim_expr_builder.Max(expr1, sym0);
+  DimExpr output = dim_expr_builder.Min(expr2, sym1);
+  return output;
+}
+}  // namespace
+
+TEST(DimExprUtil, Convert) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+
+  DimExpr dim_expr = CreateExampleDimExpr();
+  ::pir::Attribute attr = ConvertDimExprToAttribute(ctx, dim_expr);
+  std::optional<DimExpr> opt_expr = ConvertAttributeToDimExpr(attr);
+  ASSERT_TRUE(opt_expr.has_value());
+  ASSERT_EQ(opt_expr.value(), dim_expr);
+}
+
+TEST(DimExprUtil, Substitute) {
+  DimExpr dim_expr = CreateExampleDimExpr();
+  const auto& substituted_expr = SubstituteDimExpr(
+      dim_expr, [](const std::string& str) -> std::optional<DimExpr> {
+        if (str == "S0") {
+          return DimExpr("symbol0");
+        } else if (str == "S1") {
+          return DimExpr("symbol1");
+        } else {
+          return std::nullopt;
+        }
+      });
+  const auto& ret_expr = SubstituteDimExpr(
+      substituted_expr, [](const std::string& str) -> std::optional<DimExpr> {
+        if (str == "symbol0") {
+          return DimExpr("S0");
+        } else if (str == "symbol1") {
+          return DimExpr("S1");
+        } else {
+          return std::nullopt;
+        }
+      });
+  ASSERT_EQ(ret_expr, dim_expr);
+}
+
+TEST(DimExprUtil, MakeGetterDimExpr4SymbolName) {
+  cinn::dialect::GenerateShapeOp::SymbolBindings symbol_bindings{};
+  using ShapeSymbolBinding = cinn::dialect::GenerateShapeOp::ShapeSymbolBinding;
+  symbol_bindings.emplace_back(ShapeSymbolBinding{"Symbol", 0, 0});
+  const auto& dim_expr = CreateExampleDimExpr();
+
+  const auto& shape_or_data_dim_exprs = symbol::ShapeOrDataDimExprs(
+      symbol::TensorShapeOrDataDimExprs({dim_expr}));
+
+  const auto& DimExpr4SymbolName = MakeGetterDimExpr4SymbolName(
+      symbol_bindings,
+      [&](int in_tensor_idx) -> const symbol::ShapeOrDataDimExprs& {
+        return shape_or_data_dim_exprs;
+      });
+  const auto& opt_dim_expr = DimExpr4SymbolName("Symbol");
+  ASSERT_TRUE(opt_dim_expr.has_value());
+  ASSERT_EQ(opt_dim_expr.value(), dim_expr);
+}
+
+}  // namespace cinn::dialect