Noise model enhancements (#2168)

* Some prototyping code

* Prototype

* Revise the noise model

* Add example

* Update docstring:

* Code refactor

* Add tests

* Code format

* Add Python tests

* Code format

* Support custom op

- For the path via ExecutionManager: making sure the registration is
  early enough, i.e., at the point of registration.

- For the Python path: need to let the registry know about the
  registration.

- Propagate the custom op name all the way to the simulator.

* Fix a lit test

* Add docs

* Fix remote test

* Fix spelling

* Add a C++ example

docs/CMakeLists.txt

@@ -141,6 +141,7 @@ if (CUDAQ_ENABLE_PYTHON)
   add_pycudaq_test(DepolarizingNoise noise_depolarization.py)
   add_pycudaq_test(PhaseFlipNoise noise_phase_flip.py)
   add_pycudaq_test(KrausNoise noise_kraus_operator.py)
+  add_pycudaq_test(NoiseCallback noise_callback.py)
 
   if (CUTENSORNET_ROOT AND CUDA_FOUND) 
     # This example uses tensornet backend.

docs/sphinx/examples/cpp/basics/noise_callback.cpp

@@ -0,0 +1,67 @@
+// Compile and run with:
+// ```
+// nvq++ noise_callback.cpp --target density-matrix-cpu -o dyn.x
+// && ./dyn.x
+// ```
+//
+// Note: You must set the target to a density matrix backend for the noise
+// to successfully impact the system.
+
+#include <cudaq.h>
+#include <iostream>
+
+// CUDA-Q supports several different models of noise. In this
+// case, we will examine the dynamic noise channel specified as a callback
+// function.
+
+int main() {
+
+  // We will begin by defining an empty noise model that we will add
+  // our channel to.
+  cudaq::noise_model noise;
+  //  Noise model callback function
+  const auto rx_noise = [](const auto &qubits,
+                           const auto &params) -> cudaq::kraus_channel {
+    // Model a pulse-length based rotation gate:
+    // the bigger the angle, the longer the pulse, i.e., more amplitude damping.
+    auto angle = params[0];
+    // Normalize the angle into the [0, 2*pi] range
+    while (angle > 2. * M_PI)
+      angle -= 2. * M_PI;
+
+    while (angle < 0)
+      angle += 2. * M_PI;
+    // Damping rate is linearly proportional to the angle
+    const auto damping_rate = angle / (2. * M_PI);
+    std::cout << "Angle = " << params[0]
+              << ", amplitude damping rate = " << damping_rate << "\n";
+    return cudaq::amplitude_damping_channel(damping_rate);
+  };
+
+  // Bind the noise model callback function to the `rx` gate
+  noise.add_channel<cudaq::types::rx>(rx_noise);
+
+  auto kernel = [](double angle) __qpu__ {
+    cudaq::qubit q;
+    rx(angle, q);
+    mz(q);
+  };
+
+  // Now let's set the noise and we're ready to run the simulation!
+  cudaq::set_noise(noise);
+
+  // Our results should show measurements in both the |0> and |1> states,
+  // indicating that the noise has successfully impacted the system. Note: a
+  // `rx(pi)` is equivalent to a Pauli X gate, and thus, it should be in the |1>
+  // state if no noise is present.
+  auto noisy_counts = cudaq::sample(kernel, M_PI);
+  std::cout << "Noisy result:\n";
+  noisy_counts.dump();
+
+  // To confirm this, we can run the simulation again without noise.
+  cudaq::unset_noise();
+  auto noiseless_counts = cudaq::sample(kernel, M_PI);
+  std::cout << "Noiseless result:\n";
+  noiseless_counts.dump();
+  return 0;
+}

docs/sphinx/examples/python/noise_callback.py

@@ -0,0 +1,45 @@
+import cudaq
+import numpy as np
+
+# Set the target to our density matrix simulator.
+cudaq.set_target('density-matrix-cpu')
+
+noise = cudaq.NoiseModel()
+
+
+# Noise model callback function
+def rx_noise(qubits, params):
+    # Model a pulse-length based rotation gate:
+    # the bigger the angle, the longer the pulse, i.e., more amplitude damping.
+    angle = params[0]
+    angle = angle % (2 * np.pi)
+    # Damping rate is linearly proportional to the angle
+    damping_rate = np.abs(angle / (2 * np.pi))
+    print(f"Angle = {angle}, amplitude damping rate = {damping_rate}.")
+    return cudaq.AmplitudeDampingChannel(damping_rate)
+
+
+# Bind the noise model callback function to the `rx` gate
+noise.add_channel('rx', rx_noise)
+
+
+@cudaq.kernel
+def kernel(angle: float):
+    qubit = cudaq.qubit()
+    rx(angle, qubit)
+    mz(qubit)
+
+
+# Now we're ready to run the noisy simulation of our kernel.
+# Note: We must pass the noise model to sample via keyword.
+noisy_result = cudaq.sample(kernel, np.pi, noise_model=noise)
+print(noisy_result)
+
+# Our results should show measurements in both the |0> and |1> states, indicating
+# that the noise has successfully impacted the system.
+# Note: a `rx(pi)` is equivalent to a Pauli X gate, and thus, it should be
+# in the |1> state if no noise is present.
+
+# To confirm this, we can run the simulation again without noise.
+noiseless_result = cudaq.sample(kernel, np.pi)
+print(noiseless_result)

docs/sphinx/using/extending/_noise.rst

@@ -40,10 +40,67 @@ constructor should validate the completeness (CPTP) relation.
 A :code:`cudaq::noise_model` encapsulates a mapping of quantum operation names to a 
 vector of :code:`kraus_channel` that is to be applied after invocation of that 
 quantum operation. A :code:`noise_model` can be constructed with a nullary constructor, and 
-:code:`kraus_channels` can be added via a templated :code:`add_channel` method, where the 
-template type is the quantum operation the channel applies to (e.g. :code:`model.add_channel\<cudaq::types::h\>(channel)`). Clients (e.g. simulator backends) can retrieve the :code:`kraus_channel` to 
+:code:`kraus_channels` can be added via :code:`add_channel` and :code:`add_all_qubit_channel` methods with 
+the operation given as a string or as a template argument. 
+The operation name or the template type specifies the quantum operation the channel applies to 
+(e.g. :code:`model.add_channel\<cudaq::types::h\>(channel)` or :code:`model.add_channel("h", channel)`). 
+Clients (e.g. simulator backends) can retrieve the :code:`kraus_channel` to 
 apply to the simulated state via a :code:`noise_model::get_channel(...)` call. 
 
+When adding an error channel to a noise model for a quantum operation
+we can assign the noise channel to instances of that operation on specific qubit operands or 
+to any occurrence of the operation, regardless of which qubits it acts on. 
+
+.. tab:: Python
+
+    .. code-block:: python
+
+        # Add a noise channel to z gate on qubit 0
+        noise.add_channel('z', [0], noise_channel)
+        # Add a noise channel to x gate, regardless of qubit operands. 
+        noise.add_all_qubit_channel('x', noise_channel)
+
+
+.. tab:: C++
+
+    .. code-block:: cpp
+        
+        // Add a noise channel to z gate on qubit 0
+        noise.add_channel("z", {0}, noise_channel);
+        // Add a noise channel to x gate, regardless of qubit operands. 
+        noise.add_all_qubit_channel("x", noise_channel)
+
+In addition to static noise channels, users can also define a noise channel as a 
+callback function, which returns a concrete channel definition in terms of Kraus matrices 
+depending on the gate operands and gate parameters if any.
+
+.. tab:: Python
+
+    .. code-block:: python
+
+        # Noise channel callback function
+        def noise_cb(qubits, params):
+           # Construct a channel based on specific operands and parameters
+           ...
+           return noise_channel 
+        
+        # Add a dynamic noise channel to the 'rx' gate.
+        noise.add_channel('rx', noise_cb)
+
+
+.. tab:: C++
+
+    .. code-block:: cpp
+        
+        // Add a dynamic noise channel to the 'rx' gate.
+        noise.add_channel("rx",
+            [](const auto &qubits, const auto &params) -> cudaq::kraus_channel {
+                // Construct a channel based on specific operands and parameters
+                ...
+                return noiseChannel;
+            });
+
+
 Noise models can be constructed via the :code:`cudaq::noise_model` and specified for 
 execution via a public :code:`cudaq::set_noise(cudaq::noise_model&)` function. This function 
 should forward the :code:`noise_model` to the current :code:`quantum_platform` which can attach it 

lib/Optimizer/CodeGen/QuakeToLLVM.cpp

@@ -1317,6 +1317,38 @@ class CustomUnitaryOpRewrite
     StringRef generatorName = sref.getRootReference();
     auto globalOp =
         parentModule.lookupSymbol<cudaq::cc::GlobalOp>(generatorName);
+    const auto customOpName = [&]() -> std::string {
+      auto globalName = generatorName.str();
+      // IMPORTANT: this must match the logic to generate global data
+      // globalName = f'{nvqppPrefix}{opName}_generator_{numTargets}.rodata'
+      const std::string nvqppPrefix = "__nvqpp__mlirgen__";
+      const std::string generatorSuffix = "_generator";
+      if (globalName.starts_with(nvqppPrefix)) {
+        globalName = globalName.substr(nvqppPrefix.size());
+        const size_t pos = globalName.find(generatorSuffix);
+        if (pos != std::string::npos)
+          return globalName.substr(0, pos);
+      }
+
+      return "";
+    }();
+
+    // Create a global string for the op name
+    auto insertPoint = rewriter.saveInsertionPoint();
+    rewriter.setInsertionPointToStart(parentModule.getBody());
+    // Create the custom op name global
+    auto builder = cudaq::IRBuilder::atBlockEnd(parentModule.getBody());
+    auto opNameGlobal =
+        builder.genCStringLiteralAppendNul(loc, parentModule, customOpName);
+    // Shift back to the function
+    rewriter.restoreInsertionPoint(insertPoint);
+    // Get the string address and bit cast
+    auto opNameRef = rewriter.create<LLVM::AddressOfOp>(
+        loc, cudaq::opt::factory::getPointerType(opNameGlobal.getType()),
+        opNameGlobal.getSymName());
+    auto castedOpNameRef = rewriter.create<LLVM::BitcastOp>(
+        loc, cudaq::opt::factory::getPointerType(context), opNameRef);
+
     if (!globalOp)
       return op.emitOpError("global not found for custom op");
 
@@ -1335,12 +1367,13 @@ class CustomUnitaryOpRewrite
         cudaq::opt::factory::createLLVMFunctionSymbol(
             qirFunctionName, LLVM::LLVMVoidType::get(context),
             {complex64PtrTy, cudaq::opt::getArrayType(context),
-             cudaq::opt::getArrayType(context)},
+             cudaq::opt::getArrayType(context),
+             LLVM::LLVMPointerType::get(rewriter.getI8Type())},
             parentModule);
 
     rewriter.replaceOpWithNewOp<LLVM::CallOp>(
         op, TypeRange{}, customSymbolRef,
-        ValueRange{unitaryData, controlArr, targetArr});
+        ValueRange{unitaryData, controlArr, targetArr, castedOpNameRef});
 
     return success();
   }

