wehub-resource-sync
135 lines
4.5 KiB
C++
135 lines
4.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/eliminate_invariant_loop.h"
|
|
|
|
#include "paddle/cinn/common/ir_util.h"
|
|
#include "paddle/cinn/ir/ir_analyzer/ir_analyzer.h"
|
|
#include "paddle/cinn/ir/ir_mutator.h"
|
|
#include "paddle/cinn/ir/utils/ir_replace.h"
|
|
#include "paddle/cinn/optim/replace_var_with_expr.h"
|
|
|
|
namespace cinn {
|
|
namespace optim {
|
|
|
|
namespace {
|
|
|
|
// Check whether var is used in the block's store indices or value.
|
|
bool HasVarInIndicesOrValue(const ir::Expr& block, const ir::Var& var) {
|
|
auto* block_realize = block.As<ir::ScheduleBlockRealize>();
|
|
auto* schedule_block = block_realize->schedule_block.As<ir::ScheduleBlock>();
|
|
|
|
// collect iter vars that refer to `var`
|
|
std::set<std::string> ref_iter_vars;
|
|
for (int i = 0; i < block_realize->iter_values.size(); ++i) {
|
|
auto var_use = ir::ir_utils::CollectIRNodes(
|
|
block_realize->iter_values[i],
|
|
[&](const ir::Expr* x) {
|
|
return x->is_var() && x->as_var_ref() == var;
|
|
},
|
|
/* uniq_target = */ true);
|
|
if (var_use.size() > 0) {
|
|
ref_iter_vars.insert(schedule_block->iter_vars[i]->name);
|
|
}
|
|
}
|
|
|
|
ir::Expr store = ir::analyzer::GetStoreOfSBlock(block);
|
|
auto var_use = ir::ir_utils::CollectIRNodes(
|
|
store,
|
|
[&](const ir::Expr* x) {
|
|
if (x->is_var()) {
|
|
if (x->as_var_ref() == var) return true;
|
|
if (ref_iter_vars.count(x->as_var()->name) > 0) return true;
|
|
}
|
|
return false;
|
|
},
|
|
/* uniq_target = */ true);
|
|
return var_use.size() > 0;
|
|
}
|
|
|
|
// Check whether the block is writing to a buffer whose scope is smaller than
|
|
// the For node's scope.
|
|
bool HasWriteToSmallerScope(const ir::Expr& block, const ir::For* for_node) {
|
|
ir::Expr store_tensor = ir::analyzer::GetStoreTensorOfSBlock(block);
|
|
ir::MemoryType memory_type = store_tensor.as_tensor()->buffer->memory_type;
|
|
if (for_node->is_gpu_thread_binded()) {
|
|
if (memory_type == ir::MemoryType::GPULocal) {
|
|
return true;
|
|
}
|
|
} else if (for_node->is_gpu_block_binded()) {
|
|
if (memory_type == ir::MemoryType::GPULocal ||
|
|
memory_type == ir::MemoryType::GPUShared) {
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
struct InvariantLoopEliminator : public ir::IRMutator<> {
|
|
using ir::IRMutator<>::Visit;
|
|
void operator()(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
|
|
|
|
void Visit(const ir::For* op, ir::Expr* expr) override {
|
|
auto* node = expr->As<ir::For>();
|
|
ir::IRMutator<>::Visit(&node->body, &node->body);
|
|
|
|
// do nothing if we are outside the root schedule block
|
|
if (!root_) return;
|
|
|
|
// check that all child schedule blocks satisfy the four rules
|
|
std::vector<ir::Expr> child_blocks = ir::analyzer::GetChildBlocks(*expr);
|
|
ir::Var loop_var = node->loop_var;
|
|
for (auto& block : child_blocks) {
|
|
if (HasVarInIndicesOrValue(block, loop_var)) return;
|
|
if (ir::analyzer::IsReductionSBlock(block)) return;
|
|
if (node->is_binded()) {
|
|
if (HasWriteToSmallerScope(block, node)) return;
|
|
if (!ir::analyzer::GetConsumerSBlocks(block, *root_).empty()) return;
|
|
}
|
|
}
|
|
|
|
// now we can eliminate this For node
|
|
ir::Expr body = node->body;
|
|
ir::ir_utils::IrReplaceVarBroadcast(&body, ir::Expr(loop_var), ir::Expr(0));
|
|
|
|
if (node->is_binded()) {
|
|
ir::Expr cond = ir::EQ::Make(loop_var, ir::Expr(0));
|
|
ir::Expr wrapped_body = ir::IfThenElse::Make(cond, body);
|
|
node->body = ir::Block::Make({wrapped_body});
|
|
} else {
|
|
*expr = body;
|
|
}
|
|
}
|
|
|
|
void Visit(const ir::ScheduleBlockRealize* op, ir::Expr* expr) override {
|
|
auto* schedule_block = op->schedule_block.As<ir::ScheduleBlock>();
|
|
// find and set the root schedule block
|
|
if (schedule_block->name.substr(0, 4) == "root") {
|
|
root_ = expr;
|
|
ir::IRMutator<>::Visit(op, expr);
|
|
root_ = nullptr;
|
|
}
|
|
}
|
|
|
|
private:
|
|
const ir::Expr* root_{nullptr};
|
|
};
|
|
|
|
} // namespace
|
|
|
|
void EliminateInvariantLoop(ir::Expr* expr) { InvariantLoopEliminator()(expr); }
|
|
|
|
} // namespace optim
|
|
} // namespace cinn
|