@@ -2343,23 +2343,26 @@ def func():
2343
2343
names = ["func" , "outer" , "outer" , "inner" , "inner" , "outer" , "inner" ]
2344
2344
self .do_test (func , names )
2345
2345
2346
- class TestOptimizerAPI (unittest .TestCase ):
2347
2346
2348
- @contextlib .contextmanager
2349
- def temporary_optimizer (self , opt ):
2350
- _testinternalcapi .set_optimizer (opt )
2351
- try :
2352
- yield
2353
- finally :
2354
- _testinternalcapi .set_optimizer (None )
2347
+ @contextlib .contextmanager
2348
+ def temporary_optimizer (opt ):
2349
+ _testinternalcapi .set_optimizer (opt )
2350
+ try :
2351
+ yield
2352
+ finally :
2353
+ _testinternalcapi .set_optimizer (None )
2354
+
2355
+ @contextlib .contextmanager
2356
+ def clear_executors (func ):
2357
+ # Clear executors in func before and after running a block
2358
+ func .__code__ = func .__code__ .replace ()
2359
+ try :
2360
+ yield
2361
+ finally :
2362
+ func .__code__ = func .__code__ .replace ()
2355
2363
2356
- @contextlib .contextmanager
2357
- def clear_executors (self , func ):
2358
- try :
2359
- yield
2360
- finally :
2361
- #Clear executors
2362
- func .__code__ = func .__code__ .replace ()
2364
+
2365
+ class TestOptimizerAPI (unittest .TestCase ):
2363
2366
2364
2367
def test_get_set_optimizer (self ):
2365
2368
self .assertEqual (_testinternalcapi .get_optimizer (), None )
@@ -2381,9 +2384,9 @@ def loop():
2381
2384
2382
2385
for repeat in range (5 ):
2383
2386
opt = _testinternalcapi .get_counter_optimizer ()
2384
- with self . temporary_optimizer (opt ):
2387
+ with temporary_optimizer (opt ):
2385
2388
self .assertEqual (opt .get_count (), 0 )
2386
- with self . clear_executors (loop ):
2389
+ with clear_executors (loop ):
2387
2390
loop ()
2388
2391
self .assertEqual (opt .get_count (), 1000 )
2389
2392
@@ -2409,11 +2412,37 @@ def long_loop():
2409
2412
long_loop = ns ['long_loop' ]
2410
2413
2411
2414
opt = _testinternalcapi .get_counter_optimizer ()
2412
- with self . temporary_optimizer (opt ):
2415
+ with temporary_optimizer (opt ):
2413
2416
self .assertEqual (opt .get_count (), 0 )
2414
2417
long_loop ()
2415
2418
self .assertEqual (opt .get_count (), 10 )
2416
2419
2417
2420
2421
+ class TestUops (unittest .TestCase ):
2422
+
2423
+ def test_basic_loop (self ):
2424
+ def testfunc (x ):
2425
+ i = 0
2426
+ while i < x :
2427
+ i += 1
2428
+
2429
+ opt = _testinternalcapi .get_uop_optimizer ()
2430
+
2431
+ with temporary_optimizer (opt ):
2432
+ testfunc (1000 )
2433
+
2434
+ ex = None
2435
+ for offset in range (0 , len (testfunc .__code__ .co_code ), 2 ):
2436
+ try :
2437
+ ex = _testinternalcapi .get_executor (testfunc .__code__ , offset )
2438
+ break
2439
+ except ValueError :
2440
+ pass
2441
+ self .assertIsNotNone (ex )
2442
+ uops = {opname for opname , _ in ex }
2443
+ self .assertIn ("SAVE_IP" , uops )
2444
+ self .assertIn ("LOAD_FAST" , uops )
2445
+
2446
+
2418
2447
if __name__ == "__main__" :
2419
2448
unittest .main ()
0 commit comments