python/cudaq/kernel/register_op.py

@@ -13,6 +13,7 @@
 
 from .utils import globalRegisteredOperations
 from .kernel_builder import PyKernel, __generalCustomOperation
+from ..mlir._mlir_libs._quakeDialects import cudaq_runtime
 
 
 def register_operation(operation_name: str, unitary):
@@ -58,5 +59,8 @@ def kernel():
     # Make available to kernel builder object
     setattr(PyKernel, operation_name,
             partialmethod(__generalCustomOperation, operation_name))
+    # Let the runtime know about this registered operation.
+    # Note: the matrix generator/construction is not known by the ExecutionManager in this case since we don't expect the ExecutionManager to be involved.
+    cudaq_runtime.register_custom_operation(operation_name)
 
     return

python/extension/CMakeLists.txt

@@ -48,6 +48,7 @@ declare_mlir_python_extension(CUDAQuantumPythonSources.Extension
     ../runtime/common/py_NoiseModel.cpp
     ../runtime/common/py_ObserveResult.cpp
     ../runtime/common/py_SampleResult.cpp
+    ../runtime/common/py_CustomOpRegistry.cpp
     ../runtime/cudaq/algorithms/py_draw.cpp
     ../runtime/cudaq/algorithms/py_observe_async.cpp
     ../runtime/cudaq/algorithms/py_optimizer.cpp

python/extension/CUDAQuantumExtension.cpp

@@ -10,6 +10,7 @@
 #include "cudaq.h"
 #include "cudaq/Support/Version.h"
 #include "cudaq/platform/orca/orca_qpu.h"
+#include "runtime/common/py_CustomOpRegistry.h"
 #include "runtime/common/py_ExecutionContext.h"
 #include "runtime/common/py_NoiseModel.h"
 #include "runtime/common/py_ObserveResult.h"
@@ -103,6 +104,7 @@ PYBIND11_MODULE(_quakeDialects, m) {
   cudaq::bindVQE(cudaqRuntime);
   cudaq::bindAltLaunchKernel(cudaqRuntime);
   cudaq::bindTestUtils(cudaqRuntime, *holder.get());
+  cudaq::bindCustomOpRegistry(cudaqRuntime);
 
   cudaqRuntime.def("set_random_seed", &cudaq::set_random_seed,
                    "Provide the seed for backend quantum kernel simulation.");

python/runtime/common/py_CustomOpRegistry.cpp

@@ -0,0 +1,34 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates.                  *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+#include "py_CustomOpRegistry.h"
+#include "common/CustomOp.h"
+#include <pybind11/complex.h>
+#include <pybind11/functional.h>
+#include <pybind11/stl.h>
+
+namespace cudaq {
+struct py_unitary_operation : public unitary_operation {
+  std::vector<std::complex<double>>
+  unitary(const std::vector<double> &parameters =
+              std::vector<double>()) const override {
+    throw std::runtime_error("Attempt to invoke the placeholder for Python "
+                             "unitary op. This is illegal.");
+    return {};
+  }
+};
+
+void bindCustomOpRegistry(py::module &mod) {
+  mod.def(
+      "register_custom_operation",
+      [&](const std::string &opName) {
+        cudaq::customOpRegistry::getInstance()
+            .registerOperation<py_unitary_operation>(opName);
+      },
+      "Register a custom operation");
+}
+} // namespace cudaq

python/runtime/common/py_CustomOpRegistry.h

@@ -0,0 +1,16 @@
+/****************************************************************-*- C++ -*-****
+ * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates.                  *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+namespace cudaq {
+/// @brief Bind the custom operation registry to Python.
+void bindCustomOpRegistry(py::module &mod);
+} // namespace cudaq