[Python] Evaluate a negative literal exponent in floating point

python/cudaq/kernel/ast_bridge.py

@@ -650,6 +650,17 @@ def getConstantInt(self, value, width=64):
         ty = self.getIntegerType(width)
         return arith.ConstantOp(ty, self.getIntegerAttr(ty, value)).result
 
+    def __integerLiteralValue(self, node):
+        """Return the value of an AST node that is an integer literal, possibly
+        negated, and `None` for any other node.
+        """
+        if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
+            value = self.__integerLiteralValue(node.operand)
+            return None if value is None else -value
+        if isinstance(node, ast.Constant) and isinstance(node.value, int):
+            return node.value
+        return None
+
     def __arithmetic_to_bool(self, value):
         """Converts an integer or floating point value to a bool by comparing it
         to zero."""
@@ -5810,11 +5821,13 @@ def visit_Continue(self, node):
         else:
             cc.ContinueOp([])
 
-    def __process_binary_op(self, left, right, nodeType):
+    def __process_binary_op(self, left, right, nodeType, rightNode=None):
         """Process a binary operation in the AST and map them to equivalents in
         the MLIR.
 
-        This method handles arithmetic operations between values.
+        This method handles arithmetic operations between values. `rightNode` is
+        the AST node the `right` value was created from, if available; it is
+        used to detect literal operands.
         """
 
         # `measure_handle` operands in arithmetic context discriminate
@@ -5927,6 +5940,21 @@ def __process_binary_op(self, left, right, nodeType):
         if issubclass(nodeType, ast.Pow):
             if IntegerType.isinstance(left.type) and IntegerType.isinstance(
                     right.type):
+                # Python raises an integer to a negative integer power in
+                # floating point (`2 ** -1` is `0.5`), and only keeps the
+                # integer type for a non-negative exponent. `math.ipowi` is an
+                # integer power that returns 0 for a negative exponent unless
+                # the base is +/-1, so promote the base and use the
+                # floating-point power when the exponent is a negative literal.
+                # The sign of an exponent that is only known at run time cannot
+                # be taken into account here, since the type of the expression
+                # has to be determined at compile time.
+                exponent = (None if rightNode is None else
+                            self.__integerLiteralValue(rightNode))
+                if exponent is not None and exponent < 0:
+                    left = self.changeOperandToType(self.getFloatType(), left)
+                    self.pushValue(math.FPowIOp(left, right).result)
+                    return
                 # `math.ipowi` does not lower to LLVM as is
                 # workaround, use math to function conversion
                 self.pushValue(math.IPowIOp(left, right).result)
@@ -6019,7 +6047,7 @@ def visit_BinOp(self, node):
         right = self.popValue()
 
         # pushes to the value stack
-        self.__process_binary_op(left, right, type(node.op))
+        self.__process_binary_op(left, right, type(node.op), node.right)
 
     def visit_AugAssign(self, node):
         """Visit augment-assign operations (e.g. +=)."""
@@ -6054,7 +6082,7 @@ def visit_AugAssign(self, node):
         # some complexity. We hence effectively disallow using any kind of
         # assignment as expression.
         self.valueStack.pushFrame()
-        self.__process_binary_op(loaded, value, type(node.op))
+        self.__process_binary_op(loaded, value, type(node.op), node.value)
         self.valueStack.popFrame()
         res = self.popValue()
 

python/tests/kernel/test_kernel_float.py

@@ -623,3 +623,89 @@ def check(c: any):
     check([np.float32(np.pi / 2), 0])
     check([1, 0])
     check([np.pi / 2, 0, True])
+
+
+def test_negative_integer_exponent_is_float():
+    """An integer raised to a negative integer power is a float, as in Python,
+    and not an integer power that truncates to zero."""
+
+    @cudaq.kernel
+    def int_base_neg_one() -> float:
+        return 2**-1
+
+    @cudaq.kernel
+    def int_base_neg_three() -> float:
+        return 2**-3
+
+    @cudaq.kernel
+    def other_int_base_neg_one() -> float:
+        return 3**-1
+
+    @cudaq.kernel
+    def negative_int_base() -> float:
+        return (-2)**-1
+
+    @cudaq.kernel
+    def variable_base(base: int) -> float:
+        return base**-2
+
+    assert is_close(2**-1, int_base_neg_one())
+    assert is_close(2**-3, int_base_neg_three())
+    assert is_close(3**-1, other_int_base_neg_one())
+    assert is_close((-2)**-1, negative_int_base())
+    assert is_close(3**-2, variable_base(3))
+    assert is_close((-2)**-2, variable_base(-2))
+
+
+def test_non_negative_integer_exponent_is_integer():
+    """A non-negative exponent keeps the integer power and the integer type."""
+
+    @cudaq.kernel
+    def cube() -> int:
+        return 2**3
+
+    @cudaq.kernel
+    def zeroth_power() -> int:
+        return 2**0
+
+    @cudaq.kernel
+    def large_power() -> int:
+        return 2**30
+
+    @cudaq.kernel
+    def variable_exponent(base: int, exponent: int) -> int:
+        return base**exponent
+
+    @cudaq.kernel
+    def loop_bound() -> int:
+        total = 0
+        for i in range(2**3):
+            total += i
+        return total
+
+    assert cube() == 2**3
+    assert zeroth_power() == 2**0
+    assert large_power() == 2**30
+    assert variable_exponent(2, 5) == 2**5
+    assert variable_exponent(-2, 3) == (-2)**3
+    assert loop_bound() == sum(range(2**3))
+
+
+def test_float_base_with_integer_exponent():
+    """A floating-point base keeps working for either sign of the exponent."""
+
+    @cudaq.kernel
+    def float_base_neg_one() -> float:
+        return 2.0**-1
+
+    @cudaq.kernel
+    def float_base_neg_three() -> float:
+        return 2.0**-3
+
+    @cudaq.kernel
+    def float_base_cube() -> float:
+        return 2.0**3
+
+    assert is_close(2.0**-1, float_base_neg_one())
+    assert is_close(2.0**-3, float_base_neg_three())
+    assert is_close(2.0**3, float_base_cube())