codekingpro/portable-devtools

1"""
2Introspection helper functions.
3"""
4
5__all__ = ['opt_func_info']
6
7
8def opt_func_info(func_name=None, signature=None):
9    """
10    Returns a dictionary containing the currently supported CPU dispatched
11    features for all optimized functions.
12
13    Parameters
14    ----------
15    func_name : str (optional)
16        Regular expression to filter by function name.
17
18    signature : str (optional)
19        Regular expression to filter by data type.
20
21    Returns
22    -------
23    dict
24        A dictionary where keys are optimized function names and values are
25        nested dictionaries indicating supported targets based on data types.
26
27    Examples
28    --------
29    Retrieve dispatch information for functions named 'add' or 'sub' and
30    data types 'float64' or 'float32':
31
32    >>> import numpy as np
33    >>> dict = np.lib.introspect.opt_func_info(
34    ...     func_name="add|abs", signature="float64|complex64"
35    ... )
36    >>> import json
37    >>> print(json.dumps(dict, indent=2))   # may vary (architecture)
38        {
39          "absolute": {
40            "dd": {
41              "current": "SSE41",
42              "available": "SSE41 baseline(SSE SSE2 SSE3)"
43            },
44            "Ff": {
45              "current": "FMA3__AVX2",
46              "available": "AVX512F FMA3__AVX2 baseline(SSE SSE2 SSE3)"
47            },
48            "Dd": {
49              "current": "FMA3__AVX2",
50              "available": "AVX512F FMA3__AVX2 baseline(SSE SSE2 SSE3)"
51            }
52          },
53          "add": {
54            "ddd": {
55              "current": "FMA3__AVX2",
56              "available": "FMA3__AVX2 baseline(SSE SSE2 SSE3)"
57            },
58            "FFF": {
59              "current": "FMA3__AVX2",
60              "available": "FMA3__AVX2 baseline(SSE SSE2 SSE3)"
61            }
62          }
63        }
64
65    """
66    import re
67
68    from numpy._core._multiarray_umath import __cpu_targets_info__ as targets, dtype
69
70    if func_name is not None:
71        func_pattern = re.compile(func_name)
72        matching_funcs = {
73            k: v for k, v in targets.items()
74            if func_pattern.search(k)
75        }
76    else:
77        matching_funcs = targets
78
79    if signature is not None:
80        sig_pattern = re.compile(signature)
81        matching_sigs = {}
82        for k, v in matching_funcs.items():
83            matching_chars = {}
84            for chars, targets in v.items():
85                if any(
86                    sig_pattern.search(c) or sig_pattern.search(dtype(c).name)
87                    for c in chars
88                ):
89                    matching_chars[chars] = targets  # noqa: PERF403
90            if matching_chars:
91                matching_sigs[k] = matching_chars
92    else:
93        matching_sigs = matching_funcs
94    return matching_sigs
95