A thread local cache for _ThreadLocalExtraJitContext The extra_jit_context in xla_jit.thread_local_state() may get updated and thus incurring dispatch overhead for comparing this python object during jit calls. We want to duduplicate the objects that have the same hash/equality to also
| 624 | |
| 625 | |
| 626 | class _ThreadLocalStateCache(threading.local): |
| 627 | """"A thread local cache for _ThreadLocalExtraJitContext |
| 628 | |
| 629 | The extra_jit_context in xla_jit.thread_local_state() may get updated and thus |
| 630 | incurring dispatch overhead for comparing this python object during jit calls. |
| 631 | We want to duduplicate the objects that have the same hash/equality to also |
| 632 | have the same object ID, since the equality check is much faster if the object |
| 633 | IDs match. |
| 634 | """ |
| 635 | |
| 636 | def __init__(self): |
| 637 | self.canonicalize = functools.lru_cache(128)(lambda x: x) |
| 638 | |
| 639 | |
| 640 | _thread_local_state_cache = _ThreadLocalStateCache() |