chore: import upstream snapshot with attribution

python/paddle/static/python_op.py

@@ -0,0 +1,390 @@
+# 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 sys
+import types
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass
+from functools import cached_property, partial, wraps
+from typing import Any, Generic, TypeVar, overload
+
+from typing_extensions import ParamSpec
+
+import paddle
+from paddle import _C_ops
+
+HAS_VAR_ARGS_OR_KWARGS: int = inspect.CO_VARARGS | inspect.CO_VARKEYWORDS
+
+
+P1 = ParamSpec("P1")
+R1 = TypeVar("R1")
+
+
+class MissingArgument:
+    def __init__(self, fn: Callable[P1, R1], name: str):
+        self.fn = fn
+        self.name = name
+
+    def __repr__(self):
+        return f"<Required parameter '{self.name}' for function {self.fn.__name__}>"
+
+
+def extract_default(fn: Callable[P1, R1], parameter: inspect.Parameter):
+    if parameter.kind is inspect.Parameter.VAR_POSITIONAL:
+        return ()
+    elif parameter.kind is inspect.Parameter.VAR_KEYWORD:
+        return {}
+    elif parameter.default is inspect.Parameter.empty:
+        return MissingArgument(fn, parameter.name)
+    return parameter.default
+
+
+def get_fn_defaults_params(fn: Callable[P1, R1]) -> tuple:
+    fn_defaults_params = [
+        extract_default(fn, param)
+        for param in inspect.signature(fn).parameters.values()
+    ]
+    for i, default in enumerate(fn_defaults_params):
+        if not isinstance(default, MissingArgument):
+            fn_defaults_params = fn_defaults_params[i:]
+            break
+    return tuple(fn_defaults_params)
+
+
+def eliminate_positional_or_keyword_only(
+    fn: Callable[P1, R1],
+) -> Callable[P1, R1]:
+    assert isinstance(fn, types.FunctionType), "Only support regular function"
+    code = fn.__code__
+    co_flags: int = code.co_flags & ~HAS_VAR_ARGS_OR_KWARGS
+
+    argcount = (
+        code.co_argcount
+        + code.co_kwonlyargcount
+        + bool(code.co_flags & inspect.CO_VARARGS)
+        + bool(code.co_flags & inspect.CO_VARKEYWORDS)
+    )
+    if sys.version_info >= (3, 11):
+        new_code = types.CodeType(
+            argcount,  # co_argcount
+            0,  # posonlyargcount, eliminated
+            0,  # kwonlyargcount, eliminated
+            code.co_nlocals,
+            code.co_stacksize,
+            co_flags,
+            code.co_code,
+            code.co_consts,
+            code.co_names,
+            code.co_varnames,
+            code.co_filename,
+            code.co_name,
+            code.co_qualname,
+            code.co_firstlineno,
+            code.co_linetable,
+            code.co_exceptiontable,
+            code.co_freevars,
+            code.co_cellvars,
+        )
+    else:
+        new_code = types.CodeType(
+            argcount,  # co_argcount
+            0,  # posonlyargcount, eliminated
+            0,  # kwonlyargcount, eliminated
+            code.co_nlocals,
+            code.co_stacksize,
+            co_flags,
+            code.co_code,
+            code.co_consts,
+            code.co_names,
+            code.co_varnames,
+            code.co_filename,
+            code.co_name,
+            code.co_firstlineno,
+            code.co_linetable,
+            code.co_freevars,
+            code.co_cellvars,
+        )
+
+    fn_defaults_params = get_fn_defaults_params(fn)
+    new_fn = types.FunctionType(
+        new_code,
+        fn.__globals__,
+        fn.__name__,
+        fn_defaults_params,
+        fn.__closure__,
+    )
+    new_fn.__name__ = fn.__name__
+    new_fn.__doc__ = fn.__doc__
+    new_fn.__annotations__ = fn.__annotations__
+    new_fn.__kwdefaults__ = None  # already merged into defaults
+    return new_fn
+
+
+@dataclass
+class FunctionPack(Generic[P1, R1]):
+    fn: Callable[P1, R1]
+    infer_meta: Callable[..., Any]
+
+    def id(self) -> int:
+        return id(self.fn)
+
+
+class ConstantParams:
+    def __init__(self, params: dict[str, Any]):
+        self.params = params
+
+    def __hash__(self):
+        return custom_hash(self.params)
+
+    def __eq__(self, other):
+        if not isinstance(other, ConstantParams):
+            return False
+        return self.params == other.params
+
+
+@dataclass
+class OriginalFunctionPack(FunctionPack[P1, R1]):
+    def __post_init__(self):
+        self._specialized_fns: dict[ConstantParams, FunctionPack[P1, R1]] = {}
+
+    @cached_property
+    def fn_eliminated(self) -> Callable[P1, R1]:
+        return eliminate_positional_or_keyword_only(self.fn)
+
+    @cached_property
+    def infer_meta_eliminated(self) -> Callable[..., Any]:
+        return eliminate_positional_or_keyword_only(self.infer_meta)
+
+    def get_bound_args(self, /, *args: P1.args, **kwargs: P1.kwargs):
+        sig = inspect.signature(self.fn)
+        bound_args = sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+        return bound_args.arguments
+
+    def separate_mutable_and_const_params(
+        self, /, *args: P1.args, **kwargs: P1.kwargs
+    ) -> tuple[dict[str, paddle.pir.Value], dict[str, Any]]:
+        params = self.get_bound_args(*args, **kwargs)
+
+        mutable_params = {}
+        const_params = {}
+
+        # TODO: Support container types like list, dict, tuple
+        for k, v in params.items():
+            if isinstance(v, paddle.pir.Value):
+                mutable_params[k] = v
+            else:
+                const_params[k] = v
+
+        return mutable_params, const_params
+
+    def specialize(self, const_params: dict[str, Any]) -> FunctionPack[P1, R1]:
+        const_params_wrapper = ConstantParams(const_params)
+        if const_params_wrapper in self._specialized_fns:
+            return self._specialized_fns[const_params_wrapper]
+
+        specialized_fn = partial(self.fn_eliminated, **const_params)
+        specialized_infer_meta = partial(
+            self.infer_meta_eliminated, **const_params
+        )
+        specialized_fn_pack = FunctionPack(
+            specialized_fn, specialized_infer_meta
+        )
+        self._specialized_fns[const_params_wrapper] = specialized_fn_pack
+        return specialized_fn_pack
+
+
+class FunctionRegistry:
+    def __init__(self):
+        self._registry: dict[str, OriginalFunctionPack[Any, Any]] = {}
+
+    def register(
+        self,
+        name: str,
+        fn: Callable[P1, R1],
+        infer_meta: Callable[..., Any],
+    ):
+        if name not in self._registry:
+            self._registry[name] = OriginalFunctionPack(fn, infer_meta)
+            return self._registry[name]
+        fn_pack = self._registry[name]
+        if fn is not fn_pack.fn or infer_meta is not fn_pack.infer_meta:
+            raise ValueError(
+                f"Function '{name}' is already registered with a different implementation."
+            )
+        return fn_pack
+
+    def get(self, name: str) -> OriginalFunctionPack[Any, Any]:
+        if name not in self._registry:
+            raise KeyError(f"Function '{name}' is not registered.")
+        return self._registry[name]
+
+
+FUNCTION_REGISTRY = FunctionRegistry()
+
+
+def bind_constants(fn, infer_meta, *args, **kwargs):
+    sig = inspect.signature(fn)
+    bound_args = sig.bind(*args, **kwargs)
+    bound_args.apply_defaults()
+    params = bound_args.arguments
+
+    mutable_params = {}
+    const_params = {}
+
+    for k, v in params.items():
+        if isinstance(v, paddle.pir.Value):
+            mutable_params[k] = v
+        else:
+            const_params[k] = v
+
+    mutable_arg_names = list(mutable_params.keys())
+    fn = eliminate_positional_or_keyword_only(fn)
+    infer_meta = eliminate_positional_or_keyword_only(infer_meta)
+    return (
+        mutable_arg_names,
+        partial(fn, **const_params),
+        partial(infer_meta, **const_params),
+        list(mutable_params.values()),
+        const_params,
+    )
+
+
+def run_in_dynamic_mode(fn):
+    def dynamic_mode_fn(*args, **kwargs):
+        with paddle.base.dygraph.base.guard():
+            return fn(*args, **kwargs)
+
+    return dynamic_mode_fn
+
+
+def custom_hash(obj):
+    # Compute a hash for various types of objects, including unhashable ones.
+    # This may not be collision-free. For example, hash(-1) is same as hash(-2).
+    # We use dict to resolve collisions in ConstantParams.
+
+    # Handle basic types
+    if isinstance(obj, (int, float, str, bool, bytes)):
+        return hash(obj)
+
+    # Handle sequences (like list, tuple, set, frozenset)
+    if isinstance(obj, (Sequence, frozenset, set)):
+        type_id_map = {list: 1, tuple: 2, frozenset: 3, set: 4}
+        type_id = type_id_map.get(type(obj), 0)
+        return hash((type_id, *tuple(custom_hash(item) for item in obj)))
+
+    # Handle mappings (like dict)
+    if isinstance(obj, Mapping):
+        type_id = 5
+        items_hashed = tuple(
+            sorted((custom_hash(k), custom_hash(v)) for k, v in obj.items())
+        )
+        return hash((type_id, *items_hashed))
+
+    # Fallback: try to use the built-in hash, or use id() if unhashable
+    try:
+        return hash(obj)
+    except TypeError:
+        return id(obj)
+
+
+@overload
+def register_op(
+    fn: Callable[P1, R1],
+    /,
+    *,
+    name: str | None = None,
+    infer_meta: Callable[..., Any] | None = None,
+    input_names: list[str] | None = None,
+    output_names: list[str] | None = None,
+    inplace_map: dict[str, str] | None = None,
+) -> Callable[P1, R1]: ...
+
+
+@overload
+def register_op(
+    fn: None = None,
+    /,
+    *,
+    name: str | None = None,
+    infer_meta: Callable[..., Any] | None = None,
+    input_names: list[str] | None = None,
+    output_names: list[str] | None = None,
+    inplace_map: dict[str, str] | None = None,
+) -> Callable[[Callable[P1, R1]], Callable[P1, R1]]: ...
+
+
+def register_op(
+    fn: Callable[P1, R1] | None = None,
+    /,
+    *,
+    name: str | None = None,
+    infer_meta: Callable[..., Any] | None = None,
+    input_names: list[str] | None = None,
+    output_names: list[str] | None = None,
+    inplace_map: dict[str, str] | None = None,
+):
+    if input_names is None:
+        raise ValueError("Currently, input_names must be provided.")
+    if output_names is None:
+        raise ValueError("Currently, output_names must be provided.")
+    if infer_meta is None:
+        raise ValueError("Currently, infer_meta must be provided.")
+
+    def _register_op(
+        real_fn: Callable[P1, R1],
+    ) -> Callable[P1, R1]:
+        op_name = name or real_fn.__name__
+
+        @paddle.jit.marker.unified
+        @wraps(real_fn)
+        def wrapped_fn(*args: P1.args, **kwargs: P1.kwargs) -> R1:
+            if paddle.in_dynamic_mode():
+                return real_fn(*args, **kwargs)
+
+            fn_pack = FUNCTION_REGISTRY.register(op_name, real_fn, infer_meta)
+            mutable_params, const_params = (
+                fn_pack.separate_mutable_and_const_params(*args, **kwargs)
+            )
+            specialized_fn_pack = fn_pack.specialize(const_params)
+
+            assert len(mutable_params) == len(input_names), (
+                f"Number of mutable arguments ({len(mutable_params)}) does not match "
+                f"the number of input names ({len(input_names)})."
+            )
+
+            out = _C_ops._run_python_op(
+                *mutable_params.values(),
+                name=f"{op_name}_{specialized_fn_pack.id()}",
+                input_names=input_names,
+                output_names=output_names,
+                attrs={
+                    "infer_meta_fn_ptr": specialized_fn_pack.infer_meta,
+                    "fn_ptr": run_in_dynamic_mode(specialized_fn_pack.fn),
+                },
+                inplace_map=inplace_map or {},
+            )
+
+            return out[0] if len(output_names) == 1 else out
+
+        return wrapped_fn
+
+    # Handle @register_op(...)
+    if fn is None:
+        return _register_op
+    # Handle @register_op
+    return _register_op(fn)