MindSpore代码评析(六)pipeline模块的整体分析
pipeline模块负责图编译与执行,具体包括ME(Mindexpression)图编译功能、ME图到GE(Graphengine)图转换功能,以及GE 图执行功能等。其主要实现过程将之前抽象化的树转变为ANF图,对ANF图进行分析,将其转变为GE图,最后执行GE图,来实现MindSpore的图形处理。
其中,action文件负责对MindIR的转换,对ANF图的生成和优化。Pipeline_ge文件负责创建、执行和输出DF图,初始化GE图的执行数据集,生成GE图,并对其进行执行,其他模块则实现资源调配、DEBUG功能、基础运算等功能。
功能实现的主要依赖方法如下:
1.ParseAction方法作用为根据Python AST抽象语法树,生成初步的ANF图。
```cpp
bool ParseAction(const ResourcePtr &res) {
if (!res->input()) {
MS_LOG(EXCEPTION) << "Parse error";
}
py::object input = res->input();
parse::Parser::InitParserEnvironment(input);
py::module path = py::module::import("os.path");
std::string dir = path.attr("dirname")(py::globals()["__file__"]).cast<std::string>();
parse::python_adapter::set_python_env_flag(true);
parse::python_adapter::SetPythonPath(dir);
ValuePtr converted_ret = nullptr;
bool converted = parse::ConvertData(input, &converted_ret, true);
if (!converted) {
MS_LOG(EXCEPTION) << "Attribute convert error with type:" << std::string(py::str(input));
}
FuncGraphPtr top_graph = nullptr;
if (py::isinstance<Cell>(input)) {
top_graph = parse::MakeTopGraph(input, converted_ret);
} else if (converted_ret->isa<FuncGraph>()) {
top_graph = converted_ret->cast<FuncGraphPtr>();
} else {
MS_LOG(EXCEPTION) << "Object to parse " << std::string(py::str(input)) << " is not function or cell.";
}
parse::Parser::UpdateTopFuncGraph(top_graph);
res->set_func_graph(top_graph);
FuncGraphManagerPtr manager = res->manager();
if (manager == nullptr) {
MS_LOG(EXCEPTION) << "Manager is nullptr.";
}
manager->AddFuncGraph(top_graph);
return true;
}
```
2.Python AST是解析Python文件后生成的抽象语法树,ANF是MindSpore中用于表示IR (Intermediate representation)的一种形式,包含ANode和CNode。
3.SymbolResolveAction作用为解析ANF图,移除ANF图中多余的CNode节点。
```cpp
bool SymbolResolveAction(const ResourcePtr &res) {
if (res->manager() == nullptr) {
MS_LOG(EXCEPTION) << "SymbolResolve error, manager is null";
}
auto func_graph = res->func_graph();
if (func_graph == nullptr) {
MS_LOG(EXCEPTION) << "SymbolResolve error, graph is null";
}
bool ret = parse::ResolveFuncGraph(func_graph, res);
// Remove unused nodes in cnode order list.
if (func_graph) {
func_graph->EraseUnusedNodeInOrder();
for (auto fg : func_graph->func_graphs_used_total()) {
if (fg) {
fg->EraseUnusedNodeInOrder();
}
}
}
return ret;
}
```
4.InferenceOptPrepareAction主要调用GradVarPrepare类,实现优化推理操作,提供梯度参数变长参数支持。
```cpp
bool InferenceOptPrepareAction(const ResourcePtr &res) {
if (res->manager() == nullptr) {
MS_LOG(EXCEPTION) << "InferenceOptPrepare error, manager is null.";
}
if (res->func_graph() == nullptr) {
MS_LOG(EXCEPTION) << "InferenceOptPrepare error, graph is null.";
}
return InferenceOptPreparePass(res);
}
```
5.AbstractSpecializeAction作用为分析ANF图中抽象值的类型和维度,然后执行特化操作。
```cpp
bool AbstractSpecializeAction(const ResourcePtr &res) {
if (res->func_graph() == nullptr) {
MS_LOG(EXCEPTION) << "AbstractSpecialize error";
}
FuncGraphPtr func_graph = res->func_graph();
abstract::AbstractBasePtrList args_spec = res->args_spec();
auto context = parallel::ParallelContext::GetInstance();
MS_EXCEPTION_IF_NULL(parallel::ParallelContext::GetInstance());
context->ParallelParameterContextInitShape(func_graph);
// get original loaded graph to check inputs later
auto loaded_graph_ptr = GetLoadedGraph(res);
// suppose that there is not KeywordArgument for the top graph
// get the hyper parameter
for (const auto ¶m : func_graph->parameters()) {
auto param_node = std::static_pointer_cast<Parameter>(param);
if (param_node->has_default()) {
auto value = param_node->default_param();
auto abs_value = value->ToAbstract()->cast<abstract::AbstractTensorPtr>();
auto ref_key = std::make_shared<RefKey>(param_node->name());
auto abs_ref_key = ref_key->ToAbstract();
auto abs_ref = std::make_shared<abstract::AbstractRef>(abs_ref_key, abs_value);
context->ParallelParameterContextRestoreShape(func_graph, param_node, abs_ref);
args_spec.push_back(abs_ref);
context->ParallelParameterContextCkptShape(func_graph, param_node, abs_ref);
}
}
// Analyze
AnalysisResult result = AbstractAnalyze(res, func_graph, args_spec);
// The top graph may be replaced by infer, update the top graph when the infer is done
parse::Parser::UpdateTopFuncGraph(result.context->func_graph());
// Specialize
FuncGraphPtr new_fg = ProgramSpecialize(res, result.context->func_graph(), result.context);
res->set_func_graph(new_fg);
// Remove unused nodes in cnode order list, this is prepared for auto-monad.
if (new_fg) {
new_fg->EraseUnusedNodeInOrder();
for (auto fg : new_fg->func_graphs_used_total()) {
if