Request for kernel to return both state and results

[ej-psi]

Required prerequisites

• Search the issue tracker to check if your feature has already been mentioned or rejected in other issues.

Describe the feature

This is my team's single most important request. I'm available to discuss and help with this!

After running a kernel, we'd like to be able to get the result (the return value from the kernel) and also carry the resulting state forward for the next kernel.

Currently:

• If we call cudaq.run(kernel, state, args) on our kernel, we get the return value from the kernel, but not the resulting state.
• If we call cudaq.get_state(kernel, state, args) on our kernel, we get the resulting state, but not the return value.
• We'd really like both.

What we'd like, specifically:

• Call a kernel, and get the return values.
• Call another kernel, using the resulting state from the previous kernel, and get the return values.
• ...repeat thousands of times

Note that we don't actually need to get the state after each kernel call; if there's a way to just let it persist for the next kernel (like it would in the physical FTQC machine), that's great.

Why this is needed:

• In FTQC, mid-circuit measurements are essential. For example (see diagram), we compute a QROM (hundreds of gates), and then to uncompute it we measure the ancilla qubits and do some phase fixup which depends on the measurement results. This may happen thousands of times in a single program.

Our current workaround:

We have a solution, but it's not great.

1. We call state = cudaq.get_state(initial_state_kernel, num_qubits) to create the state
2. Then we call state = cudaq.get_state(ops_kernel, state, ops_list) with a list of QPU ops (X gates and such) to evolve the state.
3. Repeat step 2 as needed with more ops.
4. When the ops list includes measurements, then RIGHT after them we call results = cudaq.run(fetch_result_kernel, state, ops_list, shots_count=1) which re-measures the same qubits (this is wasteful but doesn't change the state), return the results.
5. Goto step 2

• Step 2 is potentially wasteful, if the state vector is copied from CPU to GPU, because we're not actually going to need to inspect it.
• Step 4 is definitely wasteful, because we wouldn't need to stop and do another kernel call to re-measure, if only we could get the return values (which contain the measurement results) the first time the instructions are run.

The good part

From the user's point of view, our workaround makes mid-circuit measurements easy, allowing users to branch their code and issue more gates on the same live state, based on measurement results. It's just messy under the hood at the moment.

[sacpis]

Seems like a request for an API change. @1tnguyen would you be able to review this request?

[1tnguyen]

Hi @ej-psi,

Thank you for your request.
From my understanding of the above workflow, it doesn't look like we need the state at all, i.e., we don't need to inspect it (e.g., for debugging purposes).

Have you considered an end-to-end workflow with mid-circuit measurements using CUDA-Q?
With cudaq::run, we can return arbitrary data from a kernel, such as mid-circuit measurement results.

Just as an example, I wrote this toy example, which includes mid-circuit measurements and conditional operations. The measurement results can be packed in a couple of ways, e.g., a flat buffer (like in the example), converted to integers (cudaq::to_integer), or a custom struct.

#include <iostream>
#include <cudaq.h>

__qpu__ std::vector<bool> kernel() {
  cudaq::qvector data_qubits(5);
  cudaq::qvector aux_qubits(5);
  const int num_iterations = 10;
  
  // Pre-allocate a vector to store all results from the iterations. 
  // Each iteration produces 5 measurement results, so we need num_iterations * 5 space.
  std::vector<bool> all_results(num_iterations * 5);
  
  // Some randon gates, just for testing purpose, you can replace it with any gates you want.
  for (int i = 0; i < 5; ++i) {
    h(data_qubits[i]);
    x<cudaq::ctrl>(data_qubits[i], aux_qubits[i]);
  }

  for (int iter = 0; iter < num_iterations; ++iter) {
    // Measure aux_qubits
    auto results = mz(aux_qubits);
    // Store results
    for (int i = 0; i < 5; ++i) {
      all_results[iter * 5 + i] = results[i];
    }

    // Conditional gates based on the measurement results
    for (int i = 0; i < 5; ++i) {
      // Gate operations based on the measurement results
      // This is just an example, you can replace it with any gates you want.
      if (results[i]) {
        x(data_qubits[i]);
      } else {
        z(data_qubits[i]);
      }
    }
  }

  return all_results;
}

int main() {
 
  std::vector<std::vector<bool>> results = cudaq::run(/*shots*/2, kernel);
  std::cout << "Results:" << std::endl;
  for (auto &result : results) {
    for (auto bit : result) {
      std::cout << bit << " ";
    }
    std::cout << std::endl;
  }
  return 0;
}

[ej-psi]

Hi @1tnguyen and thank you for your response!

Our applications are complex, making library calls (including 3rd party libraries) and often involve more than 100M QPU instructions. We're preparing them to run on physical FTQC hardware, and porting everything to run inside a kernel is impractical in this case.

If we could call cudaq.run() and have it continue evolving the same state as the previous kernel, that would also solve our issue.

Please lmk if I can help with more information or test cases!

[1tnguyen]

Our applications are complex, making library calls (including 3rd party libraries) and often involve more than 100M QPU instructions. We're preparing them to run on physical FTQC hardware, and porting everything to run inside a kernel is impractical in this case.

Thanks @ej-psi for the context information.
I want to add that we have cudaq::device_call utility to support making library calls from within the kernel. This would allow us to construct an entry-point kernel that lives entirely in the QPU domain.

If retrieving/reinstating the state is needed, I think we can get around the inefficiency w.r.t. 're-measure' qubits by arranging the get_state and run around the measure gates. For example, get_state before measure, then reinstate the state (in another kernel) to do the measurement. This would be really helpful if we had a concrete test case/sample to analyze.

@ej-psi Could you provide a simple test case demonstrating the application flow?

[ej-psi]

Hi @1tnguyen ,

If retrieving/reinstating the state is needed, I think we can get around the inefficiency w.r.t. 're-measure' qubits by arranging the get_state and run around the measure gates. For example, get_state before measure, then reinstate the state (in another kernel) to do the measurement. This would be really helpful if we had a concrete test case/sample to analyze.

This doesn't work, because we need both the partially-collapsed state and the measurement result in order to continue the program. So... if we do call get_state() before the measurement, This gives us the unmeasured state. If we then pass it to another kernel to perform the measurement, that kernel will, collapse the state randomly, and then we can either get the result (run()) or the state (get_state()), but not both.

Could you provide a simple test case demonstrating the application flow?

Our libraries aren't public yet, so it's tricky to post a complete example, but here's some code to demonstrate the application flow:

In this case, we're preparing a USP (just a uniform superposition of states) over a non-power-of-2 number of terms. We're going to repeat until success (RUS), and then continue afterward, using the prepared state in more computation.

This can be done in Workbench with the following code:

d = 18                        # The number of uniform nonzero terms we want
done = False
while not done:
    qreg.write(0)             # Clear the register
    qreg.had()                # Put it into a power-of-2 superposition
    with qreg < d as cond:    # Quantum-compare to the desired number of terms
        done = cond.read()    # Collapse the state and test for success

# more computation using qreg, now that it's prepared

...when we run it, the code runs a random number of times, looking like this:

So, following this case, here's what we wish we could do with CUDA-Q:

1. Create the initial state
2. Send the first batch of gates to the kernel (with the measurement) return the measurement result.
3. If it's a 1, we succeeded, so call the "more computation" using the resulting state.
4. If it's a zero, repeat the USP gates until the result is 1, then continue.

...but what we're doing now (to make it work) is:

1. Create the initial state
2. Send the first batch of gates to the kernel, stopping right after the measurement, return the resulting state.
3. Pass that state into a special kernel which repeats the measurement on the already-partially-collapsed state (so there should be no randomness here) and return the measurement result.
4. If it's a 1, we succeeded, so call the "more computation" using the resulting state.
5. If it's a zero, repeat the process until the result is 1, then continue.

Why it's important

Here, the measurement has some randomness to it. It collapses the state, and returns the result. At the moment, using CUDA-Q we can only get both the state and the result by repeating the measurement on the already-collapsed state.