Fixes for Python modulo operation

python/cudaq/kernel/ast_bridge.py

@@ -4604,9 +4604,7 @@ def __process_binary_op(self, left, right, nodeType):
         right = self.ifPointerThenLoad(right)
 
         # type promotion for anything except pow to match Python behavior
-        if not issubclass(
-                nodeType,
-            (ast.Pow, ast.Mod)):  # FIXME: remove modulo here and fix it below
+        if not issubclass(nodeType, ast.Pow):
             superiorTy = self.__get_superior_type(left.type, right.type)
             if superiorTy is not None:
                 left = self.changeOperandToType(superiorTy,
@@ -4716,18 +4714,15 @@ def __process_binary_op(self, left, right, nodeType):
                                     self.currentNode)
 
         if issubclass(nodeType, ast.Mod):
-            # FIXME: This should be revised to
-            # 1) properly fail when we have a complex number
-            # 2) use `arith.RemFOp` for floating point
-            # (these changes are split out into a separate PR
-            # per review request)
-            if F64Type.isinstance(left.type):
-                left = arith.FPToSIOp(self.getIntegerType(), left).result
-            if F64Type.isinstance(right.type):
-                right = arith.FPToSIOp(self.getIntegerType(), right).result
-
-            self.pushValue(arith.RemUIOp(left, right).result)
-            return
+            if IntegerType.isinstance(left.type):
+                self.pushValue(arith.RemUIOp(left, right).result)
+                return
+            if (F64Type.isinstance(left.type) or F32Type.isinstance(left.type)):
+                self.pushValue(arith.RemFOp(left, right).result)
+                return
+            else:
+                self.emitFatalError("unhandled BinOp.Mod types",
+                                    self.currentNode)
 
         if issubclass(nodeType, ast.LShift):
             if IntegerType.isinstance(left.type):

python/tests/mlir/ast_elif.py

@@ -19,7 +19,7 @@ def test_elif():
     def cost(thetas: np.ndarray):  # can pass 1D ndarray or list
         q = cudaq.qvector(4)
         for i, theta in enumerate(thetas):
-            if i % 2.0:  # asserting we convert 2.0 to 2
+            if i % 2.0:  # asserting we convert i to float
                 ry(theta, q[i % 4])
             else:
                 rx(theta, q[i % 4])
@@ -44,10 +44,10 @@ def cost(thetas: np.ndarray):  # can pass 1D ndarray or list
 # CHECK:             %[[VAL_10:.*]] = cc.stdvec_data %[[VAL_0]] : (!cc.stdvec<f64>) -> !cc.ptr<!cc.array<f64 x ?>>
 # CHECK:             %[[VAL_11:.*]] = cc.compute_ptr %[[VAL_10]]{{\[}}%[[VAL_9]]] : (!cc.ptr<!cc.array<f64 x ?>>, i64) -> !cc.ptr<f64>
 # CHECK:             %[[VAL_12:.*]] = arith.constant 2.000000e+00 : f64
-# CHECK:             %[[VAL_13:.*]] = arith.fptosi %[[VAL_12]] : f64 to i64
-# CHECK:             %[[VAL_14:.*]] = arith.remui %[[VAL_9]], %[[VAL_13]] : i64
-# CHECK:             %[[VAL_15:.*]] = arith.constant 0 : i64
-# CHECK:             %[[VAL_16:.*]] = arith.cmpi ne, %[[VAL_14]], %[[VAL_15]] : i64
+# CHECK:             %[[VAL_13:.*]] = cc.cast signed %[[VAL_9]] : (i64) -> f64
+# CHECK:             %[[VAL_14:.*]] = arith.remf %[[VAL_13]], %[[VAL_12]] : f64
+# CHECK:             %[[VAL_15:.*]] = arith.constant 0.000000e+00 : f64
+# CHECK:             %[[VAL_16:.*]] = arith.cmpf une, %[[VAL_14]], %[[VAL_15]] : f64
 # CHECK:             cc.if(%[[VAL_16]]) {
 # CHECK:               %[[VAL_17:.*]] = cc.load %[[VAL_11]] : !cc.ptr<f64>
 # CHECK:               %[[VAL_18:.*]] = arith.constant 4 : i64
@@ -72,7 +72,8 @@ def cost(thetas: np.ndarray):  # can pass 1D ndarray or list
 
 # CHECK-LABEL:   func.func @__nvqpp__mlirgen__cost(
 # CHECK-SAME:      %[[VAL_0:.*]]: !cc.stdvec<f64>) attributes {"cudaq-entrypoint", "cudaq-kernel"} {
-# CHECK-DAG:       %[[VAL_1:.*]] = arith.constant 2 : i64
+# CHECK-DAG:       %[[VAL:.*]] = arith.constant 0.000000e+00 : f64
+# CHECK-DAG:       %[[VAL_1:.*]] = arith.constant 2.000000e+00 : f64
 # CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 1 : i64
 # CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 0 : i64
 # CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 4 : i64
