Add operations for setting state and basis state and integrate this with skip_initial_state_prep

.dep-versions

@@ -15,4 +15,4 @@ pennylane=0.38.0.dev11
 # 'runtime/Makefile' and at all GitHub workflows, using the exact
 # commit hash corresponding to the merged PR that implements the
 # desired feature.
-lightning=0.38.0-dev26
+lightning=0.38.0-dev32

.github/workflows/build-wheel-linux-x86_64.yaml

@@ -342,7 +342,7 @@ jobs:
               -DPYTHON_EXECUTABLE=$(which python${{ matrix.python_version }}) \
               -Dpybind11_DIR=$(python${{ matrix.python_version }} -c "import pybind11; print(pybind11.get_cmake_dir())") \
               -DENABLE_LAPACK=OFF \
-              -DLIGHTNING_GIT_TAG=8f517d24c71c6c7765f5c1bf29b0264b951de96a \
+              -DLIGHTNING_GIT_TAG=c6b86a5 \
               -DENABLE_WARNINGS=OFF \
               -DENABLE_OPENQASM=ON \
               -DENABLE_OPENMP=OFF \

.github/workflows/build-wheel-macos-arm64.yaml

@@ -301,7 +301,7 @@ jobs:
               -DCMAKE_LIBRARY_OUTPUT_DIRECTORY=$GITHUB_WORKSPACE/runtime-build/lib \
               -DPYTHON_EXECUTABLE=$(which python${{ matrix.python_version }}) \
               -Dpybind11_DIR=$(python${{ matrix.python_version }} -c "import pybind11; print(pybind11.get_cmake_dir())") \
-              -DLIGHTNING_GIT_TAG=8f517d24c71c6c7765f5c1bf29b0264b951de96a \
+              -DLIGHTNING_GIT_TAG=c6b86a5 \
               -DENABLE_LAPACK=OFF \
               -DENABLE_WARNINGS=OFF \
               -DENABLE_OPENQASM=ON \

.github/workflows/build-wheel-macos-x86_64.yaml

@@ -292,7 +292,7 @@ jobs:
               -DCMAKE_LIBRARY_OUTPUT_DIRECTORY=$GITHUB_WORKSPACE/runtime-build/lib \
               -DPYTHON_EXECUTABLE=$(which python${{ matrix.python_version }}) \
               -Dpybind11_DIR=$(python${{ matrix.python_version }} -c "import pybind11; print(pybind11.get_cmake_dir())") \
-              -DLIGHTNING_GIT_TAG=8f517d24c71c6c7765f5c1bf29b0264b951de96a \
+              -DLIGHTNING_GIT_TAG=c6b86a5 \
               -DENABLE_LAPACK=OFF \
               -DENABLE_WARNINGS=OFF \
               -DENABLE_OPENQASM=ON \

.github/workflows/scripts/linux_arm64/rh8/build_catalyst.sh

@@ -45,7 +45,7 @@ cmake -S runtime -B runtime-build -G Ninja \
     -DPYTHON_INCLUDE_DIR=/opt/_internal/cpython-${PYTHON_VERSION}.${PYTHON_SUBVERSION}/include/python${PYTHON_VERSION} \
     -DPYTHON_LIBRARY=/opt/_internal/cpython-${PYTHON_VERSION}.${PYTHON_SUBVERSION}/lib \
     -Dpybind11_DIR=/opt/_internal/cpython-${PYTHON_VERSION}.${PYTHON_SUBVERSION}/lib/python${PYTHON_VERSION}/site-packages/pybind11/share/cmake/pybind11 \
-    -DLIGHTNING_GIT_TAG=8f517d24c71c6c7765f5c1bf29b0264b951de96a \
+    -DLIGHTNING_GIT_TAG=c6b86a5 \
     -DENABLE_LAPACK=OFF \
     -DENABLE_WARNINGS=OFF \
     -DENABLE_OPENQASM=ON \

demos/adaptive_circuits_demo.ipynb

@@ -70,6 +70,7 @@
     "\n",
     "import jax.numpy as jnp\n",
     "\n",
