wehub-resource-sync
94 lines
3.2 KiB
C++
94 lines
3.2 KiB
C++
// 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
|