chore: import upstream snapshot with attribution

python/paddle/incubate/cc/ap/__init__.py

@@ -0,0 +1,15 @@
+# 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 .facade_op import FacadeOp as FacadeOp

python/paddle/incubate/cc/ap/apy_to_axpr_json.py

@@ -0,0 +1,559 @@
+# 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

python/paddle/incubate/cc/ap/facade_op.py

@@ -0,0 +1,93 @@
+# 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.
+
+import warnings
+
+import paddle
+
+from .pir_attrs_serializer import PirAttrsSerializer
+
+
+class FacadeOp:
+    def __init__(self):
+        self.custom_op_name_ = self.custom_op_name()
+        self.infer_meta_ = self._check_to_str_pair(self.infer_meta())
+        self.infer_symbolic_ = self._check_to_str_pair(self.infer_symbolic())
+        self.num_inputs_ = self.num_inputs()
+        self.attrs_serializer_ = PirAttrsSerializer(self.attributes_schema)
+
+    def custom_op_name(self) -> str:
+        raise NotImplementedError(
+            "static method custom_op_name() is not overwritten"
+        )
+
+    def infer_meta(self) -> str:
+        raise NotImplementedError(
+            "static method infer_meta() is not overwritten"
+        )
+
+    def infer_symbolic(self) -> str:
+        raise NotImplementedError(
+            "static method infer_symbolic() is not overwritten"
+        )
+
+    def num_inputs(self) -> int:
+        raise NotImplementedError(
+            "static method num_inputs() is not overwritten"
+        )
+
+    def num_outputs(self, args) -> int:
+        raise NotImplementedError(
+            "static method num_outputs() is not overwritten"
+        )
+
+    def attributes_schema(self):
+        # annotations matter.
+        raise NotImplementedError(
+            "static method attributes_schema() is not overwritten"
+        )
+
+    def __call__(self, args, **kwargs):
+        if paddle.in_dynamic_mode():
+            warnings.warn("ap FacadeOp should not run in dynamic mode")
+        assert isinstance(args, (tuple, list))
+        self._check_num_inputs(len(args))
+        serialized_attrs = self.attrs_serializer_(**kwargs)
+        ret = paddle._C_ops.ap_facade(
+            args if len(args) > 0 else None,
+            self.num_outputs(args),
+            self.custom_op_name_,
+            self.infer_meta_,
+            self.infer_symbolic_,
+            serialized_attrs,
+        )
+        self._check_num_outputs(args, len(ret))
+        return ret
+
+    def _check_num_inputs(self, num_args):
+        if self.num_inputs_ >= 0:
+            assert self.num_inputs_ == num_args
+
+    def _check_num_outputs(self, args, num_rets):
+        num_outputs = self.num_outputs(args)
+        if num_outputs >= 0:
+            assert num_outputs == num_rets
+
+    def _check_to_str_pair(self, pair_str):
+        assert isinstance(pair_str, str)
+        pair = pair_str.split(".")
+        assert len(pair) == 2
+        assert pair[0] not in (None, "")
+        assert pair[1] not in (None, "")
+        return pair_str

python/paddle/incubate/cc/ap/pir_attrs_serializer.py

