chore: import upstream snapshot with attribution

python/paddle/.gitignore

@@ -0,0 +1,3 @@
+version.py
+cuda_env.py
+version

python/paddle/_C.py

@@ -0,0 +1,34 @@
+# 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 typing
+
+from paddle.base import core, libpaddle
+from paddle.base.libpaddle import (
+    _get_current_raw_stream as _cuda_getCurrentRawStream,  # noqa: F401
+)
+
+# Define _GLIBCXX_USE_CXX11_ABI based on compilation flags
+_GLIBCXX_USE_CXX11_ABI = getattr(libpaddle, '_GLIBCXX_USE_CXX11_ABI', True)
+_PYBIND11_COMPILER_TYPE = getattr(libpaddle, '_PYBIND11_COMPILER_TYPE', "")
+_PYBIND11_STDLIB = getattr(libpaddle, '_PYBIND11_STDLIB', "")
+_PYBIND11_BUILD_ABI = getattr(libpaddle, '_PYBIND11_BUILD_ABI', "")
+
+
+def _get_custom_class_python_wrapper(
+    namespace_name: str, class_name: str
+) -> typing.Any:
+    return core.torch_compat._get_custom_class_python_wrapper(
+        namespace_name, class_name
+    )

python/paddle/_C_ops.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2021 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 paddle.base import core
+
+__all__ = []
+
+
+for name in dir(core.eager.ops):
+    globals()[name] = getattr(core.eager.ops, name)
+    __all__.append(name)
+
+for name in dir(core.pir.ops):
+    globals()[name] = getattr(core.pir.ops, name)
+    if name not in __all__:
+        __all__.append(name)

python/paddle/_C_ops.pyi

@@ -0,0 +1,16 @@
+# Copyright (c) 2021 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 paddle._typing.libs.libpaddle.eager.ops import *  # noqa: F403
+from paddle._typing.libs.libpaddle.pir.ops import *  # noqa: F403

python/paddle/_classes.py

@@ -0,0 +1,63 @@
+# 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.
+
+# The file has been adapted from pytorch project
+# Licensed under BSD-style license -
+# https://github.com/pytorch/pytorch/blob/main/LICENSE
+
+from __future__ import annotations
+
+import types
+from typing import Any
+
+import paddle
+
+from ._ops import import_module, load_library
+
+PADDLE_CLASSES_MODULE_NAME = "paddle.classes"
+
+
+class ClassesNameSpace(types.ModuleType):
+    def __init__(self, name: str):
+        super().__init__(f"{PADDLE_CLASSES_MODULE_NAME}.{name}")
+        self.name = name
+
+    def __getattr__(self, name: str) -> Any:
+        if name == "__file__":
+            return PADDLE_CLASSES_MODULE_NAME  # type: ignore
+        return paddle.base.core.torch_compat._get_custom_class_python_wrapper(
+            self.name, name
+        )
+
+
+class PaddleClassesModule(types.ModuleType):
+    __file__ = "_classes.py"
+
+    def __init__(self):
+        super().__init__(PADDLE_CLASSES_MODULE_NAME)
+
+    def __getattr__(self, name: str):
+        namespace = ClassesNameSpace(name)
+        # Insert to __dict__ to avoid repeatedly __getattr__ overhead
+        setattr(self, name, namespace)
+        return namespace
+
+    def import_module(self, module):
+        return import_module(module)
+
+    def load_library(self, path):
+        return load_library(path)
+
+
+classes = PaddleClassesModule()

python/paddle/_legacy_C_ops.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2022 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 paddle.base import core
+
+__all__ = []
+
+for name in dir(core.eager.ops.legacy):
+    globals()[name] = getattr(core.eager.ops.legacy, name)
+    __all__.append(name)

python/paddle/_ops.py

@@ -0,0 +1,155 @@
+# 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.
+
+# The file has been adapted from pytorch project
+# Licensed under BSD-style license -
+# https://github.com/pytorch/pytorch/blob/main/LICENSE
+
+from __future__ import annotations
+
+import contextlib
+import ctypes
+import importlib
+import os
+import sys
+import types
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+from typing_extensions import ParamSpec
+
+import paddle
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+_InputT = ParamSpec("_InputT")
+_RetT = TypeVar("_RetT")
+
+PADDLE_OPS_MODULE_NAME = "paddle.ops"
+
+# Query `hasattr` only once.
+_SET_GLOBAL_FLAGS = hasattr(sys, "getdlopenflags") and hasattr(
+    sys, "setdlopenflags"
+)
+
+
+@contextlib.contextmanager
+def dl_open_guard():
+    """
+    Context manager to set the RTLD_GLOBAL dynamic linker flag while we open a
+    shared library to load custom operators.
+    """
+    if not _SET_GLOBAL_FLAGS:
+        yield
+        return
+    old_flags = sys.getdlopenflags()
+    sys.setdlopenflags(old_flags | ctypes.RTLD_GLOBAL)
+    try:
+        yield
+    finally:
+        sys.setdlopenflags(old_flags)
+
+
+def import_module(module: str):
+    return importlib.import_module(module)
+
+
+def load_library(path: str):
+    """
+    Load a shared library at the specified path.
+    """
+    path = os.path.realpath(path)
+    with dl_open_guard():
+        ctypes.CDLL(path)
+
+
+class PythonOpRegistry:
+    def __init__(self):
+        self._registry: dict[str, Callable[..., object]] = {}
+
+    def register(self, name: str, fn: Callable[..., object]):
+        if name in self._registry:
+            raise ValueError(f"Operator '{name}' is already registered.")
+        self._registry[name] = fn
+
+    def has_operator(self, name: str) -> bool:
+        return name in self._registry
+
+    def get_operator(self, name: str) -> Callable[..., object]:
+        if name not in self._registry:
+            raise ValueError(f"Operator '{name}' is not registered.")
+        return self._registry[name]
+
+
+PYTHON_OP_REGISTRY = PythonOpRegistry()
+
+
+class OverloadedOpFunction(Generic[_InputT, _RetT]):
+    def __init__(self, namespace: str, name: str):
+        self.namespace = namespace
+        self.name = name
+
+    @cached_property
+    def callable_fn(self) -> Callable[_InputT, _RetT]:
+        if PYTHON_OP_REGISTRY.has_operator(f"{self.namespace}::{self.name}"):
+            return PYTHON_OP_REGISTRY.get_operator(  # type: ignore
+                f"{self.namespace}::{self.name}"
+            )
+        return paddle.base.core.torch_compat._get_operation(
+            f"{self.namespace}::{self.name}"
+        )
+
+    def __getattr__(self, name: str) -> Callable[_InputT, _RetT]:
+        if name == "default":
+            return self.callable_fn
+        raise AttributeError(
+            f"'{self.namespace}.{self.name}' has no attribute '{name}'"
+        )
+
+    def __call__(self, *args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
+        return self.callable_fn(*args, **kwargs)
+
+
+class OpNameSpace(types.ModuleType):
+    def __init__(self, name):
+        super().__init__(f"{PADDLE_OPS_MODULE_NAME}.{name}")
+        self.name = name
+
+    def __getattr__(self, name: str) -> OverloadedOpFunction[..., Any]:
+        if name == "__file__":
+            return PADDLE_OPS_MODULE_NAME  # type: ignore
+        return OverloadedOpFunction(self.name, name)
+
+
+class PaddleOpsModule(types.ModuleType):
+    __file__ = "_ops.py"
+
+    def __init__(self):
+        super().__init__(PADDLE_OPS_MODULE_NAME)
+
+    def __getattr__(self, name: str):
+        namespace = OpNameSpace(name)
+        # Insert to __dict__ to avoid repeatedly __getattr__ overhead
+        setattr(self, name, namespace)
+        return namespace
+
+    def import_module(self, module):
+        return import_module(module)
+
+    def load_library(self, path):
+        return load_library(path)
+
+
+ops = PaddleOpsModule()

python/paddle/_pir_ops.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2021 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 paddle.base import core
+
+__all__ = []
+
+for name in dir(core.pir.ops):
+    globals()[name] = getattr(core.pir.ops, name)
+    __all__.append(name)

python/paddle/_typing/__init__.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2024 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.
+
+# Basic
+from .basic import (
+    NestedList as NestedList,
+    NestedNumericSequence as NestedNumericSequence,
+    NestedSequence as NestedSequence,
+    NestedStructure as NestedStructure,
+    Numeric as Numeric,
+    NumericSequence as NumericSequence,
+    ParamAttrLike as ParamAttrLike,
+    TensorIndex as TensorIndex,
+    TensorLike as TensorLike,
+    TensorOrTensors as TensorOrTensors,
+    unreached as unreached,
+)
+
+# Device
+from .device_like import (
+    PlaceLike as PlaceLike,
+)
+
+# DType
+from .dtype_like import DTypeLike as DTypeLike
+
+# DataLayout
+from .layout import (
+    DataLayout0D as DataLayout0D,
+    DataLayout1D as DataLayout1D,
+    DataLayout1DVariant as DataLayout1DVariant,
+    DataLayout2D as DataLayout2D,
+    DataLayout3D as DataLayout3D,
+    DataLayoutImage as DataLayoutImage,
+    DataLayoutND as DataLayoutND,
+)
+
+# Shape
+from .shape import (
+    ShapeLike as ShapeLike,
+    Size1 as Size1,
+    Size2 as Size2,
+    Size3 as Size3,
+    Size4 as Size4,
+    Size5 as Size5,
+    Size6 as Size6,
+    SizeN as SizeN,
+)

python/paddle/_typing/basic.py

@@ -0,0 +1,66 @@
+# Copyright (c) 2024 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
+
+from collections.abc import Sequence
+from types import EllipsisType
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    TypeAlias,
+    TypeVar,
+    Union,
+)
+
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import Never
+
+if TYPE_CHECKING:
+    from paddle import ParamAttr, Tensor
+    from paddle.nn.initializer import Initializer
+    from paddle.regularizer import WeightDecayRegularizer
+
+
+Numeric: TypeAlias = Union[int, float, bool, complex, np.number, "Tensor"]
+TensorLike: TypeAlias = Union[npt.NDArray[Any], "Tensor", Numeric]
+_TensorIndexItem: TypeAlias = Union[
+    None, bool, int, slice, "Tensor", EllipsisType
+]
+TensorIndex: TypeAlias = (
+    _TensorIndexItem | tuple[_TensorIndexItem, ...] | list[_TensorIndexItem]
+)
+
+
+_T = TypeVar("_T")
+
+NestedSequence = _T | Sequence["NestedSequence[_T]"]
+NestedList = _T | list["NestedList[_T]"]
+NestedStructure = (
+    _T | dict[str, "NestedStructure[_T]"] | Sequence["NestedStructure[_T]"]
+)
+NumericSequence = Sequence[Numeric]
+NestedNumericSequence: TypeAlias = NestedSequence[Numeric]
+TensorOrTensors: TypeAlias = Union["Tensor", Sequence["Tensor"]]
+
+ParamAttrLike: TypeAlias = Union[
+    "ParamAttr", "Initializer", "WeightDecayRegularizer", str, bool
+]
+
+
+def unreached() -> Never:
+    """Mark a code path as unreachable.
+    Refer to https://typing.readthedocs.io/en/latest/source/unreachable.html#marking-code-as-unreachable
+    """
+    raise RuntimeError("Unreachable code path")

python/paddle/_typing/device_like.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2024 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
+
+from typing import TYPE_CHECKING, TypeAlias, Union
+
+if TYPE_CHECKING:
+    from paddle.base.core import Place
+
+PlaceLike: TypeAlias = Union[
+    "Place",
+    str,  # some string like "cpu", "gpu:0", etc.
+]

python/paddle/_typing/dtype_like.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2024 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
+
+from typing import TYPE_CHECKING, Literal, TypeAlias, Union
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from paddle import dtype
+
+_DTypeLiteral: TypeAlias = Literal[
+    "uint8",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "float32",
+    "float64",
+    "float16",
+    "bfloat16",
+    "complex64",
+    "complex128",
+    "bool",
+]
+
+_DTypeNumpy: TypeAlias = (
+    type[
+        np.uint8
+        | np.int8
+        | np.int16
+        | np.int32
+        | np.int64
+        | np.float16
+        | np.float32
+        | np.float64
+        | np.complex64
+        | np.complex128
+        | np.bool_
+    ]
+    | np.dtype
+)
+
+
+DTypeLike: TypeAlias = Union["dtype", _DTypeNumpy, _DTypeLiteral]

python/paddle/_typing/layout.py

@@ -0,0 +1,29 @@
+# Copyright (c) 2024 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
+
+from typing import Literal, TypeAlias
+
+# Note: Do not conform to predefined naming style in pylint.
+DataLayout0D: TypeAlias = Literal["NC"]
+DataLayout1D: TypeAlias = Literal["NCL", "NLC"]
+DataLayout2D: TypeAlias = Literal["NCHW", "NHWC"]
+DataLayout3D: TypeAlias = Literal["NCDHW", "NDHWC"]
+
+DataLayoutND: TypeAlias = (
+    DataLayout0D | DataLayout1D | DataLayout2D | DataLayout3D
+)
+
+DataLayout1DVariant: TypeAlias = Literal["NCW", "NWC"]
+DataLayoutImage: TypeAlias = Literal["HWC", "CHW"]

python/paddle/_typing/libs/README.md

@@ -0,0 +1,5 @@
+# Stub files for pybind11 C++ APIs
+
+The files in this directory (`python/paddle/_typing/libs`) are auto generated from pybind11 C++ APIs.
+
+Please do NOT edit these files, or the changes may be lost.

python/paddle/_typing/shape.py

@@ -0,0 +1,42 @@
+# Copyright (c) 2024 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
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, TypeAlias, Union
+
+if TYPE_CHECKING:
+    from .. import Tensor
+
+_DynamicShapeLike: TypeAlias = Union[
+    Sequence[Union[int, "Tensor", None]],
+    "Tensor",
+]
+
+
+_StaticShapeLike: TypeAlias = Union[
+    Sequence[int],
+    "Tensor",
+]
+
+ShapeLike: TypeAlias = _DynamicShapeLike | _StaticShapeLike
+
+# for size parameters, eg, kernel_size, stride ...
+Size1: TypeAlias = int | tuple[int] | list[int]
+Size2: TypeAlias = int | tuple[int, int] | list[int]
+Size3: TypeAlias = int | tuple[int, int, int] | list[int]
+Size4: TypeAlias = int | tuple[int, int, int, int] | list[int]
+Size5: TypeAlias = int | tuple[int, int, int, int, int] | list[int]
+Size6: TypeAlias = int | tuple[int, int, int, int, int, int] | list[int]
+SizeN: TypeAlias = int | tuple[int, ...] | list[int]

python/paddle/amp/__init__.py

@@ -0,0 +1,113 @@
+#   Copyright (c) 2020 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
+
+from paddle.base import core
+from paddle.base.framework import (
+    _current_expected_place,
+    _get_paddle_place,
+)
+
+from . import (  # noqa: F401
+    accuracy_compare,
+    debugging,
+    grad_scaler,
+)
+from .amp_lists import (  # noqa: F401
+    black_list,
+    white_list,
+)
+from .auto_cast import (  # noqa: F401
+    amp_decorate,
+    amp_guard,
+    auto_cast,
+    autocast,
+    decorate,
+    get_autocast_dtype,
+    is_autocast_enabled,
+)
+from .grad_scaler import (  # noqa: F401
+    AmpScaler,
+    GradScaler,
+    OptimizerState,
+)
+
+__all__ = [
+    'auto_cast',
+    'GradScaler',
+    'decorate',
+    'is_float16_supported',
+    'is_bfloat16_supported',
+    'is_autocast_enabled',
+    'get_autocast_dtype',
+    'get_autocast_cpu_dtype',
+    'get_autocast_gpu_dtype',
+]
+
+get_autocast_cpu_dtype = get_autocast_dtype
+get_autocast_gpu_dtype = get_autocast_dtype
+
+
+def is_float16_supported(device: str | None = None) -> bool:
+    """
+    Determine whether the place supports float16 in the auto-mixed-precision training.
+
+    Args:
+        device (str|None, optional): Specify the running device.
+            It can be ``cpu``, ``gpu``, ``xpu``, ``gpu:x`` and ``xpu:x``,
+            where ``x`` is the index of the GPUs or XPUs. if device is None, the device is the current device. Default: None.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> paddle.amp.is_float16_supported()  # True or False
+            False
+    """
+
+    device = (
+        _current_expected_place()
+        if device is None
+        else _get_paddle_place(device)
+    )
+
+    return core.is_float16_supported(device)
+
+
+def is_bfloat16_supported(device: str | None = None) -> bool:
+    """
+    Determine whether the place supports bfloat16 in the auto-mixed-precision training.
+
+    Args:
+        device (str|None, optional): Specify the running device.
+            It can be ``cpu``, ``gpu``, ``xpu``, ``gpu:x`` and ``xpu:x``,
+            where ``x`` is the index of the GPUs or XPUs. if device is None, the device is the current device. Default: None.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> paddle.amp.is_bfloat16_supported()  # True or False
+            True
+    """
+
+    device = (
+        _current_expected_place()
+        if device is None
+        else _get_paddle_place(device)
+    )
+
+    return core.is_bfloat16_supported(device)

python/paddle/amp/accuracy_compare.py

@@ -0,0 +1,700 @@
+#   Copyright (c) 2023 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 os
+
+import numpy as np
+
+
+# Judge whether the value is within the range indicated by fp16
+def is_infinite(value, dtype=np.float16):
+    # return value > np.finfo(np.float16).max or value < np.finfo(np.float16).min
+    array = np.array([value]).astype(dtype)
+    return np.isinf(array) or np.isnan(array)
+
+
+# Judge whether the value of fp32 is equal to that of fp16
+def is_allclose(actual, expected, atol=1e-2, rtol=1e-2):
+    return np.allclose(
+        np.array([actual]), np.array([expected]), atol=atol, rtol=rtol
+    )
+
+
+class TensorInfo:
+    def __init__(self):
+        self.device = None
+        self.op_type = None
+        self.tensor_name = None
+        self.dtype = None
+        self.numel = None
+        self.max_value = None
+        self.min_value = None
+        self.mean_value = None
+        self.has_inf = None
+        self.has_nan = None
+        self.num_zero = None
+
+    def __str__(self):
+        return f"[TensorInfo] device={self.device}, op_type={self.op_type}, tensor_name={self.tensor_name}, dtype={self.dtype}, numel={self.numel}, num_inf={self.has_inf}, num_nan={self.has_nan}, num_zero={self.num_zero}, max_value={self.max_value:.6f}, min_value={self.min_value:.6f}, mean_value={self.mean_value:.6f}"
+
+    def key(
+        self,
+    ):
+        return self.op_type + "/" + self.tensor_name
+
+    def init_from_string(self, line):
+        try:
+            line_frags = line.strip().split(" ")
+            for frag in line_frags:
+                word_str = (
+                    frag.replace("[", "").replace("]", "").replace(",", "")
+                )
+                words = word_str.split("=")
+                if words[0] == "op":
+                    self.op_type = words[1]
+                elif words[0] == "device":
+                    self.device = words[1]
+                elif words[0] == "tensor":
+                    self.tensor_name = words[1]
+                elif words[0] == "dtype":
+                    self.dtype = words[1]
+                elif words[0] == "numel":
+                    self.numel = np.int64(words[1])
+                elif words[0] == "max":
+                    self.max_value = np.float32(words[1])
+                elif words[0] == "min":
+                    self.min_value = np.float32(words[1])
+                elif words[0] == "mean":
+                    self.mean_value = np.float32(words[1])
+                elif words[0] == "num_inf":
+                    self.has_inf = int(words[1])
+                elif words[0] == "num_nan":
+                    self.has_nan = int(words[1])
+                elif words[0] == "num_zero":
+                    self.num_zero = np.int64(words[1])
+        except Exception as e:
+            print(f"!! Error parsing {line}")
+        return self
+
+
+class MixedPrecisionTensorInfo:
+    def __init__(
+        self, fp32_tensor_info, fp16_tensor_info, fp32_idx=0, grad_scale=1.0
+    ):
+        self.is_normal = True
+        self.fp32_idx = fp32_idx
+
+        self.fp32_tensor_name = None
+        self.fp32_dtype = None
+        self.fp32_max_value = None
+        self.fp32_min_value = None
+        self.fp32_mean_value = None
+        self.fp32_num_zero = None
+        self.scaled_fp32_max_value = None
+        self.scaled_fp32_min_value = None
+
+        self.fp16_tensor_name = None
+        self.fp16_dtype = None
+        self.fp16_max_value = None
+        self.fp16_min_value = None
+        self.fp16_mean_value = None
+        self.fp16_num_zero = None
+        self.fp16_has_inf = None
+        self.fp16_has_nan = None
+
+        self.fp32_div_fp16_max_value = None
+        self.fp32_div_fp16_min_value = None
+        self.fp32_div_fp16_mean_value = None
+
+        if fp32_tensor_info is not None:
+            self.op_type = fp32_tensor_info.op_type
+            self.numel = fp32_tensor_info.numel
+            self.fp32_num_zero = fp32_tensor_info.num_zero
+            self.fp32_tensor_name = fp32_tensor_info.tensor_name
+            self.fp32_dtype = fp32_tensor_info.dtype
+            self.fp32_max_value = fp32_tensor_info.max_value
+            self.fp32_min_value = fp32_tensor_info.min_value
+            self.fp32_mean_value = fp32_tensor_info.mean_value
+            if "GRAD" in self.fp32_tensor_name:
+                self.scaled_fp32_max_value = (
+                    grad_scale * fp32_tensor_info.max_value
+                )
+                self.scaled_fp32_min_value = (
+                    grad_scale * fp32_tensor_info.min_value
+                )
+
+        if fp16_tensor_info is not None:
+            self.op_type = fp16_tensor_info.op_type
+            self.numel = fp16_tensor_info.numel
+            self.fp16_num_zero = fp16_tensor_info.num_zero
+            self.fp16_tensor_name = fp16_tensor_info.tensor_name
+            self.fp16_dtype = fp16_tensor_info.dtype
+            self.fp16_max_value = fp16_tensor_info.max_value
+            self.fp16_min_value = fp16_tensor_info.min_value
+            self.fp16_mean_value = fp16_tensor_info.mean_value
+            self.fp16_has_inf = fp16_tensor_info.has_inf
+            self.fp16_has_nan = fp16_tensor_info.has_nan
+
+        if fp32_tensor_info is not None and fp16_tensor_info is not None:
+            # Check whether the op name and data are equal
+            assert fp32_tensor_info.op_type == fp16_tensor_info.op_type
+            assert fp32_tensor_info.numel == fp16_tensor_info.numel, (
+                f"Error:\n\tFP32 Tensor Info:{fp32_tensor_info}\n\tFP16 Tensor Info:{fp16_tensor_info}"
+            )
+            # Fp16 divided by fp32
+            self.fp32_div_fp16_max_value = self._div(
+                self.fp16_max_value, self.fp32_max_value
+            )
+            self.fp32_div_fp16_min_value = self._div(
+                self.fp16_min_value, self.fp32_min_value
+            )
+            self.fp32_div_fp16_mean_value = self._div(
+                self.fp16_mean_value, self.fp32_mean_value
+            )
+
+        self._check_normal()
+
+    def __str__(self):
+        def _float_str(value):
+            return f"{value:.6f}" if value is not None else value
+
+        debug_str = f"[MixedPrecisionTensorInfo] op_type={self.op_type}, numel={self.numel}"
+        debug_str += f"\n  FP32: tensor_name={self.fp32_tensor_name}, dtype={self.fp32_dtype}, max_value={_float_str(self.fp32_max_value)}, min_value={_float_str(self.fp32_min_value)}, mean_value={_float_str(self.fp32_mean_value)}"
+        debug_str += f"\n  FP16: tensor_name={self.fp16_tensor_name}, dtype={self.fp16_dtype}, max_value={_float_str(self.fp16_max_value)}, min_value={_float_str(self.fp16_min_value)}, mean_value={_float_str(self.fp16_mean_value)}, has_inf={self.fp16_has_inf}, has_nan={self.fp16_has_nan}"
+        return debug_str
+
+    def _div(self, a, b):
+        if a is not None and b is not None:
+            return a / b if b != 0 else 1
+        return None
+
+    def get_tensor_name(self):
+        if self.fp32_tensor_name is None:
+            return self.fp16_tensor_name  # + "#" + str(self.idx)
+        elif self.fp16_tensor_name is None:
+            return self.fp32_tensor_name + "#" + str(self.fp32_idx)
+        else:
+            return (
+                self.fp16_tensor_name.replace(".cast_fp16", "/.cast_fp16/")
+                + "#"
+                + str(self.fp32_idx)
+            )
+
+    def _check_normal(self):
+        # When the OP meets the following conditions, it is abnormal data, and use --skip_normal_tensors to retain the data in Excel:
+        # 1. The number of OP outputs exceeds the indication range of int32
+        # 2. The output data exceeds the representation range of fp16
+        # 3. Nan or inf appears in fp16 output data
+        # 4. The maximum value of fp32 is not equal to the maximum value of fp16
+        # 5. The minimum value of fp32 is not equal to the minimum value of fp16
+        if self.numel is not None and self.numel > np.iinfo(np.int32).max:
+            self.is_normal = False
+            return
+
+        check_list = [
+            self.fp32_max_value,
+            self.fp32_min_value,
+            self.scaled_fp32_max_value,
+            self.scaled_fp32_min_value,
+            self.fp16_max_value,
+            self.fp16_min_value,
+        ]
+
+        for value in check_list:
+            if value is not None and is_infinite(value):
+                self.is_normal = False
+                return
+
+        if self.fp16_has_inf is not None and self.fp16_has_inf:
+            self.is_normal = False
+            return
+        if self.fp16_has_nan is not None and self.fp16_has_nan:
+            self.is_normal = False
+            return
+
+        if (
+            self.scaled_fp32_max_value is not None
+            and self.fp16_max_value is not None
+            and not is_allclose(self.fp16_max_value, self.scaled_fp32_max_value)
+        ):
+            self.is_normal = False
+            return
+        if (
+            self.scaled_fp32_min_value is not None
+            and self.fp16_min_value is not None
+            and not is_allclose(self.fp16_min_value, self.scaled_fp32_min_value)
+        ):
+            self.is_normal = False
+            return
+
+
+class ExcelWriter:
+    def __init__(self, log_fp32_dir, log_fp16_dir, output_path):
+        self.log_fp32_dir = log_fp32_dir
+        self.log_fp16_dir = log_fp16_dir
+
+        try:
+            import xlsxwriter as xlw
+        except ImportError:
+            print(
+                "import xlsxwriter failed. please run 'pip install xlsxwriter==3.0.9' to install it"
+            )
+
+        self.workbook = xlw.Workbook(output_path)
+        self.title_format = self.workbook.add_format(
+            {
+                'bold': True,
+                'border': 1,
+                'font_color': 'black',
+                'bg_color': '#6495ED',
+                'align': 'center',
+            }
+        )
+        self.tensor_name_format = self.workbook.add_format(
+            {'bold': True, 'bg_color': '#F5F5F5'}
+        )
+        self.red_bg_cell_format = self.workbook.add_format(
+            {'bold': True, 'bg_color': 'red'}
+        )
+        self.yellow_bg_cell_format = self.workbook.add_format(
+            {'bold': True, 'bg_color': 'yellow'}
+        )
+        self.orange_bg_cell_format = self.workbook.add_format(
+            {'bold': True, 'bg_color': 'orange'}
+        )
+
+    def close(self):
+        self.workbook.close()
+        self.workbook = None
+
+    def _write_dtype(self, worksheet, value, row, col):
+        if value is None:
+            worksheet.write(row, col, "--")
+        else:
+            if value == "fp16":
+                worksheet.write(row, col, value, self.yellow_bg_cell_format)
+            else:
+                worksheet.write(row, col, value)
+
+    def _write_tensor_name(self, worksheet, mp_tensor_info, row, col):
+        tensor_name = mp_tensor_info.get_tensor_name()
+        if (
+            mp_tensor_info.fp32_tensor_name is not None
+            and mp_tensor_info.fp16_tensor_name
+        ):
+            worksheet.write(row, col, tensor_name, self.tensor_name_format)
+        else:
+            worksheet.write(row, col, tensor_name)
+
+    def _write_maxmin_value(
+        self, worksheet, value, row, col, check_finite=True
+    ):
+        if value is None:
+            worksheet.write(row, col, "--")
+        else:
+            if abs(value) < 1e-5:
+                value_str = f"{value:.6E}"
+            else:
+                value_str = f"{value:.6f}"
+            if check_finite and is_infinite(value, np.float16):
+                worksheet.write(row, col, value_str, self.red_bg_cell_format)
+            else:
+                worksheet.write(row, col, value_str)
+
+    def _write_tensor_num_zero(
+        self, worksheet, value, row, col, check_finite=True
+    ):
+        if value is None:
+            worksheet.write(row, col, "--")
+        else:
+            value_str = f"{value:>10d}"
+            worksheet.write(row, col, value_str)
+
+    def _write_infinite_status(self, worksheet, value, row, col):
+        if value is None:
+            worksheet.write(row, col, "--")
+        else:
+            if value == 1:
+                worksheet.write(row, col, value, self.red_bg_cell_format)
+            else:
+                worksheet.write(row, col, value)
+
+    def _write_fp32divfp16_value(self, worksheet, value, row, col, loss_scale):
+        def _in_range(value, scale=1):
+            return value > scale * 0.95 and value < scale * 1.05
+
+        if value is None:
+            worksheet.write(row, col, "--")
+        else:
+            value_str = f"{value:.6f}"
+            if _in_range(value, scale=1) or _in_range(value, loss_scale):
+                worksheet.write(row, col, value_str)
+            else:
+                worksheet.write(row, col, value_str, self.orange_bg_cell_format)
+
+    def _write_titles(self, worksheet, loss_scale, row):
+        column_width_dict = {
+            "op_type": 24,
+            "tensor_name": 60,
+            "numel": 10,
+            "num_zero": 10,
+            "infinite": 8,
+            "dtype": 8,
+            "max_value": 16,
+            "min_value": 16,
+            "mean_value": 16,
+            "num_inf": 8,
+            "num_nan": 8,
+        }
+        title_names = ["op_type", "tensor_name", "numel", "infinite"]
+        if self.log_fp16_dir is None:
+            # only fp32 values
+            worksheet.merge_range("E1:H1", "fp32", self.title_format)
+            worksheet.merge_range(
+                "I1:J1", f"fp32 (scale={loss_scale})", self.title_format
+            )
+            title_names.extend(
+                [
+                    "dtype",
+                    "max_value",
+                    "min_value",
+                    "mean_value",
+                    "max_value",
+                    "min_value",
+                ]
+            )
+        elif self.log_fp32_dir is None:
+            # only fp16 values
+            worksheet.merge_range(
+                "E1:J1", f"fp16 (scale={loss_scale})", self.title_format
+            )
+            title_names.extend(
+                [
+                    "dtype",
+                    "max_value",
+                    "min_value",
+                    "mean_value",
+                    "num_zero",
+                    "num_inf",
+                    "num_nan",
+                ]
+            )
+        else:
+            # fp32 and fp16 values
+            worksheet.merge_range("E1:H1", "fp32", self.title_format)
+            worksheet.merge_range(
+                "I1:N1", f"fp16 (scale={loss_scale})", self.title_format
+            )
+            worksheet.merge_range("O1:Q1", "fp16 / fp32", self.title_format)
+            title_names.extend(
+                [
+                    "dtype",
+                    "max_value",
+                    "min_value",
+                    "mean_value",
+                    "num_zero",
+                    "dtype",
+                    "max_value",
+                    "min_value",
+                    "mean_value",
+                    "num_zero",
+                    "num_inf",
+                    "num_nan",
+                    "max_value",
+                    "min_value",
+                    "mean_value",
+                ]
+            )
+
+        for col in range(len(title_names)):
+            col_char = chr(ord("A") + col)
+            worksheet.set_column(
+                col_char + ":" + col_char, column_width_dict[title_names[col]]
+            )
+        for col in range(len(title_names)):
+            worksheet.write(row, col, title_names[col], self.title_format)
+
+    def add_worksheet(
+        self, mp_tensor_info_list, sheetname, loss_scale, skip_normal_tensors
+    ):
+        assert self.workbook is not None
+
+        worksheet = self.workbook.add_worksheet(sheetname)
+        row = 1
+
+        self._write_titles(worksheet, loss_scale, row)
+        row += 1
+
+        infinite_op_types = []
+        for tensor_info in mp_tensor_info_list:
+            if (
+                not tensor_info.is_normal
+                and tensor_info.op_type not in infinite_op_types
+            ):
+                infinite_op_types.append(tensor_info.op_type)
+
+            if skip_normal_tensors and tensor_info.is_normal:
+                continue
+
+            worksheet.write(row, 0, tensor_info.op_type)
+            self._write_tensor_name(worksheet, tensor_info, row, 1)
+
+            if tensor_info.numel > np.iinfo(np.int32).max:
+                worksheet.write(
+                    row, 2, tensor_info.numel, self.bad_value_format
+                )
+            else:
+                worksheet.write(row, 2, tensor_info.numel)
+
+            if tensor_info.is_normal:
+                worksheet.write(row, 3, "0")
+            else:
+                worksheet.write(row, 3, "1", self.red_bg_cell_format)
+
+            col = 4
+
+            if self.log_fp32_dir is not None:
+                self._write_dtype(worksheet, tensor_info.fp32_dtype, row, col)
+                self._write_maxmin_value(
+                    worksheet, tensor_info.fp32_max_value, row, col + 1
+                )
+                self._write_maxmin_value(
+                    worksheet, tensor_info.fp32_min_value, row, col + 2
+                )
+                self._write_maxmin_value(
+                    worksheet, tensor_info.fp32_mean_value, row, col + 3
+                )
+                self._write_tensor_num_zero(
+                    worksheet, tensor_info.fp32_num_zero, row, col + 4
+                )
+                col += 5
+
+                if self.log_fp16_dir is None:
+                    self._write_maxmin_value(
+                        worksheet, tensor_info.scaled_fp32_max_value, row, col
+                    )
+                    self._write_maxmin_value(
+                        worksheet,
+                        tensor_info.scaled_fp32_min_value,
+                        row,
+                        col + 1,
+                    )
+                    col += 2
+
+            if self.log_fp16_dir is not None:
+                self._write_dtype(worksheet, tensor_info.fp16_dtype, row, col)
+                self._write_maxmin_value(
+                    worksheet, tensor_info.fp16_max_value, row, col + 1
+                )
+                self._write_maxmin_value(
+                    worksheet, tensor_info.fp16_min_value, row, col + 2
+                )
+                self._write_maxmin_value(
+                    worksheet, tensor_info.fp16_mean_value, row, col + 3
+                )
+                self._write_tensor_num_zero(
+                    worksheet, tensor_info.fp32_num_zero, row, col + 4
+                )
+                col += 5
+
+                self._write_infinite_status(
+                    worksheet, tensor_info.fp16_has_inf, row, col
+                )
+                self._write_infinite_status(
+                    worksheet, tensor_info.fp16_has_nan, row, col + 1
+                )
+                col += 2
+
+            if self.log_fp32_dir is not None and self.log_fp16_dir is not None:
+                self._write_fp32divfp16_value(
+                    worksheet,
+                    tensor_info.fp32_div_fp16_max_value,
+                    row,
+                    col,
+                    loss_scale,
+                )
+                self._write_fp32divfp16_value(
+                    worksheet,
+                    tensor_info.fp32_div_fp16_min_value,
+                    row,
+                    col + 1,
+                    loss_scale,
+                )
+                self._write_fp32divfp16_value(
+                    worksheet,
+                    tensor_info.fp32_div_fp16_mean_value,
+                    row,
+                    col + 2,
+                    loss_scale,
+                )
+                col += 3
+
+            row += 1
+
+        print(f"-- OP Types produce infinite outputs: {infinite_op_types}")
+
+
+def parse_lines(lines, specified_op_list=None):
+    tensor_info_list = []
+
+    for i in range(len(lines)):
+        if i % 10 == 0:
+            print(
+                f"-- Processing {i:-8d} / {len(lines):-8d} line",
+                end="\r",
+            )
+        line = lines[i]
+        if "[PRECISION]" in line:
+            tensor_info = TensorInfo()
+            tensor_info.init_from_string(line)
+            if (
+                tensor_info.tensor_name is not None
+                and tensor_info.tensor_name != ""
+            ):
+                has_tensor_name = True
+            if (
+                specified_op_list is None
+                or tensor_info.op_type in specified_op_list
+            ):
+                tensor_info_list.append(tensor_info)
+            # print(tensor_info)
+    return tensor_info_list
+
+
+def parse_log(log_dir, filename, specified_op_list=None):
+    if log_dir is None or filename is None:
+        return None
+
+    complete_filename = log_dir + "/" + filename
+    tensor_info_list = None
+    has_tensor_name = False
+
+    try:
+        with open(complete_filename, 'r') as f:
+            lines = f.readlines()
+            tensor_info_list = parse_lines(lines, specified_op_list)
+    except FileNotFoundError:
+        print("the file ", complete_filename, "is not found")
+        return None, has_tensor_name
+    return tensor_info_list, has_tensor_name
+
+
+def merge_tensor_info_list(
+    fp32_tensor_info_list, fp16_tensor_info_list, grad_scale
+):
+    mp_tensor_info_list = []
+    if fp16_tensor_info_list is not None:
+        fp32_tensor_info_dict = {}
+        fp32_write_count = {}
+        if fp32_tensor_info_list is not None:
+            for tensor_info in fp32_tensor_info_list:
+                tensor_info_key = tensor_info.key()
+                count = fp32_write_count.get(tensor_info_key, 0)
+                fp32_write_count[tensor_info_key] = count + 1
+                fp32_tensor_info_dict[tensor_info_key + "#" + str(count)] = (
+                    tensor_info
+                )
+
+        fp32_read_count = {}
+        for i in range(len(fp16_tensor_info_list)):
+            if i % 10 == 0:
+                print(
+                    f"-- Processing {i:-8d} / {len(fp16_tensor_info_list):-8d} FP16 Tensor Info",
+                    end="\r",
+                )
+            fp16_tensor_info = fp16_tensor_info_list[i]
+            fp32_tensor_info_key = (
+                fp16_tensor_info.key()
+                .replace(".cast_fp16", "")
+                .replace(".cast_fp32", "")
+            )
+            count = fp32_read_count.get(fp32_tensor_info_key, 0)
+            fp32_tensor_info = fp32_tensor_info_dict.get(
+                fp32_tensor_info_key + "#" + str(count), None
+            )
+            if fp32_tensor_info is not None:
+                fp32_read_count[fp32_tensor_info_key] = count + 1
+            mp_tensor_info = MixedPrecisionTensorInfo(
+                fp32_tensor_info, fp16_tensor_info, count, grad_scale
+            )
+            mp_tensor_info_list.append(mp_tensor_info)
+            # print(mp_tensor_info)
+    elif fp32_tensor_info_list is not None:
+        fp32_count = {}
+        for i in range(len(fp32_tensor_info_list)):
+            if i % 10 == 0:
+                print(
+                    f"-- Processing {i:-8d} / {len(fp32_tensor_info_list):-8d} FP32 Tensor Info",
+                    end="\r",
+                )
+            tensor_info = fp32_tensor_info_list[i]
+            tensor_info_key = tensor_info.key()
+            count = fp32_count.get(tensor_info_key, 0)
+            fp32_count[tensor_info_key] = count + 1
+            mp_tensor_info = MixedPrecisionTensorInfo(
+                tensor_info, None, count, grad_scale
+            )
+            mp_tensor_info_list.append(mp_tensor_info)
+
+    return mp_tensor_info_list
+
+
+def compare_accuracy(
+    dump_path,
+    another_dump_path,
+    output_filename,
+    loss_scale=1,
+    dump_all_tensors=False,
+):
+    excel_writer = ExcelWriter(dump_path, another_dump_path, output_filename)
+    grad_scale = loss_scale
+    workerlog_filenames = []
+    filenames = os.listdir(dump_path)
+    for name in filenames:
+        if "worker_" in name:
+            workerlog_filenames.append(name)
+    print(
+        f"-- There are {len(workerlog_filenames)} workerlogs under {dump_path}: {workerlog_filenames}"
+    )
+
+    for filename in sorted(workerlog_filenames):
+        print(f"-- [Step 1/4] Parsing FP32 logs under {dump_path}/{filename}")
+        fp32_tensor_info_list, fp32_has_tensor_name = parse_log(
+            dump_path, filename, None
+        )
+        print(
+            f"-- [Step 2/4] Parsing FP16 logs under {another_dump_path}/{filename}"
+        )
+        fp16_tensor_info_list, fp16_has_tensor_name = parse_log(
+            another_dump_path, filename, None
+        )
+
+        print(f"-- [Step 3/4] Merge FP32 and FP16 tensor info for {filename}")
+        mp_tensor_info_list = merge_tensor_info_list(
+            fp32_tensor_info_list, fp16_tensor_info_list, grad_scale
+        )
+        print(
+            f"-- [Step 4/4] Add worksheet for mixed precision tensor info of {filename}"
+        )
+        excel_writer.add_worksheet(
+            mp_tensor_info_list,
+            filename,
+            loss_scale,
+            False,
+        )
+
+    print(f"-- Write to {output_filename}")
+
+    print()
+    excel_writer.close()

python/paddle/amp/amp_lists.py

@@ -0,0 +1,139 @@
+# Copyright (c) 2023 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.
+
+# The set of ops that support fp16 and bf16 calculation and are considered numerically-
+# safe and performance-critical. These ops are always converted to fp16 or bf16.
+
+from __future__ import annotations
+
+WHITE_LIST = {
+    'conv2d',
+    'einsum',
+    'matmul',
+    'matmul_v2',
+    'linear_v2',
+    'max_pool2d_with_index',
+    'mul',
+    'fused_gemm_epilogue',
+    "fused_rotary_position_embedding",
+    "flash_attn",
+}
+
+# The set of ops that support fp16, and bf16 was unsupported.
+ONLY_FP16_WHITE_LIST = {
+    'fake_quantize_dequantize_abs_max',
+    'fake_quantize_dequantize_moving_average_abs_max',
+    'fused_attention',
+    'fused_feedforward',
+}
+
+FP16_WHITE_LIST = WHITE_LIST | ONLY_FP16_WHITE_LIST
+
+# The set of ops that support fp16 calculation and are considered numerically-
+# dangerous and whose effects may also be observed in downstream ops.
+FP16_BLACK_LIST = {
+    'tan',
+    'acos',
+    'asin',
+    'sinh',
+    'cosh',
+    'atanh',
+    'tanh_shrink',
+    'erfinv',
+    'exp',
+    'expm1',
+    'log',
+    'log10',
+    'log2',
+    'reciprocal',
+    'rsqrt',
+    'pow',
+    'square',
+    'reduce_sum',
+    'mean',
+    'reduce_mean',
+    'reduce_prod',
+    'cumprod',
+    'cumsum',
+    'dist',
+    'pnorm',
+    'frobenius_norm',
+    'renorm',
+    'group_norm',
+    'layer_norm',
+    'softmax',
+    'softmin',
+    'softplus',
+    'log_softmax',
+    'softmax_with_cross_entropy',
+    'sigmoid_cross_entropy_with_logits',
+    'c_softmax_with_cross_entropy',
+    'c_softmax_with_multi_label_cross_entropy',
+    'cross_entropy',
+    'cross_entropy2',
+    'nll_loss',
+    'huber_loss',
+    'triplet_margin_loss',
+    'log_loss',
+    'hsigmoid_loss',
+    'margin_cross_entropy',
+}
+
+# FP16/BF16 performance of grad op is worse than that of FP32. Use FP32 by default.
+EXTRA_BLACK_LIST = {
+    'linear_interp_v2',
+    'nearest_interp_v2',
+    'bilinear_interp_v2',
+    'bicubic_interp_v2',
+    'trilinear_interp_v2',
+    'lookup_table',
+    'lookup_table_v2',
+    'scatter',
+}
+
+BF16_WHITE_LIST = WHITE_LIST
+BF16_BLACK_LIST = FP16_BLACK_LIST
+
+
+# At OD level, ops in WHITE_LIST will use FP16/BF16 and the others will use FP32.
+def white_list() -> dict[str, dict[str, set[str]]]:
+    white_list = {
+        "float16": {
+            "OD": FP16_WHITE_LIST,
+            "O1": FP16_WHITE_LIST,
+            "O2": FP16_WHITE_LIST,
+        },
+        "bfloat16": {
+            "OD": BF16_WHITE_LIST,
+            "O1": BF16_WHITE_LIST,
+            "O2": BF16_WHITE_LIST,
+        },
+    }
+    return white_list
+
+
+def black_list() -> dict[str, dict[str, set[str]]]:
+    black_list = {
+        "float16": {
+            "OD": set(),
+            "O1": FP16_BLACK_LIST | EXTRA_BLACK_LIST,
+            "O2": EXTRA_BLACK_LIST,
+        },
+        "bfloat16": {
+            "OD": set(),
+            "O1": BF16_BLACK_LIST | EXTRA_BLACK_LIST,
+            "O2": EXTRA_BLACK_LIST,
+        },
+    }
+    return black_list

python/paddle/amp/debugging.py

@@ -0,0 +1,731 @@
+#   Copyright (c) 2023 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 contextlib
+import random
+from enum import Enum
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    TypeVar,
+)
+
+import numpy as np
+from typing_extensions import ParamSpec
+
+import paddle
+from paddle import _C_ops
+from paddle.base import core
+
+from ..framework import LayerHelper, in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Generator, Sequence
+
+    from paddle import Tensor
+
+_InputT = ParamSpec("_InputT")
+_RetT = TypeVar("_RetT")
+__all__ = [
+    "DebugMode",
+    "TensorCheckerConfig",
+    "check_numerics",
+    "enable_operator_stats_collection",
+    "disable_operator_stats_collection",
+    "collect_operator_stats",
+    "enable_tensor_checker",
+    "disable_tensor_checker",
+    "compare_accuracy",
+    "check_layer_numerics",
+]
+
+
+class DebugMode(Enum):
+    """
+    The DebugMode is a feature that helps to present the state of the TensorCheckerConfig. Each DebugMode has a specific meaning, which is explained below:
+
+    - DebugMode.CHECK_NAN_INF_AND_ABORT: This mode prints or saves information about Tensors that contain NaN/Inf and interrupts the program.
+
+    - DebugMode.CHECK_NAN_INF: This mode prints or saves critical information about Tensors that contain NaN/Inf but allows the program to continue running.
+
+    - DebugMode.CHECK_ALL_FOR_OVERFLOW: This mode checks the output of the FP32 operator and prints or saves information about key Tensors that exceed the FP16 representation range, such as overflow or underflow.
+
+    - DebugMode.CHECK_ALL: This mode prints or saves output Tensor key information for all operators.
+
+    """
+
+    CHECK_NAN_INF_AND_ABORT = 0
+    CHECK_NAN_INF = 1
+    CHECK_ALL_FOR_OVERFLOW = 2
+    CHECK_ALL = 3
+    # CHECK_ALL_AND_ABORT = 4
+    # DUMP_ALL = 5
+
+
+def check_layer_numerics(
+    func: Callable[_InputT, _RetT],
+) -> Callable[_InputT, _RetT]:
+    """
+    This decorator is used to check the numerical values of the layer's input and output data.
+
+    Args:
+        func (callable): The function to be decorated.
+
+    Returns:
+        callable: The decorated function.
+
+    Raises:
+        None.
+
+    Example:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> class MyLayer(paddle.nn.Layer):
+            ...     def __init__(self, dtype):
+            ...         super().__init__()
+            ...         self._w = self.create_parameter([2, 3], dtype=dtype)
+            ...         self._b = self.create_parameter([2, 3], dtype=dtype)
+            ...
+            ...     @paddle.amp.debugging.check_layer_numerics
+            ...     def forward(self, x):
+            ...         # return 1/x * self._w + self._b   open it you will see the error log
+            ...         return x @ self._w + self._b
+            >>> dtype = 'float32'
+            >>> x = paddle.rand([10, 2, 2], dtype=dtype)  # type: ignore[call-overload]
+            >>> model = MyLayer(dtype)
+            >>> x[0] = float(0)
+            >>> loss = model(x)
+            >>> adam = paddle.optimizer.Adam(parameters=model.parameters())
+            >>> loss.backward()
+            >>> adam.step()
+
+            >>> # error log
+            >>> # [PRECISION] [ERROR] in [device=gpu:0, op=divide, tensor=, dtype=fp32], numel=40, num_nan=0, num_inf=4, num_zero=0, max=inf, min=1.048930e+00, mean=inf
+            >>> # Traceback (most recent call last):
+            >>> #   File "tmp.py", line 16, in <module>
+            >>> #     loss = model(x)
+            >>> #   File "/paddle/nn/layer/layers.py", line 1254, in __call__
+            >>> #     return self.forward(*inputs, **kwargs)
+            >>> #   File "/paddle/amp/debugging.py", line 116, in wrapper
+            >>> #     out_data = func(self, *modified_args, **kwargs)
+            >>> #   File "test.py", line 10, in forward
+            >>> #     return 1/x *  self._w+ self._b
+            >>> # RuntimeError: (PreconditionNotMet) There are NAN or INF (num_nan=0, num_inf=4, num_zero=0) in [device=gpu:0, op=divide, tensor=, dtype=fp32].
+    """
+
+    def wrapper(self, *args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
+        if args:
+            # Set temp data and temp.gradient = False
+            start_data = args[0]
+            if not isinstance(start_data, paddle.Tensor):
+                raise RuntimeError("First input of this layer must be tensor.")
+            start_data.stop_gradient = False
+            modified_args = list(args)  # Convert args to a mutable list
+            # Set FLAGS_check_nan_inf = 1
+            modified_args[0] = _C_ops.enable_check_model_nan_inf(start_data, 1)
+            # Call the forward function
+            out_data = func(self, *modified_args, **kwargs)
+            # Set FLAGS_check_nan_inf = 0
+            out = _C_ops.disable_check_model_nan_inf(out_data, 0)
+            return out
+        else:
+            raise RuntimeError("No elements found in args.")
+        out = func(self, *args, **kwargs)
+        return out
+
+    return wrapper
+
+
+def set_checked_op_list(checked_op_list: Sequence[str] | None) -> None:
+    # check checked_op_list
+    if checked_op_list is not None:
+        if isinstance(checked_op_list, (list, tuple)):
+            check_op_list = ",".join(value for value in checked_op_list)
+            paddle.base.core.set_checked_op_list(check_op_list)
+        else:
+            raise ValueError("checked_op_list must be list or tuple")
+
+
+def set_skipped_op_list(skipped_op_list: Sequence[str] | None) -> None:
+    # check skipped_op_list
+    if skipped_op_list is not None:
+        if isinstance(skipped_op_list, (list, tuple)):
+            skip_op_list = ",".join(value for value in skipped_op_list)
+            paddle.base.core.set_skipped_op_list(skip_op_list)
+        else:
+            raise ValueError("skipped_op_list must be list or tuple")
+
+
+class TensorCheckerConfig:
+    """
+    The purpose of this class is to collect the configuration for checking NaN and Inf values in the tensors of a module or operator. It takes the following arguments:
+
+    Args:
+        enable(bool): Indicating whether to enable the detection of NaN and Inf values in tensors. The default value is False, which means that these tools will not be used.
+
+        debug_mode(DebugMode, optional): A parameter that determines the type of debugging to be used. Default is DebugMode.CHECK_NAN_INF_AND_ABORT.
+
+        output_dir(string|None, optional): The path to store collected data. If this parameter is set to None, the data will be printed to the terminal. Default is None.
+
+        checked_op_list(list|tuple|None, optional): Specifies a list of operators that need to be checked during program execution, for example, checked_op_list=['elementwise_add', 'conv2d'], indicating that the output results of elementwise_add and conv2d should be checked for nan/inf during program execution. Default is None.
+
+        skipped_op_list(list|tuple|None, optional): Specifies a list of operators that do not need to be checked during program execution, for example, skipped_op_list=['elementwise_add', 'conv2d'], indicating that the output results of elementwise_add and conv2d should not be checked for nan/inf during program execution. None is None.
+
+        debug_step(list|tuple|None, optional): A list or tuple used primarily for nan/inf checking during model training. For example, debug_step=[1,5] indicates that nan/inf checking should only be performed on model training iterations 1 to 5. Default is None.
+
+        stack_height_limit(int, optional): An integer value specifying the maximum depth of the call stack. This feature supports printing the call stack at the error location. Currently, only enabling or disabling call stack printing is supported. If you want to print the corresponding C++ call stack when NaN is detected in GPU Kernel, set stack_height_limit to 1, otherwise set it to 0. Default is 1.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> checker_config = paddle.amp.debugging.TensorCheckerConfig(
+            ...     enable=True, debug_mode=paddle.amp.debugging.DebugMode.CHECK_NAN_INF
+            ... )
+            >>> paddle.amp.debugging.enable_tensor_checker(checker_config)
+
+            >>> x = paddle.to_tensor([1, 0, 3], place=paddle.CPUPlace(), dtype='float32', stop_gradient=False)
+            >>> y = paddle.to_tensor([0.2, 0, 0.5], place=paddle.CPUPlace(), dtype='float32')
+            >>> res = paddle.pow(x, y)
+            >>> paddle.autograd.backward(res, retain_graph=True)
+            >>> paddle.amp.debugging.disable_tensor_checker()
+
+            >>> # [PRECISION] [ERROR] in [device=cpu, op=elementwise_pow_grad, tensor=, dtype=fp32], numel=3, num_nan=1, num_inf=0, num_zero=0, max=2.886751e-01, min=2.000000e-01, mean=-nan
+
+            >>> # when DebugMode.CHECK_NAN_INF_AND_ABORT and stack_height_limit = 1
+            >>> # Traceback (most recent call last):
+            >>> #     res = paddle.pow(x, y)
+            >>> #   File "/usr/local/lib/python3.8/dist-packages/paddle/tensor/math.py", line 447, in pow
+            >>> #     return _C_ops.elementwise_pow(x, y)
+
+    """
+
+    # For module debugging
+    current_step_id: int = 0
+
+    enable: bool
+    debug_mode: DebugMode
+    output_dir: str | None
+    checked_op_list: Sequence[str] | None
+    skipped_op_list: Sequence[str] | None
+    debug_step: Sequence[int] | None
+    stack_height_limit: int
+    start_step: int | None
+    end_step: int | None
+    seed: int
+    initial_seed: int
+
+    def __init__(
+        self,
+        enable: bool,
+        debug_mode: DebugMode = DebugMode.CHECK_NAN_INF_AND_ABORT,
+        output_dir: str | None = None,
+        checked_op_list: Sequence[str] | None = None,
+        skipped_op_list: Sequence[str] | None = None,
+        debug_step: Sequence[int] | None = None,
+        stack_height_limit: int = 1,
+    ):
+        self.enable = enable
+        self.debug_mode = debug_mode
+        self.output_dir = output_dir
+
+        self.checked_op_list = checked_op_list
+        self.skipped_op_list = skipped_op_list
+
+        self.debug_step = debug_step
+        self.stack_height_limit = stack_height_limit
+
+        self.start_step = None
+        self.end_step = None
+
+        self.seed = 123
+        self.initial_seed = 123
+
+        # check debug_step
+        if debug_step is not None:
+            if isinstance(debug_step, (tuple, list)):
+                assert (
+                    len(self.debug_step) == 2
+                    and self.debug_step[1] > self.debug_step[0]
+                )
+                self.start_step, self.end_step = self.debug_step
+                self.start_step = max(self.start_step, 0)
+            else:
+                raise ValueError("debug_step must be list or tuple")
+
+        if core.is_compiled_with_cuda():
+            for i in range(core.get_cuda_device_count()):
+                self.initial_seed = core.default_cuda_generator(
+                    i
+                ).initial_seed()
+        elif core.is_compiled_with_xpu():
+            for i in range(core.get_xpu_device_count()):
+                self.initial_seed = core.default_xpu_generator(i).initial_seed()
+
+        self.initial_seed = core.default_cpu_generator().initial_seed()
+
+        # check debug_mode
+        if self.debug_mode.name not in DebugMode.__members__:
+            raise ValueError(
+                "debug_mode in DebugMode",
+                self.debug_mode,
+                DebugMode.__members__,
+            )
+
+        set_checked_op_list(self.checked_op_list)
+
+        set_skipped_op_list(self.skipped_op_list)
+
+        if self.enable:
+            self._set_seed(self.enable)
+
+    def _set_seed(self, flag: int) -> None:
+        if self.initial_seed != self.seed:
+            self.seed = self.initial_seed
+
+        if self.seed > np.iinfo(np.uint32).max or self.seed < 0:
+            print("[Warning: Seed must be between 0 and 2**32 - 1")
+            self.seed = 123
+
+        # get random seed
+        paddle.seed(self.seed)
+        np.random.seed(self.seed)
+        random.seed(self.seed)
+
+        # info
+        print("AMP Debugging TensorCheckerConfig: seed ", self.seed)
+
+        # set cudnn and cpu
+        if core.is_compiled_with_cuda():
+            paddle.set_flags({"FLAGS_cudnn_deterministic": flag})
+            print(
+                "AMP Debugging TensorCheckerConfig: FLAGS_cudnn_deterministic is ",
+                flag,
+            )
+
+    def _set_env(self, check_flag: int) -> None:
+        paddle.set_flags({"FLAGS_check_nan_inf": check_flag})
+        if check_flag:
+            # set debug level
+            paddle.set_flags(
+                {"FLAGS_check_nan_inf_level": self.debug_mode.value}
+            )
+
+            # set output_dir
+            if self.output_dir is not None:
+                paddle.base.core.set_nan_inf_debug_path(self.output_dir)
+
+            # set stack_height_limit
+            if isinstance(self.stack_height_limit, (int)):
+                paddle.base.core.set_nan_inf_stack_limit(
+                    self.stack_height_limit
+                )
+            else:
+                raise ValueError("stack_height_limit must be int")
+
+    def update_and_check_step_id(self) -> bool:
+        if self.enable:
+            if self.start_step is not None and self.end_step is not None:
+                if (
+                    self.start_step > TensorCheckerConfig.current_step_id
+                    or TensorCheckerConfig.current_step_id >= self.end_step
+                ):
+                    return False
+                else:
+                    TensorCheckerConfig.current_step_id += 1
+            return True
+        return False
+
+    def start_check_nan_inf(self) -> None:
+        if self.enable:
+            self._set_env(self.enable)
+
+    def stop_check_nan_inf(self) -> None:
+        self._set_env(False)
+
+
+def check_numerics(
+    tensor: Tensor,
+    op_type: str,
+    var_name: str,
+    debug_mode: DebugMode = DebugMode.CHECK_NAN_INF_AND_ABORT,
+) -> tuple[Tensor, Tensor]:
+    """
+    This function is used to debugging a tensor, finding the number of NaNs, Infs and zeros in the tensor.
+
+    Args:
+        tensor(Tensor): The target tensor to check.
+        op_type(str): The OP or API name which produce the target tensor.
+        var_name(str): The name of target tensor.
+        debug_mode(paddle.amp.debugging.DebugMode, optional): The mode of debugging to be used. Default is DebugMode.CHECK_NAN_INF_AND_ABORT.
+
+    Returns:
+        (Tensor, Tensor): A tuple of tensors containing
+
+        - **stats** (Tensor): Returns the number of NaNs, Infs and zeros of input tensor. The shape is [3] and dtype is int64.
+        - **values** (Tensor): Returns the maximum, minimum and mean value of input tensor. The shape is [3] and dtype is float.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> checker_config = paddle.amp.debugging.TensorCheckerConfig(
+            ...     enable=True, debug_mode=paddle.amp.debugging.DebugMode.CHECK_NAN_INF
+            ... )
+
+            >>> x = paddle.to_tensor([1, 0, 3], place=paddle.CPUPlace(), dtype='float32')
+            >>> y = paddle.to_tensor([0.2, 0, 0.5], place=paddle.CPUPlace(), dtype='float32')
+            >>> res = paddle.pow(x, y)
+            >>> paddle.amp.debugging.check_numerics(res, "pow", "res")
+
+    """
+    stack_height_limit = -1
+    output_dir = ""
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.check_numerics(
+            tensor,
+            op_type,
+            var_name,
+            debug_mode.value,
+            stack_height_limit,
+            output_dir,
+        )
+
+    helper = LayerHelper("check_numerics", **locals())
+
+    stats = helper.create_variable_for_type_inference(dtype="int64")
+    values = helper.create_variable_for_type_inference(dtype="float")
+
+    helper.append_op(
+        type='check_numerics',
+        inputs={
+            'tensor': tensor,
+        },
+        attrs={
+            'op_type': op_type,
+            'var_name': var_name,
+            'check_nan_inf_level': debug_mode.value,
+            'stack_height_limit': stack_height_limit,
+            'output_dir': output_dir,
+        },
+        outputs={'stats': [stats], 'values': [values]},
+    )
+    return stats, values
+
+
+def _get_operator_stats_flag() -> Any:
+    flags = paddle.get_flags(["FLAGS_low_precision_op_list"])
+    return flags["FLAGS_low_precision_op_list"]
+
+
+def _print_operator_stats(op_count_dict: dict[str, str | list[int]]) -> None:
+    """
+    Parse and print the stats of operators, mainly including the calls of
+    dtypes such as different fp32, fp16, bf16 and others.
+
+    Args:
+        op_count_dict(dict): a dict to record the number of calls for different
+            operator and dtype. An example is
+            {'conv2d': '1,0,0,0', 'elementwise_add': '1,0,0,0'} or
+            {'conv2d': [1, 0, 0, 0], 'elementwise_add': [1, 0, 0, 0]}.
+    """
+    print("<{:-^120}>".format(" op list "))
+    total_ops = 0
+    print(
+        "<{:-^40}".format(" Op Name "),
+        "|",
+        "{:-^17}".format(" FP16 Calls "),
+        "|",
+        "{:-^17}".format(" BF16 Calls "),
+        "|",
+        "{:-^17}".format(" FP32 Calls"),
+        "|",
+        "{:-^17}>".format(" Other Calls "),
+    )
+    if op_count_dict is not None and isinstance(op_count_dict, dict):
+        for op_type in sorted(op_count_dict):
+            # fp16, bf16, fp32, other
+            value = op_count_dict[op_type]
+            if isinstance(value, list):
+                called = value
+            elif isinstance(value, str):
+                called = value.split(",")
+            else:
+                raise ValueError(
+                    f"Input {value} is expected to be a list of str, but received {type(value)}."
+                )
+            print(
+                f"  {op_type:<40}|  {called[0]:<17}|  {called[1]:<17}|  {called[2]:<17}|  {called[3]:<17}"
+            )
+            total_ops += 1
+    print("<{:-^120}>\n".format(" op count: " + str(total_ops) + " "))
+
+
+def enable_operator_stats_collection() -> None:
+    """
+    Enable to collect the number of operators for different data types.
+    The statistical data are categorized according to four data types, namely
+    float32, float16, bfloat16 and others. This function is used in pair with
+    the corresponding disable function.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> import paddle
+            >>> paddle.device.set_device('gpu')
+
+            >>> conv = paddle.nn.Conv2D(3, 2, 3)
+            >>> x = paddle.rand([10, 3, 32, 32])
+
+            >>> paddle.amp.debugging.enable_operator_stats_collection()
+            >>> # AMP list including cast, conv2d, elementwise_add, reshape
+            >>> with paddle.amp.auto_cast(enable=True, level='O2'):
+            ...     out = conv(x)
+            >>> # Print to the standard output.
+            >>> paddle.amp.debugging.disable_operator_stats_collection()
+            >>> # <------------------------------------------------------- op list -------------------------------------------------------->
+            >>> # <--------------- Op Name ---------------- | -- FP16 Calls --- | -- BF16 Calls --- | --- FP32 Calls--- | -- Other Calls -->
+            >>> #   cast                                    |  1                |  0                |  2                |  0
+            >>> #   conv2d                                  |  1                |  0                |  0                |  0
+            >>> #   elementwise_add                         |  0                |  0                |  1                |  0
+            >>> #   reshape                                 |  0                |  0                |  1                |  0
+            >>> # <----------------------------------------------------- op count: 4 ------------------------------------------------------>
+
+    """
+    # Clear the previous stats.
+    paddle.base.core.clear_low_precision_op_list()
+    paddle.set_flags({'FLAGS_low_precision_op_list': 1})
+
+
+def disable_operator_stats_collection() -> None:
+    """
+    Disable the collection the number of operators for different data types.
+    This function is used in pair with the corresponding enable function.
+    The statistical data are categorized according to four data types, namely
+    float32, float16, bfloat16 and others, and will be printed after the
+    function call.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> conv = paddle.nn.Conv2D(3, 2, 3)
+            >>> x = paddle.rand([10, 3, 32, 32])
+
+            >>> paddle.amp.debugging.enable_operator_stats_collection()
+            >>> # AMP list including cast, conv2d, elementwise_add, reshape
+            >>> with paddle.amp.auto_cast(enable=True, level='O2'):
+            ...     out = conv(x)
+            >>> # Print to the standard output.
+            >>> paddle.amp.debugging.disable_operator_stats_collection()
+            >>> # <------------------------------------------------------- op list -------------------------------------------------------->
+            >>> # <--------------- Op Name ---------------- | -- FP16 Calls --- | -- BF16 Calls --- | --- FP32 Calls--- | -- Other Calls -->
+            >>> #   cast                                    |  1                |  0                |  2                |  0
+            >>> #   conv2d                                  |  1                |  0                |  0                |  0
+            >>> #   elementwise_add                         |  0                |  0                |  1                |  0
+            >>> #   reshape                                 |  0                |  0                |  1                |  0
+            >>> # <----------------------------------------------------- op count: 4 ------------------------------------------------------>
+
+    """
+    if not _get_operator_stats_flag():
+        return
+
+    op_count_dict = paddle.base.core.get_low_precision_op_list()
+    _print_operator_stats(op_count_dict)
+    paddle.set_flags({'FLAGS_low_precision_op_list': 0})
+
+
+@contextlib.contextmanager
+def collect_operator_stats() -> Generator[None, None, None]:
+    """
+    The context switcher to enable to collect the number of operators for
+    different data types. The statistical data are categorized according
+    to four data types, namely float32, float16, bfloat16 and others, and
+    will be printed when exiting the context.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> conv = paddle.nn.Conv2D(3, 2, 3)
+            >>> x = paddle.rand([10, 3, 32, 32])
+
+            >>> with paddle.amp.debugging.collect_operator_stats():
+            ...     # AMP list including cast, conv2d, elementwise_add, reshape
+            ...     with paddle.amp.auto_cast(enable=True, level='O2'):
+            ...         out = conv(x)
+            >>> # Print to the standard output.
+            >>> # <------------------------------------------------------- op list -------------------------------------------------------->
+            >>> # <--------------- Op Name ---------------- | -- FP16 Calls --- | -- BF16 Calls --- | --- FP32 Calls--- | -- Other Calls -->
+            >>> #   cast                                    |  1                |  0                |  2                |  0
+            >>> #   conv2d                                  |  1                |  0                |  0                |  0
+            >>> #   elementwise_add                         |  0                |  0                |  1                |  0
+            >>> #   reshape                                 |  0                |  0                |  1                |  0
+            >>> # <----------------------------------------------------- op count: 4 ------------------------------------------------------>
+
+    """
+    enable_operator_stats_collection()
+    yield
+    disable_operator_stats_collection()
+
+
+def compare_accuracy(
+    dump_path: str,
+    another_dump_path: str,
+    output_filename: str,
+    loss_scale: float = 1,
+    dump_all_tensors: bool = False,
+) -> None:
+    r"""
+    This is a precision comparison tool that can be used to compare log data of float16 and float32.
+
+    Args:
+        dump_path(str): The path of the running log, such as the log for execution using the float32 data type.
+        another_dump_path(str): the path of another running log ,such as the log for execution using the float16 data type.
+        output_filename(str): the excel file name of compare output.
+        loss_scale(float, optional): the loss_scale during the training phase. Default is 1.
+        dump_all_tensors(bool, optional): dump all tensor, It is currently not support. Default is False.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.base import core
+            >>> try:
+            ...     import xlsxwriter as xlw
+            ... except ImportError:
+            ...     import subprocess
+            ...
+            ...     subprocess.check_call(
+            ...         ['python', '-m', 'pip', 'install', 'xlsxwriter==3.0.9'],
+            ...     )
+            ...     import xlsxwriter as xlw
+            ...
+            ...     if core.is_compiled_with_cuda():
+            ...         paddle.set_flags({"FLAGS_check_nan_inf": 1, "FLAGS_check_nan_inf_level": 3})
+            ...         path = "workerlog_log_dir"
+            ...         paddle.base.core.set_nan_inf_debug_path(path)
+            ...         x = paddle.to_tensor([2, 3, 4, 0], dtype="float32")
+            ...         y = paddle.to_tensor([1, 5, 2, 0], dtype="float32")
+            ...         z1 = x + y
+            ...         out_excel = "compare_accuracy_out_excel.csv"
+            ...         paddle.amp.debugging.compare_accuracy(
+            ...             path,
+            ...             path,
+            ...             out_excel,
+            ...             loss_scale=1,
+            ...             dump_all_tensors=False,
+            ...         )
+    """
+    assert dump_all_tensors is False, "It is currently not supported."
+    paddle.amp.accuracy_compare.compare_accuracy(
+        dump_path,
+        another_dump_path,
+        output_filename,
+        loss_scale,
+        dump_all_tensors=False,
+    )
+
+
+def enable_tensor_checker(checker_config: TensorCheckerConfig) -> None:
+    """
+    The enable_tensor_checker(checker_config) function enables model-level accuracy checking and is used in combination with disables_tensor_checker() to achieve model-level precision checking by checking the output Tensors of all operators within the specified range.
+
+    Args:
+        checker_config(TensorCheckerConfig): Checker_config is to collect the configuration for checking NaN and Inf values in the tensors of a module or operator.
+
+    Note:
+        If disable_tensor_checker() is called before backward(), the gradient operator will not be checked.
+        If disable_tensor_checker() is called before optimizer.step(), the optimizer and other weight update related operators will not be checked.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> checker_config = paddle.amp.debugging.TensorCheckerConfig(
+            ...     enable=True, debug_mode=paddle.amp.debugging.DebugMode.CHECK_NAN_INF
+            ... )
+            >>> paddle.amp.debugging.enable_tensor_checker(checker_config)
+
+            >>> x = paddle.to_tensor([1, 0, 3], place=paddle.CPUPlace(), dtype='float32', stop_gradient=False)
+            >>> y = paddle.to_tensor([0.2, 0, 0.5], place=paddle.CPUPlace(), dtype='float32')
+            >>> res = paddle.pow(x, y)
+            >>> paddle.autograd.backward(res, retain_graph=True)
+            >>> paddle.amp.debugging.disable_tensor_checker()
+            >>> # [PRECISION] [ERROR] in [device=cpu, op=elementwise_pow_grad, tensor=, dtype=fp32], numel=3, num_nan=1, num_inf=0, num_zero=0, max=2.886751e-01, min=2.000000e-01, mean=-nan
+
+            >>> # when DebugMode.CHECK_NAN_INF_AND_ABORT and stack_height_limit = 1
+            >>> # Traceback (most recent call last):
+            >>> #   File "tp.py", line 8, in <module>
+            >>> #     res = paddle.pow(x, y)
+            >>> #   File "/usr/local/lib/python3.8/dist-packages/paddle/tensor/math.py", line 447, in pow
+            >>> #     return _C_ops.elementwise_pow(x, y)
+
+    """
+    if checker_config.update_and_check_step_id():
+        checker_config.start_check_nan_inf()
+    else:
+        checker_config.stop_check_nan_inf()
+
+
+def disable_tensor_checker() -> None:
+    """
+    disable_tensor_checker() is used to disable accuracy checking, and is used together with enable_tensor_checker(config) to achieve model-level precision checking by checking the output Tensors of all operators within the specified range.
+
+    Note:
+        If disable_tensor_checker() is called before backward(), the gradient operator will not be checked;
+        If disable_tensor_checker() is called before optimizer.step(), the optimizer and other weight update related operators will not be checked.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> checker_config = paddle.amp.debugging.TensorCheckerConfig(
+            ...     enable=True, debug_mode=paddle.amp.debugging.DebugMode.CHECK_NAN_INF
+            ... )
+            >>> paddle.amp.debugging.enable_tensor_checker(checker_config)
+
+            >>> x = paddle.to_tensor([1, 0, 3], place=paddle.CPUPlace(), dtype='float32', stop_gradient=False)
+            >>> y = paddle.to_tensor([0.2, 0, 0.5], place=paddle.CPUPlace(), dtype='float32')
+            >>> res = paddle.pow(x, y)
+            >>> paddle.autograd.backward(res, retain_graph=True)
+            >>> paddle.amp.debugging.disable_tensor_checker()
+            >>> # [PRECISION] [ERROR] in [device=cpu, op=elementwise_pow_grad, tensor=, dtype=fp32], numel=3, num_nan=1, num_inf=0, num_zero=0, max=2.886751e-01, min=2.000000e-01, mean=-nan
+
+            >>> # when DebugMode.CHECK_NAN_INF_AND_ABORT and stack_height_limit = 1
+            >>> # Traceback (most recent call last):
+            >>> #      res = paddle.pow(x, y)
+            >>> #    File "/usr/local/lib/python3.8/dist-packages/paddle/tensor/math.py", line 447, in pow
+            >>> #      return _C_ops.elementwise_pow(x, y)
+
+    """
+    paddle.set_flags({"FLAGS_check_nan_inf": 0})

python/paddle/api_tracer/__init__.py

@@ -0,0 +1,21 @@
+# 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 .api_tracer import start_api_tracer, stop_api_tracer
+
+__all__ = [
+    'api_tracer',
+    'start_api_tracer',
+    'stop_api_tracer',
+]

python/paddle/api_tracer/api_tracer.py

@@ -0,0 +1,297 @@
+# 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 inspect
+import math
+import threading
+
+import numpy as np
+import yaml
+
+import paddle
+
+_trace_guard = threading.local()
+_originals = {}  # {api_path: (parent_obj, method_name, original_fn)}
+_hooked_apis = []
+
+
+class HookAPIMap:
+    pass
+
+
+class ConfigDump:
+    def __init__(self):
+        pass
+
+    def open_file(self, path):
+        self.file = open(path, "a+")
+
+    def dump_config(self, api, input_args, input_kwargs, outputs):
+        result = api + "("
+        for value in input_args:
+            tmp = self.dump_item_str(api, value)
+            if tmp == "":
+                return
+            result = result + tmp + ", "
+        for key, value in input_kwargs.items():
+            tmp = self.dump_item_str(api, value)
+            if tmp == "":
+                return
+            result = result + key + "=" + tmp + ", "
+
+        result = result + ")"
+        # self.file.write(") -> ")
+        # if isinstance(outputs, (list, tuple)):
+        #     for output in outputs:
+        #         self.file.write(self.dump_item_str(api, output) + ", ")
+        # else:
+        #     self.file.write(self.dump_item_str(api, outputs) + ", ")
+
+        self.file.write(result)
+        self.file.write("\n")
+        self.file.flush()
+
+    def dump_item_str(self, api, item):
+        type_mapping = {
+            np.int16: int,
+            np.int32: int,
+            np.int64: int,
+            np.float16: float,
+            np.float32: float,
+            np.float64: float,
+            np.integer: int,
+            np.floating: float,
+            np.bool_: bool,
+            np.complexfloating: complex,
+            np.str_: str,
+            np.bytes_: bytes,
+            # np.unicode_: str,
+        }
+        for numpy_type, builtin_type in type_mapping.items():
+            if isinstance(item, numpy_type):
+                item = builtin_type(item)
+                break
+
+        if isinstance(item, paddle.Tensor):
+            result = (
+                "Tensor(" + str(item.shape) + ',"' + str(item.dtype)[7:] + '"'
+            )
+            if item.place.is_cpu_place():
+                result = result + ",place=" + str(item.place)
+            if not item.is_contiguous():
+                result = (
+                    result + ",is_contiguous=False,strides=" + str(item.strides)
+                )
+            return result + ")"
+        elif isinstance(item, paddle.base.core.DataType):
+            return "Dtype(" + str(item)[7:] + ")"
+        elif isinstance(item, paddle.base.core.VarDesc.VarType):
+            return "VarType(" + str(item)[7:] + ")"
+        elif isinstance(item, paddle.base.libpaddle.Place):
+            return str(item)
+        elif isinstance(item, list):
+            result = "list["
+            for sub_item in item:
+                tmp = self.dump_item_str(api, sub_item)
+                if tmp == "":
+                    return ""
+                result = result + tmp + ","
+            result = result + "]"
+            return result
+        elif isinstance(item, tuple):
+            result = "tuple("
+            for sub_item in item:
+                tmp = self.dump_item_str(api, sub_item)
+                if tmp == "":
+                    return ""
+                result = result + tmp + ","
+            result = result + ")"
+            return result
+        elif isinstance(item, slice):
+            start_str = (
+                str(int(item.start.numpy()))
+                if isinstance(item.start, paddle.Tensor)
+                else str(item.start)
+            )
+            stop_str = (
+                str(int(item.stop.numpy()))
+                if isinstance(item.stop, paddle.Tensor)
+                else str(item.stop)
+            )
+            step_str = (
+                str(int(item.step.numpy()))
+                if isinstance(item.step, paddle.Tensor)
+                else str(item.step)
+            )
+            return "slice(" + start_str + "," + stop_str + "," + step_str + ")"
+        elif isinstance(item, complex):
+            return (
+                "complex("
+                + self.dump_item_str(api, item.real)
+                + ","
+                + self.dump_item_str(api, item.imag)
+                + ")"
+            )
+        elif item is None:
+            return "None"
+        elif isinstance(
+            item, (paddle.base.Variable, paddle.base.libpaddle.pir.Value)
+        ):
+            return ""
+        elif item == math.inf:
+            return "math.inf"
+        elif item == -math.inf:
+            return "-math.inf"
+        elif item == math.nan:
+            return "math.nan"
+        elif item == -math.nan:
+            return "-math.nan"
+        elif isinstance(item, (bool, int, float)):
+            return str(item)
+        elif isinstance(item, str):
+            return '"' + item + '"'
+        elif isinstance(item, type):
+            return (
+                "type("
+                + str(item)[str(item).index("'") + 1 : str(item).rindex("'")]
+                + ")"
+            )
+        elif isinstance(item, np.ndarray):
+            return str(item)[1:-1]
+        elif isinstance(item, np.dtype):
+            return "Dtype(" + str(item) + ")"
+        elif item == Ellipsis:
+            return "Ellipsis"
+        else:
+            print(
+                "[api_tracer error] : dump_item_str ",
+                api,
+                ", item = ",
+                item,
+                ", type(item) = ",
+                type(item),
+            )
+            return ""
+
+
+config_dump = ConfigDump()
+
+
+class APITemplate:
+    def __init__(self, api_name):
+        self.api_name = api_name
+
+    def __call__(self, *args, **kwargs):
+        if getattr(_trace_guard, 'in_hook', False):
+            return getattr(HookAPIMap, self.api_name)(*args, **kwargs)
+        _trace_guard.in_hook = True
+        try:
+            output = getattr(HookAPIMap, self.api_name)(*args, **kwargs)
+            try:
+                config_dump.dump_config(self.api_name, args, kwargs, output)
+            except Exception as err:
+                print(
+                    "[api_tracer error] : config_dump.dump_config ",
+                    self.api_name,
+                    str(err),
+                )
+            return output
+        finally:
+            _trace_guard.in_hook = False
+
+
+def wrapped_api(api_name):
+    def api_template(*args, **kwargs):
+        return APITemplate(api_name)(*args, **kwargs)
+
+    return api_template
+
+
+def expand_wildcard(api_pattern):
+    if not api_pattern.endswith('.*'):
+        return [api_pattern]
+    module_path = api_pattern[:-2]
+    try:
+        module = eval(module_path)
+    except Exception as e:
+        print(f"[api_tracer error] : expand_wildcard {api_pattern}, {e}")
+        return []
+    apis = []
+    for name in dir(module):
+        if module_path.endswith('_C_ops'):
+            if name.startswith('__'):
+                continue
+        elif name.startswith('_'):
+            continue
+        try:
+            attr = getattr(module, name)
+            if (
+                callable(attr)
+                and not inspect.isclass(attr)
+                and not inspect.ismodule(attr)
+            ):
+                apis.append(f"{module_path}.{name}")
+        except Exception:
+            pass
+    return apis
+
+
+def start_api_tracer(api_path, save_config_path):
+    print(paddle.__version__)
+    with open(api_path, "r") as f:
+        raw_apis = yaml.safe_load(f).get("apis") or []
+
+    sample_apis = []
+    for api in raw_apis:
+        sample_apis.extend(expand_wildcard(api))
+    sample_apis = list(dict.fromkeys(sample_apis))
+    print(f"[api_tracer] Expanded to {len(sample_apis)} APIs")
+
+    for api in sample_apis:
+        parent_package, method_name = api.rsplit(".", maxsplit=1)
+        try:
+            parent = eval(parent_package)
+            original = getattr(parent, method_name)
+            _originals[api] = (parent, method_name, original)
+            setattr(HookAPIMap, api, original)
+            setattr(parent, method_name, wrapped_api(api))
+            _hooked_apis.append(api)
+        except Exception as err:
+            print(
+                "[api_tracer error] : start_api_tracer ",
+                api,
+                str(err),
+            )
+
+    config_dump.open_file(save_config_path)
+
+
+def stop_api_tracer():
+    for api in _hooked_apis:
+        entry = _originals.pop(api, None)
+        if entry:
+            parent, method_name, original = entry
+            try:
+                setattr(parent, method_name, original)
+            except Exception as err:
+                print(f"[api_tracer error] : stop_api_tracer {api} {err}")
+    _hooked_apis.clear()
+    if (
+        hasattr(config_dump, 'file')
+        and config_dump.file
+        and not config_dump.file.closed
+    ):
+        config_dump.file.flush()
+        config_dump.file.close()
+    print("[api_tracer] Stopped, all hooks removed.")

python/paddle/apy/device/cuda/compile_command_util.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2026 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 ap
+
+
+class CompileCommandGenerator:
+    def __init__(self):
+        self.file_ext = "cu"
+        self.op_type2generate_func = ap.OrderedDict(
+            [
+                ['matmul', self.generate_matmul_compile_command],
+            ]
+        )
+
+    def __call__(self, op_type, tpl_dirname, library_name):
+        return self.op_type2generate_func[op_type](tpl_dirname, library_name)
+
+    def generate_matmul_compile_command(self, tpl_dirname, library_name):
+        matmul_source_dir = f"{tpl_dirname}/matmul"
+
+        compile_cmd = "nvcc -std=c++20 -O3 -Xcompiler=-fPIC -arch=sm_80 --expt-relaxed-constexpr"
+        compile_cmd = compile_cmd + " -I ${AP_CUTLASS_DIR}/include"
+        compile_cmd = compile_cmd + " -I ${AP_CUTLASS_DIR}/tools/util/include"
+        compile_cmd = compile_cmd + " -I " + matmul_source_dir
+        compile_cmd = (
+            compile_cmd
+            + " -DCUTLASS_ENABLE_TENSOR_CORE_MMA=1 -DCUTLASS_DEBUG_TRACE_LEVEL=0"
+        )
+        compile_cmd = (
+            compile_cmd + " -DAP_ENABLE_AUTOTUNE=1 -DAP_ENABLE_DEBUG=0"
+        )
+        compile_cmd = (
+            compile_cmd
+            + f" --shared {library_name}.{self.file_ext} -o lib{library_name}.so"
+        )
+        return compile_cmd

python/paddle/apy/device/dcu/compile_command_util.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2026 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 ap
+
+
+class CompileCommandGenerator:
+    def __init__(self):
+        self.file_ext = "cu"
+        self.op_type2generate_func = ap.OrderedDict(
+            [
+                ['matmul', self.generate_matmul_compile_command],
+            ]
+        )
+
+    def __call__(self, op_type, tpl_dirname, library_name):
+        return self.op_type2generate_func[op_type](tpl_dirname, library_name)
+
+    def generate_matmul_compile_command(self, tpl_dirname, library_name):
+        matmul_source_dir = f"{tpl_dirname}/matmul"
+
+        compile_cmd = (
+            "hipcc -std=c++17 -O3 -fPIC --offload-arch=gfx928 -Wno-return-type"
+        )
+        compile_cmd = compile_cmd + " -I ${AP_CUTLASS_DIR}/include"
+        compile_cmd = compile_cmd + " -I ${AP_CUTLASS_DIR}/tools/util/include"
+        compile_cmd = compile_cmd + " -I " + matmul_source_dir
+        compile_cmd = (
+            compile_cmd + " -DAP_ENABLE_AUTOTUNE=0 -DAP_ENABLE_DEBUG=0"
+        )
+        compile_cmd = (
+            compile_cmd
+            + f" --shared {library_name}.{self.file_ext} -o lib{library_name}.so"
+        )
+        return compile_cmd

python/paddle/apy/matmul_pass/__main__.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.
+
+import matmul_variadic_ptn  # noqa: F401

python/paddle/apy/matmul_pass/abstract_drr.py

@@ -0,0 +1,41 @@
+# 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.
+
+
+class DrrPass:
+    def make_drr_ctx(self):
+        drr_ctx = DrrCtx()  # noqa: F821
+        drr_ctx.set_drr_pass_type(self.drr_pass_type())
+        drr_ctx.init_source_pattern(self.source_pattern)
+        drr_ctx.init_constraint_func(self.constraint)
+        drr_ctx.init_result_pattern(self.result_pattern)
+        return drr_ctx
+
+    def constraint(self, o, t):
+        return True
+
+    def drr_pass_type(self):
+        return "abstract_drr_pass_type"
+
+
+class register_drr_pass:
+    def __init__(self, pass_name, nice):
+        self.pass_name = pass_name
+        self.nice = nice
+
+    def __call__(self, drr_pass_cls):
+        Registry.abstract_drr_pass(  # noqa: F821
+            self.pass_name, self.nice, drr_pass_cls
+        )
+        return drr_pass_cls

python/paddle/apy/matmul_pass/access_topo_drr.py

@@ -0,0 +1,41 @@
+# 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.
+
+
+class DrrPass:
+    def make_drr_ctx(self):
+        drr_ctx = DrrCtx()  # noqa: F821
+        drr_ctx.set_drr_pass_type(self.drr_pass_type())
+        drr_ctx.init_source_pattern(self.source_pattern)
+        drr_ctx.init_constraint_func(self.constraint)
+        drr_ctx.init_result_pattern(self.result_pattern)
+        return drr_ctx
+
+    def constraint(self, o, t):
+        return True
+
+    def drr_pass_type(self):
+        return "access_topo_drr_pass_type"
+
+
+class register_drr_pass:
+    def __init__(self, pass_name, tag):
+        self.pass_name = pass_name
+        self.tag = tag
+
+    def __call__(self, drr_pass_cls):
+        Registry.access_topo_drr_pass(  # noqa: F821
+            self.pass_name, self.tag, drr_pass_cls
+        )
+        return drr_pass_cls

python/paddle/apy/matmul_pass/code_gen_value_util.py

@@ -0,0 +1,29 @@
+# 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.
+
+
+class CodeGenValue:
+    def __init__(self, pir_type, var_name):
+        self.pir_type = pir_type
+        self.var_name = var_name
+        self.const_value = None
+
+    def get_dtype(self):
+        def convert_to_dtype(pir_dtype, shape, data_layout):
+            return pir_dtype.convert_to_dtype()
+
+        return self.pir_type.match(t_dtensor=convert_to_dtype)
+
+    def is_dense_tensor_type(self):
+        return self.pir_type.get_type_name() == "t_dtensor"

python/paddle/apy/matmul_pass/index_code_gen_value_util.py

@@ -0,0 +1,19 @@
+# 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.
+
+
+class IndexCodeGenValue:
+    def __init__(self, iter_var_names):
+        self.iter_var_names = iter_var_names
+        self.const_data = None

python/paddle/apy/matmul_pass/index_drr_pass_util.py

@@ -0,0 +1,33 @@
+# 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 access_topo_drr
+import ap
+import pir
+
+
+class InsertReshapeBeforeYieldPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.yield_op = o.ap_native_op("cf.yield")
+        o.yield_op([t.output], [])
+
+    def result_pattern(self, o, t):
+        t.declare_internal_native_ir_value("reshaped_output")
+        o.reshape_op = o.ap_native_op("cinn_op.reshape")
+        o.reshape_op.shape = lambda o, t: pir.a_array(
+            [pir.a_i32(ap.DataValue.int32("-1"))]
+        )
+        o.reshape_op([t.output], [t.reshaped_output])
+        o.yield_op = o.ap_native_op("cf.yield")
+        o.yield_op([t.reshaped_output], [])

python/paddle/apy/matmul_pass/index_program_translator_util.py

@@ -0,0 +1,137 @@
+# 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 ap
+import index_drr_pass_util
+import ir_tools
+import op_index_translator_util
+
+
+class IndexProgramTranslatorMap:
+    def __init__(
+        self,
+        index_func_unique_id2index_program,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        self.kernel_arg_translator = kernel_arg_translator
+        self.anchor_iter_var_names = anchor_iter_var_names
+        items = index_func_unique_id2index_program.items()
+        self.index_func_unique_id2translator = ap.OrderedDict(
+            ap.map(
+                lambda i: [items[i][0], self.make_translator(i, items[i][1])],
+                range(len(items)),
+            )
+        )
+
+    def get_offset_var_name(
+        self,
+        index_func_unique_id,
+        mut_kernel_arg_id_registry,
+        mut_lir_code_gen_ctx,
+    ):
+        translator = self.index_func_unique_id2translator[index_func_unique_id]
+        ret = translator.translate(
+            mut_kernel_arg_id_registry=mut_kernel_arg_id_registry,
+            mut_lir_code_gen_ctx=mut_lir_code_gen_ctx,
+        )
+        return ret.iter_var_names[0]
+
+    def make_translator(self, program_id, index_program):
+        pass_manager = ir_tools.create_pass_manager()
+        drr_pass = index_drr_pass_util.InsertReshapeBeforeYieldPass()
+        pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(drr_pass)
+        )
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(index_program)
+        return IndexProgramTranslator(
+            index_program,
+            program_id=program_id,
+            kernel_arg_translator=self.kernel_arg_translator,
+            index_op_translator_maker=op_index_translator_util.OpIndexTranslatorFactory(),
+            anchor_iter_var_names=self.anchor_iter_var_names,
+        )
+
+
+class IndexProgramTranslator:
+    def __init__(
+        self,
+        index_program,
+        program_id,
+        kernel_arg_translator,
+        index_op_translator_maker,
+        anchor_iter_var_names,
+    ):
+        self.program_id = program_id
+        self.program_property = index_program.copy_to_const_program_data()
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_op_translator_maker = index_op_translator_maker
+        self.anchor_iter_var_names = anchor_iter_var_names
+        self.ir_value_index2translated_value = ap.MutableList()
+
+        def PushNone(x):
+            self.ir_value_index2translated_value.append(None)
+
+        ap.map(PushNone, self.program_property.values)
+
+    def translate(
+        self,
+        mut_kernel_arg_id_registry,
+        mut_lir_code_gen_ctx,
+    ):
+        def TranslateOp(op_property):
+            self._translate_op(
+                op_property, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+            )
+
+        ap.map(TranslateOp, self.program_property.ops)
+        return self.ir_value_index2translated_value[-1]
+
+    def _translate_op(
+        self, op_property, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        index_op_translator = self.index_op_translator_maker(
+            index_program_id=self.program_id,
+            op_property=op_property,
+            input_properties=ap.map(
+                self._get_value_property, op_property.input_value_indexes
+            ),
+            output_properties=ap.map(
+                self._get_value_property, op_property.output_value_indexes
+            ),
+            kernel_arg_translator=self.kernel_arg_translator,
+            anchor_iter_var_names=self.anchor_iter_var_names,
+        )
+        inputs = ap.map(
+            self._get_translated_value, op_property.input_value_indexes
+        )
+        outputs = index_op_translator(
+            inputs,
+            mut_kernel_arg_id_registry=mut_kernel_arg_id_registry,
+            mut_lir_code_gen_ctx=mut_lir_code_gen_ctx,
+        )
+        ap.map(
+            self._set_translated_value,
+            zip(op_property.output_value_indexes, outputs),
+        )
+
+    def _get_value_property(self, i):
+        return self.program_property.values[i]
+
+    def _get_translated_value(self, i):
+        return self.ir_value_index2translated_value[i]
+
+    def _set_translated_value(self, pair):
+        self.ir_value_index2translated_value[pair[0]] = pair[1]

python/paddle/apy/matmul_pass/kernel_arg_id_util.py

@@ -0,0 +1,88 @@
+# 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 ap
+
+
+class KernelArgIdNameRegistry:
+    def __init__(self, code_gen_ctx, tensor_match_ctx, name_prefix):
+        self.code_gen_ctx = code_gen_ctx
+        self.tensor_match_ctx = tensor_match_ctx
+        self.name_prefix = name_prefix
+        self.generated_kernel_arg_id2unique_name = ap.MutableOrderedDict()
+        self.all_kernel_arg_id2unique_name = ap.MutableOrderedDict()
+        self.in_tensor_data_ptr_seq_no = 0
+        self.out_tensor_data_ptr_seq_no = 0
+        self.dim_expr_seq_no = 0
+
+    def get_or_create_kernel_arg_id_manul_var_name(
+        self, kernel_arg_id, cpp_var_name
+    ):
+        create = lambda: cpp_var_name
+        return self.all_kernel_arg_id2unique_name.get_or_create(
+            kernel_arg_id, create
+        )
+
+    def get_in_tensor_data_ptr_var_name(self, in_ir_value_name):
+        ir_value = getattr(self.tensor_match_ctx, in_ir_value_name)
+        kernel_arg_id = self.code_gen_ctx.in_tensor_data_ptr_kernel_arg_id(
+            ir_value
+        )
+        create = self._get_creator(
+            kernel_arg_id, self._create_in_tensor_data_ptr_var_name
+        )
+        return self.generated_kernel_arg_id2unique_name.get_or_create(
+            kernel_arg_id, create
+        )
+
+    def _get_creator(self, kernel_arg_id, backend_creator):
+        return lambda: self.all_kernel_arg_id2unique_name.get_or_create(
+            kernel_arg_id, backend_creator
+        )
+
+    def _create_in_tensor_data_ptr_var_name(self):
+        name = f"{self.name_prefix}in_ptr_{self.in_tensor_data_ptr_seq_no}"
+        self.in_tensor_data_ptr_seq_no = self.in_tensor_data_ptr_seq_no + 1
+        return name
+
+    def get_out_tensor_data_ptr_var_name(self, out_ir_value_name):
+        ir_value = getattr(self.tensor_match_ctx, out_ir_value_name)
+        kernel_arg_id = self.code_gen_ctx.out_tensor_data_ptr_kernel_arg_id(
+            ir_value
+        )
+        create = self._get_creator(
+            kernel_arg_id, self._create_out_tensor_data_ptr_var_name
+        )
+        return self.generated_kernel_arg_id2unique_name.get_or_create(
+            kernel_arg_id, create
+        )
+
+    def _create_out_tensor_data_ptr_var_name(self):
+        name = f"{self.name_prefix}out_ptr_{self.out_tensor_data_ptr_seq_no}"
+        self.out_tensor_data_ptr_seq_no = self.out_tensor_data_ptr_seq_no + 1
+        return name
+
+    def get_dim_expr_var_name(self, dim_expr):
+        kernel_arg_id = self.code_gen_ctx.dim_expr_kernel_arg_id(dim_expr)
+        create = self._get_creator(
+            kernel_arg_id, self._create_dim_expr_var_name
+        )
+        return self.generated_kernel_arg_id2unique_name.get_or_create(
+            kernel_arg_id, create
+        )
+
+    def _create_dim_expr_var_name(self):
+        name = f"{self.name_prefix}dim_{self.dim_expr_seq_no}"
+        self.dim_expr_seq_no = self.dim_expr_seq_no + 1
+        return name

python/paddle/apy/matmul_pass/kernel_arg_translator_util.py

@@ -0,0 +1,30 @@
+# 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.
+
+
+class KernelArgTranslator:
+    def __init__(self, param_struct_name):
+        self.param_struct_name = param_struct_name
+
+    def get_kernel_arg_name(self, var_name):
+        return var_name
+
+    def get_param_struct_field_name(self, var_name):
+        return var_name
+
+    def get_param_struct_init_name(self, var_name):
+        return f"{self.param_struct_name}.{var_name}"
+
+    def get_use_name(self, var_name):
+        return f"{self.param_struct_name}.{var_name}"

python/paddle/apy/matmul_pass/low_level_ir_code_gen_ctx_util.py

@@ -0,0 +1,59 @@
+# 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 ap
+
+
+class CudaLikeIrCodeGenCtx:
+    def __init__(self, compute_dtype):
+        self.stmts = ap.MutableList()
+        self.dtype2type_name = ap.OrderedDict(
+            [
+                [ap.DataType.float, "float"],
+                [ap.DataType.float16, "half"],
+                [ap.DataType.bfloat16, "nv_bfloat16"],
+                [ap.DataType.int32, "int"],
+                [ap.DataType.int64, "int64_t"],
+            ]
+        )
+        self.compute_dtype = compute_dtype
+        self.compute_dtype_name = self.dtype2type_name[self.compute_dtype]
+        self.type_cast_str_list = [
+            "",
+            f"static_cast<{self.compute_dtype_name}>",
+        ]
+
+    def assign(self, dst, src):
+        self.stmts.append(f"{dst.var_name} = {src.var_name};")
+
+    def let(self, var, val_name):
+        var_dtype_name = self.dtype2type_name[var.get_dtype()]
+        is_same = self.compute_dtype == var.get_dtype()
+        type_name = (
+            f"{var_dtype_name}" if is_same else f"{self.compute_dtype_name}"
+        )
+        type_cast_str = (
+            "" if is_same else f"static_cast<{self.compute_dtype_name}>"
+        )
+        self.stmts.append(
+            f"{type_name} {var.var_name} = {type_cast_str}({val_name});"
+        )
+
+    def store(self, dtype, dst, offset_var_name, src):
+        is_same = dtype == self.dtype2type_name[self.compute_dtype]
+        type_cast_str = "" if is_same else f"static_cast<{dtype}>"
+        self.stmts.append(f"{dst}[{offset_var_name}] = {type_cast_str}({src});")
+
+    def get_stmts_joined_str(self, indent):
+        return f"\n{indent}".join([*self.stmts])

python/paddle/apy/matmul_pass/matmul/.apy_ignore

@@ -0,0 +1,3 @@
+cutlass
+generate_configs.py
+hytlass

python/paddle/apy/matmul_pass/matmul/cutlass_patch/backend.h

@@ -0,0 +1,24 @@
+// Copyright (c) 2026 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.
+
+#pragma once
+
+#ifdef __NVCC__
+#include "cutlass/cutlass.h"
+#define GPUStream_t cudaStream_t
+#elif defined(__HIPCC__)
+#include "hytlass/hytlass.h"
+namespace cutlass = hytlass;
+#define GPUStream_t hipStream_t
+#endif

python/paddle/apy/matmul_pass/matmul/cutlass_patch/batched_matrix_coord.h

@@ -0,0 +1,33 @@
+// 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.
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+namespace cutlass_patch {
+
+struct BatchedMatrixCoord {
+  int batch;
+  int row;
+  int column;
+
+  CUTLASS_HOST_DEVICE
+  BatchedMatrixCoord() : batch(0), row(0), column(0) {}
+
+  CUTLASS_HOST_DEVICE
+  BatchedMatrixCoord(int b, int r, int c) : batch(b), row(r), column(c) {}
+};
+
+};  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/cuda/all_tuning_configs.h

@@ -0,0 +1,496 @@
+// 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.
+
+// auto-generated by generate_configs.py
+
+#pragma once
+
+#include "cutlass/gemm_coord.h"
+
+namespace ap {
+
+constexpr int kNumConfigsHalf = 23;
+constexpr int kNumConfigsFloat = 13;
+
+template <int SwizzleFactor, bool Batched>
+struct SwizzleWrapper {
+  using Type =
+      cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<SwizzleFactor>;
+};
+
+// template <int SwizzleFactor>
+// struct SwizzleWrapper<SwizzleFactor, true> {
+//   using Type =
+//       cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle;
+// };
+
+#define AP_AUTOTUNE(func, stream_ptr, count, ...)                             \
+  {                                                                           \
+    using FuncType = decltype(func<0>);                                       \
+    static int selected_config_id = -1;                                       \
+    static std::vector<std::function<FuncType>> matmul_functions =            \
+        []<std::size_t... Is>(std::index_sequence<Is...>) {                   \
+      return std::vector<std::function<FuncType>>{func<Is>...};               \
+    }                                                                         \
+    (std::make_index_sequence<count>());                                      \
+                                                                              \
+    if (selected_config_id == -1) {                                           \
+      selected_config_id =                                                    \
+          ap::ProfileBestConfig(matmul_functions, stream_ptr, ##__VA_ARGS__); \
+    }                                                                         \
+                                                                              \
+    matmul_functions[selected_config_id](__VA_ARGS__);                        \
+  }
+
+#define AP_AUTOTUNE_half(func, stream_ptr, ...) \
+  AP_AUTOTUNE(func, stream_ptr, ap::kNumConfigsHalf, __VA_ARGS__)
+#define AP_AUTOTUNE_float(func, stream_ptr, ...) \
+  AP_AUTOTUNE(func, stream_ptr, ap::kNumConfigsFloat, __VA_ARGS__)
+#define AP_AUTOTUNE_bfloat16(func, stream_ptr, ...) \
+  AP_AUTOTUNE_half(func, stream_ptr, __VA_ARGS__)
+
+template <typename ElementT, int SwizzleFactor, bool Batched, int Id = 0>
+struct GemmTuningConfigs {
+  using TShape = cutlass::gemm::GemmShape<256, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 2;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = Id;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 1> {
+  using TShape = cutlass::gemm::GemmShape<64, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 1;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 2> {
+  using TShape = cutlass::gemm::GemmShape<64, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 2;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 3> {
+  using TShape = cutlass::gemm::GemmShape<128, 64, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 32, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 3;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 4> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 4;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 5> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 5;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 6> {
+  using TShape = cutlass::gemm::GemmShape<256, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 6;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 7> {
+  using TShape = cutlass::gemm::GemmShape<256, 64, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 7;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 8> {
+  using TShape = cutlass::gemm::GemmShape<256, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 8;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 9> {
+  using TShape = cutlass::gemm::GemmShape<256, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 9;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 10> {
+  using TShape = cutlass::gemm::GemmShape<128, 32, 64>;
+  using WShape = cutlass::gemm::GemmShape<32, 32, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 10;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 11> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 11;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 12> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 12;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 13> {
+  using TShape = cutlass::gemm::GemmShape<256, 64, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 13;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 14> {
+  using TShape = cutlass::gemm::GemmShape<256, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 14;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 15> {
+  using TShape = cutlass::gemm::GemmShape<32, 64, 64>;
+  using WShape = cutlass::gemm::GemmShape<16, 32, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 15;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 16> {
+  using TShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using WShape = cutlass::gemm::GemmShape<32, 32, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 16;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 17> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 17;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 18> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 18;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 19> {
+  using TShape = cutlass::gemm::GemmShape<64, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 6;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 19;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 20> {
+  using TShape = cutlass::gemm::GemmShape<128, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 32, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 6;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 20;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 21> {
+  using TShape = cutlass::gemm::GemmShape<128, 32, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 32, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 7;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 21;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 22> {
+  using TShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 32, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 10;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 22;
+};
+
+// Specialization for float
+template <int SwizzleFactor, bool Batched, int Id>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, Id> {
+  using TShape = cutlass::gemm::GemmShape<64, 64, 16>;
+  using WShape = cutlass::gemm::GemmShape<32, 32, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = Id;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 1> {
+  using TShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 32, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 1;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 2> {
+  using TShape = cutlass::gemm::GemmShape<64, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 2;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 3> {
+  using TShape = cutlass::gemm::GemmShape<64, 256, 16>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 3;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 4> {
+  using TShape = cutlass::gemm::GemmShape<64, 256, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 4;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 5> {
+  using TShape = cutlass::gemm::GemmShape<128, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 32, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 5;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 6> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 16>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 6;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 7> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 7;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 8> {
+  using TShape = cutlass::gemm::GemmShape<256, 64, 16>;
+  using WShape = cutlass::gemm::GemmShape<64, 32, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 8;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 9> {
+  using TShape = cutlass::gemm::GemmShape<256, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 32, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 9;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 10> {
+  using TShape = cutlass::gemm::GemmShape<64, 128, 16>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 10;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 11> {
+  using TShape = cutlass::gemm::GemmShape<128, 64, 16>;
+  using WShape = cutlass::gemm::GemmShape<64, 32, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 11;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 12> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 16>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 12;
+};
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/cutlass_patch/cuda/cutlass_matmul.cuh

@@ -0,0 +1,273 @@
+// 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.
+
+#pragma once
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+#include "cutlass/cutlass.h"
+#include "cutlass/gemm_coord.h"
+#include "cutlass/layout/matrix.h"
+
+#include "cutlass/epilogue/thread/linear_combination_bias_elementwise.h"
+#include "cutlass/util/device_memory.h"
+
+#include "cutlass/gemm/device/gemm_universal.h"
+#include "cutlass/gemm/device/gemm_universal_with_broadcast.h"
+
+#include "cutlass_patch/batched_matrix_coord.h"
+#include "cutlass_patch/cuda/default_config_id.h"
+#include "cutlass_patch/epilogue/thread/linear_combination_unary.h"
+#include "cutlass_patch/epilogue/thread/linear_combination_variadic.h"
+#include "cutlass_patch/gemm/device/gemm_universal_with_variadic.h"
+
+#include "params.h"  // NOLINT
+
+#define CHECK_CUTLASS(status)                                             \
+  {                                                                       \
+    cutlass::Status error = status;                                       \
+    if (error != cutlass::Status::kSuccess) {                             \
+      std::cerr << "Got cutlass error: " << cutlassGetStatusString(error) \
+                << " at: " << __LINE__ << std::endl;                      \
+      exit(EXIT_FAILURE);                                                 \
+    }                                                                     \
+  }
+
+namespace ap {
+using bfloat16 = nv_bfloat16;
+
+template <typename T, int N>
+using Array = cutlass::Array<T, N>;
+
+using MatrixCoord = cutlass_patch::BatchedMatrixCoord;
+
+// Convert CUDA data type to cutlass data type
+template <typename T>
+struct CutlassDataType {
+  using Type = T;
+};
+
+template <>
+struct CutlassDataType<half> {
+  using Type = cutlass::half_t;
+};
+
+template <>
+struct CutlassDataType<__nv_bfloat16> {
+  using Type = cutlass::bfloat16_t;
+};
+
+// Convert to cutlass layout
+template <bool Transposed>
+struct MatrixLayout {
+  using Type = cutlass::layout::RowMajor;
+};
+
+template <>
+struct MatrixLayout<true> {
+  using Type = cutlass::layout::ColumnMajor;
+};
+
+// Operation performed by GEMM
+template <typename ElementT>
+struct GemmOperation {
+  using Type = cutlass::arch::OpMultiplyAdd;
+};
+
+template <>
+struct GemmOperation<float> {
+  using Type = cutlass::arch::OpMultiplyAddFastF32;
+};
+
+static cutlass::gemm::GemmUniversalMode GetGemmMode(int batch_count) {
+  return batch_count > 1 ? cutlass::gemm::GemmUniversalMode::kBatched
+                         : cutlass::gemm::GemmUniversalMode::kGemm;
+}
+
+static void *GetWorkspace(size_t workspace_size) {
+  static cutlass::device_memory::allocation<uint8_t> workspace;
+  if (workspace.size() < workspace_size) {
+    workspace.reset(workspace_size);
+  }
+  return workspace.get();
+}
+
+template <typename GemmFunc>
+cutlass::Status SetMaxDynamicSharedMemorySize() {
+  cudaError_t cudart_result;
+
+  // If requires more than 48KB: configure for extended, dynamic shared memory
+  if constexpr (GemmFunc::kSharedStorageSize >= (48 << 10)) {
+    cudart_result =
+        cudaFuncSetAttribute(cutlass::Kernel2<typename GemmFunc::GemmKernel>,
+                             cudaFuncAttributeMaxDynamicSharedMemorySize,
+                             GemmFunc::kSharedStorageSize);
+    if (cudart_result != cudaSuccess) {
+      CUTLASS_TRACE_HOST("cudaFuncSetAttribute() returned error "
+                         << cudaGetErrorString(cudart_result));
+      return cutlass::Status::kErrorInternal;
+    }
+  }
+
+#if AP_ENABLE_DEBUG
+  // Update SM occupancy member
+  int sm_occupancy = -1;
+  cudart_result = cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(
+      &sm_occupancy,
+      cutlass::Kernel2<typename GemmFunc::GemmKernel>,
+      GemmFunc::GemmKernel::kThreadCount,
+      GemmFunc::kSharedStorageSize,
+      cudaOccupancyDisableCachingOverride);
+  if (cudart_result != cudaSuccess) {
+    CUTLASS_TRACE_HOST(
+        "cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags() returned "
+        "error "
+        << cudaGetErrorString(cudart_result));
+    return cutlass::Status::kErrorInternal;
+  }
+  CUTLASS_TRACE_HOST("sm_occupancy: (" << sm_occupancy
+                                       << ") "
+                                          "smem_size: ("
+                                       << GemmFunc::kSharedStorageSize
+                                       << ") "
+                                          "GemmKernel::kThreadCount: ("
+                                       << GemmFunc::GemmKernel::kThreadCount
+                                       << ")");
+#endif
+  return cutlass::Status::kSuccess;
+}
+
+template <typename ElementT,
+          typename ElementComputeT,
+          template <typename T>
+          class VariadicFunctor,
+          int AlignA = 128 / cutlass::sizeof_bits<ElementT>::value,
+          int AlignB = 128 / cutlass::sizeof_bits<ElementT>::value,
+          int ConfigId = DefaultConfig::kConfigId,
+          int SwizzleFactor = DefaultConfig::kSwizzleFactor,
+          bool Batched = DefaultConfig::kBatched>
+void MatmulAddVariadic(
+    const GemmEpilogueParams &params,
+    const typename VariadicFunctor<ElementComputeT>::Arguments &variadic_args) {
+  using ElementAccumulator =
+      typename CutlassDataType<ElementComputeT>::Type;  // <- data type of
+                                                        // accumulator
+  using ElementComputeEpilogue =
+      ElementAccumulator;  // <- data type of epilogue operations
+  using ElementInputA =
+      typename CutlassDataType<ElementT>::Type;  // <- data type of elements in
+                                                 // input matrix A
+  using ElementInputB =
+      typename CutlassDataType<ElementT>::Type;  // <- data type of elements in
+                                                 // input matrix B
+  using ElementOutput =
+      typename CutlassDataType<ElementT>::Type;  // <- data type of elements in
+                                                 // output matrix D
+
+  constexpr int AlignC = AlignB;
+
+  // Epilogue operation as LinearCombination:
+  //  alpha * accumulator + beta * source
+  using EpilogueOutputOp =
+      cutlass_patch::epilogue::thread::LinearCombinationVariadic<
+          VariadicFunctor,
+          ElementOutput,
+          AlignC,
+          ElementAccumulator,
+          ElementComputeEpilogue,
+          cutlass::epilogue::thread::ScaleType::NoBetaScaling>;  // <- alpha x
+                                                                 // AB + bias
+
+  using GemmFunc = cutlass_patch::gemm::device::GemmUniversalWithVariadic<
+      ElementInputA,
+      cutlass::layout::RowMajor,
+      ElementInputB,
+      cutlass::layout::RowMajor,
+      ElementOutput,
+      cutlass::layout::RowMajor,
+      ElementAccumulator,
+      cutlass::arch::OpClassTensorOp,
+      cutlass::arch::Sm80,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          TShape,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          WShape,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          IShape,
+      EpilogueOutputOp,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          SwizzleThreadBlock,
+      GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::kNumStages,
+      AlignA,
+      AlignB,
+      typename GemmOperation<ElementT>::Type>;
+
+  CHECK_CUTLASS(SetMaxDynamicSharedMemorySize<GemmFunc>());
+
+  /// Arguments
+  cutlass::gemm::GemmCoord problem_size{params.m, params.n, params.k};
+
+  const ElementInputA *input =
+      reinterpret_cast<const ElementInputA *>(params.input);
+  const ElementInputB *weight =
+      reinterpret_cast<const ElementInputB *>(params.weight);
+  const ElementOutput *bias =
+      reinterpret_cast<const ElementOutput *>(params.bias);
+  ElementOutput *output = reinterpret_cast<ElementOutput *>(params.output);
+
+  ElementComputeEpilogue alpha = static_cast<ElementComputeEpilogue>(1);
+  ElementComputeEpilogue beta = bias ? static_cast<ElementComputeEpilogue>(1)
+                                     : static_cast<ElementComputeEpilogue>(0);
+
+  typename GemmFunc::Arguments arguments{
+      GetGemmMode(params.batch_count),
+      problem_size,        // <- problem size of matrix multiplication
+      params.batch_count,  // <- batch_count or k-dimension split factor
+      {alpha, beta, variadic_args},  // <- epilogue params, alpha, beta
+      input,                         // <- input, ptr_A, A, shape={M, K}
+      weight,                        // <- input, ptr_B, B, shape={K, N}
+      bias,                          // <- input, ptr_C, shape={M, N} or {1, N}
+      output,                        // <- output, ptr_D, Z, shape={M, N}
+      params.shape_args.batch_stride_A,
+      params.shape_args.batch_stride_B,
+      params.shape_args.batch_stride_C,
+      params.shape_args.batch_stride_D,
+      params.shape_args.lda,
+      params.shape_args.ldb,
+      params.shape_args.ldc_bias,
+      params.shape_args.ldd};
+
+  size_t workspace_size = GemmFunc::get_workspace_size(arguments);
+  void *workspace = workspace_size > 0 ? GetWorkspace(workspace_size) : nullptr;
+
+  GemmFunc device_gemm;
+
+  cudaStream_t *stream_ptr =
+      reinterpret_cast<cudaStream_t *>(params.stream_ptr);
+
+  CHECK_CUTLASS(device_gemm.can_implement(arguments));
+  CHECK_CUTLASS(device_gemm.initialize(arguments, workspace, *stream_ptr));
+
+  //
+  // Run the GEMM
+  //
+  CHECK_CUTLASS(device_gemm(*stream_ptr));
+#if AP_ENABLE_DEBUG
+  CHECK_CUDA(cudaStreamSynchronize(*stream_ptr));
+#endif
+}
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/cutlass_patch/cuda/default_config_id.h

@@ -0,0 +1,27 @@
+// 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.
+
+#pragma once
+
+#include "all_tuning_configs.h"  // NOLINT
+
+namespace ap {
+
+struct DefaultConfig {
+  static constexpr int kConfigId = 0;
+  static constexpr int kSwizzleFactor = 1;
+  static constexpr bool kBatched = false;
+};
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/cutlass_patch/epilogue/thread/linear_combination_unary.h

@@ -0,0 +1,302 @@
+// 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.
+
+/*! \file
+  \brief Functor performing linear combination operations used by epilogues.
+  It is modified from LinearCombinationGeneric.
+*/
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+#ifdef __NVCC__
+#include "cutlass/array.h"
+#include "cutlass/epilogue/thread/scale_type.h"
+#include "cutlass/functional.h"
+#include "cutlass/numeric_conversion.h"
+#include "cutlass/numeric_types.h"
+#elif defined(__HIPCC__)
+#include "hytlass/array.h"
+#include "hytlass/epilogue/thread/scale_type.h"
+#include "hytlass/functional.h"
+#include "hytlass/numeric_conversion.h"
+#include "hytlass/numeric_types.h"
+#endif
+
+namespace cutlass_patch {
+namespace epilogue {
+namespace thread {
+
+template <class UnaryOp, class = void>
+struct GenericUnaryTraits {
+  static constexpr bool IsArgumentsNeeded = false;
+  struct Arguments {};
+};
+
+template <class UnaryOp>
+struct GenericUnaryTraits<UnaryOp,
+                          decltype(typename UnaryOp::Arguments(), void())> {
+  static constexpr bool IsArgumentsNeeded = true;
+  using Arguments = typename UnaryOp::Arguments;
+};
+
+/// Applies a linear combination operator followed by an unary function to an
+/// array of elements.
+///
+/// D = unary_op(alpha * accumulator + beta * source)
+///
+template <
+    template <typename T>
+    class UnaryOp,
+    typename ElementOutput_,  ///< Data type used to load and store tensors
+    int ElementsPerAccess,    ///< Number of elements computed per operation
+                            ///< Usually it is 128/sizeof_bits<ElementOutput_>,
+                            ///< but we use 64 or 32 sometimes when there are
+                            ///< not enough data to store
+    typename ElementAccumulator_ = ElementOutput_,  ///< Accumulator data type
+    typename ElementCompute_ =
+        ElementOutput_,  ///< Data type used to compute linear combination
+    cutlass::epilogue::thread::ScaleType::Kind Scale =
+        cutlass::epilogue::thread::ScaleType::Default,  ///< Control Alpha and
+                                                        ///< Beta scaling
+    cutlass::FloatRoundStyle Round = cutlass::FloatRoundStyle::round_to_nearest,
+    bool IsHeavy = false>
+class LinearCombinationUnary {
+ public:
+  using ElementOutput = ElementOutput_;
+  using ElementAccumulator = ElementAccumulator_;
+  using ElementCompute = ElementCompute_;
+  using UnaryArguments =
+      typename GenericUnaryTraits<UnaryOp<ElementCompute>>::Arguments;
+
+  static bool const kIsHeavy = IsHeavy;
+  static int const kElementsPerAccess = ElementsPerAccess;
+  static int const kCount = ElementsPerAccess;
+  static const cutlass::epilogue::thread::ScaleType::Kind kScale = Scale;
+
+  using FragmentOutput = cutlass::Array<ElementOutput, kElementsPerAccess>;
+  using FragmentAccumulator =
+      cutlass::Array<ElementAccumulator, kElementsPerAccess>;
+  using FragmentSource = cutlass::Array<ElementOutput, kElementsPerAccess>;
+  using FragmentCompute = cutlass::Array<ElementCompute, kElementsPerAccess>;
+
+  static cutlass::FloatRoundStyle const kRound = Round;
+
+  /// Host-constructable parameters structure
+  struct Params {
+    ElementCompute alpha;             ///< scales accumulators
+    ElementCompute beta;              ///< scales source tensor
+    ElementCompute const *alpha_ptr;  ///< pointer to accumulator scalar - if
+                                      ///< not null, loads it from memory
+    ElementCompute const *beta_ptr;   ///< pointer to source scalar - if not
+                                      ///< null, loads it from memory
+    UnaryArguments unary_args;
+
+    //
+    // Methods
+    //
+
+    CUTLASS_HOST_DEVICE
+    Params()
+        : alpha(ElementCompute(1)),
+          beta(ElementCompute(0)),
+          alpha_ptr(nullptr),
+          beta_ptr(nullptr) {}
+
+    CUTLASS_HOST_DEVICE
+    Params(ElementCompute alpha,
+           ElementCompute beta = ElementCompute(0),
+           UnaryArguments unary_args_ = UnaryArguments{})
+        : alpha(alpha),
+          beta(beta),
+          alpha_ptr(nullptr),
+          beta_ptr(nullptr),
+          unary_args(unary_args_) {}
+
+    CUTLASS_HOST_DEVICE
+    Params(ElementCompute const *alpha_ptr,
+           ElementCompute const *beta_ptr = nullptr)
+        : alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {}
+  };
+
+ private:
+  //
+  // Data members
+  //
+
+  Params params_;
+  bool skip_elementwise_;
+
+ public:
+  /// Constructs the function object, possibly loading from pointers in host
+  /// memory
+  CUTLASS_HOST_DEVICE
+  explicit LinearCombinationUnary(Params const &params) {
+    params_ = params;
+    params_.alpha = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
+    params_.beta = (params.beta_ptr ? *params.beta_ptr : params.beta);
+    skip_elementwise_ = false;
+  }
+
+  /// Returns true if source is needed
+  CUTLASS_HOST_DEVICE
+  bool is_source_needed() const {
+    if (Scale == cutlass::epilogue::thread::ScaleType::NoBetaScaling)
+      return params_.beta != ElementCompute(0);
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling)
+      return false;
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) return false;
+
+    return params_.beta != ElementCompute(0);
+  }
+
+  /// Functionally required for serial reduction in the epilogue
+  CUTLASS_HOST_DEVICE
+  void set_k_partition(int k_partition, int k_partition_count) {
+    if (k_partition) {
+      params_.beta = ElementCompute(1);
+    }
+
+    if (k_partition != k_partition_count - 1) {
+      skip_elementwise_ = true;
+    }
+  }
+
+  /// Computes linear scaling: D = alpha * accumulator + beta * source
+  CUTLASS_HOST_DEVICE
+  FragmentOutput operator()(FragmentAccumulator const &accumulator,
+                            FragmentOutput const &source) const {
+    // Convert source to internal compute numeric type
+    cutlass::NumericArrayConverter<ElementCompute,
+                                   ElementOutput,
+                                   kElementsPerAccess,
+                                   Round>
+        source_converter;
+    cutlass::NumericArrayConverter<ElementCompute,
+                                   ElementAccumulator,
+                                   kElementsPerAccess,
+                                   Round>
+        accumulator_converter;
+
+    FragmentCompute converted_source = source_converter(source);
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+
+    cutlass::multiplies<FragmentCompute> mul_add_source;
+    cutlass::multiply_add<FragmentCompute> mul_add_accumulator;
+    UnaryOp<ElementCompute> unary_op;
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::NoBetaScaling) {
+      intermediate = converted_source;
+      // D = alpha * Accum + X
+      intermediate = mul_add_accumulator(
+          params_.alpha, converted_accumulator, intermediate);
+    } else if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) {
+      intermediate = converted_accumulator;
+    } else {
+      // X =  beta * C + uniform
+      intermediate = mul_add_source(params_.beta, converted_source);
+      // D = alpha * Accum + X
+      intermediate = mul_add_accumulator(
+          params_.alpha, converted_accumulator, intermediate);
+    }
+
+    if constexpr (GenericUnaryTraits<
+                      UnaryOp<ElementCompute>>::IsArgumentsNeeded) {
+      if (!skip_elementwise_) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = unary_op(intermediate[i], params_.unary_args);
+        }
+      }
+    } else {
+      if (!skip_elementwise_) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = unary_op(intermediate[i]);
+        }
+      }
+    }
+
+    // Convert to destination numeric type
+    cutlass::NumericArrayConverter<ElementOutput,
+                                   ElementCompute,
+                                   kElementsPerAccess,
+                                   Round>
+        destination_converter;
+
+    return destination_converter(intermediate);
+  }
+
+  /// Computes linear scaling: D = alpha * accumulator
+  CUTLASS_HOST_DEVICE
+  FragmentOutput operator()(FragmentAccumulator const &accumulator) const {
+    // Convert source to internal compute numeric type
+    cutlass::NumericArrayConverter<ElementCompute,
+                                   ElementAccumulator,
+                                   kElementsPerAccess,
+                                   Round>
+        accumulator_converter;
+
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+
+    cutlass::multiplies<FragmentCompute> mul_add_accumulator;
+    UnaryOp<ElementCompute> unary_op;
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) {
+      intermediate = converted_accumulator;
+    } else {
+      // D = alpha * Accum
+      intermediate = mul_add_accumulator(params_.alpha, converted_accumulator);
+    }
+
+    if constexpr (GenericUnaryTraits<
+                      UnaryOp<FragmentCompute>>::IsArgumentsNeeded) {
+      if (!skip_elementwise_) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = unary_op(intermediate[i], params_.unary_args);
+        }
+      }
+    } else {
+      if (!skip_elementwise_) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = unary_op(intermediate[i]);
+        }
+      }
+    }
+
+    // Convert to destination numeric type
+    cutlass::NumericArrayConverter<ElementOutput,
+                                   ElementCompute,
+                                   kElementsPerAccess,
+                                   Round>
+        destination_converter;
+
+    return destination_converter(intermediate);
+  }
+};
+
+}  // namespace thread
+}  // namespace epilogue
+}  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/epilogue/thread/linear_combination_variadic.h

@@ -0,0 +1,337 @@
+// 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.
+
+/*! \file
+  \brief Functor performing linear combination operations used by epilogues.
+*/
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+#ifdef __NVCC__
+#include "cutlass/array.h"
+#include "cutlass/epilogue/thread/scale_type.h"
+#include "cutlass/functional.h"
+#include "cutlass/numeric_conversion.h"
+#include "cutlass/numeric_types.h"
+#elif defined(__HIPCC__)
+#include "hytlass/array.h"
+#include "hytlass/epilogue/thread/scale_type.h"
+#include "hytlass/functional.h"
+#include "hytlass/numeric_conversion.h"
+#include "hytlass/numeric_types.h"
+#endif
+
+#include "cutlass_patch/batched_matrix_coord.h"
+#include "cutlass_patch/trace_device.h"
+
+namespace cutlass_patch {
+namespace epilogue {
+namespace thread {
+
+template <class VariadicOp, class = void>
+struct GenericVariadicTraits {
+  static constexpr bool IsArgumentsNeeded = false;
+  struct Arguments {};
+};
+
+template <class VariadicOp>
+struct GenericVariadicTraits<VariadicOp,
+                             decltype(typename VariadicOp::Arguments(),
+                                      void())> {
+  static constexpr bool IsArgumentsNeeded = true;
+  using Arguments = typename VariadicOp::Arguments;
+};
+
+/// Applies a linear combination operator to an array of elements.
+///
+/// D = VariadicOp(alpha * accumulator + beta * source)
+///
+template <
+    template <typename T>
+    class VariadicOp,
+    typename ElementOutput_,  ///< Data type used to load and store tensors
+    int ElementsPerAccess,    ///< Number of elements computed per operation.
+                            ///< Usually it is 128/sizeof_bits<ElementOutput_>,
+                            ///< but we use 64 or 32 sometimes when there are
+                            ///< not enough data to store
+    typename ElementAccumulator_ = ElementOutput_,  ///< Accumulator data type
+    typename ElementCompute_ =
+        ElementOutput_,  ///< Data type used to compute linear combination
+    cutlass::epilogue::thread::ScaleType::Kind Scale =
+        cutlass::epilogue::thread::ScaleType::Default,  ///< Control Alpha and
+                                                        ///< Beta scaling
+    cutlass::FloatRoundStyle Round = cutlass::FloatRoundStyle::round_to_nearest,
+    bool IsHeavy = false>
+class LinearCombinationVariadic {
+ public:
+  using ElementOutput = ElementOutput_;
+  using ElementAccumulator = ElementAccumulator_;
+  using ElementCompute = ElementCompute_;
+  using VariadicArguments =
+      typename GenericVariadicTraits<VariadicOp<ElementCompute>>::Arguments;
+
+  static bool const kIsHeavy = IsHeavy;
+  static int const kElementsPerAccess = ElementsPerAccess;
+  static int const kCount = ElementsPerAccess;
+  static const cutlass::epilogue::thread::ScaleType::Kind kScale = Scale;
+
+  using FragmentOutput = cutlass::Array<ElementOutput, kElementsPerAccess>;
+  using FragmentAccumulator =
+      cutlass::Array<ElementAccumulator, kElementsPerAccess>;
+  using FragmentSource = cutlass::Array<ElementOutput, kElementsPerAccess>;
+  using FragmentCompute = cutlass::Array<ElementCompute, kElementsPerAccess>;
+
+  static cutlass::FloatRoundStyle const kRound = Round;
+
+  /// Host-constructable parameters structure
+  struct Params {
+    ElementCompute alpha;             ///< scales accumulators
+    ElementCompute beta;              ///< scales source tensor
+    ElementCompute const *alpha_ptr;  ///< pointer to accumulator scalar - if
+                                      ///< not null, loads it from memory
+    ElementCompute const *beta_ptr;   ///< pointer to source scalar - if not
+                                      ///< null, loads it from memory
+    VariadicArguments variadic_args;
+
+    CUTLASS_HOST_DEVICE
+    Params()
+        : alpha(ElementCompute(1)),
+          beta(ElementCompute(0)),
+          alpha_ptr(nullptr),
+          beta_ptr(nullptr) {}
+
+    CUTLASS_HOST_DEVICE
+    Params(ElementCompute alpha,
+           ElementCompute beta,
+           VariadicArguments variadic_args_ = VariadicArguments{})
+        : alpha(alpha),
+          beta(beta),
+          alpha_ptr(nullptr),
+          beta_ptr(nullptr),
+          variadic_args(variadic_args_) {}
+  };
+
+ private:
+  //
+  // Data members
+  //
+
+  Params params_;
+  bool skip_elementwise_;
+
+ public:
+  /// Constructs the function object, possibly loading from pointers in host
+  /// memory
+  CUTLASS_HOST_DEVICE
+  LinearCombinationVariadic(Params const &params) {
+    params_ = params;
+    params_.alpha = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
+    params_.beta = (params.beta_ptr ? *params.beta_ptr : params.beta);
+    skip_elementwise_ = false;
+  }
+
+  /// Returns true if source is needed
+  CUTLASS_HOST_DEVICE
+  bool is_source_needed() const {
+    if (Scale == cutlass::epilogue::thread::ScaleType::NoBetaScaling)
+      return params_.beta != ElementCompute(0);
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling)
+      return false;
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) return false;
+
+    return params_.beta != ElementCompute(0);
+  }
+
+  /// Functionally required for serial reduction in the epilogue
+  CUTLASS_HOST_DEVICE
+  void set_k_partition(int k_partition, int k_partition_count) {
+    if (k_partition) {
+      params_.beta = ElementCompute(1);
+    }
+
+    if (k_partition != k_partition_count - 1) {
+      skip_elementwise_ = true;
+    }
+  }
+
+  /// Computes linear scaling with source: D = alpha * accumulator + beta *
+  /// source
+  CUTLASS_HOST_DEVICE
+  FragmentOutput operator()(FragmentAccumulator const &accumulator,
+                            FragmentSource const &source,
+                            int row_offset,
+                            int column_offset) const {
+    CUTLASS_TRACE_DEVICE(
+        "kElementsPerAccess: %d, row_offset: %d, column_offset: %d",
+        kElementsPerAccess,
+        row_offset,
+        column_offset);
+
+    // Convert source to internal compute numeric type
+    cutlass::NumericArrayConverter<ElementCompute,
+                                   ElementOutput,
+                                   kElementsPerAccess,
+                                   Round>
+        source_converter;
+    cutlass::NumericArrayConverter<ElementCompute,
+                                   ElementAccumulator,
+                                   kElementsPerAccess,
+                                   Round>
+        accumulator_converter;
+
+    FragmentCompute converted_source = source_converter(source);
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+
+    cutlass::multiplies<FragmentCompute> mul_add_source;
+    cutlass::multiply_add<FragmentCompute> mul_add_accumulator;
+    VariadicOp<ElementCompute> variadic_op;
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::NoBetaScaling) {
+      intermediate = converted_source;
+      // D = alpha * Accum + X
+      intermediate = mul_add_accumulator(
+          params_.alpha, converted_accumulator, intermediate);
+    } else if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) {
+      intermediate = converted_accumulator;
+    } else {
+      // X =  beta * C + uniform
+      intermediate = mul_add_source(params_.beta, converted_source);
+      // D = alpha * Accum + X
+      intermediate = mul_add_accumulator(
+          params_.alpha, converted_accumulator, intermediate);
+    }
+
+    if constexpr (GenericVariadicTraits<
+                      VariadicOp<ElementCompute>>::IsArgumentsNeeded) {
+      if (!skip_elementwise_) {
+#if CUTLASS_EPILOGUE_ENABLE_VECTORIZE
+        intermediate = variadic_op.Compute<kElementsPerAccess>(
+            intermediate,
+            params_.variadic_args,
+            BatchedMatrixCoord(blockIdx.z, row_offset, column_offset));
+#else
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = variadic_op(
+              intermediate[i],
+              params_.variadic_args,
+              BatchedMatrixCoord(blockIdx.z, row_offset, column_offset + i));
+        }
+#endif
+      }
+    } else {
+      if (!skip_elementwise_) {
+#if CUTLASS_EPILOGUE_ENABLE_VECTORIZE
+        intermediate = variadic_op.Compute<kElementsPerAccess>(intermediate);
+#else
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = variadic_op(intermediate[i]);
+        }
+#endif
+      }
+    }
+
+    // Convert to destination numeric type
+    cutlass::NumericArrayConverter<ElementOutput,
+                                   ElementCompute,
+                                   kElementsPerAccess,
+                                   Round>
+        destination_converter;
+
+    return destination_converter(intermediate);
+  }
+
+  /// Computes linear scaling: D = alpha * accumulator
+  CUTLASS_HOST_DEVICE
+  FragmentOutput operator()(FragmentAccumulator const &accumulator,
+                            int row_offset,
+                            int column_offset) const {
+    CUTLASS_TRACE_DEVICE(
+        "kElementsPerAccess: %d, row_offset: %d, column_offset: %d",
+        kElementsPerAccess,
+        row_offset,
+        column_offset);
+
+    // Convert source to internal compute numeric type
+    cutlass::NumericArrayConverter<ElementCompute,
+                                   ElementAccumulator,
+                                   kElementsPerAccess,
+                                   Round>
+        accumulator_converter;
+
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+
+    cutlass::multiplies<FragmentCompute> mul_accumulator;
+    VariadicOp<ElementCompute> variadic_op;
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) {
+      intermediate = converted_accumulator;
+    } else {
+      // D = alpha * Accum
+      intermediate = mul_accumulator(params_.alpha, converted_accumulator);
+    }
+
+    if constexpr (GenericVariadicTraits<
+                      VariadicOp<FragmentCompute>>::IsArgumentsNeeded) {
+      if (!skip_elementwise_) {
+#if CUTLASS_EPILOGUE_ENABLE_VECTORIZE
+        intermediate = variadic_op.Compute<kElementsPerAccess>(
+            intermediate,
+            params_.variadic_args,
+            BatchedMatrixCoord(blockIdx.z, row_offset, column_offset));
+#else
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = variadic_op(
+              intermediate[i],
+              params_.variadic_args,
+              BatchedMatrixCoord(blockIdx.z, row_offset, column_offset + i));
+        }
+#endif
+      }
+    } else {
+      if (!skip_elementwise_) {
+#if CUTLASS_EPILOGUE_ENABLE_VECTORIZE
+        intermediate = variadic_op.Compute<kElementsPerAccess>(intermediate);
+#else
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = variadic_op(intermediate[i]);
+        }
+#endif
+      }
+    }
+
+    // Convert to destination numeric type
+    cutlass::NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round>
+        destination_converter;
+
+    return destination_converter(intermediate);
+  }
+};
+
+}  // namespace thread
+}  // namespace epilogue
+}  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/epilogue/threadblock/default_epilogue_with_variadic.h

@@ -0,0 +1,243 @@
+// 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.
+
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  The epilogue rearranges the result of a matrix product through shared memory
+  to match canonical tensor layouts in global memory. Epilogues support
+  conversion and reduction operations.
+
+*/
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+#ifdef __NVCC__
+#include "cutlass/array.h"
+#include "cutlass/numeric_types.h"
+
+#include "cutlass/gemm/gemm.h"
+
+#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h"
+#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h"
+#include "cutlass/epilogue/threadblock/epilogue.h"
+
+#include "cutlass/layout/permute.h"
+#elif defined(__HIPCC__)
+#include "hytlass/array.h"
+#include "hytlass/numeric_types.h"
+
+#include "hytlass/gemm/gemm.h"
+
+#include "hytlass/epilogue/threadblock/default_epilogue_tensor_op.h"
+#include "hytlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h"
+#include "hytlass/epilogue/threadblock/epilogue.h"
+
+#include "hytlass/layout/permute.h"
+#endif
+
+#include "cutlass_patch/epilogue/threadblock/epilogue_with_variadic.h"
+// #include "cutlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h"
+
+namespace cutlass_patch {
+namespace epilogue {
+namespace threadblock {
+
+/// Defines sensible defaults for epilogues for SimtOps.
+template <typename Shape,
+          typename WarpMmaSimt,
+          typename ElementOutput,
+          typename OutputOp,
+          int ElementsPerAccess,
+          bool ScatterD = false,
+          typename PermuteDLayout = cutlass::layout::NoPermute,
+          cutlass::conv::StrideSupport StrideSupport =
+              cutlass::conv::StrideSupport::kUnity,
+          int Rank = 4>
+struct DefaultEpilogueWithVariadicSimt {
+  static cutlass::conv::StrideSupport const kStrideSupport = StrideSupport;
+  static int const kRank = Rank;
+
+  static bool const UseCUDAStore =
+      cutlass::platform::is_same<ElementOutput, double>::value;
+
+  /// Use defaults related to the existing epilogue
+  using Base = cutlass::epilogue::threadblock::
+      DefaultEpilogueSimt<Shape, WarpMmaSimt, OutputOp, ElementsPerAccess>;
+
+  using PackedOutputTileIterator =
+      cutlass::epilogue::threadblock::PredicatedTileIterator<
+          typename Base::OutputTileThreadMap,
+          ElementOutput,
+          ScatterD,
+          PermuteDLayout,
+          UseCUDAStore>;
+
+  using StridedOutputTileIterator =
+      cutlass::epilogue::threadblock::PredicatedTileIteratorConv<
+          typename Base::OutputTileThreadMap,
+          ElementOutput,
+          ScatterD,
+          PermuteDLayout,
+          UseCUDAStore,
+          kRank>;
+
+  //
+  // Stores the result z = (y = GEMM(A, B, C), variadic)
+  //
+  using OutputTileIterator = typename cutlass::platform::conditional<
+      StrideSupport == cutlass::conv::StrideSupport::kUnity,
+      PackedOutputTileIterator,
+      StridedOutputTileIterator>::type;
+
+  //
+  // Define the epilogue
+  //
+  using Epilogue = cutlass_patch::epilogue::threadblock::EpilogueWithVariadic<
+      Shape,
+      WarpMmaSimt,
+      Base::kPartitionsK,
+      OutputTileIterator,
+      typename Base::AccumulatorFragmentIterator,
+      typename Base::WarpTileIterator,
+      typename Base::SharedLoadIterator,
+      OutputOp,
+      typename Base::Padding>;
+};
+
+/// Defines sensible defaults for strided dgrad epilogues for SimtOps.
+template <typename Shape,
+          typename WarpMmaSimt,
+          typename ElementOutput,
+          typename OutputOp,
+          int ElementsPerAccess,
+          bool ScatterD = false,
+          typename PermuteDLayout = cutlass::layout::NoPermute>
+struct DefaultEpilogueWithVariadicSimtStridedDgrad {
+  /// Use defaults related to the existing epilogue
+  using Base = cutlass::epilogue::threadblock::DefaultEpilogueSimtStridedDgrad<
+      Shape,
+      WarpMmaSimt,
+      OutputOp,
+      ElementsPerAccess>;
+
+  //
+  // Stores the result z = (y = GEMM(A, B, C), variadic)
+  //
+  using OutputTileIterator =
+      cutlass::epilogue::threadblock::PredicatedTileIteratorStridedDgrad<
+          typename Base::OutputTileThreadMap,
+          ElementOutput>;
+
+  //
+  // Define the epilogue
+  //
+  using Epilogue = cutlass_patch::epilogue::threadblock::EpilogueWithVariadic<
+      Shape,
+      WarpMmaSimt,
+      Base::kPartitionsK,
+      OutputTileIterator,
+      typename Base::AccumulatorFragmentIterator,
+      typename Base::WarpTileIterator,
+      typename Base::SharedLoadIterator,
+      OutputOp,
+      typename Base::Padding>;
+};
+
+/// Defines sensible defaults for epilogues for TensorOps.
+template <typename Shape,
+          typename WarpMmaTensorOp,
+          int PartitionsK,
+          typename ElementOutput,
+          typename OutputOp,
+          int ElementsPerAccess,
+          bool ScatterD = false,
+          typename PermuteDLayout = cutlass::layout::NoPermute>
+struct DefaultEpilogueWithVariadicTensorOp {
+  /// Use defaults related to the existing epilogue
+  using Base = cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<
+      Shape,
+      WarpMmaTensorOp,
+      PartitionsK,
+      OutputOp,
+      ElementsPerAccess>;
+
+  //
+  // Stores the result z = (y = GEMM(A, B, C), variadic)
+  //
+  using OutputTileIterator =
+      cutlass::epilogue::threadblock::PredicatedTileIterator<
+          typename Base::OutputTileThreadMap,
+          ElementOutput,
+          ScatterD,
+          PermuteDLayout>;
+
+  //
+  // Define the epilogue
+  //
+  using Epilogue = cutlass_patch::epilogue::threadblock::EpilogueWithVariadic<
+      Shape,
+      WarpMmaTensorOp,
+      PartitionsK,
+      OutputTileIterator,
+      typename Base::AccumulatorFragmentIterator,
+      typename Base::WarpTileIterator,
+      typename Base::SharedLoadIterator,
+      OutputOp,
+      typename Base::Padding,
+      Base::kFragmentsPerIteration>;
+};
+
+/// Defines sensible defaults for epilogues for VoltaTensorOps.
+template <typename Shape,
+          typename WarpMmaTensorOp,
+          int PartitionsK,
+          typename ElementOutput,
+          typename OutputOp,
+          int ElementsPerAccess>
+struct DefaultEpilogueWithVariadicVoltaTensorOp {
+  /// Use defaults related to the existing epilogue
+  using Base = cutlass::epilogue::threadblock::DefaultEpilogueVoltaTensorOp<
+      Shape,
+      WarpMmaTensorOp,
+      PartitionsK,
+      OutputOp,
+      ElementsPerAccess>;
+
+  //
+  // Stores the result z = (y = GEMM(A, B, C), variadic)
+  //
+  using OutputTileIterator = cutlass::epilogue::threadblock::
+      PredicatedTileIterator<typename Base::OutputTileThreadMap, ElementOutput>;
+
+  //
+  // Define the epilogue
+  //
+  using Epilogue = cutlass_patch::epilogue::threadblock::EpilogueWithVariadic<
+      Shape,
+      WarpMmaTensorOp,
+      PartitionsK,
+      OutputTileIterator,
+      typename Base::AccumulatorFragmentIterator,
+      typename Base::WarpTileIterator,
+      typename Base::SharedLoadIterator,
+      OutputOp,
+      typename Base::Padding>;
+};
+
+}  // namespace threadblock
+}  // namespace epilogue
+}  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/epilogue/threadblock/epilogue_with_variadic.h

@@ -0,0 +1,666 @@
+// 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.
+
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  The epilogue rearranges the result of a matrix product through shared memory
+  to match canonical tensor layouts in global memory. Epilogues support
+  conversion and reduction operations.
+
+  The shared memory resource is time-sliced across warps.
+*/
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+#ifdef __NVCC__
+#include <cuda/std/cassert>
+
+#include "cutlass/aligned_buffer.h"
+#include "cutlass/array.h"
+#include "cutlass/functional.h"
+#include "cutlass/layout/tensor.h"
+#include "cutlass/layout/vector.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/tensor_coord.h"
+
+#include "cutlass/gemm/gemm.h"
+
+#include "cutlass/transform/pitch_linear_thread_map.h"
+#include "cutlass/transform/threadblock/regular_tile_iterator.h"
+
+#include "cutlass/epilogue/threadblock/epilogue_base.h"
+#include "cutlass/epilogue/threadblock/epilogue_base_streamk.h"
+#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
+#elif defined(__HIPCC__)
+#include "hytlass/aligned_buffer.h"
+#include "hytlass/array.h"
+#include "hytlass/functional.h"
+#include "hytlass/layout/tensor.h"
+#include "hytlass/layout/vector.h"
+#include "hytlass/numeric_types.h"
+#include "hytlass/tensor_coord.h"
+
+#include "hytlass/gemm/gemm.h"
+
+#include "hytlass/transform/pitch_linear_thread_map.h"
+#include "hytlass/transform/threadblock/regular_tile_iterator.h"
+
+#include "hytlass/epilogue/threadblock/epilogue_base.h"
+#include "hytlass/epilogue/threadblock/epilogue_base_streamk.h"
+#include "hytlass/epilogue/threadblock/predicated_tile_iterator.h"
+#endif
+
+#include "cutlass_patch/trace_device.h"
+
+namespace cutlass_patch {
+namespace epilogue {
+namespace threadblock {
+
+/// Epilogue operator
+template <
+    typename Shape_,  ///< Shape of threadblock tile (concept: GemmShape)
+    typename WarpMmaOperator_,     ///< Warp-level MMA operator (concept:
+                                   ///< gemm::warp::MmaTensorOp)
+    int PartitionsK,               ///< Number of partitions of the K dimension
+    typename OutputTileIterator_,  ///< Tile iterator reading and writing
+                                   ///< output tensors
+    typename AccumulatorFragmentIterator_,  ///< Fragment iterator
+                                            ///< selecting accumulators
+    typename WarpTileIterator_,    ///< Warp-scoped tile iterator writing
+                                   ///< accumulators to SMEM
+    typename SharedLoadIterator_,  ///< Threadblock-scoped tile iterator
+                                   ///< loading from SMEM
+    typename OutputOp_,            ///< Output operator
+    typename Padding_,  ///< Padding added to SMEM allocation to avoid
+                        ///< bank conflicts (concept: MatrixShape)
+    int FragmentsPerPartition =
+        1,                  ///< Used to coarsten the epilogue granularity
+    int IterationsUnroll =  ///< Used to reduce binary size when epilogue
+                            ///< op is large
+    (!cutlass::epilogue::threadblock::IsEpilogueFunctorHeavy<OutputOp_>::value)>
+class EpilogueWithVariadic
+    : public cutlass::epilogue::threadblock::EpilogueBase<
+          Shape_,
+          typename WarpMmaOperator_::Shape,
+          PartitionsK,
+          AccumulatorFragmentIterator_,
+          WarpTileIterator_,
+          Padding_,
+          FragmentsPerPartition>,
+      public cutlass::epilogue::threadblock::EpilogueBaseStreamK<
+          Shape_,
+          PartitionsK,
+          WarpMmaOperator_,
+          AccumulatorFragmentIterator_> {
+ public:
+  using Base = cutlass::epilogue::threadblock::EpilogueBase<
+      Shape_,
+      typename WarpMmaOperator_::Shape,
+      PartitionsK,
+      AccumulatorFragmentIterator_,
+      WarpTileIterator_,
+      Padding_,
+      FragmentsPerPartition>;
+
+  using BaseStreamK = cutlass::epilogue::threadblock::EpilogueBaseStreamK<
+      Shape_,
+      PartitionsK,
+      WarpMmaOperator_,
+      AccumulatorFragmentIterator_>;
+
+  using Shape = Shape_;
+  using WarpMmaOperator = WarpMmaOperator_;
+  static int const kPartitionsK = PartitionsK;
+  using OutputTileIterator = OutputTileIterator_;
+  using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
+  using WarpTileIterator = WarpTileIterator_;
+  using SharedLoadIterator = SharedLoadIterator_;
+  using OutputOp = OutputOp_;
+  using Padding = Padding_;
+  using Layout = cutlass::layout::RowMajor;
+  using LongIndex = typename Layout::LongIndex;
+
+  /// Number of warps per block
+  using WarpCount = typename Base::WarpCount;
+
+  /// Number of threads per block
+  static int const kBlockThreads = 32 * WarpCount::kCount;
+
+  /// Per-thread accumulator tile type
+  using AccumulatorTile = typename Base::AccumulatorTile;
+
+  /// Numerical accumulation element type
+  using ElementAccumulator = typename WarpMmaOperator::ElementC;
+
+  /// Fragment type used by the accumulator tile's fragment iterator
+  using AccumulatorFragment = typename AccumulatorFragmentIterator::Fragment;
+
+  /// Output element
+  using ElementOutput = typename OutputTileIterator::Element;
+
+  /// Output access size
+  static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
+
+  /// Tensor reference to destination tensor
+  using TensorRef = typename OutputTileIterator::TensorRef;
+
+  /// Tensor reference to sync tensor
+  using SyncTensorRef =
+      typename cutlass::TensorRef<int, cutlass::layout::PackedVectorLayout>;
+
+  /// Const tensor reference to source tensor
+  using ConstTensorRef = typename OutputTileIterator::ConstTensorRef;
+
+  /// Vector type used by the global output iterator
+  using OutputAccessType =
+      cutlass::Array<typename OutputTileIterator::Element,
+                     OutputTileIterator::kElementsPerAccess>;
+
+  /// Vector type used by the shared output iterator
+  using AccumulatorAccessType =
+      cutlass::Array<typename WarpTileIterator::Element,
+                     OutputTileIterator::kElementsPerAccess>;
+
+  static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1
+                                        ? Base::kFragmentsPerIteration
+                                        : kPartitionsK;
+
+  static int constexpr kSmemPointerOffset =
+      Base::SharedStorage::StorageShape::kCount / kSmemTiles;
+
+ public:
+  static_assert(
+      SharedLoadIterator::Fragment::kElements ==
+          OutputTileIterator::Fragment::kElements,
+      "Mismatch between shared load iterator and output tile iterator.");
+
+  static_assert(OutputTileIterator::kElementsPerAccess,
+                "OutputTileIterator::kElementsPerAccess must not be zero.");
+
+  static_assert(!(OutputTileIterator::Fragment::kElements %
+                  OutputTileIterator::kElementsPerAccess),
+                "Divisibility");
+
+  static_assert(kPartitionsK == 1 || Base::kFragmentsPerIteration == 1,
+                "One of these must be exactly 1.");
+
+ public:
+  /// Aspect for when epilogue source is not needed
+  struct SourceAspectNotNeeded {
+    /// Constructor
+    CUTLASS_DEVICE
+    SourceAspectNotNeeded() {}
+
+    // No-op
+    CUTLASS_DEVICE
+    void load() {}
+
+    /// Invoke the output functor over each vector of output
+    CUTLASS_DEVICE
+    void apply_output_operator(
+        const OutputTileIterator &output_iterator,
+        typename OutputTileIterator::Fragment &output_fragment,  // NOLINT
+        OutputOp const &output_op,
+        typename SharedLoadIterator::Fragment const &aligned_accum_fragment) {
+      CUTLASS_TRACE_DEVICE("");
+
+      OutputAccessType *output_frag_ptr =
+          reinterpret_cast<OutputAccessType *>(&output_fragment);
+
+      AccumulatorAccessType const *compute_frag_ptr =
+          reinterpret_cast<AccumulatorAccessType const *>(
+              &aligned_accum_fragment);
+
+      const int32_t thread_start_row = output_iterator.thread_start_row();
+      const int32_t thread_start_column = output_iterator.thread_start_column();
+
+      const typename OutputTileIterator::Index extent_row =
+          output_iterator.extent_row();
+      const typename OutputTileIterator::Index extent_column =
+          output_iterator.extent_column();
+
+      using ThreadMap = typename OutputTileIterator::ThreadMap;
+
+      typename OutputTileIterator::Mask mask;
+      output_iterator.get_mask(mask);
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster;
+           ++cluster) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+          CUTLASS_PRAGMA_UNROLL
+          for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+            int frag_row_idx =
+                (row + ThreadMap::Iterations::kRow *
+                           (group + ThreadMap::Iterations::kGroup * cluster));
+
+            int row_offset = thread_start_row + row * ThreadMap::Delta::kRow +
+                             group * ThreadMap::Delta::kGroup +
+                             cluster * ThreadMap::Delta::kCluster;
+
+            bool row_guard = row_offset < extent_row;
+
+            CUTLASS_PRAGMA_UNROLL
+            for (int column = 0; column < ThreadMap::Iterations::kColumn;
+                 ++column) {
+              bool guard = row_guard && mask.predicates[column];
+              if (!guard) {
+                continue;
+              }
+
+              int column_offset =
+                  thread_start_column + column * ThreadMap::Delta::kColumn;
+              int frag_offset =
+                  frag_row_idx * ThreadMap::Iterations::kColumn + column;
+
+              output_frag_ptr[frag_offset] = output_op(
+                  compute_frag_ptr[frag_offset], row_offset, column_offset);
+            }
+          }
+        }
+      }
+    }
+  };
+
+  /// Aspect for when epilogue source is needed
+  struct SourceAspectNeeded {
+    OutputTileIterator source_iterator;
+
+    typename OutputTileIterator::Fragment source_fragment;
+
+    /// Invoke the output functor over each vector of output
+    CUTLASS_DEVICE
+    static void apply_output_operator(
+        const OutputTileIterator &output_iterator,
+        typename OutputTileIterator::Fragment &output_fragment,  // NOLINT
+        OutputOp const &output_op,
+        typename SharedLoadIterator::Fragment const &aligned_accum_fragment,
+        typename OutputTileIterator::Fragment const &source_fragment) {
+      CUTLASS_TRACE_DEVICE("");
+
+      OutputAccessType *output_frag_ptr =
+          reinterpret_cast<OutputAccessType *>(&output_fragment);
+
+      AccumulatorAccessType const *compute_frag_ptr =
+          reinterpret_cast<AccumulatorAccessType const *>(
+              &aligned_accum_fragment);
+
+      OutputAccessType const *source_frag_ptr =
+          reinterpret_cast<OutputAccessType const *>(&source_fragment);
+
+      typename OutputTileIterator::Element const *source_ptr =
+          reinterpret_cast<typename OutputTileIterator::Element const *>(
+              &source_fragment);
+
+      const int32_t thread_start_row = output_iterator.thread_start_row();
+      const int32_t thread_start_column = output_iterator.thread_start_column();
+
+      const typename OutputTileIterator::Index extent_row =
+          output_iterator.extent_row();
+      const typename OutputTileIterator::Index extent_column =
+          output_iterator.extent_column();
+
+      using ThreadMap = typename OutputTileIterator::ThreadMap;
+
+      typename OutputTileIterator::Mask mask;
+      output_iterator.get_mask(mask);
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster;
+           ++cluster) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+          CUTLASS_PRAGMA_UNROLL
+          for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+            int frag_row_idx =
+                (row + ThreadMap::Iterations::kRow *
+                           (group + ThreadMap::Iterations::kGroup * cluster));
+
+            int row_offset = thread_start_row + row * ThreadMap::Delta::kRow +
+                             group * ThreadMap::Delta::kGroup +
+                             cluster * ThreadMap::Delta::kCluster;
+
+            bool row_guard = row_offset < extent_row;
+
+            CUTLASS_PRAGMA_UNROLL
+            for (int column = 0; column < ThreadMap::Iterations::kColumn;
+                 ++column) {
+              bool guard = row_guard && mask.predicates[column];
+              if (!guard) {
+                continue;
+              }
+
+              int column_offset =
+                  thread_start_column + column * ThreadMap::Delta::kColumn;
+              int frag_offset =
+                  frag_row_idx * ThreadMap::Iterations::kColumn + column;
+
+              output_frag_ptr[frag_offset] =
+                  output_op(compute_frag_ptr[frag_offset],
+                            source_frag_ptr[frag_offset],
+                            row_offset,
+                            column_offset);
+            }
+          }
+        }
+      }
+    }
+
+    /// Constructor
+    CUTLASS_DEVICE
+    explicit SourceAspectNeeded(OutputTileIterator source_iterator)
+        : source_iterator(source_iterator) {
+      source_fragment.clear();
+    }
+
+    // Load addend source fragment from global memory
+    CUTLASS_DEVICE
+    void load() {
+      source_iterator.load(source_fragment);
+      ++source_iterator;
+    }
+
+    /// Invoke the output functor over each vector of output
+    CUTLASS_DEVICE
+    void apply_output_operator(
+        const OutputTileIterator &output_iterator,
+        typename OutputTileIterator::Fragment &output_fragment,  // NOLINT
+        OutputOp const &output_op,
+        typename SharedLoadIterator::Fragment const &aligned_accum_fragment) {
+      apply_output_operator(output_iterator,
+                            output_fragment,
+                            output_op,
+                            aligned_accum_fragment,
+                            source_fragment);
+    }
+  };
+
+ private:
+  /// Loads fragment from shared memory aligned with output tensor
+  SharedLoadIterator shared_load_iterator_;
+
+  /// Thread index in the threadblock
+  int thread_idx;
+
+  /// Warp index in the threadblock
+  int warp_idx;
+
+ public:
+  /// Constructor
+  CUTLASS_DEVICE
+  EpilogueWithVariadic(
+      typename Base::SharedStorage
+          &shared_storage,  // NOLINT  ///< Shared storage object
+      int thread_idx,       ///< ID of a thread within the threadblock
+      int warp_idx,         ///< ID of warp within threadblock
+      int lane_idx)         ///< Id of thread within warp
+      : Base(shared_storage, thread_idx, warp_idx, lane_idx),
+        BaseStreamK(thread_idx),
+        shared_load_iterator_(shared_storage.reference(), thread_idx),
+        thread_idx(thread_idx),
+        warp_idx(warp_idx) {}
+
+  /// Aggregates the accumulator sets shared by peer blocks in the global
+  /// workspace, performing epilogue computations, writing to output
+  CUTLASS_DEVICE
+  void reduce(int peer_idx_begin,
+              int peer_idx_end,
+              int reduce_fragment_idx,
+              void *element_workspace,
+              OutputOp const &output_op,  ///< Output operator
+              OutputTileIterator
+                  destination_iterator,  ///< Tile iterator for destination
+              OutputTileIterator
+                  source_iterator) {  ///< Threadblock tile coordinate in GEMM
+                                      ///< (in units of threadblock tiles)
+    CUTLASS_TRACE_DEVICE("");
+
+    // Reduce peer accumulator fragments into one fragment
+    AccumulatorFragment accum_fragment;
+    BaseStreamK::reduce(accum_fragment,
+                        peer_idx_begin,
+                        peer_idx_end,
+                        reduce_fragment_idx,
+                        element_workspace);
+
+    // Store fragment to shared memory
+    this->warp_tile_iterator_.store(accum_fragment);
+
+    __syncthreads();
+
+    // Initialize/load source-fragment data
+    typename OutputTileIterator::Fragment source_fragment;
+    source_fragment.clear();
+
+    if (output_op.is_source_needed()) {
+      source_iterator += reduce_fragment_idx;
+      source_iterator.load(source_fragment);
+    }
+
+    // Load fragment from shared memory
+    typename SharedLoadIterator::Fragment aligned_accum_fragment;
+    shared_load_iterator_.load(aligned_accum_fragment);
+
+    // Add fragments shared by other k partitions
+    if (kPartitionsK > 1) {
+      cutlass::plus<typename SharedLoadIterator::Fragment> add_fragments;
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 1; i < kPartitionsK; ++i) {
+        typename SharedLoadIterator::Fragment aligned_addend_fragment;
+        shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
+        shared_load_iterator_.load(aligned_addend_fragment);
+        aligned_accum_fragment =
+            add_fragments(aligned_accum_fragment, aligned_addend_fragment);
+      }
+    }
+
+    // Compute the output result
+    typename OutputTileIterator::Fragment output_fragment;
+
+    // Apply the output operator
+    SourceAspectNeeded::apply_output_operator(
+        output_fragment, output_op, aligned_accum_fragment, source_fragment);
+
+    // Store the final result
+    destination_iterator += reduce_fragment_idx;
+    destination_iterator.store(output_fragment);
+  }
+
+  /// Perform the epilogue computations and stream the result to global memory.
+  CUTLASS_DEVICE
+  void operator()(OutputOp const &output_op,  ///< Output operator
+                  OutputTileIterator
+                      destination_iterator,  ///< Tile iterator for destination
+                  AccumulatorTile const &
+                      accumulators) {  ///< Complete warp-level accumulator tile
+    CUTLASS_TRACE_DEVICE("");
+    operator()(
+        output_op, destination_iterator, accumulators, SourceAspectNotNeeded());
+  }
+
+  /// Perform the epilogue computations and stream the result to global memory.
+  /// Implements two alternative codepaths, depending on whether the output op
+  /// requires addend data to be loaded.
+  CUTLASS_DEVICE
+  void operator()(OutputOp const &output_op,  ///< Output operator
+                  OutputTileIterator
+                      destination_iterator,  ///< Tile iterator for destination
+                  AccumulatorTile const
+                      &accumulators,  ///< Complete warp-level accumulator tile
+                  OutputTileIterator
+                      source_iterator) {  ///< Tile iterator for addend source
+    CUTLASS_TRACE_DEVICE("");
+    if (output_op.is_source_needed()) {
+      operator()(output_op,
+                 destination_iterator,
+                 accumulators,
+                 SourceAspectNeeded(source_iterator));
+    } else {
+      operator()(output_op,
+                 destination_iterator,
+                 accumulators,
+                 SourceAspectNotNeeded());
+    }
+  }
+
+  /// Perform the epilogue computations and stream the result to global memory.
+  /// Implements a single codepath, regardless of whether the output op requires
+  /// addend data to be loaded
+  CUTLASS_DEVICE
+  void unified(OutputOp const &output_op,  ///< Output operator
+               OutputTileIterator
+                   destination_iterator,  ///< Tile iterator for destination
+               AccumulatorTile const
+                   &accumulators,  ///< Complete warp-level accumulator tile
+               OutputTileIterator
+                   source_iterator) {  ///< Tile iterator for addend source
+    CUTLASS_TRACE_DEVICE("");
+    if (!output_op.is_source_needed()) {
+      source_iterator.clear_mask();
+      __syncthreads();  // Dummy (CUDA 11.0)
+    }
+
+    operator()(output_op,
+               destination_iterator,
+               accumulators,
+               SourceAspectNeeded(source_iterator));
+  }
+
+  template <class Seq>
+  struct acc2smem;
+
+  template <size_t... Seq>
+  struct acc2smem<cutlass::index_sequence<Seq...>> {
+    template <int Advance>
+    CUTLASS_DEVICE static void helper(
+        AccumulatorFragmentIterator accum_fragment_iterator,
+        WarpTileIterator &warp_tile_iterator) {  // NOLINT
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < Advance; i++) {
+        ++accum_fragment_iterator;
+      }
+
+      typename AccumulatorFragmentIterator::Fragment accum_fragment;
+
+      accum_fragment_iterator.load(accum_fragment);
+      ++accum_fragment_iterator;
+      warp_tile_iterator.store(accum_fragment);
+    }
+
+    CUTLASS_DEVICE
+    static void push(size_t pos,
+                     AccumulatorFragmentIterator const &iterator_begin,
+                     WarpTileIterator &warp_tile_iterator) {  // NOLINT
+      int dummy[] = {(pos == Seq) &&
+                     (helper<Seq>(iterator_begin, warp_tile_iterator), 0)...};
+    }
+  };
+
+  /// Streams the result to global memory
+  template <typename SourceAspect>
+  CUTLASS_DEVICE void operator()(
+      OutputOp const &output_op,  ///< Output operator
+      OutputTileIterator
+          destination_iterator,  ///< Tile iterator for destination
+      AccumulatorTile const
+          &accumulators,  ///< Complete warp-level accumulator tile
+      SourceAspect source) {
+    CUTLASS_TRACE_DEVICE("");
+
+    // Iterator over warp-level accumulator fragment
+    AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
+
+    //
+    // Iterate over accumulator tile
+    //
+
+#ifdef __clang__
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wcuda-compat"
+// Turn off clangs warning about loop unroll argument using parens.
+#endif
+
+#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1)
+    for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
+      //
+      // Load the source
+      //
+
+      source.load();
+      //
+      // Convert and store fragment
+      //
+
+      __syncthreads();
+
+      acc2smem<cutlass::make_index_sequence<OutputTileIterator::kIterations>>::
+          push(iter, accum_fragment_iterator, this->warp_tile_iterator_);
+
+      __syncthreads();
+
+      //
+      // Load fragments from shared memory
+      //
+
+      typename SharedLoadIterator::Fragment
+          aligned_accum_fragment[kPartitionsK];
+      shared_load_iterator_.load(aligned_accum_fragment[0]);
+
+      if (kPartitionsK > 1) {
+        cutlass::plus<typename SharedLoadIterator::Fragment> add_fragments;
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 1; i < kPartitionsK; ++i) {
+          shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
+          shared_load_iterator_.load(aligned_accum_fragment[i]);
+          aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0],
+                                                    aligned_accum_fragment[i]);
+        }
+
+        shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) *
+                                                 kSmemPointerOffset);
+      }
+
+      //
+      // Compute the output result
+      //
+
+      typename OutputTileIterator::Fragment output_fragment;
+      source.apply_output_operator(destination_iterator,
+                                   output_fragment,
+                                   output_op,
+                                   aligned_accum_fragment[0]);
+
+      //
+      // Store the final result
+      //
+
+      destination_iterator.store(output_fragment);
+      ++destination_iterator;
+    }
+
+#ifdef __clang__
+#pragma clang diagnostic pop
+#endif
+  }
+};
+
+}  // namespace threadblock
+}  // namespace epilogue
+}  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/gemm/device/gemm_universal_with_variadic.h

@@ -0,0 +1,416 @@
+// 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.
+
+/*! \file
+    \brief
+*/
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+#ifdef __NVCC__
+#include "cutlass/arch/arch.h"
+#include "cutlass/arch/mma.h"
+#include "cutlass/device_kernel.h"
+#include "cutlass/numeric_types.h"
+
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/gemm/kernel/gemm_universal.h"
+#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
+
+#include "cutlass/gemm/device/default_gemm_configuration.h"
+#include "cutlass/gemm/device/gemm_universal_base.h"
+#include "cutlass/gemm/kernel/default_gemm_universal.h"
+
+#include "cutlass/layout/permute.h"
+#elif defined(__HIPCC__)
+#include "hytlass/arch/arch.h"
+#include "hytlass/arch/mma.h"
+#include "hytlass/device_kernel.h"
+#include "hytlass/numeric_types.h"
+
+#include "hytlass/gemm/gemm.h"
+#include "hytlass/gemm/kernel/gemm_universal.h"
+#include "hytlass/gemm/threadblock/threadblock_swizzle.h"
+
+#include "hytlass/gemm/device/default_gemm_configuration.h"
+#include "hytlass/gemm/device/gemm_universal_base.h"
+#include "hytlass/gemm/kernel/default_gemm_universal.h"
+
+#include "hytlass/layout/permute.h"
+#endif
+
+#include "cutlass_patch/gemm/kernel/default_gemm_with_variadic.h"
+
+namespace cutlass_patch {
+namespace gemm {
+namespace device {
+
+/*!
+  GemmUniversal with variadic epilogues.
+*/
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator_ = ElementC_,
+    /// Operator class tag
+    typename OperatorClass_ = cutlass::arch::OpClassSimt,
+    /// Tag indicating architecture to tune for.  This is the minimum SM that
+    /// supports the intended feature. The device kernel can be built
+    /// targeting any SM larger than this number.
+    typename ArchTag_ = cutlass::arch::Sm70,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape_ = typename cutlass::gemm::device::
+        DefaultGemmConfiguration<OperatorClass_,
+                                 ArchTag_,
+                                 ElementA_,
+                                 ElementB_,
+                                 ElementC_,
+                                 ElementAccumulator_>::ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape_ = typename cutlass::gemm::device::
+        DefaultGemmConfiguration<OperatorClass_,
+                                 ArchTag_,
+                                 ElementA_,
+                                 ElementB_,
+                                 ElementC_,
+                                 ElementAccumulator_>::WarpShape,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape_ = typename cutlass::gemm::device::
+        DefaultGemmConfiguration<OperatorClass_,
+                                 ArchTag_,
+                                 ElementA_,
+                                 ElementB_,
+                                 ElementC_,
+                                 ElementAccumulator_>::InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp_ = typename cutlass::gemm::device::
+        DefaultGemmConfiguration<OperatorClass_,
+                                 ArchTag_,
+                                 ElementA_,
+                                 ElementB_,
+                                 ElementC_,
+                                 ElementAccumulator_>::EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle_ =
+        cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
+    /// Number of stages used in the pipelined mainloop
+    int Stages = cutlass::gemm::device::DefaultGemmConfiguration<
+        OperatorClass_,
+        ArchTag_,
+        ElementA_,
+        ElementB_,
+        ElementC_,
+        ElementAccumulator_>::kStages,
+    /// Access granularity of A matrix in units of elements
+    int AlignmentA = cutlass::gemm::device::DefaultGemmConfiguration<
+        OperatorClass_,
+        ArchTag_,
+        ElementA_,
+        ElementB_,
+        ElementC_,
+        ElementAccumulator_>::kAlignmentA,
+    /// Access granularity of B matrix in units of elements
+    int AlignmentB = cutlass::gemm::device::DefaultGemmConfiguration<
+        OperatorClass_,
+        ArchTag_,
+        ElementA_,
+        ElementB_,
+        ElementC_,
+        ElementAccumulator_>::kAlignmentB,
+    /// Operation performed by GEMM
+    typename Operator_ = typename cutlass::gemm::device::
+        DefaultGemmConfiguration<OperatorClass_,
+                                 ArchTag_,
+                                 ElementA_,
+                                 ElementB_,
+                                 ElementC_,
+                                 ElementAccumulator_>::Operator,
+    /// Complex elementwise transformation on A operand
+    cutlass::ComplexTransform TransformA = cutlass::ComplexTransform::kNone,
+    /// Complex elementwise transformation on B operand
+    cutlass::ComplexTransform TransformB = cutlass::ComplexTransform::kNone>
+class GemmUniversalWithVariadic
+    : public cutlass::gemm::device::GemmUniversalBase<
+          typename cutlass_patch::gemm::kernel::DefaultGemmWithVariadic<
+              ElementA_,
+              LayoutA_,
+              TransformA,
+              AlignmentA,
+              ElementB_,
+              LayoutB_,
+              TransformB,
+              AlignmentB,
+              ElementC_,
+              LayoutC_,
+              ElementAccumulator_,
+              OperatorClass_,
+              ArchTag_,
+              ThreadblockShape_,
+              WarpShape_,
+              InstructionShape_,
+              EpilogueOutputOp_,
+              ThreadblockSwizzle_,
+              Stages,
+              Operator_>::GemmKernel> {
+ public:
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp = EpilogueOutputOp_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  using Operator = Operator_;
+  static int const kStages = Stages;
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static int const kAlignmentC = EpilogueOutputOp::kCount;
+  static cutlass::ComplexTransform const kTransformA = TransformA;
+  static cutlass::ComplexTransform const kTransformB = TransformB;
+
+  using Base = cutlass::gemm::device::GemmUniversalBase<
+      typename cutlass_patch::gemm::kernel::DefaultGemmWithVariadic<
+          ElementA_,
+          LayoutA_,
+          TransformA,
+          AlignmentA,
+          ElementB_,
+          LayoutB_,
+          TransformB,
+          AlignmentB,
+          ElementC_,
+          LayoutC_,
+          ElementAccumulator_,
+          OperatorClass_,
+          ArchTag_,
+          ThreadblockShape_,
+          WarpShape_,
+          InstructionShape_,
+          EpilogueOutputOp_,
+          ThreadblockSwizzle_,
+          Stages,
+          Operator_>::GemmKernel>;
+
+  using Arguments = typename Base::Arguments;
+  using GemmKernel = typename Base::GemmKernel;
+};
+
+/// Partial specialization for column-major output exchanges problem size and
+/// operand.
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator_,
+    /// Operator class tag
+    typename OperatorClass_,
+    /// Tag indicating architecture to tune for.  This is the minimum SM that
+    /// supports the intended feature. The device kernel can be built
+    /// targeting any SM larger than this number.
+    typename ArchTag_,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape_,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape_,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape_,
+    /// Epilogue output operator
+    typename EpilogueOutputOp_,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle_,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Access granularity of A matrix in units of elements
+    int AlignmentA,
+    /// Access granularity of B matrix in units of elements
+    int AlignmentB,
+    /// Operation performed by GEMM
+    typename Operator_,
+    /// Complex elementwise transformation on A operand
+    cutlass::ComplexTransform TransformA,
+    /// Complex elementwise transformation on B operand
+    cutlass::ComplexTransform TransformB>
+class GemmUniversalWithVariadic<ElementA_,
+                                LayoutA_,
+                                ElementB_,
+                                LayoutB_,
+                                ElementC_,
+                                cutlass::layout::ColumnMajor,  // partially
+                                                               // specialized on
+                                                               // LayoutC
+                                ElementAccumulator_,
+                                OperatorClass_,
+                                ArchTag_,
+                                ThreadblockShape_,
+                                WarpShape_,
+                                InstructionShape_,
+                                EpilogueOutputOp_,
+                                ThreadblockSwizzle_,
+                                Stages,
+                                AlignmentA,
+                                AlignmentB,
+                                Operator_,
+                                TransformA,
+                                TransformB> {
+ public:
+  using ElementA = ElementA_;
+  using LayoutA = LayoutA_;
+  using TensorRefA = cutlass::TensorRef<ElementA const, LayoutA>;
+  using ElementB = ElementB_;
+  using LayoutB = LayoutB_;
+  using TensorRefB = cutlass::TensorRef<ElementB const, LayoutB>;
+  using ElementC = ElementC_;
+  using LayoutC = cutlass::layout::ColumnMajor;
+  using TensorRefC = cutlass::TensorRef<ElementC const, LayoutC>;
+  using TensorRefD = cutlass::TensorRef<ElementC, LayoutC>;
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp = EpilogueOutputOp_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  using Operator = Operator_;
+  static int const kStages = Stages;
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static cutlass::ComplexTransform const kTransformA = TransformA;
+  static cutlass::ComplexTransform const kTransformB = TransformB;
+
+  using UnderlyingOperator = typename GemmUniversalWithVariadic<
+      ElementB,
+      typename cutlass::layout::LayoutTranspose<LayoutB>::type,
+      ElementA,
+      typename cutlass::layout::LayoutTranspose<LayoutA>::type,
+      ElementC,
+      cutlass::layout::RowMajor,
+      ElementAccumulator,
+      OperatorClass,
+      ArchTag,
+      ThreadblockShape,
+      WarpShape,
+      InstructionShape,
+      EpilogueOutputOp,
+      ThreadblockSwizzle,
+      Stages,
+      kAlignmentB,
+      kAlignmentA,
+      Operator,
+      kTransformB,
+      kTransformA>::Base;
+
+  using GemmKernel = typename UnderlyingOperator::GemmKernel;
+  static int const kAlignmentC = EpilogueOutputOp::kCount;
+
+  /// Argument structure
+  using Arguments = typename UnderlyingOperator::Arguments;
+
+ private:
+  UnderlyingOperator underlying_operator_;
+
+ public:
+  /// Constructs the GEMM.
+  GemmUniversalWithVariadic() {}
+
+  /// Helper to construct a transposed equivalent for the underlying GEMM
+  /// operator
+  static Arguments to_underlying_arguments(Arguments const &args) {
+    return args.transposed_problem();
+  }
+
+  /// Determines whether the GEMM can execute the given problem.
+  static cutlass::Status can_implement(Arguments const &args) {
+    return UnderlyingOperator::can_implement(to_underlying_arguments(args));
+  }
+
+  /// Gets the workspace size
+  static size_t get_workspace_size(Arguments const &args) {
+    return UnderlyingOperator::get_workspace_size(
+        to_underlying_arguments(args));
+  }
+
+  /// Computes the grid shape
+  static dim3 get_grid_shape(Arguments const &args) {
+    return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args));
+  }
+
+  /// Computes the maximum number of active blocks per multiprocessor
+  static int maximum_active_blocks(int smem_capacity = -1) {
+    return UnderlyingOperator::maximum_active_blocks(smem_capacity);
+  }
+
+  /// Initializes GEMM state from arguments.
+  cutlass::Status initialize(Arguments const &args,
+                             void *workspace = nullptr,
+                             GPUStream_t stream = nullptr) {
+    return underlying_operator_.initialize(
+        to_underlying_arguments(args), workspace, stream);
+  }
+
+  /// Lightweight update given a subset of arguments
+  cutlass::Status update(Arguments const &args, void *workspace = nullptr) {
+    return underlying_operator_.update(to_underlying_arguments(args),
+                                       workspace);
+  }
+
+  /// Runs the kernel using initialized state.
+  cutlass::Status run(GPUStream_t stream = nullptr) {
+    return underlying_operator_.run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  cutlass::Status operator()(GPUStream_t stream = nullptr) {
+    return run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  cutlass::Status operator()(Arguments const &args,
+                             void *workspace = nullptr,
+                             GPUStream_t stream = nullptr) {
+    cutlass::Status status = initialize(args, workspace, stream);
+
+    if (status == cutlass::Status::kSuccess) {
+      status = run(stream);
+    }
+
+    return status;
+  }
+};
+
+}  // namespace device
+}  // namespace gemm
+}  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/gemm/kernel/default_gemm_with_variadic.h

@@ -0,0 +1,227 @@
+// 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.
+
+/*! \file
+  \brief
+    Defines a GEMM with Reduction based on an existing UniversalGemm kernel.
+
+*/
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+#ifdef __NVCC__
+#include "cutlass/gemm/kernel/default_gemm_universal.h"
+#include "cutlass/gemm/kernel/gemm_universal.h"
+#elif defined(__HIPCC__)
+#include "hytlass/gemm/kernel/default_gemm_universal.h"
+#include "hytlass/gemm/kernel/gemm_universal.h"
+#endif
+
+#include "cutlass_patch/epilogue/threadblock/default_epilogue_with_variadic.h"
+#include "cutlass_patch/epilogue/threadblock/epilogue_with_variadic.h"
+
+namespace cutlass_patch {
+namespace gemm {
+namespace kernel {
+
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Complex elementwise transformation on A operand
+    cutlass::ComplexTransform TransformA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Complex elementwise transformation on B operand
+    cutlass::ComplexTransform TransformB,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Operator class tag
+    typename OperatorClass,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator      - must satisfy concept of
+    /// 'EpilogueWithVariadicOp'
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Operation performed by GEMM
+    typename Operator,
+    ///
+    typename Enable = void>
+struct DefaultGemmWithVariadic {
+  using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal<
+      ElementA_,
+      LayoutA_,
+      TransformA,
+      kAlignmentA,
+      ElementB_,
+      LayoutB_,
+      TransformB,
+      kAlignmentB,
+      ElementC_,
+      LayoutC_,
+      ElementAccumulator,
+      OperatorClass,
+      ArchTag,
+      ThreadblockShape,
+      WarpShape,
+      InstructionShape,
+      EpilogueOutputOp,
+      ThreadblockSwizzle,
+      Stages,
+      Operator>::GemmKernel;
+
+  // Define epilogue
+  using Epilogue = typename cutlass_patch::epilogue::threadblock::
+      DefaultEpilogueWithVariadicTensorOp<
+          typename GemmBase::Epilogue::Shape,
+          typename GemmBase::Epilogue::WarpMmaOperator,
+          GemmBase::Epilogue::kPartitionsK,
+          ElementC_,
+          EpilogueOutputOp,
+          GemmBase::Epilogue::kElementsPerAccess>::Epilogue;
+
+  // Compose the GEMM kernel
+  using GemmKernel = cutlass::gemm::kernel::
+      GemmUniversal<typename GemmBase::Mma, Epilogue, ThreadblockSwizzle>;
+};
+
+/// Partial specialization: ArchTag = cutlass::arch::Sm70
+///
+///
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Complex elementwise transformation on A operand
+    cutlass::ComplexTransform TransformA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Complex elementwise transformation on B operand
+    cutlass::ComplexTransform TransformB,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Operator class tag
+    typename OperatorClass,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator      - must satisfy concept of
+    /// 'EpilogueWithVariadicOp'
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Operation performed by GEMM
+    typename Operator,
+    ///
+    typename Enable>
+struct DefaultGemmWithVariadic<ElementA_,
+                               LayoutA_,
+                               TransformA,
+                               kAlignmentA,
+                               ElementB_,
+                               LayoutB_,
+                               TransformB,
+                               kAlignmentB,
+                               ElementC_,
+                               LayoutC_,
+                               ElementAccumulator,
+                               OperatorClass,
+                               cutlass::arch::Sm70,
+                               ThreadblockShape,
+                               WarpShape,
+                               InstructionShape,
+                               EpilogueOutputOp,
+                               ThreadblockSwizzle,
+                               Stages,
+                               Operator,
+                               Enable> {
+  using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal<
+      ElementA_,
+      LayoutA_,
+      TransformA,
+      kAlignmentA,
+      ElementB_,
+      LayoutB_,
+      TransformB,
+      kAlignmentB,
+      ElementC_,
+      LayoutC_,
+      ElementAccumulator,
+      OperatorClass,
+      cutlass::arch::Sm70,
+      ThreadblockShape,
+      WarpShape,
+      InstructionShape,
+      EpilogueOutputOp,
+      ThreadblockSwizzle,
+      Stages,
+      Operator>::GemmKernel;
+
+  // Define epilogue
+  using Epilogue = typename cutlass_patch::epilogue::threadblock::
+      DefaultEpilogueWithVariadicVoltaTensorOp<
+          typename GemmBase::Epilogue::Shape,
+          typename GemmBase::Epilogue::WarpMmaOperator,
+          GemmBase::Epilogue::kPartitionsK,
+          ElementC_,
+          EpilogueOutputOp,
+          GemmBase::Epilogue::kElementsPerAccess>::Epilogue;
+
+  // Compose the GEMM kernel
+  using GemmKernel = cutlass::gemm::kernel::
+      GemmUniversal<typename GemmBase::Mma, Epilogue, ThreadblockSwizzle>;
+};
+
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/hip/all_tuning_configs.h

@@ -0,0 +1,552 @@
+// Copyright (c) 2026 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.
+
+#pragma once
+
+#include "hytlass/gemm_coord.h"
+
+namespace ap {
+
+constexpr int kNumConfigsHalf = 28;
+constexpr int kNumConfigsFloat = 13;
+
+template <int SwizzleFactor, bool Batched>
+struct SwizzleWrapper {
+  using Type =
+      hytlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<SwizzleFactor>;
+};
+
+#define AP_AUTOTUNE(func, stream_ptr, count, ...)                             \
+  {                                                                           \
+    using FuncType = decltype(func<0>);                                       \
+    static int selected_config_id = -1;                                       \
+    static std::vector<std::function<FuncType>> matmul_functions =            \
+        []<std::size_t... Is>(std::index_sequence<Is...>) {                   \
+      return std::vector<std::function<FuncType>>{func<Is>...};               \
+    }                                                                         \
+    (std::make_index_sequence<count>());                                      \
+    if (selected_config_id == -1) {                                           \
+      selected_config_id =                                                    \
+          ap::ProfileBestConfig(matmul_functions, stream_ptr, ##__VA_ARGS__); \
+    }                                                                         \
+    matmul_functions[selected_config_id](__VA_ARGS__);                        \
+  }
+
+#define AP_AUTOTUNE_half(func, stream_ptr, ...) \
+  AP_AUTOTUNE(func, stream_ptr, ap::kNumConfigsHalf, __VA_ARGS__)
+#define AP_AUTOTUNE_float(func, stream_ptr, ...) \
+  AP_AUTOTUNE(func, stream_ptr, ap::kNumConfigsFloat, __VA_ARGS__)
+#define AP_AUTOTUNE_bfloat16(func, stream_ptr, ...) \
+  AP_AUTOTUNE_half(func, stream_ptr, __VA_ARGS__)
+
+template <typename ElementT, int SwizzleFactor, bool Batched, int Id = 0>
+struct GemmTuningConfigs {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 2;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = Id;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 1> {
+  using TShape = hytlass::gemm::GemmShape<64, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 1;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 2> {
+  using TShape = hytlass::gemm::GemmShape<64, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 2;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 3> {
+  using TShape = hytlass::gemm::GemmShape<128, 64, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 3;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 4> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 4;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 5> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 5;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 6> {
+  using TShape = hytlass::gemm::GemmShape<256, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 6;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 7> {
+  using TShape = hytlass::gemm::GemmShape<256, 64, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 7;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 8> {
+  using TShape = hytlass::gemm::GemmShape<256, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 8;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 9> {
+  using TShape = hytlass::gemm::GemmShape<256, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 9;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 10> {
+  using TShape = hytlass::gemm::GemmShape<128, 32, 64>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 10;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 11> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 11;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 12> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 12;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 13> {
+  using TShape = hytlass::gemm::GemmShape<256, 64, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 13;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 14> {
+  using TShape = hytlass::gemm::GemmShape<256, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 14;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 15> {
+  using TShape = hytlass::gemm::GemmShape<32, 64, 64>;
+  using WShape = hytlass::gemm::GemmShape<16, 32, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 15;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 16> {
+  using TShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 16;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 17> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 17;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 18> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 18;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 19> {
+  using TShape = hytlass::gemm::GemmShape<64, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 6;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 19;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 20> {
+  using TShape = hytlass::gemm::GemmShape<128, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 6;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 20;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 21> {
+  using TShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 10;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 21;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 22> {
+  using TShape = hytlass::gemm::GemmShape<128, 256, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 2;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 22;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 23> {
+  using TShape = hytlass::gemm::GemmShape<128, 256, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 23;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 24> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 24;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 25> {
+  using TShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 25;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 26> {
+  using TShape = hytlass::gemm::GemmShape<64, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 26;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 27> {
+  using TShape = hytlass::gemm::GemmShape<128, 64, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 27;
+};
+
+// Specialization for float
+template <int SwizzleFactor, bool Batched, int Id>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, Id> {
+  using TShape = hytlass::gemm::GemmShape<64, 64, 16>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = Id;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 1> {
+  using TShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 1;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 2> {
+  using TShape = hytlass::gemm::GemmShape<64, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 2;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 3> {
+  using TShape = hytlass::gemm::GemmShape<64, 256, 16>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 3;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 4> {
+  using TShape = hytlass::gemm::GemmShape<64, 256, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 4;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 5> {
+  using TShape = hytlass::gemm::GemmShape<128, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 5;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 6> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 16>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 6;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 7> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 7;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 8> {
+  using TShape = hytlass::gemm::GemmShape<256, 64, 16>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 8;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 9> {
+  using TShape = hytlass::gemm::GemmShape<256, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 9;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 10> {
+  using TShape = hytlass::gemm::GemmShape<64, 128, 16>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 10;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 11> {
+  using TShape = hytlass::gemm::GemmShape<128, 64, 16>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 11;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 12> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 16>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 12;
+};
+
+struct DefaultConfig {
+  static constexpr int kConfigId = 0;
+  static constexpr int kSwizzleFactor = 1;
+  static constexpr bool kBatched = false;
+};
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/cutlass_patch/hip/hytlass_matmul.h

@@ -0,0 +1,254 @@
+// Copyright (c) 2026 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.
+
+#pragma once
+
+#include "hytlass/epilogue/thread/linear_combination_bias_elementwise.h"
+#include "hytlass/gemm/device/gemm_universal.h"
+#include "hytlass/gemm/device/gemm_universal_with_broadcast.h"
+#include "hytlass/util/device_memory.h"
+
+#include "cutlass_patch/batched_matrix_coord.h"
+#include "cutlass_patch/epilogue/thread/linear_combination_unary.h"
+#include "cutlass_patch/epilogue/thread/linear_combination_variadic.h"
+#include "cutlass_patch/gemm/device/gemm_universal_with_variadic.h"
+#include "cutlass_patch/hip/all_tuning_configs.h"
+
+#include "params.h"  // NOLINT
+
+#define CHECK_HYTLASS(status)                               \
+  {                                                         \
+    auto error = status;                                    \
+    if (error != hytlass::Status::kSuccess) {               \
+      std::cerr << "HYTLASS error = " << int(error) << " (" \
+                << hytlassGetStatusString(error) << ")"     \
+                << " at line " << __LINE__ << std::endl;    \
+      std::abort();                                         \
+    }                                                       \
+  }
+
+namespace ap {
+
+using MatrixCoord = cutlass_patch::BatchedMatrixCoord;
+using bfloat16 = __hip_bfloat16;
+
+// Operation performed by GEMM
+template <typename ElementT>
+struct GemmOperation {
+  using Type = hytlass::arch::OpMultiplyAdd;
+};
+
+template <>
+struct GemmOperation<float> {
+  using Type = hytlass::arch::OpMultiplyAddFastF32;
+};
+
+static hytlass::gemm::GemmUniversalMode GetGemmMode(int batch_count) {
+  return batch_count > 1 ? hytlass::gemm::GemmUniversalMode::kBatched
+                         : hytlass::gemm::GemmUniversalMode::kGemm;
+}
+
+static void *GetWorkspace(size_t workspace_size) {
+  static hytlass::device_memory::allocation<uint8_t> workspace;
+  if (workspace.size() < workspace_size) {
+    workspace.reset(workspace_size);
+  }
+  return workspace.get();
+}
+
+template <typename GemmFunc>
+hytlass::Status SetMaxDynamicSharedMemorySize() {
+  hipError_t hiprt_result;
+
+  // If requires more than 48KB: configure for extended, dynamic shared memory
+  if constexpr (GemmFunc::kSharedStorageSize >= (48 << 10)) {
+    hiprt_result = hipFuncSetAttribute(
+        (const void *)hytlass::Kernel2<typename GemmFunc::GemmKernel>,
+        hipFuncAttributeMaxDynamicSharedMemorySize,
+        GemmFunc::kSharedStorageSize);
+    if (hiprt_result != hipSuccess) {
+      HYTLASS_TRACE_HOST("hipFuncSetAttribute() returned error "
+                         << hipGetErrorString(hiprt_result));
+      return hytlass::Status::kErrorInternal;
+    }
+  }
+
+#if AP_ENABLE_DEBUG
+  // Update SM occupancy member
+  int sm_occupancy = -1;
+  hiprt_result = hipOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(
+      &sm_occupancy,
+      hytlass::Kernel2<typename GemmFunc::GemmKernel>,
+      GemmFunc::GemmKernel::kThreadCount,
+      GemmFunc::kSharedStorageSize,
+      hipOccupancyDisableCachingOverride);
+  if (hiprt_result != hipSuccess) {
+    HYTLASS_TRACE_HOST(
+        "hipOccupancyMaxActiveBlocksPerMultiprocessorWithFlags() returned "
+        "error "
+        << hipGetErrorString(hiprt_result));
+    return hytlass::Status::kErrorInternal;
+  }
+  HYTLASS_TRACE_HOST("sm_occupancy: (" << sm_occupancy
+                                       << ") "
+                                          "smem_size: ("
+                                       << GemmFunc::kSharedStorageSize
+                                       << ") "
+                                          "GemmKernel::kThreadCount: ("
+                                       << GemmFunc::GemmKernel::kThreadCount
+                                       << ")");
+#endif
+  return hytlass::Status::kSuccess;
+}
+
+// Convert HIP data type to hytlass data type
+template <typename T>
+struct HytlassDataType {
+  using Type = T;
+};
+
+template <>
+struct HytlassDataType<half> {
+  using Type = hytlass::half_t;
+};
+
+template <>
+struct HytlassDataType<__hip_bfloat16> {
+  using Type = hytlass::bfloat16_t;
+};
+
+// Convert to hytlass layout
+template <bool Transposed>
+struct MatrixLayout {
+  using Type = hytlass::layout::RowMajor;
+};
+
+template <>
+struct MatrixLayout<true> {
+  using Type = hytlass::layout::ColumnMajor;
+};
+
+template <typename ElementT,
+          typename ElementComputeT,
+          template <typename T>
+          class VariadicFunctor,
+          int AlignA = 128 / hytlass::sizeof_bits<ElementT>::value,
+          int AlignB = 128 / hytlass::sizeof_bits<ElementT>::value,
+          int ConfigId = DefaultConfig::kConfigId,
+          int SwizzleFactor = DefaultConfig::kSwizzleFactor,
+          bool Batched = DefaultConfig::kBatched>
+void MatmulAddVariadic(
+    const GemmEpilogueParams &params,
+    const typename VariadicFunctor<ElementComputeT>::Arguments &variadic_args) {
+  // <- data type of accumulator
+  using ElementAccumulator = typename HytlassDataType<ElementComputeT>::Type;
+  // <- data type of epilogue operations
+  using ElementComputeEpilogue = ElementAccumulator;
+  // <- data type of elements in input matrix A
+  using ElementInputA = typename HytlassDataType<ElementT>::Type;
+  // <- data type of elements in input matrix B
+  using ElementInputB = typename HytlassDataType<ElementT>::Type;
+  // <- data type of elements in output matrix D
+  using ElementOutput = typename HytlassDataType<ElementT>::Type;
+
+  constexpr int AlignC = AlignB;
+
+  // Epilogue operation as LinearCombination:
+  //  alpha * accumulator + beta * source
+  using EpilogueOutputOp =
+      cutlass_patch::epilogue::thread::LinearCombinationVariadic<
+          VariadicFunctor,
+          ElementOutput,
+          AlignC,
+          ElementAccumulator,
+          ElementComputeEpilogue,
+          hytlass::epilogue::thread::ScaleType::NoBetaScaling>;
+
+  using GemmFunc = cutlass_patch::gemm::device::GemmUniversalWithVariadic<
+      ElementInputA,
+      hytlass::layout::RowMajor,
+      ElementInputB,
+      hytlass::layout::RowMajor,
+      ElementOutput,
+      hytlass::layout::RowMajor,
+      ElementAccumulator,
+      hytlass::arch::OpClassTensorOp,
+      hytlass::arch::Gfx928,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          TShape,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          WShape,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          IShape,
+      EpilogueOutputOp,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          SwizzleThreadBlock,
+      GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::kNumStages,
+      AlignA,
+      AlignB,
+      typename GemmOperation<ElementT>::Type>;
+
+  CHECK_HYTLASS(SetMaxDynamicSharedMemorySize<GemmFunc>());
+
+  /// Arguments
+  hytlass::gemm::GemmCoord problem_size{params.m, params.n, params.k};
+
+  const ElementInputA *input =
+      reinterpret_cast<const ElementInputA *>(params.input);
+  const ElementInputB *weight =
+      reinterpret_cast<const ElementInputB *>(params.weight);
+  const ElementOutput *bias =
+      reinterpret_cast<const ElementOutput *>(params.bias);
+  ElementOutput *output = reinterpret_cast<ElementOutput *>(params.output);
+
+  ElementComputeEpilogue alpha = static_cast<ElementComputeEpilogue>(1);
+  ElementComputeEpilogue beta = bias ? static_cast<ElementComputeEpilogue>(1)
+                                     : static_cast<ElementComputeEpilogue>(0);
+
+  typename GemmFunc::Arguments arguments{
+      GetGemmMode(params.batch_count),
+      problem_size,        // <- problem size of matrix multiplication
+      params.batch_count,  // <- batch_count or k-dimension split factor
+      {alpha, beta, variadic_args},  // <- epilogue params, alpha, beta
+      input,                         // <- input, ptr_A, A, shape={M, K}
+      weight,                        // <- input, ptr_B, B, shape={K, N}
+      bias,                          // <- input, ptr_C, shape={M, N} or {1, N}
+      output,                        // <- output, ptr_D, Z, shape={M, N}
+      params.shape_args.batch_stride_A,
+      params.shape_args.batch_stride_B,
+      params.shape_args.batch_stride_C,
+      params.shape_args.batch_stride_D,
+      params.shape_args.lda,
+      params.shape_args.ldb,
+      params.shape_args.ldc_bias,
+      params.shape_args.ldd};
+
+  size_t workspace_size = GemmFunc::get_workspace_size(arguments);
+  void *workspace = workspace_size > 0 ? GetWorkspace(workspace_size) : nullptr;
+
+  GemmFunc device_gemm;
+
+  hipStream_t *stream_ptr = reinterpret_cast<hipStream_t *>(params.stream_ptr);
+
+  CHECK_HYTLASS(device_gemm.can_implement(arguments));
+  CHECK_HYTLASS(device_gemm.initialize(arguments, workspace, *stream_ptr));
+
+  // Run the GEMM
+  CHECK_HYTLASS(device_gemm(*stream_ptr));
+#if AP_ENABLE_DEBUG
+  CHECK_HIP(hipStreamSynchronize(*stream_ptr));
+#endif
+}
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/cutlass_patch/trace_device.h

@@ -0,0 +1,73 @@
+// 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.
+
+#pragma once
+
+#if CUTLASS_DEBUG_TRACE_LEVEL
+
+#ifndef CUTLASS_TRACE_DEVICE
+#define CUTLASS_TRACE_DEVICE(format, ...)                        \
+  {                                                              \
+    if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 && \
+        threadIdx.x == 0 && threadIdx.y == 0) {                  \
+      printf("[DEVICE][%s:%d, %s]" format "\n",                  \
+             __FILE__,                                           \
+             __LINE__,                                           \
+             __FUNCTION__,                                       \
+             ##__VA_ARGS__);                                     \
+    }                                                            \
+  }
+#endif
+
+#ifndef CUTLASS_TRACE_DEVICE_TID_DETAIL
+#define CUTLASS_TRACE_DEVICE_TID_DETAIL(bidz, bidx, tidx, format, ...)    \
+  {                                                                       \
+    if (blockIdx.x == bidx && blockIdx.y == 0 && blockIdx.z == bidz &&    \
+        threadIdx.x == tidx && threadIdx.y == 0) {                        \
+      printf("[DEVICE][%s:%d, %s][bid={%d,%d,%d}, tid={%d,%d,%d}]" format \
+             "\n",                                                        \
+             __FILE__,                                                    \
+             __LINE__,                                                    \
+             __FUNCTION__,                                                \
+             blockIdx.x,                                                  \
+             blockIdx.y,                                                  \
+             blockIdx.z,                                                  \
+             threadIdx.x,                                                 \
+             threadIdx.y,                                                 \
+             threadIdx.z,                                                 \
+             ##__VA_ARGS__);                                              \
+    }                                                                     \
+  }
+#endif
+
+#ifndef CUTLASS_TRACE_DEVICE_TID
+#define CUTLASS_TRACE_DEVICE_TID(format, ...)                       \
+  {                                                                 \
+    CUTLASS_TRACE_DEVICE_TID_DETAIL(0, 0, 0, format, ##__VA_ARGS__) \
+    CUTLASS_TRACE_DEVICE_TID_DETAIL(0, 0, 1, format, ##__VA_ARGS__) \
+    CUTLASS_TRACE_DEVICE_TID_DETAIL(0, 1, 0, format, ##__VA_ARGS__) \
+  }
+#endif
+
+#else
+
+#ifndef CUTLASS_TRACE_DEVICE
+#define CUTLASS_TRACE_DEVICE(format, ...)
+#endif
+
+#ifndef CUTLASS_TRACE_DEVICE_TID
+#define CUTLASS_TRACE_DEVICE_TID(format, ...)
+#endif
+
+#endif

python/paddle/apy/matmul_pass/matmul/generate_configs.py

@@ -0,0 +1,298 @@
+# 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.
+
+head_template = """// auto-generated by generate_configs.py
+
+#pragma once
+
+#include "cutlass/gemm_coord.h"
+
+namespace ap {
+
+constexpr int kNumConfigsHalf = ${num_configs_fp16};
+constexpr int kNumConfigsFloat = ${num_configs_fp32};
+
+template <int SwizzleFactor, bool Batched> struct SwizzleWrapper {
+  using Type =
+      cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<SwizzleFactor>;
+};
+
+// template <int SwizzleFactor>
+// struct SwizzleWrapper<SwizzleFactor, true> {
+//   using Type =
+//       cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle;
+// };
+"""
+
+autotune_wrapper_template = """
+#define AP_AUTOTUNE_${datatype}(func, stream, ...) { \\
+  using FuncType = decltype(func<0>); \\
+  static int selected_config_id = -1; \\
+  static std::vector<std::function<FuncType>> \\
+      matmul_functions = { \\
+          ${repeat_functions} \\
+          };  \\
+  if (selected_config_id == -1) { \\
+    selected_config_id = ap::ProfileBestConfig(matmul_functions, stream, ##__VA_ARGS__); \\
+  } \\
+  matmul_functions[selected_config_id](__VA_ARGS__); \\
+}
+"""
+
+fp16_config_template_0 = """
+template <typename ElementT, int SwizzleFactor, bool Batched, int Id = 0>
+struct GemmTuningConfigs {
+  using TShape = cutlass::gemm::GemmShape<${tshape}>;
+  using WShape = cutlass::gemm::GemmShape<${wshape}>;
+  using IShape = cutlass::gemm::GemmShape<${ishape}>;
+  static constexpr int kNumStages = ${stages};
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = Id;
+};
+"""
+
+fp16_config_template = """
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ${config_id}> {
+  using TShape = cutlass::gemm::GemmShape<${tshape}>;
+  using WShape = cutlass::gemm::GemmShape<${wshape}>;
+  using IShape = cutlass::gemm::GemmShape<${ishape}>;
+  static constexpr int kNumStages = ${stages};
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = ${config_id};
+};
+"""
+
+fp32_config_template_0 = """
+// Specialization for float
+template <int SwizzleFactor, bool Batched, int Id>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, Id> {
+  using TShape = cutlass::gemm::GemmShape<${tshape}>;
+  using WShape = cutlass::gemm::GemmShape<${wshape}>;
+  using IShape = cutlass::gemm::GemmShape<${ishape}>;
+  static constexpr int kNumStages = ${stages};
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = Id;
+};
+"""
+
+fp32_config_template = """
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, ${config_id}> {
+  using TShape = cutlass::gemm::GemmShape<${tshape}>;
+  using WShape = cutlass::gemm::GemmShape<${wshape}>;
+  using IShape = cutlass::gemm::GemmShape<${ishape}>;
+  static constexpr int kNumStages = ${stages};
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = ${config_id};
+};
+"""
+
+tail_code_str = """
+} // namespace ap
+"""
+
+
+class GemmTuningConfig:
+    def __init__(self, tshape, wshape, ishape, stages, level):
+        self.tshape = tshape
+        self.wshape = wshape
+        self.ishape = ishape
+        self.stages = stages
+        self.level = level
+
+    def __eq__(self, other):
+        def check_shape(s1, s2):
+            assert len(s1) == len(s2)
+            res = True
+            for i in range(len(s1)):
+                if s1[i] != s2[i]:
+                    res = False
+                    break
+            return res
+
+        res = check_shape(self.tshape, other.tshape)
+        res = res and check_shape(self.wshape, other.wshape)
+        res = res and check_shape(self.ishape, other.ishape)
+        res = res and self.stages == other.stages
+        return res
+
+
+def all_configs_sm80_fp16():
+    all_tuning_configs_fp16 = [
+        GemmTuningConfig([16, 128, 64], [16, 32, 64], [16, 8, 16], 2, 2),
+        GemmTuningConfig([32, 128, 64], [32, 32, 64], [16, 8, 16], 2, 2),
+        GemmTuningConfig([64, 128, 64], [64, 64, 64], [16, 8, 16], 2, 2),
+        GemmTuningConfig([128, 128, 64], [64, 64, 64], [16, 8, 16], 2, 2),
+        GemmTuningConfig([128, 128, 64], [128, 32, 64], [16, 8, 16], 2, 2),
+        GemmTuningConfig([128, 256, 64], [64, 64, 64], [16, 8, 16], 2, 2),
+        GemmTuningConfig([256, 128, 64], [64, 64, 64], [16, 8, 16], 2, 1),
+        GemmTuningConfig([16, 128, 64], [16, 32, 64], [16, 8, 16], 3, 3),
+        GemmTuningConfig([32, 128, 64], [32, 32, 64], [16, 8, 16], 3, 2),
+        GemmTuningConfig([64, 128, 64], [32, 64, 64], [16, 8, 16], 3, 1),
+        GemmTuningConfig([64, 128, 64], [64, 64, 64], [16, 8, 16], 3, 1),
+        GemmTuningConfig([64, 256, 64], [64, 64, 64], [16, 8, 16], 3, 3),
+        GemmTuningConfig([128, 64, 64], [64, 32, 64], [16, 8, 16], 3, 1),
+        GemmTuningConfig([128, 128, 32], [64, 64, 32], [16, 8, 16], 3, 1),
+        GemmTuningConfig([128, 128, 64], [64, 64, 64], [16, 8, 16], 3, 1),
+        GemmTuningConfig([128, 128, 64], [128, 32, 64], [16, 8, 16], 3, 2),
+        GemmTuningConfig([128, 256, 32], [64, 64, 32], [16, 8, 16], 3, 2),
+        GemmTuningConfig([128, 256, 64], [64, 64, 64], [16, 8, 16], 3, 2),
+        GemmTuningConfig([256, 64, 32], [64, 64, 32], [16, 8, 16], 3, 1),
+        GemmTuningConfig([256, 64, 64], [64, 64, 64], [16, 8, 16], 3, 1),
+        GemmTuningConfig([256, 128, 32], [64, 64, 32], [16, 8, 16], 3, 1),
+        GemmTuningConfig([256, 128, 64], [64, 64, 64], [16, 8, 16], 3, 1),
+        GemmTuningConfig([16, 128, 64], [16, 32, 64], [16, 8, 16], 4, 3),
+        GemmTuningConfig([32, 128, 64], [32, 32, 64], [16, 8, 16], 4, 2),
+        GemmTuningConfig([64, 128, 64], [64, 64, 64], [16, 8, 16], 4, 2),
+        GemmTuningConfig([64, 256, 32], [64, 64, 32], [16, 8, 16], 4, 2),
+        GemmTuningConfig([64, 256, 64], [64, 64, 64], [16, 8, 16], 4, 3),
+        GemmTuningConfig([128, 32, 64], [32, 32, 64], [16, 8, 16], 4, 1),
+        GemmTuningConfig([128, 128, 32], [64, 64, 32], [16, 8, 16], 4, 1),
+        GemmTuningConfig([128, 128, 64], [64, 64, 64], [16, 8, 16], 4, 1),
+        GemmTuningConfig([128, 128, 64], [128, 32, 64], [16, 8, 16], 4, 2),
+        GemmTuningConfig([256, 64, 64], [64, 64, 64], [16, 8, 16], 4, 1),
+        GemmTuningConfig([256, 64, 32], [64, 64, 32], [16, 8, 16], 4, 1),
+        GemmTuningConfig([16, 64, 64], [16, 32, 64], [16, 8, 16], 5, 2),
+        GemmTuningConfig([16, 128, 64], [16, 32, 64], [16, 8, 16], 5, 3),
+        GemmTuningConfig([32, 64, 64], [16, 32, 64], [16, 8, 16], 5, 1),
+        GemmTuningConfig([32, 128, 64], [32, 32, 64], [16, 8, 16], 5, 3),
+        GemmTuningConfig([64, 64, 64], [32, 32, 64], [16, 8, 16], 5, 1),
+        GemmTuningConfig([64, 128, 64], [64, 64, 64], [16, 8, 16], 5, 3),
+        GemmTuningConfig([128, 128, 32], [64, 64, 32], [16, 8, 16], 5, 1),
+        GemmTuningConfig([128, 128, 64], [64, 64, 64], [16, 8, 16], 5, 1),
+        GemmTuningConfig([128, 128, 64], [128, 32, 64], [16, 8, 16], 5, 2),
+        GemmTuningConfig([64, 128, 32], [32, 64, 32], [16, 8, 16], 6, 1),
+        GemmTuningConfig([128, 64, 32], [64, 32, 32], [16, 8, 16], 6, 1),
+        GemmTuningConfig([128, 32, 32], [32, 32, 32], [16, 8, 16], 7, 1),
+        GemmTuningConfig([64, 64, 32], [32, 32, 32], [16, 8, 16], 10, 1),
+    ]
+    return all_tuning_configs_fp16
+
+
+def all_configs_sm80_fp32():
+    all_tuning_configs_fp32 = [
+        GemmTuningConfig([64, 64, 16], [32, 32, 16], [16, 8, 8], 3, 1),
+        GemmTuningConfig([64, 64, 32], [32, 32, 32], [16, 8, 8], 3, 1),
+        GemmTuningConfig([64, 128, 32], [32, 64, 32], [16, 8, 8], 3, 1),
+        GemmTuningConfig([64, 256, 16], [32, 64, 16], [16, 8, 8], 3, 1),
+        GemmTuningConfig([64, 256, 32], [32, 64, 32], [16, 8, 8], 3, 1),
+        GemmTuningConfig([128, 64, 32], [64, 32, 32], [16, 8, 8], 3, 1),
+        GemmTuningConfig([128, 128, 16], [32, 64, 16], [16, 8, 8], 3, 1),
+        GemmTuningConfig([128, 128, 32], [32, 64, 32], [16, 8, 8], 3, 1),
+        GemmTuningConfig([256, 64, 16], [64, 32, 16], [16, 8, 8], 3, 1),
+        GemmTuningConfig([256, 64, 32], [64, 32, 32], [16, 8, 8], 3, 1),
+        GemmTuningConfig([64, 128, 16], [32, 64, 16], [16, 8, 8], 4, 1),
+        GemmTuningConfig([128, 64, 16], [64, 32, 16], [16, 8, 8], 4, 1),
+        GemmTuningConfig([128, 128, 16], [32, 64, 16], [16, 8, 8], 4, 1),
+    ]
+    return all_tuning_configs_fp32
+
+
+def generate_autotune_wrapper(datatype, num_configs):
+    repeat_func_strs = []
+    for i in range(num_configs):
+        repeat_func_strs.append(f"func<{i}>")
+    code_str = autotune_wrapper_template.replace(
+        "${datatype}", datatype
+    ).replace("${repeat_functions}", ", \\\n          ".join(repeat_func_strs))
+    return code_str
+
+
+def get_configs(all_configs_list, level=3):
+    consigs_list = []
+    for i in range(len(all_configs_list)):
+        if all_configs_list[i].level <= level:
+            already_have = False
+            for config in consigs_list:
+                if config == all_configs_list[i]:
+                    print(f"-- The {i}-th config is repeat.")
+                    already_have = True
+                    break
+            if not already_have:
+                consigs_list.append(all_configs_list[i])
+    return consigs_list
+
+
+def generate_configs(configs_list, config_template_0, config_template):
+    code_str = ""
+    config_id = 0
+    for config in configs_list:
+        if config_id == 0:
+            config_code_str = (
+                config_template_0.replace(
+                    "${tshape}", ", ".join(map(str, config.tshape))
+                )
+                .replace("${wshape}", ", ".join(map(str, config.wshape)))
+                .replace("${ishape}", ", ".join(map(str, config.ishape)))
+                .replace("${stages}", str(config.stages))
+            )
+        else:
+            config_code_str = (
+                config_template.replace(
+                    "${tshape}", ", ".join(map(str, config.tshape))
+                )
+                .replace("${wshape}", ", ".join(map(str, config.wshape)))
+                .replace("${ishape}", ", ".join(map(str, config.ishape)))
+                .replace("${stages}", str(config.stages))
+                .replace("${config_id}", str(config_id))
+            )
+        code_str += config_code_str
+        config_id += 1
+    return len(configs_list), code_str
+
+
+def main():
+    level = 1
+    num_fp16_configs, fp16_configs_code_str = generate_configs(
+        configs_list=get_configs(all_configs_sm80_fp16(), level=level),
+        config_template_0=fp16_config_template_0,
+        config_template=fp16_config_template,
+    )
+    num_fp32_configs, fp32_configs_code_str = generate_configs(
+        configs_list=get_configs(all_configs_sm80_fp32(), level=level),
+        config_template_0=fp32_config_template_0,
+        config_template=fp32_config_template,
+    )
+    print(
+        f"-- Total {num_fp16_configs} fp16 configs, {num_fp32_configs} fp32 configs"
+    )
+    head_code_str = head_template.replace(
+        "${num_configs_fp16}", str(num_fp16_configs)
+    ).replace("${num_configs_fp32}", str(num_fp32_configs))
+    fp16_autotune_wrapper_code_str = generate_autotune_wrapper(
+        "half", num_fp16_configs
+    )
+    fp32_autotune_wrapper_code_str = generate_autotune_wrapper(
+        "float", num_fp32_configs
+    )
+    with open("all_tuning_configs.h", "w") as f:
+        f.write(head_code_str)
+        f.write(fp16_autotune_wrapper_code_str)
+        f.write(fp32_autotune_wrapper_code_str)
+        f.write(fp16_configs_code_str)
+        f.write(fp32_configs_code_str)
+        f.write(tail_code_str)
+
+
+if __name__ == "__main__":
+    main()

python/paddle/apy/matmul_pass/matmul/math_function.h

@@ -0,0 +1,39 @@
+// 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.
+
+#pragma once
+
+#ifdef __NVCC__
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#endif
+#ifdef __HIPCC__
+#include <hip/hip_bfloat16.h>
+#include <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+#endif
+
+namespace ap {
+
+template <typename T>
+__forceinline__ __host__ __device__ T ComputePow(T base, T exponent) {
+  T res = (exponent == static_cast<T>(2))
+              ? (base * base)
+              : ((exponent == static_cast<T>(3)) ? (base * base * base)
+                                                 : (powf(base, exponent)));
+  return res;
+}
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/matmul.h

@@ -0,0 +1,63 @@
+// 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.
+
+#pragma once
+
+#include <iostream>
+#include <map>
+#include <vector>
+
+#ifdef __NVCC__
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+#define CHECK_CUDA(func)                                                      \
+  {                                                                           \
+    cudaError_t err = func;                                                   \
+    if (err != cudaSuccess) {                                                 \
+      std::cerr << "[" << __FILE__ << ":" << __LINE__ << ", " << __FUNCTION__ \
+                << "] "                                                       \
+                << "CUDA error(" << err << "), " << cudaGetErrorString(err)   \
+                << " when call " << #func << std::endl;                       \
+      exit(EXIT_FAILURE);                                                     \
+    }                                                                         \
+  }
+
+#include "cutlass_patch/cuda/cutlass_matmul.cuh"  // NOLINT
+#include "math_function.h"                        // NOLINT
+#include "profile.h"                              // NOLINT
+#endif
+
+#ifdef __HIPCC__
+#include <hip/hip_bfloat16.h>
+#include <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+
+#define CHECK_HIP(func)                                                       \
+  {                                                                           \
+    hipError_t err = func;                                                    \
+    if (err != hipSuccess) {                                                  \
+      std::cerr << "[" << __FILE__ << ":" << __LINE__ << ", " << __FUNCTION__ \
+                << "] "                                                       \
+                << "HIP error(" << err << "), " << hipGetErrorString(err)     \
+                << " when call " << #func << std::endl;                       \
+      exit(EXIT_FAILURE);                                                     \
+    }                                                                         \
+  }
+
+#include "cutlass_patch/hip/hytlass_matmul.h"  // NOLINT
+#include "math_function.h"                     // NOLINT
+#include "profile.h"                           // NOLINT
+#endif

python/paddle/apy/matmul_pass/matmul/params.h

@@ -0,0 +1,194 @@
+// Copyright (c) 2026 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.
+
+#pragma once
+
+#include <cstdint>
+#include <iostream>
+#include <limits>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#define ASSERT_CHECK(__cond)                                            \
+  do {                                                                  \
+    const bool __cond_var = (__cond);                                   \
+    if (!__cond_var) {                                                  \
+      ::std::string __err_msg = ::std::string("`") + #__cond +          \
+                                "` check failed at " + __FILE__ + ":" + \
+                                ::std::to_string(__LINE__);             \
+      throw std::runtime_error(__err_msg);                              \
+    }                                                                   \
+  } while (0)
+
+namespace ap {
+
+inline int CheckedCastToInt(int64_t value) {
+  ASSERT_CHECK(value >= 0);
+  ASSERT_CHECK(value <= static_cast<int64_t>(std::numeric_limits<int>::max()));
+  return static_cast<int>(value);
+}
+
+template <typename T, int Dim>
+struct Alignment {
+  static constexpr int kValue =
+      ((Dim % 8) == 0) ? 8
+                       : (((Dim % 4) == 0) ? 4 : (((Dim % 2) == 0) ? 2 : 1));
+};
+
+template <int Dim>
+struct Alignment<float, Dim> {
+  static constexpr int kValue =
+      ((Dim % 4) == 0) ? 4 : (((Dim % 2) == 0) ? 2 : 1);
+};
+
+struct GemmEpilogueParams {
+  int batch_count;
+  int m;
+  int n;
+  int k;
+
+  bool transpose_a;
+  bool transpose_b;
+
+  // Shape related aruguments
+  struct ShapeArguments {
+    int64_t batch_stride_A;
+    int64_t batch_stride_B;
+    int64_t batch_stride_C;
+    int64_t batch_stride_D;
+    int64_t lda;
+    int64_t ldb;
+    int64_t ldc_bias;
+    int64_t ldd;
+  };
+
+  ShapeArguments shape_args;
+
+  const void *input;
+  const void *weight;
+  const void *bias;
+  void *output;
+
+  void *stream_ptr;
+
+  std::vector<int64_t> input0_shape;
+  std::vector<int64_t> input1_shape;
+  std::vector<const void *> epilogue_in_ptrs;
+  std::vector<void *> epilogue_out_ptrs;
+  std::vector<std::vector<int64_t>> epilogue_in_shapes;
+  std::vector<std::vector<int64_t>> epilogue_out_shapes;
+
+  GemmEpilogueParams() {}
+  GemmEpilogueParams(void *stream_ptr,
+                     const void *input,
+                     const void *weight,
+                     const void *bias,
+                     void *output,
+                     const std::vector<int64_t> &input_shape,
+                     const std::vector<int64_t> &weight_shape,
+                     const std::vector<int64_t> &bias_shape,
+                     bool transpose_a = false,
+                     bool transpose_b = false)
+      : stream_ptr(stream_ptr),
+        input(input),
+        weight(weight),
+        bias(bias),
+        output(output),
+        transpose_a(transpose_a),
+        transpose_b(transpose_b) {
+    ASSERT_CHECK(input_shape.size() >= 2U);
+    ASSERT_CHECK(weight_shape.size() >= 2U);
+
+    input0_shape = input_shape;
+    input1_shape = weight_shape;
+
+    int64_t batch_count_i64 = 1;
+    for (size_t i = 0; i < input_shape.size() - 2; ++i) {
+      batch_count_i64 *= input_shape[i];
+    }
+    batch_count = CheckedCastToInt(batch_count_i64);
+
+    int64_t m_i64;
+    int64_t n_i64;
+    int64_t k_i64;
+    if (transpose_a) {
+      m_i64 = input_shape[input_shape.size() - 1];
+      k_i64 = input_shape[input_shape.size() - 2];
+    } else {
+      m_i64 = input_shape[input_shape.size() - 2];
+      k_i64 = input_shape[input_shape.size() - 1];
+    }
+    if (transpose_b) {
+      ASSERT_CHECK(weight_shape[weight_shape.size() - 1] == k_i64);
+      n_i64 = weight_shape[weight_shape.size() - 2];
+    } else {
+      ASSERT_CHECK(weight_shape[weight_shape.size() - 2] == k_i64);
+      n_i64 = weight_shape[weight_shape.size() - 1];
+    }
+    m = CheckedCastToInt(m_i64);
+    n = CheckedCastToInt(n_i64);
+    k = CheckedCastToInt(k_i64);
+
+    if (bias) {
+      ASSERT_CHECK(bias_shape.size() >= 1U);
+      ASSERT_CHECK(bias_shape[bias_shape.size() - 1] == n_i64);
+    }
+
+#if AP_ENABLE_DEBUG
+    std::cout << "-- [GemmEpilogueParams] batch_count: " << batch_count
+              << ", m: " << m << ", n: " << n << ", k: " << k << std::endl;
+    std::cout << "-- [GemmEpilogueParams] input: " << input << std::endl;
+    std::cout << "-- [GemmEpilogueParams] weight: " << weight << std::endl;
+    std::cout << "-- [GemmEpilogueParams] bias: " << bias << std::endl;
+    std::cout << "-- [GemmEpilogueParams] output: " << output << std::endl;
+    std::cout << "-- [GemmEpilogueParams] stream: " << stream << std::endl;
+#endif
+
+    shape_args.batch_stride_A = m_i64 * k_i64;
+    shape_args.batch_stride_B = (weight_shape.size() == 2) ? 0 : n_i64 * k_i64;
+    shape_args.batch_stride_D = m_i64 * n_i64;
+
+    shape_args.lda = transpose_a ? m_i64 : k_i64;
+    shape_args.ldb = transpose_b ? k_i64 : n_i64;
+    shape_args.ldd = n_i64;
+
+    bool is_C_bias = bias_shape.size() == 1UL;
+
+    /// Only available in RRR format
+    shape_args.batch_stride_C = (!bias || is_C_bias) ? 0 : m_i64 * n_i64;
+    shape_args.ldc_bias = (!bias || is_C_bias) ? 0 : n_i64;
+  }
+
+  void SetEpilogues(const std::vector<const void *> &in_ptrs,
+                    const std::vector<void *> &out_ptrs) {
+    epilogue_in_ptrs = in_ptrs;
+    epilogue_out_ptrs = out_ptrs;
+  }
+
+  void SetEpilogueAndShapes(
+      const std::vector<const void *> &in_ptrs,
+      const std::vector<std::vector<int64_t>> &in_shapes,
+      const std::vector<void *> &out_ptrs,
+      const std::vector<std::vector<int64_t>> &out_shapes) {
+    ASSERT_CHECK(in_ptrs.size() == in_shapes.size());
+    epilogue_in_ptrs = in_ptrs;
+    epilogue_in_shapes = in_shapes;
+    ASSERT_CHECK(out_ptrs.size() == out_shapes.size());
+    epilogue_out_ptrs = out_ptrs;
+    epilogue_out_shapes = out_shapes;
+  }
+};
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/profile.h

@@ -0,0 +1,151 @@
+// 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.
+
+#pragma once
+
+#include <functional>
+#include <utility>
+
+#ifdef __NVCC__
+#include <cuda_profiler_api.h>
+
+#define GPUEvent_t cudaEvent_t
+
+#ifndef GPUStream_t
+#define GPUStream_t cudaStream_t
+#endif
+
+#define GPUEventCreate(e) cudaEventCreate(e)
+#define GPUEventDestroy(e) cudaEventDestroy(e)
+#define GPUEventRecord(e, s) cudaEventRecord(e, s)
+#define GPUEventSynchronize(e) cudaEventSynchronize(e)
+#define GPUEventElapsedTime(ms, s, e) cudaEventElapsedTime(ms, s, e)
+#define GPUProfilerStart() cudaProfilerStart()
+#define GPUProfilerStop() cudaProfilerStop()
+#define GPUStreamSynchronize(s) cudaStreamSynchronize(s)
+#define CHECK_GPU CHECK_CUDA
+#endif
+
+#ifdef __HIPCC__
+#include <hip/hip_runtime.h>
+#include <hip/hip_runtime_api.h>
+
+#define GPUEvent_t hipEvent_t
+
+#ifndef GPUStream_t
+#define GPUStream_t hipStream_t
+#endif
+
+#define GPUEventCreate(e) hipEventCreate(e)
+#define GPUEventDestroy(e) hipEventDestroy(e)
+#define GPUEventRecord(e, s) hipEventRecord(e, s)
+#define GPUEventSynchronize(e) hipEventSynchronize(e)
+#define GPUEventElapsedTime(ms, s, e) hipEventElapsedTime(ms, s, e)
+#define GPUProfilerStart() hipProfilerStart()
+#define GPUProfilerStop() hipProfilerStop()
+#define GPUStreamSynchronize(s) hipStreamSynchronize(s)
+#define CHECK_GPU CHECK_HIP
+#endif
+
+namespace ap {
+
+class GpuTimer {
+ public:
+  explicit GpuTimer(bool profile) : profile_(profile) {
+    CHECK_GPU(GPUEventCreate(&start_));
+    CHECK_GPU(GPUEventCreate(&stop_));
+  }
+
+  ~GpuTimer() {
+    CHECK_GPU(GPUEventDestroy(start_));
+    CHECK_GPU(GPUEventDestroy(stop_));
+  }
+
+  void Start(GPUStream_t stream) {
+    CHECK_GPU(GPUEventRecord(start_, stream));
+    if (profile_) {
+      CHECK_GPU(GPUProfilerStart());
+    }
+  }
+
+  void Stop(GPUStream_t stream) {
+    CHECK_GPU(GPUEventRecord(stop_, stream));
+    if (profile_) {
+      CHECK_GPU(GPUProfilerStop());
+    }
+  }
+
+  float ElapsedTime() {
+    float milliseconds = 0;
+    CHECK_GPU(GPUEventSynchronize(stop_));
+    CHECK_GPU(GPUEventElapsedTime(&milliseconds, start_, stop_));
+    return milliseconds;
+  }
+
+ private:
+  bool profile_{false};
+  GPUEvent_t start_{nullptr};
+  GPUEvent_t stop_{nullptr};
+};
+
+template <typename FuncType, typename... Args>
+int ProfileBestConfig(const std::vector<FuncType> &funcs,
+                      void *stream_ptr,
+                      Args &&...args) {
+  std::cout
+      << "=================================================================="
+      << std::endl;
+
+  constexpr int kWarmupIters = 1;
+  constexpr int kRepeatIters = 100;
+
+  GpuTimer gpu_timer(false);
+  float min_time_ms = 100000.f;
+  int min_time_idx = -1;
+
+  GPUStream_t stream = *reinterpret_cast<GPUStream_t *>(stream_ptr);
+
+  for (int idx = 0; idx < funcs.size(); ++idx) {
+    auto func = funcs[idx];
+    for (int i = 0; i < kWarmupIters; i++) {
+      func(std::forward<Args>(args)...);
+    }
+    if (stream) {
+      CHECK_GPU(GPUStreamSynchronize(stream));
+    }
+
+    gpu_timer.Start(stream);
+    for (int i = 0; i < kRepeatIters; i++) {
+      func(std::forward<Args>(args)...);
+    }
+    gpu_timer.Stop(stream);
+
+    float elapsed_time_ms = gpu_timer.ElapsedTime();
+    std::cout << "-- [ProfileBestConfig] No " << idx
+              << ", elapsed_time: " << elapsed_time_ms << " ms" << std::endl;
+    if (elapsed_time_ms < min_time_ms) {
+      min_time_ms = elapsed_time_ms;
+      min_time_idx = idx;
+    }
+  }
+
+  std::cout << "-- [ProfileBestConfig] best config idx: " << min_time_idx
+            << std::endl;
+  std::cout
+      << "=================================================================="
+      << std::endl;
+  return min_time_idx;
+}
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul_epilogue_pass.py

@@ -0,0 +1,140 @@
+# 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 access_topo_drr
+import pir
+
+
+class RemoveDataOpPairPass(access_topo_drr.DrrPass):
+    def __init__(self, src_data_op_name, dst_data_op_name):
+        self.src_data_op_name = pir.a_str(src_data_op_name)
+        self.dst_data_op_name = pir.a_str(dst_data_op_name)
+
+    def source_pattern(self, o, t):
+        o.src_data_op = o.ap_native_op("pd_op.data")
+        o.src_data_op([], [t.input0])
+        o.dst_data_op = o.ap_native_op("pd_op.data")
+        o.dst_data_op([], [t.input1])
+        o.up_spider_op = o.ap_native_op("ap_op.up_spider")
+        o.up_spider_op([t.input0, t.input1], [])
+
+    def constraint(self, o, t):
+        return [o.src_data_op.name, o.dst_data_op.name] == [
+            self.src_data_op_name,
+            self.dst_data_op_name,
+        ]
+
+    def result_pattern(self, o, t):
+        pass
+
+
+class RemoveDataOp2SumOp2DataOpPass(access_topo_drr.DrrPass):
+    def __init__(self, src_data_op_name, dst_data_op_name):
+        self.src_data_op_name = pir.a_str(src_data_op_name)
+        self.dst_data_op_name = pir.a_str(dst_data_op_name)
+
+    def source_pattern(self, o, t):
+        o.src_data_op = o.ap_native_op("pd_op.data")
+        o.src_data_op.name = self.src_data_op_name
+        o.src_data_op([], [t.input0])
+        o.full_int_array_op = o.ap_native_op("pd_op.full_int_array")
+        o.full_int_array_op([], [t.axis])
+        o.sum_op = o.ap_native_op("pd_op.sum")
+        o.sum_op([t.input0, t.axis], [t.sum_out])
+        o.dst_data_op = o.ap_native_op("pd_op.data")
+        o.dst_data_op.name = self.dst_data_op_name
+        o.dst_data_op([], [t.input1])
+        o.up_spider_op = o.ap_native_op("ap_op.up_spider")
+        o.up_spider_op([t.sum_out, t.input1], [])
+
+    def result_pattern(self, o, t):
+        pass
+
+
+class RemoveElementInputIndexPass(access_topo_drr.DrrPass):
+    def __init__(self, src_data_op_name, dst_load_from_global_op_name):
+        self.src_data_op_name = pir.a_str(src_data_op_name)
+        self.dst_load_from_global_op_name = pir.a_str(
+            dst_load_from_global_op_name
+        )
+
+    def source_pattern(self, o, t):
+        o.src_data_op = o.ap_native_op("pd_op.data")
+        o.src_data_op.name = self.src_data_op_name
+        o.src_data_op([], [t.src_input])
+
+        o.dst_load_from_global_op = o.ap_native_op("ap_op.load_from_global")
+        o.dst_load_from_global_op.index_func_unique_id = (
+            self.dst_load_from_global_op_name
+        )
+        o.dst_load_from_global_op(
+            [t.dst_input], [t.dst_load_from_global_output]
+        )
+        o.up_spider_op = o.ap_native_op("ap_op.up_spider")
+        o.up_spider_op([t.src_input, t.dst_load_from_global_output], [])
+
+    def result_pattern(self, o, t):
+        pass
+
+
+class RemoveBroadcastInputIndexPass(access_topo_drr.DrrPass):
+    def __init__(self, src_data_op_name, dst_load_from_global_op_name):
+        self.src_data_op_name = pir.a_str(src_data_op_name)
+        self.dst_load_from_global_op_name = pir.a_str(
+            dst_load_from_global_op_name
+        )
+
+    def source_pattern(self, o, t):
+        o.src_data_op = o.ap_native_op("pd_op.data")
+        o.src_data_op.name = self.src_data_op_name
+        o.src_data_op([], [t.input0])
+        o.full_int_array_op = o.ap_native_op("pd_op.full_int_array")
+        o.full_int_array_op([], [t.axis])
+        o.sum_op = o.ap_native_op("pd_op.sum")
+        o.sum_op([t.input0, t.axis], [t.sum_out])
+        o.dst_load_from_global_op = o.ap_native_op("ap_op.load_from_global")
+        o.dst_load_from_global_op.index_func_unique_id = (
+            self.dst_load_from_global_op_name
+        )
+        o.dst_load_from_global_op(
+            [t.dst_input], [t.dst_load_from_global_output]
+        )
+        o.up_spider_op = o.ap_native_op("ap_op.up_spider")
+        o.up_spider_op([t.sum_out, t.dst_load_from_global_output], [])
+
+    def result_pattern(self, o, t):
+        pass
+
+
+class RemoveOutputIndexPass(access_topo_drr.DrrPass):
+    def __init__(self, src_data_op_name, dst_store_to_global_op_name):
+        self.src_data_op_name = pir.a_str(src_data_op_name)
+        self.dst_store_to_global_op_name = pir.a_str(
+            dst_store_to_global_op_name
+        )
+
+    def source_pattern(self, o, t):
+        o.src_data_op = o.ap_native_op("pd_op.data")
+        o.src_data_op.name = self.src_data_op_name
+        o.src_data_op([], [t.src_input])
+        o.down_spider_op = o.ap_native_op("ap_op.down_spider")
+        o.down_spider_op([t.src_input], [t.dst_output_val])
+        o.dst_store_to_global_op = o.ap_native_op("ap_op.store_to_global")
+        o.dst_store_to_global_op.index_func_unique_id = (
+            self.dst_store_to_global_op_name
+        )
+        o.dst_store_to_global_op([t.dst_output, t.dst_output_val], [])
+
+    def result_pattern(self, o, t):
+        pass

python/paddle/apy/matmul_pass/matmul_variadic_ptn.py

@@ -0,0 +1,673 @@
+# 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 abstract_drr
+import access_topo_drr  # noqa: F401
+import ap
+import index_program_translator_util
+import ir_tools
+import kernel_arg_id_util
+import kernel_arg_translator_util  # noqa: F401
+import low_level_ir_code_gen_ctx_util  # noqa: F401
+import matmul_epilogue_pass
+import matmul_variadic_tpl
+import op_compute_translator_util
+import op_conversion_drr_pass  # noqa: F401
+import pir  # noqa: F401
+import program_translator_util
+import topo_drr_pass
+import umprime  # noqa: F401
+
+
+class MatmulEpilogueFusion(abstract_drr.DrrPass):
+    def source_pattern(self, o, t):
+        in_num = self.number_of_inputs()
+        out_num = self.number_of_outputs()
+        o.matmul_op = o.ap_native_op("pd_op.matmul")
+        o.matmul_op([t.input0, t.input1], [t.mm_out])
+        o.trivial_op = o.ap_trivial_fusion_op()
+        o.trivial_op(
+            [
+                t.mm_out,
+                *ap.map(
+                    lambda index: getattr(t, f"input{index + 2}"),
+                    range(in_num - 2),
+                ),
+            ],
+            ap.map(lambda index: getattr(t, f"output{index}"), range(out_num)),
+        )
+
+    def result_pattern(self, o, t):
+        in_num = self.number_of_inputs()
+        out_num = self.number_of_outputs()
+        o.fustion_op = o.ap_pattern_fusion_op(self.code_gen)
+        o.fustion_op(
+            ap.map(lambda index: getattr(t, f"input{index}"), range(in_num)),
+            ap.map(lambda index: getattr(t, f"output{index}"), range(out_num)),
+        )
+
+    def constraint(self, o, t):
+        program = ir_tools.copy_fused_ops_to_program(
+            o.trivial_op, tensor_match_ctx=t
+        )
+        print("before-umprime: ", program)
+        # umprime passes
+        pass_manager = ir_tools.create_pass_manager()
+        pass_manager.add_pass(ir_tools.create_access_topo_drr_pass("umprime"))
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(program)
+        print("before-access_topo_pass", program)
+        init_pass_manager = ir_tools.create_pass_manager()
+        init_down_spider = topo_drr_pass.InitDownSpiderAccessTopoPass("mm_out")
+        init_pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(init_down_spider)
+        )
+        outputs_name_list = ap.map(
+            lambda i: f"output{i}", range(self.number_of_outputs())
+        )
+        inputs_name_list = (
+            ap.map(
+                lambda i: f"input{i + 2}", range(self.number_of_inputs() - 2)
+            )
+            if self.number_of_inputs() > 2
+            else []
+        )
+        print('inputs_name_list: ', ', '.join(inputs_name_list))
+        init_fake_data_for_yield_input = (
+            topo_drr_pass.FakeDataForYieldAccessTopoPass(outputs_name_list)
+        )
+        init_pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(
+                init_fake_data_for_yield_input
+            )
+        )
+        init_pass_manager.run(program)
+        print("after-init-access_topo_pass", program)
+        pass_manager = ir_tools.create_pass_manager()
+        pass_manager.add_pass(ir_tools.create_access_topo_drr_pass("default"))
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(program)
+        print("after-apply-access_topo_pass", program)
+        pass_manager = ir_tools.create_pass_manager()
+        ap.map(
+            lambda dst_name: pass_manager.add_pass(
+                ir_tools.create_access_topo_drr_one_step_pass(
+                    matmul_epilogue_pass.RemoveDataOpPairPass(
+                        src_data_op_name="mm_out", dst_data_op_name=dst_name
+                    )
+                )
+            ),
+            inputs_name_list,
+        )
+        ap.map(
+            lambda dst_name: pass_manager.add_pass(
+                ir_tools.create_access_topo_drr_one_step_pass(
+                    matmul_epilogue_pass.RemoveDataOp2SumOp2DataOpPass(
+                        src_data_op_name="mm_out", dst_data_op_name=dst_name
+                    )
+                )
+            ),
+            inputs_name_list,
+        )
+
+        ap.map(
+            lambda dst_name: pass_manager.add_pass(
+                ir_tools.create_access_topo_drr_one_step_pass(
+                    matmul_epilogue_pass.RemoveDataOpPairPass(
+                        src_data_op_name="mm_out", dst_data_op_name=dst_name
+                    )
+                )
+            ),
+            outputs_name_list,
+        )
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(program)
+        print("after-remove-input-output-access_topo_pass", program)
+        return program.empty()
+
+    def _insert_load_from_global(self, program, input_names):
+        init_pass_manager = ir_tools.create_pass_manager()
+
+        def AddPass(input_name):
+            ir_pass = topo_drr_pass.InitNaiveLoadFromGlobalAccessTopoPass(
+                input_name
+            )
+            init_pass_manager.add_pass(
+                ir_tools.create_access_topo_drr_one_step_pass(ir_pass)
+            )
+
+        ap.map(AddPass, input_names)
+        init_pass_manager.run(program)
+
+    def _insert_store_to_global(self, program, output_names):
+        init_pass_manager = ir_tools.create_pass_manager()
+        ir_pass = topo_drr_pass.FakeDataStoreToGlobalForYieldAccessTopoPass(
+            output_names
+        )
+        init_pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(ir_pass)
+        )
+        init_pass_manager.run(program)
+
+    def _make_kernel_arg_translator(self):
+        return matmul_variadic_tpl.make_kernel_arg_translator()
+
+    def _apply_topo_access_passes(self, mut_program, anchor_data_op_name):
+        init_pass_manager = ir_tools.create_pass_manager()
+        init_down_spider = topo_drr_pass.InitDownSpiderAccessTopoPass(
+            anchor_data_op_name
+        )
+        init_pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(init_down_spider)
+        )
+        init_pass_manager.run(mut_program)
+        pass_manager = ir_tools.create_pass_manager()
+        pass_manager.add_pass(ir_tools.create_access_topo_drr_pass("default"))
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(mut_program)
+
+    def _simplify_index_program(self, mut_program):
+        pass_manager = ir_tools.create_pass_manager()
+        drr_pass = topo_drr_pass.ConvertUpSpiderStoreDataOpToYieldOpPass()
+        pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(drr_pass)
+        )
+        drr_pass = topo_drr_pass.ConvertDownSpiderStoreDataOpToYieldOpPass()
+        pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(drr_pass)
+        )
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(mut_program)
+        return mut_program
+
+    def _make_index_func_unique_id2index_program(
+        self, compute_program, anchor_data_op_name, input_names, output_names
+    ):
+        full_index_program = compute_program.clone()
+        self._apply_topo_access_passes(full_index_program, anchor_data_op_name)
+        print('full_index_program: ', full_index_program)
+
+        def MatchAndCopyInputIndex(dst_input_name):
+            pass_manager = ir_tools.create_pass_manager()
+            removed_programs = ap.MutableList()
+            rm_elementwise_drr_pass = (
+                matmul_epilogue_pass.RemoveElementInputIndexPass(
+                    src_data_op_name=anchor_data_op_name,
+                    dst_load_from_global_op_name=dst_input_name,
+                )
+            )
+            rm_elementwise_ir_pass = (
+                ir_tools.create_access_topo_drr_one_step_pass(
+                    rm_elementwise_drr_pass,
+                    matched_pattern_mut_list=removed_programs,
+                )
+            )
+            pass_manager.add_pass(rm_elementwise_ir_pass)
+            rm_broadcast_drr_pass = (
+                matmul_epilogue_pass.RemoveBroadcastInputIndexPass(
+                    src_data_op_name=anchor_data_op_name,
+                    dst_load_from_global_op_name=dst_input_name,
+                )
+            )
+            rm_broadcast_ir_pass = (
+                ir_tools.create_access_topo_drr_one_step_pass(
+                    rm_broadcast_drr_pass,
+                    matched_pattern_mut_list=removed_programs,
+                )
+            )
+            pass_manager.add_pass(rm_broadcast_ir_pass)
+            pass_manager.run(full_index_program)
+
+            def Converter(program):
+                return [dst_input_name, self._simplify_index_program(program)]
+
+            return ap.map(Converter, removed_programs)
+
+        input_and_index_programs = ap.flat_map(
+            MatchAndCopyInputIndex, input_names
+        )
+
+        def MatchAndCopyOutputIndex(dst_output_name):
+            print('full_index_program output: ', full_index_program)
+            pass_manager = ir_tools.create_pass_manager()
+            removed_programs = ap.MutableList()
+            drr_pass = matmul_epilogue_pass.RemoveOutputIndexPass(
+                src_data_op_name=anchor_data_op_name,
+                dst_store_to_global_op_name=dst_output_name,
+            )
+            ir_pass = ir_tools.create_access_topo_drr_one_step_pass(
+                drr_pass, matched_pattern_mut_list=removed_programs
+            )
+            pass_manager.add_pass(ir_pass)
+            pass_manager.run(full_index_program)
+
+            def Converter(program):
+                return [dst_output_name, self._simplify_index_program(program)]
+
+            print('len removed of output: ', len(removed_programs))
+            return ap.map(Converter, removed_programs)
+
+        output_and_index_programs = ap.flat_map(
+            MatchAndCopyOutputIndex, output_names
+        )
+        return ap.OrderedDict(
+            [*input_and_index_programs, *output_and_index_programs]
+        )
+
+    def _replace_with_load_from_register(
+        self, mut_program, load_ir_value_name, register_var_name
+    ):
+        pass_manager = ir_tools.create_pass_manager()
+        drr_pass = topo_drr_pass.ReplaceWithLoadFromRegisterPass(
+            name=load_ir_value_name, register_var_name=register_var_name
+        )
+        pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(drr_pass)
+        )
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(mut_program)
+        return mut_program
+
+    def _replace_with_store_to_register(
+        self, mut_program, store_ir_value_name, register_var_name
+    ):
+        pass_manager = ir_tools.create_pass_manager()
+        drr_pass = topo_drr_pass.ReplaceWithStoreToRegisterPass(
+            name=store_ir_value_name, register_var_name=register_var_name
+        )
+        pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(drr_pass)
+        )
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(mut_program)
+        return mut_program
+
+    def _get_program_translator(self, ctx, o, t):
+        outputs_name_list = ap.map(
+            lambda i: f"output{i}", range(self.number_of_outputs())
+        )
+        other_outputs_name_list = ap.map(
+            lambda i: f"output{i + 1}", range(self.number_of_outputs() - 1)
+        )
+        local_outputs_name_list = ap.map(
+            lambda i: f"out{i}", range(self.number_of_outputs())
+        )
+        inputs_name_list = (
+            ap.map(
+                lambda i: f"input{i + 2}", range(self.number_of_inputs() - 2)
+            )
+            if self.number_of_inputs() > 2
+            else []
+        )
+        mut_program = ir_tools.copy_fused_ops_to_program(
+            o.trivial_op, tensor_match_ctx=t
+        )
+        print("before-umprime: ", mut_program)
+        pass_manager = ir_tools.create_pass_manager()
+        pass_manager.add_pass(ir_tools.create_access_topo_drr_pass("umprime"))
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(mut_program)
+        self._insert_load_from_global(mut_program, input_names=["mm_out"])
+        self._insert_load_from_global(mut_program, input_names=inputs_name_list)
+        self._insert_store_to_global(
+            mut_program, output_names=outputs_name_list
+        )
+        kernel_arg_translator = self._make_kernel_arg_translator()
+        index_func_unique_id2index_program = (
+            self._make_index_func_unique_id2index_program(
+                mut_program,
+                anchor_data_op_name="mm_out",
+                input_names=inputs_name_list,
+                output_names=other_outputs_name_list,
+            )
+        )
+        print(
+            "index_func_unique_id2index_program:\n",
+            index_func_unique_id2index_program,
+        )
+        index_program_translator_map = index_program_translator_util.IndexProgramTranslatorMap(
+            index_func_unique_id2index_program=index_func_unique_id2index_program,
+            kernel_arg_translator=kernel_arg_translator,
+            anchor_iter_var_names=matmul_variadic_tpl.get_anchor_iter_var_names(),
+        )
+        self._replace_with_load_from_register(
+            mut_program, load_ir_value_name="mm_out", register_var_name="x"
+        )
+        self._replace_with_store_to_register(mut_program, "output0", "out")
+        print("mut_program:", mut_program)
+        op_compute_translator_maker = (
+            op_compute_translator_util.OpComputeTranslatorFactory()
+        )
+        program_translator = program_translator_util.ProgramTranslator(
+            program_property=mut_program.copy_to_const_program_data(),
+            kernel_arg_translator=kernel_arg_translator,
+            index_program_translator_map=index_program_translator_map,
+            op_translator_maker=op_compute_translator_maker,
+        )
+
+        return program_translator
+
+    def code_gen(self, ctx, o, t):
+        program_translator = self._get_program_translator(ctx, o, t)
+        mut_kernel_arg_id_registry = kernel_arg_id_util.KernelArgIdNameRegistry(
+            code_gen_ctx=ctx, tensor_match_ctx=t, name_prefix=""
+        )
+
+        template_module = matmul_variadic_tpl.MatmulVariadicTemplate(
+            program_translator=program_translator,
+            mut_kernel_arg_id_registry=mut_kernel_arg_id_registry,
+        )
+
+        def get_symbolic_shape_args_list(sym_dim):
+            return ctx.dim_expr_kernel_arg_id(sym_dim)
+
+        input0_shape_kargs = ap.map(
+            get_symbolic_shape_args_list, t.input0.symbolic_shape_to_list()
+        )
+        input1_shape_kargs = ap.map(
+            get_symbolic_shape_args_list, t.input1.symbolic_shape_to_list()
+        )
+        return template_module.compile(
+            input0_karg=ctx.in_tensor_data_ptr_kernel_arg_id(t.input0),
+            input1_karg=ctx.in_tensor_data_ptr_kernel_arg_id(t.input1),
+            output_karg=ctx.out_tensor_data_ptr_kernel_arg_id(t.output0),
+            input0_shape_kargs=input0_shape_kargs,
+            input1_shape_kargs=input1_shape_kargs,
+        )
+
+
+class NumberOfInputsTrait0:
+    def number_of_inputs(self):
+        return 0
+
+
+class NumberOfInputsTrait1:
+    def number_of_inputs(self):
+        return 1
+
+
+class NumberOfInputsTrait2:
+    def number_of_inputs(self):
+        return 2
+
+
+class NumberOfInputsTrait3:
+    def number_of_inputs(self):
+        return 3
+
+
+class NumberOfInputsTrait4:
+    def number_of_inputs(self):
+        return 4
+
+
+class NumberOfInputsTrait5:
+    def number_of_inputs(self):
+        return 5
+
+
+class NumberOfInputsTrait6:
+    def number_of_inputs(self):
+        return 6
+
+
+class NumberOfInputsTrait7:
+    def number_of_inputs(self):
+        return 7
+
+
+class NumberOfInputsTrait8:
+    def number_of_inputs(self):
+        return 8
+
+
+class NumberOfInputsTrait9:
+    def number_of_inputs(self):
+        return 9
+
+
+class NumberOfInputsTrait10:
+    def number_of_inputs(self):
+        return 10
+
+
+class NumberOfInputsTrait11:
+    def number_of_inputs(self):
+        return 11
+
+
+class NumberOfInputsTrait12:
+    def number_of_inputs(self):
+        return 12
+
+
+class NumberOfInputsTrait13:
+    def number_of_inputs(self):
+        return 13
+
+
+class NumberOfInputsTrait14:
+    def number_of_inputs(self):
+        return 14
+
+
+class NumberOfInputsTrait15:
+    def number_of_inputs(self):
+        return 15
+
+
+class NumberOfInputsTrait16:
+    def number_of_inputs(self):
+        return 16
+
+
+class NumberOfInputsTrait17:
+    def number_of_inputs(self):
+        return 17
+
+
+class NumberOfOutputsTrait0:
+    def number_of_outputs(self):
+        return 0
+
+
+class NumberOfOutputsTrait1:
+    def number_of_outputs(self):
+        return 1
+
+
+class NumberOfOutputsTrait2:
+    def number_of_outputs(self):
+        return 2
+
+
+class NumberOfOutputsTrait3:
+    def number_of_outputs(self):
+        return 3
+
+
+class NumberOfOutputsTrait4:
+    def number_of_outputs(self):
+        return 4
+
+
+class NumberOfOutputsTrait5:
+    def number_of_outputs(self):
+        return 5
+
+
+class NumberOfOutputsTrait6:
+    def number_of_outputs(self):
+        return 6
+
+
+class NumberOfOutputsTrait7:
+    def number_of_outputs(self):
+        return 7
+
+
+class NumberOfOutputsTrait8:
+    def number_of_outputs(self):
+        return 8
+
+
+class NumberOfOutputsTrait9:
+    def number_of_outputs(self):
+        return 9
+
+
+class NumberOfOutputsTrait10:
+    def number_of_outputs(self):
+        return 10
+
+
+class NumberOfOutputsTrait11:
+    def number_of_outputs(self):
+        return 11
+
+
+class NumberOfOutputsTrait12:
+    def number_of_outputs(self):
+        return 12
+
+
+class NumberOfOutputsTrait13:
+    def number_of_outputs(self):
+        return 13
+
+
+class NumberOfOutputsTrait14:
+    def number_of_outputs(self):
+        return 14
+
+
+class NumberOfOutputsTrait15:
+    def number_of_outputs(self):
+        return 15
+
+
+class NumberOfOutputsTrait16:
+    def number_of_outputs(self):
+        return 16
+
+
+class NumberOfOutputsTrait17:
+    def number_of_outputs(self):
+        return 17
+
+
+class NumberOfOutputsTrait18:
+    def number_of_outputs(self):
+        return 18
+
+
+class NumberOfOutputsTrait19:
+    def number_of_outputs(self):
+        return 19
+
+
+class NumberOfOutputsTrait20:
+    def number_of_outputs(self):
+        return 20
+
+
+class NumberOfOutputsTrait21:
+    def number_of_outputs(self):
+        return 21
+
+
+class NumberOfOutputsTrait22:
+    def number_of_outputs(self):
+        return 22
+
+
+def get_mixin_class(base_class, number_of_inputs, number_of_outputs):
+    num_inputs_to_input_trait_class = [
+        None,
+        NumberOfInputsTrait1,
+        NumberOfInputsTrait2,
+        NumberOfInputsTrait3,
+        NumberOfInputsTrait3,
+        NumberOfInputsTrait4,
+        NumberOfInputsTrait5,
+        NumberOfInputsTrait6,
+        NumberOfInputsTrait7,
+        NumberOfInputsTrait8,
+        NumberOfInputsTrait9,
+        NumberOfInputsTrait10,
+        NumberOfInputsTrait11,
+        NumberOfInputsTrait12,
+        NumberOfInputsTrait13,
+        NumberOfInputsTrait14,
+        NumberOfInputsTrait15,
+        NumberOfInputsTrait16,
+        NumberOfInputsTrait17,
+    ]
+    num_outputs_to_output_trait_class = [
+        None,
+        NumberOfOutputsTrait1,
+        NumberOfOutputsTrait2,
+        NumberOfOutputsTrait3,
+        NumberOfOutputsTrait4,
+        NumberOfOutputsTrait5,
+        NumberOfOutputsTrait6,
+        NumberOfOutputsTrait7,
+        NumberOfOutputsTrait8,
+        NumberOfOutputsTrait9,
+        NumberOfOutputsTrait10,
+        NumberOfOutputsTrait11,
+        NumberOfOutputsTrait12,
+        NumberOfOutputsTrait13,
+        NumberOfOutputsTrait14,
+        NumberOfOutputsTrait15,
+        NumberOfOutputsTrait16,
+        NumberOfOutputsTrait17,
+        NumberOfOutputsTrait18,
+        NumberOfOutputsTrait19,
+        NumberOfOutputsTrait20,
+        NumberOfOutputsTrait21,
+        NumberOfOutputsTrait22,
+    ]
+    return type(
+        f"MatmulEpilogueFusion{number_of_inputs}_{number_of_outputs}",
+        [
+            base_class,
+            num_inputs_to_input_trait_class[number_of_inputs],
+            num_outputs_to_output_trait_class[number_of_outputs],
+        ],
+        ap.SerializableAttrMap(),
+    )
+
+
+# abstract_drr.register_drr_pass("matmul_binary_outs_fusion", nice=0)(get_mixin_class(MatmulEpilogueFusion, 3, 2))
+
+
+def register_class(base_class, max_num_inputs, max_num_outputs):
+    def register_drr_class(num_inputs, num_outputs):
+        abstract_drr.register_drr_pass(
+            f"matmul_binary_in{num_inputs}_out{num_outputs}_fusion", nice=0
+        )(get_mixin_class(base_class, num_inputs, num_outputs))
+
+    def register_num_inputs_drr_classes(num_inputs):
+        def register_num_outputs_drr_classes(num_outputs):
+            return register_drr_class(num_inputs + 2, num_outputs + 1)
+
+        ap.map(register_num_outputs_drr_classes, range(max_num_outputs))
+
+    ap.map(register_num_inputs_drr_classes, range(max_num_inputs))
+
+
+register_class(
+    base_class=MatmulEpilogueFusion, max_num_inputs=10, max_num_outputs=10
+)

python/paddle/apy/matmul_pass/matmul_variadic_tpl.py

@@ -0,0 +1,363 @@
+# 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 ap
+import compile_command_util
+import kernel_arg_translator_util
+import low_level_ir_code_gen_ctx_util
+
+
+def make_kernel_arg_translator():
+    return kernel_arg_translator_util.KernelArgTranslator(
+        param_struct_name="args"
+    )
+
+
+def get_anchor_iter_var_names():
+    return ["coord.batch", "coord.row", "coord.column"]
+
+
+class MatmulVariadicTemplate:
+    def __init__(
+        self,
+        program_translator,
+        mut_kernel_arg_id_registry,
+    ):
+        self.program_translator = program_translator
+        self.mut_kernel_arg_id_registry = mut_kernel_arg_id_registry
+        self.kernel_arg_translator = make_kernel_arg_translator()
+        self.dtype2type_name = ap.OrderedDict(
+            [
+                [ap.PointerType.const_float_ptr, "const float*"],
+                [ap.PointerType.const_float16_ptr, "const half*"],
+                [ap.PointerType.const_bfloat16_ptr, "const bfloat16*"],
+                [ap.PointerType.float_ptr, "float*"],
+                [ap.PointerType.float16_ptr, "half*"],
+                [ap.PointerType.bfloat16_ptr, "bfloat16*"],
+                [ap.DataType.float, "float"],
+                [ap.DataType.float16, "half"],
+                [ap.DataType.bfloat16, "bfloat16"],
+                [ap.DataType.int64_t, "int64_t"],
+            ]
+        )
+        self.input_dim_karg_to_shape_access = ap.MutableOrderedDict()
+        self.kernel_name = "MatmulVariadicKernel"
+        self.library_name = "matmul_variadic_kernel"
+
+    def _register_name(self, pair):
+        registry = self.mut_kernel_arg_id_registry
+        registry.get_or_create_kernel_arg_id_manul_var_name(
+            kernel_arg_id=pair[0], cpp_var_name=pair[1]
+        )
+
+    def compile(
+        self,
+        input0_karg,
+        input1_karg,
+        output_karg,
+        input0_shape_kargs,
+        input1_shape_kargs,
+    ):
+        kargs_name_pair_list = [
+            [input0_karg, "input0"],
+            [input1_karg, "input1"],
+            [output_karg, "output"],
+            *ap.map(
+                lambda i: [input0_shape_kargs[i], f"input0_dim{i}"],
+                range(len(input0_shape_kargs)),
+            ),
+            *ap.map(
+                lambda i: [input1_shape_kargs[i], f"input1_dim{i}"],
+                range(len(input1_shape_kargs)),
+            ),
+        ]
+
+        ap.map(self._register_name, kargs_name_pair_list)
+        mut_lir_code_gen_ctx = (
+            low_level_ir_code_gen_ctx_util.CudaLikeIrCodeGenCtx(
+                compute_dtype=ap.DataType.float
+            )
+        )
+        self.program_translator.translate(
+            mut_kernel_arg_id_registry=self.mut_kernel_arg_id_registry,
+            mut_lir_code_gen_ctx=mut_lir_code_gen_ctx,
+        )
+        trivial_code_str = mut_lir_code_gen_ctx.get_stmts_joined_str(
+            indent="    "
+        )
+
+        project_module = self.make_project(
+            trivial_code_str,
+            input0_karg,
+            input1_karg,
+            output_karg,
+            input0_shape_kargs,
+            input1_shape_kargs,
+        )
+        return CodeGenResult(  # noqa: F821
+            module=project_module,
+            kernel_dispatch_func=KernelDispatch,
+            kernel_dispatch_const_data=ap.SerializableAttrMap(
+                kernel_args_getters=self.get_kernel_arg_runtime_getters()
+            ),
+        )
+
+    def get_kernel_arg_runtime_getters(self):
+        all_kernel_arg_id_and_unique_names = self.mut_kernel_arg_id_registry.all_kernel_arg_id2unique_name.items()
+        return ap.map(
+            lambda pair: pair[0].runtime_getter,
+            all_kernel_arg_id_and_unique_names,
+        )
+
+    def get_kernel_arg_types(self):
+        all_kernel_arg_id_and_unique_names = self.mut_kernel_arg_id_registry.all_kernel_arg_id2unique_name.items()
+        return ap.map(
+            lambda pair: pair[0].type, all_kernel_arg_id_and_unique_names
+        )
+
+    def get_kernel_arg_id_var_name(self, kernel_arg_id):
+        all_kernel_arg_id2unique_name = (
+            self.mut_kernel_arg_id_registry.all_kernel_arg_id2unique_name
+        )
+        return all_kernel_arg_id2unique_name[kernel_arg_id]
+
+    def get_kernel_arg_list_str(self, for_declare):
+        def declare_epilogue_arguments_field(pair):
+            kernel_arg_id = pair[0]
+            var_name = pair[1]
+            field_name = self.kernel_arg_translator.get_param_struct_field_name(
+                var_name
+            )
+            dtype = kernel_arg_id.type
+            type_name = self.dtype2type_name[dtype]
+            return (
+                f"{type_name} {field_name}" if for_declare else f"{field_name}"
+            )
+
+        all_kernel_arg_id_and_names = self.mut_kernel_arg_id_registry.all_kernel_arg_id2unique_name.items()
+        return ", ".join(
+            ap.map(
+                declare_epilogue_arguments_field, all_kernel_arg_id_and_names
+            )
+        )
+
+    def get_epilogue_arguments_fields_str(self, indent):
+        def declare_epilogue_arguments_field(pair):
+            kernel_arg_id = pair[0]
+            var_name = pair[1]
+            field_name = self.kernel_arg_translator.get_param_struct_field_name(
+                var_name
+            )
+            dtype = kernel_arg_id.type
+            type_name = self.dtype2type_name[dtype]
+            return f"{type_name} {field_name};"
+
+        generated_kernel_arg_id_and_names = self.mut_kernel_arg_id_registry.generated_kernel_arg_id2unique_name.items()
+        return f"\n{indent}".join(
+            ap.map(
+                declare_epilogue_arguments_field,
+                generated_kernel_arg_id_and_names,
+            )
+        )
+
+    def get_epilogue_arguments_init_str(self, param_obj_name, indent):
+        def declare_epilogue_arguments_assign(pair):
+            kernel_arg_id = pair[0]
+            var_name = pair[1]
+            field_name = self.kernel_arg_translator.get_param_struct_field_name(
+                var_name
+            )
+            return f"{param_obj_name}.{field_name} = {var_name};"
+
+        generated_kernel_arg_id_and_names = self.mut_kernel_arg_id_registry.generated_kernel_arg_id2unique_name.items()
+        return f"\n{indent}".join(
+            ap.map(
+                declare_epilogue_arguments_assign,
+                generated_kernel_arg_id_and_names,
+            )
+        )
+
+    def get_params_input_shape_init_str(
+        self, input_name, input_shape_kargs, indent
+    ):
+        def init_input_shape_with_args(i):
+            def get_creator():
+                return f"{input_name}_shape[{i}]"
+
+            karg_var_name = self.get_kernel_arg_id_var_name(
+                input_shape_kargs[i]
+            )
+            self.input_dim_karg_to_shape_access.get_or_create(
+                karg_var_name, get_creator
+            )
+            return f"{indent}{input_name}_shape[{i}] = {karg_var_name};"
+
+        shape_vector_init_str = (
+            f"{input_name}_shape.resize({len(input_shape_kargs)});\n"
+        )
+        return shape_vector_init_str + "\n".join(
+            ap.map(init_input_shape_with_args, range(len(input_shape_kargs)))
+        )
+
+    def make_project(
+        self,
+        trivial_code_str,
+        input0_karg,
+        input1_karg,
+        output_karg,
+        input0_shape_kargs,
+        input1_shape_kargs,
+    ):
+        code_template = """
+// auto generated codes
+#include "matmul.h"
+#include <vector>
+
+namespace ap {
+
+template <typename T>
+struct VariadicEpilogueFunctor {
+  struct Arguments {
+    ${AP_EPILOGUE_ARGUMENTS_FIELDS}
+  };
+
+  // Note: need to support vectorized operation
+  __forceinline__ __host__ __device__
+  T operator()(T x, const Arguments& args, const MatrixCoord& coord) const {
+    T out;
+    ${AP_EPILOGUE_COMPUTATION_STATEMENTS}
+    return out;
+  }
+};
+
+template <int TuningConfigId>
+static void RunMatmulWithVariadicKernel(const GemmEpilogueParams &params, ${AP_KERNEL_ARGS_DECLARE}) {
+  using ElementT = ${output_dtype};
+  using ElementComputeT = float;
+
+  typename VariadicEpilogueFunctor<ElementComputeT>::Arguments epilogue_args;
+
+  ${AP_EPILOGUE_ARGUMENTS_INIT}
+
+  constexpr int AlignA = Alignment<ElementT, ${k_value}>::kValue;
+  constexpr int AlignB = Alignment<ElementT, ${n_value}>::kValue;
+
+  MatmulAddVariadic<ElementT, ElementComputeT, VariadicEpilogueFunctor,
+                           AlignA, AlignB, TuningConfigId>(params, epilogue_args);
+}
+
+} // namespace ap
+
+extern "C" {
+
+void ${kernel_name}(void* stream_ptr, ${AP_KERNEL_ARGS_DECLARE}) {
+  std::vector<int64_t> ${input0}_shape;
+  ${AP_PARAMS_INPUT0_SHAPE_INIT}
+
+  std::vector<int64_t> ${input1}_shape;
+  ${AP_PARAMS_INPUT1_SHAPE_INIT}
+
+  ap::GemmEpilogueParams params(
+      stream_ptr, ${input0}, ${input1}, nullptr, ${output}, ${input0}_shape, ${input1}_shape, std::vector<int64_t>{});
+
+#if AP_ENABLE_AUTOTUNE
+  AP_AUTOTUNE_${output_dtype}(ap::RunMatmulWithVariadicKernel, stream_ptr, params, ${AP_KERNEL_ARGS_CALL});
+#else
+  ap::RunMatmulWithVariadicKernel<ap::DefaultConfig::kConfigId>(params, ${AP_KERNEL_ARGS_CALL});
+#endif
+}
+}
+  """
+
+        output_dtype = self.dtype2type_name[output_karg.type.data_type]
+        code = (
+            code_template.replace(
+                "${AP_EPILOGUE_COMPUTATION_STATEMENTS}", trivial_code_str
+            )
+            .replace(
+                "${AP_KERNEL_ARGS_DECLARE}",
+                self.get_kernel_arg_list_str(for_declare=True),
+            )
+            .replace(
+                "${AP_KERNEL_ARGS_CALL}",
+                self.get_kernel_arg_list_str(for_declare=False),
+            )
+            .replace(
+                "${AP_PARAMS_INPUT0_SHAPE_INIT}",
+                self.get_params_input_shape_init_str(
+                    "${input0}", input0_shape_kargs, indent="  "
+                ),
+            )
+            .replace(
+                "${AP_PARAMS_INPUT1_SHAPE_INIT}",
+                self.get_params_input_shape_init_str(
+                    "${input1}", input1_shape_kargs, indent="  "
+                ),
+            )
+            .replace(
+                "${AP_EPILOGUE_ARGUMENTS_FIELDS}",
+                self.get_epilogue_arguments_fields_str(indent="    "),
+            )
+            .replace(
+                "${AP_EPILOGUE_ARGUMENTS_INIT}",
+                self.get_epilogue_arguments_init_str(
+                    "epilogue_args", indent="  "
+                ),
+            )
+            .replace("${kernel_name}", self.kernel_name)
+            .replace("${input0}", self.get_kernel_arg_id_var_name(input0_karg))
+            .replace("${input1}", self.get_kernel_arg_id_var_name(input1_karg))
+            .replace("${output}", self.get_kernel_arg_id_var_name(output_karg))
+            .replace("${output_dtype}", output_dtype)
+            .replace("${k_value}", f"{input0_shape_kargs[-1].value}")
+            .replace("${n_value}", f"{input1_shape_kargs[-1].value}")
+        )
+
+        dir_name = ap.dirname(__file__)
+        compile_command_generator = (
+            compile_command_util.CompileCommandGenerator()
+        )
+        compile_cmd = compile_command_generator(
+            "matmul", dir_name, self.library_name
+        )
+        file_ext = compile_command_generator.file_ext
+
+        return CodeModule(  # noqa: F821
+            FuncDeclare(  # noqa: F821
+                ap.DataType.void,
+                self.kernel_name,
+                [ap.PointerType.void_ptr, *self.get_kernel_arg_types()],
+            ),
+            Project(  # noqa: F821
+                nested_files=Project.Directory(  # noqa: F821
+                    [
+                        f"{self.library_name}.{file_ext}",
+                        Project.FileContent(code),  # noqa: F821
+                    ],
+                    ["make.sh", Project.FileContent(compile_cmd)],  # noqa: F821
+                ),
+                compile_cmd="sh make.sh",
+                so_relative_path=f"lib{self.library_name}.so",
+            ),
+        )
+
+
+def KernelDispatch(ctx):
+    import ap
+
+    so_func = ctx.get_so_function("MatmulVariadicKernel")
+    stream_ptr = ctx.device_ctx.get_stream_addr_as_void_ptr()
+    getters = ctx.kernel_dispatch_const_data.kernel_args_getters
+    args = [stream_ptr, *ap.map(lambda getter: getter(ctx), getters)]
+    ap.apply(so_func, args)

python/paddle/apy/matmul_pass/op_compute_translator_util.py

@@ -0,0 +1,836 @@
+# 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 ap
+import code_gen_value_util
+
+
+class ApOpLoadFromRegisterCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        out = self.get_out_cg_val(0)
+        return [out]
+
+    def get_out_cg_val(self, i):
+        register_var_name_attr = self.op_property.attributes.register_var_name
+        register_var_name = register_var_name_attr.match(a_str=lambda x: x)
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type, register_var_name
+        )
+
+
+class ApOpLoadFromGlobalCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        index_func_unique_id_attr = (
+            self.op_property.attributes.index_func_unique_id
+        )
+        index_func_unique_id = index_func_unique_id_attr.match(
+            a_str=lambda x: x
+        )
+        offset_var_name = self.index_program_translator_map.get_offset_var_name(
+            index_func_unique_id=index_func_unique_id,
+            mut_kernel_arg_id_registry=mut_kernel_arg_id_registry,
+            mut_lir_code_gen_ctx=mut_lir_code_gen_ctx,
+        )
+        data_op_name = inputs[0].var_name
+        arg_name = mut_kernel_arg_id_registry.get_in_tensor_data_ptr_var_name(
+            data_op_name
+        )
+        ptr_var_name = self.kernel_arg_translator.get_use_name(arg_name)
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"{ptr_var_name}[{offset_var_name}]")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class ApOpStoreToRegisterCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        mut_lir_code_gen_ctx.stmts.append(
+            f"{self.get_out_var_name()} = {inputs[0].var_name};"
+        )
+        return []
+
+    def get_out_var_name(self):
+        register_var_name_attr = self.op_property.attributes.register_var_name
+        return register_var_name_attr.match(a_str=lambda x: x)
+
+
+class ApOpStoreToGlobalCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+        self.ptr_type2data_type = ap.OrderedDict(
+            [
+                [ap.PointerType.const_float_ptr, "const float"],
+                [ap.PointerType.const_float16_ptr, "const half"],
+                [ap.PointerType.float_ptr, "float"],
+                [ap.PointerType.float16_ptr, "half"],
+            ]
+        )
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        index_func_unique_id_attr = (
+            self.op_property.attributes.index_func_unique_id
+        )
+        index_func_unique_id = index_func_unique_id_attr.match(
+            a_str=lambda x: x
+        )
+        offset_var_name = self.index_program_translator_map.get_offset_var_name(
+            index_func_unique_id=index_func_unique_id,
+            mut_kernel_arg_id_registry=mut_kernel_arg_id_registry,
+            mut_lir_code_gen_ctx=mut_lir_code_gen_ctx,
+        )
+        arg_name = mut_kernel_arg_id_registry.get_out_tensor_data_ptr_var_name(
+            index_func_unique_id
+        )
+        glb_data_type = self.get_glb_type(
+            mut_kernel_arg_id_registry, index_func_unique_id
+        )
+        ptr_var_name = self.kernel_arg_translator.get_use_name(arg_name)
+        mut_lir_code_gen_ctx.store(
+            glb_data_type, ptr_var_name, offset_var_name, inputs[1].var_name
+        )
+        return []
+
+    def get_glb_type(self, mut_kernel_arg_id_registry, index_func_unique_id):
+        arg_name = mut_kernel_arg_id_registry.get_out_tensor_data_ptr_var_name(
+            index_func_unique_id
+        )
+        kernel_arg_name2ids = ap.OrderedDict(
+            ap.map(
+                lambda item: [item[1], item[0]],
+                mut_kernel_arg_id_registry.generated_kernel_arg_id2unique_name.items(),
+            )
+        )
+        kernel_arg_id = kernel_arg_name2ids[arg_name]
+        return self.ptr_type2data_type[kernel_arg_id.type]
+
+    def get_out_var_name(self):
+        register_var_name_attr = self.op_property.attributes.register_var_name
+        return register_var_name_attr.match(a_str=lambda x: x)
+
+
+class PdOpDataCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        out = self.get_out_cg_val(0)
+        return [out]
+
+    def get_out_cg_val(self, i):
+        name = self.op_property.attributes.name.match(a_str=lambda x: x)
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type, name
+        )
+
+
+class PdOpFullCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        value = self.op_property.attributes.value.match(a_f64=lambda x: x)
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"{value}")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpCastCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+        self.dtype2type_name = ap.OrderedDict(
+            [
+                [ap.DataType.float, "float"],
+                [ap.DataType.float16, "half"],
+                [ap.DataType.bfloat16, "nv_bfloat16"],
+                [ap.DataType.int32, "int"],
+                [ap.DataType.int64, "int64_t"],
+            ]
+        )
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        dtype = self.op_property.attributes.dtype.match(a_dtype=lambda x: x)
+        dtype_name = self.dtype2type_name[dtype]
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(
+            out, f"static_cast<{dtype_name}>({inputs[0].var_name})"
+        )
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpExpCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"expf({inputs[0].var_name})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpReluCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(
+            out, f"({inputs[0].var_name} > 0 ? {inputs[0].var_name} : 0) "
+        )
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpErfCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        var_name = inputs[0].var_name
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"erf({var_name})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpElementwisePowCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        exponent = inputs[1].var_name
+        var_name = inputs[0].var_name
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"ComputePow({var_name}, {exponent})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpSinCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        var_name = inputs[0].var_name
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"sin({var_name})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpTanhCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        var_name = inputs[0].var_name
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"tanh({var_name})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpFloorCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        var_name = inputs[0].var_name
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"floor({var_name})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class CinnOpScaleCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        scale = self.op_property.attributes.scale.match(a_f32=lambda x: x)
+        bias = self.op_property.attributes.bias.match(a_f32=lambda x: x)
+        bias_after_scale = self.op_property.attributes.bias_after_scale.match(
+            a_bool=lambda x: x
+        )
+        in_name = inputs[0].var_name
+        true_str = f"{scale} * {in_name} + {bias}"
+        false_str = f"{scale} * ({in_name} + {bias})"
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(
+            out, true_str if bias_after_scale else false_str
+        )
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpSubtractCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        a = inputs[0]
+        b = inputs[1]
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"({a.var_name} - {b.var_name})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpAddCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        a = inputs[0]
+        b = inputs[1]
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"{a.var_name} + {b.var_name}")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpMultiplyCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        a = inputs[0]
+        b = inputs[1]
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"{a.var_name} * {b.var_name}")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpDivideCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        a = inputs[0]
+        b = inputs[1]
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"{a.var_name} / {b.var_name}")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpMaximumCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        a = inputs[0]
+        b = inputs[1]
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(
+            out,
+            f"(({a.var_name} >= {b.var_name}) ? ({a.var_name}) : ({b.var_name}))",
+        )
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class CinnOpYieldStoreCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        return inputs
+
+
+class CinnOpBroadcastCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        return inputs
+
+
+class CinnOpExpandCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        return [inputs[0]]
+
+
+class CinnOpGenerateShapeCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        out = self.get_out_cg_val(0)
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class OpComputeTranslatorFactory:
+    def __init__(self):
+        self.op_name2class = ap.OrderedDict(
+            [
+                ["ap_op.load_from_register", ApOpLoadFromRegisterCodeGen],
+                ["ap_op.load_from_global", ApOpLoadFromGlobalCodeGen],
+                ["ap_op.store_to_register", ApOpStoreToRegisterCodeGen],
+                ["ap_op.store_to_global", ApOpStoreToGlobalCodeGen],
+                ["pd_op.data", PdOpDataCodeGen],
+                ["pd_op.full", PdOpFullCodeGen],
+                ["pd_op.cast", PdOpCastCodeGen],
+                ["pd_op.exp", PdOpExpCodeGen],
+                ["pd_op.relu", PdOpReluCodeGen],
+                ["pd_op.sin", PdOpSinCodeGen],
+                ["pd_op.tanh", PdOpTanhCodeGen],
+                ["pd_op.floor", PdOpFloorCodeGen],
+                ["pd_op.erf", PdOpErfCodeGen],
+                ["pd_op.elementwise_pow", PdOpElementwisePowCodeGen],
+                ["cinn_op.scale", CinnOpScaleCodeGen],
+                ["pd_op.subtract", PdOpSubtractCodeGen],
+                ["pd_op.add", PdOpAddCodeGen],
+                ["pd_op.multiply", PdOpMultiplyCodeGen],
+                ["pd_op.divide", PdOpDivideCodeGen],
+                ["pd_op.maximum", PdOpMaximumCodeGen],
+                ["cinn_op.yield_store", CinnOpYieldStoreCodeGen],
+                ["cinn_op.broadcast", CinnOpBroadcastCodeGen],
+                ["pd_op.expand", CinnOpExpandCodeGen],
+                ["cinn_op.generate_shape", CinnOpGenerateShapeCodeGen],
+            ]
+        )
+
+    def __call__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        cls = self._get_class(op_property.op_name)
+        return cls(
+            op_property=op_property,
+            input_properties=input_properties,
+            output_properties=output_properties,
+            kernel_arg_translator=kernel_arg_translator,
+            index_program_translator_map=index_program_translator_map,
+        )
+
+    def _get_class(self, op_name):
+        return self.op_name2class[op_name]

python/paddle/apy/matmul_pass/op_conversion_drr_pass.py

@@ -0,0 +1,222 @@
+# 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 access_topo_drr
+
+
+@access_topo_drr.register_drr_pass("pd_op_cast", tag="default")
+class PdOpCastAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.cast_op = o.ap_native_op("pd_op.cast")
+        o.cast_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_tanh", tag="default")
+class PdOpTanhAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.tanh_op = o.ap_native_op("pd_op.tanh")
+        o.tanh_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_floor", tag="default")
+class PdOpFloorAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.floor_op = o.ap_native_op("pd_op.floor")
+        o.floor_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_erf", tag="default")
+class PdOpErfAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.erf_op = o.ap_native_op("pd_op.erf")
+        o.erf_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_elementwise_pow", tag="default")
+class PdOpElementwisePowAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.source_op = o.ap_native_op("pd_op.elementwise_pow")
+        o.source_op([t.input0, t.input1], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.add")
+        o.result_op([t.input0, t.input1], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_exp", tag="default")
+class PdOpExpAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.exp_op = o.ap_native_op("pd_op.exp")
+        o.exp_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("cinn_op_scale", tag="default")
+class CinnOpScaleAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.scale_op = o.ap_native_op("cinn_op.scale")
+        o.scale_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_sin", tag="default")
+class PdOpSinAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.sin_op = o.ap_native_op("pd_op.sin")
+        o.sin_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("cinn_op_yield_store", tag="default")
+class CinnOpYieldStoreAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.yield_op = o.ap_native_op("cinn_op.yield_store")
+        o.yield_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_subtract", tag="default")
+class PdOpSubtractAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.source_op = o.ap_native_op("pd_op.subtract")
+        o.source_op([t.input0, t.input1], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.add")
+        o.result_op([t.input0, t.input1], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_divide", tag="default")
+class PdOpDivideAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.source_op = o.ap_native_op("pd_op.divide")
+        o.source_op([t.input0, t.input1], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.add")
+        o.result_op([t.input0, t.input1], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_multiply", tag="default")
+class PdOpMultiplyAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.source_op = o.ap_native_op("pd_op.multiply")
+        o.source_op([t.input0, t.input1], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.add")
+        o.result_op([t.input0, t.input1], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_maximum", tag="default")
+class PdOpMaximumAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.source_op = o.ap_native_op("pd_op.maximum")
+        o.source_op([t.input0, t.input1], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.add")
+        o.result_op([t.input0, t.input1], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_left_full_add", tag="default")
+class PdOpLeftFullAddAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.full_op = o.ap_native_op("pd_op.full")
+        o.full_op([], [t.intermediate])
+        o.source_op = o.ap_native_op("pd_op.add")
+        o.source_op([t.intermediate, t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_right_full_add", tag="default")
+class PdOpRightFullAddAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.full_op = o.ap_native_op("pd_op.full")
+        o.full_op([], [t.intermediate])
+        o.source_op = o.ap_native_op("pd_op.add")
+        o.source_op([t.input, t.intermediate], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass(
+    "full_generate_shape_expand_left_add", tag="default"
+)
+class FullGenerateShapeExpandLeftAddAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.full = o.ap_native_op("pd_op.full")
+        o.full([], [t.intermediate0])
+        o.generate_shape = o.ap_native_op("cinn_op.generate_shape")
+        o.generate_shape([t.input0], [t.intermediate1])
+        o.expand = o.ap_native_op("pd_op.expand")
+        o.expand([t.intermediate0, t.intermediate1], [t.expanded_input])
+        o.add = o.ap_native_op("pd_op.add")
+        o.add([t.expanded_input, t.input0], [t.output])
+
+    def result_pattern(self, o, t):
+        o.relu = o.ap_native_op("pd_op.relu")
+        o.relu([t.input0], [t.output])
+
+
+@access_topo_drr.register_drr_pass(
+    "full_generate_shape_expand_right_add", tag="default"
+)
+class FullGenerateShapeExpandRightAddAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.full = o.ap_native_op("pd_op.full")
+        o.full([], [t.intermediate0])
+        o.generate_shape = o.ap_native_op("cinn_op.generate_shape")
+        o.generate_shape([t.input0], [t.intermediate1])
+        o.expand = o.ap_native_op("pd_op.expand")
+        o.expand([t.intermediate0, t.intermediate1], [t.expanded_input])
+        o.add = o.ap_native_op("pd_op.add")
+        o.add([t.input0, t.expanded_input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.relu = o.ap_native_op("pd_op.relu")
+        o.relu([t.input0], [t.output])

python/paddle/apy/matmul_pass/op_index_translator_util.py

@@ -0,0 +1,226 @@
+# 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 ap
+import index_code_gen_value_util
+
+
+class PdOpDataCodeGen:
+    def __init__(
+        self,
+        index_program_id,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        self.index_program_id = index_program_id
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.anchor_iter_var_names = anchor_iter_var_names
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        return [
+            index_code_gen_value_util.IndexCodeGenValue(
+                self.anchor_iter_var_names
+            )
+        ]
+
+
+class PdOpFullIntArrayCodeGen:
+    def __init__(
+        self,
+        index_program_id,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        self.index_program_id = index_program_id
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.anchor_iter_var_names = anchor_iter_var_names
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        out = index_code_gen_value_util.IndexCodeGenValue(None)
+
+        def get_int64(attr):
+            return attr.match(a_i64=lambda x: x)
+
+        def convert_list(lst):
+            return ap.map(get_int64, lst)
+
+        out.const_data = self.op_property.attributes.value.match(
+            a_array=convert_list
+        )
+        return [out]
+
+
+class PdOpSumCodeGen:
+    def __init__(
+        self,
+        index_program_id,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        self.index_program_id = index_program_id
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.anchor_iter_var_names = anchor_iter_var_names
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        input_iter_var_names = inputs[0].iter_var_names
+        reduced_axes_set = ap.OrderedDict(
+            ap.map(lambda x: [int(x), True], inputs[1].const_data)
+        )
+        non_reduced_axes = ap.filter(
+            lambda x: reduced_axes_set.contains(x) == False,  # noqa: E712
+            range(len(input_iter_var_names)),
+        )
+        output_iter_var_names = ap.map(
+            lambda i: input_iter_var_names[i], non_reduced_axes
+        )
+        return [
+            index_code_gen_value_util.IndexCodeGenValue(output_iter_var_names)
+        ]
+
+
+class CinnOpReshapeCodeGen:
+    def __init__(
+        self,
+        index_program_id,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        self.index_program_id = index_program_id
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.anchor_iter_var_names = anchor_iter_var_names
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        symbolic_shape = self.input_properties[0].symbolic_shape
+
+        def get_or_create_dim_var_name(dim_expr):
+            arg_var_name = mut_kernel_arg_id_registry.get_dim_expr_var_name(
+                dim_expr
+            )
+            return self.kernel_arg_translator.get_use_name(arg_var_name)
+
+        def get_dim_var_name(i):
+            dim_expr = symbolic_shape[i]
+            return get_or_create_dim_var_name(dim_expr)
+
+        rank = len(symbolic_shape)
+        stride_dims_list = ap.map(
+            lambda num_dims: ap.map(
+                lambda i: get_dim_var_name(num_dims + i + 1),
+                range(rank - 1 - num_dims),
+            ),
+            range(rank),
+        )
+        var_name_and_dims_list = ap.map(
+            lambda pair: [pair[0], *pair[1]],
+            zip(inputs[0].iter_var_names, stride_dims_list),
+        )
+        offset_expr = " + ".join(
+            ap.map(lambda elts: " * ".join(elts), var_name_and_dims_list)
+        )
+        assert len(self.output_properties[0].symbolic_shape) == 1, (
+            "len(self.output_properties[0]) should be 1"
+        )
+        return [
+            index_code_gen_value_util.IndexCodeGenValue([f"({offset_expr})"])
+        ]
+
+
+class CfYieldCodeGen:
+    def __init__(
+        self,
+        index_program_id,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        self.index_program_id = index_program_id
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.anchor_iter_var_names = anchor_iter_var_names
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        return []
+
+
+class OpIndexTranslatorFactory:
+    def __init__(self):
+        self.op_name2class = ap.OrderedDict(
+            [
+                ["pd_op.data", PdOpDataCodeGen],
+                ["pd_op.full_int_array", PdOpFullIntArrayCodeGen],
+                ["pd_op.sum", PdOpSumCodeGen],
+                ["cinn_op.reshape", CinnOpReshapeCodeGen],
+                ["cf.yield", CfYieldCodeGen],
+            ]
+        )
+
+    def __call__(
+        self,
+        index_program_id,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        cls = self._get_class(op_property.op_name)
+        return cls(
+            index_program_id=index_program_id,
+            op_property=op_property,
+            input_properties=input_properties,
+            output_properties=output_properties,
+            kernel_arg_translator=kernel_arg_translator,
+            anchor_iter_var_names=anchor_iter_var_names,
+        )
+
+    def _get_class(self, op_name):
+        return self.op_name2class[op_name]

python/paddle/apy/matmul_pass/program_translator_util.py

@@ -0,0 +1,81 @@
+# 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 ap
+
+
+class ProgramTranslator:
+    def __init__(
+        self,
+        program_property,
+        kernel_arg_translator,
+        index_program_translator_map,
+        op_translator_maker,
+    ):
+        self.program_property = program_property
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+        self.op_translator_maker = op_translator_maker
+        self.ir_value_index2translated_value = ap.MutableList()
+
+        def PushNone(x):
+            self.ir_value_index2translated_value.append(None)
+
+        map(PushNone, self.program_property.values)
+
+    # mut_kernel_arg_id_registry: mutable KernelArgIdLazyContext
+    # mut_lir_code_gen_ctx: mutable low level ir code generation context
+    def translate(self, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx):
+        def TranslateOp(op_property):
+            self._translate_op(
+                op_property, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+            )
+
+        map(TranslateOp, self.program_property.ops)
+
+    def _translate_op(
+        self, op_property, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        op_translator = self.op_translator_maker(
+            op_property=op_property,
+            input_properties=map(
+                self._get_value_property, op_property.input_value_indexes
+            ),
+            output_properties=map(
+                self._get_value_property, op_property.output_value_indexes
+            ),
+            kernel_arg_translator=self.kernel_arg_translator,
+            index_program_translator_map=self.index_program_translator_map,
+        )
+        inputs = map(
+            self._get_translated_value, op_property.input_value_indexes
+        )
+        outputs = op_translator(
+            inputs,
+            mut_kernel_arg_id_registry=mut_kernel_arg_id_registry,
+            mut_lir_code_gen_ctx=mut_lir_code_gen_ctx,
+        )
+        map(
+            self._set_translated_value,
+            zip(op_property.output_value_indexes, outputs),
+        )
+
+    def _get_value_property(self, i):
+        return self.program_property.values[i]
+
+    def _get_translated_value(self, i):
+        return self.ir_value_index2translated_value[i]
+
+    def _set_translated_value(self, pair):
+        self.ir_value_index2translated_value[pair[0]] = pair[1]

python/paddle/apy/matmul_pass/topo_drr_pass.py

@@ -0,0 +1,489 @@
+# 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 access_topo_drr
+import ap
+import pir
+
+
+class FakeDataForYieldAccessTopoPass(access_topo_drr.DrrPass):
+    def __init__(self, fake_data_names):
+        self.num_outputs = len(fake_data_names)
+        self.fake_data_names = fake_data_names
+        self.undefined_place = pir.a_place(pir.UndefinedPlace())
+
+    def source_pattern(self, o, t):
+        o.yield_op = o.ap_native_op("cf.yield")
+
+        def get_yield_input(i):
+            return getattr(t, f"output{i}")
+
+        o.yield_op(ap.map(get_yield_input, range(self.num_outputs)), [])
+
+    def result_pattern(self, o, t):
+        self.result_pattern_data_op(o, t)
+        self.result_pattern_up_spider(o, t)
+
+    def result_pattern_data_op(self, o, t):
+        ap.map(
+            lambda i: self.data_op_for_output(o, t, i), range(self.num_outputs)
+        )
+
+    def data_op_for_output(self, o, t, i):
+        t.declare_internal_native_ir_value(f"data_out{i}")
+        data_op_unique_name = f"data_op_for_output{i}"
+        setattr(o, data_op_unique_name, o.ap_native_op("pd_op.data"))
+        data_op = getattr(o, data_op_unique_name)
+        data_op.name = lambda o, t: self.get_name(o, t, i)
+        data_op.shape = lambda o, t: self.get_shape(o, t, i)
+        data_op.dtype = lambda o, t: self.get_dtype(o, t, i)
+        data_op.place = lambda o, t: self.get_place(o, t, i)
+        data_op([], [getattr(t, f"data_out{i}")])
+
+    def get_name(self, o, t, i):
+        return pir.a_str(self.fake_data_names[i])
+
+    def get_shape(self, o, t, i):
+        ir_tensor = getattr(t, f"output{i}")
+
+        def GetDims(dtype, dims, data_layout):
+            return dims
+
+        return pir.a_intarray(ir_tensor.type.match(t_dtensor=GetDims))
+
+    def get_dtype(self, o, t, i):
+        ir_tensor = getattr(t, f"output{i}")
+
+        def GetDtype(dtype, dims, data_layout):
+            return dtype
+
+        return pir.a_dtype(ir_tensor.type.match(t_dtensor=GetDtype))
+
+    def get_place(self, o, t, i):
+        return self.undefined_place
+
+    def result_pattern_up_spider(self, o, t):
+        ap.map(
+            lambda i: self.up_spider_for_output(o, t, i),
+            range(self.num_outputs),
+        )
+
+    def up_spider_for_output(self, o, t, i):
+        t.declare_internal_native_ir_value(f"add_out{i}")
+        up_spider_op_name = f"up_spider_op{i}"
+        setattr(o, up_spider_op_name, o.ap_native_op("ap_op.up_spider"))
+        getattr(o, up_spider_op_name)(
+            [getattr(t, f"output{i}"), getattr(t, f"data_out{i}")], []
+        )
+
+
+class FakeDataStoreToGlobalForYieldAccessTopoPass(access_topo_drr.DrrPass):
+    def __init__(self, fake_data_names):
+        self.num_outputs = len(fake_data_names)
+        self.fake_data_names = fake_data_names
+        self.undefined_place = pir.a_place(pir.UndefinedPlace())
+
+    def source_pattern(self, o, t):
+        o.yield_op = o.ap_native_op("cf.yield")
+
+        def get_yield_input(i):
+            return getattr(t, f"output{i}")
+
+        o.yield_op(ap.map(get_yield_input, range(self.num_outputs)), [])
+
+    def result_pattern(self, o, t):
+        self.result_pattern_data_op(o, t)
+        self.result_pattern_store_to_global_op(o, t)
+
+    def result_pattern_data_op(self, o, t):
+        ap.map(
+            lambda i: self.data_op_for_output(o, t, i), range(self.num_outputs)
+        )
+
+    def data_op_for_output(self, o, t, i):
+        t.declare_internal_native_ir_value(f"data_out{i}")
+        data_op_unique_name = f"data_op_for_output{i}"
+        setattr(o, data_op_unique_name, o.ap_native_op("pd_op.data"))
+        data_op = getattr(o, data_op_unique_name)
+        data_op.name = lambda o, t: self.get_name(o, t, i)
+        data_op.shape = lambda o, t: self.get_shape(o, t, i)
+        data_op.dtype = lambda o, t: self.get_dtype(o, t, i)
+        data_op.place = lambda o, t: self.get_place(o, t, i)
+        data_op([], [getattr(t, f"data_out{i}")])
+
+    def get_name(self, o, t, i):
+        return pir.a_str(self.fake_data_names[i])
+
+    def get_shape(self, o, t, i):
+        ir_tensor = getattr(t, f"output{i}")
+
+        def GetDims(dtype, dims, data_layout):
+            return dims
+
+        return pir.a_intarray(ir_tensor.type.match(t_dtensor=GetDims))
+
+    def get_dtype(self, o, t, i):
+        ir_tensor = getattr(t, f"output{i}")
+
+        def GetDtype(dtype, dims, data_layout):
+            return dtype
+
+        return pir.a_dtype(ir_tensor.type.match(t_dtensor=GetDtype))
+
+    def get_place(self, o, t, i):
+        return self.undefined_place
+
+    def result_pattern_store_to_global_op(self, o, t):
+        ap.map(
+            lambda i: self.store_to_global_op_for_output(o, t, i),
+            range(self.num_outputs),
+        )
+
+    def store_to_global_op_for_output(self, o, t, i):
+        store_to_global_op_name = f"store_to_global_op{i}"
+        setattr(
+            o, store_to_global_op_name, o.ap_native_op("ap_op.store_to_global")
+        )
+        store_to_global_op = getattr(o, store_to_global_op_name)
+        store_to_global_op.index_func_unique_id = lambda o, t: pir.a_str(
+            self.fake_data_names[i]
+        )
+        store_to_global_op(
+            [getattr(t, f"data_out{i}"), getattr(t, f"output{i}")], []
+        )
+
+
+class ConvertUpSpiderStoreDataOpToYieldOpPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.data_op = o.ap_native_op("pd_op.data")
+        o.data_op([], [t.input1])
+        o.load_from_global_op = o.ap_native_op("ap_op.load_from_global")
+        o.load_from_global_op([t.input1], [t.tmp1])
+        o.up_spider_op = o.ap_native_op("ap_op.up_spider")
+        o.up_spider_op([t.input0, t.tmp1], [])
+
+    def result_pattern(self, o, t):
+        o.yield_op = o.ap_native_op("cf.yield")
+        o.yield_op([t.input0], [])
+
+
+class ConvertDownSpiderStoreDataOpToYieldOpPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.data_mm_op = o.ap_native_op("pd_op.data")
+        o.data_mm_op([], [t.input1])
+        o.down_spider_op = o.ap_native_op("ap_op.down_spider")
+        o.down_spider_op([t.input1], [t.tmp1])
+        o.store_to_global = o.ap_native_op("ap_op.store_to_global")
+        o.store_to_global([t.input0, t.tmp1], [])
+
+    def result_pattern(self, o, t):
+        o.yield_op = o.ap_native_op("cf.yield")
+        o.yield_op([t.input0], [])
+
+
+class InitDownSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def __init__(self, data_input_name):
+        self.data_input_name_attr = pir.a_str(data_input_name)
+
+    def source_pattern(self, o, t):
+        o.data_op = o.ap_native_op("pd_op.data")
+        o.data_op([], [t.output])
+
+    def constraint(self, o, t):
+        return o.data_op.name == self.data_input_name_attr
+
+    def result_pattern(self, o, t):
+        t.declare_internal_native_ir_value("input")
+        o.new_data_op = o.ap_native_op("pd_op.data")
+        o.new_data_op.name = lambda o, t: o.data_op.name
+        o.new_data_op.shape = lambda o, t: o.data_op.shape
+        o.new_data_op.dtype = lambda o, t: o.data_op.dtype
+        o.new_data_op.place = lambda o, t: o.data_op.place
+        o.new_data_op([], [t.input])
+        o.down_spider = o.ap_native_op("ap_op.down_spider")
+        o.down_spider([t.input], [t.output])
+
+
+class InitNaiveLoadFromGlobalAccessTopoPass(access_topo_drr.DrrPass):
+    def __init__(self, data_input_name):
+        self.data_input_name_attr = pir.a_str(data_input_name)
+
+    def source_pattern(self, o, t):
+        o.data_op = o.ap_native_op("pd_op.data")
+        o.data_op([], [t.output])
+
+    def constraint(self, o, t):
+        return o.data_op.name == self.data_input_name_attr
+
+    def result_pattern(self, o, t):
+        t.declare_internal_native_ir_value("input")
+        o.new_data_op = o.ap_native_op("pd_op.data")
+        o.new_data_op.name = lambda o, t: o.data_op.name
+        o.new_data_op.shape = lambda o, t: o.data_op.shape
+        o.new_data_op.dtype = lambda o, t: o.data_op.dtype
+        o.new_data_op.place = lambda o, t: o.data_op.place
+        o.new_data_op([], [t.input])
+        o.load_from_global = o.ap_native_op("ap_op.load_from_global")
+        o.load_from_global.index_func_unique_id = lambda o, t: (
+            self.data_input_name_attr
+        )
+        o.load_from_global([t.input], [t.output])
+
+
+class ReplaceWithLoadFromRegisterPass(access_topo_drr.DrrPass):
+    def __init__(self, name, register_var_name):
+        self.name = pir.a_str(name)
+        self.register_var_name = pir.a_str(register_var_name)
+
+    def source_pattern(self, o, t):
+        o.data_op = o.ap_native_op("pd_op.data")
+        o.data_op([], [t.input])
+        o.load_from_global = o.ap_native_op("ap_op.load_from_global")
+        o.load_from_global.index_func_unique_id = self.name
+        o.load_from_global([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.load_from_register = o.ap_native_op("ap_op.load_from_register")
+        o.load_from_register.name = lambda o, t: self.name
+        o.load_from_register.register_var_name = lambda o, t: (
+            self.register_var_name
+        )
+        o.load_from_register.type = lambda o, t: pir.a_type(t.output.type)
+        o.load_from_register.symbolic_shape_or_data = lambda o, t: pir.a_symbol(
+            t.output.get_symbolic_shape_or_data()
+        )
+        o.load_from_register([], [t.output])
+
+
+class ReplaceWithStoreToRegisterPass(access_topo_drr.DrrPass):
+    def __init__(self, name, register_var_name):
+        self.name = pir.a_str(name)
+        self.register_var_name = pir.a_str(register_var_name)
+
+    def source_pattern(self, o, t):
+        o.data_op = o.ap_native_op("pd_op.data")
+        o.data_op([], [t.output])
+        o.store_to_global_op = o.ap_native_op("ap_op.store_to_global")
+        o.store_to_global_op.index_func_unique_id = self.name
+        o.store_to_global_op([t.output, t.output_val], [])
+
+    def result_pattern(self, o, t):
+        o.store_to_register_op = o.ap_native_op("ap_op.store_to_register")
+        o.store_to_register_op.name = lambda o, t: self.name
+        o.store_to_register_op.register_var_name = lambda o, t: (
+            self.register_var_name
+        )
+        o.store_to_register_op([t.output_val], [])
+
+
+@access_topo_drr.register_drr_pass("down_spider_relu", tag="default")
+class DownSpiderReluAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.spider0 = o.ap_native_op("ap_op.down_spider")
+        o.spider0([t.input], [t.tmp])
+        o.relu1 = o.ap_native_op("pd_op.relu")
+        o.relu1([t.tmp], [t.output])
+
+    def result_pattern(self, o, t):
+        o.fustion_op = o.ap_native_op("ap_op.down_spider")
+        o.fustion_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass(
+    "down_spider_load_from_global", tag="default"
+)
+class DownSpiderLoadFromGlobalAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.spider0 = o.ap_native_op("ap_op.down_spider")
+        o.spider0([t.input], [t.tmp])
+        o.load_from_global_op = o.ap_native_op("ap_op.load_from_global")
+        o.load_from_global_op([t.tmp], [t.output])
+
+    def result_pattern(self, o, t):
+        o.fustion_op = o.ap_native_op("ap_op.down_spider")
+        o.fustion_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("down_spider_up_spider", tag="default")
+class DownSpiderUpSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.down_spider_op = o.ap_native_op("ap_op.down_spider")
+        o.down_spider_op([t.input], [t.tmp0])
+        o.up_spider_op = o.ap_native_op("ap_op.up_spider")
+        o.up_spider_op([t.tmp0, t.input], [])
+
+    def result_pattern(self, o, t):
+        pass
+
+
+@access_topo_drr.register_drr_pass("left_down_spider_add", tag="default")
+class LeftDownSpiderAddAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.spider = o.ap_native_op("ap_op.down_spider")
+        o.spider([t.input0], [t.tmp0])
+        o.add = o.ap_native_op("pd_op.add")
+        o.add([t.tmp0, t.input1], [t.output])
+
+    def result_pattern(self, o, t):
+        o.down_spider = o.ap_native_op("ap_op.down_spider")
+        o.down_spider([t.input0], [t.output])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.input1], [])
+
+
+@access_topo_drr.register_drr_pass("right_down_spider_add", tag="default")
+class RightDownSpiderAddAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.spider = o.ap_native_op("ap_op.down_spider")
+        o.spider([t.input0], [t.tmp0])
+        o.add = o.ap_native_op("pd_op.add")
+        o.add([t.input1, t.tmp0], [t.output])
+
+    def result_pattern(self, o, t):
+        o.down_spider = o.ap_native_op("ap_op.down_spider")
+        o.down_spider([t.input0], [t.output])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.input1], [])
+
+
+@access_topo_drr.register_drr_pass("expand_up_spider", tag="default")
+class ExpandUpSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.expand = o.ap_native_op("pd_op.expand")
+        o.expand([t.input1, t.input2], [t.expanded_input])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.expanded_input], [])
+
+    def constraint(self, o, t):
+        input_shape = t.input1.symbolic_shape_to_list()
+        output_shape = t.expanded_input.symbolic_shape_to_list()
+        rank_diff = len(output_shape) - len(input_shape)
+        return rank_diff > 0
+
+        # TODO: Get Inner Expanded Axes
+        def GetInnerExpanded_axes(i):
+            if input_shape[i] == output_shape[i + rank_diff]:
+                return []
+            else:
+                return [i]
+
+        input_rank = len(t.input1.symbolic_shape_to_list())
+        inner_expanded_axes = ap.flat_map(
+            GetInnerExpanded_axes, range(input_rank)
+        )
+        return rank_diff > 0 and len(inner_expanded_axes) == 0
+
+    def result_pattern(self, o, t):
+        t.declare_internal_native_ir_value("reduced_input")
+        o.sum = o.ap_native_op("pd_op.sum")
+        o.sum.axis = self.get_axis
+        o.sum.keepdim = self.get_keepdim
+        o.sum([t.input0], [t.reduced_input])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.reduced_input, t.input1], [])
+
+    def get_keepdim(self, o, t):
+        return pir.a_bool(False)
+
+    def get_axis(self, o, t):
+        input_rank = len(t.input1.symbolic_shape_to_list())
+        output_rank = len(t.expanded_input.symbolic_shape_to_list())
+        axes = range(output_rank - input_rank)
+        return pir.a_intarray(axes)
+
+
+@access_topo_drr.register_drr_pass("cinn_broadcast_up_spider", tag="default")
+class CinnBroadcastUpSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.broadcast_op = o.ap_native_op("cinn_op.broadcast")
+        o.broadcast_op([t.input1], [t.expanded_input])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.expanded_input], [])
+
+    def constraint(self, o, t):
+        input_shape = t.input1.symbolic_shape_to_list()
+        output_shape = t.expanded_input.symbolic_shape_to_list()
+        rank_diff = len(output_shape) - len(input_shape)
+        return rank_diff > 0
+
+        # TODO: Get Inner Expanded Axes
+        def GetInnerExpanded_axes(i):
+            if input_shape[i] == output_shape[i + rank_diff]:
+                return []
+            else:
+                return [i]
+
+        input_rank = len(t.input1.symbolic_shape_to_list())
+        inner_expanded_axes = ap.flat_map(
+            GetInnerExpanded_axes, range(input_rank)
+        )
+        return rank_diff > 0 and len(inner_expanded_axes) == 0
+
+    def result_pattern(self, o, t):
+        t.declare_internal_native_ir_value("reduced_input")
+        o.sum = o.ap_native_op("pd_op.sum")
+        o.sum.axis = self.get_axis
+        o.sum.keepdim = self.get_keepdim
+        o.sum([t.input0], [t.reduced_input])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.reduced_input, t.input1], [])
+
+    def get_keepdim(self, o, t):
+        return pir.a_bool(False)
+
+    def get_axis(self, o, t):
+        input_rank = len(t.input1.symbolic_shape_to_list())
+        output_rank = len(t.expanded_input.symbolic_shape_to_list())
+        axes = range(output_rank - input_rank)
+        return pir.a_intarray(axes)
+
+
+@access_topo_drr.register_drr_pass("right_down_spider_up_spider", tag="default")
+class RightDownSpiderUpSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.expand = o.ap_native_op("ap_op.down_spider")
+        o.expand([t.input1], [t.output1])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.output1], [])
+
+    def result_pattern(self, o, t):
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.input1], [])
+
+
+@access_topo_drr.register_drr_pass("left_down_spider_up_spider", tag="default")
+class LeftDownSpiderUpSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.expand = o.ap_native_op("ap_op.down_spider")
+        o.expand([t.input0], [t.output0])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.output0, t.input1], [])
+
+    def result_pattern(self, o, t):
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.input1], [])
+
+
+@access_topo_drr.register_drr_pass(
+    "triangle_left_down_spider_up_spider", tag="default"
+)
+class TriangleLeftDownSpiderUpSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.expand = o.ap_native_op("ap_op.down_spider")
+        o.expand([t.input0], [t.output0])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.output0], [])
+
+    def result_pattern(self, o, t):
+        pass

python/paddle/apy/matmul_pass/umprime.py

@@ -0,0 +1,59 @@
+# 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 access_topo_drr
+import ap
+import pir
+
+
+@access_topo_drr.register_drr_pass("pd_op_static_relu", tag="umprime")
+class PdOpReluAccessTopoPass(access_topo_drr.DrrPass):
+    def __init__(self):
+        self.zero = pir.a_f64(ap.DataValue.float64("0"))
+
+    def source_pattern(self, o, t):
+        o.full_op = o.ap_native_op("pd_op.full")
+        o.full_op([], [t.intermediate])
+        o.maximum_op = o.ap_native_op("pd_op.maximum")
+        o.maximum_op([t.input, t.intermediate], [t.output])
+
+    def constraint(self, o, t):
+        return o.full_op.value == self.zero
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_dynamic_relu", tag="umprime")
+class PdOpDynReluAccessTopoPass(access_topo_drr.DrrPass):
+    def __init__(self):
+        self.zero = pir.a_f64(ap.DataValue.float64("0"))
+
+    def source_pattern(self, o, t):
+        o.full_op = o.ap_native_op("pd_op.full")
+        o.full_op([], [t.intermediate0])
+        o.generate_shape_op = o.ap_native_op("cinn_op.generate_shape")
+        o.generate_shape_op([t.input0], [t.intermediate1])
+        o.expand_op = o.ap_native_op("pd_op.expand")
+        o.expand_op([t.intermediate0, t.intermediate1], [t.intermediate2])
+        o.maximum_op = o.ap_native_op("pd_op.maximum")
+        o.maximum_op([t.input1, t.intermediate2], [t.output])
+
+    def constraint(self, o, t):
+        return o.full_op.value == self.zero
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input1], [t.output])

python/paddle/apy/sys/__builtin__.py

@@ -0,0 +1,40 @@
+# 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.
+
+DataType = __builtin__DataType  # noqa: F821
+DataValue = __builtin__DataValue  # noqa: F821
+PointerType = __builtin__PointerType  # noqa: F821
+PointerValue = __builtin__PointerValue  # noqa: F821
+MutableList = __builtin__MutableList  # noqa: F821
+OrderedDict = __builtin__OrderedDict  # noqa: F821
+MutableOrderedDict = __builtin__MutableOrderedDict  # noqa: F821
+AttrMap = __builtin__AttrMap  # noqa: F821
+SerializableAttrMap = __builtin__BuiltinSerializableAttrMap  # noqa: F821
+
+_raise = __builtin__raise  # noqa: F821
+
+foreach = __builtin__foreach  # noqa: F821
+range = __builtin__range  # noqa: F821
+map = __builtin__map  # noqa: F821
+reduce = __builtin__reduce  # noqa: F821
+filter = __builtin__filter  # noqa: F821
+zip = __builtin__zip  # noqa: F821
+flat_map = __builtin__flat_map  # noqa: F821
+apply = __builtin__apply  # noqa: F821
+replace_or_trim_left_comma = __builtin__replace_or_trim_left_comma  # noqa: F821
+
+sorted = __builtin__sorted  # noqa: F821
+
+dirname = __builtin__dirname  # noqa: F821
+basename = __builtin__basename  # noqa: F821

python/paddle/apy/sys/__builtin_registry_item__.py

@@ -0,0 +1,62 @@
+# 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 __builtin__
+
+
+class RegistryEntry:
+    def __init__(self):
+        self.__tag_name__ = None
+        self.__nice__ = None
+        self.__values__ = __builtin__.MutableList()
+        self.__child_register_item_name2value__ = (
+            __builtin__.MutableOrderedDict()
+        )
+
+    # tag_name: str
+    # nice: int
+    def __getattr__(self, attrname):
+        def contains():
+            return self.__child_register_item_name2value__.contains(attrname)
+
+        def find():
+            return self.__child_register_item_name2value__[attrname]
+
+        def create():
+            register_entry = RegistryEntry()
+            self.__child_register_item_name2value__[attrname] = register_entry
+            return register_entry
+
+        return find() if contains() else create()
+
+    def __call__(self, tag_name, nice):
+        registry_obj = RegistryObject(tag_name, nice)
+        self.__values__.append(registry_obj)
+        return RegisterItemDecorator(registry_obj)
+
+
+class RegistryObject:
+    def __init__(self, tag_name, nice):
+        self.tag_name = tag_name
+        self.nice = nice
+        self.value = None
+
+
+class RegisterItemDecorator:
+    def __init__(self, register_obj):
+        self.register_obj = register_obj
+
+    def __call__(self, value):
+        self.register_obj.value = value
+        return value

python/paddle/apy/sys/__builtin_trivial_op_names__.py

@@ -0,0 +1,40 @@
+# 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.
+
+
+def GetGroupedTrivialOpNames():
+    return [
+        "pd_op.sin",
+        "pd_op.add",
+        "pd_op.relu",
+        "pd_op.data",
+        "pd_op.full",
+        "pd_op.cast",
+        "pd_op.exp",
+        "pd_op.relu",
+        "pd_op.tanh",
+        "pd_op.floor",
+        "pd_op.erf",
+        "pd_op.elementwise_pow",
+        "cinn_op.scale",
+        "pd_op.subtract",
+        "pd_op.add",
+        "pd_op.multiply",
+        "pd_op.divide",
+        "pd_op.maximum",
+        "cinn_op.yield_store",
+        "cinn_op.broadcast",
+        "pd_op.expand",
+        "cinn_op.generate_shape",
+    ]

python/paddle/apy/sys/ap.py

@@ -0,0 +1,44 @@
+# 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 __builtin__
+
+DataType = __builtin__.DataType
+DataValue = __builtin__.DataValue
+PointerType = __builtin__.PointerType
+PointerValue = __builtin__.PointerValue
+MutableList = __builtin__.MutableList
+OrderedDict = __builtin__.OrderedDict
+MutableOrderedDict = __builtin__.MutableOrderedDict
+AttrMap = __builtin__.AttrMap
+SerializableAttrMap = __builtin__.SerializableAttrMap
+
+_raise = __builtin__._raise
+
+foreach = __builtin__.foreach
+range = __builtin__.range
+map = __builtin__.map
+reduce = __builtin__.reduce
+filter = __builtin__.filter
+zip = __builtin__.zip
+flat_map = __builtin__.flat_map
+apply = __builtin__.apply
+replace_or_trim_left_comma = __builtin__.replace_or_trim_left_comma
+
+registry = __builtin__registry  # noqa: F821
+
+sorted = __builtin__.sorted
+
+dirname = __builtin__.dirname
+basename = __builtin__.basename

python/paddle/audio/__init__.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2022 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 . import backends, datasets, features, functional
+from .backends.backend import info, load, save
+
+__all__ = [
+    "functional",
+    "features",
+    "datasets",
+    "backends",
+    "load",
+    "info",
+    "save",
+]

python/paddle/audio/backends/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2022 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 . import init_backend
+from .init_backend import (
+    get_current_backend,
+    list_available_backends,
+    set_backend,
+)
+
+init_backend._init_set_audio_backend()
+
+__all__ = [
+    'get_current_backend',
+    'list_available_backends',
+    'set_backend',
+]

python/paddle/audio/backends/backend.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2022 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
+
+from typing import TYPE_CHECKING, BinaryIO
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from paddle import Tensor
+
+
+class AudioInfo:
+    """Audio info, return type of backend info function"""
+
+    sample_rate: int
+    num_samples: int
+    num_channels: int
+    bits_per_sample: int
+    encoding: str
+
+    def __init__(
+        self,
+        sample_rate: int,
+        num_samples: int,
+        num_channels: int,
+        bits_per_sample: int,
+        encoding: str,
+    ) -> None:
+        self.sample_rate = sample_rate
+        self.num_samples = num_samples
+        self.num_channels = num_channels
+        self.bits_per_sample = bits_per_sample
+        self.encoding = encoding
+
+
+def info(filepath: str | BinaryIO) -> AudioInfo:
+    """Get signal information of input audio file.
+
+    Args:
+        filepath: audio path or file object.
+
+    Returns:
+        AudioInfo: info of the given audio.
+
+    Example:
+        .. code-block:: pycon
+
+            >>> import os
+            >>> import paddle
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+            >>> base_dir = os.getcwd()
+            >>> filepath = os.path.join(base_dir, "test.wav")
+
+            >>> paddle.audio.save(filepath, waveform, sample_rate)
+            >>> wav_info = paddle.audio.info(filepath)
+    """
+    # for API doc
+    raise NotImplementedError("please set audio backend")
+
+
+def load(
+    filepath: str | Path,
+    frame_offset: int = 0,
+    num_frames: int = -1,
+    normalize: bool = True,
+    channels_first: bool = True,
+) -> tuple[Tensor, int]:
+    """Load audio data from file.Load the audio content start form frame_offset, and get num_frames.
+
+    Args:
+        frame_offset: from 0 to total frames,
+        num_frames: from -1 (means total frames) or number frames which want to read,
+        normalize:
+            if True: return audio which norm to (-1, 1), dtype=float32
+            if False: return audio with raw data, dtype=int16
+
+        channels_first:
+            if True: return audio with shape (channels, time)
+
+    Return:
+        Tuple[paddle.Tensor, int]: (audio_content, sample rate)
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import os
+            >>> import paddle
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+            >>> base_dir = os.getcwd()
+            >>> filepath = os.path.join(base_dir, "test.wav")
+
+            >>> paddle.audio.save(filepath, waveform, sample_rate)
+            >>> wav_data_read, sr = paddle.audio.load(filepath)
+    """
+    # for API doc
+    raise NotImplementedError("please set audio backend")
+
+
+def save(
+    filepath: str,
+    src: Tensor,
+    sample_rate: int,
+    channels_first: bool = True,
+    encoding: str | None = None,
+    bits_per_sample: int | None = 16,
+) -> None:
+    """
+    Save audio tensor to file.
+
+    Args:
+        filepath: saved path
+        src: the audio tensor
+        sample_rate: the number of samples of audio per second.
+        channels_first: src channel information
+            if True, means input tensor is (channels, time)
+            if False, means input tensor is (time, channels)
+        encoding:encoding format, wave_backend only support PCM16 now.
+        bits_per_sample: bits per sample, wave_backend only support 16 bits now.
+
+    Returns:
+        None
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+            >>> filepath = "./test.wav"
+
+            >>> paddle.audio.save(filepath, waveform, sample_rate)
+    """
+    # for API doc
+    raise NotImplementedError("please set audio backend")

python/paddle/audio/backends/init_backend.py

@@ -0,0 +1,192 @@
+# Copyright (c) 2022 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 sys
+import warnings
+
+import paddle
+
+from . import backend, wave_backend
+
+
+def _check_version(version: str) -> bool:
+    # require paddleaudio >= 1.0.2
+    ver_arr = version.split('.')
+    v0 = int(ver_arr[0])
+    v1 = int(ver_arr[1])
+    v2 = int(ver_arr[2])
+    if v0 < 1:
+        return False
+    if v0 == 1 and v1 == 0 and v2 <= 1:
+        return False
+    return True
+
+
+def list_available_backends() -> list[str]:
+    """List available backends, the backends in paddleaudio and the default backend.
+
+    Returns:
+        list[str]: The list of available backends.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+            >>> wav_path = "./test.wav"
+
+            >>> current_backend = paddle.audio.backends.get_current_backend()
+            >>> print(current_backend)
+            wave_backend
+
+            >>> backends = paddle.audio.backends.list_available_backends()
+            >>> # default backends is ['wave_backend']
+            >>> # backends is ['wave_backend', 'soundfile'], if have installed paddleaudio >= 1.0.2
+            >>> if 'soundfile' in backends:
+            ...     paddle.audio.backends.set_backend('soundfile')
+            >>> paddle.audio.save(wav_path, waveform, sample_rate)
+
+    """
+    backends = []
+    try:
+        import paddleaudio
+    except ImportError:
+        package = "paddleaudio"
+        warn_msg = (
+            f"Failed importing {package}. \n"
+            "only wave_backend(only can deal with PCM16 WAV) supported.\n"
+            "if want soundfile_backend(more audio type supported),\n"
+            f"please manually installed (usually with `pip install {package} >= 1.0.2`). "
+        )
+        warnings.warn(warn_msg)
+
+    if "paddleaudio" in sys.modules:
+        version = paddleaudio.__version__
+        if not _check_version(version):
+            err_msg = (
+                f"the version of paddleaudio installed is {version},\n"
+                "please ensure the paddleaudio >= 1.0.2."
+            )
+            raise ImportError(err_msg)
+        backends = paddleaudio.backends.list_audio_backends()
+    backends.append("wave_backend")
+    return backends
+
+
+def get_current_backend() -> str:
+    """Get the name of the current audio backend
+
+    Returns:
+        str: The name of the current backend,
+        the wave_backend or backend imported from paddleaudio
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+            >>> wav_path = "./test.wav"
+
+            >>> current_backend = paddle.audio.backends.get_current_backend()
+            >>> print(current_backend)
+            wave_backend
+
+            >>> backends = paddle.audio.backends.list_available_backends()
+            >>> # default backends is ['wave_backend']
+            >>> # backends is ['wave_backend', 'soundfile'], if have installed paddleaudio >= 1.0.2
+
+            >>> if 'soundfile' in backends:
+            ...     paddle.audio.backends.set_backend('soundfile')
+            >>> paddle.audio.save(wav_path, waveform, sample_rate)
+
+    """
+    current_backend = None
+    if "paddleaudio" in sys.modules:
+        import paddleaudio
+
+        current_backend = paddleaudio.backends.get_audio_backend()
+        if paddle.audio.load == paddleaudio.load:
+            return current_backend
+    return "wave_backend"
+
+
+def set_backend(backend_name: str) -> None:
+    """Set the backend by one of the list_audio_backend return.
+
+    Args:
+        backend (str): one of the list_audio_backend. "wave_backend" is the default. "soundfile" imported from paddleaudio.
+
+    Returns:
+        None
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+            >>> wav_path = "./test.wav"
+
+            >>> current_backend = paddle.audio.backends.get_current_backend()
+            >>> print(current_backend)
+            wave_backend
+
+            >>> backends = paddle.audio.backends.list_available_backends()
+            >>> # default backends is ['wave_backend']
+            >>> # backends is ['wave_backend', 'soundfile'], if have installed paddleaudio >= 1.0.2
+
+            >>> if 'soundfile' in backends:
+            ...     paddle.audio.backends.set_backend('soundfile')
+            >>> paddle.audio.save(wav_path, waveform, sample_rate)
+
+    """
+    if backend_name not in list_available_backends():
+        raise NotImplementedError
+
+    if backend_name == "wave_backend":
+        module = wave_backend
+    else:
+        import paddleaudio
+
+        paddleaudio.backends.set_audio_backend(backend_name)
+        module = paddleaudio
+
+    for func in ["save", "load", "info"]:
+        setattr(backend, func, getattr(module, func))
+        setattr(paddle.audio, func, getattr(module, func))
+
+
+def _init_set_audio_backend() -> None:
+    # init the default wave_backend.
+    for func in ["save", "load", "info"]:
+        setattr(backend, func, getattr(wave_backend, func))

python/paddle/audio/backends/wave_backend.py

@@ -0,0 +1,236 @@
+# Copyright (c) 2022 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 wave
+from typing import TYPE_CHECKING, BinaryIO
+
+import numpy as np
+
+import paddle
+
+from .backend import AudioInfo
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from paddle import Tensor
+
+
+def _error_message():
+    package = "paddleaudio"
+    warn_msg = (
+        "only PCM16 WAV supported. \n"
+        "if want support more other audio types, please "
+        f"manually installed (usually with `pip install {package}`). \n "
+        "and use paddle.audio.backends.set_backend('soundfile') to set audio backend"
+    )
+    return warn_msg
+
+
+def info(filepath: str | BinaryIO) -> AudioInfo:
+    """Get signal information of input audio file.
+
+    Args:
+        filepath: audio path or file object.
+
+    Returns:
+        AudioInfo: info of the given audio.
+
+    Example:
+        .. code-block:: pycon
+
+            >>> import os
+            >>> import paddle
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+            >>> base_dir = os.getcwd()
+            >>> filepath = os.path.join(base_dir, "test.wav")
+
+            >>> paddle.audio.save(filepath, waveform, sample_rate)
+            >>> wav_info = paddle.audio.info(filepath)
+    """
+
+    if hasattr(filepath, 'read'):
+        file_obj = filepath
+    else:
+        file_obj = open(filepath, 'rb')
+
+    try:
+        file_ = wave.open(file_obj)
+    except wave.Error:
+        file_obj.seek(0)
+        file_obj.close()
+        err_msg = _error_message()
+        raise NotImplementedError(err_msg)
+
+    channels = file_.getnchannels()
+    sample_rate = file_.getframerate()
+    sample_frames = file_.getnframes()  # audio frame
+    bits_per_sample = file_.getsampwidth() * 8
+    encoding = "PCM_S"  # default WAV encoding, only support
+    file_obj.close()
+    return AudioInfo(
+        sample_rate, sample_frames, channels, bits_per_sample, encoding
+    )
+
+
+def load(
+    filepath: str | Path,
+    frame_offset: int = 0,
+    num_frames: int = -1,
+    normalize: bool = True,
+    channels_first: bool = True,
+) -> tuple[Tensor, int]:
+    """Load audio data from file. load the audio content start form frame_offset, and get num_frames.
+
+    Args:
+        frame_offset: from 0 to total frames,
+        num_frames: from -1 (means total frames) or number frames which want to read,
+        normalize:
+            if True: return audio which norm to (-1, 1), dtype=float32
+            if False: return audio with raw data, dtype=int16
+
+        channels_first:
+            if True: return audio with shape (channels, time)
+
+    Return:
+        Tuple[paddle.Tensor, int]: (audio_content, sample rate)
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import os
+            >>> import paddle
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+            >>> base_dir = os.getcwd()
+            >>> filepath = os.path.join(base_dir, "test.wav")
+
+            >>> paddle.audio.save(filepath, waveform, sample_rate)
+            >>> wav_data_read, sr = paddle.audio.load(filepath)
+    """
+    if hasattr(filepath, 'read'):
+        file_obj = filepath
+    else:
+        file_obj = open(filepath, 'rb')
+
+    try:
+        file_ = wave.open(file_obj)
+    except wave.Error:
+        file_obj.seek(0)
+        file_obj.close()
+        err_msg = _error_message()
+        raise NotImplementedError(err_msg)
+
+    channels = file_.getnchannels()
+    sample_rate = file_.getframerate()
+    frames = file_.getnframes()  # audio frame
+
+    audio_content = file_.readframes(frames)
+    file_obj.close()
+
+    # default_subtype = "PCM_16", only support PCM16 WAV
+    audio_as_np16 = np.frombuffer(audio_content, dtype=np.int16)
+    audio_as_np32 = audio_as_np16.astype(np.float32)
+    if normalize:
+        # dtype = "float32"
+        audio_norm = audio_as_np32 / (2**15)
+    else:
+        # dtype = "int16"
+        audio_norm = audio_as_np32
+
+    waveform = np.reshape(audio_norm, (frames, channels))
+    if num_frames != -1:
+        waveform = waveform[frame_offset : frame_offset + num_frames, :]
+    waveform = paddle.to_tensor(waveform)
+    if channels_first:
+        waveform = paddle.transpose(waveform, perm=[1, 0])
+    return waveform, sample_rate
+
+
+def save(
+    filepath: str,
+    src: Tensor,
+    sample_rate: int,
+    channels_first: bool = True,
+    encoding: str | None = None,
+    bits_per_sample: int | None = 16,
+) -> None:
+    """
+    Save audio tensor to file.
+
+    Args:
+        filepath: saved path
+        src: the audio tensor
+        sample_rate: the number of samples of audio per second.
+        channels_first: src channel information
+            if True, means input tensor is (channels, time)
+            if False, means input tensor is (time, channels)
+        encoding: audio encoding format, wave_backend only support PCM16 now.
+        bits_per_sample: bits per sample, wave_backend only support 16 bits now.
+
+    Returns:
+        None
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+            >>> filepath = "./test.wav"
+
+            >>> paddle.audio.save(filepath, waveform, sample_rate)
+    """
+    assert src.ndim == 2, "Expected 2D tensor"
+
+    audio_numpy = src.numpy()
+
+    # change src shape to (time, channels)
+    if channels_first:
+        audio_numpy = np.transpose(audio_numpy)
+
+    channels = audio_numpy.shape[1]
+
+    # only support PCM16
+    if bits_per_sample not in (None, 16):
+        raise ValueError("Invalid bits_per_sample, only support 16 bit")
+
+    sample_width = int(bits_per_sample / 8)  # 2
+
+    if src.dtype == paddle.float32:
+        audio_numpy = (audio_numpy * (2**15)).astype("<h")
+
+    with wave.open(filepath, 'w') as f:
+        f.setnchannels(channels)
+        f.setsampwidth(sample_width)
+        f.setframerate(sample_rate)
+        f.writeframes(audio_numpy.tobytes())

python/paddle/audio/datasets/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2022 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 .esc50 import ESC50
+from .tess import TESS
+
+__all__ = ["ESC50", "TESS"]

python/paddle/audio/datasets/dataset.py

@@ -0,0 +1,98 @@
+# Copyright (c) 2022 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 paddle
+
+from ..features import MFCC, LogMelSpectrogram, MelSpectrogram, Spectrogram
+
+feat_funcs = {
+    'raw': None,
+    'melspectrogram': MelSpectrogram,
+    'mfcc': MFCC,
+    'logmelspectrogram': LogMelSpectrogram,
+    'spectrogram': Spectrogram,
+}
+
+
+class AudioClassificationDataset(paddle.io.Dataset):
+    """
+    Base class of audio classification dataset.
+    """
+
+    def __init__(
+        self,
+        files: list[str],
+        labels: list[int],
+        feat_type: str = 'raw',
+        sample_rate: int | None = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            files (:obj:`List[str]`): A list of absolute path of audio files.
+            labels (:obj:`List[int]`): Labels of audio files.
+            feat_type (:obj:`str`, `optional`, defaults to `raw`):
+                It identifies the feature type that user wants to extract an audio file.
+        """
+        super().__init__()
+
+        if feat_type not in feat_funcs.keys():
+            raise RuntimeError(
+                f"Unknown feat_type: {feat_type}, it must be one in {list(feat_funcs.keys())}"
+            )
+
+        self.files = files
+        self.labels = labels
+
+        self.feat_type = feat_type
+        self.sample_rate = sample_rate
+        self.feat_config = (
+            kwargs  # Pass keyword arguments to customize feature config
+        )
+
+    def _get_data(self, input_file: str):
+        raise NotImplementedError
+
+    def _convert_to_record(self, idx):
+        file, label = self.files[idx], self.labels[idx]
+        waveform, sample_rate = paddle.audio.load(file)
+        self.sample_rate = sample_rate
+
+        feat_func = feat_funcs[self.feat_type]
+
+        record = {}
+        if len(waveform.shape) == 2:
+            waveform = waveform.squeeze(0)  # 1D input
+        waveform = paddle.to_tensor(waveform, dtype=paddle.float32)
+        if feat_func is not None:
+            waveform = waveform.unsqueeze(0)  # (batch_size, T)
+            if self.feat_type != 'spectrogram':
+                feature_extractor = feat_func(
+                    sr=self.sample_rate, **self.feat_config
+                )
+            else:
+                feature_extractor = feat_func(**self.feat_config)
+            record['feat'] = feature_extractor(waveform).squeeze(0)
+        else:
+            record['feat'] = waveform
+        record['label'] = label
+        return record
+
+    def __getitem__(self, idx):
+        record = self._convert_to_record(idx)
+        return record['feat'], record['label']
+
+    def __len__(self):
+        return len(self.files)

python/paddle/audio/datasets/esc50.py

@@ -0,0 +1,227 @@
+# Copyright (c) 2021 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 os
+from typing import TYPE_CHECKING, Any, Literal, NamedTuple, TypeAlias
+
+from paddle.dataset.common import DATA_HOME
+from paddle.utils import download
+
+from .dataset import AudioClassificationDataset
+
+if TYPE_CHECKING:
+    _ModeLiteral: TypeAlias = Literal[
+        'train',
+        'dev',
+    ]
+    _FeatTypeLiteral: TypeAlias = Literal[
+        'raw',
+        'melspectrogram',
+        'mfcc',
+        'logmelspectrogram',
+        'spectrogram',
+    ]
+
+
+__all__ = []
+
+
+class ESC50(AudioClassificationDataset):
+    """
+    The ESC-50 dataset is a labeled collection of 2000 environmental audio recordings
+    suitable for benchmarking methods of environmental sound classification. The dataset
+    consists of 5-second-long recordings organized into 50 semantical classes (with
+    40 examples per class)
+
+    Reference:
+        ESC: Dataset for Environmental Sound Classification
+        http://dx.doi.org/10.1145/2733373.2806390
+
+    Args:
+       mode (str, optional): It identifies the dataset mode (train or dev). Default:train.
+       split (int, optional): It specify the fold of dev dataset. Default:1.
+       feat_type (str, optional): It identifies the feature type that user wants to extract of an audio file. Default:raw.
+       archive(dict, optional): it tells where to download the audio archive. Default:None.
+
+    Returns:
+        :ref:`api_paddle_io_Dataset`. An instance of ESC50 dataset.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +TIMEOUT(60)
+            >>> import paddle
+
+            >>> esc50_dataset = paddle.audio.datasets.ESC50(
+            ...     mode='dev',
+            ...     feat_type='raw',
+            ... )
+            >>> for idx in range(5):
+            ...     audio, label = esc50_dataset[idx]
+            ...     # do something with audio, label
+            ...     print(audio.shape, label)
+            ...     # [audio_data_length] , label_id
+            paddle.Size([220500]) 0
+            paddle.Size([220500]) 14
+            paddle.Size([220500]) 36
+            paddle.Size([220500]) 36
+            paddle.Size([220500]) 19
+
+            >>> esc50_dataset = paddle.audio.datasets.ESC50(
+            ...     mode='dev',
+            ...     feat_type='mfcc',
+            ...     n_mfcc=40,
+            ... )
+            >>> for idx in range(5):
+            ...     audio, label = esc50_dataset[idx]
+            ...     # do something with mfcc feature, label
+            ...     print(audio.shape, label)
+            ...     # [feature_dim, length] , label_id
+            paddle.Size([40, 1723]) 0
+            paddle.Size([40, 1723]) 14
+            paddle.Size([40, 1723]) 36
+            paddle.Size([40, 1723]) 36
+            paddle.Size([40, 1723]) 19
+
+    """
+
+    archive: dict[str, str] = {
+        'url': 'https://paddleaudio.bj.bcebos.com/datasets/ESC-50-master.zip',
+        'md5': '7771e4b9d86d0945acce719c7a59305a',
+    }
+
+    label_list: list[str] = [
+        # Animals
+        'Dog',
+        'Rooster',
+        'Pig',
+        'Cow',
+        'Frog',
+        'Cat',
+        'Hen',
+        'Insects (flying)',
+        'Sheep',
+        'Crow',
+        # Natural soundscapes & water sounds
+        'Rain',
+        'Sea waves',
+        'Crackling fire',
+        'Crickets',
+        'Chirping birds',
+        'Water drops',
+        'Wind',
+        'Pouring water',
+        'Toilet flush',
+        'Thunderstorm',
+        # Human, non-speech sounds
+        'Crying baby',
+        'Sneezing',
+        'Clapping',
+        'Breathing',
+        'Coughing',
+        'Footsteps',
+        'Laughing',
+        'Brushing teeth',
+        'Snoring',
+        'Drinking, sipping',
+        # Interior/domestic sounds
+        'Door knock',
+        'Mouse click',
+        'Keyboard typing',
+        'Door, wood creaks',
+        'Can opening',
+        'Washing machine',
+        'Vacuum cleaner',
+        'Clock alarm',
+        'Clock tick',
+        'Glass breaking',
+        # Exterior/urban noises
+        'Helicopter',
+        'Chainsaw',
+        'Siren',
+        'Car horn',
+        'Engine',
+        'Train',
+        'Church bells',
+        'Airplane',
+        'Fireworks',
+        'Hand saw',
+    ]
+    meta: str = os.path.join('ESC-50-master', 'meta', 'esc50.csv')
+    audio_path: str = os.path.join('ESC-50-master', 'audio')
+
+    class meta_info(NamedTuple):
+        filename: str
+        fold: str
+        target: str
+        category: str
+        esc10: str
+        src_file: str
+        take: str
+
+    def __init__(
+        self,
+        mode: _ModeLiteral = 'train',
+        split: int = 1,
+        feat_type: _FeatTypeLiteral = 'raw',
+        archive: dict[str, str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        assert split in range(1, 6), (
+            f'The selected split should be integer, and 1 <= split <= 5, but got {split}'
+        )
+        if archive is not None:
+            self.archive = archive
+        files, labels = self._get_data(mode, split)
+        super().__init__(
+            files=files, labels=labels, feat_type=feat_type, **kwargs
+        )
+
+    def _get_meta_info(self) -> list[meta_info]:
+        ret = []
+        with open(os.path.join(DATA_HOME, self.meta), 'r') as rf:
+            for line in rf.readlines()[1:]:
+                ret.append(self.meta_info(*line.strip().split(',')))
+        return ret
+
+    def _get_data(
+        self, mode: _ModeLiteral, split: int
+    ) -> tuple[list[str], list[int]]:
+        if not os.path.isdir(
+            os.path.join(DATA_HOME, self.audio_path)
+        ) or not os.path.isfile(os.path.join(DATA_HOME, self.meta)):
+            download.get_path_from_url(
+                self.archive['url'],
+                DATA_HOME,
+                self.archive['md5'],
+                decompress=True,
+            )
+
+        meta_info = self._get_meta_info()
+
+        files = []
+        labels = []
+        for sample in meta_info:
+            filename, fold, target, _, _, _, _ = sample
+            if mode == 'train' and int(fold) != split:
+                files.append(os.path.join(DATA_HOME, self.audio_path, filename))
+                labels.append(int(target))
+
+            if mode != 'train' and int(fold) == split:
+                files.append(os.path.join(DATA_HOME, self.audio_path, filename))
+                labels.append(int(target))
+
+        return files, labels

python/paddle/audio/datasets/tess.py

@@ -0,0 +1,166 @@
+# Copyright (c) 2022 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 os
+from typing import TYPE_CHECKING, Any, NamedTuple
+
+from paddle.dataset.common import DATA_HOME
+from paddle.utils import download
+
+from .dataset import AudioClassificationDataset
+
+if TYPE_CHECKING:
+    from .esc50 import _FeatTypeLiteral, _ModeLiteral
+
+__all__ = []
+
+
+class TESS(AudioClassificationDataset):
+    """
+    TESS is a set of 200 target words were spoken in the carrier phrase
+    "Say the word _____' by two actresses (aged 26 and 64 years) and
+    recordings were made of the set portraying each of seven emotions(anger,
+    disgust, fear, happiness, pleasant surprise, sadness, and neutral).
+    There are 2800 stimuli in total.
+
+    Reference:
+        Toronto emotional speech set (TESS) https://tspace.library.utoronto.ca/handle/1807/24487
+        https://doi.org/10.5683/SP2/E8H2MF
+
+    Args:
+       mode (str, optional): It identifies the dataset mode (train or dev). Defaults to train.
+       n_folds (int, optional): Split the dataset into n folds. 1 fold for dev dataset and n-1 for train dataset. Defaults to 5.
+       split (int, optional): It specify the fold of dev dataset. Defaults to 1.
+       feat_type (str, optional): It identifies the feature type that user wants to extract of an audio file. Defaults to raw.
+       archive(dict): it tells where to download the audio archive. Defaults to None.
+
+    Returns:
+        :ref:`api_paddle_io_Dataset`. An instance of TESS dataset.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +TIMEOUT(60)
+            >>> import paddle
+
+            >>> tess_dataset = paddle.audio.datasets.TESS(
+            ...     mode='dev',
+            ...     feat_type='raw',
+            ... )
+            >>> for idx in range(5):
+            ...     audio, label = tess_dataset[idx]
+            ...     # do something with audio, label
+            ...     print(audio.shape, label)
+            ...     # [audio_data_length] , label_id
+
+            >>> tess_dataset = paddle.audio.datasets.TESS(
+            ...     mode='dev',
+            ...     feat_type='mfcc',
+            ...     n_mfcc=40,
+            ... )
+            >>> for idx in range(5):
+            ...     audio, label = tess_dataset[idx]
+            ...     # do something with mfcc feature, label
+            ...     print(audio.shape, label)
+            ...     # [feature_dim, num_frames] , label_id
+    """
+
+    archive: dict[str, str] = {
+        'url': 'https://bj.bcebos.com/paddleaudio/datasets/TESS_Toronto_emotional_speech_set.zip',
+        'md5': '1465311b24d1de704c4c63e4ccc470c7',
+    }
+
+    label_list: list[str] = [
+        'angry',
+        'disgust',
+        'fear',
+        'happy',
+        'neutral',
+        'ps',  # pleasant surprise
+        'sad',
+    ]
+
+    audio_path: str = 'TESS_Toronto_emotional_speech_set'
+
+    class meta_info(NamedTuple):
+        speaker: str
+        word: str
+        emotion: str
+
+    def __init__(
+        self,
+        mode: _ModeLiteral = 'train',
+        n_folds: int = 5,
+        split: int = 1,
+        feat_type: _FeatTypeLiteral = 'raw',
+        archive: dict[str, str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        assert isinstance(n_folds, int) and (n_folds >= 1), (
+            f'the n_folds should be integer and n_folds >= 1, but got {n_folds}'
+        )
+        assert split in range(1, n_folds + 1), (
+            f'The selected split should be integer and should be 1 <= split <= {n_folds}, but got {split}'
+        )
+        if archive is not None:
+            self.archive = archive
+        files, labels = self._get_data(mode, n_folds, split)
+        super().__init__(
+            files=files, labels=labels, feat_type=feat_type, **kwargs
+        )
+
+    def _get_meta_info(self, files) -> list[meta_info]:
+        ret = []
+        for file in files:
+            basename_without_extend = os.path.basename(file)[:-4]
+            ret.append(self.meta_info(*basename_without_extend.split('_')))
+        return ret
+
+    def _get_data(
+        self, mode: str, n_folds: int, split: int
+    ) -> tuple[list[str], list[int]]:
+        if not os.path.isdir(os.path.join(DATA_HOME, self.audio_path)):
+            download.get_path_from_url(
+                self.archive['url'],
+                DATA_HOME,
+                self.archive['md5'],
+                decompress=True,
+            )
+
+        wav_files = []
+        for root, _, files in os.walk(os.path.join(DATA_HOME, self.audio_path)):
+            for file in files:
+                if file.endswith('.wav'):
+                    wav_files.append(os.path.join(root, file))
+
+        meta_info = self._get_meta_info(wav_files)
+
+        files = []
+        labels = []
+        for idx, sample in enumerate(meta_info):
+            _, _, emotion = sample
+            target = self.label_list.index(emotion)
+            fold = idx % n_folds + 1
+
+            if mode == 'train' and int(fold) != split:
+                files.append(wav_files[idx])
+                labels.append(target)
+
+            if mode != 'train' and int(fold) == split:
+                files.append(wav_files[idx])
+                labels.append(target)
+
+        return files, labels

python/paddle/audio/features/__init__.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2022 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 .layers import (
+    MFCC,
+    LogMelSpectrogram,
+    MelSpectrogram,
+    Spectrogram,
+)
+
+__all__ = [
+    'LogMelSpectrogram',
+    'MelSpectrogram',
+    'MFCC',
+    'Spectrogram',
+]

python/paddle/audio/features/layers.py

@@ -0,0 +1,448 @@
+# Copyright (c) 2022 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
+
+from functools import partial
+from typing import TYPE_CHECKING, Literal, TypeAlias
+
+import paddle
+from paddle import nn
+
+from ..functional import compute_fbank_matrix, create_dct, power_to_db
+from ..functional.window import get_window
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+_WindowLiteral: TypeAlias = Literal[
+    'hamming',
+    'hann',
+    'kaiser',
+    'bartlett',
+    'nuttall',
+    'gaussian',
+    'exponential',
+    'triang',
+    'bohman',
+    'blackman',
+    'cosine',
+    'tukey',
+    'taylor',
+]
+
+
+class Spectrogram(nn.Layer):
+    """Compute spectrogram of given signals, typically audio waveforms.
+    The spectrogram is defined as the complex norm of the short-time Fourier transformation.
+
+    Args:
+        n_fft (int, optional): The number of frequency components of the discrete Fourier transform. Defaults to 512.
+        hop_length (Optional[int], optional): The hop length of the short time FFT. If `None`, it is set to `win_length//4`. Defaults to None.
+        win_length (Optional[int], optional): The window length of the short time FFT. If `None`, it is set to same as `n_fft`. Defaults to None.
+        window (str, optional): The window function applied to the signal before the Fourier transform. Supported window functions: 'hamming', 'hann', 'gaussian', 'exponential', 'triang', 'bohman', 'blackman', 'cosine', 'tukey', 'taylor', 'bartlett', 'kaiser', 'nuttall'. Defaults to 'hann'.
+        power (float, optional): Exponent for the magnitude spectrogram. Defaults to 2.0.
+        center (bool, optional): Whether to pad `x` to make that the :math:`t \times hop\\_length` at the center of `t`-th frame. Defaults to True.
+        pad_mode (str, optional): Choose padding pattern when `center` is `True`. Defaults to 'reflect'.
+        dtype (str, optional): Data type of input and window. Defaults to 'float32'.
+
+    Returns:
+        :ref:`api_paddle_nn_Layer`. An instance of Spectrogram.
+
+
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.audio.features import Spectrogram
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+
+            >>> feature_extractor = Spectrogram(n_fft=512, window='hann', power=1.0)
+            >>> feats = feature_extractor(waveform)
+    """
+
+    power: float
+    fft_window: Tensor
+
+    def __init__(
+        self,
+        n_fft: int = 512,
+        hop_length: int | None = 512,
+        win_length: int | None = None,
+        window: _WindowLiteral = 'hann',
+        power: float = 1.0,
+        center: bool = True,
+        pad_mode: Literal['reflect'] = 'reflect',
+        dtype: str = 'float32',
+    ) -> None:
+        super().__init__()
+
+        assert power > 0, 'Power of spectrogram must be > 0.'
+        self.power = power
+
+        if win_length is None:
+            win_length = n_fft
+
+        self.fft_window = get_window(
+            window, win_length, fftbins=True, dtype=dtype
+        )
+        self._stft = partial(
+            paddle.signal.stft,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            window=self.fft_window,
+            center=center,
+            pad_mode=pad_mode,
+        )
+        self.register_buffer('fft_window', self.fft_window)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x (Tensor): Tensor of waveforms with shape `(N, T)`
+
+        Returns:
+            Tensor: Spectrograms with shape `(N, n_fft//2 + 1, num_frames)`.
+        """
+        stft = self._stft(x)
+        spectrogram = paddle.pow(paddle.abs(stft), self.power)
+        return spectrogram
+
+
+class MelSpectrogram(nn.Layer):
+    """Compute the melspectrogram of given signals, typically audio waveforms. It is computed by multiplying spectrogram with Mel filter bank matrix.
+
+    Args:
+        sr (int, optional): Sample rate. Defaults to 22050.
+        n_fft (int, optional): The number of frequency components of the discrete Fourier transform. Defaults to 512.
+        hop_length (Optional[int], optional): The hop length of the short time FFT. If `None`, it is set to `win_length//4`. Defaults to None.
+        win_length (Optional[int], optional): The window length of the short time FFT. If `None`, it is set to same as `n_fft`. Defaults to None.
+        window (str, optional): The window function applied to the signal before the Fourier transform. Supported window functions: 'hamming', 'hann', 'gaussian', 'exponential', 'triang', 'bohman', 'blackman', 'cosine', 'tukey', 'taylor', 'bartlett', 'kaiser', 'nuttall'. Defaults to 'hann'.
+        power (float, optional): Exponent for the magnitude spectrogram. Defaults to 2.0.
+        center (bool, optional): Whether to pad `x` to make that the :math:`t \times hop\\_length` at the center of `t`-th frame. Defaults to True.
+        pad_mode (str, optional): Choose padding pattern when `center` is `True`. Defaults to 'reflect'.
+        n_mels (int, optional): Number of mel bins. Defaults to 64.
+        f_min (float, optional): Minimum frequency in Hz. Defaults to 50.0.
+        f_max (Optional[float], optional): Maximum frequency in Hz. Defaults to None.
+        htk (bool, optional): Use HTK formula in computing fbank matrix. Defaults to False.
+        norm (Union[str, float], optional): Type of normalization in computing fbank matrix. Slaney-style is used by default. You can specify norm=1.0/2.0 to use customized p-norm normalization. Defaults to 'slaney'.
+        dtype (str, optional): Data type of input and window. Defaults to 'float32'.
+
+    Returns:
+        :ref:`api_paddle_nn_Layer`. An instance of MelSpectrogram.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.audio.features import MelSpectrogram
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+
+            >>> feature_extractor = MelSpectrogram(sr=sample_rate, n_fft=512, window='hann', power=1.0)
+            >>> feats = feature_extractor(waveform)
+    """
+
+    n_mels: int
+    f_min: float
+    f_max: float
+    htk: bool
+    norm: Literal['slaney'] | float
+    fbank_matrix: Tensor
+
+    def __init__(
+        self,
+        sr: int = 22050,
+        n_fft: int = 2048,
+        hop_length: int | None = 512,
+        win_length: int | None = None,
+        window: _WindowLiteral = 'hann',
+        power: float = 2.0,
+        center: bool = True,
+        pad_mode: Literal['reflect'] = 'reflect',
+        n_mels: int = 64,
+        f_min: float = 50.0,
+        f_max: float | None = None,
+        htk: bool = False,
+        norm: Literal['slaney'] | float = 'slaney',
+        dtype: str = 'float32',
+    ) -> None:
+        super().__init__()
+
+        self._spectrogram = Spectrogram(
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            window=window,
+            power=power,
+            center=center,
+            pad_mode=pad_mode,
+            dtype=dtype,
+        )
+        self.n_mels = n_mels
+        self.f_min = f_min
+        self.f_max = f_max
+        self.htk = htk
+        self.norm = norm
+        if f_max is None:
+            f_max = sr // 2
+        self.fbank_matrix = compute_fbank_matrix(
+            sr=sr,
+            n_fft=n_fft,
+            n_mels=n_mels,
+            f_min=f_min,
+            f_max=f_max,
+            htk=htk,
+            norm=norm,
+            dtype=dtype,
+        )
+        self.register_buffer('fbank_matrix', self.fbank_matrix)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x (Tensor): Tensor of waveforms with shape `(N, T)`
+
+        Returns:
+            Tensor: Mel spectrograms with shape `(N, n_mels, num_frames)`.
+        """
+        spect_feature = self._spectrogram(x)
+        mel_feature = paddle.matmul(self.fbank_matrix, spect_feature)
+        return mel_feature
+
+
+class LogMelSpectrogram(nn.Layer):
+    """Compute log-mel-spectrogram feature of given signals, typically audio waveforms.
+
+    Args:
+        sr (int, optional): Sample rate. Defaults to 22050.
+        n_fft (int, optional): The number of frequency components of the discrete Fourier transform. Defaults to 512.
+        hop_length (Optional[int], optional): The hop length of the short time FFT. If `None`, it is set to `win_length//4`. Defaults to None.
+        win_length (Optional[int], optional): The window length of the short time FFT. If `None`, it is set to same as `n_fft`. Defaults to None.
+        window (str, optional): The window function applied to the signal before the Fourier transform. Supported window functions: 'hamming', 'hann', 'gaussian', 'exponential', 'triang', 'bohman', 'blackman', 'cosine', 'tukey', 'taylor', 'bartlett', 'kaiser', 'nuttall'. Defaults to 'hann'.
+        power (float, optional): Exponent for the magnitude spectrogram. Defaults to 2.0.
+        center (bool, optional): Whether to pad `x` to make that the :math:`t \times hop\\_length` at the center of `t`-th frame. Defaults to True.
+        pad_mode (str, optional): Choose padding pattern when `center` is `True`. Defaults to 'reflect'.
+        n_mels (int, optional): Number of mel bins. Defaults to 64.
+        f_min (float, optional): Minimum frequency in Hz. Defaults to 50.0.
+        f_max (Optional[float], optional): Maximum frequency in Hz. Defaults to None.
+        htk (bool, optional): Use HTK formula in computing fbank matrix. Defaults to False.
+        norm (Union[str, float], optional): Type of normalization in computing fbank matrix. Slaney-style is used by default. You can specify norm=1.0/2.0 to use customized p-norm normalization. Defaults to 'slaney'.
+        ref_value (float, optional): The reference value. If smaller than 1.0, the db level of the signal will be pulled up accordingly. Otherwise, the db level is pushed down. Defaults to 1.0.
+        amin (float, optional): The minimum value of input magnitude. Defaults to 1e-10.
+        top_db (Optional[float], optional): The maximum db value of spectrogram. Defaults to None.
+        dtype (str, optional): Data type of input and window. Defaults to 'float32'.
+
+    Returns:
+        :ref:`api_paddle_nn_Layer`. An instance of LogMelSpectrogram.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.audio.features import LogMelSpectrogram
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+
+            >>> feature_extractor = LogMelSpectrogram(sr=sample_rate, n_fft=512, window='hann', power=1.0)
+            >>> feats = feature_extractor(waveform)
+    """
+
+    ref_value: float
+    amin: float
+    top_db: float | None
+
+    def __init__(
+        self,
+        sr: int = 22050,
+        n_fft: int = 512,
+        hop_length: int | None = None,
+        win_length: int | None = None,
+        window: _WindowLiteral = 'hann',
+        power: float = 2.0,
+        center: bool = True,
+        pad_mode: Literal['reflect'] = 'reflect',
+        n_mels: int = 64,
+        f_min: float = 50.0,
+        f_max: float | None = None,
+        htk: bool = False,
+        norm: Literal['slaney'] | float = 'slaney',
+        ref_value: float = 1.0,
+        amin: float = 1e-10,
+        top_db: float | None = None,
+        dtype: str = 'float32',
+    ) -> None:
+        super().__init__()
+
+        self._melspectrogram = MelSpectrogram(
+            sr=sr,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            window=window,
+            power=power,
+            center=center,
+            pad_mode=pad_mode,
+            n_mels=n_mels,
+            f_min=f_min,
+            f_max=f_max,
+            htk=htk,
+            norm=norm,
+            dtype=dtype,
+        )
+
+        self.ref_value = ref_value
+        self.amin = amin
+        self.top_db = top_db
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x (Tensor): Tensor of waveforms with shape `(N, T)`
+
+        Returns:
+            Tensor: Log mel spectrograms with shape `(N, n_mels, num_frames)`.
+        """
+        mel_feature = self._melspectrogram(x)
+        log_mel_feature = power_to_db(
+            mel_feature,
+            ref_value=self.ref_value,
+            amin=self.amin,
+            top_db=self.top_db,
+        )
+        return log_mel_feature
+
+
+class MFCC(nn.Layer):
+    """Compute mel frequency cepstral coefficients(MFCCs) feature of given waveforms.
+
+    Args:
+        sr (int, optional): Sample rate. Defaults to 22050.
+        n_mfcc (int, optional): [description]. Defaults to 40.
+        n_fft (int, optional): The number of frequency components of the discrete Fourier transform. Defaults to 512.
+        hop_length (Optional[int], optional): The hop length of the short time FFT. If `None`, it is set to `win_length//4`. Defaults to None.
+        win_length (Optional[int], optional): The window length of the short time FFT. If `None`, it is set to same as `n_fft`. Defaults to None.
+        window (str, optional): The window function applied to the signal before the Fourier transform. Supported window functions: 'hamming', 'hann', 'gaussian', 'exponential', 'triang', 'bohman', 'blackman', 'cosine', 'tukey', 'taylor', 'bartlett', 'kaiser', 'nuttall'. Defaults to 'hann'.
+        power (float, optional): Exponent for the magnitude spectrogram. Defaults to 2.0.
+        center (bool, optional): Whether to pad `x` to make that the :math:`t \times hop\\_length` at the center of `t`-th frame. Defaults to True.
+        pad_mode (str, optional): Choose padding pattern when `center` is `True`. Defaults to 'reflect'.
+        n_mels (int, optional): Number of mel bins. Defaults to 64.
+        f_min (float, optional): Minimum frequency in Hz. Defaults to 50.0.
+        f_max (Optional[float], optional): Maximum frequency in Hz. Defaults to None.
+        htk (bool, optional): Use HTK formula in computing fbank matrix. Defaults to False.
+        norm (Union[str, float], optional): Type of normalization in computing fbank matrix. Slaney-style is used by default. You can specify norm=1.0/2.0 to use customized p-norm normalization. Defaults to 'slaney'.
+        ref_value (float, optional): The reference value. If smaller than 1.0, the db level of the signal will be pulled up accordingly. Otherwise, the db level is pushed down. Defaults to 1.0.
+        amin (float, optional): The minimum value of input magnitude. Defaults to 1e-10.
+        top_db (Optional[float], optional): The maximum db value of spectrogram. Defaults to None.
+        dtype (str, optional): Data type of input and window. Defaults to 'float32'.
+
+    Returns:
+        :ref:`api_paddle_nn_Layer`. An instance of MFCC.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.audio.features import MFCC
+
+            >>> sample_rate = 16000
+            >>> wav_duration = 0.5
+            >>> num_channels = 1
+            >>> num_frames = sample_rate * wav_duration
+            >>> wav_data = paddle.linspace(-1.0, 1.0, int(num_frames)) * 0.1
+            >>> waveform = wav_data.tile([num_channels, 1])
+
+            >>> feature_extractor = MFCC(sr=sample_rate, n_fft=512, window='hann')
+            >>> feats = feature_extractor(waveform)
+    """
+
+    dct_matrix: Tensor
+
+    def __init__(
+        self,
+        sr: int = 22050,
+        n_mfcc: int = 40,
+        n_fft: int = 512,
+        hop_length: int | None = None,
+        win_length: int | None = None,
+        window: _WindowLiteral = 'hann',
+        power: float = 2.0,
+        center: bool = True,
+        pad_mode: Literal['reflect'] = 'reflect',
+        n_mels: int = 64,
+        f_min: float = 50.0,
+        f_max: float | None = None,
+        htk: bool = False,
+        norm: Literal['slaney'] | float = 'slaney',
+        ref_value: float = 1.0,
+        amin: float = 1e-10,
+        top_db: float | None = None,
+        dtype: str = 'float32',
+    ) -> None:
+        super().__init__()
+        assert n_mfcc <= n_mels, (
+            f'n_mfcc cannot be larger than n_mels: {n_mfcc} vs {n_mels}'
+        )
+        self._log_melspectrogram = LogMelSpectrogram(
+            sr=sr,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            window=window,
+            power=power,
+            center=center,
+            pad_mode=pad_mode,
+            n_mels=n_mels,
+            f_min=f_min,
+            f_max=f_max,
+            htk=htk,
+            norm=norm,
+            ref_value=ref_value,
+            amin=amin,
+            top_db=top_db,
+            dtype=dtype,
+        )
+        self.dct_matrix = create_dct(n_mfcc=n_mfcc, n_mels=n_mels, dtype=dtype)
+        self.register_buffer('dct_matrix', self.dct_matrix)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x (Tensor): Tensor of waveforms with shape `(N, T)`
+
+        Returns:
+            Tensor: Mel frequency cepstral coefficients with shape `(N, n_mfcc, num_frames)`.
+        """
+        log_mel_feature = self._log_melspectrogram(x)
+        mfcc = paddle.matmul(
+            log_mel_feature.transpose((0, 2, 1)), self.dct_matrix
+        ).transpose((0, 2, 1))  # (B, n_mels, L)
+        return mfcc

python/paddle/audio/functional/__init__.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2026 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 .functional import (
+    compute_fbank_matrix,
+    create_dct,
+    fft_frequencies,
+    hz_to_mel,
+    mel_frequencies,
+    mel_to_hz,
+    power_to_db,
+    resample,
+)
+from .window import get_window
+
+__all__ = [
+    'compute_fbank_matrix',
+    'create_dct',
+    'fft_frequencies',
+    'hz_to_mel',
+    'mel_frequencies',
+    'mel_to_hz',
+    'power_to_db',
+    'get_window',
+    'resample',
+]

python/paddle/audio/functional/functional.py

@@ -0,0 +1,611 @@
+# Copyright (c) 2022 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.
+# Modified from librosa(https://github.com/librosa/librosa)
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, Literal, TypeVar
+
+import paddle
+from paddle import Tensor
+from paddle.base.framework import Variable
+from paddle.pir import Value
+
+if TYPE_CHECKING:
+    _TensorOrFloat = TypeVar("_TensorOrFloat", Tensor, float)
+
+
+def hz_to_mel(freq: _TensorOrFloat, htk: bool = False) -> _TensorOrFloat:
+    """Convert Hz to Mels.
+
+    Args:
+        freq (Union[Tensor, float]): The input tensor with arbitrary shape.
+        htk (bool, optional): Use htk scaling. Defaults to False.
+
+    Returns:
+        Union[Tensor, float]: Frequency in mels.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> val = 3.0
+            >>> htk_flag = True
+            >>> mel_paddle_tensor = paddle.audio.functional.hz_to_mel(paddle.to_tensor(val), htk_flag)
+    """
+
+    if htk:
+        if isinstance(freq, (Tensor, Variable, Value)):
+            return 2595.0 * paddle.log10(1.0 + freq / 700.0)
+        else:
+            return 2595.0 * math.log10(1.0 + freq / 700.0)
+
+    # Fill in the linear part
+    f_min = 0.0
+    f_sp = 200.0 / 3
+
+    mels = (freq - f_min) / f_sp
+
+    # Fill in the log-scale part
+
+    min_log_hz = 1000.0  # beginning of log region (Hz)
+    min_log_mel = (min_log_hz - f_min) / f_sp  # same (Mels)
+    logstep = math.log(6.4) / 27.0  # step size for log region
+
+    if isinstance(freq, (Tensor, Variable, Value)):
+        target = (
+            min_log_mel + paddle.log(freq / min_log_hz + 1e-10) / logstep
+        )  # prevent nan with 1e-10
+        mask = (freq > min_log_hz).astype(freq.dtype)
+        mels = target * mask + mels * (
+            1 - mask
+        )  # will replace by masked_fill OP in future
+    else:
+        if freq >= min_log_hz:
+            mels = min_log_mel + math.log(freq / min_log_hz + 1e-10) / logstep
+
+    return mels
+
+
+def mel_to_hz(mel: _TensorOrFloat, htk: bool = False) -> _TensorOrFloat:
+    """Convert mel bin numbers to frequencies.
+
+    Args:
+        mel (Union[float, Tensor]): The mel frequency represented as a tensor with arbitrary shape.
+        htk (bool, optional): Use htk scaling. Defaults to False.
+
+    Returns:
+        Union[float, Tensor]: Frequencies in Hz.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> val = 3.0
+            >>> htk_flag = True
+            >>> mel_paddle_tensor = paddle.audio.functional.mel_to_hz(paddle.to_tensor(val), htk_flag)
+    """
+    if htk:
+        return 700.0 * (10.0 ** (mel / 2595.0) - 1.0)
+
+    f_min = 0.0
+    f_sp = 200.0 / 3
+    freqs = f_min + f_sp * mel
+    # And now the nonlinear scale
+    min_log_hz = 1000.0  # beginning of log region (Hz)
+    min_log_mel = (min_log_hz - f_min) / f_sp  # same (Mels)
+    logstep = math.log(6.4) / 27.0  # step size for log region
+    if isinstance(mel, (Tensor, Variable, Value)):
+        target = min_log_hz * paddle.exp(logstep * (mel - min_log_mel))
+        mask = (mel > min_log_mel).astype(mel.dtype)
+        freqs = target * mask + freqs * (
+            1 - mask
+        )  # will replace by masked_fill OP in future
+    else:
+        if mel >= min_log_mel:
+            freqs = min_log_hz * math.exp(logstep * (mel - min_log_mel))
+    return freqs
+
+
+def mel_frequencies(
+    n_mels: int = 64,
+    f_min: float = 0.0,
+    f_max: float = 11025.0,
+    htk: bool = False,
+    dtype: str = 'float32',
+) -> Tensor:
+    """Compute mel frequencies.
+
+    Args:
+        n_mels (int, optional): Number of mel bins. Defaults to 64.
+        f_min (float, optional): Minimum frequency in Hz. Defaults to 0.0.
+        fmax (float, optional): Maximum frequency in Hz. Defaults to 11025.0.
+        htk (bool, optional): Use htk scaling. Defaults to False.
+        dtype (str, optional): The data type of the return frequencies. Defaults to 'float32'.
+
+    Returns:
+        Tensor: Tensor of n_mels frequencies in Hz with shape `(n_mels,)`.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> n_mels = 64
+            >>> f_min = 0.5
+            >>> f_max = 10000
+            >>> htk_flag = True
+
+            >>> paddle_mel_freq = paddle.audio.functional.mel_frequencies(n_mels, f_min, f_max, htk_flag, 'float64')
+    """
+    # 'Center freqs' of mel bands - uniformly spaced between limits
+    min_mel = hz_to_mel(f_min, htk=htk)
+    max_mel = hz_to_mel(f_max, htk=htk)
+    mels = paddle.linspace(min_mel, max_mel, n_mels, dtype=dtype)
+    freqs = mel_to_hz(mels, htk=htk)
+    return freqs
+
+
+def fft_frequencies(sr: int, n_fft: int, dtype: str = 'float32') -> Tensor:
+    """Compute fourier frequencies.
+
+    Args:
+        sr (int): Sample rate.
+        n_fft (int): Number of fft bins.
+        dtype (str, optional): The data type of the return frequencies. Defaults to 'float32'.
+
+    Returns:
+        Tensor: FFT frequencies in Hz with shape `(n_fft//2 + 1,)`.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> sr = 16000
+            >>> n_fft = 128
+            >>> fft_freq = paddle.audio.functional.fft_frequencies(sr, n_fft)
+    """
+    return paddle.linspace(0, float(sr) / 2, int(1 + n_fft // 2), dtype=dtype)
+
+
+def compute_fbank_matrix(
+    sr: int,
+    n_fft: int,
+    n_mels: int = 64,
+    f_min: float = 0.0,
+    f_max: float | None = None,
+    htk: bool = False,
+    norm: Literal['slaney'] | float = 'slaney',
+    dtype: str = 'float32',
+) -> Tensor:
+    """Compute fbank matrix.
+
+    Args:
+        sr (int): Sample rate.
+        n_fft (int): Number of fft bins.
+        n_mels (int, optional): Number of mel bins. Defaults to 64.
+        f_min (float, optional): Minimum frequency in Hz. Defaults to 0.0.
+        f_max (Optional[float], optional): Maximum frequency in Hz. Defaults to None.
+        htk (bool, optional): Use htk scaling. Defaults to False.
+        norm (Union[str, float], optional): Type of normalization. Defaults to 'slaney'.
+        dtype (str, optional): The data type of the return matrix. Defaults to 'float32'.
+
+    Returns:
+        Tensor: Mel transform matrix with shape `(n_mels, n_fft//2 + 1)`.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> sr = 23
+            >>> n_fft = 51
+            >>> fbank = paddle.audio.functional.compute_fbank_matrix(sr, n_fft)
+    """
+
+    if f_max is None:
+        f_max = float(sr) / 2
+
+    # Initialize the weights
+    weights = paddle.zeros((n_mels, int(1 + n_fft // 2)), dtype=dtype)
+
+    # Center freqs of each FFT bin
+    fftfreqs = fft_frequencies(sr=sr, n_fft=n_fft, dtype=dtype)
+
+    # 'Center freqs' of mel bands - uniformly spaced between limits
+    mel_f = mel_frequencies(
+        n_mels + 2, f_min=f_min, f_max=f_max, htk=htk, dtype=dtype
+    )
+
+    fdiff = mel_f[1:] - mel_f[:-1]  # np.diff(mel_f)
+    ramps = mel_f.unsqueeze(1) - fftfreqs.unsqueeze(0)
+    # ramps = np.subtract.outer(mel_f, fftfreqs)
+
+    for i in range(n_mels):
+        # lower and upper slopes for all bins
+        lower = -ramps[i] / fdiff[i]
+        upper = ramps[i + 2] / fdiff[i + 1]
+
+        # .. then intersect them with each other and zero
+        weights[i] = paddle.maximum(
+            paddle.zeros_like(lower), paddle.minimum(lower, upper)
+        )
+
+    # Slaney-style mel is scaled to be approx constant energy per channel
+    if norm == 'slaney':
+        enorm = 2.0 / (mel_f[2 : n_mels + 2] - mel_f[:n_mels])
+        weights *= enorm.unsqueeze(1)
+    elif isinstance(norm, (int, float)):
+        weights = paddle.nn.functional.normalize(weights, p=norm, axis=-1)
+
+    return weights
+
+
+def power_to_db(
+    spect: Tensor,
+    ref_value: float = 1.0,
+    amin: float = 1e-10,
+    top_db: float | None = 80.0,
+) -> Tensor:
+    """Convert a power spectrogram (amplitude squared) to decibel (dB) units. The function computes the scaling `10 * log10(x / ref)` in a numerically stable way.
+
+    Args:
+        spect (Tensor): STFT power spectrogram.
+        ref_value (float, optional): The reference value. If smaller than 1.0, the db level of the signal will be pulled up accordingly. Otherwise, the db level is pushed down. Defaults to 1.0.
+        amin (float, optional): Minimum threshold. Defaults to 1e-10.
+        top_db (Optional[float], optional): Threshold the output at `top_db` below the peak. Defaults to None.
+
+    Returns:
+        Tensor: Power spectrogram in db scale.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> val = 3.0
+            >>> decibel_paddle = paddle.audio.functional.power_to_db(paddle.to_tensor(val))
+    """
+    if amin <= 0:
+        raise Exception("amin must be strictly positive")
+
+    if ref_value <= 0:
+        raise Exception("ref_value must be strictly positive")
+
+    ones = paddle.ones_like(spect)
+    log_spec = 10.0 * paddle.log10(paddle.maximum(ones * amin, spect))
+    log_spec -= 10.0 * math.log10(max(ref_value, amin))
+
+    if top_db is not None:
+        if top_db < 0:
+            raise Exception("top_db must be non-negative")
+        log_spec = paddle.maximum(log_spec, ones * (log_spec.max() - top_db))
+
+    return log_spec
+
+
+def create_dct(
+    n_mfcc: int,
+    n_mels: int,
+    norm: Literal['ortho'] | None = 'ortho',
+    dtype: str = 'float32',
+) -> Tensor:
+    """Create a discrete cosine transform(DCT) matrix.
+
+    Args:
+        n_mfcc (int): Number of mel frequency cepstral coefficients.
+        n_mels (int): Number of mel filterbanks.
+        norm (Optional[str], optional): Normalization type. Defaults to 'ortho'.
+        dtype (str, optional): The data type of the return matrix. Defaults to 'float32'.
+
+    Returns:
+        Tensor: The DCT matrix with shape `(n_mels, n_mfcc)`.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> n_mfcc = 23
+            >>> n_mels = 257
+            >>> dct = paddle.audio.functional.create_dct(n_mfcc, n_mels)
+    """
+    n = paddle.arange(n_mels, dtype=dtype)
+    k = paddle.arange(n_mfcc, dtype=dtype).unsqueeze(1)
+    dct = paddle.cos(
+        math.pi / float(n_mels) * (n + 0.5) * k
+    )  # size (n_mfcc, n_mels)
+    if norm is None:
+        dct *= 2.0
+    else:
+        assert norm == "ortho"
+        dct[0] *= 1.0 / math.sqrt(2.0)
+        dct *= math.sqrt(2.0 / float(n_mels))
+    return dct.T
+
+
+def _get_sinc_resample_kernel(
+    orig_freq: int,
+    new_freq: int,
+    gcd: int,
+    lowpass_filter_width: int = 6,
+    rolloff: float = 0.99,
+    resampling_method: Literal[
+        "sinc_interp_hann", "sinc_interp_kaiser"
+    ] = "sinc_interp_hann",
+    beta: float | None = None,
+    dtype: paddle.dtype | None = None,
+):
+    """
+    Generate the sinc interpolation kernel for resampling.
+
+    This internal function computes the resampling kernel based on the sinc
+    interpolation formula with windowing. The kernel is used by
+    _apply_sinc_resample_kernel to perform the actual resampling.
+
+    Args:
+        orig_freq (int): Original sampling frequency.
+        new_freq (int): Target sampling frequency.
+        gcd (int): Greatest common divisor of orig_freq and new_freq.
+        lowpass_filter_width (int, optional): Controls the sharpness of the filter,
+            larger value means sharper but less efficient. Default: 6.
+        rolloff (float, optional): Roll-off frequency as a fraction of the Nyquist.
+            Lower values reduce anti-aliasing but also attenuate high frequencies.
+            Default: 0.99.
+        resampling_method (str, optional): Window method for filter design.
+            Options: ["sinc_interp_hann", "sinc_interp_kaiser"]. Default: "sinc_interp_hann".
+        beta (float, optional): Shape parameter for Kaiser window. Required only
+            when resampling_method="sinc_interp_kaiser". Default: None.
+        dtype (paddle.dtype, optional): Data type for kernel computation.
+            If None, uses float64 for computation and converts to float32 for output.
+            Default: None.
+
+    Returns:
+        tuple: (kernel, width)
+            - kernel (Tensor): Resampling kernel of shape (1, 1, kernel_width)
+            - width (int): Half-width of the filter in terms of input samples
+
+    Raises:
+        Exception: If frequencies are not integers.
+        ValueError: If resampling_method is invalid or lowpass_filter_width <= 0.
+    """
+    if not (int(orig_freq) == orig_freq and int(new_freq) == new_freq):
+        raise ValueError(
+            "Frequencies must be of integer type to ensure quality resampling computation. "
+            "To work around this, manually convert both frequencies to integer values "
+            "that maintain their resampling rate ratio before passing them into the function. "
+            "Example: To downsample a 44100 hz waveform by a factor of 8, use "
+            "`orig_freq=8` and `new_freq=1` instead of `orig_freq=44100` and `new_freq=5512.5`. "
+        )
+
+    if resampling_method not in ["sinc_interp_hann", "sinc_interp_kaiser"]:
+        raise ValueError(f"Invalid resampling method: {resampling_method}")
+
+    orig_freq = int(orig_freq) // gcd
+    new_freq = int(new_freq) // gcd
+
+    if lowpass_filter_width <= 0:
+        raise ValueError("Low pass filter width should be positive.")
+    base_freq = min(orig_freq, new_freq)
+
+    # Perform antialiasing filtering by removing the highest frequencies.
+    base_freq *= rolloff
+
+    # Calculate filter width based on lowpass_filter_width and frequency ratio
+    width = math.ceil(lowpass_filter_width * orig_freq / base_freq)
+    idx_dtype = dtype if dtype is not None else paddle.float64
+
+    idx = (
+        paddle.arange(-width, width + orig_freq, dtype=idx_dtype)[None, None]
+        / orig_freq
+    )
+
+    t = (
+        paddle.arange(0, -new_freq, -1, dtype=dtype)[:, None, None] / new_freq
+        + idx
+    )
+    t *= base_freq
+    t = t.clip_(-lowpass_filter_width, lowpass_filter_width)
+
+    # we do not use built-in paddle windows here as we need to evaluate the window
+    # at specific positions, not over a regular grid.
+    if resampling_method == "sinc_interp_hann":
+        window = paddle.cos(t * math.pi / lowpass_filter_width / 2) ** 2
+    else:
+        # sinc_interp_kaiser
+        if beta is None:
+            beta = 14.769656459379492
+        beta_tensor = paddle.to_tensor(float(beta))
+        window = paddle.i0(
+            beta_tensor * paddle.sqrt(1 - (t / lowpass_filter_width) ** 2),
+        ) / paddle.i0(beta_tensor)
+
+    t *= math.pi
+
+    scale = base_freq / orig_freq
+    kernels = paddle.where(
+        t == 0, paddle.to_tensor(1.0).cast(t.dtype), t.sin() / t
+    )
+    kernels *= window * scale
+
+    if dtype is None:  # pragma: no cover
+        kernels = kernels.cast(paddle.float32)
+
+    return kernels, width
+
+
+def _apply_sinc_resample_kernel(
+    waveform: Tensor,
+    orig_freq: int,
+    new_freq: int,
+    gcd: int,
+    kernel: Tensor,
+    width: int,
+):
+    """
+    Apply sinc interpolation resampling using precomputed kernel.
+
+    This internal function performs the actual resampling operation using the
+    kernel generated by _get_sinc_resample_kernel. It handles batch processing
+    and ensures correct output length.
+
+    Args:
+        waveform (Tensor): Input waveform of shape (..., time). Must be floating point.
+        orig_freq (int): Original sampling frequency.
+        new_freq (int): Target sampling frequency.
+        gcd (int): Greatest common divisor of orig_freq and new_freq.
+        kernel (Tensor): Resampling kernel from _get_sinc_resample_kernel.
+        width (int): Half-width of the filter from _get_sinc_resample_kernel.
+
+    Returns:
+        Tensor: Resampled waveform of shape (..., new_time).
+
+    """
+
+    orig_freq = int(orig_freq) // gcd
+    new_freq = int(new_freq) // gcd
+
+    # pack batch
+    shape = waveform.shape
+    waveform = waveform.reshape([-1, shape[-1]])
+
+    num_wavs, length = waveform.shape
+    waveform = paddle.nn.functional.pad(waveform, (width, width + orig_freq))
+    resampled = paddle.nn.functional.conv1d(
+        waveform[:, None], kernel, stride=orig_freq
+    )
+    resampled = resampled.transpose([0, 2, 1]).reshape((num_wavs, -1))
+    target_length = paddle.ceil(
+        paddle.to_tensor(new_freq * length / orig_freq)
+    ).astype(paddle.int64)
+    resampled = resampled[..., :target_length]
+
+    # unpack batch
+    resampled = resampled.reshape(shape[:-1] + resampled.shape[-1:])
+    return resampled
+
+
+def resample(
+    waveform: Tensor,
+    orig_freq: int,
+    new_freq: int,
+    lowpass_filter_width: int = 6,
+    rolloff: float = 0.99,
+    resampling_method: Literal[
+        "sinc_interp_hann", "sinc_interp_kaiser"
+    ] = "sinc_interp_hann",
+    beta: float | None = None,
+) -> Tensor:
+    """
+    Resample the waveform from orig_freq to new_freq using bandlimited interpolation.
+
+    This function implements resampling through sinc interpolation with windowing.
+    It first computes a resampling kernel based on the specified parameters, then
+    applies it to the input waveform using convolution. The algorithm handles both
+    upsampling and downsampling while minimizing aliasing artifacts.
+
+    Args:
+        waveform (Tensor): The input signal of dimension (..., time). Must be
+            floating point type (float32 or float64).
+        orig_freq (int): The original frequency of the signal. Must be positive.
+        new_freq (int): The desired target frequency. Must be positive.
+        lowpass_filter_width (int, optional): Controls the sharpness of the filter.
+            Larger values give sharper filtering but are less efficient.
+            Default: 6.
+        rolloff (float, optional): The roll-off frequency of the filter as a fraction
+            of the Nyquist frequency. Lower values reduce anti-aliasing but also
+            attenuate some high frequencies. Default: 0.99.
+        resampling_method (str, optional): The windowing method to use for filter
+            design. Options: "sinc_interp_hann" (Hann window) or "sinc_interp_kaiser"
+            (Kaiser window). Default: "sinc_interp_hann".
+        beta (float, optional): Shape parameter for the Kaiser window. Required only
+            when resampling_method="sinc_interp_kaiser". If not provided for Kaiser,
+            a default value of 14.769656459379492 is used. Default: None.
+
+    Returns:
+        Tensor: The waveform resampled to new_freq, with dimension (..., new_time).
+
+    Raises:
+        ValueError: If orig_freq or new_freq are not positive.
+        Exception: If frequencies are not integers (see note below).
+        TypeError: If waveform is not floating point.
+
+    Note:
+        - orig_freq and new_freq must be integers. For non-integer frequencies,
+          convert them to integers while maintaining the ratio.
+        - For repeated resampling with same parameters, use
+          :class:`paddle.audio.transforms.Resample` for better efficiency.
+        - Uses windowed sinc interpolation for high-quality audio resampling.
+        - This function does not support ONNX export now.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.audio.functional import resample
+
+            >>> # Create a sample waveform (1 channel, 1000 samples at 16000 Hz)
+            >>> waveform = paddle.randn([1, 1000])
+
+            >>> # Downsample from 16000 Hz to 8000 Hz
+            >>> resampled = resample(waveform, 16000, 8000)
+            >>> print(resampled.shape)
+            paddle.Size([1, 500])
+
+            >>> # Upsample from 16000 Hz to 48000 Hz with custom filter width
+            >>> resampled = resample(waveform, 16000, 48000, lowpass_filter_width=12)
+            >>> print(resampled.shape)
+            paddle.Size([1, 3000])
+
+            >>> # Use Kaiser window resampling
+            >>> resampled = resample(waveform, 16000, 8000, resampling_method="sinc_interp_kaiser", beta=12.0)
+            >>> print(resampled.shape)
+            paddle.Size([1, 500])
+
+            >>> # Batch processing: multiple waveforms
+            >>> batch_waveforms = paddle.randn([4, 1, 1000])  # [batch, channels, time]
+            >>> resampled_batch = resample(batch_waveforms, 16000, 8000)
+            >>> print(resampled_batch.shape)
+            paddle.Size([4, 1, 500])
+    """
+    if orig_freq <= 0.0 or new_freq <= 0.0:
+        raise ValueError(
+            "Original frequency and desired frequency should be positive integers"
+        )
+    if not waveform.is_floating_point():
+        raise TypeError(
+            f"Expected floating point type for waveform tensor, but received {waveform.dtype}."
+        )
+
+    if orig_freq == new_freq:
+        return waveform
+
+    gcd = math.gcd(int(orig_freq), int(new_freq))
+
+    kernel, width = _get_sinc_resample_kernel(
+        orig_freq,
+        new_freq,
+        gcd,
+        lowpass_filter_width,
+        rolloff,
+        resampling_method,
+        beta,
+        waveform.dtype,
+    )
+    resampled = _apply_sinc_resample_kernel(
+        waveform, orig_freq, new_freq, gcd, kernel, width
+    )
+    return resampled

python/paddle/audio/functional/window.py

@@ -0,0 +1,731 @@
+# Copyright (c) 2022 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
+from __future__ import annotations
+
+import math
+import warnings
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+import paddle
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle._typing import PlaceLike
+
+    from ..features.layers import _WindowLiteral
+
+from paddle.base.framework import (
+    _current_expected_place,
+    _get_paddle_place,
+    _to_pinned_place,
+    in_dynamic_or_pir_mode,
+)
+
+
+class WindowFunctionRegister:
+    def __init__(self):
+        self._functions_dict = {}
+
+    def register(self, func=None):
+        def add_subfunction(func):
+            name = func.__name__
+            self._functions_dict[name] = func
+            return func
+
+        return add_subfunction
+
+    def get(self, name):
+        return self._functions_dict[name]
+
+
+window_function_register = WindowFunctionRegister()
+
+
+@window_function_register.register()
+def _cat(x: list[Tensor], data_type: str) -> Tensor:
+    l = []
+    for t in x:
+        if np.isscalar(t) and not isinstance(t, str):
+            l.append(paddle.to_tensor([t], data_type))
+        else:
+            l.append(paddle.to_tensor(t, data_type))
+    return paddle.concat(l)
+
+
+@window_function_register.register()
+def _bartlett(M: int, sym: bool = True, dtype: str = 'float64') -> Tensor:
+    """
+    Computes the Bartlett window.
+    This function is consistent with scipy.signal.windows.bartlett().
+    """
+    if _len_guards(M):
+        return paddle.ones((M,), dtype=dtype)
+    M, needs_trunc = _extend(M, sym)
+
+    n = paddle.arange(0, M, dtype=dtype)
+    M = paddle.to_tensor(M, dtype=dtype)
+    w = paddle.where(
+        paddle.less_equal(n, (M - 1) / 2.0),
+        2.0 * n / (M - 1),
+        2.0 - 2.0 * n / (M - 1),
+    )
+
+    return _truncate(w, needs_trunc)
+
+
+@window_function_register.register()
+def _kaiser(
+    M: int, beta: float, sym: bool = True, dtype: str = 'float64'
+) -> Tensor:
+    """Compute the Kaiser window.
+    This function is consistent with scipy.signal.windows.kaiser().
+    """
+    if _len_guards(M):
+        return paddle.ones((M,), dtype=dtype)
+    M, needs_trunc = _extend(M, sym)
+
+    beta = paddle.to_tensor(beta, dtype=dtype)
+
+    n = paddle.arange(0, M, dtype=dtype)
+    M = paddle.to_tensor(M, dtype=dtype)
+    alpha = (M - 1) / 2.0
+    w = paddle.i0(
+        beta * paddle.sqrt(1 - ((n - alpha) / alpha) ** 2.0)
+    ) / paddle.i0(beta)
+
+    return _truncate(w, needs_trunc)
+
+
+@window_function_register.register()
+def _nuttall(M: int, sym: bool = True, dtype: str = 'float64') -> Tensor:
+    """Nuttall window.
+    This function is consistent with scipy.signal.windows.nuttall().
+    """
+    a = paddle.to_tensor(
+        [0.3635819, 0.4891775, 0.1365995, 0.0106411], dtype=dtype
+    )
+    return _general_cosine(M, a=a, sym=sym, dtype=dtype)
+
+
+@window_function_register.register()
+def _acosh(x: Tensor | float) -> Tensor:
+    if isinstance(x, float):
+        return math.log(x + math.sqrt(x**2 - 1))
+    return paddle.log(x + paddle.sqrt(paddle.square(x) - 1))
+
+
+@window_function_register.register()
+def _extend(M: int, sym: bool) -> bool:
+    """Extend window by 1 sample if needed for DFT-even symmetry."""
+    if not sym:
+        return M + 1, True
+    else:
+        return M, False
+
+
+@window_function_register.register()
+def _len_guards(M: int) -> bool:
+    """Handle small or incorrect window lengths."""
+    if int(M) != M or M < 0:
+        raise ValueError('Window length M must be a non-negative integer')
+
+    return M <= 1
+
+
+@window_function_register.register()
+def _truncate(w: Tensor, needed: bool) -> Tensor:
+    """Truncate window by 1 sample if needed for DFT-even symmetry."""
+    if needed:
+        return w[:-1]
+    else:
+        return w
+
+
+@window_function_register.register()
+def _general_gaussian(
+    M: int, p, sig, sym: bool = True, dtype: str = 'float64'
+) -> Tensor:
+    """Compute a window with a generalized Gaussian shape.
+    This function is consistent with scipy.signal.windows.general_gaussian().
+    """
+    if _len_guards(M):
+        return paddle.ones((M,), dtype=dtype)
+    M, needs_trunc = _extend(M, sym)
+
+    n = paddle.arange(0, M, dtype=dtype) - (M - 1.0) / 2.0
+    w = paddle.exp(-0.5 * paddle.abs(n / sig) ** (2 * p))
+
+    return _truncate(w, needs_trunc)
+
+
+@window_function_register.register()
+def _general_cosine(
+    M: int, a: list[float], sym: bool = True, dtype: str = 'float64'
+) -> Tensor:
+    """Compute a generic weighted sum of cosine terms window.
+    This function is consistent with scipy.signal.windows.general_cosine().
+    """
+    if _len_guards(M):
+        return paddle.ones((M,), dtype=dtype)
+    M, needs_trunc = _extend(M, sym)
+    fac = paddle.linspace(-math.pi, math.pi, M, dtype=dtype)
+    w = paddle.zeros((M,), dtype=dtype)
+    for k in range(len(a)):
+        w += a[k] * paddle.cos(k * fac)
+    return _truncate(w, needs_trunc)
+
+
+@window_function_register.register()
+def _general_hamming(
+    M: int, alpha: float, sym: bool = True, dtype: str = 'float64'
+) -> Tensor:
+    """Compute a generalized Hamming window.
+    This function is consistent with scipy.signal.windows.general_hamming()
+    """
+    return _general_cosine(M, [alpha, 1.0 - alpha], sym, dtype=dtype)
+
+
+@window_function_register.register()
+def _taylor(
+    M: int, nbar=4, sll=30, norm=True, sym: bool = True, dtype: str = 'float64'
+) -> Tensor:
+    """Compute a Taylor window.
+    The Taylor window taper function approximates the Dolph-Chebyshev window's
+    constant sidelobe level for a parameterized number of near-in sidelobes.
+    """
+    if _len_guards(M):
+        return paddle.ones((M,), dtype=dtype)
+    M, needs_trunc = _extend(M, sym)
+    # Original text uses a negative sidelobe level parameter and then negates
+    # it in the calculation of B. To keep consistent with other methods we
+    # assume the sidelobe level parameter to be positive.
+    B = 10 ** (sll / 20)
+    A = _acosh(B) / math.pi
+    s2 = nbar**2 / (A**2 + (nbar - 0.5) ** 2)
+    ma = paddle.arange(1, nbar, dtype=dtype)
+
+    Fm = paddle.empty((nbar - 1,), dtype=dtype)
+    signs = paddle.empty_like(ma)
+    signs[::2] = 1
+    signs[1::2] = -1
+    m2 = ma * ma
+    for mi in range(len(ma)):
+        number = signs[mi] * paddle.prod(
+            1 - m2[mi] / s2 / (A**2 + (ma - 0.5) ** 2)
+        )
+        if mi == 0:
+            denom = 2 * paddle.prod(1 - m2[mi] / m2[mi + 1 :])
+        elif mi == len(ma) - 1:
+            denom = 2 * paddle.prod(1 - m2[mi] / m2[:mi])
+        else:
+            denom = (
+                2
+                * paddle.prod(1 - m2[mi] / m2[:mi])
+                * paddle.prod(1 - m2[mi] / m2[mi + 1 :])
+            )
+
+        Fm[mi] = number / denom
+
+    def W(n):
+        return 1 + 2 * paddle.matmul(
+            Fm.unsqueeze(0),
+            paddle.cos(2 * math.pi * ma.unsqueeze(1) * (n - M / 2.0 + 0.5) / M),
+        )
+
+    w = W(paddle.arange(0, M, dtype=dtype))
+
+    # normalize (Note that this is not described in the original text [1])
+    if norm:
+        scale = 1.0 / W((M - 1) / 2)
+        w *= scale
+    w = w.squeeze()
+    return _truncate(w, needs_trunc)
+
+
+@window_function_register.register()
+def _hamming(M: int, sym: bool = True, dtype: str = 'float64') -> Tensor:
+    """Compute a Hamming window.
+    The Hamming window is a taper formed by using a raised cosine with
+    non-zero endpoints, optimized to minimize the nearest side lobe.
+    """
+    return _general_hamming(M, 0.54, sym, dtype=dtype)
+
+
+@window_function_register.register()
+def _hann(M: int, sym: bool = True, dtype: str = 'float64') -> Tensor:
+    """Compute a Hann window.
+    The Hann window is a taper formed by using a raised cosine or sine-squared
+    with ends that touch zero.
+    """
+    return _general_hamming(M, 0.5, sym, dtype=dtype)
+
+
+@window_function_register.register()
+def _tukey(
+    M: int, alpha=0.5, sym: bool = True, dtype: str = 'float64'
+) -> Tensor:
+    """Compute a Tukey window.
+    The Tukey window is also known as a tapered cosine window.
+    """
+    if _len_guards(M):
+        return paddle.ones((M,), dtype=dtype)
+
+    if alpha <= 0:
+        return paddle.ones((M,), dtype=dtype)
+    elif alpha >= 1.0:
+        return _hann(M, sym=sym)
+
+    M, needs_trunc = _extend(M, sym)
+
+    n = paddle.arange(0, M, dtype=dtype)
+    width = int(alpha * (M - 1) / 2.0)
+    n1 = n[0 : width + 1]
+    n2 = n[width + 1 : M - width - 1]
+    n3 = n[M - width - 1 :]
+
+    w1 = 0.5 * (1 + paddle.cos(math.pi * (-1 + 2.0 * n1 / alpha / (M - 1))))
+    w2 = paddle.ones(n2.shape, dtype=dtype)
+    w3 = 0.5 * (
+        1
+        + paddle.cos(math.pi * (-2.0 / alpha + 1 + 2.0 * n3 / alpha / (M - 1)))
+    )
+    w = paddle.concat([w1, w2, w3])
+
+    return _truncate(w, needs_trunc)
+
+
+@window_function_register.register()
+def _gaussian(
+    M: int, std: float, sym: bool = True, dtype: str = 'float64'
+) -> Tensor:
+    """Compute a Gaussian window.
+    The Gaussian widows has a Gaussian shape defined by the standard deviation(std).
+    """
+    if _len_guards(M):
+        return paddle.ones((M,), dtype=dtype)
+    M, needs_trunc = _extend(M, sym)
+
+    n = paddle.arange(0, M, dtype=dtype) - (M - 1.0) / 2.0
+    sig2 = 2 * std * std
+    w = paddle.exp(-(n**2) / sig2)
+
+    return _truncate(w, needs_trunc)
+
+
+@window_function_register.register()
+def _exponential(
+    M: int, center=None, tau=1.0, sym: bool = True, dtype: str = 'float64'
+) -> Tensor:
+    """Compute an exponential (or Poisson) window."""
+    if sym and center is not None:
+        raise ValueError("If sym==True, center must be None.")
+    if _len_guards(M):
+        return paddle.ones((M,), dtype=dtype)
+    M, needs_trunc = _extend(M, sym)
+
+    if center is None:
+        center = (M - 1) / 2
+
+    n = paddle.arange(0, M, dtype=dtype)
+    w = paddle.exp(-paddle.abs(n - center) / tau)
+
+    return _truncate(w, needs_trunc)
+
+
+@window_function_register.register()
+def _triang(M: int, sym: bool = True, dtype: str = 'float64') -> Tensor:
+    """Compute a triangular window."""
+    if _len_guards(M):
+        return paddle.ones((M,), dtype=dtype)
+    M, needs_trunc = _extend(M, sym)
+
+    n = paddle.arange(1, (M + 1) // 2 + 1, dtype=dtype)
+    if M % 2 == 0:
+        w = (2 * n - 1.0) / M
+        w = paddle.concat([w, w[::-1]])
+    else:
+        w = 2 * n / (M + 1.0)
+        w = paddle.concat([w, w[-2::-1]])
+
+    return _truncate(w, needs_trunc)
+
+
+@window_function_register.register()
+def _bohman(M: int, sym: bool = True, dtype: str = 'float64') -> Tensor:
+    """Compute a Bohman window.
+    The Bohman window is the autocorrelation of a cosine window.
+    """
+    if _len_guards(M):
+        return paddle.ones((M,), dtype=dtype)
+    M, needs_trunc = _extend(M, sym)
+
+    fac = paddle.abs(paddle.linspace(-1, 1, M, dtype=dtype)[1:-1])
+    w = (1 - fac) * paddle.cos(math.pi * fac) + 1.0 / math.pi * paddle.sin(
+        math.pi * fac
+    )
+    w = _cat([0, w, 0], dtype)
+
+    return _truncate(w, needs_trunc)
+
+
+@window_function_register.register()
+def _blackman(M: int, sym: bool = True, dtype: str = 'float64') -> Tensor:
+    """Compute a Blackman window.
+    The Blackman window is a taper formed by using the first three terms of
+    a summation of cosines. It was designed to have close to the minimal
+    leakage possible.  It is close to optimal, only slightly worse than a
+    Kaiser window.
+    """
+    return _general_cosine(M, [0.42, 0.50, 0.08], sym, dtype=dtype)
+
+
+@window_function_register.register()
+def _cosine(M: int, sym: bool = True, dtype: str = 'float64') -> Tensor:
+    """Compute a window with a simple cosine shape."""
+    if _len_guards(M):
+        return paddle.ones((M,), dtype=dtype)
+    M, needs_trunc = _extend(M, sym)
+    w = paddle.sin(math.pi / M * (paddle.arange(0, M, dtype=dtype) + 0.5))
+
+    return _truncate(w, needs_trunc)
+
+
+def get_window(
+    window: _WindowLiteral | tuple[_WindowLiteral, float],
+    win_length: int,
+    fftbins: bool = True,
+    dtype: str | None = 'float64',
+) -> Tensor:
+    """Return a window of a given length and type.
+
+    Args:
+        window (Union[str, Tuple[str, float]]): The window function applied to the signal before the Fourier transform. Supported window functions: 'hamming', 'hann', 'gaussian', 'general_gaussian', 'exponential', 'triang', 'bohman', 'blackman', 'cosine', 'tukey', 'taylor', 'bartlett', 'kaiser', 'nuttall'.
+        win_length (int): Number of samples.
+        fftbins (bool, optional): If True, create a "periodic" window. Otherwise, create a "symmetric" window, for use in filter design. Defaults to True.
+        dtype (str, optional): The data type of the return window. Defaults to 'float64'.
+
+    Returns:
+        Tensor: The window represented as a tensor.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> n_fft = 512
+            >>> cosine_window = paddle.audio.functional.get_window('cosine', n_fft)
+
+            >>> std = 7
+            >>> gaussian_window = paddle.audio.functional.get_window(('gaussian', std), n_fft)
+    """
+    if dtype is None:
+        dtype = 'float32'
+    sym = not fftbins
+    args = ()
+    if isinstance(window, tuple):
+        winstr = window[0]
+        if len(window) > 1:
+            args = window[1:]
+    elif isinstance(window, str):
+        if window in ['gaussian', 'exponential', 'kaiser']:
+            raise ValueError(
+                "The '" + window + "' window needs one or "
+                "more parameters -- pass a tuple."
+            )
+        else:
+            winstr = window
+    else:
+        raise ValueError(f"{type(window)} as window type is not supported.")
+
+    try:
+        winfunc = window_function_register.get('_' + winstr)
+    except KeyError as e:
+        raise ValueError("Unknown window type.") from e
+    params = (win_length, *args)
+    kwargs = {'sym': sym}
+    return winfunc(*params, dtype=dtype, **kwargs)
+
+
+def _apply_window_postprocess(
+    w: Tensor,
+    *,
+    layout: str | None = None,
+    device: PlaceLike | None = None,
+    pin_memory: bool = False,
+    requires_grad: bool = False,
+) -> Tensor:
+    if layout not in (None, 'strided'):
+        raise RuntimeError(
+            "Window functions only support layout='strided' or None"
+        )
+    if layout is not None:
+        warnings.warn("layout only supports 'strided' in Paddle; ignored")
+
+    if in_dynamic_or_pir_mode():
+        device = (
+            _get_paddle_place(device)
+            if device is not None
+            else _current_expected_place()
+        )
+        if pin_memory and paddle.in_dynamic_mode() and device is not None:
+            device = _to_pinned_place(device)
+    w = w.to(device=device)
+    if pin_memory and paddle.in_dynamic_mode():
+        w = w.pin_memory()
+    if requires_grad is True:
+        w.stop_gradient = False
+    return w
+
+
+def hamming_window(
+    window_length: int,
+    periodic: bool = True,
+    alpha: float = 0.54,
+    beta: float = 0.46,
+    *,
+    dtype: str = 'float32',
+    layout: str | None = None,
+    device: PlaceLike | None = None,
+    pin_memory: bool = False,
+    requires_grad: bool = False,
+):
+    """
+    Compute a generalized Hamming window.
+
+    Args:
+        window_length (int): The size of the returned window. Must be positive.
+        periodic (bool, optional): If True, returns a window for use as a periodic function; if False, returns a symmetric window. Defaults to True.
+        alpha (float, optional): The coefficient α in the equation above. Defaults to 0.54.
+        beta (float, optional): The coefficient β in the equation above. Defaults to 0.46.
+        dtype (str, optional): The data type of the returned tensor. Defaults to 'float32'.
+        layout (str, optional): Only included for API consistency with PyTorch; ignored in Paddle. Defaults to None.
+        device(PlaceLike|None, optional): The desired device of returned tensor.
+            if None, uses the current device for the default tensor type (see paddle.device.set_device()).
+            device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. Default: None.
+        pin_memory(bool, optional): If set, return tensor would be allocated in the pinned memory. Works only for CPU tensors. Default: False
+        requires_grad(bool, optional):  If autograd should record operations on the returned tensor. Default: False.
+
+    Returns:
+        Tensor: A 1-D tensor of shape `(window_length,)` containing the Hamming window.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> win = paddle.hamming_window(400, requires_grad=True)
+            >>> win = paddle.hamming_window(256, alpha=0.5, beta=0.5)
+    """
+    w0 = get_window('hamming', window_length, fftbins=periodic, dtype=dtype)
+    alpha0, beta0 = 0.54, 0.46
+    B = beta / beta0
+    A = alpha - B * alpha0
+    w = A + B * w0
+    return _apply_window_postprocess(
+        w,
+        layout=layout,
+        device=device,
+        pin_memory=pin_memory,
+        requires_grad=requires_grad,
+    )
+
+
+def hann_window(
+    window_length: int,
+    periodic: bool = True,
+    *,
+    dtype: str = 'float32',
+    layout: str | None = None,
+    device: PlaceLike | None = None,
+    pin_memory: bool = False,
+    requires_grad: bool = False,
+):
+    """
+    Compute a Hann window.
+
+    Args:
+        window_length (int): The size of the returned window. Must be positive.
+        periodic (bool, optional): If True, returns a window for use as a periodic function; if False, returns a symmetric window. Defaults to True.
+        dtype (str, optional): The data type of the returned tensor. Defaults to 'float32'.
+        layout (str, optional): Only included for API consistency with PyTorch; ignored in Paddle. Defaults to None.
+        device(PlaceLike|None, optional): The desired device of returned tensor.
+            if None, uses the current device for the default tensor type (see paddle.device.set_device()).
+            device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. Default: None.
+        pin_memory(bool, optional): If set, return tensor would be allocated in the pinned memory. Works only for CPU tensors. Default: False
+        requires_grad(bool, optional):  If autograd should record operations on the returned tensor. Default: False.
+
+    Returns:
+        Tensor: A 1-D tensor of shape `(window_length,)` containing the Hann window.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> win = paddle.hann_window(512)
+            >>> win = paddle.hann_window(512, requires_grad=True)
+    """
+    w = get_window('hann', window_length, fftbins=periodic, dtype=dtype)
+    return _apply_window_postprocess(
+        w,
+        layout=layout,
+        device=device,
+        pin_memory=pin_memory,
+        requires_grad=requires_grad,
+    )
+
+
+def kaiser_window(
+    window_length: int,
+    periodic: bool = True,
+    beta: float = 12.0,
+    *,
+    dtype: str | None = 'float32',
+    layout: str | None = None,
+    device: PlaceLike | None = None,
+    pin_memory: bool = False,
+    requires_grad: bool = False,
+    out: Tensor | None = None,
+):
+    """
+    Compute a Kaiser window.
+
+    Args:
+        window_length (int): The size of the returned window. Must be positive.
+        periodic (bool, optional): If True, returns a window for use as a periodic function; if False, returns a symmetric window. Defaults to True.
+        beta (float, optional): Shape parameter for the window. Defaults to 12.0.
+        dtype (str, optional): The data type of the returned tensor. Defaults to 'float32'.
+        layout (str, optional): Only included for API consistency with PyTorch; ignored in Paddle. Defaults to None.
+        device(PlaceLike|None, optional): The desired device of returned tensor.
+            if None, uses the current device for the default tensor type (see paddle.device.set_device()).
+            device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. Default: None.
+        pin_memory(bool, optional): If set, return tensor would be allocated in the pinned memory. Works only for CPU tensors. Default: False
+        requires_grad(bool, optional):  If autograd should record operations on the returned tensor. Default: False.
+        out(Tensor|None, optional): The output tensor. Default: None.
+
+    Returns:
+        Tensor: A 1-D tensor of shape `(window_length,)` containing the Kaiser window.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> win = paddle.kaiser_window(400, beta=8.6)
+            >>> win = paddle.kaiser_window(400, requires_grad=True)
+    """
+    w = get_window(
+        ('kaiser', beta), window_length, fftbins=periodic, dtype=dtype
+    )
+    w = _apply_window_postprocess(
+        w,
+        layout=layout,
+        device=device,
+        pin_memory=pin_memory,
+        requires_grad=requires_grad,
+    )
+    return paddle.assign(w, out) if out is not None else w
+
+
+def blackman_window(
+    window_length: int,
+    periodic: bool = True,
+    *,
+    dtype: str = 'float32',
+    layout: str | None = None,
+    device: PlaceLike | None = None,
+    pin_memory: bool = False,
+    requires_grad: bool = False,
+):
+    """
+    Compute a Blackman window.
+
+    Args:
+        window_length (int): The size of the returned window. Must be positive.
+        periodic (bool, optional): If True, returns a window for use as a periodic function; if False, returns a symmetric window. Defaults to True.
+        dtype (str, optional): The data type of the returned tensor. Defaults to 'float32'.
+        layout (str, optional): Only included for API consistency with PyTorch; ignored in Paddle. Defaults to None.
+        device(PlaceLike|None, optional): The desired device of returned tensor.
+            if None, uses the current device for the default tensor type (see paddle.device.set_device()).
+            device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. Default: None.
+        pin_memory(bool, optional): If set, return tensor would be allocated in the pinned memory. Works only for CPU tensors. Default: False
+        requires_grad(bool, optional):  If autograd should record operations on the returned tensor. Default: False.
+
+    Returns:
+        Tensor: A 1-D tensor of shape `(window_length,)` containing the Blackman window.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> win = paddle.blackman_window(256)
+            >>> win = paddle.blackman_window(256, requires_grad=True)
+    """
+    w = get_window('blackman', window_length, fftbins=periodic, dtype=dtype)
+    return _apply_window_postprocess(
+        w,
+        layout=layout,
+        device=device,
+        pin_memory=pin_memory,
+        requires_grad=requires_grad,
+    )
+
+
+def bartlett_window(
+    window_length: int,
+    periodic: bool = True,
+    *,
+    dtype: str = 'float32',
+    layout: str | None = None,
+    device: PlaceLike | None = None,
+    pin_memory: bool = False,
+    requires_grad: bool = False,
+):
+    """
+    Compute a Bartlett window.
+
+    Args:
+        window_length (int): The size of the returned window. Must be positive.
+        periodic (bool, optional): If True, returns a window for use as a periodic function; if False, returns a symmetric window. Defaults to True.
+        dtype (str, optional): The data type of the returned tensor. Defaults to 'float32'.
+        layout (str, optional): Only included for API consistency with PyTorch; ignored in Paddle. Defaults to None.
+        device(PlaceLike|None, optional): The desired device of returned tensor.
+            if None, uses the current device for the default tensor type (see paddle.device.set_device()).
+            device will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types. Default: None.
+        pin_memory(bool, optional): If set, return tensor would be allocated in the pinned memory. Works only for CPU tensors. Default: False
+        requires_grad(bool, optional):  If autograd should record operations on the returned tensor. Default: False.
+
+    Returns:
+        Tensor: A 1-D tensor of shape `(window_length,)` containing the Bartlett window.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> n_fft = 512
+            >>> win = paddle.bartlett_window(n_fft)
+
+            >>> win = paddle.bartlett_window(n_fft, requires_grad=True)
+    """
+    w = get_window('bartlett', window_length, fftbins=periodic, dtype=dtype)
+    return _apply_window_postprocess(
+        w,
+        layout=layout,
+        device=device,
+        pin_memory=pin_memory,
+        requires_grad=requires_grad,
+    )

python/paddle/autograd/__init__.py

@@ -0,0 +1,44 @@
+# Copyright (c) 2021 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 ..base.dygraph.base import (  # noqa: F401
+    enable_grad,
+    grad,
+    inference_mode,
+    is_grad_enabled,
+    no_grad_ as no_grad,
+    set_grad_enabled,
+)
+from . import (  # noqa: F401
+    backward_mode,
+    function,
+    grad_mode,
+    ir_backward,
+)
+from .autograd import hessian, jacobian
+from .backward_mode import backward
+from .py_layer import PyLayer, PyLayerContext
+from .saved_tensors_hooks import saved_tensors_hooks
+
+Function = PyLayer
+
+__all__ = [
+    'jacobian',
+    'hessian',
+    'backward',
+    'PyLayer',
+    'Function',
+    'PyLayerContext',
+    'saved_tensors_hooks',
+]

python/paddle/autograd/autograd.py

@@ -0,0 +1,775 @@
+# Copyright (c) 2023 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
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, overload
+
+import paddle
+from paddle.base import framework
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+def as_tensors(xs):
+    if isinstance(xs, framework.Variable):
+        return xs
+    elif isinstance(xs, Sequence):
+        return tuple(xs)
+    else:
+        return xs
+
+
+class Jacobian:
+    r"""Computes the Jacobian matrix of given xs and ys.
+
+    Once the Jacobian ``J`` is constructed, you can use a multidimensional index
+    to retrieve the submatrix of ``J``, as same as slicing a Tensor. The
+    submatrix is lazily evaluated along row axis, and will be cached once
+    evaluated.
+
+    you can retrieve the submatrix by
+    following methods:
+
+        * J[:], retrieving the full matrix.
+        * J[:, :, j], retrieving the partial derivatives w.r.t. the j'th input
+          variable.
+        * J[:, i, :], retrieving the partial derivatives w.r.t. the i'th output
+          variable.
+        * J[:, i, j], retrieving the partial derivatives w.r.t. the i'th output
+          variable and the j'th input variable.
+
+    Notes:
+
+        Ellipsis index is not supported currently.
+
+    Args:
+
+        ys (Tensor|Tuple[Tensor, ...]): The output derived from xs .
+        xs (Tensor|Tuple[Tensor, ...]): The input tensor(s) .
+        is_batched (bool): If true, the first axis is batch axis. Defaults to
+            False.
+
+    Returns:
+
+        Jacobian (Object): A python object retains the Jacobian matrix.
+
+    """
+
+    def __init__(
+        self,
+        ys: Tensor,
+        xs: Tensor,
+        is_batched: bool = False,
+    ) -> None:
+        if not is_batched:
+            if not 0 <= len(xs.shape) <= 1:
+                raise ValueError(
+                    f"xs.ndim should be 0 or 1 when is_batched=False"
+                    f" but got {len(xs.shape)}"
+                )
+            if not 0 <= len(ys.shape) <= 1:
+                raise ValueError(
+                    f"ys.ndim should be 0 or 1 when is_batched=False"
+                    f" but got {len(ys.shape)}"
+                )
+            self._jacobian = _JacobianNoBatch(ys, xs)
+        else:
+            if not 1 <= len(ys.shape) <= 2:
+                raise ValueError(
+                    f"ys.ndim should be 1 or 2 when is_batched=True"
+                    f" but got {len(ys.shape)}"
+                )
+            if not 1 <= len(xs.shape) <= 2:
+                raise ValueError(
+                    f"xs.ndim should be 1 or 2 when is_batched=True"
+                    f" but got {len(xs.shape)}"
+                )
+            self._jacobian = _JacobianBatchFirst(ys, xs)
+
+    @property
+    def shape(self) -> list[int]:
+        """The shape of flattened Jacobian matrix."""
+        return self._jacobian.shape
+
+    def __getitem__(self, indexes):
+        return self._jacobian[indexes]
+
+    def __getattr__(self, __name: str):  # noqa: PYI063
+        if __name == "shape":
+            return getattr(self._jacobian, __name)
+        if __name == "_evaluate_all":
+            return getattr(self._jacobian, __name)
+        return getattr(self._jacobian._evaluate_all(), __name)
+
+    def __add__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs + rhs
+
+    def __sub__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs - rhs
+
+    def __mul__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs * rhs
+
+    def __div__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs / rhs
+
+    def __truediv__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs / rhs
+
+    def __pow__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs**rhs
+
+    def __mod__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs % rhs
+
+    def __floordiv__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs // rhs
+
+    def __matmul__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs @ rhs
+
+    def __eq__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs == rhs
+
+    def __ne__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs != rhs
+
+    def __lt__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs < rhs
+
+    def __le__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs <= rhs
+
+    def __gt__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs > rhs
+
+    def __ge__(self, other):
+        lhs = self._evaluate_all()
+        rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
+        return lhs >= rhs
+
+
+class Hessian(Jacobian):
+    pass
+
+
+class _Jacobian:
+    """The base class for computing Jacobian matrix.
+
+    ``_Jacobian`` implements the core logic of multidimensional index and lazy
+    evaluation for Jacobian matrix, subclass only need to overwrite following
+    methods:
+
+        * ``_lazy_axis()``,  return the axis along which will be lazy
+            evaluating.
+        * ``_flatten(xs)``, flattens the inputs ``xs``.
+        * ``_evaluate(index)``, evaluates one slice along ``_lazy_axis`` .
+
+    Notes:
+
+        Because currently PaddlePaddle only support reverse differentiation by
+        ``paddle.grad``, so lazy evaluation is only supported along the row of
+        Jacobian matrix, which means that slicing along row will get better
+        performance.
+
+    """
+
+    def __init__(self, ys, xs):
+        self.original_xs_shape = xs.shape
+        self.original_ys_shape = ys.shape
+        self._xs = xs
+        self._ys = ys
+        if len(self._ys.shape) == 0 and not self.is_batched:
+            self._ys = self._ys.reshape(
+                [
+                    -1,
+                ]
+            )
+        if len(self._ys.shape) == 1 and self.is_batched:
+            self._ys = self._ys.reshape([-1, 1])
+
+        self._flatten_xs = self._flatten(as_tensors(self._xs))
+        self._flatten_ys = self._flatten(as_tensors(self._ys))
+        self._cache = {}
+
+    @property
+    def _lazy_axis(self):
+        """ "The axis of lazily evaluated."""
+        raise NotImplementedError
+
+    def _lazy_indexes(self, indexes):
+        idx = indexes[self._lazy_axis]
+        return (
+            (idx,)
+            if isinstance(idx, int)
+            else tuple(range(idx.start, idx.stop, idx.step))
+        )
+
+    def _flatten(self, xs):
+        raise NotImplementedError
+
+    def _shifted_indexes(self, indexes, lazy_axis_size=0):
+        idx = indexes[self._lazy_axis]
+        shifted_lazy_axis_idx = (
+            0 if isinstance(idx, int) else slice(0, lazy_axis_size, 1)
+        )
+        return (
+            *indexes[: self._lazy_axis],
+            shifted_lazy_axis_idx,
+            *indexes[self._lazy_axis + 1 :],
+        )
+
+    def __getitem__(self, indexes):
+        if self.is_batched is False:
+            if len(self.shape) == 0:
+                # xs and ys are both 0-D tensor
+                raise IndexError("0-D tensor can not be indexed.")
+            elif len(self.shape) == 1:
+                # either ys or xs is 0-D tensor
+                indexes = (
+                    (0, indexes)
+                    if len(self.original_ys_shape) == 0
+                    else (indexes, 0)
+                )
+        else:
+            if len(self.shape) == 1:
+                # xs and ys are both 1-D tensor
+                indexes = (indexes, 0, 0)
+            elif len(self.shape) == 2:
+                # either xs or ys is 1-D tensor
+                if isinstance(indexes, slice):
+                    indexes = (indexes, slice(None, None, None))
+                else:
+                    indexes = (
+                        (indexes[0], 0, indexes[1])
+                        if len(self.original_ys_shape) == 1
+                        else (indexes[0], indexes[1], 0)
+                    )
+
+        indexes = _multi_index(indexes, self.inner_shape)
+
+        if isinstance(indexes[self._lazy_axis], int):
+            other_indexes = (
+                indexes[: self._lazy_axis] + indexes[self._lazy_axis + 1 :]
+            )
+            return self._cached_evaluate(indexes[self._lazy_axis])[
+                other_indexes
+            ]
+        lazy_indexes = self._lazy_indexes(indexes)
+        # Using concat and reshape to replace stack operator temporarily, as
+        # it is not a primitive operator.
+        shape = list(self.inner_shape)
+        shape[self._lazy_axis] = len(lazy_indexes)
+        part_jac = paddle.concat(
+            [self._cached_evaluate(i) for i in lazy_indexes],
+            axis=self._lazy_axis,
+        ).reshape(shape)
+        result = part_jac[self._shifted_indexes(indexes, len(lazy_indexes))]
+
+        # squeeze redundant 1 in shape
+        if len(result.shape) > len(self.shape):
+            for _ in range(len(result.shape) - len(self.shape)):
+                result = result.squeeze(-1)
+
+        return result
+
+    def _cached_evaluate(self, k):
+        if k is None:
+            return self._cached_evaluate(0).reshape([])
+        v = self._cache.get(k)
+        if v is None:
+            v = self._evaluate(k)
+            self._cache[k] = v
+        return v
+
+    def _evaluate(self, index):
+        """Evaluate one slice at along lazy axis."""
+        raise NotImplementedError
+
+    def _evaluate_all(self):
+        if len(self.shape) == 0:
+            return self._cached_evaluate(None)
+        else:
+            return self[:]
+
+
+class _JacobianNoBatch(_Jacobian):
+    """Compute Jacobian matrix without batch dimension.
+    Suppose the mapping is :math:`f: R^M \to R^N`, the output shape is
+    ``(N, M)`` .
+    """
+
+    def __init__(self, ys, xs):
+        self.is_batched = False
+        super().__init__(ys, xs)
+        # inner_shape is for convenient, it will regard 0-D tensor as 1-D tensor
+        self.inner_shape = [
+            *(self._flatten_ys.shape[0:1]),
+            *(self._flatten_xs.shape[0:1]),
+        ]
+        self.shape = [
+            *(self.original_ys_shape[0:1]),
+            *(self.original_xs_shape[0:1]),
+        ]
+
+    @property
+    def _lazy_axis(self):
+        return 0
+
+    def _flatten(self, xs):
+        if not isinstance(xs, Sequence):
+            return xs.reshape((-1,))
+        return paddle.concat(tuple(x.reshape((-1,)) for x in xs))
+
+    def _evaluate(self, row_index):
+        return self._flatten(
+            _grad_for_jacobian(
+                self._flatten_ys[row_index],
+                self._xs,
+            )
+        )
+
+
+class _JacobianBatchFirst(_Jacobian):
+    """Compute Jacobian matrix with batch at first axis.
+    Suppose the mapping is :math:`f: R^{B,M} \to R^{B,N}`, the output shape is
+    ``(B, N, M)`` .
+    """
+
+    def __init__(self, ys, xs):
+        self.is_batched = True
+        super().__init__(ys, xs)
+        # inner_shape is for convenient, it will regard 0-D tensor as 1-D tensor
+        self.inner_shape = [
+            *(self._flatten_xs.shape[0:1]),
+            *(self._flatten_ys.shape[1:2]),
+            *(self._flatten_xs.shape[1:2]),
+        ]
+        self.shape = [
+            *(self._flatten_xs.shape[0:1]),
+            *(self.original_ys_shape[1:2]),
+            *(self.original_xs_shape[1:2]),
+        ]
+
+    @property
+    def _lazy_axis(self):
+        return 1
+
+    def _flatten(self, xs):
+        if not isinstance(xs, Sequence):
+            return xs.reshape((xs.shape[0], -1))
+        return paddle.concat(
+            tuple(x.reshape((x.shape[0], -1)) for x in as_tensors(xs)), 1
+        )
+
+    def _evaluate(self, row_index):
+        return self._flatten(
+            _grad_for_jacobian(self._flatten_ys[:, row_index], self._xs)
+        )
+
+
+def _multi_index(indexes, shape):
+    """A tool for parsing N-dimensional index into a standard format.
+
+    Currently supporting following input format:
+        * ([positive|negative|slice], ...), the right-most elements can be
+            omitted.
+
+    The standard format after converted is slice tuple which contains N elements:
+        * ([positive|slice], ..., [positive|slice])
+
+    Notes:
+        Ellipsis indexes such as ``(..., i), (i, ...)`` is not supported.
+
+    Args:
+        indexes (tuple): The input indexes.
+        shape (tuple): The input shape.
+
+    Returns:
+        tuple: The standard format index as the above description.
+    """
+    indexes = indexes if isinstance(indexes, Sequence) else (indexes,)
+    if any(isinstance(i, type(Ellipsis)) for i in indexes):
+        raise IndexError('Ellipsis index currently is not supported.')
+    # Fill the right-most elements.
+    indexes = indexes + (slice(0, None, None),) * (len(shape) - len(indexes))
+    # Convert to positive index.
+    positive_indexes = []
+    for i, index in enumerate(indexes):
+        if isinstance(index, slice):
+            index = slice(
+                index.start or 0, index.stop or shape[i], index.step or 1
+            )
+            positive_indexes.append(
+                slice(
+                    index.start + shape[i] if index.start < 0 else index.start,
+                    index.stop + shape[i] if index.stop < 0 else index.stop,
+                    # Negative step means index backward, no need to convert to
+                    # positive integer.
+                    index.step,
+                )
+            )
+        elif isinstance(index, int):
+            positive_indexes.append(index + shape[i] if index < 0 else index)
+        else:
+            raise TypeError(f'Not supported index type {index}.')
+    return tuple(positive_indexes)
+
+
+@overload
+def jacobian(
+    ys: Tensor,
+    xs: Tensor,
+    batch_axis: int | None = ...,
+) -> Jacobian: ...
+
+
+@overload
+def jacobian(
+    ys: Sequence[Tensor],
+    xs: Sequence[Tensor],
+    batch_axis: int | None = ...,
+) -> tuple[tuple[Jacobian, ...], ...]: ...
+
+
+@overload
+def jacobian(
+    ys: Tensor,
+    xs: Sequence[Tensor],
+    batch_axis: int | None = ...,
+) -> tuple[Jacobian, ...]: ...
+
+
+@overload
+def jacobian(
+    ys: Sequence[Tensor],
+    xs: Tensor,
+    batch_axis: int | None = ...,
+) -> tuple[Jacobian, ...]: ...
+
+
+def jacobian(
+    ys,
+    xs,
+    batch_axis=None,
+):
+    r"""
+    Computes the Jacobian of the dependent variable ``ys`` versus the independent
+    variable ``xs``.
+
+    Where ``ys`` represents the output of ``xs`` after a certain operation, ``ys`` and
+    ``xs`` can be Tensor or tuple of Tensors, ``batch_axis`` indicates the position of
+    the batch dimension of the parameter data.
+
+    When the input is a tuple Tensors, the returned result is a ``Jacobian`` object with
+    the same number of nesting levels as ``xs``, and each Jacobian has the same shape as
+    The ``xs`` tuples are identical in one-to-one correspondence.
+
+    - When ``batch_axis=None``, only 0-dimensional Tensor or 1-dimensional Tensor is
+      supported, assuming the shape of ``xs`` is ``[N, ]``, the shape of ``ys`` is
+      ``[M, ]``, then the output Jacobian matrix shape is ``[M, N]``.
+
+    - When ``batch_axis=0``, only 1-dimensional Tensor or 2-dimensional Tensor is
+      supported, assuming the shape of ``xs`` is ``[B, N]``, The shape of ``ys`` is
+      ``[B, M]``, then the output Jacobian matrix shape is ``[B, M, N]``.
+
+    After the ``Jacobian`` object is created, the actual calculation process does not
+    occur, but the lazy evaluation method is used for calculation. It can be
+    multi-dimensional indexed to obtain the entire Jacobian matrix or sub-matrix, and
+    the actual calculation will be performed at this time the value is calculated and
+    the result is returned. At the same time, in the actual evaluation process, the
+    calculated sub-matrix will be cached to avoid duplicate calculations in the
+    subsequent indexing process.
+
+    For example, assuming ``Jacobian`` instance ``J`` has shape ``[B, M, N]``, assuming
+    ``M > 4`` , then ``J[:, 1:4:1, :]`` means to get the values from row ``1`` to row
+    ``3`` of ``J``. In actual calculation, only the rows ``1`` to ``3`` are evaluated,
+    and the calculation results of ``1`` to ``3`` will be cached at the granularity of
+    the row, and will be used next time. When obtaining one or more rows of results
+    above, the already calculated parts will not be recalculated.
+
+    Args:
+
+        ys (Union[paddle.Tensor, Tuple[paddle.Tensor, ...]]): Output or tuple of outputs derived from xs.
+        xs (Union[paddle.Tensor, Tuple[paddle.Tensor, ...]]): Input or tuple of inputs.
+        batch_axis (Optional[int], optional): Index of batch axis. Defaults to None.
+
+    Returns:
+
+        Union[Tuple[Tuple[Jacobian, ...], ...], Tuple[Jacobian, ...], Jacobian]: Jacobian(s) of ys derived from xs.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> x1 = paddle.randn([3])
+            >>> x2 = paddle.randn([3])
+            >>> x1.stop_gradient = False
+            >>> x2.stop_gradient = False
+
+            >>> y = x1 + x2
+
+            >>> J = paddle.autograd.jacobian(y, (x1, x2))
+            >>> J_y_x1 = J[0][:]  # evaluate result of dy/dx1
+            >>> J_y_x2 = J[1][:]  # evaluate result of dy/dx2
+
+            >>> print(J_y_x1.shape)
+            paddle.Size([3, 3])
+            >>> print(J_y_x2.shape)
+            paddle.Size([3, 3])
+    """
+
+    if batch_axis is not None and batch_axis != 0:
+        raise ValueError(
+            f"batch_axis should be None or 0, but got {batch_axis}."
+        )
+
+    # TODO(HydrogenSulfate): support batch_axis > 0
+    is_batched = batch_axis is not None
+    if isinstance(ys, Sequence) and isinstance(xs, Sequence):
+        _jacobian = tuple(
+            tuple(Jacobian(_ys, _xs, is_batched) for _xs in xs) for _ys in ys
+        )
+    elif isinstance(ys, Sequence) and not isinstance(xs, Sequence):
+        _jacobian = tuple(Jacobian(_ys, xs, is_batched) for _ys in ys)
+    elif not isinstance(ys, Sequence) and isinstance(xs, Sequence):
+        _jacobian = tuple(Jacobian(ys, _xs, is_batched) for _xs in xs)
+    else:
+        _jacobian = Jacobian(ys, xs, is_batched)
+
+    return _jacobian
+
+
+@overload
+def hessian(
+    ys: Tensor,
+    xs: Tensor,
+    batch_axis: int | None = ...,
+) -> Hessian: ...
+
+
+@overload
+def hessian(
+    ys: Tensor,
+    xs: Sequence[Tensor],
+    batch_axis: int | None = ...,
+) -> tuple[tuple[Hessian, ...], ...]: ...
+
+
+def hessian(
+    ys,
+    xs,
+    batch_axis=None,
+):
+    r"""
+    Computes the Jacobian of the dependent variable ``ys`` versus the independent
+    variable ``xs``.
+
+    Among them, ``ys`` means the output of ``xs`` after a certain operation, ``ys`` can
+    only be a single Tensor, ``xs`` can be a Tensor or a Tensor tuple, and
+    ``batch_axis`` means The position of the batch dimension of the parameter data.
+
+    When the input ``xs`` is a Tensor tuple, the returned result is a ``Hessian`` tuple,
+    assuming that the internal shape of the ``xs`` tuple is composed of ``([M1, ], [M2, ])``, the shape of the returned
+    result consists of ``(([M1, M1], [M1, M2]), ([M2, M1], [M2, M2]))``
+
+    - When ``batch_axis=None``, only 0-dimensional Tensor or 1-dimensional Tensor is
+      supported, assuming that the shape of ``xs`` is ``[N, ]``, and the shape of ``ys`` is ``[ ]`` (0-dimensional Tensor), the final output is a single Hessian matrix whose shape is ``[N, N]``.
+
+    - When ``batch_axis=0``, only 1-dimensional Tensor or 2-dimensional Tensor is
+      supported, assuming that the shape of ``xs`` is ``[B, N]``, and the shape of ``ys`` is ``[B, ]``, the final output Jacobian matrix shape is ``[B, N, N]``.
+
+    After the ``Hessian`` object is created, the complete calculation process does not
+    occur, but a partial lazy evaluation method is used for calculation. It can be
+    multi-dimensionally indexed to obtain the entire Hessian matrix or sub-matrix. At
+    this time, the actual Evaluates the computation and returns the result. At the same
+    time, in the actual evaluation process, the calculated sub-matrix will be cached to
+    avoid repeated calculations in the subsequent indexing process.
+
+    Args:
+
+        ys (paddle.Tensor): Output derived from xs which contain one element.
+        xs (Union[paddle.Tensor, Tuple[paddle.Tensor, ...]]): Input or tuple of inputs.
+        batch_axis (Optional[int], optional): Index of batch axis. Defaults to None.
+
+    Returns:
+
+        Union[Tuple[Tuple[Hessian, ...], ...], Tuple[Hessian, ...], Hessian]: Hessian(s) of ys derived from xs.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> x1 = paddle.randn([3])
+            >>> x2 = paddle.randn([4])
+            >>> x1.stop_gradient = False
+            >>> x2.stop_gradient = False
+
+            >>> y = x1.sum() + x2.sum()
+
+            >>> H = paddle.autograd.hessian(y, (x1, x2))
+            >>> H_y_x1_x1 = H[0][0][:]  # evaluate result of ddy/dx1x1
+            >>> H_y_x1_x2 = H[0][1][:]  # evaluate result of ddy/dx1x2
+            >>> H_y_x2_x1 = H[1][0][:]  # evaluate result of ddy/dx2x1
+            >>> H_y_x2_x2 = H[1][1][:]  # evaluate result of ddy/dx2x2
+
+            >>> print(H_y_x1_x1.shape)
+            paddle.Size([3, 3])
+            >>> print(H_y_x1_x2.shape)
+            paddle.Size([3, 4])
+            >>> print(H_y_x2_x1.shape)
+            paddle.Size([4, 3])
+            >>> print(H_y_x2_x2.shape)
+            paddle.Size([4, 4])
+    """
+
+    if batch_axis is None:
+        if ys.numel() > 1:
+            raise ValueError(
+                f"Only support ys.numel()({ys.numel()})==1 when batch_axis is None."
+            )
+        ys = ys.reshape(())
+    elif isinstance(batch_axis, int):
+        if ys[0].numel() > 1:
+            raise ValueError(
+                f"Only support ys[0].numel()({ys.numel()})==1 when batch_axis is int"
+            )
+        # TODO(HydrogenSulfate): support batch_axis > 0
+        if batch_axis != 0:
+            raise ValueError("Only support batch_axis=0 yet.")
+        ys = ys.reshape((-1,))
+    else:
+        raise ValueError(
+            f"batch_axis should be None or int, but got {type(batch_axis)}."
+        )
+
+    _jacobian = jacobian(ys, xs, batch_axis)
+    if not isinstance(xs, Sequence):
+        hessian = jacobian(_jacobian, xs, batch_axis)
+
+        # change classname to Hessian instead of Jacobian.
+        hessian.__class__ = Hessian
+    else:
+        hessian = tuple(jacobian(_j, xs, batch_axis) for _j in _jacobian)
+
+        # change classname to Hessian instead of Jacobian.
+        for i in range(len(hessian)):
+            for j in range(len(hessian[0])):
+                hessian[i][j].__class__ = Hessian
+
+    return hessian
+
+
+def _replace_none_with_zero_tensor(xs, refs):
+    if xs is None:
+        xs = paddle.zeros_like(refs)
+        xs.stop_gradient = refs.stop_gradient
+        return xs
+    elif isinstance(xs, Sequence):
+        return tuple(
+            _replace_none_with_zero_tensor(x, refs[i]) for i, x in enumerate(xs)
+        )
+    else:
+        return xs
+
+
+def _grad_for_jacobian(ys, xs, v=None):
+    """A gradient function that can be used in dynamic graph and static graph.
+
+    The ``grad`` combines ``paddle.grad`` used in dynamic graph and
+    ``paddle.static.gradients`` used in static graph, and do following changes:
+
+    * The ``allow_unused`` flag is removed and set defaults to true internally,
+        none in outputs will be replaced by zero tensor.
+    * The ``create_graph`` flag is removed and set defaults to true internally,
+        only makes sense in dynamic graph.
+    * When xs is a single Tensor, ``paddle.grad`` returns a list which only
+        contains one Tensor. It may confuse users, thus in this case we improve
+        to return a single Tensor in _grad_for_jacobian interface.
+
+    Args:
+        ys (Tensor|Sequence[Tensor]): The output tensor or tensor sequence of
+            the graph to compute gradients.
+        xs (Tensor|Sequence[Tensor]): The input tensor or tensor sequence of the graph to
+            compute gradients. The returned values of this API are the
+            gradients of inputs .
+        v (Tensor|Sequence[Tensor]|None,optional): The initial gradient values
+            of outputs . If grad_outputs is None, the initial gradient values of
+            outputs would be Tensors filled with 1; if grad_outputs is not None,
+            it must have the same length as outputs , and in this case, the
+            initial gradient value of the i-th outputs would be: (1) a Tensor
+            filled with 1 when the i-th element of grad_outputs is None;
+            (2) the i-th element of grad_outputs when the i-th element of
+            grad_outputs is a Tensor. Default None.
+
+    Returns:
+        Tensor|tuple[Tensor]: Tensor or a tuple of Tensors, whose length is the
+            same as the Tensor number inside inputs, and the i-th returned
+            Tensor is the sum of gradients of outputs with respect to the i-th
+            inputs.
+    """
+    if paddle.in_dynamic_mode():
+        # paddle.grad returns a list though the inputs is a single Tensor. The
+        # follow code snippet fixes the problem by return the first element of
+        # xs_grad when the xs is a single Tensor.
+        xs_grad = paddle.grad(ys, xs, v, create_graph=True, allow_unused=True)
+        if (
+            isinstance(xs, paddle.base.framework.Variable)
+            and isinstance(xs_grad, Sequence)
+            and len(xs_grad) > 0
+        ):
+            xs_grad = xs_grad[0]
+    else:
+        xs_grad = paddle.static.gradients(ys, xs, v)
+        if (
+            isinstance(xs, framework.Variable)
+            and isinstance(xs_grad, Sequence)
+            and len(xs_grad) > 0
+        ):
+            xs_grad = xs_grad[0]
+    return _replace_none_with_zero_tensor(xs_grad, xs)

python/paddle/autograd/backward_mode.py

@@ -0,0 +1,152 @@
+# Copyright (c) 2021 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
+
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle.base import core, framework
+from paddle.base.backward import gradients_with_optimizer  # noqa: F401
+from paddle.utils.download import check_and_create_dir
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from paddle import Tensor
+
+
+__all__ = []
+
+
+@framework.dygraph_only
+def backward(
+    tensors: Tensor | Sequence[Tensor],
+    grad_tensors: Tensor | Sequence[Tensor | None] | None = None,
+    retain_graph: bool = False,
+    create_graph: bool = False,
+    *,
+    dump_backward_graph_path: str | None = None,
+) -> None:
+    """
+    Compute the backward gradients of given tensors.
+
+    Args:
+        tensors(list of Tensors): the tensors which the gradient to be computed. The tensors can not contain the same tensor.
+
+        grad_tensors(list of Tensors of None, optional): the init gradients of the `tensors`` .If not None, it must have the same length with ``tensors`` ,
+            and if any of the elements is None, then the init gradient is the default value which is filled with 1.0.
+            If None, all the gradients of the ``tensors`` is the default value which is filled with 1.0.
+            Defaults to None.
+
+        retain_graph(bool, optional): If False, the graph used to compute grads will be freed. If you would
+            like to add more ops to the built graph after calling this method( :code:`backward` ), set the parameter
+            :code:`retain_graph` to True, then the grads will be retained. Thus, setting it to False is much more memory-efficient.
+            Defaults to False.
+        dump_backward_graph_path(str, optional): Specifies the directory path for storing the debug file.
+            If this parameter is specified, the backward-related graph (in dot format)
+            and the debugging call stack information will be generated in this directory.
+    Returns:
+        NoneType: None
+
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> x = paddle.to_tensor([[1, 2], [3, 4]], dtype='float32', stop_gradient=False)
+            >>> y = paddle.to_tensor([[3, 2], [3, 4]], dtype='float32')
+
+            >>> grad_tensor1 = paddle.to_tensor([[1, 2], [2, 3]], dtype='float32')
+            >>> grad_tensor2 = paddle.to_tensor([[1, 1], [1, 1]], dtype='float32')
+
+            >>> z1 = paddle.matmul(x, y)
+            >>> z2 = paddle.matmul(x, y)
+
+            >>> paddle.autograd.backward([z1, z2], [grad_tensor1, grad_tensor2], True)
+            >>> print(x.grad)
+            Tensor(shape=[2, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
+            [[12., 18.],
+             [17., 25.]])
+
+
+            >>> x.clear_grad()
+
+            >>> paddle.autograd.backward([z1, z2], [grad_tensor1, None], True)
+            >>> print(x.grad)
+            Tensor(shape=[2, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
+            [[12., 18.],
+             [17., 25.]])
+
+            >>> x.clear_grad()
+
+            >>> paddle.autograd.backward([z1, z2])
+            >>> print(x.grad)
+            Tensor(shape=[2, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
+            [[10., 14.],
+             [10., 14.]])
+
+
+    """
+
+    def check_tensors(
+        in_out_list: Sequence[Tensor] | Tensor, name: str
+    ) -> Sequence[Tensor]:
+        assert in_out_list is not None, f"{name} should not be None"
+
+        if isinstance(in_out_list, (list, tuple)):
+            assert len(in_out_list) > 0, f"{name} cannot be empty"
+            for each_var in in_out_list:
+                assert isinstance(each_var, paddle.Tensor), (
+                    f"Elements of {name} must be paddle.Tensor"
+                )
+            return in_out_list
+        else:
+            assert isinstance(in_out_list, paddle.Tensor), (
+                f"{name} must be Tensor or list of Tensor"
+            )
+            return [in_out_list]
+
+    tensors = check_tensors(tensors, "tensors")
+
+    assert len(tensors) == len(set(tensors)), (
+        "The argument 'tensors' of paddle.autograd.backward contains duplicate paddle.Tensor object."
+    )
+
+    if grad_tensors is not None:
+        if not isinstance(grad_tensors, (list, tuple)):
+            grad_tensors = [grad_tensors]
+
+        for each_tensor in grad_tensors:
+            if each_tensor is not None:
+                assert isinstance(each_tensor, paddle.Tensor), (
+                    "The argument 'grad_tensors' of paddle.autograd.backward is invalid, it can be 'None', 'paddle.Tensor' or 'list[None/paddle.Tensor]'."
+                )
+    else:
+        grad_tensors = []
+
+    if len(grad_tensors) > 0:
+        assert len(tensors) == len(grad_tensors), (
+            "The length of grad_tensors must be equal to tensors"
+        )
+
+    assert isinstance(retain_graph, bool), "retain_graph must be True or False"
+    check_and_create_dir(dump_backward_graph_path)
+    core.eager.run_backward(
+        tensors,
+        grad_tensors,
+        retain_graph,
+        create_graph,
+        dump_backward_graph_path,
+    )

python/paddle/autograd/backward_utils.py

@@ -0,0 +1,810 @@
+# Copyright (c) 2023 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,tes
+# 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 collections
+import logging
+import warnings
+from collections.abc import Sequence
+from functools import lru_cache
+from typing import Any
+
+from paddle import pir
+from paddle.base import core
+from paddle.base.libpaddle.pir import (
+    get_used_external_value,
+)
+from paddle.base.wrapped_decorator import signature_safe_contextmanager
+
+# TODO(CZ): to be removed when we support dynamic shape by default.
+ALLOW_DYNAMIC_SHAPE_VJP_OPS = [
+    "pd_op.abs",
+    "pd_op.add",
+    "pd_op.amax",
+    "pd_op.amin",
+    "pd_op.angle",
+    "pd_op.argsort",
+    "pd_op.assign",
+    "pd_op.batch_norm_",
+    "pd_op.cast",
+    "pd_op.ceil",
+    "pd_op.concat",
+    "pd_op.cos",
+    "pd_op.cumprod",
+    "pd_op.cumsum",
+    "pd_op.divide",
+    "pd_op.dot",
+    "pd_op.dropout",
+    "pd_op.elementwise_pow",
+    "pd_op.erf",
+    "pd_op.exp",
+    "pd_op.expand",
+    "pd_op.floor",
+    "pd_op.fmax",
+    "pd_op.fmin",
+    "pd_op.gather",
+    "pd_op.gather_nd",
+    "pd_op.gelu",
+    "pd_op.group_norm",
+    "pd_op.hardsigmoid",
+    "pd_op.hardswish",
+    "pd_op.kron",
+    "pd_op.kthvalue",
+    "pd_op.layer_norm",
+    "pd_op.leaky_relu",
+    "pd_op.log",
+    "pd_op.logcumsumexp",
+    "pd_op.logsumexp",
+    "pd_op.linear_v2",
+    "pd_op.matmul",
+    "pd_op.max",
+    "pd_op.maximum",
+    "pd_op.mean",
+    "pd_op.minimum",
+    "pd_op.multiply",
+    "pd_op.pad",
+    "pd_op.pow",
+    "pd_op.prod",
+    "pd_op.reduce_as",
+    "pd_op.relu",
+    "pd_op.relu6",
+    "pd_op.reshape",
+    "pd_op.roll",
+    "pd_op.rsqrt",
+    "pd_op.scale",
+    "pd_op.scatter",
+    "pd_op.scatter_nd_add",
+    "pd_op.sigmoid",
+    "pd_op.silu",
+    "pd_op.sin",
+    "pd_op.softmax",
+    "pd_op.softsign",
+    "pd_op.split",
+    "pd_op.sqrt",
+    "pd_op.square",
+    "pd_op.squeeze",
+    "pd_op.stack",
+    "pd_op.subtract",
+    "pd_op.sum",
+    "pd_op.swiglu",
+    "pd_op.swish",
+    "pd_op.take_along_axis",
+    "pd_op.tanh",
+    "pd_op.tile",
+    "pd_op.topk",
+    "pd_op.transpose",
+    "pd_op.trunc",
+    "pd_op.unsqueeze",
+    "pd_op.where",
+    "pd_op.p_norm",
+    "pd_op.index_put",
+    "pd_op.index_add",
+    "pd_op.elu",
+    "pd_op.masked_fill",
+    "pd_op.masked_select",
+    "pd_op.var",
+]
+
+
+class ValueWrapper:
+    def __init__(self, value) -> None:
+        if isinstance(value, ValueWrapper):
+            assert isinstance(value._value, (type(None), pir.Value))
+        else:
+            if not isinstance(value, (type(None), pir.Value)):
+                raise TypeError(
+                    "Value Wrapper is only support None and pir.Value"
+                )
+        self._value = value._value if isinstance(value, ValueWrapper) else value
+
+    def __hash__(self) -> int:
+        if isinstance(self._value, pir.Value):
+            return self._value.hash()
+        else:
+            return hash(self._value)
+
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, ValueWrapper):
+            warnings.warn(
+                f'In ValueWrapper.__eq__ expected type of `other` is ValueWrapper but received {other.__class__}.'
+            )
+            return False
+
+        if self._value is None or other._value is None:
+            return self._value is None and other._value is None
+        return self._value.is_same(other._value)
+
+
+class ValueDict:
+    def __init__(
+        self,
+        iter=None,
+        *,
+        default_factory=None,
+    ):
+        self._items: dict[ValueWrapper] = {}
+        self._default_factory = default_factory
+        if iter is not None:
+            for key, val in iter.items():
+                self[key] = val
+
+    def copy(self):
+        ret = ValueDict()
+        ret._items = self._items.copy()
+        ret._default_factory = self._default_factory
+        return ret
+
+    def update(self, other_dict):
+        for key, val in other_dict.items():
+            self[key] = val
+
+    def keys(self):
+        for key in self._items.keys():
+            yield key._value
+
+    def values(self):
+        return self._items.values()
+
+    def items(self):
+        for key, val in self._items.items():
+            yield key._value, val
+
+    def get(self, key, default=None):
+        if not self.__contains__(key):
+            return default
+        return self._items[ValueWrapper(key)]
+
+    def pop(self, key):
+        if not self.__contains__(key):
+            raise KeyError(f'{key} is not in ValueDict')
+        return self._items.pop(ValueWrapper(key))
+
+    def setdefault(self, key, default=None):
+        if not self.__contains__(key):
+            self[key] = default
+        return self[key]
+
+    def __setitem__(self, key, val: Any):
+        self._items[ValueWrapper(key)] = val
+
+    def __getitem__(self, key):
+        if not self.__contains__(key):
+            if self._default_factory is not None:
+                self[key] = self._default_factory()
+            else:
+                raise KeyError(f'{key} is not in ValueDict')
+        return self._items[ValueWrapper(key)]
+
+    def __bool__(self):
+        return bool(self._items)
+
+    def __len__(self):
+        return len(self._items)
+
+    def __iter__(self):
+        return self.keys()
+
+    def __contains__(self, key):
+        return ValueWrapper(key) in self._items
+
+    def __repr__(self) -> str:
+        items_str = ", ".join(f"{key}: {val}" for key, val in self.items())
+        return f'ValueDict({items_str})'
+
+
+class ValueSet:
+    def __init__(
+        self, iter: Sequence[ValueWrapper] | set[ValueWrapper] | None = None
+    ):
+        self._set: set[ValueWrapper] = set()
+        if iter is not None:
+            for val in iter:
+                self.add(val)
+
+    def copy(self):
+        ret = ValueSet()
+        ret._set = self._set.copy()
+        return ret
+
+    def add(self, val):
+        if not self.__contains__(val):
+            self._set.add(ValueWrapper(val))
+
+    def update(self, other: set):
+        for val in other:
+            self.add(val)
+
+    def pop(self):
+        return self._set.pop()._value
+
+    def remove(self, val):
+        self._set.remove(ValueWrapper(val))
+
+    def discard(self, val):
+        self._set.discard(ValueWrapper(val))
+
+    def __and__(self, other: ValueSet):
+        return ValueSet(self._set & other._set)
+
+    def __sub__(self, other: ValueSet):
+        return ValueSet(self._set - other._set)
+
+    def __or__(self, other: ValueSet):
+        return ValueSet(self._set | other._set)
+
+    def __bool__(self):
+        return bool(self._set)
+
+    def __len__(self):
+        return len(self._set)
+
+    def __iter__(self):
+        for val in self._set:
+            yield val._value
+
+    def __contains__(self, val):
+        return ValueWrapper(val) in self._set
+
+    def __repr__(self) -> str:
+        items_str = ", ".join(repr(item) for item in self)
+        return f'ValueSet({items_str})'
+
+
+class State:
+    """
+    record relationship of forward op/value and backward op/value
+    one state must be binding with a block, if block has parent block,
+    state will include parent block info.
+
+    """
+
+    def __init__(self, block):
+        self.block = block
+        # value -> list(list(value))
+        self.value_to_valuegrad = ValueDict(default_factory=list)
+        self.value_to_sumvaluegrad = ValueDict(default_factory=list)
+        # operation -> list(operation)
+        self.op_to_opgrad = collections.defaultdict(list)
+
+        # value -> list(value)
+        self.valuegrad_to_value = ValueDict(default_factory=list)
+        self.sumvaluegrad_to_value = ValueDict(default_factory=list)
+        # operation -> list(operation)
+        self.opgrad_to_op = collections.defaultdict(list)
+        # only for controlflow
+        # inside_value is sub block value, which will yield to parent block,
+        # parent block value is outside_value
+        self.inside_value_to_outside_value_map = ValueDict()
+
+    def turn_map(self) -> None:
+        self.valuegrad_to_value = ValueDict(default_factory=list)
+        self.sumvaluegrad_to_value = ValueDict(default_factory=list)
+        self.opgrad_to_op = collections.defaultdict(list)
+
+        for k, v in self.value_to_valuegrad.items():
+            if v != []:
+                for value in v[0]:
+                    self.valuegrad_to_value[value] = [k]
+        for k, v in self.value_to_sumvaluegrad.items():
+            if v != []:
+                for value in v[0]:
+                    self.sumvaluegrad_to_value[value] = [k]
+        for k, v in self.op_to_opgrad.items():
+            if v != []:
+                self.opgrad_to_op[v[0]] = [k]
+
+    def copy(self, new_block):
+        state = State(new_block)
+        state.value_to_valuegrad = self.value_to_valuegrad.copy()
+        state.value_to_sumvaluegrad = self.value_to_sumvaluegrad.copy()
+
+        # operation -> list(operation)
+        state.op_to_opgrad = self.op_to_opgrad.copy()
+
+        # value -> list(value)
+        state.valuegrad_to_value = self.valuegrad_to_value.copy()
+        state.sumvaluegrad_to_value = self.sumvaluegrad_to_value.copy()
+        # operation -> list(operation)
+        state.opgrad_to_op = self.opgrad_to_op.copy()
+
+        # only for controlflow
+        state.inside_value_to_outside_value_map = (
+            self.inside_value_to_outside_value_map.copy()
+        )
+
+        return state
+
+
+def _check_vjp_dynamic_shape(op, inputs):
+    for items in inputs:
+        for item in items:
+            if (
+                item.is_dense_tensor_type()
+                and item.initialized()
+                and -1 in item.shape
+            ):
+                return True
+
+
+# Prim currently does not support dynamic shape, when dynamic shape exits in shape of op inputs, prim will be skipped its vjp op.
+@signature_safe_contextmanager
+def dynamic_shape_prim_vjp_guard(op, inputs):
+    origin_prim = core._is_bwd_prim_enabled()
+    if op.name() == "cf.tuple_push":
+        skip_prim = True
+    else:
+        skip_prim = (
+            origin_prim
+            and core._enable_prim_skip_dynamic_shape()
+            and _check_vjp_dynamic_shape(op, inputs)
+            and op.name() not in ALLOW_DYNAMIC_SHAPE_VJP_OPS
+        )
+
+    try:
+        if origin_prim and skip_prim:
+            core._set_prim_backward_enabled(False)
+        yield
+    finally:
+        if origin_prim:
+            core._set_prim_backward_enabled(True)
+
+
+def check_type(input, input_name, expected_type, op_name, extra_message=''):
+    if not isinstance(input, expected_type):
+        raise TypeError(
+            f"The type of '{input_name}' in {op_name} must be {expected_type}, but received {type(input)}. {extra_message}"
+        )
+
+
+def _as_list(x):
+    if x is None:
+        return []
+    return list(x) if isinstance(x, Sequence) else [x]
+
+
+def some_in_set(value_list, value_set):
+    return any(v in value_set for v in value_list)
+
+
+def is_control_flow(op):
+    return op.name() == "pd_op.if" or op.name() == "pd_op.while"
+
+
+def is_builtin_op(op):
+    dialect_name, opname = op.name().split(".")
+    return dialect_name == "builtin"
+
+
+def update_no_grad_set_by_stopgradient(block, no_grad_set):
+    for op in block.ops:
+        if is_control_flow(op):
+            for sub_block in op.blocks():
+                update_no_grad_set_by_stopgradient(sub_block, no_grad_set)
+        for value in op.results():
+            if value.stop_gradient and value not in no_grad_set:
+                no_grad_set.add(value)
+
+
+def get_real_op_inputs(op):
+    if op.name() == "pd_op.if":
+        return get_used_external_value(op)
+    elif op.name() == "pd_op.while":
+        return op.operands_source() + get_used_external_value(
+            op.as_while_op().body()
+        )
+    elif op.name() == "pd_op.pylayer":
+        return get_used_external_value(op)
+    else:
+        return op.operands_source()
+
+
+def get_real_op_outputs(op):
+    outputs = op.results()
+    if op.name() == "pd_op.array_write_":
+        for x in op.operands():
+            outputs.append(x.source())
+    if op.name() == "pd_op.while":
+        for internal_op in op.as_while_op().body().ops:
+            if internal_op.name() == "pd_op.array_write_":
+                for x in internal_op.operands():
+                    outputs.append(x.source())
+    return outputs
+
+
+def inverse_sort_op(old_ops):
+    '''
+    if topo graph is op1 -> op2 -> op3
+    return [op3, op2, op1]
+
+    '''
+
+    # init pending_count[op] which describes number of
+    # pending edges for its grad_op
+
+    pending_count = collections.defaultdict(int)
+    ops = []
+    [ops.append(x) for x in old_ops if x not in ops]
+    ops_set = set(ops)
+    sorted_list = []
+    for op in ops:
+        for x in get_real_op_inputs(op):
+            if not pir.is_fake_value(x) and x.get_defining_op() in ops_set:
+                pending_count[x.get_defining_op()] += 1
+
+    queue = collections.deque()
+
+    for op in ops:
+        if pending_count[op] == 0:
+            queue.append(op)
+
+    while queue:
+        op = queue.popleft()
+        sorted_list.append(op)
+        for x in get_real_op_inputs(op):
+            x_op = x.get_defining_op()
+            pending_count[x_op] -= 1
+            if pending_count[x_op] == 0:
+                queue.append(x_op)
+
+    if len(sorted_list) != len(ops):
+        raise ValueError(
+            "inverse_sort_op wrong, sorted_list size is not equal to origin_list size"
+        )
+    change_list = []
+    # true  %0 = op1, 1% = increment(0%), 3% = op2(0%), tuple_push(%0, 1%, 3%),
+    # no one use 1% so increment be the first op, actually op2 use 1% ,
+    # sorted_list = [increment, op2, op1] should be [op2, increment, op1],
+    # tuple_push(0%) must be forward last op, backward first op, so skip it.
+    for op in reversed(sorted_list):
+        if op.name() == 'pd_op.increment_':
+            idx_1 = sorted_list.index(op)
+            idx_2 = sorted_list.index(op)
+
+            for op_in in reversed(sorted_list[: sorted_list.index(op)]):
+                if (
+                    some_in_set(
+                        op.operands_source(),
+                        ValueSet(get_real_op_inputs(op_in)),
+                    )
+                    and op_in.name() != "cf.tuple_push"
+                ):
+                    idx_2 = sorted_list.index(op_in)
+            if idx_1 != idx_2:
+                change_list.append((idx_1, idx_2))
+    for idx_1, idx_2 in change_list:
+        sorted_list[idx_1], sorted_list[idx_2] = (
+            sorted_list[idx_2],
+            sorted_list[idx_1],
+        )
+
+    return sorted_list
+
+
+def is_inplace_net(op_list):
+    '''
+    when program has inplace op , it's difficult to find the actual pending_count.
+    '''
+    for op in op_list:
+        if op.name() in ["pd_op.array_write_", "pd_op.assign_out_"]:
+            return True
+        if is_control_flow(op):
+            for block in op.blocks():
+                if is_inplace_net(block.ops):
+                    return True
+
+    return False
+
+
+def remove_op(block, op, state):
+    '''
+    remove op from block
+    '''
+    if state.opgrad_to_op[op] != []:
+        fwd_op = state.opgrad_to_op[op][0]
+        state.op_to_opgrad[fwd_op].remove(op)
+
+    for valuegrad in op.results():
+        if state.valuegrad_to_value[valuegrad] != []:
+            value = state.valuegrad_to_value[valuegrad][0]
+            state.value_to_valuegrad[value] = []
+
+            if value in state.sumvaluegrad_to_value:
+                raise ValueError(
+                    f'input_grad in [%s] is value which need to sum {op.name()}'
+                )
+    # NOTE(SigureMo): Ensure access to the op's results before removing it.
+    # Otherwise, the op will be deconstructed and access the num_results
+    # will be undefined behavior, it always cause hanging on the macOS.
+    block.remove_op(op)
+
+
+def while_prune_check(while_tuple_ops):
+    if len(while_tuple_ops) != 0:
+        for opresult in while_tuple_ops[0].results():
+            if not opresult.use_empty():
+                return False
+        return True
+    return False
+
+
+def remove_useless_full_like_ops(block, ops, state):
+    '''
+    remove ops which are not in use recursively,
+
+    '''
+    remove_ops = []
+    inverse_ops = inverse_sort_op(list(ops))
+    # from output to input
+    for op in inverse_ops:
+        if op.name() == "pd_op.full_like":
+            if op.result(0).use_empty():
+                full_op = op.operand_source(1).get_defining_op()
+                remove_ops.append(op)
+                remove_ops.append(full_op)
+        elif is_control_flow(op):
+            for sub_block in op.blocks():
+                remove_useless_full_like_ops(sub_block, sub_block.ops, state)
+
+    for op in remove_ops:
+        remove_op(block, op, state)
+
+
+def all_stop_gradient_true(block):
+    for op in block.ops:
+        for value in op.results():
+            if value.stop_gradient is False:
+                return False
+    return True
+
+
+def all_input_stop_gradient_true(list_of_list):
+    for list_ in list_of_list:
+        for stop_gradient in list_:
+            if stop_gradient is False:
+                return False
+    return True
+
+
+def all_output_grad_none(list_of_list):
+    for list_ in list_of_list:
+        for value in list_:
+            if value is not None:
+                return False
+    return True
+
+
+def op_has_vjp(op):
+    # NOTE(MarioLulab): In PIR mode, even though the `PyLayer` op does
+    # not have a vjp interface, we still need to generate the backward
+    # block based on its registered backward function. To achieve this,
+    # we add more handling logic for `PyLayer` Op in the `call_vjp` function
+    return core.has_vjp(op) or op.name() == "pd_op.pylayer"
+
+
+def parent_total_ops(block):
+    '''
+    when block is sub_block, forward op should include its parent block ops
+    (sub block nest should Add on demand to avoid block copy)
+    '''
+    total_ops = []
+    if block.parent_block is not None:
+        if block.parent_block.parent_block:
+            total_ops += block.parent_block.parent_block.ops
+        total_ops += block.parent_block.ops
+    total_ops += block.ops
+
+    return total_ops
+
+
+# only for control_flow to find corresponding value or value_list
+def return_map_value(value, map):
+    output = value
+    while output in map:
+        output = map[output]
+    return output
+
+
+def return_map_value_list(value, map):
+    output = []
+    for i in range(len(value)):
+        if value[i] in map:
+            output.append(return_map_value(value[i], map))
+        else:
+            output.append(value[i])
+    return output
+
+
+def argument_to_value(while_op):
+    '''
+    return while op's relationship of (block_argument to input value) and (input value to block_argument).
+    '''
+    if while_op.name() != "pd_op.while":
+        return ValueDict(), ValueDict()
+
+    assert len(while_op.as_while_op().block_arguments()) + 1 == len(
+        while_op.operands_source()
+    ), (
+        "while op's block_arguments size + 1 should same to while op's operands_source size"
+    )
+    arg_to_value_map = ValueDict()
+    value_to_arg_map = ValueDict()
+    for arg, value in zip(
+        while_op.as_while_op().block_arguments(),
+        while_op.operands_source()[1:],
+    ):
+        arg_to_value_map[arg] = value
+        value_to_arg_map[value] = arg
+    return arg_to_value_map, value_to_arg_map
+
+
+def get_grad_semantic_info(op):
+    '''
+    return whether op's inputs has grad, usually handled from yaml.
+    some op has uncertain inputs need special handling.
+    '''
+    if op.name() in [
+        "builtin.combine",
+        "pd_op.if",
+        "pd_op.while",
+        "pd_op.pylayer",
+        "cf.tuple_push",
+        "dist_op.moe_global_mesh_tensor",
+        "dist_op.moe_sub_mesh_tensors",
+        "dist_op.dist_reshape",
+    ]:
+        grad_semantic_info = [True for _ in range(len(get_real_op_inputs(op)))]
+        if op.name() == "pd_op.if":
+            grad_semantic_info[0] = False
+    else:
+        grad_semantic_info = op.get_input_grad_semantics()
+    return grad_semantic_info
+
+
+def get_split_op(value):
+    for op in value.all_used_ops():
+        if op.name() == "builtin.split":
+            return op
+    return None
+
+
+@lru_cache
+def warning_once(message: str):
+    logging.warning(message)
+
+
+def update_if_output_stopgradient(if_op, true_yield_op, false_yield_op):
+    """
+    Update if_op's stop_gradient based on true_yield_op and false_yield_op.
+
+    Args:
+    true_yield_op: true block of if_op's last op.
+    false_yield_op: false block of if_op's last op.
+    if_op: update it's op_results()'s stop_gradient.
+    """
+    if (
+        true_yield_op.name() != 'cf.yield'
+        or false_yield_op.name() != 'cf.yield'
+    ):
+        raise ValueError("param is not yield op")
+
+    # Check if operands_source sizes match
+    if len(true_yield_op.operands_source()) != len(
+        false_yield_op.operands_source()
+    ):
+        raise ValueError("Mismatched yield operands_source sizes")
+
+    # Check if op_results size matches operands_source
+    if len(if_op.results()) != len(true_yield_op.operands_source()):
+        raise ValueError(
+            "Mismatched if op_results size with yield operands_source"
+        )
+
+    # Update if_op's stop_gradient
+    for i in range(len(true_yield_op.operands_source())):
+        stop_grad1 = true_yield_op.operand_source(i).stop_gradient
+        stop_grad2 = false_yield_op.operand_source(i).stop_gradient
+
+        # Set to False if either stop_gradient is False
+        if not stop_grad1 or not stop_grad2:
+            if_op.result(i).stop_gradient = False
+
+
+def update_while_output_stopgradient(while_op, yield_op):
+    """
+    Update while_op's stop_gradient based on yield_op.
+
+    Args:
+    yield_op: The yield operation associated with the while loop.
+    while_op: The while operation whose op_results()'s stop_gradient needs to be updated.
+    """
+    # Check if yield_op is indeed a yield operation
+    if yield_op.name() != 'cf.yield':
+        raise ValueError("yield_op is not a yield operation")
+
+    # Check if operands_source size of yield_op matches op_results size of while_op
+    if len(while_op.results()) + 1 != len(yield_op.operands_source()):
+        raise ValueError(
+            f"Mismatched while op_results  size %d with yield operands_source %d. {len(while_op.results()) + 1, len(yield_op.operands_source())}"
+        )
+
+    # Update while_op's stop_gradient
+    for i in range(1, len(yield_op.operands_source())):
+        stop_grad = yield_op.operand_source(i).stop_gradient
+
+        # Set to False if stop_gradient is False
+        if not stop_grad:
+            while_op.result(i - 1).stop_gradient = False
+
+
+def find_index_of_yield(value, yield_op):
+    for i, v in enumerate(yield_op.operands_source()):
+        if v.is_same(value):
+            return i
+    return -1
+
+
+def update_tuple_pop_origin_inputs(tuple_pop_outputs):
+    if tuple_pop_outputs == []:
+        return tuple_pop_outputs
+    op = tuple_pop_outputs[0][0].get_defining_op()
+    assert op.name() == "cf.tuple_pop"
+    stack_op = op.operand_source(0).get_defining_op()
+    tuple_push_inputs = stack_op.result(1).first_use().owner().operands_source()
+    tuple_push_inputs_with_if = []
+    for input in tuple_push_inputs:
+        if input.first_use().owner().name() == "cf.yield":
+            yield_op = input.first_use().owner()
+            index = find_index_of_yield(input, yield_op)
+            assert index != -1
+            tuple_push_inputs_with_if.append(
+                yield_op.get_parent_block().parent_op.result(index)
+            )
+        else:
+            tuple_push_inputs_with_if.append(input)
+
+    # pass inlets
+    return tuple_push_inputs_with_if[1:]
+
+
+def value_in_block(value, block):
+    value_block = value.get_defining_op().get_parent_block()
+    while block.parent_op.name() != "builtin.module":
+        if block == value_block:
+            return True
+        block = block.parent_block
+    # now block is module op's block
+    if block == value_block:
+        return True
+
+    return False

python/paddle/autograd/function.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2026 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 .py_layer import (
+    PyLayerContext as FunctionCtx,  # noqa: F401
+    once_differentiable,  # noqa: F401
+)

python/paddle/autograd/grad_mode.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 paddle.base.dygraph.base import set_grad_enabled  # noqa: F401

python/paddle/autograd/py_layer.py

@@ -0,0 +1,504 @@
+# Copyright (c) 2021 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
+
+from typing import TYPE_CHECKING, Any, Concatenate, TypeVar
+
+import paddle
+from paddle.base import core
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    from paddle import Tensor
+
+
+__all__ = []
+
+
+_RetT = TypeVar('_RetT')
+
+
+class PyLayerContext:
+    """
+    ``PyLayerContext`` can assist the :ref:`api_paddle_autograd_PyLayer` in implementing certain functionalities.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.autograd import PyLayer
+
+            >>> class cus_tanh(PyLayer):
+            ...     @staticmethod
+            ...     def forward(ctx, x):
+            ...         # ctx is a object of PyLayerContext.
+            ...         y = paddle.tanh(x)
+            ...         ctx.save_for_backward(y)
+            ...         return y
+            ...
+            ...     @staticmethod
+            ...     def backward(ctx, dy):
+            ...         # ctx is a object of PyLayerContext.
+            ...         (y,) = ctx.saved_tensor()
+            ...         grad = dy * (1 - paddle.square(y))
+            ...         return grad
+    """
+
+    container: tuple[Tensor, ...]
+    not_inplace_tensors: tuple[Tensor, ...]
+    non_differentiable: tuple[Tensor, ...]
+    materialize_grads: bool
+    grad_in_dtype_consistent: bool
+
+    def set_grad_in_dtype_consistent(self, flag: bool) -> None:
+        """
+        Set whether to maintain gradient input dtype consistency between forward output and backward input.
+
+        Note:
+            This API should be called only inside `forward`.
+            By default, backward input gradients are automatically cast to match the dtype of forward outputs.
+            Set this to `False` to disable automatic casting and maintain original gradient dtypes in backward.
+
+        Args:
+            flag (bool): Whether to enable automatic dtype conversion in backward.
+                - `True`:  Cast backward input gradient to match forward output dtype (default behavior)
+                - `False`: Preserve original dtype of backward input gradient
+
+        Returns:
+            None
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.autograd import PyLayer
+                >>> paddle.seed(2025)
+                >>> class cus_tanh(PyLayer):
+                ...     @staticmethod
+                ...     def forward(ctx, x):
+                ...         y = paddle.tanh(x)
+                ...         # Pass tensors to backward.
+                ...         ctx.save_for_backward(y)
+                ...         # The gradient input in the backward process
+                ...         # will not be automatically cast to the dtype of the forward output.
+                ...         ctx.set_grad_in_dtype_consistent(False)
+                ...         return y
+                ...
+                ...     @staticmethod
+                ...     def backward(ctx, dy):
+                ...
+                ...         # Get the tensors passed by forward.
+                ...         (y,) = ctx.saved_tensor()
+                ...         grad = dy * (1 - paddle.square(y))
+                ...         return grad
+                >>> class cus_tanh_cast_grad(PyLayer):
+                ...     @staticmethod
+                ...     def forward(ctx, x):
+                ...         y = paddle.tanh(x)
+                ...         # Pass tensors to backward.
+                ...         ctx.save_for_backward(y)
+                ...         return y
+                ...
+                ...     @staticmethod
+                ...     def backward(ctx, dy):
+                ...         # Get the tensors passed by forward.
+                ...         (y,) = ctx.saved_tensor()
+                ...         grad = dy * (1 - paddle.square(y))
+                ...         # The gradient input in cus_tanh be cast to bfloat16 manually,
+                ...         # and cus_tanh will not cast the gradient to the dtype of the forward output.
+                ...         grad = paddle.cast(grad, paddle.float16)
+                ...         return grad
+                >>> x = paddle.randn([3, 3]).astype("float32")
+                >>> x.stop_gradient = False
+                >>> y = cus_tanh.apply(x)
+                >>> z = cus_tanh_cast_grad.apply(y)
+                >>> z.sum().backward()
+
+        """
+        self.grad_in_dtype_consistent = flag
+
+    def save_for_backward(self, *tensors: Tensor) -> None:
+        """
+        Saves given tensors that backward need. Use ``saved_tensor`` in the `backward` to get the saved tensors.
+
+        Note:
+            This API should be called at most once, and only inside `forward`.
+
+        Args:
+            tensors(list of Tensors): Tensors to be stored.
+
+        Returns:
+            None
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.autograd import PyLayer
+
+                >>> class cus_tanh(PyLayer):
+                ...     @staticmethod
+                ...     def forward(ctx, x):
+                ...         # ctx is a context object that store some objects for backward.
+                ...         y = paddle.tanh(x)
+                ...         # Pass tensors to backward.
+                ...         ctx.save_for_backward(y)
+                ...         return y
+                ...
+                ...     @staticmethod
+                ...     def backward(ctx, dy):
+                ...         # Get the tensors passed by forward.
+                ...         (y,) = ctx.saved_tensor()
+                ...         grad = dy * (1 - paddle.square(y))
+                ...         return grad
+
+        """
+        self.container = tensors
+
+    def saved_tensor(self) -> tuple[Tensor, ...]:
+        """
+        Get the tensors stored by ``save_for_backward``.
+
+        Returns:
+            list of Tensors or None: If context contains tensors stored by `save_for_backward`,
+            then return these tensors, otherwise return None.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.autograd import PyLayer
+
+                >>> class cus_tanh(PyLayer):
+                ...     @staticmethod
+                ...     def forward(ctx, x):
+                ...         # ctx is a context object that store some objects for backward.
+                ...         y = paddle.tanh(x)
+                ...         # Pass tensors to backward.
+                ...         ctx.save_for_backward(y)
+                ...         return y
+                ...
+                ...     @staticmethod
+                ...     def backward(ctx, dy):
+                ...         # Get the tensors passed by forward.
+                ...         (y,) = ctx.saved_tensor()
+                ...         grad = dy * (1 - paddle.square(y))
+                ...         return grad
+        """
+        return self.container
+
+    @property
+    def saved_tensors(self):
+        """
+        Get the tensors stored by ``save_for_backward``. This attribute is an alias for the method ``saved_tensor()``.
+
+        Returns:
+            list of Tensors or None: If context contains tensors stored by `save_for_backward`,
+            then return these tensors, otherwise return None.
+        """
+        return self.saved_tensor()
+
+    def mark_not_inplace(self, *args: Tensor) -> None:
+        """
+        Marks inputs as not inplace.
+        This should be called at most once, only from inside the `forward` method,
+        and all arguments should be Tensor inputs.
+
+        If the Tensor returned by `forward` method is the same as the Tensor input of forward,
+        and this Tensor is marked as not_inplace, then Paddle will help the user create a new Tensor as output.
+        Thereby preventing the auto grad information of the input Tensor from being overwritten.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+
+                >>> class Exp(paddle.autograd.PyLayer):
+                ...     @staticmethod
+                ...     def forward(ctx, x):
+                ...         ctx.mark_not_inplace(x)
+                ...         return x
+                ...
+                ...     @staticmethod
+                ...     def backward(ctx, grad_output):
+                ...         out = grad_output.exp()
+                ...         return out
+
+                >>> paddle.seed(2023)
+                >>> x = paddle.randn((1, 1))
+                >>> x.stop_gradient = False
+                >>> attn_layers = []
+                >>> for idx in range(0, 2):
+                ...     attn_layers.append(Exp())
+
+                >>> for step in range(0, 2):
+                ...     a = x
+                ...     for j in range(0, 2):
+                ...         a = attn_layers[j].apply(x)
+                ...     a.backward()
+        """
+        self.not_inplace_tensors = args
+
+    def mark_non_differentiable(self, *args: Tensor) -> None:
+        """
+        Marks outputs as non-differentiable.
+        This should be called at most once, only from inside the `forward` method,
+        and all arguments should be tensor outputs.
+
+        This will mark outputs as not requiring gradients, increasing the
+        efficiency of backward computation. You still need to accept a gradient
+        for each output in `backward`, but it's always going to
+        be a zero tensor with the same shape as the shape of a corresponding
+        output.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.autograd import PyLayer
+                >>> import numpy as np
+
+                >>> class Tanh(PyLayer):
+                ...     @staticmethod
+                ...     def forward(ctx, x):
+                ...         a = x + x
+                ...         b = x + x + x
+                ...         ctx.mark_non_differentiable(a)
+                ...         return a, b
+                ...
+                ...     @staticmethod
+                ...     def backward(ctx, grad_a, grad_b):
+                ...         assert np.equal(grad_a.numpy(), paddle.zeros([1]).numpy())
+                ...         assert np.equal(grad_b.numpy(), paddle.ones([1], dtype="float64").numpy())
+                ...         return grad_b
+
+                >>> x = paddle.ones([1], dtype="float64")
+                >>> x.stop_gradient = False
+                >>> a, b = Tanh.apply(x)
+                >>> b.sum().backward()
+        """
+        self.non_differentiable = args
+
+    def set_materialize_grads(self, value: bool) -> None:
+        """
+        Sets whether to materialize output grad tensors. Default is True.
+
+        This should be called only from inside the `forward` method.
+
+        If True, undefined output grad tensors will be expanded to tensors full
+        of zeros prior to calling the `backward` method.
+
+        If False, undefined output grad tensors will be None.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.autograd import PyLayer
+                >>> import numpy as np
+
+                >>> class Tanh(PyLayer):
+                ...     @staticmethod
+                ...     def forward(ctx, x):
+                ...         return x + x + x, x + x
+                ...
+                ...     @staticmethod
+                ...     def backward(ctx, grad, grad2):
+                ...         assert np.equal(grad2.numpy(), paddle.zeros([1]).numpy())
+                ...         return grad
+
+                >>> class Tanh2(PyLayer):
+                ...     @staticmethod
+                ...     def forward(ctx, x):
+                ...         ctx.set_materialize_grads(False)
+                ...         return x + x + x, x + x
+                ...
+                ...     @staticmethod
+                ...     def backward(ctx, grad, grad2):
+                ...         assert grad2 == None
+                ...         return grad
+
+                >>> x = paddle.ones([1], dtype="float64")
+                >>> x.stop_gradient = False
+                >>> Tanh.apply(x)[0].backward()
+
+                >>> x2 = paddle.ones([1], dtype="float64")
+                >>> x2.stop_gradient = False
+                >>> Tanh2.apply(x2)[0].backward()
+        """
+        self.materialize_grads = value
+
+
+class PyLayerBackward(core.eager.PyLayer, PyLayerContext):
+    def backward(self, *args):
+        return self._forward_cls.backward(self, *args)
+
+
+class PyLayerMeta(type):
+    def __init__(cls, name, bases, attrs):
+        cls._backward_function = type(
+            name + '_backward', (PyLayerBackward,), {"_forward_cls": cls}
+        )
+
+        super().__init__(name, bases, attrs)
+
+
+class PyLayer(core.eager.PyLayer, PyLayerContext, metaclass=PyLayerMeta):
+    """
+    Paddle implements Python custom operators on the PaddlePaddle framework by creating a subclass of
+    ``PyLayer``, which must comply with the following rules:
+
+    1. The subclass must contain static ``forward`` and ``backward`` functions, with the first argument being
+    :ref:`api_paddle_autograd_PyLayerContext`. If a returned value in ``backward`` corresponds to a ``Tensor`` that
+    requires gradients in ``forward``, the returned value must be a ``Tensor``.
+
+    2. Except for the first argument, other arguments of ``backward`` are gradients of the output ``Tensors``
+    of ``forward``. Therefore, the number of input ``Tensor`` in ``backward`` must be the same as the number
+    of output ``Tensor`` in ``forward``. If you need to use input ``Tensor`` from ``forward`` in ``backward``,
+    you can save these ``Tensors`` by inputting them into :ref:`api_paddle_autograd_PyLayerContext`'s
+    ``save_for_backward`` method and use them in ``backward`` later.
+
+    3. The output of ``backward`` can be ``Tensor`` or ``list/tuple(Tensor)``, which are gradients of the
+    output ``Tensor`` of ``forward``. Therefore, the number of output ``Tensor`` in ``backward`` is the same
+    as the number of input ``Tensor`` in ``forward``.
+
+    After building the custom operator, apply it by running the ``apply`` method.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.autograd import PyLayer
+
+            >>> class cus_tanh(PyLayer):
+            ...     @staticmethod
+            ...     def forward(ctx, x):
+            ...         y = paddle.tanh(x)
+            ...         # Pass tensors to backward.
+            ...         ctx.save_for_backward(y)
+            ...         return y
+            ...
+            ...     @staticmethod
+            ...     def backward(ctx, dy):
+            ...         # Get the tensors passed by forward.
+            ...         (y,) = ctx.saved_tensor()
+            ...         grad = dy * (1 - paddle.square(y))
+            ...         return grad
+
+            >>> paddle.seed(2023)
+            >>> data = paddle.randn([2, 3], dtype="float64")
+            >>> data.stop_gradient = False
+            >>> z = cus_tanh.apply(data)
+            >>> z.mean().backward()
+
+            >>> print(data.grad)
+            Tensor(shape=[2, 3], dtype=float64, place=Place(cpu), stop_gradient=True,
+            [[0.16604150, 0.05858341, 0.14051214],
+             [0.15677770, 0.01564609, 0.02991660]])
+    """
+
+    @staticmethod
+    def forward(
+        ctx: PyLayerContext, *args: Any, **kwargs: Any
+    ) -> Tensor | Sequence[Tensor]:
+        """
+        It is to be overloaded by subclasses. It must accept a object of :ref:`api_paddle_autograd_PyLayerContext` as
+        the first argument, followed by any number of arguments (tensors or other types).
+        `None` can not be included in the returned result.
+
+        Args:
+            *args(tuple): input of PyLayer.
+            **kwargs(dict): input of PyLayer.
+
+        Returns:
+            tensors or other types : output of PyLayer.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.autograd import PyLayer
+
+                >>> class cus_tanh(PyLayer):
+                ...     @staticmethod
+                ...     def forward(ctx, x):
+                ...         y = paddle.tanh(x)
+                ...         # Pass tensors to backward.
+                ...         ctx.save_for_backward(y)
+                ...         return y
+                ...
+                ...     @staticmethod
+                ...     def backward(ctx, dy):
+                ...         # Get the tensors passed by forward.
+                ...         (y,) = ctx.saved_tensor()
+                ...         grad = dy * (1 - paddle.square(y))
+                ...         return grad
+        """
+        raise NotImplementedError(
+            "You must implement the forward function for PyLayer."
+        )
+
+    @staticmethod
+    def backward(ctx: PyLayerContext, *args: Any) -> Tensor | Sequence[Tensor]:
+        """
+        This is a function to calculate the gradient. It is to be overloaded by subclasses.
+        It must accept a object of :ref:`api_paddle_autograd_PyLayerContext` as the first
+        argument, and the rest arguments are the gradient of forward's output tensors.
+        Output tensors of backward are the gradient of forward's input tensors.
+
+        Args:
+            *args(tuple): The gradient of forward's output tensor(s).
+            **kwargs(dict): The gradient of forward's output tensor(s).
+
+        Returns:
+            Tensor or list of Tensors: The gradient of forward's input tensor(s).
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.autograd import PyLayer
+
+                >>> class cus_tanh(PyLayer):
+                ...     @staticmethod
+                ...     def forward(ctx, x):
+                ...         y = paddle.tanh(x)
+                ...         # Pass tensors to backward.
+                ...         ctx.save_for_backward(y)
+                ...         return y
+                ...
+                ...     @staticmethod
+                ...     def backward(ctx, dy):
+                ...         # Get the tensors passed by forward.
+                ...         (y,) = ctx.saved_tensor()
+                ...         grad = dy * (1 - paddle.square(y))
+                ...         return grad
+        """
+
+        raise NotImplementedError(
+            "You must implement the backward function for PyLayer."
+        )
+
+
+def once_differentiable(
+    backward: Callable[Concatenate[PyLayerContext, ...], _RetT],
+) -> Callable[Concatenate[PyLayerContext, ...], _RetT]:
+    def wrapper(ctx: PyLayerContext, *args: Any) -> _RetT:
+        with paddle.base.dygraph.no_grad():
+            outputs = backward(ctx, *args)
+        return outputs
+
+    return wrapper

python/paddle/autograd/saved_tensors_hooks.py

@@ -0,0 +1,127 @@
+# Copyright (c) 2022 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
+
+from typing import TYPE_CHECKING, Any
+
+from paddle.base import core
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from paddle import Tensor
+__all__ = []
+
+
+class saved_tensors_hooks:
+    """
+    Dynamic graph, registers a pair of pack / unpack hooks for saved tensors.
+
+    Parameters:
+        pack_hook (function): The pack hook will be called every time the forward
+            operation inputs/outputs tensors need be saved for backward. Then you
+            can save it to CPU or Disk. The input of `pack_hook` is a tensor need
+            be saved. The output of `pack_hook` is then stored information instead
+            of the original tensor. `pack_hook` will also be called while any
+            tensor need be saved by `PyLayerContext.save_for_backward`. If a tensor
+            saved for backward is no need buffer, `pack_hook` will not be called.
+            Only the tensor saved for backward is DenseTensor, `pack_hook` will be
+            called.
+        unpack_hook (function): The unpack hook will be called every time the
+            backward need use the saved inputs/outputs tensors. Then you can reload
+            the tensor and return it to paddle framework. The input of `unpack_hook`
+            is the information returned by `pack_hook`. The output of `unpack_hook`
+            is a tensor reloaded by the information, and the tensor must has the same
+            content as the original tensor passed as input to the corresponding
+            `pack_hook`.
+
+    Returns:
+            None
+
+    Examples:
+        .. code-block:: pycon
+            :name: code-example1
+
+            >>> # Example1
+            >>> import paddle
+
+            >>> def pack_hook(x):
+            ...     print("Packing", x)
+            ...     return x.numpy()
+
+            >>> def unpack_hook(x):
+            ...     print("UnPacking", x)
+            ...     return paddle.to_tensor(x)
+
+            >>> a = paddle.ones([3, 3])
+            >>> b = paddle.ones([3, 3]) * 2
+            >>> a.stop_gradient = False
+            >>> b.stop_gradient = False
+            >>> with paddle.autograd.saved_tensors_hooks(pack_hook, unpack_hook):
+            ...     y = paddle.multiply(a, b)
+            >>> y.sum().backward()
+
+        .. code-block:: pycon
+            :name: code-example2
+
+            >>> # Example2
+            >>> import paddle
+            >>> from paddle.autograd import PyLayer
+
+            >>> class cus_multiply(PyLayer):
+            ...     @staticmethod
+            ...     def forward(ctx, a, b):
+            ...         y = paddle.multiply(a, b)
+            ...         ctx.save_for_backward(a, b)
+            ...         return y
+            ...
+            ...     @staticmethod
+            ...     def backward(ctx, dy):
+            ...         a, b = ctx.saved_tensor()
+            ...         grad_a = dy * a
+            ...         grad_b = dy * b
+            ...         return grad_a, grad_b
+
+            >>> def pack_hook(x):
+            ...     print("Packing", x)
+            ...     return x.numpy()
+
+            >>> def unpack_hook(x):
+            ...     print("UnPacking", x)
+            ...     return paddle.to_tensor(x)
+
+            >>> a = paddle.ones([3, 3])
+            >>> b = paddle.ones([3, 3]) * 2
+            >>> a.stop_gradient = False
+            >>> b.stop_gradient = False
+            >>> with paddle.autograd.saved_tensors_hooks(pack_hook, unpack_hook):
+            ...     y = cus_multiply.apply(a, b)
+            >>> y.sum().backward()
+    """
+
+    def __init__(
+        self,
+        pack_hook: Callable[[Tensor], Any | None],
+        unpack_hook: Callable[[Any], Tensor | None],
+    ) -> None:
+        self.pack_hook = pack_hook
+        self.unpack_hook = unpack_hook
+
+    def __enter__(self) -> None:
+        core.eager.register_saved_tensors_hooks(
+            self.pack_hook, self.unpack_hook
+        )
+
+    def __exit__(self, *args: object) -> None:
+        core.eager.reset_saved_tensors_hooks()

python/paddle/base/.gitignore

@@ -0,0 +1,3 @@
+proto
+core.so
+*.so

python/paddle/base/__init__.py

@@ -0,0 +1,213 @@
+#   Copyright (c) 2018 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 atexit
+import os
+import platform
+import sys
+
+# The legacy core need to be removed before "import core",
+# in case of users installing paddlepaddle without -U option
+core_suffix = 'so'
+if os.name == 'nt':
+    core_suffix = 'pyd'
+
+legacy_core = (
+    os.path.abspath(os.path.dirname(__file__)) + os.sep + 'core.' + core_suffix
+)
+if os.path.exists(legacy_core):
+    sys.stderr.write('Deleting legacy file ' + legacy_core + '\n')
+    try:
+        os.remove(legacy_core)
+    except Exception as e:
+        raise e
+
+# import all class inside framework into base module
+# import all class inside executor into base module
+from . import (  # noqa: F401
+    backward,
+    compiler,
+    core,
+    data_feed_desc,
+    dataset,
+    dygraph,
+    executor,
+    framework,
+    incubate,
+    initializer,
+    io,
+    layers,
+    trainer_desc,
+    unique_name,
+)
+from .backward import (  # noqa: F401
+    append_backward,
+    gradients,
+)
+from .compiler import (  # noqa: F401
+    BuildStrategy,
+    CompiledProgram,
+    IpuCompiledProgram,
+    IpuStrategy,
+)
+from .core import (  # noqa: F401
+    CPUPlace,
+    CUDAPinnedPlace,
+    CUDAPlace,
+    CustomPlace,
+    DenseTensor,
+    DenseTensorArray,
+    IPUPlace,
+    Scope,
+    XPUPinnedPlace,
+    XPUPlace,
+    _check_last_cuda_error,
+    _cuda_synchronize,
+    _Scope,
+    _set_warmup,
+)
+from .data_feed_desc import DataFeedDesc  # noqa: F401
+from .data_feeder import DataFeeder  # noqa: F401
+from .dataset import (  # noqa: F401
+    DatasetFactory,
+    InMemoryDataset,
+)
+from .dygraph.base import disable_dygraph, enable_dygraph
+from .dygraph.tensor_patch_methods import monkey_patch_tensor
+from .executor import (  # noqa: F401
+    Executor,
+    global_scope,
+    scope_guard,
+)
+from .framework import (  # noqa: F401
+    Program,
+    Variable,
+    cpu_places,
+    cuda_pinned_places,
+    cuda_places,
+    default_main_program,
+    default_startup_program,
+    device_guard,
+    get_flags,
+    in_dygraph_mode,
+    in_dynamic_or_pir_mode,
+    in_pir_mode,
+    ipu_shard_guard,
+    is_compiled_with_cinn,
+    is_compiled_with_cuda,
+    is_compiled_with_rocm,
+    is_compiled_with_xpu,
+    name_scope,
+    process_type_promotion,
+    program_guard,
+    require_version,
+    set_flags,
+    set_ipu_shard,
+    xpu_places,
+)
+from .initializer import set_global_initializer  # noqa: F401
+from .layers.math_op_patch import monkey_patch_variable
+from .lod_tensor import (  # noqa: F401
+    create_lod_tensor,
+    create_random_int_lodtensor,
+)
+from .param_attr import ParamAttr, WeightNormParamAttr  # noqa: F401
+from .trainer_desc import (  # noqa: F401
+    MultiTrainer,
+    TrainerDesc,
+)
+
+Tensor = DenseTensor
+enable_imperative = enable_dygraph
+disable_imperative = disable_dygraph
+
+__all__ = []
+
+
+def __bootstrap__():
+    """
+    Enable reading gflags from environment variables.
+
+    Returns:
+        None
+    """
+    try:
+        num_threads = int(os.getenv('OMP_NUM_THREADS', '1'))
+    except ValueError:
+        num_threads = 1
+
+    if num_threads > 1:
+        print(
+            f'WARNING: OMP_NUM_THREADS set to {num_threads}, not 1. The computation '
+            'speed will not be optimized if you use data parallel. It will '
+            'fail if this PaddlePaddle binary is compiled with OpenBlas since'
+            ' OpenBlas does not support multi-threads.',
+            file=sys.stderr,
+        )
+        print('PLEASE USE OMP_NUM_THREADS WISELY.', file=sys.stderr)
+
+    os.environ['OMP_NUM_THREADS'] = str(num_threads)
+
+    flag_prefix = "FLAGS_"
+    read_env_flags = [
+        key[len(flag_prefix) :]
+        for key in core.globals().keys()
+        if key.startswith(flag_prefix)
+    ]
+
+    def remove_flag_if_exists(name):
+        if name in read_env_flags:
+            read_env_flags.remove(name)
+
+    sysstr = platform.system()
+    if 'Darwin' in sysstr:
+        remove_flag_if_exists('use_pinned_memory')
+
+    if core.is_compiled_with_ipu():
+        # Currently we request all ipu available for training and testing
+        #   finer control of pod of IPUs will be added later
+        read_env_flags += []
+
+    core.init_gflags(["--tryfromenv=" + ",".join(read_env_flags)])
+    # Note(zhouwei25): sys may not have argv in some cases,
+    # Such as: use Python/C API to call Python from C++
+    try:
+        core.init_glog(sys.argv[0])
+    except Exception:
+        sys.argv = [""]
+        core.init_glog(sys.argv[0])
+    # don't init_p2p when in unittest to save time.
+    core.init_memory_method()
+    core.init_devices()
+    core.init_gflags_from_env()
+    core.init_tensor_operants()
+    core.init_default_kernel_signatures()
+
+
+# TODO(panyx0718): Avoid doing complex initialization logic in __init__.py.
+# Consider paddle.init(args) or paddle.main(args)
+monkey_patch_variable()
+__bootstrap__()
+monkey_patch_tensor()
+
+# NOTE(Aurelius84): clean up ExecutorCacheInfo in advance manually.
+atexit.register(core.clear_executor_cache)
+atexit.register(core.pir.clear_cinn_compilation_cache)
+
+# NOTE(Aganlengzi): clean up KernelFactory in advance manually.
+# NOTE(wangran16): clean up DeviceManager in advance manually.
+# Keep clear_kernel_factory running before clear_device_manager
+atexit.register(core.clear_device_manager)
+atexit.register(core.clear_kernel_factory)
+atexit.register(core.ProcessGroupIdMap.destroy)

python/paddle/base/core.py

@@ -0,0 +1,652 @@
+# Copyright (c) 2019 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 os
+import platform
+import site
+import sys
+import warnings
+
+has_paddle_dy_lib = False
+
+dy_lib_name = 'libpaddle'
+dy_lib_suffix = 'so'
+if os.name == 'nt':
+    dy_lib_suffix = 'pyd'
+
+current_path = os.path.abspath(os.path.dirname(__file__))
+if os.path.exists(current_path + os.sep + dy_lib_name + '.' + dy_lib_suffix):
+    has_paddle_dy_lib = True
+
+
+try:
+    if os.name == 'nt':
+        third_lib_path = current_path + os.sep + '..' + os.sep + 'libs'
+        # Will load shared library from 'path' on windows
+        os.environ['path'] = (
+            current_path + ';' + third_lib_path + ';' + os.environ['path']
+        )
+        sys.path.insert(0, third_lib_path)
+        # Note: from python3.8, PATH will not take effect
+        # https://github.com/python/cpython/pull/12302
+        # Use add_dll_directory to specify dll resolution path
+        os.add_dll_directory(third_lib_path)
+
+except ImportError as e:
+    if os.name == 'nt':
+        executable_path = os.path.abspath(os.path.dirname(sys.executable))
+        raise ImportError(
+            f"""NOTE: You may need to run \"set PATH={executable_path};%PATH%\"
+        if you encounters \"DLL load failed\" errors. If you have python
+        installed in other directory, replace \"{executable_path}\" with your own
+        directory. The original error is: \n {e}"""
+        )
+    else:
+        raise ImportError(
+            """NOTE: You may need to run \"export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH\"
+        if you encounters \"libmkldnn.so not found\" errors. If you have python
+        installed in other directory, replace \"/usr/local/lib\" with your own
+        directory. The original error is: \n"""
+            + str(e)
+        )
+except Exception as e:
+    raise e
+
+
+def avx_supported():
+    """
+    Whether current system(Linux, MacOS, Windows) is supported with AVX.
+    """
+    sysstr = platform.system().lower()
+    has_avx = False
+    if sysstr == 'linux':
+        try:
+            pipe = os.popen('cat /proc/cpuinfo | grep -i avx')
+            has_avx = pipe.read() != ''
+            pipe.close()
+        except Exception as e:
+            sys.stderr.write(
+                'Can not get the AVX flag from /proc/cpuinfo.\n'
+                f'The original error is: {e}\n'
+            )
+        return has_avx
+    elif sysstr == 'darwin':
+        try:
+            pipe = os.popen('sysctl machdep.cpu.features | grep -i avx')
+            has_avx = pipe.read() != ''
+            pipe.close()
+        except Exception as e:
+            sys.stderr.write(
+                'Can not get the AVX flag from machdep.cpu.features.\n'
+                f'The original error is: {e}\n'
+            )
+        if not has_avx:
+            import subprocess
+
+            pipe = subprocess.Popen(
+                'sysctl machdep.cpu.leaf7_features | grep -i avx',
+                shell=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            _ = pipe.communicate()
+            has_avx = True if pipe.returncode == 0 else False
+        return has_avx
+    elif sysstr == 'windows':
+        import ctypes
+
+        ONE_PAGE = ctypes.c_size_t(0x1000)
+
+        def asm_func(code_str, restype=ctypes.c_uint32, argtypes=()):
+            # Call the code_str as a function
+            # Alloc 1 page to ensure the protection
+            pfnVirtualAlloc = ctypes.windll.kernel32.VirtualAlloc
+            pfnVirtualAlloc.restype = ctypes.c_void_p
+            MEM_COMMIT = ctypes.c_ulong(0x1000)
+            PAGE_READWRITE = ctypes.c_ulong(0x4)
+            address = pfnVirtualAlloc(
+                None, ONE_PAGE, MEM_COMMIT, PAGE_READWRITE
+            )
+            if not address:
+                raise Exception("Failed to VirtualAlloc")
+
+            # Copy the code into the memory segment
+            memmove = ctypes.CFUNCTYPE(
+                ctypes.c_void_p,
+                ctypes.c_void_p,
+                ctypes.c_void_p,
+                ctypes.c_size_t,
+            )(ctypes._memmove_addr)
+            if memmove(address, code_str, len(code_str)) < 0:
+                raise Exception("Failed to memmove")
+
+            # Enable execute permissions
+            PAGE_EXECUTE = ctypes.c_ulong(0x10)
+            pfnVirtualProtect = ctypes.windll.kernel32.VirtualProtect
+            res = pfnVirtualProtect(
+                ctypes.c_void_p(address),
+                ONE_PAGE,
+                PAGE_EXECUTE,
+                ctypes.byref(ctypes.c_ulong(0)),
+            )
+            if not res:
+                raise Exception("Failed VirtualProtect")
+
+            # Flush instruction cache
+            pfnGetCurrentProcess = ctypes.windll.kernel32.GetCurrentProcess
+            pfnGetCurrentProcess.restype = ctypes.c_void_p
+            prochandle = ctypes.c_void_p(pfnGetCurrentProcess())
+            res = ctypes.windll.kernel32.FlushInstructionCache(
+                prochandle, ctypes.c_void_p(address), ONE_PAGE
+            )
+            if not res:
+                raise Exception("Failed FlushInstructionCache")
+
+            # Cast the memory to function
+            functype = ctypes.CFUNCTYPE(restype, *argtypes)
+            func = functype(address)
+            return func, address
+
+        # http://en.wikipedia.org/wiki/CPUID#EAX.3D1:_Processor_Info_and_Feature_Bits
+        # mov eax,0x1; cpuid; mov cx, ax; ret
+        code_str = b"\xb8\x01\x00\x00\x00\x0f\xa2\x89\xc8\xc3"
+        avx_bit = 28
+        retval = 0
+        try:
+            # Convert the code_str into a function that returns uint
+            func, address = asm_func(code_str)
+            retval = func()
+            ctypes.windll.kernel32.VirtualFree(
+                ctypes.c_void_p(address), ctypes.c_size_t(0), ONE_PAGE
+            )
+        except Exception as e:
+            sys.stderr.write(
+                'Failed getting the AVX flag on Windows.\n'
+                f'The original error is: {e}\n'
+            )
+        return (retval & (1 << avx_bit)) > 0
+    else:
+        sys.stderr.write(f'Do not get AVX flag on {sysstr}\n')
+        return False
+
+
+def run_shell_command(cmd):
+    import subprocess
+
+    out, err = subprocess.Popen(
+        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True
+    ).communicate()
+    if err:
+        return None
+    else:
+        return out.decode('utf-8').strip()
+
+
+def get_dso_path(core_so, dso_name):
+    if core_so and dso_name:
+        return run_shell_command(
+            f"ldd {core_so}|grep {dso_name}|awk '{{print $3}}'"
+        )
+    else:
+        return None
+
+
+def load_dso(dso_absolute_path):
+    if dso_absolute_path:
+        try:
+            from ctypes import cdll
+
+            cdll.LoadLibrary(dso_absolute_path)
+        except:
+            warnings.warn(f"Load {dso_absolute_path} failed")
+
+
+def pre_load(dso_name):
+    if has_paddle_dy_lib:
+        core_so = current_path + os.sep + dy_lib_name + '.' + dy_lib_suffix
+    else:
+        core_so = None
+    dso_path = get_dso_path(core_so, dso_name)
+    load_dso(dso_path)
+
+
+def get_libc_ver():
+    ldd_glibc = run_shell_command("ldd --version | awk '/ldd/{print $NF}'")
+    if ldd_glibc is not None:
+        return ("glibc", ldd_glibc)
+
+    ldd_musl = run_shell_command("ldd 2>&1 | awk '/Version/{print $NF}'")
+    if ldd_musl is not None:
+        return ("musl", ldd_musl)
+    return (None, None)
+
+
+def less_than_ver(a, b):
+    if a is None or b is None:
+        return False
+
+    import operator
+    import re
+
+    def to_list(s):
+        s = re.sub(r'(\.0+)+$', '', s)
+        return [int(x) for x in s.split('.')]
+
+    return operator.lt(to_list(a), to_list(b))
+
+
+# NOTE(zhiqiu): An error may occurs when import paddle in linux platform with glibc < 2.22,
+# the error message of which is "dlopen: cannot load any more object with static TLS".
+# This happens when:
+# (1) the number of dynamic shared libraries (DSO) loaded > 14,
+# (2) after that, load a dynamic shared library (DSO) with static TLS.
+# For paddle, the problem is that 'libgomp' is a DSO with static TLS, and it is loaded after 14 DSOs.
+# So, here is a tricky way to solve the problem by pre load 'libgomp' before 'libpaddle.so'.
+# The final solution is to upgrade glibc to > 2.22 on the target system.
+if platform.system().lower() == 'linux':
+    libc_type, libc_ver = get_libc_ver()
+    if libc_type == 'glibc' and less_than_ver(libc_ver, '2.23'):
+        try:
+            pre_load('libgomp')
+        except Exception as e:
+            # NOTE(zhiqiu): do not abort if failed, since it may success when import libpaddle.so
+            sys.stderr.write('Error: Can not preload libgomp.so')
+
+try:
+    from . import libpaddle
+
+    if avx_supported() and not libpaddle.is_compiled_with_avx():
+        sys.stderr.write(
+            "Hint: Your machine support AVX, but the installed paddlepaddle doesn't have avx core. "
+            "Hence, no-avx core with worse performance will be imported.\nIf you like, you could "
+            "reinstall paddlepaddle by 'python -m pip install --force-reinstall paddlepaddle-gpu[==version]' "
+            "to get better performance.\n"
+        )
+
+    # assign tensor alias
+    libpaddle.LoDTensor = libpaddle.DenseTensor
+    libpaddle.Tensor = libpaddle.DenseTensor
+    libpaddle.VarDesc.VarType.LOD_TENSOR = (
+        libpaddle.VarDesc.VarType.DENSE_TENSOR
+    )
+    libpaddle.VarDesc.VarType.LOD_TENSOR_ARRAY = (
+        libpaddle.VarDesc.VarType.DENSE_TENSOR_ARRAY
+    )
+
+    from .libpaddle import *  # noqa: F403
+    from .libpaddle import (  # noqa: F401
+        __doc__,
+        __file__,
+        __name__,
+        __package__,
+        __unittest_throw_exception__,
+        _append_python_callable_object_and_return_id,
+        _check_last_cuda_error,
+        _cleanup,
+        _create_loaded_parameter,
+        _cuda_synchronize,
+        _device_synchronize,
+        _dygraph_debug_level,
+        _get_all_register_op_kernels,
+        _get_amp_attrs,
+        _get_amp_op_list,
+        _get_current_stream,
+        _get_eager_deletion_vars,
+        _get_legacy_default_stream,
+        _get_phi_kernel_name,
+        _get_registered_phi_kernels,
+        _get_stream_from_external,
+        _get_use_default_grad_op_desc_maker_ops,
+        _has_grad,
+        _is_compiled_with_heterps,
+        _is_dygraph_debug_enabled,
+        _is_program_version_supported,
+        _Profiler,
+        _ProfilerResult,
+        _promote_types_if_complex_exists,
+        _RecordEvent,
+        _Scope,
+        _set_amp_op_list,
+        _set_current_stream,
+        _set_eager_deletion_mode,
+        _set_fuse_parameter_group_size,
+        _set_fuse_parameter_memory_size,
+        _set_has_grad,
+        _set_paddle_lib_path,
+        _set_warmup,
+        _switch_tracer,
+        _test_enforce_gpu_success,
+        _xpu_device_synchronize,
+        _xpu_get_current_stream,
+        _xpu_set_current_stream,
+    )
+
+    # isort: off
+
+    # custom device
+    from .libpaddle import (  # noqa: F401
+        CustomDeviceEvent,
+        CustomDeviceStream,
+        _get_current_custom_device_stream,
+        _set_current_custom_device_stream,
+        _synchronize_custom_device,
+    )
+
+    # prim controller flags
+    from .libpaddle import (  # noqa: F401
+        __set_all_prim_enabled,
+        __set_bwd_prim_enabled,
+        __set_eager_prim_enabled,
+        __set_fwd_prim_enabled,
+        _add_skip_comp_ops,
+        _is_bwd_prim_enabled,
+        _is_eager_prim_enabled,
+        _is_fwd_prim_enabled,
+        _is_all_prim_enabled,
+        _remove_skip_comp_ops,
+        _set_bwd_prim_blacklist,
+        _set_prim_target_grad_name,
+    )
+
+    # type promotion
+
+    # isort: on
+    if sys.platform != 'win32':
+        from .libpaddle import (  # noqa: F401
+            _array_to_share_memory_tensor,
+            _cleanup_mmap_fds,
+            _convert_to_tensor_list,
+            _erase_process_pids,
+            _remove_tensor_list_mmap_fds,
+            _set_max_memory_map_allocation_pool_size,
+            _set_process_pids,
+            _set_process_signal_handler,
+            _throw_error_if_process_failed,
+        )
+
+except Exception as e:
+    if has_paddle_dy_lib:
+        sys.stderr.write(
+            'Error: Can not import paddle core while this file exists: '
+            + current_path
+            + os.sep
+            + 'libpaddle.'
+            + dy_lib_suffix
+            + '\n'
+        )
+    if not avx_supported() and libpaddle.is_compiled_with_avx():
+        sys.stderr.write(
+            "Error: Your machine doesn't support AVX, but the installed PaddlePaddle is avx core, "
+            "you should reinstall paddlepaddle with no-avx core.\n"
+        )
+    raise e
+
+
+def set_paddle_custom_device_lib_path(lib_dir):
+    if os.environ.get('CUSTOM_DEVICE_ROOT', None) is not None:
+        # use set environment value
+        return
+    path1 = os.path.normpath(
+        os.path.join(lib_dir, '..', 'paddle_custom_device')
+    )
+    if os.path.exists(path1):
+        # set CUSTOM_DEVICE_ROOT default path (lib_dir/../paddle_custom_device)
+        os.environ['CUSTOM_DEVICE_ROOT'] = path1
+    else:
+        path2 = os.path.normpath(
+            os.path.join(lib_dir, '..', '..', 'paddle_custom_device')
+        )
+        if os.path.exists(path2):
+            # set CUSTOM_DEVICE_ROOT default path (lib_dir/../../paddle_custom_device)
+            os.environ['CUSTOM_DEVICE_ROOT'] = path2
+        else:
+            os.environ['CUSTOM_DEVICE_ROOT'] = ''
+
+
+# set paddle lib path
+def set_paddle_lib_path():
+    site_dirs = site.getsitepackages()
+    for site_dir in site_dirs:
+        lib_dir = os.path.sep.join([site_dir, 'paddle', 'libs'])
+        if os.path.exists(lib_dir):
+            _set_paddle_lib_path(lib_dir)
+            set_paddle_custom_device_lib_path(lib_dir)
+            return
+    if hasattr(site, 'USER_SITE') and site.USER_SITE:
+        lib_dir = os.path.sep.join([site.USER_SITE, 'paddle', 'libs'])
+        if os.path.exists(lib_dir):
+            _set_paddle_lib_path(lib_dir)
+            set_paddle_custom_device_lib_path(lib_dir)
+
+
+set_paddle_lib_path()
+
+
+# This api is used for check of model output.
+# In some cases, model does not straightly return data which can be used for check.
+# When this flag is set true, required data should be returned in model.
+def _model_return_data():
+    flag = os.getenv("FLAGS_model_return_data")
+    if flag and flag.lower() in ("1", "true"):
+        return True
+    else:
+        return False
+
+
+# This api is used for check whether prim is on
+def _prim_return_log():
+    flag = os.getenv("FLAGS_prim_log")
+    if flag and flag.lower() in ("1", "true"):
+        return True
+    else:
+        return False
+
+
+# ops in forward_blacklist will not be replaced by composite ops.
+prim_config = {
+    "forward_blacklist": set(),
+    "composite_ops_record": set(),
+    "backward_blacklist": set(),
+}
+
+
+def _get_batch_norm_none_var(op):
+    """Some outputs of batch_norm's replaced composite rule are not needed and will be removed."""
+    use_run_stat = (
+        op.attr("is_test") and (not op.attr("trainable_statistics"))
+    ) or op.attr("use_global_stats")
+    if use_run_stat:
+        return ["ReserveSpace", "SavedMean", "SavedVariance"]
+    else:
+        return ["ReserveSpace"]
+
+
+# In some case, inputs and outputs of composite op or its replaced composite rule might be None.
+# It means such arg will be no longer required in processed program by composite mechanism.
+# Therefore, such special ops should be recorded in advance and be released in args check.
+ops_contain_none = {
+    "batch_norm": _get_batch_norm_none_var,
+    "flatten_contiguous_range": ["XShape"],
+    "squeeze2": ["XShape"],
+    "unsqueeze2": ["XShape"],
+}
+
+
+# some intermediate outputs like xshape will no longer used after decomp, but return none to keep output num the same as origin op
+# key is the name of op, and value is the index of output in op.outputs
+decomp_ops_contain_unused_output = {
+    "pd_op.squeeze": [1],
+    "pd_op.unsqueeze": [1],
+    "pd_op.batch_norm": [5],
+}
+
+
+# This api is used for development for dynamic shape in prim, and will be removed in future.
+def _enable_prim_skip_dynamic_shape():
+    from paddle.base.framework import get_flags
+
+    return get_flags("FLAGS_prim_skip_dynamic")["FLAGS_prim_skip_dynamic"]
+
+
+def _enable_prim_dynamic_shape():
+    from paddle.base.framework import get_flags
+
+    return get_flags("FLAGS_prim_enable_dynamic")["FLAGS_prim_enable_dynamic"]
+
+
+def _enable_dist_prim_all():
+    flag = os.getenv("FLAGS_dist_prim_all")
+    if flag and flag.lower() in ("1", "true"):
+        return True
+    else:
+        return False
+
+
+def _enable_auto_recompute():
+    flag = os.getenv("FLAGS_enable_auto_recompute")
+
+    # NOTE(chenxi67): open recompute when cinn is enabled
+    from paddle.base.framework import in_cinn_mode
+
+    if in_cinn_mode():
+        if flag and flag.lower() in ("0", "false"):
+            return False
+        else:
+            return True
+    if flag and flag.lower() in ("1", "true"):
+        return True
+    else:
+        return False
+
+
+def _set_prim_forward_blacklist(*args):
+    for item in args:
+        if not isinstance(item, str):
+            raise TypeError("ops set in forward_blacklist must belong to str")
+        else:
+            prim_config["forward_blacklist"].add(item)
+
+
+# Currently, this function is not utilized anywhere in the codebase.
+# It may be intended for future use or could be removed if unnecessary.
+# def _reset_prim_forward_blacklist():
+#     prim_config["forward_blacklist"] = set()
+
+
+def _set_prim_backward_blacklist(*args):
+    ops = set(args)
+    new_ops = set()
+    for item in ops:
+        if not isinstance(item, str):
+            raise TypeError("All items in set must be strings.")
+        item = item.removeprefix("pd_op.")
+        prim_config["backward_blacklist"].add(item)
+        new_ops.add(item)
+    _set_bwd_prim_blacklist(new_ops)
+
+
+def _set_prim_backward_enabled(value: bool, print_flag: bool = False):
+    assert isinstance(value, bool), (
+        f"value should be bool, but got {type(value)}"
+    )
+    __set_bwd_prim_enabled(value)
+    if _prim_return_log() or print_flag:
+        print("backward prim enabled: ", bool(_is_bwd_prim_enabled()))
+
+
+def _set_prim_forward_enabled(value: bool, print_flag: bool = False):
+    assert isinstance(value, bool), (
+        f"value should be bool, but got {type(value)}"
+    )
+    __set_fwd_prim_enabled(value)
+    if _prim_return_log() or print_flag:
+        print("forward prim enabled: ", bool(_is_fwd_prim_enabled()))
+
+
+def set_prim_eager_enabled(value: bool, print_flag: bool = False):
+    assert isinstance(value, bool), (
+        f"value should be bool, but got {type(value)}"
+    )
+    __set_eager_prim_enabled(value)
+    if _prim_return_log() or print_flag:
+        print("eager prim enabled: ", bool(_is_eager_prim_enabled()))
+
+
+def _set_prim_all_enabled(value: bool, print_flag: bool = False):
+    assert isinstance(value, bool), (
+        f"value should be bool, but got {type(value)}"
+    )
+    __set_all_prim_enabled(value)
+    if _prim_return_log() or print_flag:
+        print(
+            "all prim enabled: ",
+            bool(_is_all_prim_enabled()),
+        )
+
+
+def __check_and_set_prim_all_enabled(print_flag=False):
+    from paddle.utils.environments import strtobool
+
+    prim_all_env = os.getenv("FLAGS_prim_all")
+    prim_fwd_env = os.getenv("FLAGS_prim_forward")
+    prim_bwd_env = os.getenv("FLAGS_prim_backward")
+    if prim_all_env is not None:
+        prim_all_flag = strtobool(prim_all_env)
+        _set_prim_all_enabled(prim_all_flag, print_flag)
+
+    if prim_fwd_env is not None:
+        prim_fwd_flag = strtobool(prim_fwd_env)
+        _set_prim_forward_enabled(prim_fwd_flag, print_flag)
+
+    if prim_bwd_env is not None:
+        prim_bwd_flag = strtobool(prim_bwd_env)
+        _set_prim_backward_enabled(prim_bwd_flag, print_flag)
+
+
+__check_and_set_prim_all_enabled(print_flag=True)
+
+
+SKIPPED_PRIM_VJP_DEFAULT_OPS = ["matmul_grad"]
+
+
+def _clear_prim_vjp_skip_default_ops():
+    for item in SKIPPED_PRIM_VJP_DEFAULT_OPS:
+        _remove_skip_comp_ops(item)
+
+
+# Since some decomposition of special ops like matmul_grad will reduce performance and is difficult to optimize currently by CINN.
+# This api is used for development for in prim and cinn, and will be removed in future.
+def _check_and_set_prim_vjp_skip_default_ops():
+    flag = os.getenv("FLAGS_prim_vjp_skip_default_ops", "1")
+    if flag and flag.lower() in ("1", "true"):
+        _set_prim_backward_blacklist(*SKIPPED_PRIM_VJP_DEFAULT_OPS)
+        return True
+    else:
+        _clear_prim_vjp_skip_default_ops()
+        return False
+
+
+_check_and_set_prim_vjp_skip_default_ops()
+
+
+def _check_prim_vjp_ops():
+    ops_org = os.getenv("FLAGS_prim_backward_blacklist", "")
+    if ops_org:
+        ops = []
+        for item in ops_org.split(";"):
+            ops.append(item.strip())
+        _set_prim_backward_blacklist(*ops)
+
+
+_check_prim_vjp_ops()

python/paddle/base/core.pyi

@@ -0,0 +1,15 @@
+# Copyright (c) 2024 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 paddle._typing.libs.libpaddle import *  # noqa: F403

python/paddle/base/data_feed_desc.py

@@ -0,0 +1,252 @@
+#   Copyright (c) 2018 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 google.protobuf import text_format
+
+from paddle.base.proto import data_feed_pb2
+
+__all__ = []
+
+
+class DataFeedDesc:
+    r"""
+    :api_attr: Static Graph
+
+    Datafeed descriptor, describing input training data format.
+
+    DataFeedDesc shall be initialized from a valid protobuf message from disk.
+
+    See :code:`paddle/base/framework/data_feed.proto` for message definition.
+    A typical message might look like:
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle.base as base
+            >>> with open("data.proto", "w") as f:
+            ...     f.write('name: "MultiSlotDataFeed"\n')
+            ...     f.write('batch_size: 2\n')
+            ...     f.write('multi_slot_desc {\n')
+            ...     f.write('    slots {\n')
+            ...     f.write('        name: "words"\n')
+            ...     f.write('        type: "uint64"\n')
+            ...     f.write('        is_dense: false\n')
+            ...     f.write('        is_used: true\n')
+            ...     f.write('    }\n')
+            ...     f.write('    slots {\n')
+            ...     f.write('        name: "label"\n')
+            ...     f.write('        type: "uint64"\n')
+            ...     f.write('        is_dense: false\n')
+            ...     f.write('        is_used: true\n')
+            ...     f.write('    }\n')
+            ...     f.write('}')
+            >>> data_feed = base.DataFeedDesc('data.proto')
+
+        However, users usually shouldn't care about the message format; instead,
+        they are encouraged to use :code:`Data Generator` as a tool to generate a
+        valid data description, in the process of converting their raw log files to
+        training files acceptable to Executor.
+
+        DataFeedDesc can also be changed during runtime. Once you got familiar with
+        what each field mean, you can modify it to better suit your need. E.g.:
+
+        .. code-block:: pycon
+
+            >>> import paddle.base as base
+            >>> data_feed = base.DataFeedDesc('data.proto')
+            >>> data_feed.set_batch_size(128)
+            >>> data_feed.set_dense_slots(['words'])  # The slot named 'words' will be dense
+            >>> data_feed.set_use_slots(['words'])  # The slot named 'words' will be used
+
+            >>> # Finally, the content can be dumped out for debugging purpose:
+
+            >>> print(data_feed.desc())
+
+    Args:
+        proto_file(string): Disk file containing a data feed description.
+
+    """
+
+    def __init__(self, proto_file):
+        self.proto_desc = data_feed_pb2.DataFeedDesc()
+        self.proto_desc.pipe_command = "cat"
+        with open(proto_file, 'r') as f:
+            text_format.Parse(f.read(), self.proto_desc)
+        if self.proto_desc.name == "MultiSlotDataFeed":
+            self.__name_to_index = {
+                slot.name: i
+                for i, slot in enumerate(self.proto_desc.multi_slot_desc.slots)
+            }
+
+    def set_batch_size(self, batch_size):
+        r"""
+        Set :attr:`batch_size` in ``paddle.base.DataFeedDesc`` . :attr:`batch_size` can be changed during training.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle.base as base
+                >>> with open("data.proto", "w") as f:
+                ...     f.write('name: "MultiSlotDataFeed"\n')
+                ...     f.write('batch_size: 2\n')
+                ...     f.write('multi_slot_desc {\n')
+                ...     f.write('    slots {\n')
+                ...     f.write('        name: "words"\n')
+                ...     f.write('        type: "uint64"\n')
+                ...     f.write('        is_dense: false\n')
+                ...     f.write('        is_used: true\n')
+                ...     f.write('    }\n')
+                ...     f.write('    slots {\n')
+                ...     f.write('        name: "label"\n')
+                ...     f.write('        type: "uint64"\n')
+                ...     f.write('        is_dense: false\n')
+                ...     f.write('        is_used: true\n')
+                ...     f.write('    }\n')
+                ...     f.write('}')
+                >>> data_feed = base.DataFeedDesc('data.proto')
+                >>> data_feed.set_batch_size(128)
+
+        Args:
+            batch_size (int): The number of batch size.
+
+        Returns:
+            None.
+
+        """
+        self.proto_desc.batch_size = batch_size
+
+    def set_dense_slots(self, dense_slots_name):
+        r"""
+        Set slots in :attr:`dense_slots_name` as dense slots. **Note: In default, all slots are sparse slots.**
+
+        Features for a dense slot will be fed into a Tensor, while those for a
+        sparse slot will be fed into a DenseTensor.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle.base as base
+                >>> with open("data.proto", "w") as f:
+                ...     f.write('name: "MultiSlotDataFeed"\n')
+                ...     f.write('batch_size: 2\n')
+                ...     f.write('multi_slot_desc {\n')
+                ...     f.write('    slots {\n')
+                ...     f.write('        name: "words"\n')
+                ...     f.write('        type: "uint64"\n')
+                ...     f.write('        is_dense: false\n')
+                ...     f.write('        is_used: true\n')
+                ...     f.write('    }\n')
+                ...     f.write('    slots {\n')
+                ...     f.write('        name: "label"\n')
+                ...     f.write('        type: "uint64"\n')
+                ...     f.write('        is_dense: false\n')
+                ...     f.write('        is_used: true\n')
+                ...     f.write('    }\n')
+                ...     f.write('}')
+                >>> data_feed = base.DataFeedDesc('data.proto')
+                >>> data_feed.set_dense_slots(['words'])
+
+        Args:
+            dense_slots_name (list(str)): a list of slot names which will be set dense.
+
+        Returns:
+            None.
+
+        """
+        if self.proto_desc.name != "MultiSlotDataFeed":
+            raise ValueError(
+                "Only MultiSlotDataFeed needs set_dense_slots, please check your datafeed.proto"
+            )
+        for name in dense_slots_name:
+            self.proto_desc.multi_slot_desc.slots[
+                self.__name_to_index[name]
+            ].is_dense = True
+
+    def set_use_slots(self, use_slots_name):
+        r"""
+        Set if a specific slot will be used for training. A dataset shall
+        contain a lot of features, through this function one can select which
+        ones will be used for a specific model.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle.base as base
+                >>> with open("data.proto", "w") as f:
+                ...     f.write('name: "MultiSlotDataFeed"\n')
+                ...     f.write('batch_size: 2\n')
+                ...     f.write('multi_slot_desc {\n')
+                ...     f.write('    slots {\n')
+                ...     f.write('        name: "words"\n')
+                ...     f.write('        type: "uint64"\n')
+                ...     f.write('        is_dense: false\n')
+                ...     f.write('        is_used: true\n')
+                ...     f.write('    }\n')
+                ...     f.write('    slots {\n')
+                ...     f.write('        name: "label"\n')
+                ...     f.write('        type: "uint64"\n')
+                ...     f.write('        is_dense: false\n')
+                ...     f.write('        is_used: true\n')
+                ...     f.write('    }\n')
+                ...     f.write('}')
+                >>> data_feed = base.DataFeedDesc('data.proto')
+                >>> data_feed.set_use_slots(['words'])
+
+        Args:
+            use_slots_name: a list of slot names which will be used in training
+
+        Note:
+            Default is not used for all slots
+        """
+        if self.proto_desc.name != "MultiSlotDataFeed":
+            raise ValueError(
+                "Only MultiSlotDataFeed needs set_use_slots, please check your datafeed.proto"
+            )
+        for name in use_slots_name:
+            self.proto_desc.multi_slot_desc.slots[
+                self.__name_to_index[name]
+            ].is_used = True
+
+    def desc(self):
+        r"""
+        Returns a protobuf message for this DataFeedDesc
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle.base as base
+                >>> with open("data.proto", "w") as f:
+                ...     f.write('name: "MultiSlotDataFeed"\n')
+                ...     f.write('batch_size: 2\n')
+                ...     f.write('multi_slot_desc {\n')
+                ...     f.write('    slots {\n')
+                ...     f.write('        name: "words"\n')
+                ...     f.write('        type: "uint64"\n')
+                ...     f.write('        is_dense: false\n')
+                ...     f.write('        is_used: true\n')
+                ...     f.write('    }\n')
+                ...     f.write('    slots {\n')
+                ...     f.write('        name: "label"\n')
+                ...     f.write('        type: "uint64"\n')
+                ...     f.write('        is_dense: false\n')
+                ...     f.write('        is_used: true\n')
+                ...     f.write('    }\n')
+                ...     f.write('}')
+                >>> data_feed = base.DataFeedDesc('data.proto')
+                >>> print(data_feed.desc())
+
+        Returns:
+            A string message
+        """
+        return text_format.MessageToString(self.proto_desc)