Add option to control unwind lowering

cudaq/include/cudaq/Optimizer/Transforms/Passes.h

@@ -26,7 +26,8 @@ namespace cudaq::opt {
 /// Add a pass pipeline to transform call between kernels to direct calls that
 /// do not go through the runtime layers, inline all calls, and detect if calls
 /// to kernels remain in the fully inlined into entry point kernel.
-void addAggressiveInlining(mlir::OpPassManager &pm, bool fatalCheck = false);
+void addAggressiveInlining(mlir::OpPassManager &pm, bool fatalCheck = false,
+                           bool lowerUnwind = true);
 void registerAggressiveInliningPipeline();
 
 void registerUnrollingPipeline();

cudaq/include/cudaq/Target/CompileTarget.h

@@ -44,6 +44,10 @@ struct CompileTarget {
     /// Optional pass pipeline to run after code generation.
     std::string postCodeGenPasses;
 
+    /// How the AOT pipeline handles unwind operations. The `none` mode causes
+    /// compilation to fail if unwind operations remain in the IR.
+    std::string aotUnwindMode = "cfg";
+
     /// Whether to disable qubit mapping.
     bool disableQubitMapping = false;
 

cudaq/include/cudaq/Target/TargetConfig.h

@@ -105,6 +105,10 @@ struct BackendEndConfigEntry {
   std::string JITHighLevelPipeline;
   std::string JITMidLevelPipeline;
   std::string JITLowLevelPipeline;
+  /// How the AOT pipeline handles unwind operations: `cfg`, `dataflow`,
+  /// `none`. The `none` mode performs no unwind lowering, causing compilation
+  /// to fail if unwind operations remain in the IR.
+  std::string AOTUnwindMode;
   /// Exact cudaq-opt passes for pseudo-targets
   std::string TargetPassPipeline;
   /// Codegen emission configuration (hardware REST QPU)

cudaq/include/cudaq/Target/TargetConfigYaml.h

@@ -52,6 +52,8 @@ struct MappingTraits<cudaq::config::ConditionalBuildConfig> {
 template <>
 struct MappingTraits<cudaq::config::BackendEndConfigEntry> {
   static void mapping(IO &io, cudaq::config::BackendEndConfigEntry &info);
+  static std::string validate(IO &io,
+                              cudaq::config::BackendEndConfigEntry &info);
 };
 template <>
 struct MappingTraits<cudaq::config::BackendFeatureMap> {

cudaq/lib/Optimizer/Transforms/AggressiveInlining.cpp

@@ -160,16 +160,18 @@ class CheckKernelCalls
 static void defaultInlinerOptPipeline(OpPassManager &pm) {}
 
 /// Run the passes in the correct order.
-/// 1) Lower unwind control flow before creating call-site scopes.
+/// 1) Optionally lower unwind control flow before creating call-site scopes.
 /// 2) Convert calls between kernels to direct calls (on the QPU).
 /// 3) Aggressively inline all calls.
 /// 4) Detect if kernel inlining has failed and left behind calls to kernels.
 /// Such a failure is most likely a sign that there is a cycle in the call
 /// graph. [This check is a bad idea: this should be deferred to final codegen
 /// when translating the final Quake IR.]
-void cudaq::opt::addAggressiveInlining(OpPassManager &pm, bool fatalChecks) {
+void cudaq::opt::addAggressiveInlining(OpPassManager &pm, bool fatalChecks,
+                                       bool lowerUnwind) {
   llvm::StringMap<OpPassManager> opPipelines;
-  pm.addNestedPass<func::FuncOp>(cudaq::opt::createUnwindLowering());
+  if (lowerUnwind)
+    pm.addNestedPass<func::FuncOp>(cudaq::opt::createUnwindLowering());
   pm.addPass(cudaq::opt::createConvertToDirectCalls());
   pm.addPass(createInlinerPass(opPipelines, defaultInlinerOptPipeline));
   if (fatalChecks)
@@ -192,6 +194,10 @@ struct AggressiveInliningPipelineOptions
       *this, "fatal-check",
       llvm::cl::desc("run checker and produce fatal errors immediately"),
       llvm::cl::init(false)};
+  PassOptions::Option<bool> lowerUnwind{
+      *this, "lower-unwind",
+      llvm::cl::desc("lower unwind operations before inlining"),
+      llvm::cl::init(true)};
 };
 } // namespace
 
@@ -200,6 +206,6 @@ void cudaq::opt::registerAggressiveInliningPipeline() {
       "aggressive-inlining",
       "Convert calls between kernels to direct calls and inline functions.",
       [](OpPassManager &pm, const AggressiveInliningPipelineOptions &opt) {
-        addAggressiveInlining(pm, opt.runFatalChecker);
+        addAggressiveInlining(pm, opt.runFatalChecker, opt.lowerUnwind);
       });
 }

cudaq/lib/Optimizer/Transforms/Pipelines.cpp

@@ -13,6 +13,8 @@
 using namespace mlir;
 
 namespace {
+enum class AOTUnwindMode { CFG, Dataflow, None };
+
 struct TargetPrepPipelineOptions
     : public PassPipelineOptions<TargetPrepPipelineOptions> {
   PassOptions::Option<bool> eraseNoise{
@@ -50,6 +52,17 @@ struct TargetFinalizationPipelineOptions
 };
 
 struct PythonAOTOptions : public PassPipelineOptions<PythonAOTOptions> {
+  PassOptions::Option<AOTUnwindMode> unwindMode{
+      *this, "unwind-mode",
+      llvm::cl::desc("Select how the AOT pipeline handles unwind operations."),
+      llvm::cl::init(AOTUnwindMode::CFG),
+      llvm::cl::values(
+          clEnumValN(AOTUnwindMode::CFG, "cfg",
+                     "Lower unwind operations to control-flow operations."),
+          clEnumValN(AOTUnwindMode::Dataflow, "dataflow",
+                     "Lower unwind operations using structured data flow."),
+          clEnumValN(AOTUnwindMode::None, "none",
+                     "Do not lower unwind operations."))};
   PassOptions::Option<bool> autoGenRunStack{
       *this, "gen-run-stack",
       llvm::cl::desc("Autogenerate the cudaq::run dispatch stack."),
@@ -298,7 +311,16 @@ static void createPythonAOTPipeline(OpPassManager &pm,
                                     const PythonAOTOptions &options) {
   // NB: This pipeline should be kept in synch with the pipeline in nvq++.
   pm.addNestedPass<func::FuncOp>(cudaq::opt::createVariableCoalesce());
-  pm.addNestedPass<func::FuncOp>(cudaq::opt::createUnwindLowering());
+  switch (options.unwindMode) {
+  case AOTUnwindMode::CFG:
+    pm.addNestedPass<func::FuncOp>(cudaq::opt::createUnwindLowering());
+    break;
+  case AOTUnwindMode::Dataflow:
+    pm.addNestedPass<func::FuncOp>(cudaq::opt::createUnwindByDataFlow());
+    break;
+  case AOTUnwindMode::None:
+    break;
+  }
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<func::FuncOp>(cudaq::opt::createInjectImplicitOutput());
   pm.addNestedPass<func::FuncOp>(cudaq::opt::createAddDeallocs());
@@ -315,7 +337,10 @@ static void createPythonAOTPipeline(OpPassManager &pm,
   pm.addPass(cudaq::opt::createGenerateKernelExecution(gkeOpts));
   if (options.autoGenRunStack)
     pm.addPass(cudaq::opt::createRunSemanticsHackery());
-  cudaq::opt::addAggressiveInlining(pm);
+  // `none` leaves unwind operations in the IR so downstream lowering causes
+  // compilation to fail if the source contains unwinding control flow.
+  cudaq::opt::addAggressiveInlining(pm, /*fatalCheck=*/false,
+                                    options.unwindMode != AOTUnwindMode::None);
   pm.addNestedPass<func::FuncOp>(cudaq::opt::createQuakeAddMetadata());
   pm.addNestedPass<func::FuncOp>(cudaq::opt::createConstantPropagation());
   pm.addNestedPass<func::FuncOp>(cudaq::opt::createLiftArrayAlloc());
@@ -339,7 +364,7 @@ void cudaq::opt::createPythonAOTPipeline(OpPassManager &pm,
 
 static void registerPythonAOTPipeline() {
   PassPipelineRegistration<PythonAOTOptions>(
-      "aot-prep-pipeline",
+      "python-aot-pipeline",
       "Pipeline to lower code for simulation or JIT compilation.",
       [](OpPassManager &pm, const PythonAOTOptions &options) {
         ::createPythonAOTPipeline(pm, options);

cudaq/lib/Optimizer/Transforms/UnwindByDataFlow.cpp

@@ -217,6 +217,11 @@ class UnwindByDataFlowPass
       auto i8Ty = rewriter.getI8Type();
       rewriter.setInsertionPoint(&func.getBody().front().front());
       dfJump = cudaq::cc::AllocaOp::create(rewriter, loc, i8Ty);
+      // Generated guards may load `dfJump` before an unwind occurs, so start
+      // the control state at `JumpKind::None` rather than leaving it undefined.
+      const Value noJump = arith::ConstantIntOp::create(
+          rewriter, loc, static_cast<int>(JumpKind::None), 8);
+      cudaq::cc::StoreOp::create(rewriter, loc, noJump, dfJump);
 
       for (auto *unwind : analysis.unwindOps)
         if (failed(genDominatingVars(rewriter, unwind, landingPadMap,

cudaq/lib/Target/CompileTarget.cpp

@@ -67,6 +67,9 @@ cudaq::CompileTarget cudaq::CompileTarget::createFromConfig(
 
   const auto &backendConfig =
       targetConfig.BackendConfig.value_or(defaultConfig);
+  if (!backendConfig.AOTUnwindMode.empty())
+    target.pipelineConfig.aotUnwindMode = backendConfig.AOTUnwindMode;
+
   auto prepPipeline = [&](const std::string &stage,
                           const std::string &stageName) {
     std::string pipeline = stage;
@@ -128,5 +131,5 @@ std::size_t std::hash<cudaq::CompileTarget::PipelineConfig>::operator()(
   return cudaq::detail::hashVal(
       pc.overridePassPipeline, pc.highLevelPipeline, pc.midLevelPipeline,
       pc.lowLevelPipeline, pc.codegenTranslation, pc.postCodeGenPasses,
-      pc.disableQubitMapping, pc.replaceStateWithKernel);
+      pc.aotUnwindMode, pc.disableQubitMapping, pc.replaceStateWithKernel);
 }

cudaq/lib/Target/Yaml/TargetConfigYaml.cpp

@@ -64,6 +64,9 @@ static std::string processSimBackendConfig(
     output << "JIT_LOW_LEVEL_PIPELINE=\"" << configValue.JITLowLevelPipeline
            << "\"\n";
 
+  if (!configValue.AOTUnwindMode.empty())
+    output << "AOT_UNWIND_MODE=\"" << configValue.AOTUnwindMode << "\"\n";
+
   if (!configValue.TargetPassPipeline.empty())
     output << "TARGET_PASS_PIPELINE=\"" << configValue.TargetPassPipeline
            << "\"\n";
@@ -284,8 +287,11 @@ cudaq::config::parseTargetConfig(std::string yamlContent,
   auto substitutedYamlContent =
       cudaq::config::substitutePluginRoot(std::move(yamlContent), pluginRoot);
   cudaq::config::TargetConfig config;
-  llvm::yaml::Input Input(substitutedYamlContent.c_str());
-  Input >> config;
+  llvm::yaml::Input input(substitutedYamlContent.c_str());
+  input >> config;
+  if (const auto error = input.error())
+    throw std::runtime_error("Invalid target configuration: " +
+                             error.message());
   return config;
 }
 
@@ -381,6 +387,7 @@ void MappingTraits<cudaq::config::BackendEndConfigEntry>::mapping(
   io.mapOptional("jit-high-level-pipeline", info.JITHighLevelPipeline);
   io.mapOptional("jit-mid-level-pipeline", info.JITMidLevelPipeline);
   io.mapOptional("jit-low-level-pipeline", info.JITLowLevelPipeline);
+  io.mapOptional("unwind-mode", info.AOTUnwindMode);
   io.mapOptional("target-pass-pipeline", info.TargetPassPipeline);
   io.mapOptional("codegen-emission", info.CodegenEmission);
   io.mapOptional("post-codegen-passes", info.PostCodeGenPasses);
@@ -397,6 +404,14 @@ void MappingTraits<cudaq::config::BackendEndConfigEntry>::mapping(
   io.mapOptional("rules", info.ConditionalBuildConfigs);
 }
 
+std::string MappingTraits<cudaq::config::BackendEndConfigEntry>::validate(
+    IO &io, cudaq::config::BackendEndConfigEntry &info) {
+  if (info.AOTUnwindMode.empty() || info.AOTUnwindMode == "cfg" ||
+      info.AOTUnwindMode == "dataflow" || info.AOTUnwindMode == "none")
+    return {};
+  return "'unwind-mode' must be one of: cfg, dataflow, none.";
+}
+
 void MappingTraits<cudaq::config::BackendFeatureMap>::mapping(
     IO &io, cudaq::config::BackendFeatureMap &info) {
   io.mapRequired("name", info.Name);

cudaq/test/Transforms/aot_run_vector_copy.qke

@@ -5,7 +5,8 @@
 // the terms of the Apache License 2.0 which accompanies this distribution.   //
 // ========================================================================== //
 
-// RUN: cudaq-opt --pass-pipeline='builtin.module(aot-prep-pipeline)' %s | FileCheck %s
+// RUN: cudaq-opt --pass-pipeline='builtin.module(python-aot-pipeline)' %s \
+// RUN:   | FileCheck %s
 
 // Regression test for the Python AOT pipeline's vector-return handling.
 // RunSemanticsHackery must remove the heap copy from the generated `.run`

cudaq/test/Transforms/aot_unwind_modes.qke

@@ -0,0 +1,52 @@
+// ========================================================================== //
+// Copyright (c) 2026 NVIDIA Corporation & Affiliates.                        //
+//                                                                            //
+// This source code and the accompanying materials are made available under   //
+// the terms of the Apache License 2.0 which accompanies this distribution.   //
+// ========================================================================== //
+
+// RUN: cudaq-opt --dump-pass-pipeline \
+// RUN:   --pass-pipeline='builtin.module(python-aot-pipeline{unwind-mode=dataflow})' \
+// RUN:   %s 2>&1 | FileCheck %s --check-prefix=DATAFLOW
+// RUN: cudaq-opt --dump-pass-pipeline \
+// RUN:   --pass-pipeline='builtin.module(python-aot-pipeline{unwind-mode=cfg})' \
+// RUN:   %s 2>&1 | FileCheck %s --check-prefix=CFG
+// RUN: cudaq-opt --dump-pass-pipeline \
+// RUN:   --pass-pipeline='builtin.module(python-aot-pipeline{unwind-mode=none})' \
+// RUN:   %s 2>&1 | FileCheck %s --check-prefix=NONE
+// RUN: not cudaq-opt \
+// RUN:   --pass-pipeline='builtin.module(python-aot-pipeline{unwind-mode=invalid})' \
+// RUN:   %s 2>&1 | FileCheck %s --check-prefix=INVALID
+// RUN: not cudaq-opt \
+// RUN:   --pass-pipeline='builtin.module(python-aot-pipeline{unwind-mode=none},lower-to-cfg)' \
+// RUN:   %s 2>&1 | FileCheck %s --check-prefix=NONE-ERROR
+
+module {
+  func.func @early_return(%arg0: i1) -> i32 {
+    %c1 = arith.constant 1 : i32
+    cc.if(%arg0) {
+      cc.unwind_return %c1 : i32
+    }
+    %c0 = arith.constant 0 : i32
+    return %c0 : i32
+  }
+}
+
+// DATAFLOW: Pass Manager with 28 passes:
+// DATAFLOW: variable-coalesce
+// DATAFLOW: unwind-by-dataflow
+// DATAFLOW: canonicalize
+
+// CFG: Pass Manager with 28 passes:
+// CFG: variable-coalesce
+// CFG: unwind-lowering
+// CFG: canonicalize
+
+// NONE: Pass Manager with 26 passes:
+// NONE-NOT: unwind-lowering
+// NONE-NOT: unwind-by-dataflow
+// NONE: get-concrete-matrix
+
+
+// INVALID: Cannot find option named 'invalid'
+// NONE-ERROR: error: 'cc.unwind_return' op

cudaq/tools/nvq++/nvq++.in

@@ -584,6 +584,7 @@ ENABLE_ARRAY_CONVERSION=true
 ENABLE_VARIABLE_COALESCE=true
 ENABLE_LEGACY_ADJOINT=false
 ENABLE_UNWIND_LOWERING=true
+AOT_UNWIND_MODE=cfg
 ENABLE_DEVICE_CODE_LOADER=true
 ENABLE_REALTIME_LOWERING=false
 ENABLE_KERNEL_EXECUTION=true
@@ -1036,8 +1037,21 @@ if ${ENABLE_VARIABLE_COALESCE}; then
 	OPT_PASSES="func.func(variable-coalesce)"
 fi
 if ${ENABLE_UNWIND_LOWERING}; then
-	RUN_OPT=true
-	OPT_PASSES=$(add_pass_to_pipeline "${OPT_PASSES}" "func.func(unwind-lowering)")
+	case "${AOT_UNWIND_MODE}" in
+	cfg)
+		RUN_OPT=true
+		OPT_PASSES=$(add_pass_to_pipeline "${OPT_PASSES}" "func.func(unwind-lowering)")
+		;;
+	dataflow)
+		RUN_OPT=true
+		OPT_PASSES=$(add_pass_to_pipeline "${OPT_PASSES}" "func.func(unwind-by-dataflow)")
+		;;
+	none)
+		;;
+	*)
+		error_exit "Invalid AOT unwind mode: (${AOT_UNWIND_MODE})"
+		;;
+	esac
 fi
 if ${ENABLE_INJECT_IMPLICIT_OUTPUT}; then
 	RUN_OPT=true
@@ -1078,9 +1092,15 @@ if ${ENABLE_KERNEL_EXECUTION}; then
 	fi
 fi
 if ${ENABLE_AGGRESSIVE_INLINE}; then
+	AGGRESSIVE_INLINE_OPTIONS=
+	# Do not let aggressive-inlining lower unwind operations after the AOT
+	# unwind stage deliberately left them in the IR.
+	if [[ "${AOT_UNWIND_MODE}" == "none" ]] || [[ "${ENABLE_UNWIND_LOWERING}" == "false" ]]; then
+		AGGRESSIVE_INLINE_OPTIONS="{lower-unwind=false}"
+	fi
 	RUN_OPT=true
 	if ${DO_LINK}; then
-		OPT_PASSES=$(add_pass_to_pipeline "${OPT_PASSES}" "aggressive-inlining")
+		OPT_PASSES=$(add_pass_to_pipeline "${OPT_PASSES}" "aggressive-inlining${AGGRESSIVE_INLINE_OPTIONS}")
 	else
 		OPT_PASSES=$(add_pass_to_pipeline "${OPT_PASSES}" "indirect-to-direct-calls,inline")
 	fi

python/cudaq/kernel/ast_bridge.py

@@ -34,7 +34,7 @@
 from .utils import (Color, boundaryDiagnostic, containsMeasureHandle,
                     globalRegisteredOperations, globalRegisteredTypes,
                     nvqppPrefix, mlirTypeFromAnnotation, mlirTypeFromPyType,
-                    getMLIRContext, is_recovered_value_ok,
+                    getAOTPassPipeline, getMLIRContext, is_recovered_value_ok,
                     recover_annotation_of_or_none, recover_value_of_or_none,
                     cudaq__unique_attr_name, mlirTryCreateStructType)
 
@@ -6284,8 +6284,7 @@ def compile_to_mlir(uniqueId, astModule, signature: KernelSignature, defFrame,
     # The `cudaq.pipeline.aot` span is the marker tooling uses to identify
     # pass events as AOT-pipeline (paired with `cudaq.pipeline.jit` emitted
     # from `QPU.cpp` `lower_to_qir_llvm`).
-    pm = PassManager.parse("builtin.module(aot-prep-pipeline)",
-                           context=bridge.ctx)
+    pm = PassManager.parse(getAOTPassPipeline(), context=bridge.ctx)
     try:
         with trace.span("cudaq.pipeline.aot"):
             cudaq_runtime.runPassManager(pm, bridge.module)