diff --git a/python/cudaq/kernel/ast_bridge.py b/python/cudaq/kernel/ast_bridge.py index e18809da75c..0bcba5e4a15 100644 --- a/python/cudaq/kernel/ast_bridge.py +++ b/python/cudaq/kernel/ast_bridge.py @@ -5788,12 +5788,12 @@ def visit_Break(self, node): self.emitFatalError("break statement outside of for loop body.", node) + # Get the innermost enclosing `for` or `while` loop + inArgs = [b for b in self.inForBodyStack[-1]] if self.isInIfStmtBlock(): - # Get the innermost enclosing `for` or `while` loop - inArgs = [b for b in self.inForBodyStack[-1]] cc.UnwindBreakOp(inArgs) else: - cc.BreakOp([]) + cc.BreakOp(inArgs) return @@ -5803,12 +5803,12 @@ def visit_Continue(self, node): self.emitFatalError("continue statement outside of for loop body.", node) + # Get the innermost enclosing `for` or `while` loop + inArgs = [b for b in self.inForBodyStack[-1]] if self.isInIfStmtBlock(): - # Get the innermost enclosing `for` or `while` loop - inArgs = [b for b in self.inForBodyStack[-1]] cc.UnwindContinueOp(inArgs) else: - cc.ContinueOp([]) + cc.ContinueOp(inArgs) def __process_binary_op(self, left, right, nodeType): """Process a binary operation in the AST and map them to equivalents in diff --git a/python/tests/mlir/ast_break.py b/python/tests/mlir/ast_break.py index 2770fd229c6..d5bcc9494b7 100644 --- a/python/tests/mlir/ast_break.py +++ b/python/tests/mlir/ast_break.py @@ -60,3 +60,43 @@ def kernel(x: float): # CHECK: quake.dealloc %[[VAL_8]] : !quake.veq<4> # CHECK: return # CHECK: } + + +def test_bare_break(): + # A `break` that is *not* nested in an `if` must still forward the + # loop-carried block arguments to the `cc.break` terminator. + + @cudaq.kernel(verbose=False) + def bare_break_kernel(x: float): + q = cudaq.qvector(4) + for i in range(4): + ry(x, q[i]) + break + + print(bare_break_kernel) + bare_break_kernel(1.2) + + +# CHECK-LABEL: func.func @__nvqpp__mlirgen__bare_break_kernel.. +# CHECK-SAME: %[[VAL_30:.*]]: f64) attributes {"cudaq-entrypoint", "cudaq-kernel"} { +# CHECK-DAG: %[[VAL_31:.*]] = arith.constant 1 : i64 +# CHECK-DAG: %[[VAL_32:.*]] = arith.constant 0 : i64 +# CHECK-DAG: %[[VAL_33:.*]] = arith.constant 4 : i64 +# CHECK-DAG: %[[VAL_34:.*]] = cc.undef i64 +# CHECK-DAG: %[[VAL_35:.*]] = quake.alloca !quake.veq<4> +# CHECK: %[[VAL_36:.*]]:2 = cc.loop while ((%[[VAL_37:.*]] = %[[VAL_32]], %[[VAL_38:.*]] = %[[VAL_34]]) -> (i64, i64)) { +# CHECK: %[[VAL_39:.*]] = arith.cmpi slt, %[[VAL_37]], %[[VAL_33]] : i64 +# CHECK: cc.condition %[[VAL_39]](%[[VAL_37]], %[[VAL_38]] : i64, i64) +# CHECK: } do { +# CHECK: ^bb0(%[[VAL_40:.*]]: i64, %[[VAL_41:.*]]: i64): +# CHECK: %[[VAL_42:.*]] = quake.extract_ref %[[VAL_35]]{{\[}}%[[VAL_40]]] : (!quake.veq<4>, i64) -> !quake.ref +# CHECK: quake.ry (%[[VAL_30]]) %[[VAL_42]] : (f64, !quake.ref) -> () +# CHECK: cc.break %[[VAL_40]], %[[VAL_40]] : i64, i64 +# CHECK: } step { +# CHECK: ^bb0(%[[VAL_43:.*]]: i64, %[[VAL_44:.*]]: i64): +# CHECK: %[[VAL_45:.*]] = arith.addi %[[VAL_43]], %[[VAL_31]] : i64 +# CHECK: cc.continue %[[VAL_45]], %[[VAL_44]] : i64, i64 +# CHECK: } +# CHECK: quake.dealloc %[[VAL_35]] : !quake.veq<4> +# CHECK: return +# CHECK: } diff --git a/python/tests/mlir/ast_continue.py b/python/tests/mlir/ast_continue.py index 08e39fc7f13..b3a262a9d33 100644 --- a/python/tests/mlir/ast_continue.py +++ b/python/tests/mlir/ast_continue.py @@ -84,3 +84,43 @@ def kernel(x: float): # CHECK-LABEL: kernel: # CHECK-NEXT: object is not callable # CHECK-NEXT: offending source -> x(q + + +def test_bare_continue(): + # A `continue` that is *not* nested in an `if` must still forward the + # loop-carried block arguments to the `cc.continue` terminator. + + @cudaq.kernel + def bare_continue_kernel(x: float): + q = cudaq.qvector(4) + for i in range(4): + ry(x, q[i]) + continue + + print(bare_continue_kernel) + bare_continue_kernel(1.2) + + +# CHECK-LABEL: func.func @__nvqpp__mlirgen__bare_continue_kernel.. +# CHECK-SAME: %[[VAL_40:.*]]: f64) attributes {"cudaq-entrypoint", "cudaq-kernel"} { +# CHECK-DAG: %[[VAL_41:.*]] = arith.constant 1 : i64 +# CHECK-DAG: %[[VAL_42:.*]] = arith.constant 0 : i64 +# CHECK-DAG: %[[VAL_43:.*]] = arith.constant 4 : i64 +# CHECK-DAG: %[[VAL_44:.*]] = cc.undef i64 +# CHECK-DAG: %[[VAL_45:.*]] = quake.alloca !quake.veq<4> +# CHECK: %[[VAL_46:.*]]:2 = cc.loop while ((%[[VAL_47:.*]] = %[[VAL_42]], %[[VAL_48:.*]] = %[[VAL_44]]) -> (i64, i64)) { +# CHECK: %[[VAL_49:.*]] = arith.cmpi slt, %[[VAL_47]], %[[VAL_43]] : i64 +# CHECK: cc.condition %[[VAL_49]](%[[VAL_47]], %[[VAL_48]] : i64, i64) +# CHECK: } do { +# CHECK: ^bb0(%[[VAL_50:.*]]: i64, %[[VAL_51:.*]]: i64): +# CHECK: %[[VAL_52:.*]] = quake.extract_ref %[[VAL_45]]{{\[}}%[[VAL_50]]] : (!quake.veq<4>, i64) -> !quake.ref +# CHECK: quake.ry (%[[VAL_40]]) %[[VAL_52]] : (f64, !quake.ref) -> () +# CHECK: cc.continue %[[VAL_50]], %[[VAL_50]] : i64, i64 +# CHECK: } step { +# CHECK: ^bb0(%[[VAL_53:.*]]: i64, %[[VAL_54:.*]]: i64): +# CHECK: %[[VAL_55:.*]] = arith.addi %[[VAL_53]], %[[VAL_41]] : i64 +# CHECK: cc.continue %[[VAL_55]], %[[VAL_54]] : i64, i64 +# CHECK: } +# CHECK: quake.dealloc %[[VAL_45]] : !quake.veq<4> +# CHECK: return +# CHECK: }