Noise model enhancements

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/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

@@ -1316,6 +1316,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");
 
@@ -1334,12 +1366,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

python/runtime/common/py_NoiseModel.cpp

@@ -11,6 +11,7 @@
 #include "cudaq.h"
 #include <iostream>
 #include <pybind11/complex.h>
+#include <pybind11/functional.h>
 #include <pybind11/stl.h>
 
 namespace cudaq {
@@ -70,6 +71,36 @@ of the specified quantum operation.
   qubits (List[int]): The qubit/s to apply the noise channel to.
   channel (cudaq.KrausChannel): The :class:`KrausChannel` to apply 
     to the specified `operator` on the specified `qubits`.)#")
+      .def(
+          "add_channel",
+          [](noise_model &self, std::string &opName,
+             const noise_model::PredicateFuncTy &pre) {
+            self.add_channel(opName, pre);
+          },
+          py::arg("operator"), py::arg("pre"),
+          R"#(Add the given :class:`KrausChannel` generator callback to be applied after invocation 
+of the specified quantum operation.
+
+Args:
+  operator (str): The quantum operator to apply the noise channel to.
+  pre (Callable): The callback which takes qubits operands and gate parameters and returns a concrete :class:`KrausChannel` to apply 
+    to the specified `operator`.)#")
+      .def(
+          "add_all_qubit_channel",
+          [](noise_model &self, std::string &opName, kraus_channel &channel,
+             std::size_t num_controls = 0) {
+            self.add_all_qubit_channel(opName, channel, num_controls);
+          },
+          py::arg("operator"), py::arg("channel"), py::arg("num_controls") = 0,
+
+          R"#(Add the given :class:`KrausChannel` to be applied after invocation 
+of the specified quantum operation on arbitrary qubits.
+
+Args:
+  operator (str): The quantum operator to apply the noise channel to.
+  channel (cudaq.KrausChannel): The :class:`KrausChannel` to apply 
+    to the specified `operator` on any arbitrary qubits.
+  num_controls: Number of control bits. Default is 0 (no control bits).)#")
       .def(
           "get_channels",
           [](noise_model self, const std::string &op,
@@ -112,6 +143,7 @@ void bindNoiseChannels(py::module &mod) {
                             "The `KrausChannel` is composed of a list of "
                             ":class:`KrausOperator`'s and "
                             "is applied to a specific qubit or set of qubits.")
+      .def(py::init<>(), "Create an empty :class:`KrausChannel`")
       .def(py::init<std::vector<kraus_op>>(),
            "Create a :class:`KrausChannel` composed of a list of "
            ":class:`KrausOperator`'s.")
@@ -162,7 +194,7 @@ void bindNoiseChannels(py::module &mod) {
 
       For `probability = 0.0`, the channel will behave noise-free. 
       For `probability = 0.75`, the channel will fully depolarize the state.
-      For `proability = 1.0`, the channel will be uniform.)#")
+      For `probability = 1.0`, the channel will be uniform.)#")
       .def(py::init<double>(), py::arg("probability"),
            "Initialize the `DepolarizationChannel` with the provided "
            "`probability`.");