Avoid segfault in ExecutionContext.__exit__ when kernel throws (#4066)

When a kernel raises inside a with ExecutionContext block (cudaq.observe
called with a list argument that is too short), the __exit__ exception
was still calling platform.finalizeExecutionContext(). On the custatevec
GPU backend this triggers an API call against an uninitialized simulator
state, causing a segfault.

Update: Have put back the `finalizeExecutionContext` in __exit__ and
handling the 0 allocated qubits inside simulator.

Found second issue. When there are 0 qubit allocated, calling
setToZeroState (in endExecution) on an unitialized GPU state causes
segfault as was evident from the failing jobs in the first run of this
PR.

Adding a fix to handle a case for the observe path via
handleObservation.

Fixes #4062

---------

Signed-off-by: Sachin Pisal <spisal@nvidia.com>

Author: sacpis

python/tests/builder/test_observe.py

@@ -576,6 +576,38 @@ def test_observe_numpy_array(angles, want_state, want_expectation):
         cudaq.observe(kernel, hamiltonian, bad_params, qpu_id=0, shots_count=10)
 
 
+def test_observe_kernel_exception_no_segfault():
+    """
+    When cudaq.observe is called with list arguments that are
+    too short, the kernel raises a RuntimeError inside the ExecutionContext.
+    Before this change, ExecutionContext.__exit__ would call finalizeExecutionContext on
+    an uninitialized simulator state causing a segfault on GPU backend.
+    """
+    qubit_count = 4
+    kernel, thetas = cudaq.make_kernel(list)
+    qreg = kernel.qalloc(qubit_count)
+    for i in range(qubit_count):
+        kernel.rx(thetas[i], qreg[i])
+    hamiltonian = spin.z(0) + spin.z(1) + spin.z(2) + spin.z(3)
+
+    with pytest.raises(Exception) as exception:
+        cudaq.observe(kernel,
+                      hamiltonian,
+                      np.random.uniform(size=(3, 2)),
+                      shots_count=10)
+    assert "Invalid runtime list argument" in str(exception.value)
+
+    with pytest.raises(Exception) as exception:
+        cudaq.observe(kernel,
+                      hamiltonian,
+                      np.random.uniform(size=(2,)),
+                      shots_count=10)
+    assert "Invalid runtime list argument" in str(exception.value)
+
+    result = cudaq.observe(kernel, hamiltonian, [np.pi] * qubit_count)
+    assert np.isclose(result.expectation(), -4.0, atol=1e-12)
+
+
 def test_observe_n():
     """
     Test that we can broadcast the observe call over a number of argument sets

runtime/nvqir/CircuitSimulator.h

@@ -966,6 +966,9 @@ class CircuitSimulatorBase : public CircuitSimulator {
 
   /// @brief Reset the current execution context.
   void finalizeExecutionContext(cudaq::ExecutionContext &context) override {
+    if (nQubitsAllocated == 0 && context.name != "sample")
+      return;
+
     // Get the ExecutionContext name
     auto execContextName = context.name;
 
@@ -1045,6 +1048,11 @@ class CircuitSimulatorBase : public CircuitSimulator {
 
   /// @brief Clean up state after execution ends
   void endExecution() override {
+    if (nQubitsAllocated == 0) {
+      tracker = {};
+      return;
+    }
+
     bool shouldSetToZero = cudaq::isInBatchMode() && !cudaq::isLastBatch();
 
     // Reset the state if we've deallocated all qubits.
@@ -1286,6 +1294,12 @@ class CircuitSimulatorBase : public CircuitSimulator {
   void measureSpinOp(const cudaq::spin_op &op) override {
     auto executionContext = cudaq::getExecutionContext();
 
+    if (nQubitsAllocated == 0) {
+      if (executionContext)
+        executionContext->expectationValue = 0.0;
+      return;
+    }
+
     flushGateQueue();
 
     if (executionContext->canHandleObserve) {