@@ -561,6 +561,16 @@ def testfunc(n):
561561
562562class TestUopsOptimization (unittest .TestCase ):
563563
564+ def _run_with_optimizer (self , testfunc , arg ):
565+ res = None
566+ opt = _testinternalcapi .get_uop_optimizer ()
567+ with temporary_optimizer (opt ):
568+ res = testfunc (arg )
569+
570+ ex = get_first_executor (testfunc )
571+ return res , ex
572+
573+
564574 def test_int_type_propagation (self ):
565575 def testfunc (loops ):
566576 num = 0
@@ -570,12 +580,7 @@ def testfunc(loops):
570580 num += 1
571581 return a
572582
573- opt = _testinternalcapi .get_uop_optimizer ()
574- res = None
575- with temporary_optimizer (opt ):
576- res = testfunc (32 )
577-
578- ex = get_first_executor (testfunc )
583+ res , ex = self ._run_with_optimizer (testfunc , 32 )
579584 self .assertIsNotNone (ex )
580585 self .assertEqual (res , 63 )
581586 binop_count = [opname for opname , _ , _ in ex if opname == "_BINARY_OP_ADD_INT" ]
@@ -642,12 +647,7 @@ def testfunc(loops):
642647 num += 1
643648 return a
644649
645- opt = _testinternalcapi .get_uop_optimizer ()
646- res = None
647- with temporary_optimizer (opt ):
648- res = testfunc (64 )
649-
650- ex = get_first_executor (testfunc )
650+ res , ex = self ._run_with_optimizer (testfunc , 64 )
651651 self .assertIsNotNone (ex )
652652 binop_count = [opname for opname , _ , _ in ex if opname == "_BINARY_OP_ADD_INT" ]
653653 self .assertGreaterEqual (len (binop_count ), 3 )
@@ -659,11 +659,7 @@ def dummy(x):
659659 for i in range (n ):
660660 dummy (i )
661661
662- opt = _testinternalcapi .get_uop_optimizer ()
663- with temporary_optimizer (opt ):
664- testfunc (32 )
665-
666- ex = get_first_executor (testfunc )
662+ res , ex = self ._run_with_optimizer (testfunc , 32 )
667663 self .assertIsNotNone (ex )
668664 uops = {opname for opname , _ , _ in ex }
669665 self .assertIn ("_PUSH_FRAME" , uops )
@@ -677,11 +673,7 @@ def testfunc(n):
677673 x = i + i
678674 return x
679675
680- opt = _testinternalcapi .get_uop_optimizer ()
681- with temporary_optimizer (opt ):
682- res = testfunc (32 )
683-
684- ex = get_first_executor (testfunc )
676+ res , ex = self ._run_with_optimizer (testfunc , 32 )
685677 self .assertEqual (res , 62 )
686678 self .assertIsNotNone (ex )
687679 uops = {opname for opname , _ , _ in ex }
@@ -699,11 +691,7 @@ def testfunc(n):
699691 res = x + z + a + b
700692 return res
701693
702- opt = _testinternalcapi .get_uop_optimizer ()
703- with temporary_optimizer (opt ):
704- res = testfunc (32 )
705-
706- ex = get_first_executor (testfunc )
694+ res , ex = self ._run_with_optimizer (testfunc , 32 )
707695 self .assertEqual (res , 4 )
708696 self .assertIsNotNone (ex )
709697 uops = {opname for opname , _ , _ in ex }
@@ -716,11 +704,8 @@ def testfunc(n):
716704 for _ in range (n ):
717705 return [i for i in range (n )]
718706
719- opt = _testinternalcapi .get_uop_optimizer ()
720- with temporary_optimizer (opt ):
721- testfunc (32 )
722-
723- ex = get_first_executor (testfunc )
707+ res , ex = self ._run_with_optimizer (testfunc , 32 )
708+ self .assertEqual (res , list (range (32 )))
724709 self .assertIsNotNone (ex )
725710 uops = {opname for opname , _ , _ in ex }
726711 self .assertNotIn ("_BINARY_OP_ADD_INT" , uops )
@@ -785,6 +770,56 @@ def testfunc(n):
785770 """ ))
786771 self .assertEqual (result [0 ].rc , 0 , result )
787772
773+ def test_float_add_constant_propagation (self ):
774+ def testfunc (n ):
775+ a = 1.0
776+ for _ in range (n ):
777+ a = a + 0.1
778+ return a
779+
780+ res , ex = self ._run_with_optimizer (testfunc , 32 )
781+ self .assertAlmostEqual (res , 4.2 )
782+ self .assertIsNotNone (ex )
783+ uops = {opname for opname , _ , _ in ex }
784+ guard_both_float_count = [opname for opname , _ , _ in ex if opname == "_GUARD_BOTH_FLOAT" ]
785+ self .assertLessEqual (len (guard_both_float_count ), 1 )
786+ # TODO gh-115506: this assertion may change after propagating constants.
787+ # We'll also need to verify that propagation actually occurs.
788+ self .assertIn ("_BINARY_OP_ADD_FLOAT" , uops )
789+
790+ def test_float_subtract_constant_propagation (self ):
791+ def testfunc (n ):
792+ a = 1.0
793+ for _ in range (n ):
794+ a = a - 0.1
795+ return a
796+
797+ res , ex = self ._run_with_optimizer (testfunc , 32 )
798+ self .assertAlmostEqual (res , - 2.2 )
799+ self .assertIsNotNone (ex )
800+ uops = {opname for opname , _ , _ in ex }
801+ guard_both_float_count = [opname for opname , _ , _ in ex if opname == "_GUARD_BOTH_FLOAT" ]
802+ self .assertLessEqual (len (guard_both_float_count ), 1 )
803+ # TODO gh-115506: this assertion may change after propagating constants.
804+ # We'll also need to verify that propagation actually occurs.
805+ self .assertIn ("_BINARY_OP_SUBTRACT_FLOAT" , uops )
806+
807+ def test_float_multiply_constant_propagation (self ):
808+ def testfunc (n ):
809+ a = 1.0
810+ for _ in range (n ):
811+ a = a * 2.0
812+ return a
813+
814+ res , ex = self ._run_with_optimizer (testfunc , 32 )
815+ self .assertAlmostEqual (res , 2 ** 32 )
816+ self .assertIsNotNone (ex )
817+ uops = {opname for opname , _ , _ in ex }
818+ guard_both_float_count = [opname for opname , _ , _ in ex if opname == "_GUARD_BOTH_FLOAT" ]
819+ self .assertLessEqual (len (guard_both_float_count ), 1 )
820+ # TODO gh-115506: this assertion may change after propagating constants.
821+ # We'll also need to verify that propagation actually occurs.
822+ self .assertIn ("_BINARY_OP_MULTIPLY_FLOAT" , uops )
788823
789824
790825if __name__ == "__main__" :
0 commit comments