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