def testFirstMostSpecificFunctionCacheKeyIsLookedUp(self):
ctx = function_cache.ExecutionContext(1, 1, 1, 1, 1, 1)
cache = function_cache.FunctionCache()
cache.add(function_cache.FunctionCacheKey(MockShape(1, 2, None), ctx),
trace_type.WeakrefDeletionObserver(), "a")
cache.add(function_cache.FunctionCacheKey(MockShape(1, None, 3), ctx),
trace_type.WeakrefDeletionObserver(), "b")
self.assertEqual(
cache.lookup(
function_cache.FunctionCacheKey(MockShape(1, 2, 3), ctx),
True), "a")
def testMostSpecificFunctionCacheKeyIsOrderAgnostic(self):
ctx = function_cache.FunctionContext(0)
keys = [(function_cache.FunctionCacheKey(MockShape(1, 1, 1),
MockEmptyCaptureSnapshot(),
ctx), "a"),
(function_cache.FunctionCacheKey(MockShape(1, None, 1),
MockEmptyCaptureSnapshot(),
ctx), "b"),
(function_cache.FunctionCacheKey(MockShape(None, None, 1),
MockEmptyCaptureSnapshot(),
ctx), "c"),
(function_cache.FunctionCacheKey(MockShape(None, None, None),
MockEmptyCaptureSnapshot(),
ctx), "d")]
for permutation in itertools.permutations(keys):
cache = function_cache.FunctionCache()
cache.add(permutation[0][0], trace_type.WeakrefDeletionObserver(),
permutation[0][1])
cache.add(permutation[1][0], trace_type.WeakrefDeletionObserver(),
permutation[1][1])
cache.add(permutation[2][0], trace_type.WeakrefDeletionObserver(),
permutation[2][1])
cache.add(permutation[3][0], trace_type.WeakrefDeletionObserver(),
permutation[3][1])
self.assertEqual(
cache.lookup(
function_cache.FunctionCacheKey(MockShape(1, 1, 1),
MockEmptyCaptureSnapshot(),
ctx), True),
"a")
self.assertEqual(
cache.lookup(
function_cache.FunctionCacheKey(MockShape(1, 2, 1),
MockEmptyCaptureSnapshot(),
ctx), True),
"b")
self.assertEqual(
cache.lookup(
function_cache.FunctionCacheKey(MockShape(2, 2, 1),
MockEmptyCaptureSnapshot(),
ctx), True),
"c")
self.assertEqual(
cache.lookup(
function_cache.FunctionCacheKey(MockShape(2, 2, 2),
MockEmptyCaptureSnapshot(),
ctx), True),
"d")
def setup():
cache.clear()
for key in keys:
cache.add(*key, "testing")
cache.add(
function_cache.FunctionCacheKey(MockSubtypeOf2(3),
MockEmptyCaptureSnapshot(), None),
trace_type.WeakrefDeletionObserver(), "testing")
def testMostSpecificFunctionCacheKeyIsLookedUp(self):
ctx = function_cache.FunctionContext(0)
cache = function_cache.FunctionCache()
cache.add(
function_cache.FunctionCacheKey(MockShape(1, 2, None),
MockEmptyCaptureSnapshot(), ctx),
trace_type.WeakrefDeletionObserver(), "a")
cache.add(
function_cache.FunctionCacheKey(MockShape(1, 2, 3),
MockEmptyCaptureSnapshot(), ctx),
trace_type.WeakrefDeletionObserver(), "b")
self.assertEqual(
cache.lookup(
function_cache.FunctionCacheKey(MockShape(1, 2, 3),
MockEmptyCaptureSnapshot(),
ctx), True),
"b")
def testDeleteRemovesConcreteFunctions(self):
cache = function_cache.FunctionCache()
key_1, deletion_observer_1 = function_cache.make_cache_key(1)
cache.add(key_1, deletion_observer_1, "test_1")
self.assertEqual(cache.lookup(key_1, False), "test_1")
cache.delete(key_1)
self.assertIsNone(cache.lookup(key_1, False))
key_2 = function_cache.FunctionCacheKey(MockSubtypeOf2(2), None)
cache.add(key_2, trace_type.WeakrefDeletionObserver(), "test_2")
self.assertEqual(cache.lookup(key_2, False), "test_2")
key_3 = function_cache.FunctionCacheKey(MockSubtypeOf2(3), None)
self.assertEqual(cache.lookup(key_3, True), "test_2")
cache.delete(key_2)
self.assertIsNone(cache.lookup(key_2, False))
self.assertIsNone(cache.lookup(key_3, True))
def benchmarkCacheHit50thKeyKnownSubtype(self):
# If there are 50 keys and we get a key that has a subtype in cache and
# the cache has observed the key before (to memorize the subtype).
cache = function_cache.FunctionCache()
args_per_call = 5
num_total_checks = 50
keys = []
for i in range(num_total_checks - 1):
args = []
for j in range(args_per_call):
args.append(array_ops.zeros([i, j]))
keys.append(function_context.make_cache_key(args))
for key in keys:
cache.add(*key, "testing")
cache.add(
function_cache.FunctionCacheKey(MockSubtypeOf2(2),
MockEmptyCaptureSnapshot(), None),
trace_type.WeakrefDeletionObserver(), "testing")
cache.lookup(function_cache.FunctionCacheKey(MockSubtypeOf2(3),
MockEmptyCaptureSnapshot(),
None), True)
iterations = 10000
lookup_key = function_cache.FunctionCacheKey(MockSubtypeOf2(2),
MockEmptyCaptureSnapshot(),
None)
subtyping_time = timeit.timeit(
lambda: cache.lookup(lookup_key, True), number=iterations)
self.report_benchmark(
name="cache_hit_50th_key_known_subtype",
iters=iterations,
wall_time=subtyping_time,
metrics=[{
"name": "cache_hit_50th_key_known_subtype_avg_ms",
"value": subtyping_time / iterations * 1000
}])