165 lines
9.8 KiB
Python
165 lines
9.8 KiB
Python
# Copyright 2022 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.
|
||
# ============================================================================
|
||
"""Ast optimizer for flatten recursive call."""
|
||
from typing import Any, Tuple
|
||
import ast
|
||
from ast import FunctionDef
|
||
from mindspore import log as logger
|
||
|
||
|
||
class FlattenRecursiveStmt(ast.NodeTransformer):
|
||
"""Ast optimizer for flatten recursive call."""
|
||
|
||
def __init__(self):
|
||
"""
|
||
Constructor of FlattenRecursiveStmt.
|
||
|
||
Returns:
|
||
An instance of ast optimizer for flatten recursive call.
|
||
"""
|
||
self._flatten_table: dict = {
|
||
ast.Return: ["value"],
|
||
ast.Call: ["args"],
|
||
ast.BinOp: ["left", "right"],
|
||
ast.BoolOp: ["values"],
|
||
ast.unaryop: ["operand"],
|
||
}
|
||
|
||
@staticmethod
|
||
def _generate_target_name(node: ast.AST, target_names):
|
||
"""Generate unique target name."""
|
||
if isinstance(node, ast.Call): #如果节点是ast.Call类型,它将检查节点的属性（表示函数调用的目标）的类型
|
||
func = node.func #如果是func类型，则使用其属性作为目标名称
|
||
if isinstance(func, ast.Name): #如果是ast.Name类型，则将使用其属性id作为目标名称
|
||
target_name = func.id
|
||
elif isinstance(func, ast.Attribute): #如果是ast.Attribute类型，则使用其属性attr作为目标名称
|
||
target_name = func.attr
|
||
else: #如果不是这两种类型，则会记录一个警告，并将目标名称设置为"function"
|
||
logger.warning("unhandled type of func of ast.Call while generating new target name: %s ", type(func))
|
||
target_name = "function"
|
||
elif isinstance(node, ast.Return): #如果节点是ast.Return类型，则将目标名称设置为"return_value"。
|
||
target_name = "return_value"
|
||
elif isinstance(node, (ast.BinOp, ast.boolop, ast.UnaryOp)): #如果节点是ast.BinOpast、ast.boolop或ast.UnaryOp类型，它将使用相应操作符的类型名称作为目标名称。
|
||
target_name = type(node.op).__name__
|
||
else: #对于其他类型的节点，记录一个警告，并将目标名称设置为该节点的类型名称
|
||
logger.warning("unhandled type of node while generating new target name: %s ", type(node))
|
||
target_name = type(node).__name__
|
||
suffix = 0
|
||
result = target_name
|
||
while result in target_names: #在生成的目标名称上添加数字后缀，以确保它在目标名称列表中是唯一的
|
||
suffix += 1
|
||
result = f"{target_name}_{suffix}"
|
||
target_names.append(result)
|
||
return result #返回生成的目标名称并将其添加到列表中
|
||
|
||
@staticmethod
|
||
def _fill_in_original_target_names(target_names, node):
|
||
"""Fill in original target names before getting unique names."""
|
||
for function_index in range(len(node.body)): #遍历AST（Abstract Syntax Tree）节点的node.body
|
||
child = node.body[function_index]
|
||
if not isinstance(child, ast.Assign): #在每个节点中，检查是否是一个赋值语句(ast.Assign)
|
||
continue #如果不是，继续处理下一个节点。
|
||
targets = child.targets #如果节点是赋值语句，则获取赋值语句的目标(targets)。
|
||
for target in targets: #如果目标是ast.Name类型，则获取变量名(target.id)并检查该名称是否已经存在于target_names列表中。
|
||
if not isinstance(target, ast.Name):
|
||
raise RuntimeError("currently only support ast.Name targets")
|
||
target_name = target.id
|
||
if target_name not in target_names: #如果变量名不在target_names列表中，就将其添加进去
|
||
target_names.append(target_name)
|
||
|
||
@staticmethod
|
||
# 检查node是否是以下类型之一：ast.Name, ast.Constant, ast.Num, ast.Str, ast.NameConstant, ast.Bytes, ast.Ellipsis
|
||
# 如果是，就直接返回一个空字符串和原始node，因为这些节点是不需要生成新赋值节点的。
|
||
def _create_new_assign_node(node: ast.AST, target_names) -> Tuple[str, ast.AST]:
|
||
"""Create new assign node to be inserted into ast.FunctionDef."""
|
||
if isinstance(node, (ast.Name, ast.Constant, ast.Num, ast.Str, ast.NameConstant, ast.Bytes, ast.Ellipsis)):
|
||
return "", node
|
||
#对于其他类型的节点，调用FlattenRecursiveStmt._generate_target_name函数来生成一个新的目标名称。
|
||
new_target_name = FlattenRecursiveStmt._generate_target_name(node, target_names)
|
||
return new_target_name, ast.Assign(targets=[ast.Name(id=new_target_name, ctx=ast.Store())], value=node) #创建一个新的赋值节点ast. Assign.node作为赋值的值(value)。 返回这个新的目标名称和创建的赋值节点
|
||
|
||
def _flatten_statement(self, node: ast.AST, target_names) -> [ast.AST]:
|
||
"""Flatten recursive statement according to different node type."""
|
||
flatten_config = self._flatten_table.get(type(node)) #通过查找_flatten_table字典获取node类型对应的展开配置flatten_config。如果没有找到对应的配置，它将返回一个空列表
|
||
if flatten_config is None:
|
||
return []
|
||
results = []
|
||
for todo_name in flatten_config:
|
||
todos = getattr(node, todo_name) #遍历展开配置中的每个待展开项(todo_name)
|
||
#对于列表类型的属性，遍历列表中的每个元素，调用FlattenRecursiveStmt._create_new_assign_node函数来生成一个新的赋值节点，并将生成的新目标名称替换原始元素的位置。
|
||
#如果生成的新节点与原始节点相同，则保留原始节点。 否则，将新节点添加到结果列表中。
|
||
if isinstance(todos, list):
|
||
new_list = []
|
||
for todo in todos:
|
||
new_target_name, new_node = FlattenRecursiveStmt._create_new_assign_node(todo, target_names)
|
||
if id(new_node) == id(todo):
|
||
new_list.append(todo)
|
||
else:
|
||
new_list.append(ast.Name(id=new_target_name, ctx=ast.Load()))
|
||
results.append(new_node)
|
||
setattr(node, todo_name, new_list)
|
||
#对于字典类型的属性，它会遍历字典中的每个键值对，调用FlattenRecursiveStmt._create_new_assign_node函数来生成新的赋值节点，并将生成的新目标名称替换原始值的位置。
|
||
#同样，如果生成的新节点与原始节点相同，则保留原始节点。 否则，将新节点添加到结果列表中。
|
||
elif isinstance(todos, dict):
|
||
new_dict = []
|
||
for key, value in todos:
|
||
new_target_name, new_node = FlattenRecursiveStmt._create_new_assign_node(value, target_names)
|
||
if id(new_node) == id(value):
|
||
new_dict[key] = value
|
||
else:
|
||
new_dict[key] = ast.Name(id=new_target_name, ctx=ast.Load())
|
||
results.append(new_node)
|
||
setattr(node, todo_name, new_dict)
|
||
else:
|
||
new_target_name, new_node = FlattenRecursiveStmt._create_new_assign_node(todos, target_names)
|
||
if id(new_node) != id(todos):
|
||
setattr(node, todo_name, ast.Name(id=new_target_name, ctx=ast.Load()))
|
||
results.append(new_node)
|
||
return results #函数返回结果列表，其中包含所有生成的新赋值节点
|
||
|
||
#检查node的名称是否为"construct",如果不是，直接返回原始node
|
||
def visit_FunctionDef(self, node: FunctionDef) -> Any:
|
||
"""Traverse construct node and flatten recursive nodes."""
|
||
if node.name != "construct":
|
||
return node
|
||
|
||
target_names = [] #创建一个空的target_names列表，并调用_fill_in_original_target_names函数来填充原始目标名称
|
||
self._fill_in_original_target_names(target_names, node)
|
||
index = len(node.body) - 1 #从函数体的最后一个语句开始，向前遍历每个语句（由node.body列表表示）。
|
||
while index >= 0:
|
||
child = node.body[index] #对于每个语句，函数检查是否是一个赋值语句（ast.Assign），或者是一个表达式语句（ast.Expr）。
|
||
if isinstance(child, ast.Assign):
|
||
stmt = child.value #如果是赋值语句,获取赋值语句的值
|
||
elif isinstance(child, ast.Expr):
|
||
stmt = child.value #如果是表达式语句，获取表达式的值
|
||
else:
|
||
stmt = child #否则，将语句本身作为要处理的节点
|
||
results = self._flatten_statement(stmt, target_names) #调用_flatten_statement函数来展开语句中的递归节点，并传递target_names列表。 该函数将返回一个包含生成的新赋值节点的列表。
|
||
#如果_flatten_statement返回了结果，函数将这些结果逆序遍历，并将每个新节点插入到原始节点的前面，以确保展开的赋值节点按正确的顺序插入。
|
||
if results:
|
||
results.reverse()
|
||
for result in results:
|
||
node.body.insert(index, result)
|
||
index += 1
|
||
index -= 1
|
||
return node
|
||
|
||
def transform(self, ast_root): #这是FlattenRecursiveStmt类的transform方法，它是FlattenRecursiveStmt的接口方法。
|
||
#这个方法接受一个AST（Abstract Syntax Tree）的根节点ast_root作为输入，并返回经过展开递归节点处理后的AST根节点。
|
||
"""Interface of FlattenRecursiveStmt."""
|
||
ast_root = self.visit(ast_root) #调用self.visit(ast_root)，这是一个递归的AST遍历过程，它会调用visit_FunctionDef等方法来处理AST中的各种节点。
|
||
ast_root = ast.fix_missing_locations(ast_root) #使用ast.fix_missing_locations(ast_root)来修复AST节点中缺失的位置信息。 在处理AST时，可能会对节点进行插入、删除等操作，导致节点的位置信息丢失。ast.fix_missing_locations会遍历AST并为缺失位置的节点添加默认的位置信息。
|
||
return ast_root #返回经过展开递归节点处理和位置修复后的AST根节点
|