+    "import functools\n",
     "import warnings\n",
     "warnings.filterwarnings('ignore')"
    ]
@@ -150,6 +151,13 @@
     "# Define the device, using lightning.qubit device\n",
     "dev = qml.device(\"lightning.qubit\", wires=qubits)\n",
     "\n",
+    "def stopping_condition(obj):\n",
+    "    return obj.name not in [\"BasisState\"]\n",
+    "\n",
+    "def decompose(stopping_condition, func):\n",
+    "    return qml.devices.preprocess.decompose(func, stopping_condition, skip_initial_state_prep=False)\n",
+    "\n",
+    "@functools.partial(decompose, stopping_condition)\n",
     "@qml.qnode(dev, diff_method=\"adjoint\")\n",
     "def cost_func(params):\n",
     "    qml.BasisState(hf, wires=range(qubits))\n",
@@ -237,6 +245,7 @@
     "hf = qml.qchem.hf_state(electrons=2, orbitals=6)\n",
     "print(f\"The Hartree-Fock State: {hf}\")\n",
     "\n",
+    "@functools.partial(decompose, stopping_condition)\n",
     "@qml.qnode(qml.device(\"lightning.qubit\", wires=qubits))\n",
     "def catalyst_cost_func(params):\n",
     "    qml.BasisState(hf, wires=range(qubits))\n",
@@ -571,22 +580,22 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Excitation: [0, 1, 2, 3], Gradient: -0.012782168180745743\n",
+      "Excitation: [0, 1, 2, 3], Gradient: -0.012782168180745672\n",
       "Excitation: [0, 1, 2, 5], Gradient: 0.0\n",
       "Excitation: [0, 1, 2, 7], Gradient: 0.0\n",
-      "Excitation: [0, 1, 2, 9], Gradient: 0.03426450359905824\n",
+      "Excitation: [0, 1, 2, 9], Gradient: 0.03426450359905814\n",
       "Excitation: [0, 1, 3, 4], Gradient: 0.0\n",
       "Excitation: [0, 1, 3, 6], Gradient: 0.0\n",
-      "Excitation: [0, 1, 3, 8], Gradient: -0.034264503599058235\n",
-      "Excitation: [0, 1, 4, 5], Gradient: -0.023581524379111044\n",
+      "Excitation: [0, 1, 3, 8], Gradient: -0.03426450359905814\n",
+      "Excitation: [0, 1, 4, 5], Gradient: -0.02358152437911096\n",
       "Excitation: [0, 1, 4, 7], Gradient: 0.0\n",
       "Excitation: [0, 1, 4, 9], Gradient: 0.0\n",
       "Excitation: [0, 1, 5, 6], Gradient: 0.0\n",
       "Excitation: [0, 1, 5, 8], Gradient: 0.0\n",
-      "Excitation: [0, 1, 6, 7], Gradient: -0.02358152437912195\n",
+      "Excitation: [0, 1, 6, 7], Gradient: -0.023581524379121893\n",
       "Excitation: [0, 1, 6, 9], Gradient: 0.0\n",
       "Excitation: [0, 1, 7, 8], Gradient: 0.0\n",
-      "Excitation: [0, 1, 8, 9], Gradient: -0.12362273289626559\n"
+      "Excitation: [0, 1, 8, 9], Gradient: -0.12362273289626545\n"
      ]
     }
    ],
@@ -606,6 +615,7 @@
     "# Define the device and the cost function.\n",
     "dev = qml.device(\"lightning.qubit\", wires=qubits)\n",
     "cost_fn = qml.QNode(circuit_1, dev, diff_method=\"adjoint\")\n",
+    "cost_fn = decompose(stopping_condition, cost_fn)\n",
     "\n",
     "circuit_gradient = qml.grad(cost_fn, argnum=0)\n",
     "\n",
@@ -705,14 +715,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Excitation : [0, 2], Gradient: -0.005062544629162063\n",
+      "Excitation : [0, 2], Gradient: 0.005062544629162047\n",
       "Excitation : [0, 4], Gradient: 0.0\n",
       "Excitation : [0, 6], Gradient: 0.0\n",
