gh-115480: Type / constant propagation for float binary uops (GH-115550)

Co-authored-by: Ken Jin <kenjin@python.org>

3 files changed, 184 insertions, 38 deletions

diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py
index 1a8ed34..66860c6 100644
--- a/Lib/test/test_capi/test_opt.py
+++ b/Lib/test/test_capi/test_opt.py
@@ -561,6 +561,16 @@ class TestUops(unittest.TestCase):
 
 class TestUopsOptimization(unittest.TestCase):
 
+    def _run_with_optimizer(self, testfunc, arg):
+        res = None
+        opt = _testinternalcapi.get_uop_optimizer()
+        with temporary_optimizer(opt):
+            res = testfunc(arg)
+
+        ex = get_first_executor(testfunc)
+        return res, ex
+
+
     def test_int_type_propagation(self):
         def testfunc(loops):
             num = 0
@@ -570,12 +580,7 @@ class TestUopsOptimization(unittest.TestCase):
                 num += 1
             return a
 
-        opt = _testinternalcapi.get_uop_optimizer()
-        res = None
-        with temporary_optimizer(opt):
-            res = testfunc(32)
-
-        ex = get_first_executor(testfunc)
+        res, ex = self._run_with_optimizer(testfunc, 32)
         self.assertIsNotNone(ex)
         self.assertEqual(res, 63)
         binop_count = [opname for opname, _, _ in ex if opname == "_BINARY_OP_ADD_INT"]
@@ -642,12 +647,7 @@ class TestUopsOptimization(unittest.TestCase):
                 num += 1
             return a
 
-        opt = _testinternalcapi.get_uop_optimizer()
-        res = None
-        with temporary_optimizer(opt):
-            res = testfunc(64)
-
-        ex = get_first_executor(testfunc)
+        res, ex = self._run_with_optimizer(testfunc, 64)
         self.assertIsNotNone(ex)
         binop_count = [opname for opname, _, _ in ex if opname == "_BINARY_OP_ADD_INT"]
         self.assertGreaterEqual(len(binop_count), 3)
@@ -659,11 +659,7 @@ class TestUopsOptimization(unittest.TestCase):
             for i in range(n):
                 dummy(i)
 
-        opt = _testinternalcapi.get_uop_optimizer()
-        with temporary_optimizer(opt):
-            testfunc(32)
-
-        ex = get_first_executor(testfunc)
+        res, ex = self._run_with_optimizer(testfunc, 32)
         self.assertIsNotNone(ex)
         uops = {opname for opname, _, _ in ex}
         self.assertIn("_PUSH_FRAME", uops)
@@ -677,11 +673,7 @@ class TestUopsOptimization(unittest.TestCase):
                 x = i + i
             return x
 
-        opt = _testinternalcapi.get_uop_optimizer()
-        with temporary_optimizer(opt):
-            res = testfunc(32)
-
-        ex = get_first_executor(testfunc)
+        res, ex = self._run_with_optimizer(testfunc, 32)
         self.assertEqual(res, 62)
         self.assertIsNotNone(ex)
         uops = {opname for opname, _, _ in ex}
@@ -699,11 +691,7 @@ class TestUopsOptimization(unittest.TestCase):
                 res = x + z + a + b
             return res
 
-        opt = _testinternalcapi.get_uop_optimizer()
-        with temporary_optimizer(opt):
-            res = testfunc(32)
-
-        ex = get_first_executor(testfunc)
+        res, ex = self._run_with_optimizer(testfunc, 32)
         self.assertEqual(res, 4)
         self.assertIsNotNone(ex)
         uops = {opname for opname, _, _ in ex}
@@ -716,11 +704,8 @@ class TestUopsOptimization(unittest.TestCase):
             for _ in range(n):
                 return [i for i in range(n)]
 
-        opt = _testinternalcapi.get_uop_optimizer()
-        with temporary_optimizer(opt):
-            testfunc(32)
-
-        ex = get_first_executor(testfunc)
+        res, ex = self._run_with_optimizer(testfunc, 32)
+        self.assertEqual(res, list(range(32)))
         self.assertIsNotNone(ex)
         uops = {opname for opname, _, _ in ex}
         self.assertNotIn("_BINARY_OP_ADD_INT", uops)
