wehub-resource-sync
337 lines
11 KiB
C++
337 lines
11 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_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 {
|
|
|
|
/*
|
|
serial for (i, 0ll, 32768ll) {
|
|
serial for (j, 0ll, 16ll) {
|
|
serial for (reduce_k_0, 0ll, 128ll) {
|
|
var_18[i, j] = select((var_18[i, j] > var_17[i, j, reduce_k_0]),
|
|
var_18[i, j], var_17[i, j, reduce_k_0])
|
|
}
|
|
}
|
|
}
|
|
}
|
|
*/
|
|
TEST(IRSimplifySelect, SimplifySelectToMax) {
|
|
Context::Global().ResetNameId();
|
|
|
|
// Create input IR matching the specified pattern
|
|
const std::vector<ir::Expr> shape_2d = {ir::Expr(32768), ir::Expr(16)};
|
|
const std::vector<ir::Expr> shape_3d = {
|
|
ir::Expr(32768), ir::Expr(16), ir::Expr(128)};
|
|
|
|
ir::Tensor var_17 =
|
|
ir::_Tensor_::Make("var_17", ir::Float(32), shape_3d, shape_3d);
|
|
var_17->WithBuffer("global", "var_17_buffer");
|
|
|
|
ir::Tensor var_18 =
|
|
ir::_Tensor_::Make("var_18", ir::Float(32), shape_2d, shape_2d);
|
|
var_18->WithBuffer("global", "var_18_buffer");
|
|
|
|
// Define loop variables
|
|
ir::Var var_i = ir::Var(ir::Expr(0), ir::Expr(32768), "i");
|
|
ir::Var var_j = ir::Var(ir::Expr(0), ir::Expr(16), "j");
|
|
ir::Var var_reduce_k_0 = ir::Var(ir::Expr(0), ir::Expr(128), "reduce_k_0");
|
|
|
|
// Create innermost reduction loop body
|
|
ir::Expr reduce_body = ir::Store::Make(
|
|
var_18,
|
|
ir::Select::Make(
|
|
ir::GT::Make(ir::Load::Make(var_18, {var_i, var_j}),
|
|
ir::Load::Make(var_17, {var_i, var_j, var_reduce_k_0})),
|
|
ir::Load::Make(var_18, {var_i, var_j}),
|
|
ir::Load::Make(var_17, {var_i, var_j, var_reduce_k_0})),
|
|
{var_i, var_j});
|
|
|
|
// Create reduction loop
|
|
ir::Expr reduce_loop = ir::For::Make(var_reduce_k_0,
|
|
ir::Expr(0),
|
|
ir::Expr(128),
|
|
ir::ForType::Serial,
|
|
ir::DeviceAPI::Host,
|
|
ir::Block::Make({reduce_body}));
|
|
|
|
// Create j loop
|
|
ir::Expr j_loop = ir::For::Make(var_j,
|
|
ir::Expr(0),
|
|
ir::Expr(16),
|
|
ir::ForType::Serial,
|
|
ir::DeviceAPI::Host,
|
|
ir::Block::Make({reduce_loop}));
|
|
|
|
// Create i loop
|
|
ir::Expr i_loop = ir::For::Make(var_i,
|
|
ir::Expr(0),
|
|
ir::Expr(32768),
|
|
ir::ForType::Serial,
|
|
ir::DeviceAPI::Host,
|
|
ir::Block::Make({j_loop}));
|
|
|
|
// Final expression
|
|
ir::Expr expr = ir::Block::Make({i_loop});
|
|
|
|
VLOG(6) << "Before Simplify: " << expr;
|
|
Simplify(&expr);
|
|
VLOG(6) << "After Simplify: " << expr;
|
|
|
|
// Expected output verification
|
|
std::string expected_ir = R"ROC({
|
|
serial for (i, 0, 32768)
|
|
{
|
|
serial for (j, 0, 16)
|
|
{
|
|
serial for (reduce_k_0, 0, 128)
|
|
{
|
|
var_18[i, j] = cinn_max(var_17[i, j, reduce_k_0], var_18[i, j])
|
|
}
|
|
}
|
|
}
|
|
})ROC";
|
|
|
|
EXPECT_EQ(utils::GetStreamCnt(expr), utils::Trim(expected_ir));
|
|
}
|
|
|
|
/*
|
|
serial for (i, 0ll, 32768ll) {
|
|
serial for (j, 0ll, 16ll) {
|
|
serial for (reduce_k_0, 0ll, 128ll) {
|
|
var_18[i, j] = select((var_18[i, j] < var_17[i, j, reduce_k_0]),
|
|
var_18[i, j], var_17[i, j, reduce_k_0])
|
|
}
|
|
}
|
|
}
|
|
}
|
|
*/
|
|
TEST(IRSimplifySelect, SimplifySelectToMin) {
|
|
Context::Global().ResetNameId();
|
|
|
|
// Create input IR matching the specified pattern
|
|
const std::vector<ir::Expr> shape_2d = {ir::Expr(32768), ir::Expr(16)};
|
|
const std::vector<ir::Expr> shape_3d = {
|
|
ir::Expr(32768), ir::Expr(16), ir::Expr(128)};
|
|
|
|
ir::Tensor var_17 =
|
|
ir::_Tensor_::Make("var_17", ir::Float(32), shape_3d, shape_3d);
|
|
var_17->WithBuffer("global", "var_17_buffer");
|
|
|
|
ir::Tensor var_18 =
|
|
ir::_Tensor_::Make("var_18", ir::Float(32), shape_2d, shape_2d);
|
|
var_18->WithBuffer("global", "var_18_buffer");
|
|
|
|
// Define loop variables
|
|
ir::Var var_i = ir::Var(ir::Expr(0), ir::Expr(32768), "i");
|
|
ir::Var var_j = ir::Var(ir::Expr(0), ir::Expr(16), "j");
|
|
ir::Var var_reduce_k_0 = ir::Var(ir::Expr(0), ir::Expr(128), "reduce_k_0");
|
|
|
|
// Create innermost reduction loop body
|
|
ir::Expr reduce_body = ir::Store::Make(
|
|
var_18,
|
|
ir::Select::Make(
|
|
ir::LT::Make(ir::Load::Make(var_18, {var_i, var_j}),
|
|
ir::Load::Make(var_17, {var_i, var_j, var_reduce_k_0})),
|
|
ir::Load::Make(var_18, {var_i, var_j}),
|
|
ir::Load::Make(var_17, {var_i, var_j, var_reduce_k_0})),
|
|
{var_i, var_j});
|
|
|
|
// Create reduction loop
|
|
ir::Expr reduce_loop = ir::For::Make(var_reduce_k_0,
|
|
ir::Expr(0),
|
|
ir::Expr(128),
|
|
ir::ForType::Serial,
|
|
ir::DeviceAPI::Host,
|
|
ir::Block::Make({reduce_body}));
|
|
|
|
// Create j loop
|
|
ir::Expr j_loop = ir::For::Make(var_j,
|
|
ir::Expr(0),
|
|
ir::Expr(16),
|
|
ir::ForType::Serial,
|
|
ir::DeviceAPI::Host,
|
|
ir::Block::Make({reduce_loop}));
|
|
|
|
// Create i loop
|
|
ir::Expr i_loop = ir::For::Make(var_i,
|
|
ir::Expr(0),
|
|
ir::Expr(32768),
|
|
ir::ForType::Serial,
|
|
ir::DeviceAPI::Host,
|
|
ir::Block::Make({j_loop}));
|
|
|
|
// Final expression
|
|
ir::Expr expr = ir::Block::Make({i_loop});
|
|
|
|
VLOG(6) << "Before Simplify: " << expr;
|
|
Simplify(&expr);
|
|
VLOG(6) << "After Simplify: " << expr;
|
|
|
|
// Expected output verification
|
|
std::string expected_ir = R"ROC({
|
|
serial for (i, 0, 32768)
|
|
{
|
|
serial for (j, 0, 16)
|
|
{
|
|
serial for (reduce_k_0, 0, 128)
|
|
{
|
|
var_18[i, j] = cinn_min(var_18[i, j], var_17[i, j, reduce_k_0])
|
|
}
|
|
}
|
|
}
|
|
})ROC";
|
|
|
|
EXPECT_EQ(utils::GetStreamCnt(expr), utils::Trim(expected_ir));
|
|
}
|
|
|
|
/*
|
|
serial for (i, 0ll, 32768ll)
|
|
{
|
|
serial for (j, 0, 16)
|
|
{
|
|
serial for (j_0, 0, 128)
|
|
{
|
|
var_45[i, j, j_0)] = select(
|
|
(var_18[i, ((((j * 128ll) + j_0) / 128ll) + 0ll)] <=
|
|
float32(3.4028234663852886e+38)),
|
|
select(
|
|
(var_18[i, ((((j * 128ll) + j_0) / 128ll) + 0ll)] >=
|
|
float32(9.9999997473787516e-05)),
|
|
var_18[i, ((((j * 128ll) + j_0) / 128ll) + 0ll)],
|
|
float32(9.9999997473787516e-05)
|
|
),
|
|
float32(3.4028234663852886e+38)
|
|
)
|
|
}
|
|
}
|
|
}
|
|
*/
|
|
TEST(IRSimplifySelect, SimplifySelectToMinMax) {
|
|
Context::Global().ResetNameId();
|
|
|
|
// Create input IR matching the specified pattern
|
|
const std::vector<ir::Expr> shape_2d = {ir::Expr(32768), ir::Expr(16)};
|
|
const std::vector<ir::Expr> shape_3d = {
|
|
ir::Expr(32768), ir::Expr(16), ir::Expr(128)};
|
|
|
|
ir::Tensor var_18 =
|
|
ir::_Tensor_::Make("var_18", ir::Float(32), shape_2d, shape_2d);
|
|
var_18->WithBuffer("global", "var_18_buffer");
|
|
|
|
ir::Tensor var_45 =
|
|
ir::_Tensor_::Make("var_45", ir::Float(32), shape_3d, shape_3d);
|
|
var_45->WithBuffer("global", "var_45_buffer");
|
|
|
|
// Define loop variables
|
|
ir::Var var_i = ir::Var(ir::Expr(0), ir::Expr(32768), "i");
|
|
ir::Var var_j = ir::Var(ir::Expr(0), ir::Expr(16), "j");
|
|
ir::Var var_j_0 = ir::Var(ir::Expr(0), ir::Expr(128), "j_0");
|
|
|
|
// Create innermost loop body
|
|
ir::Expr body = ir::Store::Make(
|
|
var_45,
|
|
ir::Select::Make(
|
|
ir::LE::Make(
|
|
ir::Load::Make(
|
|
var_18,
|
|
{var_i,
|
|
ir::Div::Make(
|
|
ir::Add::Make(ir::Mul::Make(var_j, ir::Expr(128)),
|
|
var_j_0),
|
|
ir::Expr(128))}),
|
|
ir::Expr(3.4028234663852886e+38f)),
|
|
ir::Select::Make(
|
|
ir::GE::Make(
|
|
ir::Load::Make(
|
|
var_18,
|
|
{var_i,
|
|
ir::Div::Make(
|
|
ir::Add::Make(ir::Mul::Make(var_j, ir::Expr(128)),
|
|
var_j_0),
|
|
ir::Expr(128))}),
|
|
ir::Expr(9.9999997473787516e-05f)),
|
|
ir::Load::Make(
|
|
var_18,
|
|
{var_i,
|
|
ir::Div::Make(
|
|
ir::Add::Make(ir::Mul::Make(var_j, ir::Expr(128)),
|
|
var_j_0),
|
|
ir::Expr(128))}),
|
|
ir::Expr(9.9999997473787516e-05f)),
|
|
ir::Expr(3.4028234663852886e+38f)),
|
|
{var_i, var_j, var_j_0});
|
|
|
|
// Create j_0 loop
|
|
ir::Expr j_0_loop = ir::For::Make(var_j_0,
|
|
ir::Expr(0),
|
|
ir::Expr(128),
|
|
ir::ForType::Serial,
|
|
ir::DeviceAPI::Host,
|
|
ir::Block::Make({body}));
|
|
|
|
// Create j loop
|
|
ir::Expr j_loop = ir::For::Make(var_j,
|
|
ir::Expr(0),
|
|
ir::Expr(16),
|
|
ir::ForType::Serial,
|
|
ir::DeviceAPI::Host,
|
|
ir::Block::Make({j_0_loop}));
|
|
|
|
// Create i loop
|
|
ir::Expr i_loop = ir::For::Make(var_i,
|
|
ir::Expr(0),
|
|
ir::Expr(32768),
|
|
ir::ForType::Serial,
|
|
ir::DeviceAPI::Host,
|
|
ir::Block::Make({j_loop}));
|
|
|
|
// Final expression
|
|
ir::Expr expr = ir::Block::Make({i_loop});
|
|
|
|
VLOG(6) << "Before Simplify: " << expr;
|
|
Simplify(&expr);
|
|
VLOG(6) << "After Simplify: " << expr;
|
|
|
|
// Expected output verification
|
|
std::string expected_ir = R"ROC({
|
|
serial for (i, 0, 32768)
|
|
{
|
|
serial for (j, 0, 16)
|
|
{
|
|
serial for (j_0, 0, 128)
|
|
{
|
|
var_45[i, j, j_0] = cinn_min(cinn_max(var_18[i, (((j * 128) + j_0) / 128)], 9.99999975e-05f), 3.40282347e+38f)
|
|
}
|
|
}
|
|
}
|
|
})ROC";
|
|
|
|
EXPECT_EQ(utils::GetStreamCnt(expr), utils::Trim(expected_ir));
|
|
}
|
|
} // namespace optim
|
|
} // namespace cinn
|