wehub-resource-sync
103 lines
3.2 KiB
C++
103 lines
3.2 KiB
C++
// Copyright (c) 2021 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/replace_cross_thread_reduction.h"
|
|
|
|
#include <gtest/gtest.h>
|
|
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#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/stmt_converter.h"
|
|
#include "paddle/cinn/utils/string.h"
|
|
|
|
namespace cinn {
|
|
namespace optim {
|
|
|
|
TEST(CrossThreadReductionReplacer, basic) {
|
|
#ifdef CINN_WITH_CUDA
|
|
Context::Global().ResetNameId();
|
|
Placeholder<float> A("A", {Expr(64), Expr(128)});
|
|
Target target = cinn::common::DefaultNVGPUTarget();
|
|
Module::Builder builder("reduce_sum", target);
|
|
Var reduce_j(128, "reduce_j");
|
|
ir::Tensor B = Compute(
|
|
{Expr(64)},
|
|
[&](Var i) { return lang::ReduceSum(A(i, reduce_j), {reduce_j}); },
|
|
"B");
|
|
ast_gen_ius::TensorGroup tensor_group({A, B});
|
|
auto func = lang::LowerToAst("reduce_sum", {A, B}, &tensor_group);
|
|
VLOG(6) << "original func\n" << func;
|
|
|
|
ir::Expr expr_func_body = ir::ConvertStmtBlockToExprBlock(func->body_block);
|
|
ir::ModuleExpr mod_expr({expr_func_body});
|
|
ir::IRSchedule ir_sch(mod_expr);
|
|
|
|
ir_sch.Bind(ir_sch.GetLoops("B")[0], "blockIdx.x");
|
|
ir_sch.Bind(ir_sch.GetLoops("B")[1], "threadIdx.x");
|
|
|
|
ir::Expr block = ir_sch.GetBlock("B");
|
|
block.As<ir::ScheduleBlockRealize>()
|
|
->schedule_block.As<ir::ScheduleBlock>()
|
|
->reduce_method = ir::BlockReduceMethod();
|
|
|
|
ir::Expr func_body = ir_sch.GetModule().GetExprs()[0];
|
|
std::vector<ir::Argument> args{
|
|
ir::Argument(ir::Var("A"), ir::Argument::IO::kInput),
|
|
ir::Argument(ir::Var("B"), ir::Argument::IO::kOutput)};
|
|
auto new_func = ir::_LoweredFunc_::Make("test_func", args, func_body, {});
|
|
VLOG(6) << "After Bind: " << new_func->body;
|
|
|
|
ReplaceCrossThreadReduction(new_func);
|
|
VLOG(6) << "After ReplaceCrossThreadReduction: " << new_func->body;
|
|
|
|
EXPECT_EQ(utils::GetStreamCnt(new_func->body), utils::Trim(R"ROC({
|
|
ScheduleBlock(root)
|
|
{
|
|
{
|
|
thread_bind[blockIdx.x] for (i, 0, 64)
|
|
{
|
|
ScheduleBlock(B__reduce_init)
|
|
{
|
|
i0 = axis.bind(i)
|
|
{
|
|
B__reduce_init[i0] = 0.00000000f
|
|
}
|
|
}
|
|
thread_bind[threadIdx.x] for (reduce_j, 0, 128)
|
|
{
|
|
ScheduleBlock(B)
|
|
{
|
|
i0_0, i1 = axis.bind(i, reduce_j)
|
|
{
|
|
B[i0_0] = cinn_block_reduce_sum_fp32(A[i0_0, i1], _Buffer_<cinn_buffer_t*: 32>(shm32__fp32_reduce), false)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
)ROC"));
|
|
#endif
|
|
}
|
|
|
|
} // namespace optim
|
|
} // namespace cinn
|