mindspore2022/mindspore/ccsrc/pipeline/jit/parse/function_block.h

125 lines
5.0 KiB
C++

/**
* This is the C++ adaptation and derivative work of Myia (https://github.com/mila-iqia/myia/).
*
* Copyright 2019-2021 Huawei Technologies Co., Ltd
*
* 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.
*/
#ifndef MINDSPORE_CCSRC_PIPELINE_JIT_PARSE_FUNCTION_BLOCK_H_
#define MINDSPORE_CCSRC_PIPELINE_JIT_PARSE_FUNCTION_BLOCK_H_
#include <vector>
#include <string>
#include <map>
#include <set>
#include <unordered_map>
#include <memory>
#include <utility>
#include "pipeline/jit/parse/parse_base.h"
#include "utils/log_adapter.h"
#include "utils/ordered_set.h"
namespace mindspore {
namespace parse {
class Parser;
class NameSpace;
class Symbol;
class FunctionBlock;
using FunctionBlockPtr = std::shared_ptr<FunctionBlock>;
// A function block is a straight-line code sequence with no branches, every block has one one exit point
// which is return. When parsing function, loop or branch , we use function block to track the structure of
// the original source code.
class FunctionBlock : public std::enable_shared_from_this<FunctionBlock> {
public:
explicit FunctionBlock(const Parser &parser);
virtual ~FunctionBlock() {}
FuncGraphPtr func_graph() { return func_graph_; }
void WriteVariable(const std::string &var_name, const AnfNodePtr &node);
AnfNodePtr ReadVariable(const std::string &var_name);
void AddPrevBlock(const FunctionBlockPtr &block);
void SetPhiArgument(const ParameterPtr &phi);
bool CollectRemovablePhi(const ParameterPtr &phi);
// A block is matured if all its predecessors is generated
void Mature();
CNodePtr ForceToBoolNode(const AnfNodePtr &cond);
CNodePtr ForceToWhileCond(const AnfNodePtr &cond);
void Jump(const FunctionBlockPtr &block, const AnfNodePtr &node);
AnfNodePtr SearchReplaceNode(const std::string &var, const ParameterPtr &phi);
void ConditionalJump(AnfNodePtr condNode, const FunctionBlockPtr &trueBlock, const FunctionBlockPtr &falseBlock,
bool unroll_loop = true);
// Create cnode for the assign statement like self.target = source.
void SetStateAssign(const AnfNodePtr &target, const AnfNodePtr &source);
void AddGlobalVar(const std::string &var_name) { (void)global_vars_.insert(var_name); }
bool IsGlobalVar(const std::string &var_name) { return global_vars_.find(var_name) != global_vars_.end(); }
AnfNodePtr MakeResolveAstOp(const py::object &op);
AnfNodePtr MakeResolveClassMember(const std::string &attr);
AnfNodePtr MakeResolveSymbol(const std::string &value);
AnfNodePtr MakeResolveOperation(const std::string &value);
AnfNodePtr MakeResolve(const std::shared_ptr<NameSpace> &name_space, const std::shared_ptr<Symbol> &resolve_symbol);
const std::unordered_map<ParameterPtr, AnfNodePtr> &removable_phis() const { return removable_phis_; }
void FindIsolatedNodes();
void AddIsolatedNode(const AnfNodePtr &target);
void AttachIsolatedNodesBeforeReturn();
private:
// Block graph
FuncGraphPtr func_graph_;
// Block parser
const Parser &parser_;
// A block is matured if all its prev_blocks is processed
bool matured_;
// Store the nest-level block.
// Refer to comments in Parser::func_block_list_;
std::vector<FunctionBlock *> prev_blocks_;
// Store args and variable's node, use a bool flag to indicate if the variable is used.
std::map<std::string, std::pair<AnfNodePtr, bool>> vars_;
// Map the parameter node to variable, it can be resolved if the block's predecessors are processed
std::map<ParameterPtr, std::string> phi_nodes_;
// Jumps map the successor block and the function call that perform jump
// Refer to comments in Parser::func_block_list_ that how to break the cyclic reference
std::map<FunctionBlock *, CNodePtr> jumps_;
// Keep all removable phis which will be removed in one pass.
std::unordered_map<ParameterPtr, AnfNodePtr> removable_phis_;
// Keep the map for the resolve node to the removable phi node.
// For the case that ReadVariable returns a phi node although this phi node
// generated in the prev block is identified as removable. The other blocks
// should find this phi node.
std::unordered_map<AnfNodePtr, ParameterPtr> resolve_to_removable_phis_;
// Hold declared global variables in function
std::set<std::string> global_vars_;
// Keep new made resolve symbol for the variable not found in vars_.
std::unordered_map<std::string, AnfNodePtr> var_to_resolve_;
// Isolated nodes.
OrderedSet<AnfNodePtr> isolated_nodes_;
};
} // namespace parse
} // namespace mindspore
#endif // MINDSPORE_CCSRC_PIPELINE_JIT_PARSE_FUNCTION_BLOCK_H_