wehub-resource-sync
296 lines
9.5 KiB
C++
296 lines
9.5 KiB
C++
// Copyright (c) 2024 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_block_reduction.h"
|
|
#include <vector>
|
|
|
|
#include "paddle/cinn/adt/adt.h"
|
|
#include "paddle/cinn/common/common.h"
|
|
#include "paddle/cinn/hlir/pe/reduction.h"
|
|
#include "paddle/cinn/ir/ir.h"
|
|
#include "paddle/cinn/ir/ir_mutator.h"
|
|
#include "paddle/cinn/ir/schedule/ir_schedule_util.h"
|
|
#include "paddle/cinn/lang/compute.h"
|
|
|
|
namespace cinn {
|
|
|
|
namespace optim {
|
|
namespace {
|
|
|
|
ir::Expr CalcBufferSizeInBytes(const ir::Buffer& buffer) {
|
|
const ir::Expr numel = buffer->SymbolicNumel();
|
|
return optim::ArithSimplify(numel * buffer->dtype.bytes());
|
|
}
|
|
|
|
std::unordered_set<std::string> GetReduceVarNames(
|
|
const ir::ScheduleBlockRealize* block_realize) {
|
|
const ir::ScheduleBlock* schedule_block =
|
|
block_realize->schedule_block.As<ir::ScheduleBlock>();
|
|
const std::vector<ir::Expr>& iter_values = block_realize->iter_values;
|
|
const std::vector<ir::Var>& iter_vars = schedule_block->iter_vars;
|
|
|
|
std::unordered_set<std::string> reduce_var_names;
|
|
for (int i = 0; i < iter_values.size(); ++i) {
|
|
if (!iter_vars[i]->is_reduce_axis) {
|
|
continue;
|
|
}
|
|
ir::ir_utils::CollectIRNodesWithoutTensor(
|
|
iter_values[i], [&](const ir::Expr* x) {
|
|
if (x->as_var()) {
|
|
reduce_var_names.insert(x->as_var()->name);
|
|
}
|
|
return false;
|
|
});
|
|
}
|
|
return reduce_var_names;
|
|
}
|
|
|
|
ir::Expr GetRightOperand(const ir::Expr& expr) {
|
|
#define GET_RIGHT_OPERAND(OpT) \
|
|
if (expr.As<OpT>()) { \
|
|
return expr.As<OpT>()->b(); \
|
|
}
|
|
|
|
GET_RIGHT_OPERAND(ir::Add);
|
|
GET_RIGHT_OPERAND(ir::Mul);
|
|
GET_RIGHT_OPERAND(ir::Max);
|
|
GET_RIGHT_OPERAND(ir::Min);
|
|
GET_RIGHT_OPERAND(ir::And);
|
|
GET_RIGHT_OPERAND(ir::Or);
|
|
|
|
#undef GET_RIGHT_OPERAND
|
|
PADDLE_THROW(
|
|
::common::errors::InvalidArgument("Not a supported reduce op: %s", expr));
|
|
}
|
|
|
|
struct BaseMutator : public ir::IRMutator<> {
|
|
using ir::IRMutator<>::Visit;
|
|
void operator()(ir::LoweredFunc fn) { Visit(fn.As<ir::_LoweredFunc_>()); }
|
|
|
|
protected:
|
|
bool IsGridReduce(const ir::ScheduleBlockRealize* block_realize) {
|
|
if (cur_loops_.empty()) {
|
|
return false;
|
|
}
|
|
auto* innermost_loop = cur_loops_.back();
|
|
if (!innermost_loop->is_gpu_block_binded()) {
|
|
return false;
|
|
}
|
|
const std::unordered_set<std::string> reduce_var_names =
|
|
GetReduceVarNames(block_realize);
|
|
return reduce_var_names.count(innermost_loop->loop_var->name) > 0;
|
|
}
|
|
|
|
void Visit(const ir::For* expr, ir::Expr* op) override {
|
|
cur_loops_.push_back(expr);
|
|
IRMutator::Visit(expr, op);
|
|
cur_loops_.pop_back();
|
|
}
|
|
|
|
void Visit(ir::Expr* expr) { IRMutator::Visit(expr, expr); }
|
|
|
|
protected:
|
|
std::vector<const ir::For*> cur_loops_;
|
|
};
|
|
|
|
struct CrossBlockReductionReplacer : public BaseMutator {
|
|
private:
|
|
void InsertTempSpaceToFuncArgs(ir::_LoweredFunc_* func_node,
|
|
const ir::Buffer& buffer,
|
|
bool need_zero_init) {
|
|
// insert the temp space after the last tensor argument and before the
|
|
// first scalar argument
|
|
auto insert_pos =
|
|
std::find_if(func_node->args.begin(),
|
|
func_node->args.end(),
|
|
[](const ir::Argument& arg) { return arg.is_var(); });
|
|
|
|
int arg_idx = std::distance(func_node->args.begin(), insert_pos);
|
|
func_node->temp_spaces.emplace_back(
|
|
CalcBufferSizeInBytes(buffer), arg_idx, need_zero_init);
|
|
|
|
ir::Argument temp_space_arg(buffer, ir::Argument::IO::kOutput);
|
|
func_node->args.insert(insert_pos, temp_space_arg);
|
|
}
|
|
|
|
void ConvertHeapBuffersToFuncArgs(ir::_LoweredFunc_* func_node) {
|
|
std::vector<ir::Buffer> global_bufs;
|
|
std::vector<ir::Buffer> local_bufs;
|
|
|
|
for (auto& buf : func_node->temp_bufs) {
|
|
if (buf->memory_type == ir::MemoryType::Heap) {
|
|
global_bufs.push_back(buf);
|
|
} else {
|
|
local_bufs.push_back(buf);
|
|
}
|
|
}
|
|
|
|
for (auto& buf : global_bufs) {
|
|
InsertTempSpaceToFuncArgs(func_node, buf, false);
|
|
}
|
|
func_node->temp_bufs = local_bufs;
|
|
}
|
|
|
|
ir::Expr GetBlockBindedSpatialLoopExtend(
|
|
const ir::ScheduleBlockRealize* block_realize) {
|
|
const std::unordered_set<std::string> reduce_var_names =
|
|
GetReduceVarNames(block_realize);
|
|
std::vector<ir::Expr> loop_extends;
|
|
for (auto* for_node : cur_loops_) {
|
|
if (reduce_var_names.count(for_node->loop_var->name) == 0 &&
|
|
for_node->is_gpu_block_binded()) {
|
|
loop_extends.push_back(for_node->extent);
|
|
}
|
|
}
|
|
PADDLE_ENFORCE_EQ(
|
|
loop_extends.size(),
|
|
1UL,
|
|
::common::errors::PreconditionNotMet(
|
|
"There should be exactly one spatial loop binded on gpu block."));
|
|
return loop_extends[0];
|
|
}
|
|
|
|
ir::Expr GetThreadBindedSpatialLoopExtend(
|
|
const ir::ScheduleBlockRealize* block_realize) {
|
|
const std::unordered_set<std::string> reduce_var_names =
|
|
GetReduceVarNames(block_realize);
|
|
std::vector<ir::Expr> loop_extends;
|
|
for (auto* for_node : cur_loops_) {
|
|
if (reduce_var_names.count(for_node->loop_var->name) == 0 &&
|
|
for_node->is_gpu_thread_binded()) {
|
|
loop_extends.push_back(for_node->extent);
|
|
}
|
|
}
|
|
PADDLE_ENFORCE_LE(
|
|
loop_extends.size(),
|
|
1UL,
|
|
::common::errors::PreconditionNotMet(
|
|
"There could be at most one spatial loop binded on gpu thread."));
|
|
if (loop_extends.empty()) {
|
|
return ir::Expr(1);
|
|
}
|
|
return loop_extends[0];
|
|
}
|
|
|
|
void ReplaceByGridReduceExternCall(const ir::ScheduleBlock* schedule_block,
|
|
const ir::Expr num_spatial_threads) {
|
|
ir::Expr update_stmt = schedule_block->body;
|
|
if (update_stmt.As<ir::Block>()) {
|
|
PADDLE_ENFORCE_EQ(
|
|
update_stmt.As<ir::Block>()->stmts.size(),
|
|
1UL,
|
|
::common::errors::InvalidArgument(
|
|
"There should be exactly one statement inside schedule_block."));
|
|
update_stmt = update_stmt.As<ir::Block>()->stmts[0];
|
|
}
|
|
PADDLE_ENFORCE_NOT_NULL(
|
|
update_stmt.As<ir::Store>(),
|
|
::common::errors::InvalidArgument(
|
|
"The top-level statement in schedule_block must be a store."));
|
|
|
|
auto* store_node = update_stmt.As<ir::Store>();
|
|
ir::Expr rvalue = GetRightOperand(store_node->value);
|
|
PADDLE_ENFORCE_NOT_NULL(rvalue.As<ir::Load>(),
|
|
::common::errors::InvalidArgument(
|
|
"The rvalue of reduce is not a load."));
|
|
|
|
std::string func_name = hlir::pe::GridReduceExternalFuncName(
|
|
store_node->value, store_node->tensor->type());
|
|
auto* load_node = rvalue.As<ir::Load>();
|
|
ir::Tensor rf_tensor = load_node->tensor.as_tensor_ref();
|
|
|
|
// The load's indices are like [ blockIdx.y, <spatial_index>... ].
|
|
// The loaded tensor's shape is like [ gridDim.y, <spatial_size>... ].
|
|
ir::Expr spatial_index = [&]() {
|
|
load_node->indices[0] = ir::Expr(0);
|
|
return load_node->index();
|
|
}();
|
|
ir::Expr spatial_size = [&]() {
|
|
load_node->indices[0] = ir::Expr(1);
|
|
for (int i = 1; i < load_node->indices.size(); i++) {
|
|
load_node->indices[i] = ir::Expr(0);
|
|
}
|
|
return load_node->index();
|
|
}();
|
|
|
|
store_node->value =
|
|
lang::CallExtern(func_name, {rf_tensor, spatial_size, spatial_index});
|
|
}
|
|
|
|
void Visit(ir::_LoweredFunc_* fn) override {
|
|
has_grid_reduce_ = false;
|
|
func_arg_buffer_names_.clear();
|
|
for (auto& arg : fn->args) {
|
|
if (arg.is_buffer()) {
|
|
func_arg_buffer_names_.insert(arg.buffer_arg()->name);
|
|
}
|
|
}
|
|
|
|
IRMutator::Visit(fn);
|
|
if (!has_grid_reduce_) {
|
|
return;
|
|
}
|
|
|
|
ConvertHeapBuffersToFuncArgs(fn);
|
|
}
|
|
|
|
void Visit(const ir::ScheduleBlockRealize* expr, ir::Expr* op) override {
|
|
const ir::ScheduleBlock* schedule_block =
|
|
expr->schedule_block.As<ir::ScheduleBlock>();
|
|
|
|
if (schedule_block->name.substr(0, 4) == "root") {
|
|
IRMutator::Visit(expr, op);
|
|
return;
|
|
}
|
|
|
|
if (!IsGridReduce(expr)) {
|
|
return;
|
|
}
|
|
|
|
if (!has_grid_reduce_) {
|
|
has_grid_reduce_ = true;
|
|
}
|
|
|
|
ir::Expr num_spatial_threads = GetThreadBindedSpatialLoopExtend(expr);
|
|
ReplaceByGridReduceExternCall(schedule_block, num_spatial_threads);
|
|
}
|
|
|
|
void Visit(const ir::Block* block, ir::Expr* op) override {
|
|
// We override the Block visitor to facilitate statement insertion.
|
|
std::vector<ir::Expr> old_parent_block_stmts;
|
|
old_parent_block_stmts.swap(cur_parent_block_stmts_);
|
|
auto* node = op->As<ir::Block>();
|
|
for (auto& stmt : node->stmts) {
|
|
IRMutator::Visit(&stmt, &stmt);
|
|
cur_parent_block_stmts_.push_back(stmt);
|
|
}
|
|
node->stmts = std::move(cur_parent_block_stmts_);
|
|
cur_parent_block_stmts_ = std::move(old_parent_block_stmts);
|
|
}
|
|
|
|
private:
|
|
std::vector<ir::Expr> cur_parent_block_stmts_;
|
|
std::unordered_set<std::string> func_arg_buffer_names_;
|
|
bool has_grid_reduce_{false};
|
|
};
|
|
|
|
} // namespace
|
|
|
|
void ReplaceCrossBlockReduction(ir::LoweredFunc fn) {
|
|
CrossBlockReductionReplacer()(fn);
|
|
}
|
|
|
|
} // namespace optim
|
|
} // namespace cinn
|