chore: import upstream snapshot with attribution

python/paddle/static/amp/function_overload.py

@@ -0,0 +1,142 @@
+# 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 implementation refers to https://arpitbhayani.me/blogs/function-overloading.
+# Note: it is customed for paddle.static.amp.decorate function.
+
+import inspect
+import logging
+from enum import Enum
+
+from paddle.base.log_helper import get_logger
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+
+class FunctionType(Enum):
+    FP16_ONLY = 0
+    COMMON = 1
+
+
+class Function:
+    """
+    Function is a wrap over standard python function
+    An instance of this Function class is also callable
+    just like the python function that it wrapped.
+    When the instance is "called" like a function it fetches
+    the function to be invoked from the virtual namespace and then
+    invokes the same.
+    """
+
+    def __init__(self, fn):
+        self.fn = fn
+
+    def __call__(self, *args, **kwargs):
+        """
+        Overriding the __call__ function which makes the
+        instance callable.
+        """
+        # fetching the function to be invoked from the virtual namespace
+        # through the arguments.
+        fn = Namespace.get_instance().get(*args, **kwargs)
+        # invoking the wrapped function and returning the value.
+        return fn(*args, **kwargs)
+
+
+class Namespace:
+    """
+    Namespace is the singleton class that is responsible
+    for holding all the functions.
+    """
+
+    __instance = None
+
+    def __init__(self):
+        if self.__instance is None:
+            self.function_map = {}
+            Namespace.__instance = self
+        else:
+            raise Exception("cannot instantiate Namespace again.")
+
+    @staticmethod
+    def get_instance():
+        if Namespace.__instance is None:
+            Namespace()
+        return Namespace.__instance
+
+    def register(self, fn, key):
+        """
+        Register the function in the virtual namespace and return
+        an instance of callable Function that wraps the function fn.
+
+        Args:
+            fn (function): the native python function handle.
+            key (FunctionType): the specified type.
+        """
+        assert isinstance(key, FunctionType), (
+            f"The type of  key is expected to be FunctionType, but received {type(key)}."
+        )
+        func = Function(fn)
+        self.function_map[key] = fn
+        return func
+
+    def get(self, *args, **kwargs):
+        """
+        Get the matching function from the virtual namespace according to the actual arguments.
+        Return None if it did not find any matching function.
+        """
+        _logger.debug(f"get function: args={args}, kwargs={kwargs}")
+        satisfied_function_keys = set(self.function_map.keys())
+        num_actual_args = len(args) + len(kwargs)
+        for func_key in self.function_map.keys():
+            if func_key not in satisfied_function_keys:
+                continue
+            fn = self.function_map[func_key]
+            specs = inspect.getfullargspec(fn)
+            if len(specs) < len(args) + len(kwargs):
+                # Remove the not satisfied function according to the number of actual arguments.
+                _logger.debug(
+                    f"fn={fn} (key={func_key}) is not satisfied and removed."
+                )
+                satisfied_function_keys.remove(func_key)
+                continue
+            if len(kwargs) > 0:
+                # Remove the not satisfied function according to argument keys in kwargs.
+                for arg_name, value in kwargs.items():
+                    if arg_name not in specs.args:
+                        _logger.debug(
+                            f"fn={fn} (key={func_key}) is not satisfied and removed."
+                        )
+                        satisfied_function_keys.remove(func_key)
+                        break
+        if len(satisfied_function_keys) == 1:
+            key = next(iter(satisfied_function_keys))
+        elif len(args) >= 3 and isinstance(args[2], float):
+            key = FunctionType.FP16_ONLY
+        else:
+            key = FunctionType.COMMON
+        return self.function_map.get(key)
+
+
+def overload(key):
+    """overload is the decorator that wraps the function
+    and returns a callable object of type Function.
+    """
+
+    def decorator(fn):
+        return Namespace.get_instance().register(fn, key)
+
+    return decorator