@@ -785,6 +770,56 @@ class TestUopsOptimization(unittest.TestCase):
         """))
         self.assertEqual(result[0].rc, 0, result)
 
+    def test_float_add_constant_propagation(self):
+        def testfunc(n):
+            a = 1.0
+            for _ in range(n):
+                a = a + 0.1
+            return a
+
+        res, ex = self._run_with_optimizer(testfunc, 32)
+        self.assertAlmostEqual(res, 4.2)
+        self.assertIsNotNone(ex)
+        uops = {opname for opname, _, _ in ex}
+        guard_both_float_count = [opname for opname, _, _ in ex if opname == "_GUARD_BOTH_FLOAT"]
+        self.assertLessEqual(len(guard_both_float_count), 1)
+        # TODO gh-115506: this assertion may change after propagating constants.
+        # We'll also need to verify that propagation actually occurs.
+        self.assertIn("_BINARY_OP_ADD_FLOAT", uops)
+
+    def test_float_subtract_constant_propagation(self):
+        def testfunc(n):
+            a = 1.0
+            for _ in range(n):
+                a = a - 0.1
+            return a
+
+        res, ex = self._run_with_optimizer(testfunc, 32)
+        self.assertAlmostEqual(res, -2.2)
+        self.assertIsNotNone(ex)
+        uops = {opname for opname, _, _ in ex}
+        guard_both_float_count = [opname for opname, _, _ in ex if opname == "_GUARD_BOTH_FLOAT"]
+        self.assertLessEqual(len(guard_both_float_count), 1)
+        # TODO gh-115506: this assertion may change after propagating constants.
+        # We'll also need to verify that propagation actually occurs.
+        self.assertIn("_BINARY_OP_SUBTRACT_FLOAT", uops)
+
+    def test_float_multiply_constant_propagation(self):
+        def testfunc(n):
+            a = 1.0
+            for _ in range(n):
+                a = a * 2.0
+            return a
+
+        res, ex = self._run_with_optimizer(testfunc, 32)
+        self.assertAlmostEqual(res, 2 ** 32)
+        self.assertIsNotNone(ex)
+        uops = {opname for opname, _, _ in ex}
+        guard_both_float_count = [opname for opname, _, _ in ex if opname == "_GUARD_BOTH_FLOAT"]
+        self.assertLessEqual(len(guard_both_float_count), 1)
+        # TODO gh-115506: this assertion may change after propagating constants.
+        # We'll also need to verify that propagation actually occurs.
+        self.assertIn("_BINARY_OP_MULTIPLY_FLOAT", uops)
 
 
 if __name__ == "__main__":
diff --git a/Python/tier2_redundancy_eliminator_bytecodes.c b/Python/tier2_redundancy_eliminator_bytecodes.c
index 6aae590..3f6e8ce 100644
--- a/Python/tier2_redundancy_eliminator_bytecodes.c
+++ b/Python/tier2_redundancy_eliminator_bytecodes.c
@@ -132,6 +132,63 @@ dummy_func(void) {
         }
     }
 
+    op(_BINARY_OP_ADD_FLOAT, (left, right -- res)) {
+        if (is_const(left) && is_const(right)) {
+            assert(PyFloat_CheckExact(get_const(left)));
+            assert(PyFloat_CheckExact(get_const(right)));
+            PyObject *temp = PyFloat_FromDouble(
+                PyFloat_AS_DOUBLE(get_const(left)) +
+                PyFloat_AS_DOUBLE(get_const(right)));
+            if (temp == NULL) {
+                goto error;
+            }
+            res = sym_new_const(ctx, temp);
+            // TODO gh-115506:
+            // replace opcode with constant propagated one and update tests!
+        }
+        else {
+            OUT_OF_SPACE_IF_NULL(res = sym_new_known_type(ctx, &PyFloat_Type));
+        }
+    }
+
+    op(_BINARY_OP_SUBTRACT_FLOAT, (left, right -- res)) {
+        if (is_const(left) && is_const(right)) {
+            assert(PyFloat_CheckExact(get_const(left)));
+            assert(PyFloat_CheckExact(get_const(right)));
+            PyObject *temp = PyFloat_FromDouble(
+                PyFloat_AS_DOUBLE(get_const(left)) -
+                PyFloat_AS_DOUBLE(get_const(right)));
+            if (temp == NULL) {
+                goto error;
+            }
+            res = sym_new_const(ctx, temp);
+            // TODO gh-115506:
+            // replace opcode with constant propagated one and update tests!
+        }
+        else {
+            OUT_OF_SPACE_IF_NULL(res = sym_new_known_type(ctx, &PyFloat_Type));
+        }
+    }
+
+    op(_BINARY_OP_MULTIPLY_FLOAT, (left, right -- res)) {
+        if (is_const(left) && is_const(right)) {
+            assert(PyFloat_CheckExact(get_const(left)));
+            assert(PyFloat_CheckExact(get_const(right)));
+            PyObject *temp = PyFloat_FromDouble(
+                PyFloat_AS_DOUBLE(get_const(left)) *
+                PyFloat_AS_DOUBLE(get_const(right)));
+            if (temp == NULL) {
+                goto error;
+            }
+            res = sym_new_const(ctx, temp);
+            // TODO gh-115506:
+            // replace opcode with constant propagated one and update tests!
+        }
+        else {
+            OUT_OF_SPACE_IF_NULL(res = sym_new_known_type(ctx, &PyFloat_Type));
+        }
+    }
+
     op(_LOAD_CONST, (-- value)) {
         // There should be no LOAD_CONST. It should be all
         // replaced by peephole_opt.
diff --git a/Python/tier2_redundancy_eliminator_cases.c.h b/Python/tier2_redundancy_eliminator_cases.c.h
index d1301ff..be2fbb9 100644
--- a/Python/tier2_redundancy_eliminator_cases.c.h
+++ b/Python/tier2_redundancy_eliminator_cases.c.h
@@ -270,27 +270,81 @@
         }
 
         case _BINARY_OP_MULTIPLY_FLOAT: {
+            _Py_UOpsSymType *right;
+            _Py_UOpsSymType *left;
             _Py_UOpsSymType *res;
-            res = sym_new_unknown(ctx);
-            if (res == NULL) goto out_of_space;
+            right = stack_pointer[-1];
+            left = stack_pointer[-2];
+            if (is_const(left) && is_const(right)) {
+                assert(PyFloat_CheckExact(get_const(left)));
+                assert(PyFloat_CheckExact(get_const(right)));
+                PyObject *temp = PyFloat_FromDouble(
+                    PyFloat_AS_DOUBLE(get_const(left)) *
+                    PyFloat_AS_DOUBLE(get_const(right)));
+                if (temp == NULL) {
+                    goto error;
+                }
+                res = sym_new_const(ctx, temp);
+                // TODO gh-115506:
+                // replace opcode with constant propagated one and update tests!
+            }
+            else {
+                OUT_OF_SPACE_IF_NULL(res = sym_new_known_type(ctx, &PyFloat_Type));
+            }
             stack_pointer[-2] = res;
             stack_pointer += -1;
             break;
         }
 
         case _BINARY_OP_ADD_FLOAT: {
+            _Py_UOpsSymType *right;
+            _Py_UOpsSymType *left;
             _Py_UOpsSymType *res;
-            res = sym_new_unknown(ctx);
-            if (res == NULL) goto out_of_space;
+            right = stack_pointer[-1];
+            left = stack_pointer[-2];
+            if (is_const(left) && is_const(right)) {
+                assert(PyFloat_CheckExact(get_const(left)));
+                assert(PyFloat_CheckExact(get_const(right)));
+                PyObject *temp = PyFloat_FromDouble(
+                    PyFloat_AS_DOUBLE(get_const(left)) +
+                    PyFloat_AS_DOUBLE(get_const(right)));
+                if (temp == NULL) {
+                    goto error;
+                }
+                res = sym_new_const(ctx, temp);
+                // TODO gh-115506:
+                // replace opcode with constant propagated one and update tests!
+            }
+            else {
+                OUT_OF_SPACE_IF_NULL(res = sym_new_known_type(ctx, &PyFloat_Type));
+            }
             stack_pointer[-2] = res;
             stack_pointer += -1;
             break;
         }
 
         case _BINARY_OP_SUBTRACT_FLOAT: {
+            _Py_UOpsSymType *right;
+            _Py_UOpsSymType *left;
             _Py_UOpsSymType *res;
-            res = sym_new_unknown(ctx);
-            if (res == NULL) goto out_of_space;
+            right = stack_pointer[-1];
+            left = stack_pointer[-2];
+            if (is_const(left) && is_const(right)) {
+                assert(PyFloat_CheckExact(get_const(left)));
+                assert(PyFloat_CheckExact(get_const(right)));
+                PyObject *temp = PyFloat_FromDouble(
+                    PyFloat_AS_DOUBLE(get_const(left)) -
+                    PyFloat_AS_DOUBLE(get_const(right)));
+                if (temp == NULL) {
+                    goto error;
+                }
+                res = sym_new_const(ctx, temp);
+                // TODO gh-115506:
+                // replace opcode with constant propagated one and update tests!
+            }
+            else {
+                OUT_OF_SPACE_IF_NULL(res = sym_new_known_type(ctx, &PyFloat_Type));
+            }
             stack_pointer[-2] = res;
             stack_pointer += -1;
             break;