wehub-resource-sync
707 lines
27 KiB
C++
707 lines
27 KiB
C++
// Copyright (c) 2025 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 "paddle/cinn/optim/realize_composite_reduce_pass.h"
|
|
#include "paddle/cinn/hlir/pe/reduction.h"
|
|
#include "paddle/cinn/ir/ir_mutator.h"
|
|
#include "paddle/cinn/ir/ir_printer.h"
|
|
#include "paddle/cinn/ir/stmt_visitors.h"
|
|
#include "paddle/cinn/ir/utils/ir_copy.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
|
|
namespace cinn {
|
|
namespace optim {
|
|
|
|
using ir::stmt::Alloc;
|
|
using ir::stmt::BlockRef;
|
|
using ir::stmt::Evaluate;
|
|
using ir::stmt::For;
|
|
using ir::stmt::Free;
|
|
using ir::stmt::IfThenElse;
|
|
using ir::stmt::Let;
|
|
using ir::stmt::Schedule;
|
|
using ir::stmt::StmtRef;
|
|
using ir::stmt::Store;
|
|
using ReduceType = ir::Reduce::ReduceType;
|
|
|
|
namespace {
|
|
ReduceType GetReduceType(const ir::Expr& expr) {
|
|
if (auto it = expr.As<ir::Call>()) {
|
|
if (it->name == hlir::pe::kVarianceFuncName) {
|
|
return ReduceType::kVariance;
|
|
} else if (it->name == hlir::pe::kArgmaxFuncName) {
|
|
return ReduceType::kArgmax;
|
|
} else if (it->name == hlir::pe::kArgminFuncName) {
|
|
return ReduceType::kArgmin;
|
|
}
|
|
}
|
|
return ReduceType::kNone;
|
|
}
|
|
|
|
struct CompositeTypes : public std::vector<common::Type> {
|
|
ReduceType type = ReduceType::kNone;
|
|
explicit CompositeTypes(ReduceType _type = ReduceType::kNone) : type(_type) {
|
|
this->reserve(2);
|
|
}
|
|
|
|
bool operator==(const CompositeTypes& other) const {
|
|
if (this->type != other.type || other.size() != this->size()) return false;
|
|
for (size_t i = 0; i < other.size(); i++) {
|
|
if (this->at(i) != other.at(i)) return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
void Print() const {
|
|
VLOG(4) << "[CompositeTypes]: " << static_cast<int>(this->type);
|
|
for (auto _t : *this) {
|
|
VLOG(4) << _t;
|
|
}
|
|
}
|
|
};
|
|
|
|
CompositeTypes GetArgReduceUnderlyingType(const ir::Expr& expr) {
|
|
if (auto it = expr.As<ir::Call>()) {
|
|
if (it->name == hlir::pe::kArgmaxFuncName ||
|
|
it->name == hlir::pe::kArgminFuncName) {
|
|
// for cinn_argxxx func, the arg1 is the argidx
|
|
// we need to check the type of the input
|
|
auto argidx_call = it->read_args[1].As<ir::Call>();
|
|
if (argidx_call != nullptr && argidx_call->name.find("argidx_") == 0) {
|
|
CompositeTypes comp_types(it->name == hlir::pe::kArgminFuncName
|
|
? ReduceType::kArgmin
|
|
: ReduceType::kArgmax);
|
|
comp_types.push_back(argidx_call->read_args[0]->type());
|
|
comp_types.push_back(expr->type());
|
|
return comp_types;
|
|
}
|
|
} else if (it->name == hlir::pe::kVarianceFuncName) {
|
|
return CompositeTypes(ReduceType::kVariance);
|
|
}
|
|
}
|
|
return CompositeTypes();
|
|
}
|
|
|
|
void SetInitValue(Store store_stmt,
|
|
common::Type new_type,
|
|
const CompositeTypes& comp_type,
|
|
std::string prefix = "") {
|
|
// prefix: if target is x86, we can not call constructor for POD struct
|
|
// the intrinsic function for creating struct is usually "create_" + typename
|
|
ir::Expr init_value = store_stmt->value();
|
|
auto call_op = init_value.As<ir::Call>();
|
|
// if the type is already a call
|
|
if (call_op != nullptr) {
|
|
call_op->set_type(new_type);
|
|
if (call_op->name.find("argidx_") == 0 ||
|
|
call_op->name.find("welford_") == 0) {
|
|
call_op->name = prefix + call_op->name;
|
|
}
|
|
return;
|
|
}
|
|
if (comp_type.type == ReduceType::kVariance) {
|
|
store_stmt->set_value(ir::Call::Make(new_type,
|
|
prefix + new_type.customized_type(),
|
|
{init_value, init_value, init_value},
|
|
{},
|
|
ir::CallType::Intrinsic));
|
|
} else if (comp_type.type == ReduceType::kArgmax ||
|
|
comp_type.type == ReduceType::kArgmin) {
|
|
ir::Expr index_init = ir::Expr(0);
|
|
index_init->set_type(common::Int(32));
|
|
if (comp_type.at(1).is_int(64)) {
|
|
index_init->set_type(common::Int(64));
|
|
}
|
|
store_stmt->set_value(ir::Call::Make(new_type,
|
|
prefix + new_type.customized_type(),
|
|
{init_value, index_init},
|
|
{},
|
|
ir::CallType::Intrinsic));
|
|
} else {
|
|
PADDLE_THROW(::common::errors::Unimplemented(
|
|
"reduce_type '%d' not allowed.", static_cast<int>(comp_type.type)));
|
|
}
|
|
}
|
|
|
|
/**
|
|
* This function resolves undefined argidx type, for example:
|
|
* \code
|
|
* spatial inner loop (argidx type defined)
|
|
* tensor_0[...] = cinn_argmax(tensor_1[...], argidx_f32_i64(tensor_2[...],
|
|
* index))
|
|
*
|
|
* follow up cross thread reduce (argidx type undefined)
|
|
* tensor_3[...] = cinn_argmax(tensor_4[...], tensor_5[...])
|
|
* \endcode
|
|
* In the above undefined case, we can not extract value type, since both
|
|
* tensors (4 and 5) in the arguments will be of index type, which lefts the
|
|
* argidx type undefined. So this function basically checks whether tensor_5 is
|
|
* in the typed_buffers map. Since cross thread reduction usually follows
|
|
* spatial inner loop reduction, so normally, tensor_0 and tensor_5 will
|
|
* normally be the same in one reduce block. and since tensor_0's type is
|
|
* defined, we can use it to resolve tensor_5 (and thus, the undefined
|
|
* tensor_3)'s type.
|
|
*/
|
|
std::map<ir::Buffer, CompositeTypes> ResolveUndefinedArgIdxType(
|
|
std::map<ir::Buffer, CompositeTypes>&& typed_buffers,
|
|
std::vector<Store>&& stores) {
|
|
for (const auto& store_stmt : stores) {
|
|
if (auto call_stmt = store_stmt->value().As<ir::Call>()) {
|
|
if (call_stmt->name != hlir::pe::kArgmaxFuncName &&
|
|
call_stmt->name != hlir::pe::kArgminFuncName)
|
|
continue;
|
|
auto load_stmt = call_stmt->read_args[1].As<ir::Load>();
|
|
PADDLE_ENFORCE_NOT_NULL(load_stmt,
|
|
::common::errors::PreconditionNotMet(
|
|
"Non-spatial inner loop arg reduce func call "
|
|
"second argument must be load."));
|
|
auto it = typed_buffers.find(load_stmt->tensor.as_tensor()->buffer);
|
|
PADDLE_ENFORCE_NE(it,
|
|
typed_buffers.end(),
|
|
::common::errors::PreconditionNotMet(
|
|
"Referenced buffer '%s' should be defined.",
|
|
load_stmt->tensor.as_tensor()->buffer->name));
|
|
auto composite_type = it->second;
|
|
typed_buffers.emplace(store_stmt->tensor().as_tensor()->buffer,
|
|
composite_type);
|
|
}
|
|
}
|
|
return typed_buffers;
|
|
}
|
|
|
|
std::map<ir::Buffer, CompositeTypes> CollectTypedReduceBuffers(
|
|
const BlockRef& body, std::vector<Store>* arg_stores) {
|
|
std::map<ir::Buffer, CompositeTypes> typed_buffers;
|
|
const auto VisitFn = [&](const StmtRef& stmt) {
|
|
if (!stmt.isa<Store>()) return;
|
|
Store store_stmt = stmt.as<Store>();
|
|
if (GetReduceType(store_stmt->value()) != ReduceType::kNone) {
|
|
auto it = typed_buffers.find(store_stmt->tensor().as_tensor()->buffer);
|
|
if (it == typed_buffers.end()) {
|
|
auto composite_type = GetArgReduceUnderlyingType(store_stmt->value());
|
|
if (composite_type.type == ReduceType::kNone) {
|
|
arg_stores->emplace_back(store_stmt);
|
|
} else {
|
|
// defined composite type can be immediately stored
|
|
typed_buffers.emplace(store_stmt->tensor().as_tensor()->buffer,
|
|
composite_type);
|
|
}
|
|
} else {
|
|
// check whether we will have conflicted store types
|
|
PADDLE_ENFORCE_EQ(
|
|
it->second == GetArgReduceUnderlyingType(store_stmt->value()),
|
|
true,
|
|
::common::errors::PreconditionNotMet(
|
|
"Composite type conflict detected in the buffer map."));
|
|
}
|
|
}
|
|
};
|
|
|
|
ir::stmt::Visit(body, VisitFn, [](auto) {});
|
|
return typed_buffers;
|
|
}
|
|
|
|
void ReplaceOutputBufferX86(
|
|
const BlockRef& body,
|
|
const std::set<ir::Buffer>& out_buffer_map,
|
|
const std::map<ir::Buffer, CompositeTypes>& typed_buffers) {
|
|
// re-route the reduce_init buffer to the local staging buffer
|
|
// and set the type for the buffers correctly
|
|
struct BufferRelationRecorder {
|
|
Store reduce_init;
|
|
Store write_back;
|
|
};
|
|
std::map<ir::Buffer, BufferRelationRecorder> buffer_relations;
|
|
for (auto buffer : out_buffer_map) {
|
|
buffer_relations.emplace(buffer, BufferRelationRecorder());
|
|
}
|
|
const auto VisitFn = [&](const StmtRef& stmt) {
|
|
if (!stmt.isa<Store>()) return;
|
|
Store store_stmt = stmt.as<Store>();
|
|
|
|
auto* tensor = store_stmt->tensor().as_tensor();
|
|
auto& buffer = tensor->buffer;
|
|
auto buffer_it = buffer_relations.find(buffer);
|
|
// check whether the buffer is related to output args
|
|
if (buffer_it == buffer_relations.end()) return;
|
|
if (ir::IsReduceInitTensorName(tensor->name)) {
|
|
buffer_it->second.reduce_init = store_stmt;
|
|
} else {
|
|
buffer_it->second.write_back = store_stmt;
|
|
}
|
|
};
|
|
|
|
ir::stmt::Visit(body, VisitFn, [](auto) {});
|
|
|
|
for (auto& [_, buffer_rel] : buffer_relations) {
|
|
// both should be defined
|
|
if (!buffer_rel.reduce_init.defined() || !buffer_rel.write_back.defined())
|
|
continue;
|
|
auto wb_value = buffer_rel.write_back->value();
|
|
if (auto load_node = wb_value.As<ir::Load>()) {
|
|
auto wb_load_buffer = load_node->tensor.as_tensor()->buffer;
|
|
auto wb_load_it = typed_buffers.find(wb_load_buffer);
|
|
PADDLE_ENFORCE_NE(wb_load_it,
|
|
typed_buffers.end(),
|
|
::common::errors::Fatal(
|
|
"Buffer '%s' should be defined in typed_buffers.",
|
|
wb_load_buffer->name));
|
|
// set the buffer of the reduce_init to write back buffer
|
|
ir::Expr new_tensor =
|
|
ir::ir_utils::IRCopy(buffer_rel.reduce_init->tensor());
|
|
new_tensor.as_tensor()->buffer = wb_load_buffer;
|
|
buffer_rel.reduce_init->set_tensor(new_tensor);
|
|
}
|
|
}
|
|
}
|
|
|
|
Store GetStoreOfSchedule(const Schedule& stmt) {
|
|
Store store_stmt;
|
|
bool found = false;
|
|
const auto VisitFn = [&](StmtRef stmt) {
|
|
if (!found && stmt.isa<Store>()) {
|
|
store_stmt = stmt.as<Store>();
|
|
found = true;
|
|
}
|
|
};
|
|
ir::stmt::Visit(stmt->body(), VisitFn, [](auto) {});
|
|
PADDLE_ENFORCE_EQ(found,
|
|
true,
|
|
::common::errors::PreconditionNotMet(
|
|
"One Schedule should have exactly one Store."));
|
|
return store_stmt;
|
|
}
|
|
|
|
Type GetCompositeReduceType(const Type& elem_type,
|
|
const CompositeTypes& composite_reduce) {
|
|
int type_bits = 0;
|
|
std::string rtype_name = "";
|
|
if (composite_reduce.type == ReduceType::kVariance) {
|
|
type_bits = elem_type.bits() * 3;
|
|
rtype_name = "welford" + hlir::pe::Type2StrForReduce(elem_type);
|
|
} else if (composite_reduce.type == ReduceType::kArgmax ||
|
|
composite_reduce.type == ReduceType::kArgmin) {
|
|
PADDLE_ENFORCE_GT(
|
|
composite_reduce.size(),
|
|
1,
|
|
::common::errors::InvalidArgument("CompositeTypes for arg reduce "
|
|
"must have at least two types"));
|
|
int max_bits =
|
|
std::max(composite_reduce[0].bits(), composite_reduce[1].bits());
|
|
type_bits = max_bits * 2;
|
|
rtype_name = "argidx" +
|
|
hlir::pe::Type2StrForArgReduce(composite_reduce[0]) +
|
|
hlir::pe::Type2StrForArgReduce(composite_reduce[1]);
|
|
} else {
|
|
PADDLE_THROW(::common::errors::InvalidArgument(
|
|
"Unsupported composite reduce type: %d",
|
|
static_cast<int>(composite_reduce.type)));
|
|
}
|
|
Type customized_type(ir::Type::type_t::Customized,
|
|
/* bits = */ type_bits,
|
|
/* width = */ 1);
|
|
customized_type.set_customized_type(rtype_name);
|
|
customized_type.set_cpp_const(false);
|
|
return customized_type;
|
|
}
|
|
|
|
struct StageReduceResultMutator : public ir::stmt::StmtMutator<> {
|
|
explicit StageReduceResultMutator(ir::LoweredFunc func) : func_(func) {
|
|
for (auto& arg : func->args) {
|
|
if (arg.is_buffer()) arg_buffers_.insert(arg.buffer_arg());
|
|
}
|
|
}
|
|
|
|
void operator()(BlockRef block) { VisitBlock(block); }
|
|
|
|
private:
|
|
void VisitStmt(Schedule stmt) override {
|
|
if (stmt->name().substr(0, 4) == "root") {
|
|
ir::stmt::StmtMutator<>::VisitBlock(stmt->body());
|
|
return;
|
|
}
|
|
Store store_stmt = GetStoreOfSchedule(stmt.as<Schedule>());
|
|
auto* store_tensor = store_stmt->tensor().as_tensor();
|
|
if (GetReduceType(store_stmt->value()) == ReduceType::kNone) return;
|
|
if (arg_buffers_.count(store_tensor->buffer) == 0) return;
|
|
|
|
// Create the staging buffer.
|
|
// We only need one element for this buffer, so its shape is {1}.
|
|
const std::vector<ir::Expr> shape = {ir::Expr(1)};
|
|
const std::vector<ir::Expr> indices = {ir::Expr(0)};
|
|
ir::Tensor staging_tensor =
|
|
ir::_Tensor_::Make(common::UniqName(store_tensor->name + "_local"),
|
|
store_tensor->buffer->dtype,
|
|
shape,
|
|
shape);
|
|
staging_tensor->WithBuffer("local", staging_tensor->name + "_buffer");
|
|
func_->temp_bufs.push_back(staging_tensor->buffer);
|
|
|
|
// Create the staging Schedule.
|
|
Schedule staging_schedule(stmt->iter_vars(),
|
|
stmt->iter_values(),
|
|
stmt->read_buffers(),
|
|
stmt->write_buffers(),
|
|
staging_tensor->name,
|
|
ir::ir_utils::IRCopy(stmt->body()),
|
|
stmt->attrs(),
|
|
stmt->reduce_method());
|
|
sibling_stmts_.push_back(staging_schedule);
|
|
|
|
// Replace all uses of the composite reduce buffer with the staging buffer.
|
|
Store staging_store = GetStoreOfSchedule(staging_schedule);
|
|
staging_store->set_tensor(staging_tensor);
|
|
staging_store->set_indices(indices);
|
|
ir::Expr staging_value = staging_store->value();
|
|
staging_value.As<ir::Call>()->read_args[0] =
|
|
ir::Load::Make(staging_tensor, indices);
|
|
staging_store->set_value(staging_value);
|
|
store_stmt->set_value(ir::Load::Make(staging_tensor, indices));
|
|
|
|
// Remove the reduction flags in the current Schedule, because reduction
|
|
// has been done in the staging Schedule.
|
|
std::vector<ir::Var> new_iter_vars;
|
|
for (auto& var : stmt->iter_vars()) {
|
|
ir::Var new_var = var->Copy().as_var_ref();
|
|
new_var->is_reduce_axis = false;
|
|
new_iter_vars.push_back(new_var);
|
|
}
|
|
stmt->set_iter_vars(new_iter_vars);
|
|
}
|
|
|
|
void VisitBlock(BlockRef block) override {
|
|
std::vector<StmtRef> old_stmts;
|
|
old_stmts.swap(sibling_stmts_);
|
|
|
|
for (StmtRef stmt : block->stmts()) {
|
|
ir::stmt::StmtMutator<>::VisitStmt(stmt);
|
|
sibling_stmts_.push_back(stmt);
|
|
}
|
|
|
|
block->set_stmts(sibling_stmts_);
|
|
sibling_stmts_ = std::move(old_stmts);
|
|
}
|
|
|
|
void VisitStmt(For stmt) override { VisitBlock(stmt->body()); }
|
|
|
|
void VisitStmt(IfThenElse stmt) override {
|
|
ir::stmt::BlockRef true_case = stmt->true_case();
|
|
VisitBlock(true_case);
|
|
stmt->set_true_case(true_case);
|
|
if (stmt->false_case().defined()) {
|
|
ir::stmt::BlockRef false_case = stmt->false_case();
|
|
VisitBlock(false_case);
|
|
stmt->set_false_case(false_case);
|
|
}
|
|
}
|
|
|
|
void VisitStmt(Let stmt) override { return; }
|
|
void VisitStmt(Store stmt) override { return; }
|
|
void VisitStmt(Alloc stmt) override { return; }
|
|
void VisitStmt(Free stmt) override { return; }
|
|
void VisitStmt(Evaluate stmt) override { return; }
|
|
|
|
private:
|
|
ir::LoweredFunc func_;
|
|
// buffers in the function's argument list
|
|
std::set<ir::Buffer> arg_buffers_;
|
|
// stmts at the same level with the currently visiting stmt
|
|
std::vector<StmtRef> sibling_stmts_;
|
|
};
|
|
|
|
struct LoadTypeMutator : public ir::IRMutator<> {
|
|
explicit LoadTypeMutator(
|
|
const std::map<ir::Buffer, std::pair<ir::Type, CompositeTypes>>&
|
|
buffer2type)
|
|
: buffer2type_(buffer2type) {}
|
|
|
|
void operator()(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
|
|
|
|
private:
|
|
void Visit(const ir::Load* op, ir::Expr* expr) override {
|
|
ir::IRMutator<>::Visit(op, expr);
|
|
auto* node = expr->As<ir::Load>();
|
|
auto& buffer = node->tensor.as_tensor()->buffer;
|
|
auto it = buffer2type_.find(buffer);
|
|
if (it != buffer2type_.end()) {
|
|
const auto& [buffer_type, composite_type] = it->second;
|
|
ir::Type new_type = GetCompositeReduceType(buffer_type, composite_type);
|
|
node->tensor.as_tensor()->set_type(new_type);
|
|
buffer->dtype = new_type;
|
|
*expr = ir::Cast::Make(buffer_type, *expr);
|
|
}
|
|
}
|
|
|
|
void UncastType(ir::Expr* expr) {
|
|
auto* cast_node = expr->As<ir::Cast>();
|
|
if (!cast_node) return;
|
|
auto* load_node = cast_node->v().As<ir::Load>();
|
|
if (!load_node) return;
|
|
if (buffer2type_.count(load_node->tensor.as_tensor()->buffer) > 0) {
|
|
*expr = cast_node->v();
|
|
}
|
|
}
|
|
|
|
void Visit(const ir::Call* op, ir::Expr* expr) override {
|
|
// this function will cast the buffer from composite type
|
|
// to an underlying type, for example welford_fp32 -> float
|
|
// uncast will undo this process
|
|
ir::IRMutator<>::Visit(op, expr);
|
|
// By default, all tensors are casted back to their element type
|
|
// before doing other computation. However, for the composite reduction
|
|
// calls, we shouldn't cast the arguments back because they hold the
|
|
// intermediate status.
|
|
if (GetReduceType(*expr) != ReduceType::kNone) {
|
|
auto* node = expr->As<ir::Call>();
|
|
UncastType(&(node->read_args[0]));
|
|
UncastType(&(node->read_args[1]));
|
|
}
|
|
}
|
|
|
|
const std::map<ir::Buffer, std::pair<ir::Type, CompositeTypes>>& buffer2type_;
|
|
};
|
|
|
|
void SetBufferType(ir::LoweredFunc func,
|
|
const std::map<ir::Buffer, CompositeTypes>& typed_buffers,
|
|
bool is_x86_arch) {
|
|
// Make a map from the buffers to their element and composite reduce types,
|
|
// otherwise it's hard to know a buffer's original type. The original type
|
|
// must be known to perform casting (back) in LoadTypeMutator::Visit()
|
|
std::map<ir::Buffer, std::pair<ir::Type, CompositeTypes>> buffer2type;
|
|
for (auto& [buffer, reduce_type] : typed_buffers) {
|
|
buffer2type.emplace(buffer, std::make_pair(buffer->dtype, reduce_type));
|
|
}
|
|
|
|
// Set function's temp_bufs type
|
|
for (auto& buffer : func->temp_bufs) {
|
|
auto it = buffer2type.find(buffer);
|
|
if (it != buffer2type.end()) {
|
|
const auto& [buffer_type, composite_type] = it->second;
|
|
buffer->dtype = GetCompositeReduceType(buffer_type, composite_type);
|
|
}
|
|
}
|
|
|
|
const auto VisitFn = [&](StmtRef stmt) {
|
|
if (!stmt.isa<Store>()) return;
|
|
Store store_stmt = stmt.as<Store>();
|
|
auto* tensor = store_stmt->tensor().as_tensor();
|
|
auto& buffer = tensor->buffer;
|
|
|
|
// Set store buffer type
|
|
auto it = buffer2type.find(buffer);
|
|
if (it != buffer2type.end()) {
|
|
ir::Expr new_tensor = ir::ir_utils::IRCopy(store_stmt->tensor());
|
|
const auto& [buffer_type, composite_type] = it->second;
|
|
ir::Type new_type = GetCompositeReduceType(buffer_type, composite_type);
|
|
new_tensor.as_tensor()->set_type(new_type);
|
|
new_tensor.as_tensor()->buffer->dtype = new_type;
|
|
store_stmt->set_tensor(new_tensor);
|
|
stmt->set_type(new_type);
|
|
if (ir::IsReduceInitTensorName(new_tensor.as_tensor()->name)) {
|
|
std::string call_prefix = is_x86_arch ? "create_" : "";
|
|
SetInitValue(store_stmt, new_type, composite_type, call_prefix);
|
|
}
|
|
}
|
|
|
|
// Set load buffer type
|
|
ir::Expr new_value = ir::ir_utils::IRCopy(store_stmt->value());
|
|
LoadTypeMutator load_type_mutator(buffer2type);
|
|
load_type_mutator(&new_value);
|
|
store_stmt->set_value(new_value);
|
|
};
|
|
|
|
ir::stmt::Mutate(func->body_block, VisitFn, [](auto) {});
|
|
}
|
|
|
|
struct ReduceExternCallMutator : public ir::IRMutator<> {
|
|
void operator()(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
|
|
|
|
private:
|
|
void Visit(const ir::Call* op, ir::Expr* expr) override {
|
|
ir::IRMutator<>::Visit(op, expr);
|
|
auto reduce_type_ = GetReduceType(*expr);
|
|
if (reduce_type_ == ReduceType::kNone) return;
|
|
ir::Expr lhs = op->read_args[0];
|
|
ir::Expr rhs = op->read_args[1];
|
|
if (lhs.type() != rhs.type()) {
|
|
if (auto call_op = rhs.As<ir::Call>()) {
|
|
// for argidx type, avoid redundant type casting, but this is ugly
|
|
if (call_op->name.find("argidx") != std::string::npos) {
|
|
rhs->set_type(lhs.type());
|
|
}
|
|
} else {
|
|
rhs = ir::Cast::Make(lhs.type(), rhs);
|
|
}
|
|
}
|
|
if (reduce_type_ == ReduceType::kVariance) {
|
|
// replace cinn_reduce_variance to operator+
|
|
*expr = ir::Add::Make(lhs, rhs);
|
|
} else if (reduce_type_ == ReduceType::kArgmax ||
|
|
reduce_type_ == ReduceType::kArgmin) {
|
|
// replace cinn_argmxx_iyy to max or min (overloaded)
|
|
if (op->name.find("argmax") != std::string::npos) {
|
|
*expr = ir::Max::Make(lhs, rhs);
|
|
} else {
|
|
*expr = ir::Min::Make(lhs, rhs);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
struct ReduceExternCallMutatorX86 : public ir::IRMutator<> {
|
|
// unlike non x86 counterpart, we do not replace the call
|
|
// by a arithmetic IR node, but instead call x86-exclusive funcs
|
|
void operator()(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
|
|
|
|
private:
|
|
void Visit(const ir::Call* op, ir::Expr* expr) override {
|
|
ir::IRMutator<>::Visit(op, expr);
|
|
auto reduce_type_ = GetReduceType(*expr);
|
|
if (reduce_type_ == ReduceType::kNone) return;
|
|
ir::Expr lhs = op->read_args[0];
|
|
ir::Expr rhs = op->read_args[1];
|
|
std::string lhs_type = lhs.type().to_string();
|
|
if (lhs.type() != rhs.type()) {
|
|
if (auto call_op = rhs.As<ir::Call>()) {
|
|
// for argidx type, avoid redundant type casting, but this is ugly
|
|
if (call_op->name.find("argidx") == 0) {
|
|
call_op->name = "create_" + call_op->name;
|
|
rhs->set_type(lhs.type());
|
|
}
|
|
} else {
|
|
// welford pod type call create function on x86
|
|
ir::Expr m2_init(0.f), weight_init(1.f);
|
|
if (lhs_type == "welford_fp64") {
|
|
m2_init->set_type(common::F64());
|
|
weight_init->set_type(common::F64());
|
|
}
|
|
rhs = ir::Call::Make(lhs.type(),
|
|
"create_" + lhs_type,
|
|
{rhs, m2_init, weight_init},
|
|
{},
|
|
ir::CallType::Intrinsic);
|
|
}
|
|
}
|
|
std::string call_prefix = "";
|
|
switch (reduce_type_) {
|
|
case ReduceType::kVariance:
|
|
call_prefix = "sum_";
|
|
break;
|
|
case ReduceType::kArgmax:
|
|
call_prefix = "max_";
|
|
break;
|
|
case ReduceType::kArgmin:
|
|
call_prefix = "min_";
|
|
break;
|
|
default:
|
|
break;
|
|
}
|
|
*expr = ir::Call::Make(lhs.type(),
|
|
call_prefix + lhs_type,
|
|
{lhs, rhs},
|
|
{},
|
|
ir::CallType::Intrinsic);
|
|
}
|
|
};
|
|
|
|
void ReplaceReduceExternCall(const BlockRef& body, bool is_x86_arch = false) {
|
|
const auto VisitFn = [&](StmtRef stmt) {
|
|
if (!stmt.isa<Store>()) return;
|
|
Store store_stmt = stmt.as<Store>();
|
|
ir::Expr new_value = ir::ir_utils::IRCopy(store_stmt->value());
|
|
if (is_x86_arch) {
|
|
ReduceExternCallMutatorX86()(&new_value);
|
|
} else {
|
|
ReduceExternCallMutator()(&new_value);
|
|
}
|
|
store_stmt->set_value(new_value);
|
|
};
|
|
|
|
ir::stmt::Mutate(body, VisitFn, [](auto) {});
|
|
}
|
|
|
|
} // namespace
|
|
|
|
LogicalResult RealizeCompositeReducePass::Run(ir::LoweredFunc func) {
|
|
BlockRef body = func->body_block;
|
|
|
|
// Step 1. Create a staging buffer for composite reduction result if it is
|
|
// directly written to the function's argument. This is because the
|
|
// result and the argument have different data types, and we need a staging
|
|
// buffer to do casting properly.
|
|
// Note: theoretically, we don't need this mutator if all reduction results
|
|
// are explicitly written back to global memory by yield_stores. However,
|
|
// current CINN frontend cannot guarantee this, so we need to do staging by
|
|
// ourself if the expected yield_store is missing.
|
|
StageReduceResultMutator mutator(func);
|
|
mutator(body);
|
|
|
|
// Step 2. Collect buffers that are used for reduce computation.
|
|
std::vector<Store> arg_stores;
|
|
auto typed_buffers = CollectTypedReduceBuffers(body, &arg_stores);
|
|
if (typed_buffers.empty()) {
|
|
// not a composite reduce func
|
|
return LogicalResult::success();
|
|
}
|
|
typed_buffers = ResolveUndefinedArgIdxType(std::move(typed_buffers),
|
|
std::move(arg_stores));
|
|
|
|
bool is_x86_arch = false;
|
|
target_.arch.Match(
|
|
[&](std::variant<common::X86Arch>) {
|
|
/**
|
|
* trace the CPU buffer for reduce init. For x86 pass, schedule pass
|
|
* will not be applied, therefore, the reduce_init buffer will be the
|
|
* same as the output buffer, which leads to incorrect buffer type and
|
|
* op type for codegen
|
|
*
|
|
* (1) we first extract the buffer for each output arg
|
|
* (2) find all stores to the corresponding output buffer, this op is
|
|
* prior to the output type cast, for x86 IR, reduce_init and the
|
|
* writing back op uses the same buffer (output tensor buffer). (3)
|
|
* create a mapping. if the buffer of a store (the value of the store)
|
|
* is in the typed_buffer, we try finding the reduce_init related op,
|
|
* and change the the buffer and op type of the reduce_init
|
|
*/
|
|
is_x86_arch = true;
|
|
std::set<ir::Buffer> output_buffers;
|
|
for (auto& arg : func->args) {
|
|
if (!arg.is_output()) continue;
|
|
output_buffers.emplace(arg.buffer_arg());
|
|
}
|
|
ReplaceOutputBufferX86(body, output_buffers, typed_buffers);
|
|
},
|
|
[&](std::variant<common::NVGPUArch,
|
|
common::CustomDeviceArch,
|
|
common::HygonDCUArchHIP,
|
|
common::HygonDCUArchSYCL,
|
|
common::ARMArch,
|
|
common::UnknownArch>) {});
|
|
// Step 3. Change the data type of buffers to the corresponding type.
|
|
SetBufferType(func, typed_buffers, is_x86_arch);
|
|
|
|
// Step 4. Replace the `cinn_reduce_variance` and `cinn_argmax` calls
|
|
// in order to reuse the cross-thread/block reduction templates.
|
|
ReplaceReduceExternCall(body, is_x86_arch);
|
|
|
|
return LogicalResult::success();
|
|
}
|
|
|
|
std::unique_ptr<FuncPass> CreateRealizeCompositeReducePass(Target target) {
|
|
return std::make_unique<RealizeCompositeReducePass>(target);
|
|
}
|
|
|
|
} // namespace optim
|
|
} // namespace cinn
|