@@ -85,24 +86,25 @@ def cost(thetas: np.ndarray):  # can pass 1D ndarray or list
 # CHECK:           ^bb0(%[[VAL_10:.*]]: i64):
 # CHECK:             %[[VAL_11:.*]] = cc.stdvec_data %[[VAL_0]] : (!cc.stdvec<f64>) -> !cc.ptr<!cc.array<f64 x ?>>
 # CHECK:             %[[VAL_12:.*]] = cc.compute_ptr %[[VAL_11]]{{\[}}%[[VAL_10]]] : (!cc.ptr<!cc.array<f64 x ?>>, i64) -> !cc.ptr<f64>
-# CHECK:             %[[VAL_13:.*]] = arith.remui %[[VAL_10]], %[[VAL_1]] : i64
-# CHECK:             %[[VAL_14:.*]] = arith.cmpi ne, %[[VAL_13]], %[[VAL_3]] : i64
-# CHECK:             cc.if(%[[VAL_14]]) {
-# CHECK:               %[[VAL_15:.*]] = cc.load %[[VAL_12]] : !cc.ptr<f64>
-# CHECK:               %[[VAL_16:.*]] = arith.remui %[[VAL_10]], %[[VAL_4]] : i64
-# CHECK:               %[[VAL_17:.*]] = quake.extract_ref %[[VAL_5]]{{\[}}%[[VAL_16]]] : (!quake.veq<4>, i64) -> !quake.ref
-# CHECK:               quake.ry (%[[VAL_15]]) %[[VAL_17]] : (f64, !quake.ref) -> ()
+# CHECK:             %[[VAL_13:.*]] = cc.cast signed %[[VAL_10]] : (i64) -> f64
+# CHECK:             %[[VAL_14:.*]] = arith.remf %[[VAL_13]], %[[VAL_1]] : f64
+# CHECK:             %[[VAL_15:.*]] = arith.cmpf une, %[[VAL_14]], %[[VAL]] : f64
+# CHECK:             cc.if(%[[VAL_15]]) {
+# CHECK:               %[[VAL_16:.*]] = cc.load %[[VAL_12]] : !cc.ptr<f64>
+# CHECK:               %[[VAL_17:.*]] = arith.remui %[[VAL_10]], %[[VAL_4]] : i64
+# CHECK:               %[[VAL_18:.*]] = quake.extract_ref %[[VAL_5]]{{\[}}%[[VAL_17]]] : (!quake.veq<4>, i64) -> !quake.ref
+# CHECK:               quake.ry (%[[VAL_16]]) %[[VAL_18]] : (f64, !quake.ref) -> ()
 # CHECK:             } else {
-# CHECK:               %[[VAL_18:.*]] = cc.load %[[VAL_12]] : !cc.ptr<f64>
-# CHECK:               %[[VAL_19:.*]] = arith.remui %[[VAL_10]], %[[VAL_4]] : i64
-# CHECK:               %[[VAL_20:.*]] = quake.extract_ref %[[VAL_5]]{{\[}}%[[VAL_19]]] : (!quake.veq<4>, i64) -> !quake.ref
-# CHECK:               quake.rx (%[[VAL_18]]) %[[VAL_20]] : (f64, !quake.ref) -> ()
+# CHECK:               %[[VAL_19:.*]] = cc.load %[[VAL_12]] : !cc.ptr<f64>
+# CHECK:               %[[VAL_20:.*]] = arith.remui %[[VAL_10]], %[[VAL_4]] : i64
+# CHECK:               %[[VAL_21:.*]] = quake.extract_ref %[[VAL_5]]{{\[}}%[[VAL_20]]] : (!quake.veq<4>, i64) -> !quake.ref
+# CHECK:               quake.rx (%[[VAL_19]]) %[[VAL_21]] : (f64, !quake.ref) -> ()
 # CHECK:             }
 # CHECK:             cc.continue %[[VAL_10]] : i64
 # CHECK:           } step {
-# CHECK:           ^bb0(%[[VAL_21:.*]]: i64):
-# CHECK:             %[[VAL_22:.*]] = arith.addi %[[VAL_21]], %[[VAL_2]] : i64
-# CHECK:             cc.continue %[[VAL_22]] : i64
+# CHECK:           ^bb0(%[[VAL_22:.*]]: i64):
+# CHECK:             %[[VAL_23:.*]] = arith.addi %[[VAL_22]], %[[VAL_2]] : i64
+# CHECK:             cc.continue %[[VAL_23]] : i64
 # CHECK:           } {invariant}
 # CHECK:           return
-# CHECK:         }
+# CHECK:         }