-      "Excitation : [0, 8], Gradient: -0.0009448055879247612\n",
-      "Excitation : [1, 3], Gradient: 0.004926625112976861\n",
+      "Excitation : [0, 8], Gradient: 0.0009448055879244589\n",
+      "Excitation : [1, 3], Gradient: -0.004926625112976867\n",
       "Excitation : [1, 5], Gradient: 0.0\n",
       "Excitation : [1, 7], Gradient: 0.0\n",
-      "Excitation : [1, 9], Gradient: 0.0014535553867889794\n",
+      "Excitation : [1, 9], Gradient: -0.0014535553867887066\n",
       "singles_select: [[0, 2], [0, 8], [1, 3], [1, 9]]\n"
      ]
     }
@@ -739,6 +749,7 @@
     "    return qml.expval(hamiltonian)\n",
     "\n",
     "cost_fn = qml.QNode(circuit_2, dev, diff_method=\"adjoint\")\n",
+    "cost_fn = decompose(stopping_condition, cost_fn)\n",
     "circuit_gradient = qml.grad(cost_fn, argnum=0)\n",
     "params = [0.0] * len(singles)\n",
     "\n",
@@ -820,6 +831,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "@functools.partial(decompose, stopping_condition)\n",
     "@qml.qnode(qml.device(\"lightning.qubit\", wires=qubits))\n",
     "def catalyst_avqe_circuit(\n",
     "    num_gates, num_selected_gates,\n",
@@ -884,22 +896,22 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Excitation: [0 1 2 3], Gradient: -0.012782148672840776\n",
-      "Excitation: [0 1 2 5], Gradient: 1.7763568394002505e-08\n",
-      "Excitation: [0 1 2 7], Gradient: 1.7763568394002505e-08\n",
-      "Excitation: [0 1 2 9], Gradient: 0.034264520110127705\n",
-      "Excitation: [0 1 3 4], Gradient: 1.7763568394002505e-08\n",
-      "Excitation: [0 1 3 6], Gradient: 1.7763568394002505e-08\n",
-      "Excitation: [0 1 3 8], Gradient: -0.03426448458299092\n",
-      "Excitation: [0 1 4 5], Gradient: -0.023581492314406205\n",
-      "Excitation: [0 1 4 7], Gradient: 1.7763568394002505e-08\n",
-      "Excitation: [0 1 4 9], Gradient: 1.7763568394002505e-08\n",
-      "Excitation: [0 1 5 6], Gradient: 1.7763568394002505e-08\n",
-      "Excitation: [0 1 5 8], Gradient: 1.7763568394002505e-08\n",
-      "Excitation: [0 1 6 7], Gradient: -0.023581501196190402\n",
-      "Excitation: [0 1 6 9], Gradient: 8.881784197001252e-09\n",
-      "Excitation: [0 1 7 8], Gradient: 8.881784197001252e-09\n",
-      "Excitation: [0 1 8 9], Gradient: -0.1236227031853332\n"
+      "Excitation: [0 1 2 3], Gradient: -0.012782095382135594\n",
+      "Excitation: [0 1 2 5], Gradient: 7.105427357601002e-08\n",
+      "Excitation: [0 1 2 7], Gradient: 7.105427357601002e-08\n",
+      "Excitation: [0 1 2 9], Gradient: 0.03426457340083289\n",
+      "Excitation: [0 1 3 4], Gradient: 7.105427357601002e-08\n",
+      "Excitation: [0 1 3 6], Gradient: 9.769962616701378e-08\n",
+      "Excitation: [0 1 3 8], Gradient: -0.034264404646933144\n",
+      "Excitation: [0 1 4 5], Gradient: -0.023581412378348432\n",
+      "Excitation: [0 1 4 7], Gradient: 9.769962616701378e-08\n",
+      "Excitation: [0 1 4 9], Gradient: 7.105427357601002e-08\n",
+      "Excitation: [0 1 5 6], Gradient: 7.105427357601002e-08\n",
+      "Excitation: [0 1 5 8], Gradient: 7.105427357601002e-08\n",
+      "Excitation: [0 1 6 7], Gradient: -0.023581412378348432\n",
+      "Excitation: [0 1 6 9], Gradient: 6.217248937900877e-08\n",
+      "Excitation: [0 1 7 8], Gradient: 6.217248937900877e-08\n",
+      "Excitation: [0 1 8 9], Gradient: -0.12362264989462801\n"
      ]
     }
    ],
