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

@@ -152,6 +152,7 @@
     "\n",
     "@qml.qnode(dev, diff_method=\"adjoint\")\n",
     "def cost_func(params):\n",
+    "    qml.RX(0.0, wires=[0])\n",
     "    qml.BasisState(hf, wires=range(qubits))\n",
     "    qml.DoubleExcitation(params[0], wires=[0, 1, 2, 3])\n",
     "    qml.DoubleExcitation(params[1], wires=[0, 1, 4, 5])\n",
@@ -239,6 +240,7 @@
     "\n",
     "@qml.qnode(qml.device(\"lightning.qubit\", wires=qubits))\n",
     "def catalyst_cost_func(params):\n",
+    "    qml.RX(0.0, wires=[0])\n",
     "    qml.BasisState(hf, wires=range(qubits))\n",
     "    qml.DoubleExcitation(params[0], wires=[0, 1, 2, 3])\n",
     "    qml.DoubleExcitation(params[1], wires=[0, 1, 4, 5])\n",
@@ -594,6 +596,7 @@
     "# Create a circuit that applies a selected group of gates\n",
     "# to a reference Hartree-Fock state.\n",
     "def circuit_1(params, excitations):\n",
+    "    qml.RX(0.0, wires=[0])\n",
     "    qml.BasisState(hf, wires=range(qubits))\n",
     "\n",
     "    for i, excitation in enumerate(excitations):\n",
@@ -723,6 +726,7 @@
     "# those that have a non-negligible gradient.\n",
     "# Repeat steps 1 and 2 for the single excitations.\n",
     "def circuit_2(params, excitations, gates_select, params_select):\n",
+    "    qml.RX(0.0, wires=[0])\n",
     "    qml.BasisState(hf, wires=range(qubits))\n",
     "\n",
     "    for i, gate in enumerate(gates_select):\n",
@@ -828,6 +832,7 @@
     "    selected_double_gates, selected_single_gates,\n",
     "    gates_verifier, selected_verifier,\n",
     "    apply_selected):\n",
+    "    qml.RX(0.0, wires=[0])\n",
     "    qml.BasisState(hf, wires=range(qubits))\n",
     "\n",
     "    # Used in steps 4-5\n",

doc/changelog.md

@@ -249,6 +249,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
@@ -265,6 +269,13 @@
   `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`
+* 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

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):

frontend/catalyst/jax_primitives.py

@@ -85,6 +85,8 @@
     ProbsOp,
     QubitUnitaryOp,
     SampleOp,
+    SetBasisStateOp,
+    SetStateOp,
     StateOp,
     TensorOp,
     VarianceOp,
@@ -143,6 +145,9 @@ class AbstractQreg(AbstractValue):
 
     hash_value = hash("AbstractQreg")
 
+    def __init__(self, length):
+        self.length = length
+
     def __eq__(self, other):
         return isinstance(other, AbstractQreg)
 
@@ -286,6 +291,10 @@ class Folding(Enum):
 apply_registered_pass_p = core.Primitive("apply_registered_pass")
 transform_named_sequence_p = core.Primitive("transform_named_sequence")
 transform_named_sequence_p.multiple_results = True
+set_state_p = jax.core.Primitive("state_prep")
+set_state_p.multiple_results = True
+set_basis_state_p = jax.core.Primitive("set_basis_state")
+set_basis_state_p.multiple_results = True
 
 
 def _assert_jaxpr_without_constants(jaxpr: ClosedJaxpr):
@@ -968,17 +977,17 @@ def _qdevice_lowering(jax_ctx: mlir.LoweringRuleContext, rtd_lib, rtd_name, rtd_
 # qalloc
 #
 @qalloc_p.def_impl
-def _qalloc_def_impl(ctx, size_value):  # pragma: no cover
+def _qalloc_def_impl(ctx, size_value, static_size=None):  # pragma: no cover
     raise NotImplementedError()
 
 
 @qalloc_p.def_abstract_eval
-def _qalloc_abstract_eval(size):
+def _qalloc_abstract_eval(size, static_size=None):
     """This function is called with abstract arguments for tracing."""
-    return AbstractQreg()
+    return AbstractQreg(static_size)
 
 
-def _qalloc_lowering(jax_ctx: mlir.LoweringRuleContext, size_value: ir.Value):
+def _qalloc_lowering(jax_ctx: mlir.LoweringRuleContext, size_value: ir.Value, static_size=None):
     ctx = jax_ctx.module_context.context
     ctx.allow_unregistered_dialects = True
 
@@ -1068,7 +1077,7 @@ def _qinsert_abstract_eval(qreg_old, qubit_idx, qubit):
     """This function is called with abstract arguments for tracing."""
     assert isinstance(qreg_old, AbstractQreg)
     assert isinstance(qubit, AbstractQbit)
-    return AbstractQreg()
+    return AbstractQreg(qreg_old.length)
 
 
 def _qinsert_lowering(
@@ -2053,6 +2062,58 @@ def _assert_lowering(jax_ctx: mlir.LoweringRuleContext, assertion, error):
     return ()
 
 
+#
+# state_prep
+#
+@set_state_p.def_impl
+def set_state_impl(ctx, *qubits_or_params):  # pragma: no cover
+    """Concrete evaluation"""
+    raise NotImplementedError()
+
+
+@set_state_p.def_abstract_eval
+def set_state_abstract(*qubits_or_params):
+    """Abstract evaluation"""
+    length = len(qubits_or_params)
+    qubits_length = length - 1
+    return (AbstractQbit(),) * qubits_length
+
+
+def _set_state_lowering(jax_ctx: mlir.LoweringRuleContext, *qubits_or_params):
+    """Lowering of set state"""
+    qubits_or_params = list(qubits_or_params)
+    param = qubits_or_params.pop()
+    qubits = qubits_or_params
+    out_qubits = [qubit.type for qubit in qubits]
+    return SetStateOp(out_qubits, param, qubits).results
+
+
+#
+# set_basis_state
+#
+@set_basis_state_p.def_impl
+def set_basis_state_impl(ctx, *qubits_or_params):  # pragma: no cover
+    """Concrete evaluation"""
+    raise NotImplementedError()
+
+
+@set_basis_state_p.def_abstract_eval
+def set_basis_state_abstract(*qubits_or_params):
+    """Abstract evaluation"""
+    length = len(qubits_or_params)
+    qubits_length = length - 1
+    return (AbstractQbit(),) * qubits_length
+
+
+def _set_basis_state_lowering(jax_ctx: mlir.LoweringRuleContext, *qubits_or_params):
+    """Lowering of set basis state"""
+    qubits_or_params = list(qubits_or_params)
+    param = qubits_or_params.pop()
+    qubits = qubits_or_params
+    out_qubits = [qubit.type for qubit in qubits]
+    return SetBasisStateOp(out_qubits, param, qubits).results
+
+
 #
 # adjoint
 #
@@ -2151,6 +2212,8 @@ def _adjoint_lowering(
 mlir.register_lowering(value_and_grad_p, _value_and_grad_lowering)
 mlir.register_lowering(apply_registered_pass_p, _apply_registered_pass_lowering)
 mlir.register_lowering(transform_named_sequence_p, _transform_named_sequence_lowering)
+mlir.register_lowering(set_state_p, _set_state_lowering)
+mlir.register_lowering(set_basis_state_p, _set_basis_state_lowering)
 
 
 def _scalar_abstractify(t):