wehub-resource-sync
192 lines
5.4 KiB
C++
192 lines
5.4 KiB
C++
// Copyright (c) 2025 CINN 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/optim/ir_simplify.h"
|
|
|
|
#include <gtest/gtest.h>
|
|
|
|
#include "paddle/cinn/cinn.h"
|
|
#include "paddle/cinn/ir/ir.h"
|
|
#include "paddle/cinn/ir/ir_base.h"
|
|
#include "paddle/cinn/ir/ir_printer.h"
|
|
#include "paddle/cinn/ir/op/ir_operators.h"
|
|
#include "paddle/cinn/ir/schedule/ir_schedule.h"
|
|
#include "paddle/cinn/ir/utils/ir_nodes_collector.h"
|
|
#include "paddle/cinn/ir/utils/stmt_converter.h"
|
|
#include "paddle/cinn/utils/string.h"
|
|
|
|
namespace cinn {
|
|
namespace optim {
|
|
|
|
/*
|
|
i_j_fused: [0ll, 524288ll)
|
|
j_0: [0, 128)
|
|
Before Normalize:
|
|
(j_0 % 128)
|
|
After Normalize:
|
|
j_0
|
|
*/
|
|
TEST(IRSimplifyBound, SimplifyMod) {
|
|
Context::Global().ResetNameId();
|
|
|
|
// Create input IR matching the specified pattern
|
|
// Define loop variable
|
|
ir::Var var_j_0 = ir::Var(ir::Expr(0), ir::Expr(128), "j_0");
|
|
|
|
// Final expression
|
|
ir::Expr expr = ir::Mod::Make(var_j_0, ir::Expr(128));
|
|
|
|
VLOG(6) << "Before Simplify: " << expr;
|
|
auto res = expr.as_index().ir::IndexExpr::Normalize(
|
|
ir::IndexExpr::OptLevel::kLevel3);
|
|
VLOG(6) << "After Simplify: " << res;
|
|
|
|
// Expected output verification
|
|
std::string expected_ir = R"ROC(j_0)ROC";
|
|
|
|
EXPECT_EQ(utils::GetStreamCnt(res), utils::Trim(expected_ir));
|
|
}
|
|
|
|
/*
|
|
i_j_fused: [0ll, 524288ll)
|
|
j_0: [0, 128)
|
|
Before Normalize:
|
|
(j_0 / 128)
|
|
After Normalize:
|
|
0
|
|
*/
|
|
TEST(IRSimplifyBound, SimplifyDiv) {
|
|
Context::Global().ResetNameId();
|
|
|
|
// Create input IR matching the specified pattern
|
|
// Define loop variable
|
|
ir::Var var_j_0 = ir::Var(ir::Expr(0), ir::Expr(128), "j_0");
|
|
|
|
// Final expression
|
|
ir::Expr expr = ir::Div::Make(var_j_0, ir::Expr(128));
|
|
|
|
VLOG(6) << "Before Normalize: " << expr;
|
|
auto res = expr.as_index().ir::IndexExpr::Normalize(
|
|
ir::IndexExpr::OptLevel::kLevel3);
|
|
VLOG(6) << "After Normalize: " << res;
|
|
|
|
// Expected output verification
|
|
std::string expected_ir = R"ROC(0)ROC";
|
|
|
|
EXPECT_EQ(utils::GetStreamCnt(res), utils::Trim(expected_ir));
|
|
}
|
|
|
|
/*
|
|
i_j_fused: [0ll, 524288ll)
|
|
j_0: [0, 128)
|
|
Before Normalize:
|
|
((((i_j_fused % 16) * 128) + j_0) / 128)
|
|
After Normalize:
|
|
(i_j_fused % 16)
|
|
*/
|
|
TEST(IRSimplifyBound, SimplifyLinearDiv) {
|
|
Context::Global().ResetNameId();
|
|
|
|
// Create input IR matching the specified pattern
|
|
// Define loop variables
|
|
ir::Var var_i_j_fused = ir::Var(ir::Expr(0), ir::Expr(524288), "i_j_fused");
|
|
ir::Var var_j_0 = ir::Var(ir::Expr(0), ir::Expr(128), "j_0");
|
|
|
|
// Final expression
|
|
ir::Expr expr = ir::Div::Make(
|
|
ir::Add::Make(ir::Mul::Make(ir::Mod::Make(var_i_j_fused, ir::Expr(16)),
|
|
ir::Expr(128)),
|
|
var_j_0),
|
|
ir::Expr(128));
|
|
|
|
VLOG(6) << "Before Normalize: " << expr;
|
|
auto res = expr.as_index().ir::IndexExpr::Normalize(
|
|
ir::IndexExpr::OptLevel::kLevel3);
|
|
VLOG(6) << "After Normalize: " << res;
|
|
|
|
// Expected output verification
|
|
std::string expected_ir = R"ROC((i_j_fused % 16))ROC";
|
|
|
|
EXPECT_EQ(utils::GetStreamCnt(res), utils::Trim(expected_ir));
|
|
}
|
|
|
|
/*
|
|
i_j_fused: [0ll, 524288ll)
|
|
j_0: [0, 128)
|
|
Before Normalize:
|
|
((((i_j_fused % 16) * 128) + j_0) % 128)
|
|
After Normalize:
|
|
j_0
|
|
*/
|
|
TEST(IRSimplifyBound, SimplifyLinearMod) {
|
|
Context::Global().ResetNameId();
|
|
|
|
// Create input IR matching the specified pattern
|
|
// Define loop variables
|
|
ir::Var var_i_j_fused = ir::Var(ir::Expr(0), ir::Expr(524288), "i_j_fused");
|
|
ir::Var var_j_0 = ir::Var(ir::Expr(0), ir::Expr(128), "j_0");
|
|
|
|
// Final expression
|
|
ir::Expr expr = ir::Mod::Make(
|
|
ir::Add::Make(ir::Mul::Make(ir::Mod::Make(var_i_j_fused, ir::Expr(16)),
|
|
ir::Expr(128)),
|
|
var_j_0),
|
|
ir::Expr(128));
|
|
|
|
VLOG(6) << "Before Normalize: " << expr;
|
|
auto res = expr.as_index().ir::IndexExpr::Normalize(
|
|
ir::IndexExpr::OptLevel::kLevel3);
|
|
VLOG(6) << "After Normalize: " << res;
|
|
|
|
// Expected output verification
|
|
std::string expected_ir = R"ROC(j_0)ROC";
|
|
|
|
EXPECT_EQ(utils::GetStreamCnt(res), utils::Trim(expected_ir));
|
|
}
|
|
|
|
/*
|
|
loop_var_2: [0, 32)
|
|
loop_var_3: [0, 4)
|
|
Before Normalize:
|
|
(((loop_var_3 * 32ll) + loop_var_2) / 128ll)
|
|
After Normalize:
|
|
0
|
|
*/
|
|
TEST(IRSimplifyBound, SimplifyLinearDiv2) {
|
|
Context::Global().ResetNameId();
|
|
|
|
// Create input IR matching the specified pattern
|
|
// Define loop variables
|
|
ir::Var loop_var_2 = ir::Var(ir::Expr(0), ir::Expr(32), "loop_var_2");
|
|
ir::Var loop_var_3 = ir::Var(ir::Expr(0), ir::Expr(4), "loop_var_3");
|
|
|
|
// Final expression
|
|
ir::Expr expr = ir::Div::Make(
|
|
ir::Add::Make(ir::Mul::Make(loop_var_3, ir::Expr(32)), loop_var_2),
|
|
ir::Expr(128));
|
|
|
|
VLOG(6) << "Before Normalize: " << expr;
|
|
auto res = expr.as_index().ir::IndexExpr::Normalize(
|
|
ir::IndexExpr::OptLevel::kLevel3);
|
|
VLOG(6) << "After Normalize: " << res;
|
|
|
|
// Expected output verification
|
|
std::string expected_ir = R"ROC(0)ROC";
|
|
|
|
EXPECT_EQ(utils::GetStreamCnt(res), utils::Trim(expected_ir));
|
|
}
|
|
|
|
} // namespace optim
|
|
} // namespace cinn
|