Skip to content

Commit 76fac7b

Browse files
authored
gh-104584: Clean up and fix uops tests and fix crash (#106492)
The uops test wasn't testing anything by default, and was failing when run with -Xuops. Made the two executor-related context managers global, so TestUops can use them (notably `with temporary_optimizer(opt)`). Made clear_executor() a little more thorough. Fixed a crash upon finalizing a uop optimizer, by adding a `tp_dealloc` handler.
1 parent 67a7988 commit 76fac7b

File tree

2 files changed

+38
-26
lines changed

2 files changed

+38
-26
lines changed

Lib/test/test_capi/test_misc.py

Lines changed: 30 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2343,23 +2343,26 @@ def func():
23432343
names = ["func", "outer", "outer", "inner", "inner", "outer", "inner"]
23442344
self.do_test(func, names)
23452345

2346-
class TestOptimizerAPI(unittest.TestCase):
23472346

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)
23552354

2356-
@contextlib.contextmanager
2357-
def clear_executors(self, func):
2358-
try:
2359-
yield
2360-
finally:
2361-
#Clear executors
2362-
func.__code__ = func.__code__.replace()
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()
2363+
2364+
2365+
class TestOptimizerAPI(unittest.TestCase):
23632366

23642367
def test_get_set_optimizer(self):
23652368
self.assertEqual(_testinternalcapi.get_optimizer(), None)
@@ -2381,9 +2384,9 @@ def loop():
23812384

23822385
for repeat in range(5):
23832386
opt = _testinternalcapi.get_counter_optimizer()
2384-
with self.temporary_optimizer(opt):
2387+
with temporary_optimizer(opt):
23852388
self.assertEqual(opt.get_count(), 0)
2386-
with self.clear_executors(loop):
2389+
with clear_executors(loop):
23872390
loop()
23882391
self.assertEqual(opt.get_count(), 1000)
23892392

@@ -2409,7 +2412,7 @@ def long_loop():
24092412
long_loop = ns['long_loop']
24102413

24112414
opt = _testinternalcapi.get_counter_optimizer()
2412-
with self.temporary_optimizer(opt):
2415+
with temporary_optimizer(opt):
24132416
self.assertEqual(opt.get_count(), 0)
24142417
long_loop()
24152418
self.assertEqual(opt.get_count(), 10)
@@ -2418,24 +2421,27 @@ def long_loop():
24182421
class TestUops(unittest.TestCase):
24192422

24202423
def test_basic_loop(self):
2421-
24222424
def testfunc(x):
24232425
i = 0
24242426
while i < x:
24252427
i += 1
24262428

2427-
testfunc(1000)
2429+
opt = _testinternalcapi.get_uop_optimizer()
2430+
2431+
with temporary_optimizer(opt):
2432+
testfunc(1000)
24282433

24292434
ex = None
2430-
for offset in range(0, 100, 2):
2435+
for offset in range(0, len(testfunc.__code__.co_code), 2):
24312436
try:
24322437
ex = _testinternalcapi.get_executor(testfunc.__code__, offset)
24332438
break
24342439
except ValueError:
24352440
pass
2436-
if ex is None:
2437-
return
2438-
self.assertIn("SAVE_IP", str(ex))
2441+
self.assertIsNotNone(ex)
2442+
uops = {opname for opname, _ in ex}
2443+
self.assertIn("SAVE_IP", uops)
2444+
self.assertIn("LOAD_FAST", uops)
24392445

24402446

24412447
if __name__ == "__main__":

Python/optimizer.c

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -532,7 +532,7 @@ uop_optimize(
532532
return trace_length;
533533
}
534534
OBJECT_STAT_INC(optimization_traces_created);
535-
_PyUOpExecutorObject *executor = (_PyUOpExecutorObject *)_PyObject_New(&UOpExecutor_Type);
535+
_PyUOpExecutorObject *executor = PyObject_New(_PyUOpExecutorObject, &UOpExecutor_Type);
536536
if (executor == NULL) {
537537
return -1;
538538
}
@@ -542,18 +542,24 @@ uop_optimize(
542542
return 1;
543543
}
544544

545+
static void
546+
uop_opt_dealloc(PyObject *self) {
547+
PyObject_Free(self);
548+
}
549+
545550
static PyTypeObject UOpOptimizer_Type = {
546551
PyVarObject_HEAD_INIT(&PyType_Type, 0)
547552
.tp_name = "uop_optimizer",
548553
.tp_basicsize = sizeof(_PyOptimizerObject),
549554
.tp_itemsize = 0,
550555
.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_DISALLOW_INSTANTIATION,
556+
.tp_dealloc = uop_opt_dealloc,
551557
};
552558

553559
PyObject *
554560
PyUnstable_Optimizer_NewUOpOptimizer(void)
555561
{
556-
_PyOptimizerObject *opt = (_PyOptimizerObject *)_PyObject_New(&UOpOptimizer_Type);
562+
_PyOptimizerObject *opt = PyObject_New(_PyOptimizerObject, &UOpOptimizer_Type);
557563
if (opt == NULL) {
558564
return NULL;
559565
}

0 commit comments

Comments
 (0)