@@ -986,7 +998,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "params_doubles: [ 0.04758599 -0.09855459  0.09865966  0.05330077  0.05331789  0.22938115]\n"
+      "params_doubles: [ 0.04758601 -0.09855455  0.09865969  0.0533007   0.0533178   0.22938115]\n"
      ]
     }
    ],

doc/changelog.md

@@ -283,6 +283,10 @@
   Catalyst variant jaxpr.
   [(#837)](https://github.com/PennyLaneAI/catalyst/pull/837)
 
+* On devices that support it, initial state preparation routines `qml.StatePrep` and `qml.BasisState`
+  are no longer decomposed when using Catalyst, improving compilation & runtime performance.
+  [(#955)](https://github.com/PennyLaneAI/catalyst/pull/955)
+
 <h3>Breaking changes</h3>
 
 * Return values of qjit-compiled functions that were previously `numpy.ndarray` are now of type
@@ -299,6 +303,12 @@
   `extrapolate_kwargs` keyword argument in `mitigate_with_zne`.
   [(#806)](https://github.com/PennyLaneAI/catalyst/pull/806)
 
+* The QuantumDevice API has now added the functions `SetState` and `SetBasisState`
+  for simulators that may benefit from instructions that directly set the state.
+  Implementing these methods is optional, and device support can be indicated via
+  the `initial_state_prep` flag in the TOML configuration file.
+  [(#955)](https://github.com/PennyLaneAI/catalyst/pull/955)
+
 <h3>Bug fixes</h3>
 
 * Catalyst no longer generates a `QubitUnitary` operation during decomposition if a device doesn't

doc/requirements.txt

@@ -30,6 +30,6 @@ lxml_html_clean
 
 # Pre-install PL development wheels
 --extra-index-url https://test.pypi.org/simple/
-pennylane-lightning-kokkos==0.38.0-dev26
-pennylane-lightning==0.38.0-dev26
+pennylane-lightning-kokkos==0.38.0-dev32
+pennylane-lightning==0.38.0-dev32
 pennylane==0.38.0.dev11

frontend/catalyst/api_extensions/control_flow.py

@@ -1107,7 +1107,9 @@ def trace_quantum(self, ctx, device, trace, qrp) -> QRegPromise:
         op = self
         for region in op.regions:
             with EvaluationContext.frame_tracing_context(ctx, region.trace):
-                qreg_in = _input_type_to_tracers(region.trace.new_arg, [AbstractQreg()])[0]
+                reg_len = qrp.base.length
+                new_qreg = AbstractQreg(reg_len)
+                qreg_in = _input_type_to_tracers(region.trace.new_arg, [new_qreg])[0]
                 qreg_out = trace_quantum_operations(
                     region.quantum_tape, device, qreg_in, ctx, region.trace
                 ).actualize()
@@ -1165,7 +1167,9 @@ def trace_quantum(self, ctx, device, trace, qrp) -> QRegPromise:
         expansion_strategy = self.expansion_strategy
 
         with EvaluationContext.frame_tracing_context(ctx, inner_trace):
-            qreg_in = _input_type_to_tracers(inner_trace.new_arg, [AbstractQreg()])[0]
+            reg_len = qrp.base.length
+            new_qreg = AbstractQreg(reg_len)
+            qreg_in = _input_type_to_tracers(inner_trace.new_arg, [new_qreg])[0]
             qrp_out = trace_quantum_operations(inner_tape, device, qreg_in, ctx, inner_trace)
             qreg_out = qrp_out.actualize()
 
@@ -1245,7 +1249,7 @@ def trace_quantum(self, ctx, device, trace, qrp) -> QRegPromise:
                 res_classical_tracers,
                 expansion_strategy=expansion_strategy,
             )
-            _input_type_to_tracers(cond_trace.new_arg, [AbstractQreg()])
+            _input_type_to_tracers(cond_trace.new_arg, [AbstractQreg(qrp.base.length)])
             cond_jaxpr, _, cond_consts = trace_to_jaxpr(
                 cond_trace, arg_expanded_classical_tracers, res_expanded_classical_tracers
             )
@@ -1256,7 +1260,7 @@ def trace_quantum(self, ctx, device, trace, qrp) -> QRegPromise:
         with EvaluationContext.frame_tracing_context(ctx, body_trace):
             region = self.regions[1]
             res_classical_tracers = region.res_classical_tracers
-            qreg_in = _input_type_to_tracers(body_trace.new_arg, [AbstractQreg()])[0]
+            qreg_in = _input_type_to_tracers(body_trace.new_arg, [AbstractQreg(qrp.base.length)])[0]
             qrp_out = trace_quantum_operations(body_tape, device, qreg_in, ctx, body_trace)
             qreg_out = qrp_out.actualize()
             arg_expanded_tracers = expand_args(

frontend/catalyst/api_extensions/quantum_operators.py

@@ -465,7 +465,7 @@ def trace_quantum(self, ctx, device, _trace, qrp) -> QRegPromise:
             frame_ctx = EvaluationContext.frame_tracing_context(ctx, body_trace)
 
         with frame_ctx as body_trace:
-            qreg_in = _input_type_to_tracers(body_trace.new_arg, [AbstractQreg()])[0]
+            qreg_in = _input_type_to_tracers(body_trace.new_arg, [AbstractQreg(qrp.base.length)])[0]
             qrp_out = trace_quantum_operations(body_tape, device, qreg_in, ctx, body_trace)
             qreg_out = qrp_out.actualize()
             body_jaxpr, _, body_consts = ctx.frames[body_trace].to_jaxpr2(

frontend/catalyst/device/decomposition.py

@@ -108,7 +108,7 @@ def catalyst_decompose(
     (toplevel_tape,), _ = decompose(
         tape,
         stopping_condition,
-        skip_initial_state_prep=False,
+        skip_initial_state_prep=capabilities.initial_state_prep_flag,
         decomposer=partial(catalyst_decomposer, capabilities=capabilities),
         max_expansion=max_expansion,
         name="catalyst on this device",

frontend/catalyst/device/verification.py

@@ -30,7 +30,7 @@
     VnEntropyMP,
 )
 from pennylane.measurements.shots import Shots
-from pennylane.operation import Operation, Tensor
+from pennylane.operation import Operation, StatePrepBase, Tensor
 from pennylane.ops import (
     Adjoint,
     CompositeOp,
@@ -209,7 +209,13 @@ def _op_checker(op, state):
         # is handled in _inv_op_checker and _ctrl_op_checker.
         # Specialed control op classes (e.g. CRZ) should be checked directly though, which is why we
         # can't use isinstance(op, Controlled).
-        if type(op) in (Controlled, ControlledOp) or isinstance(op, Adjoint):
+        if type(op) in (Controlled, ControlledOp) or isinstance(op, (Adjoint)):
+            pass
+        # Don't check StatePrep since StatePrep is not in the list of device capabilities.
+        # It is only valid when the TOML file has the initial_state_prep_flag.
+        elif (
+            isinstance(op, StatePrepBase) and qjit_device.qjit_capabilities.initial_state_prep_flag
+        ):
             pass
         elif not qjit_device.qjit_capabilities.native_ops.get(op.name):
             raise CompileError(

frontend/catalyst/from_plxpr.py

@@ -235,7 +235,7 @@ def setup(self):
         resets the wire map.
         """
         qdevice_p.bind(**_get_device_kwargs(self._device))
-        self.qreg = qalloc_p.bind(len(self._device.wires))
+        self.qreg = qalloc_p.bind(len(self._device.wires), static_size=len(self._device.wires))
         self.wire_map = {}
 
     def cleanup(self):