diff --git a/python/cudaq/kernel/ast_bridge.py b/python/cudaq/kernel/ast_bridge.py index efb7bb78033..11b02dcef5c 100644 --- a/python/cudaq/kernel/ast_bridge.py +++ b/python/cudaq/kernel/ast_bridge.py @@ -2672,11 +2672,33 @@ def checkModule(obj, moduleNames): node.func.value.id) and node.func.attr == 'kernel': return + def isExactCudaqDbgAstCall(func_node: ast.AST) -> bool: + """Return True iff `func_node` is the exact AST shape for + ``.dbg.ast.``. + + Runtime attribute lookup follows lazy aliases (e.g. ``cudaq.ast`` + resolves to ``cudaq.dbg.ast`` via ``_LAZY_SUBMODULES``), so + `devKey` is not a sufficient check. Walk the literal node + structure instead.""" + if not isinstance(func_node, ast.Attribute): + return False + if not isinstance( + func_node.value, + ast.Attribute) or func_node.value.attr != 'ast': + return False + if not isinstance( + func_node.value.value, + ast.Attribute) or func_node.value.value.attr != 'dbg': + return False + root = func_node.value.value.value + return isinstance(root, ast.Name) and self.isCudaqName(root.id) + devKey, name = resolveQualifiedName(node.func) if devKey: # Handle debug functions - if devKey == 'cudaq.dbg.ast': + if devKey == 'cudaq.dbg.ast' and isExactCudaqDbgAstCall( + node.func): # Handle a debug print statement arg = self.__groupValues(node.args, [1]) self.__insertDbgStmt(arg, name) diff --git a/python/tests/kernel/test_kernel_features.py b/python/tests/kernel/test_kernel_features.py index b3b99086028..856c2ea8311 100644 --- a/python/tests/kernel/test_kernel_features.py +++ b/python/tests/kernel/test_kernel_features.py @@ -570,6 +570,66 @@ def test2(myList: List[int]): assert '1010' in counts +def test_dbg_ast_strict_path(): + """Only cudaq.dbg.ast.print_i64/f64 is valid inside a kernel. + + cudaq.ast is a lazy alias for cudaq.dbg.ast (via _LAZY_SUBMODULES), so + without an explicit AST structure check it would pass the devKey guard. + Component reorderings like dbg.cudaq.ast are also rejected. + + See https://github.com/NVIDIA/cuda-quantum/issues/2342 + """ + + @cudaq.kernel + def valid_kernel(n: int, f: float): + q = cudaq.qvector(n) + h(q[0]) + cudaq.dbg.ast.print_i64(n) + cudaq.dbg.ast.print_f64(f) + mz(q) + + counts = cudaq.sample(valid_kernel, 2, 2.0) + assert len(counts) > 0 + + # cudaq.ast resolves to cudaq.dbg.ast at runtime via _LAZY_SUBMODULES; + # isExactCudaqDbgAstCall rejects it because the AST node chain is wrong. + with pytest.raises(RuntimeError): + + @cudaq.kernel + def invalid_cudaq_ast(n: int): + q = cudaq.qvector(n) + h(q[0]) + cudaq.ast.print_i64(n) + mz(q) + + cudaq.sample(invalid_cudaq_ast, 2) + + # dbg.cudaq.ast - resolveQualifiedName returns 'dbg.cudaq.ast', which + # never matches the devKey guard. + with pytest.raises(RuntimeError): + + @cudaq.kernel + def invalid_dbg_cudaq_ast(n: int): + q = cudaq.qvector(n) + h(q[0]) + dbg.cudaq.ast.print_i64(n) + mz(q) + + cudaq.sample(invalid_dbg_cudaq_ast, 2) + + # ast.cudaq.dbg - same: resolveQualifiedName returns 'ast.cudaq.dbg'. + with pytest.raises(RuntimeError): + + @cudaq.kernel + def invalid_ast_cudaq_dbg(n: int): + q = cudaq.qvector(n) + h(q[0]) + ast.cudaq.dbg.print_i64(n) + mz(q) + + cudaq.sample(invalid_ast_cudaq_dbg, 2) + + def test_no_dynamic_Lists(): with pytest.raises(RuntimeError) as error: