Repository navigation
perf: hoist the randint dtype-bounds table out of the per-call hot pat
#199
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
72e5b1f
b92b39f
88c3bc5
151162e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -1794,6 +1794,22 @@ def _brng_id_to_name(int brng_id): | |
| return brng_name | ||
|
|
||
|
|
||
| # Largest half-open interval [low, high) per randint dtype. | ||
| _randint_bounds = { | ||
| np.dtype(name): (name, low, high) for name, low, high in ( | ||
| ("bool", 0, 2), | ||
| ("int8", -2**7, 2**7), | ||
| ("int16", -2**15, 2**15), | ||
| ("int32", -2**31, 2**31), | ||
| ("int64", -2**63, 2**63), | ||
| ("uint8", 0, 2**8), | ||
| ("uint16", 0, 2**16), | ||
| ("uint32", 0, 2**32), | ||
| ("uint64", 0, 2**64), | ||
| ) | ||
| } | ||
|
|
||
|
|
||
| cdef class _MKLRandomState: | ||
| cdef irk_state *internal_state | ||
| cdef object lock | ||
|
|
@@ -2143,35 +2159,18 @@ cdef class _MKLRandomState: | |
| self.internal_state, irk_double_vec, size, self.lock | ||
| ) | ||
|
|
||
| # Set up dictionary of integer types and relevant functions. | ||
| # | ||
| # The dictionary is keyed by dtype(...).name and the values | ||
| # are a tuple (low, high, function), where low and high are | ||
| # the bounds of the largest half open interval `[low, high)` | ||
| # and the function is the relevant function to call for | ||
| # that precision. | ||
| # | ||
| # The functions are all the same except for changed types in | ||
| # a few places. It would be easy to template them. | ||
| # Returns randint bounds and the matching `_rand_*`/`_rand_*_broadcast` | ||
| # methods for dtype. | ||
|
|
||
| def _choose_randint_type(self, dtype): | ||
| _randint_bounds = { | ||
| "bool": (0, 2), | ||
| "int8": (-2**7, 2**7), | ||
| "int16": (-2**15, 2**15), | ||
| "int32": (-2**31, 2**31), | ||
| "int64": (-2**63, 2**63), | ||
| "uint8": (0, 2**8), | ||
| "uint16": (0, 2**16), | ||
| "uint32": (0, 2**32), | ||
| "uint64": (0, 2**64), | ||
| } | ||
|
|
||
| key = np.dtype(dtype).name | ||
| if key not in _randint_bounds: | ||
| raise TypeError(f'Unsupported dtype "{key}" for randint') | ||
|
|
||
| lowbnd, highbnd = _randint_bounds[key] | ||
| _dtype = np.dtype(dtype) | ||
| try: | ||
| key, lowbnd, highbnd = _randint_bounds[_dtype] | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| except KeyError: | ||
| raise TypeError( | ||
| f'Unsupported dtype "{_dtype.name}" for randint' | ||
| ) from None | ||
|
|
||
| return (lowbnd, highbnd, | ||
| getattr(self, f"_rand_{key}"), | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. As a possible possible follow-up: each call still builds two method names with f-strings and looks them up with getattr, so storing the method names in the table would remove that. |
||
| getattr(self, f"_rand_{key}_broadcast")) | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
missing a changelog entry