Files
paddlepaddle--paddle/python/paddle/incubate/cc/ap/apy_to_axpr_json.py
T
2026-07-13 12:40:42 +08:00

560 lines
19 KiB
Python

# Copyright (c) 2025 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.
from __future__ import annotations
import ast
import functools
import itertools
import json
import operator
import typing as t
from dataclasses import dataclass
def convert_python_stmts_to_axpr_json(python_code_stmts_str):
tree = ast.parse(python_code_stmts_str)
parser = PyToAnfParser()
return parser(tree).ConvertToAnfExpr().JsonDump()
@dataclass
class AnfExpr:
def DumpToFileAsJson(self, file_name):
with open(file_name, "w") as f:
json.dump(self.value, f, indent=2)
def JsonDump(self):
return json.dumps(self.value)
@dataclass
class AtomicAnfExpr(AnfExpr):
value: t.Any
@dataclass
class CombinedAnfExpr(AnfExpr):
value: t.Any
@dataclass
class AnfParseResult:
bindings: list[str]
body_atomic_anf_expr: AtomicAnfExpr
def __add__(self, other):
return AnfParseResult(
bindings=[*self.bindings, *other.bindings],
body_atomic_anf_expr=other.body_atomic_anf_expr,
)
def ConvertToAnfExpr(self):
ret = self.body_atomic_anf_expr
if len(self.bindings) == 0:
return ret
assert isinstance(ret, AtomicAnfExpr)
ret = CombinedAnfExpr(
["__builtin_identity__", self.body_atomic_anf_expr.value]
)
return CombinedAnfExpr(["__builtin_let__", self.bindings, ret.value])
class PyToAnfParser:
def __init__(self, seq_no_counter=None, return_count_constraint=None):
self.bindings = []
self.seq_no_counter = (
seq_no_counter if seq_no_counter is not None else itertools.count()
)
self.return_count_constraint = (
return_count_constraint
if return_count_constraint is not None
else ReturnCounterConstraint(limits=1)
)
def __call__(self, tree):
ret = self.Parse(tree)
return AnfParseResult(bindings=self.bindings, body_atomic_anf_expr=ret)
def Parse(self, tree):
method_name = f"Parse{type(tree).__name__}"
return getattr(self, method_name)(tree)
def ParseImport(self, tree):
for alias in tree.names:
assert isinstance(alias, ast.alias)
name = alias.name
asname = alias.asname if alias.asname is not None else name
self.Bind(asname, ["import", {"str": name}])
return AtomicAnfExpr(None)
def ParseClassDef(self, tree: ast.ClassDef):
assert len(tree.keywords) == 0
class_name = tree.name
def GetBases():
bases = [self.Parse(base) for base in tree.bases]
return self.BindToTmpVar(
['__builtin_list__', *[x.value for x in bases]]
)
def GetFunctions():
body_name_and_method_pair = []
for func_def in tree.body:
if isinstance(func_def, ast.Pass):
continue
assert isinstance(func_def, ast.FunctionDef), (
f"only method supported in class definition, {type(func_def)} were given."
)
func_code = self.BindToTmpVar(
[
'__builtin_getattr__',
self.Parse(func_def).value,
{"str": '__function__'},
]
)
pair = self.BindToTmpVar(
[
"__builtin_list__",
{"str": func_def.name},
func_code.value,
]
)
body_name_and_method_pair.append(pair)
positional_args = self.BindToTmpVar(['__builtin_list__'])
keyword_args = self.BindToTmpVar(
[
'__builtin_list__',
*[x.value for x in body_name_and_method_pair],
]
)
packed_args = self.BindToTmpVar(
[
'__builtin_PackedArgs__',
positional_args.value,
keyword_args.value,
]
)
return self.BindToTmpVar(
['BuiltinSerializableAttrMap', packed_args.value]
)
class_anf_expr = self.BindToTmpVar(
[
'type',
{"str": class_name},
GetBases().value,
GetFunctions().value,
]
)
for elt in reversed(tree.decorator_list):
decorator = self.Parse(elt)
class_anf_expr = self.BindToTmpVar(
[decorator.value, class_anf_expr.value]
)
self.Bind(class_name, class_anf_expr)
return class_anf_expr
def Parsekeyword(self, tree):
value = self.Parse(tree.value)
return self.BindToTmpVar(
["__builtin_list__", {"str": tree.arg}, value.value]
)
def ParseBinOp(self, tree):
left = self.Parse(tree.left)
op = self.Parse(tree.op)
right = self.Parse(tree.right)
return self.BindToTmpVar([op.value, left.value, right.value])
def ParseUnaryOp(self, tree):
op = self.Parse(tree.op)
operand = self.Parse(tree.operand)
return self.BindToTmpVar([op.value, operand.value])
def ParseCompare(self, tree):
assert len(tree.ops) == 1
op = self.Parse(tree.ops[0])
left = self.Parse(tree.left)
assert len(tree.comparators) == 1
right = self.Parse(tree.comparators[0])
return self.BindToTmpVar([op.value, left.value, right.value])
def ParseAdd(self, tree):
return AtomicAnfExpr("__builtin_Add__")
def ParseSub(self, tree):
return AtomicAnfExpr("__builtin_Sub__")
def ParseMult(self, tree):
return AtomicAnfExpr("__builtin_Mul__")
def ParseDiv(self, tree):
return AtomicAnfExpr("__builtin_Div__")
def ParseFloorDiv(self, tree):
return AtomicAnfExpr("__builtin_FloorDiv__")
def ParseMod(self, tree):
return AtomicAnfExpr("__builtin_Mod__")
def ParseUSub(self, tree):
return AtomicAnfExpr("__builtin_Neg__")
def ParseEq(self, tree):
return AtomicAnfExpr("__builtin_EQ__")
def ParseNotEq(self, tree):
return AtomicAnfExpr("__builtin_NE__")
def ParseGt(self, tree):
return AtomicAnfExpr("__builtin_GT__")
def ParseGtE(self, tree):
return AtomicAnfExpr("__builtin_GE__")
def ParseLt(self, tree):
return AtomicAnfExpr("__builtin_LT__")
def ParseLtE(self, tree):
return AtomicAnfExpr("__builtin_LE__")
def ParseModule(self, module: ast.Module):
parse_result = AnfParseResult(
bindings=[], body_atomic_anf_expr=AtomicAnfExpr(None)
)
if len(module.body) > 0:
seq_no_counter = itertools.count()
return_count_constraint = ReturnCounterConstraint(limits=0)
parse_result = functools.reduce(
operator.add,
(
PyToAnfParser(seq_no_counter, return_count_constraint)(tree)
for tree in module.body
),
)
return parse_result.ConvertToAnfExpr()
def ParseFunctionDef(self, function_def: ast.FunctionDef):
if len(function_def.body) > 0:
return_count_constraint = ReturnCounterConstraint(limits=1)
return_stmt_idx = self.GetStmtSizeUntilReturn(function_def.body)
parse_result = functools.reduce(
operator.add,
[
PyToAnfParser(self.seq_no_counter, return_count_constraint)(
tree
)
for tree in function_def.body[0:return_stmt_idx]
if not isinstance(tree, ast.Pass)
]
+ [
AnfParseResult(
bindings=[], body_atomic_anf_expr=AtomicAnfExpr(None)
)
],
)
else:
parse_result = AnfParseResult(
bindings=[], body_atomic_anf_expr=AtomicAnfExpr(None)
)
args = [arg.arg for arg in function_def.args.args]
lmbd = AtomicAnfExpr(
['lambda', args, parse_result.ConvertToAnfExpr().value]
)
for elt in reversed(function_def.decorator_list):
decorator = self.Parse(elt)
lmbd = self.BindToTmpVar([decorator.value, lmbd.value])
func_name = function_def.name
self.Bind(func_name, lmbd)
return AtomicAnfExpr(func_name)
def ParseLambda(self, function_def: ast.Lambda):
return_count_constraint = ReturnCounterConstraint(limits=0)
parser = PyToAnfParser(self.seq_no_counter, return_count_constraint)
parse_result = parser(function_def.body)
args = [arg.arg for arg in function_def.args.args]
return AtomicAnfExpr(
['lambda', args, parse_result.ConvertToAnfExpr().value]
)
def ParseIfExp(self, if_expr: ast.IfExp):
test_value = self.Parse(if_expr.test)
true_value = self.ParseExprTo0ArgLambda(if_expr.body)
false_value = self.ParseExprTo0ArgLambda(if_expr.orelse)
ret = self.BindToTmpVar(
[
'__builtin_if__',
test_value.value,
true_value.value,
false_value.value,
]
)
return ret
def ParseBoolOp(self, bool_op: ast.BoolOp):
name = type(bool_op.op).__name__
method = f"Parse{name}"
return getattr(self, method)(bool_op)
def ParseOr(self, bool_op: ast.BoolOp):
assert len(bool_op.values) == 2
test_value = self.Parse(bool_op.values[0])
true_value = AtomicAnfExpr(['lambda', [], AtomicAnfExpr(True).value])
false_value = self.ParseExprTo0ArgLambda(bool_op.values[1])
ret = self.BindToTmpVar(
[
'__builtin_if__',
test_value.value,
true_value.value,
false_value.value,
]
)
return ret
def ParseAnd(self, bool_op: ast.BoolOp):
assert len(bool_op.values) == 2
test_value = self.Parse(bool_op.values[0])
true_value = self.ParseExprTo0ArgLambda(bool_op.values[1])
false_value = AtomicAnfExpr(['lambda', [], AtomicAnfExpr(False).value])
ret = self.BindToTmpVar(
[
'__builtin_if__',
test_value.value,
true_value.value,
false_value.value,
]
)
return ret
def ParseNot(self, unary_op: ast.UnaryOp):
return AtomicAnfExpr('__builtin_not__')
def ParseExprTo0ArgLambda(self, expr):
return_count_constraint = ReturnCounterConstraint(limits=0)
parser = PyToAnfParser(self.seq_no_counter, return_count_constraint)
parse_result = parser(expr)
return AtomicAnfExpr(
['lambda', [], parse_result.ConvertToAnfExpr().value]
)
def ParseAssert(self, expr: ast.Assert):
test_value = self.Parse(expr.test)
true_value = AtomicAnfExpr(['lambda', [], AtomicAnfExpr(None).value])
# handle lambda: rase(msg)
return_count_constraint = ReturnCounterConstraint(limits=0)
parser = PyToAnfParser(self.seq_no_counter, return_count_constraint)
if expr.msg is None:
msg = parser.BindToTmpVar(AtomicAnfExpr({"str": ""}))
else:
msg = parser.Parse(expr.msg)
exception = parser.BindToTmpVar(['AssertionError', msg.value])
raise_ret = parser.BindToTmpVar(['raise', exception.value])
false_value = AtomicAnfExpr(
[
'lambda',
[],
AnfParseResult(
bindings=parser.bindings, body_atomic_anf_expr=raise_ret
)
.ConvertToAnfExpr()
.value,
]
)
ret = self.BindToTmpVar(
[
'__builtin_if__',
test_value.value,
true_value.value,
false_value.value,
]
)
return ret
def ParseAssign(self, tree):
assert len(tree.targets) == 1
if isinstance(tree.targets[0], ast.Name):
val = self.Parse(tree.value)
var = tree.targets[0].id
self.Bind(var, val)
return AtomicAnfExpr(var)
elif isinstance(tree.targets[0], ast.Attribute):
val = self.Parse(tree.value)
attr = tree.targets[0]
f = self.BindToTmpVar(
[
'__builtin_setattr__',
self.Parse(attr.value).value,
{"str": attr.attr},
]
)
return self.BindToTmpVar([f.value, {"str": attr.attr}, val.value])
elif isinstance(tree.targets[0], ast.Subscript):
val = self.Parse(tree.value)
subscript = tree.targets[0]
slice_val = self.Parse(subscript.slice).value
f = self.BindToTmpVar(
[
'__builtin_setitem__',
self.Parse(subscript.value).value,
slice_val,
]
)
return self.BindToTmpVar([f.value, slice_val, val.value])
else:
raise NotImplementedError(tree.targets)
def ParseSubscript(self, tree):
val = self.Parse(tree.value)
slc = self.Parse(tree.slice)
return self.BindToTmpVar(["__builtin_getitem__", val.value, slc.value])
def ParseExpr(self, tree):
return self.BindToTmpVar(self.Parse(tree.value))
def BindToTmpVar(self, value):
tmp_var = self.get_tmp_var()
self.Bind(tmp_var, value)
return AtomicAnfExpr(tmp_var)
def GetStmtSizeUntilReturn(self, stmts):
for idx, stmt in enumerate(stmts):
if isinstance(stmt, ast.Return):
return idx + 1
return len(stmts)
def ParseReturn(self, tree: ast.Return):
self.return_count_constraint.CountAndCheck()
value = self.Parse(tree.value)
return self.BindToTmpVar(["__builtin_return__", value.value])
def ParseStarred(self, tree: ast.Starred):
value = self.Parse(tree.value)
return self.BindToTmpVar(["__builtin_starred__", value.value])
def ParseCall(self, tree: ast.Call):
func = self.Parse(tree.func)
assert isinstance(func, AtomicAnfExpr)
def ParseArg(arg):
parsed_arg = self.Parse(arg)
assert isinstance(parsed_arg, AtomicAnfExpr)
return parsed_arg
args = [ParseArg(arg).value for arg in tree.args]
kwargs = None
if len(tree.keywords) > 0:
keywords = [ParseArg(arg).value for arg in tree.keywords]
kwargs = self.BindToTmpVar(["__builtin_list__", *keywords])
if kwargs is None:
if any(isinstance(arg, ast.Starred) for arg in tree.args):
l = self.BindToTmpVar(["__builtin_list__", *args])
return self.BindToTmpVar(
["__builtin_apply__", func.value, l.value]
)
else:
return self.BindToTmpVar([func.value, *args])
else:
args = self.BindToTmpVar(["__builtin_list__", *args])
packed_args = self.BindToTmpVar(
["__builtin_PackedArgs__", args.value, kwargs.value]
)
return self.BindToTmpVar([func.value, packed_args.value])
def ParseList(self, lst: ast.List):
return self._ParseCall('__builtin_list__', lst.elts)
def _ParseCall(self, func, ast_args):
def ParseArg(arg):
parsed_arg = self.Parse(arg)
assert isinstance(parsed_arg, AtomicAnfExpr)
return parsed_arg
args = [ParseArg(arg).value for arg in ast_args]
ret_var = self.get_tmp_var()
self.Bind(ret_var, [func, *args])
return AtomicAnfExpr(ret_var)
def ParseAttribute(self, attr: ast.Attribute):
ret_var = self.get_tmp_var()
self.Bind(
ret_var,
[
'__builtin_getattr__',
self.Parse(attr.value).value,
{"str": attr.attr},
],
)
return AtomicAnfExpr(ret_var)
def ParseName(self, name: ast.Name):
return AtomicAnfExpr(name.id)
def ParseConstant(self, constant: ast.Constant):
if isinstance(constant.value, str):
return AtomicAnfExpr({"str": constant.value})
if isinstance(constant.value, (bool, int, float)):
return AtomicAnfExpr(constant.value)
if constant.value is None:
return AtomicAnfExpr(None)
raise NotImplementedError(f"{constant} not supported by anf_expr")
def ParseJoinedStr(self, tree: ast.JoinedStr):
if len(tree.values) == 0:
return AtomicAnfExpr({"str": ""})
def ToString(elt):
parsed_elt = self.Parse(elt)
parsed_elt = self.BindToTmpVar(
['__builtin_ToString__', parsed_elt.value]
)
return parsed_elt
ret = ToString(tree.values[0])
for elt in tree.values[1:]:
parsed_elt = ToString(elt)
ret = self.BindToTmpVar(
['__builtin_Add__', ret.value, parsed_elt.value]
)
return ret
def ParseFormattedValue(self, tree: ast.FormattedValue):
return self.Parse(tree.value)
def Bind(self, var_name, anf_expr):
return getattr(self, f"Bind{type(anf_expr).__name__}")(
var_name, anf_expr
)
def BindAtomicAnfExpr(self, var_name, anf_expr):
self.bindings.append(
[var_name, ["__builtin_identity__", anf_expr.value]]
)
def Bindlist(self, var_name, anf_expr):
self.bindings.append([var_name, anf_expr])
def get_tmp_var(self):
return f"___{next(self.seq_no_counter)}"
class ReturnCounterConstraint:
def __init__(self, limits):
self.counter = itertools.count()
self.limits = limits
def CountAndCheck(self):
return_stmt_id = next(self.counter)
assert return_stmt_id < self.limits