wehub-resource-sync
254 lines
10 KiB
C++
254 lines
10 KiB
C++
// Copyright (c) 2024 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.
|
|
|
|
#pragma once
|
|
|
|
#include "paddle/cinn/operator_fusion/fusion_tracker/interpreter.h"
|
|
#include "glog/logging.h"
|
|
#include "paddle/cinn/hlir/framework/pir/trivial_op_util.h"
|
|
|
|
namespace cinn::fusion {
|
|
|
|
void RunCopyInstr(const std::shared_ptr<CopyInstr>& instr,
|
|
FusionInterpreter* interpreter) {
|
|
interpreter->scope[instr->new_name_] =
|
|
interpreter->scope.at(instr->origin_name_);
|
|
}
|
|
|
|
void RunCombineInstr(const std::shared_ptr<CombineInstr>& instr,
|
|
FusionInterpreter* interpreter) {
|
|
// TODO(@wuzhanfei)
|
|
ScopeElementPtr new_pattern = std::make_shared<ScopeElement>();
|
|
for (const auto& name : instr->names_) {
|
|
const auto& to_insert = interpreter->scope.at(name);
|
|
new_pattern->Extend(to_insert->fusion_ops);
|
|
}
|
|
VLOG(4) << "After CombineInstr Pattern: \n"
|
|
<< GetFusibleOpsExpr(new_pattern->fusion_ops);
|
|
interpreter->scope[instr->result_] = new_pattern;
|
|
}
|
|
|
|
void RunInitPatternInstr(const std::shared_ptr<InitPatternInstr>& instr,
|
|
FusionInterpreter* interpreter) {
|
|
ScopeElementPtr new_pattern = std::make_shared<ScopeElement>();
|
|
new_pattern->fusion_ops.emplace_back(
|
|
interpreter->initialized_lowered_op[(instr->get_idx())]);
|
|
interpreter->scope[instr->result_] = new_pattern;
|
|
}
|
|
|
|
void RunTrivialInlineInstr(const std::shared_ptr<TrivialInlineInstr>& instr,
|
|
FusionInterpreter* interpreter) {
|
|
PADDLE_ENFORCE_EQ(interpreter->scope[instr->upstream_]->fusion_ops.size(),
|
|
1,
|
|
::common::errors::InvalidArgument(
|
|
"Upstream op must have only one fusion_op."));
|
|
auto upstream_op = std::get<TrivialOp>(
|
|
interpreter->scope[instr->upstream_]->fusion_ops.front());
|
|
ScopeElementPtr new_pattern = std::make_shared<ScopeElement>();
|
|
|
|
auto DoTrivialFusion = [upstream_op](const auto& downstream_op) -> FusibleOp {
|
|
return cinn::hlir::framework::pir::trivial_fusion_detail::
|
|
TrivalxOther_Fusion(upstream_op, downstream_op);
|
|
};
|
|
|
|
for (auto downstream_op :
|
|
interpreter->scope[instr->downstream_]->fusion_ops) {
|
|
new_pattern->fusion_ops.emplace_back(
|
|
std::visit(DoTrivialFusion, downstream_op));
|
|
}
|
|
interpreter->scope[instr->result_] = new_pattern;
|
|
}
|
|
|
|
void RunTmpTransformInstr(const std::shared_ptr<TmpTransformInstr>& instr,
|
|
FusionInterpreter* interpreter) {
|
|
PADDLE_ENFORCE_GT(
|
|
interpreter->scope.count(instr->upstream_),
|
|
0,
|
|
::common::errors::NotFound("Can not find TmpTransformInstr upstream."));
|
|
PADDLE_ENFORCE_GT(
|
|
interpreter->scope.count(instr->downstream_),
|
|
0,
|
|
::common::errors::NotFound("Can not find TmpTransformInstr downstream."));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
interpreter->scope[instr->downstream_]->fusion_ops.size(),
|
|
1,
|
|
::common::errors::InvalidArgument(
|
|
"Downstream %s must have only one fusion_op.", instr->downstream_));
|
|
auto downstream_op =
|
|
interpreter->scope[instr->downstream_]->fusion_ops.front();
|
|
std::vector<FusibleOp> changed_upstreams;
|
|
for (auto fusion_op : interpreter->scope[instr->upstream_]->fusion_ops) {
|
|
auto upstream_op = std::get<ReduceOp>(fusion_op);
|
|
changed_upstreams = ConcatVector(
|
|
changed_upstreams,
|
|
cinn::hlir::framework::pir::trivial_fusion_detail::
|
|
TransformReduceLoopRange(
|
|
upstream_op, &downstream_op, instr->fake_reduce_iter_idx_));
|
|
}
|
|
|
|
// inplace set the upstream
|
|
interpreter->scope[instr->out_upstream_] = std::make_shared<ScopeElement>();
|
|
interpreter->scope[instr->out_upstream_]->fusion_ops = changed_upstreams;
|
|
// inplace set the downstream
|
|
interpreter->scope[instr->out_downstream_] = std::make_shared<ScopeElement>();
|
|
interpreter->scope[instr->out_downstream_]->fusion_ops.push_back(
|
|
downstream_op);
|
|
}
|
|
|
|
void RunTrivialLoopAlignInstr(
|
|
const std::shared_ptr<TrivialLoopAlignInstr>& instr,
|
|
FusionInterpreter* interpreter) {
|
|
PADDLE_ENFORCE_EQ(interpreter->scope[instr->downstream_]->fusion_ops.size(),
|
|
1,
|
|
::common::errors::InvalidArgument(
|
|
"Downstream op must have only one fusion_op."));
|
|
auto upstream_op = std::get<ReduceOp>(
|
|
interpreter->scope[instr->upstream_]->fusion_ops.front());
|
|
auto downstream_op = std::get<TrivialOp>(
|
|
interpreter->scope[instr->downstream_]->fusion_ops.front());
|
|
ScopeElementPtr new_pattern = std::make_shared<ScopeElement>();
|
|
new_pattern->fusion_ops.emplace_back(
|
|
cinn::hlir::framework::pir::trivial_fusion_detail::SinkTrivialLoopAlign(
|
|
downstream_op, upstream_op, instr->fake_reduce_iter_idx_));
|
|
interpreter->scope[instr->result_] = new_pattern;
|
|
}
|
|
|
|
void RunAxisTransformInstr(const std::shared_ptr<AxisTransformInstr>& instr,
|
|
FusionInterpreter* interpreter) {
|
|
auto substitute_dimexpr_for_shape =
|
|
[&](const std::vector<symbol::DimExpr>& shape) {
|
|
std::vector<symbol::DimExpr> result;
|
|
for (const auto& dim_expr : shape) {
|
|
symbol::DimExpr substituted = dim_expr;
|
|
while (true) {
|
|
if (substituted.isa<std::int64_t>()) break;
|
|
auto tmp_substituted = symbol::SubstituteDimExpr(
|
|
substituted, interpreter->substitute_dimexpr_map);
|
|
if (tmp_substituted == substituted) break;
|
|
substituted = tmp_substituted;
|
|
}
|
|
result.emplace_back(substituted);
|
|
}
|
|
return result;
|
|
};
|
|
auto substitute_dimexpr_for_transform = adt::match{
|
|
[&](const AppendAxisTransformPtr& trans) -> AxisTransform {
|
|
auto substituted_shape = substitute_dimexpr_for_shape(trans->shape);
|
|
return std::make_shared<AppendAxisTransform>(trans->axis,
|
|
substituted_shape);
|
|
},
|
|
[&](const ReshapeTransformPtr& trans) -> AxisTransform {
|
|
auto substituted_in_shape =
|
|
substitute_dimexpr_for_shape(trans->in_shape);
|
|
auto substituted_out_shape =
|
|
substitute_dimexpr_for_shape(trans->out_shape);
|
|
return std::make_shared<ReshapeTransform>(substituted_in_shape,
|
|
substituted_out_shape);
|
|
},
|
|
[&](const auto& trans) -> AxisTransform { return trans; }};
|
|
auto axis_transform = [&](ir::Expr op_expr) -> ir::Expr {
|
|
for (auto trans : instr->axis_transform_route_) {
|
|
auto new_trans = std::visit(substitute_dimexpr_for_transform, trans);
|
|
op_expr = std::visit(ApplyAxisTransform(op_expr), new_trans);
|
|
}
|
|
return op_expr;
|
|
};
|
|
|
|
auto new_pattern = std::make_shared<ScopeElement>();
|
|
auto fusion_ops = interpreter->scope[instr->source_]->fusion_ops;
|
|
VLOG(4) << "[AxisTransform] transform route: "
|
|
<< instr->axis_transform_route_;
|
|
for (const auto& fusion_op : fusion_ops) {
|
|
ir::Expr op_expr = std::visit(FusibleOp2Expr(), fusion_op).back();
|
|
VLOG(4) << "[AxisTransform] expr before transform: \n" << op_expr;
|
|
ir::Expr transformed_expr = axis_transform(op_expr);
|
|
if (cinn::hlir::framework::pir::trivial_fusion_detail::IsReduceBody(
|
|
transformed_expr)) {
|
|
new_pattern->fusion_ops.emplace_back(ReduceOp(transformed_expr));
|
|
} else {
|
|
new_pattern->fusion_ops.emplace_back(TrivialOp(transformed_expr));
|
|
}
|
|
}
|
|
interpreter->scope[instr->target_] = new_pattern;
|
|
}
|
|
|
|
void RunReturnInstr(const std::shared_ptr<ReturnInstr>& instr,
|
|
FusionInterpreter* interpreter) {
|
|
using namespace cinn::hlir::framework::pir::trivial_fusion_detail; // NOLINT
|
|
std::vector<ir::Expr> result;
|
|
// Insert if for append loop
|
|
for (auto fusion_op : interpreter->scope[instr->target_]->fusion_ops) {
|
|
auto exprs = std::visit(FusibleOp2Expr(), fusion_op);
|
|
for (auto expr : exprs) {
|
|
std::string output_var_name = GetOutputTensor(expr)->name;
|
|
if (!interpreter->global_var_names.count(output_var_name)) {
|
|
expr = ExprTransformerUtils::RemoveAllAppendIfTransformer()(expr);
|
|
}
|
|
expr = ExprTransformerUtils::EliminateUselessIfTransformer()(expr);
|
|
result.push_back(expr);
|
|
}
|
|
}
|
|
// Inline global vars
|
|
InlineGlobalVarCompute(result, interpreter->global_var_names);
|
|
interpreter->ret_expr = result;
|
|
}
|
|
|
|
std::vector<ir::Expr> FusionInterpreter::Run() {
|
|
/*
|
|
* instruction can't inplace change.
|
|
*/
|
|
for (auto instr : tracker->instructions_) {
|
|
VLOG(4) << "FusionInterpreter Start Run " << instr->DebugStr();
|
|
switch (instr->type()) {
|
|
case T_Copy:
|
|
RunCopyInstr(dynamic_cast_instr_with_err<CopyInstr>(instr), this);
|
|
break;
|
|
case T_Combine:
|
|
RunCombineInstr(dynamic_cast_instr_with_err<CombineInstr>(instr), this);
|
|
break;
|
|
case T_InitPattern:
|
|
RunInitPatternInstr(
|
|
dynamic_cast_instr_with_err<InitPatternInstr>(instr), this);
|
|
break;
|
|
case T_TrivialInline:
|
|
RunTrivialInlineInstr(
|
|
dynamic_cast_instr_with_err<TrivialInlineInstr>(instr), this);
|
|
break;
|
|
case T_TmpTransform:
|
|
RunTmpTransformInstr(
|
|
dynamic_cast_instr_with_err<TmpTransformInstr>(instr), this);
|
|
break;
|
|
case T_Return:
|
|
RunReturnInstr(dynamic_cast_instr_with_err<ReturnInstr>(instr), this);
|
|
break;
|
|
case T_TrivialLoopAlign:
|
|
RunTrivialLoopAlignInstr(
|
|
dynamic_cast_instr_with_err<TrivialLoopAlignInstr>(instr), this);
|
|
break;
|
|
case T_AxisTransform:
|
|
RunAxisTransformInstr(
|
|
dynamic_cast_instr_with_err<AxisTransformInstr>(instr), this);
|
|
break;
|
|
default:
|
|
PADDLE_THROW(
|
|
::common::errors::Unavailable("Unsupported Fusion Instruction"));
|
|
}
|
|
}
|
|
|
|
return TopoSort(ret_expr);
|
|
}
|
|
|
|
} // namespace cinn::fusion
|