introspect.py 2.7 KB

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