@@ -0,0 +1,234 @@
+# 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 inspect
+
+import paddle
+
+from ..data_type_util import get_dtype_lower_case_name
+from ..typing import DType
+from .apy_to_axpr_json import convert_python_stmts_to_axpr_json
+
+
+class PirAttrsSerializer:
+    def __init__(self, func):
+        self.attributes_schema = self._get_attributes_schema(func)
+        self._check_attributes_schema(self.attributes_schema)
+        self.attr_name2serializer = {
+            attr_name: serializer
+            for attr_name, schema_item in self.attributes_schema
+            for serializer in [self._get_serializer(attr_name, schema_item)]
+        }
+
+    def __call__(self, **attributes):
+        print(attributes)
+        attributes_names = {name for name, _ in attributes.items()}
+        attr_names = {name for name, _ in self.attributes_schema}
+        assert attributes_names == attr_names, (
+            f"expected attr_names: {attr_names}, but actual attr_names are {attributes_names}"
+        )
+        py_assigns = "\n".join(
+            py_stmt
+            for attr_name, attr_val in attributes.items()
+            for py_stmt in self.attr_name2serializer[attr_name](attr_val)
+        )
+        py_stmts_str = f"{py_assigns}\n{self._get_attr_map_ctor_str(self.attributes_schema)}"
+        return convert_python_stmts_to_axpr_json(py_stmts_str)
+
+    def _get_attr_map_ctor_str(self, attributes_schema):
+        kwargs = ", ".join(f"{name}={name}" for name, _ in attributes_schema)
+        return f"__builtin__AttrMap({kwargs})"
+
+    def _get_attributes_schema(self, obj):
+        if isinstance(obj, (list, tuple)):
+            return obj
+        func = obj
+        assert inspect.isfunction(func) or inspect.ismethod(func)
+        full_arg_spec = inspect.getfullargspec(func)
+        args = (
+            full_arg_spec.args[1:]
+            if inspect.ismethod(func)
+            else full_arg_spec.args
+        )
+        return [
+            (arg_name, annotation)
+            for arg_name in args
+            for annotation in [full_arg_spec.annotations[arg_name]]
+        ]
+
+    def _check_attributes_schema(self, attributes_schema):
+        for _, attr_type in attributes_schema:
+            self._check_attributes_schema_item_is_valid(attr_type)
+
+    def _check_attributes_schema_item_is_valid(self, attr_type):
+        if attr_type in self._supported_basic_types():
+            return
+        assert isinstance(attr_type, list), (
+            f"attribute type {attr_type} is not supported."
+        )
+        assert len(attr_type) == 1, (
+            "only syntax like [bool], [int], [float], [str] supported."
+        )
+        assert attr_type[0] in self._supported_basic_types(), (
+            f"supported list element types are bool/int/float/str, not include {attr_type[0]}."
+        )
+
+    def _supported_basic_types(self):
+        return (bool, int, float, str, DType)
+
+    def _get_serializer(self, attr_name, schema_item):
+        assert attr_name not in (
+            "custom_op_name",
+            "infer_meta_func_name",
+            "infer_symbolic_func_name",
+        )
+        schema_item_as_key = self._get_schema_item_as_key(schema_item)
+        return _get_serializer_factory[schema_item_as_key](attr_name)
+
+    def _get_schema_item_as_key(self, schema_item):
+        if schema_item in self._supported_basic_types():
+            return schema_item
+        assert isinstance(schema_item, list)
+        return tuple(schema_item)
+
+
+class PirAttributeSerializer:
+    def __init__(self, attr_name):
+        self.attr_name = attr_name
+
+    def __call__(self, value):
+        yield from []
+        raise NotImplementedError
+
+
+class BoolAttributeSerializer(PirAttributeSerializer):
+    def __init__(self, attr_name):
+        self.attr_name = attr_name
+
+    def __call__(self, value):
+        assert isinstance(value, bool)
+        yield f"{self.attr_name} = {value}"
+
+
+class IntAttributeSerializer(PirAttributeSerializer):
+    def __init__(self, attr_name):
+        self.attr_name = attr_name
+
+    def __call__(self, value):
+        assert isinstance(value, int)
+        yield f"{self.attr_name} = {value}"
+
+
+class FloatAttributeSerializer(PirAttributeSerializer):
+    def __init__(self, attr_name):
+        self.attr_name = attr_name
+
+    def __call__(self, value):
+        assert isinstance(value, float)
+        yield f"{self.attr_name} = {value}"
+
+
+class StrAttributeSerializer(PirAttributeSerializer):
+    def __init__(self, attr_name):
+        self.attr_name = attr_name
+
+    def __call__(self, value):
+        assert isinstance(value, str)
+        yield f"{self.attr_name} = {value}"
+
+
+class DTypeAttributeSerializer(PirAttributeSerializer):
+    def __init__(self, attr_name):
+        self.attr_name = attr_name
+
+    def __call__(self, value):
+        assert isinstance(value, paddle.dtype)
+        name = get_dtype_lower_case_name(value)
+        yield f"{self.attr_name} = __builtin__DataType.{name}"
+
+
+class BoolArrayAttributeSerializer(PirAttributeSerializer):
+    def __init__(self, attr_name):
+        self.attr_name = attr_name
+
+    def __call__(self, value):
+        assert isinstance(value, list)
+        for elt in value:
+            assert isinstance(elt, bool)
+        yield f"{self.attr_name} = {value}"
+
+
+class IntArrayAttributeSerializer(PirAttributeSerializer):
+    def __init__(self, attr_name):
+        self.attr_name = attr_name
+
+    def __call__(self, value):
+        assert isinstance(value, list)
+        for elt in value:
+            assert isinstance(elt, int)
+        yield f"{self.attr_name} = {value}"
+
+
+class FloatArrayAttributeSerializer(PirAttributeSerializer):
+    def __init__(self, attr_name):
+        self.attr_name = attr_name
+
+    def __call__(self, value):
+        assert isinstance(value, list)
+        for elt in value:
+            assert isinstance(elt, float)
+        yield f"{self.attr_name} = {value}"
+
+
+class StrArrayAttributeSerializer(PirAttributeSerializer):
+    def __init__(self, attr_name):
+        self.attr_name = attr_name
+
+    def __call__(self, value):
+        assert isinstance(value, list)
+        for elt in value:
+            assert isinstance(elt, str)
+        yield f"{self.attr_name} = {value}"
+
+
+class DTypeArrayAttributeSerializer(PirAttributeSerializer):
+    def __init__(self, attr_name):
+        self.attr_name = attr_name
+
+    def __call__(self, value):
+        assert isinstance(value, list)
+        for elt in value:
+            assert isinstance(elt, paddle.dtype)
+        value_str = ", ".join(
+            f"__builtin__DataType.{name}"
+            for dtype in value
+            for name in [get_dtype_lower_case_name(dtype)]
+        )
+        yield f"{self.attr_name} = [{value_str}]"
+
+
+_get_serializer_factory = {
+    bool: BoolAttributeSerializer,
+    int: IntAttributeSerializer,
+    float: FloatAttributeSerializer,
+    str: StrAttributeSerializer,
+    DType: DTypeAttributeSerializer,
+    (bool,): BoolArrayAttributeSerializer,
+    (int,): IntArrayAttributeSerializer,
+    (float,): FloatArrayAttributeSerializer,
+    (str,): StrArrayAttributeSerializer,
+    (DType,): DTypeArrayAttributeSerializer,
+}