Lightning Kokkos device

The lightning.kokkos device can run using a variety of HPC-focused backends, including GPUs, enabling accelerated simulation of quantum state-vector evolution.

A lightning.kokkos device can be loaded using:

import pennylane as qml
dev = qml.device("lightning.kokkos", wires=2)

The lightning.kokkos device also directly supports quantum circuit gradients using the adjoint differentiation method. By default, this method is enabled. It can also be explicitly specified using the diff_method argument when creating a device:

qml.qnode(dev, diff_method="adjoint")
def circuit(params):
    ...

Check out the Lightning-Kokkos installation guide for more information.

Supported operations and observables

Supported operations:

BasisState	Prepares a single computational basis state.
BlockEncode	Construct a unitary \(U(A)\) such that an arbitrary matrix \(A\) is encoded in the top-left block.
CNOT	The controlled-NOT operator
ControlledPhaseShift	A qubit controlled phase shift.
ControlledQubitUnitary	Apply an arbitrary fixed unitary matrix U to wires.
CRot	The controlled-Rot operator
CRX	The controlled-RX operator
CRY	The controlled-RY operator
CRZ	The controlled-RZ operator
CSWAP	The controlled-swap operator
CY	The controlled-Y operator
CZ	The controlled-Z operator
DiagonalQubitUnitary	Apply an arbitrary diagonal unitary matrix with a dimension that is a power of two.
DoubleExcitation	Double excitation rotation.
DoubleExcitationMinus	Double excitation rotation with negative phase-shift outside the rotation subspace.
DoubleExcitationPlus	Double excitation rotation with positive phase-shift outside the rotation subspace.
ECR	An echoed RZX(\(\pi/2\)) gate.
GlobalPhase	A global phase operation that multiplies all components of the state by \(e^{-i \phi}\).
Hadamard	The Hadamard operator
Identity	The Identity operator
IsingXX	Ising XX coupling gate
IsingXY	Ising (XX + YY) coupling gate
IsingYY	Ising YY coupling gate
IsingZZ	Ising ZZ coupling gate
ISWAP	The i-swap operator
MultiControlledX	Apply a PauliX gate controlled on an arbitrary computational basis state.
MultiRZ	Arbitrary multi Z rotation.
OrbitalRotation	Spin-adapted spatial orbital rotation.
PauliX	The Pauli X operator
PauliY	The Pauli Y operator
PauliZ	The Pauli Z operator
PhaseShift	Arbitrary single qubit local phase shift
PSWAP	Phase SWAP gate
QubitCarry	Apply the QubitCarry operation to four input wires.
QubitSum	Apply a QubitSum operation on three input wires.
QubitUnitary	Apply an arbitrary unitary matrix with a dimension that is a power of two.
Rot	Arbitrary single qubit rotation
RX	The single qubit X rotation
RY	The single qubit Y rotation
RZ	The single qubit Z rotation
S	The single-qubit phase gate
SingleExcitation	Single excitation rotation.
SingleExcitationMinus	Single excitation rotation with negative phase-shift outside the rotation subspace.
SingleExcitationPlus	Single excitation rotation with positive phase-shift outside the rotation subspace.
SISWAP	The square root of i-swap operator.
SQISW	alias of SISWAP
SWAP	The swap operator
SX	The single-qubit Square-Root X operator.
T	The single-qubit T gate
Toffoli	Toffoli (controlled-controlled-X) gate.

Supported observables:

Identity	The Identity operator
Hadamard	The Hadamard operator
PauliX	The Pauli X operator
PauliY	The Pauli Y operator
PauliZ	The Pauli Z operator
Projector	Observable corresponding to the state projector \(P=\ket{\phi}\bra{\phi}\).
Hermitian	An arbitrary Hermitian observable.
Hamiltonian	alias of LinearCombination
SparseHamiltonian	A Hamiltonian represented directly as a sparse matrix in Compressed Sparse Row (CSR) format.
Exp	A symbolic operator representing the exponential of a operator.
Prod	Symbolic operator representing the product of operators.
SProd	Arithmetic operator representing the scalar product of an operator with the given scalar.
Sum	Symbolic operator representing the sum of operators.

Distributed simulation with MPI:

The lightning.kokkos device supports distributed simulation using the Message Passing Interface (MPI). This enables the simulation of larger quantum circuits by distributing the workload across multiple CPU or GPU compute nodes.

To utilize distributed simulation, lightning.kokkos must be compiled with MPI support. Check out the Lightning Kokkos installation guide for more information.

With lightning.kokkos installed with MPI support, distributed simulation can be enabled in Pennylane by setting the mpi keyword argument to True when creating the device. For example:

from mpi4py import MPI
import pennylane as qml

dev = qml.device('lightning.kokkos', wires=30, mpi=True)
@qml.qnode(dev)
def circuit_mpi():
    qml.PauliX(wires=[0])
    return qml.state()
local_state_vector = circuit_mpi()

Note

The total number of MPI processes must be powers of 2. If using Kokkos with GPUs, we recommend using one GPU per MPI process.

Currently, the lightning.kokkos device with MPI supports all the gate operations and observables that a single process lightning.kokkos device supports, excluding SparseHamiltonian.

By default, each MPI process will return the overall simulation results, except for the qml.state() and qml.probs() methods for which each MPI process only returns the local simulation results for the qml.state() and qml.probs() methods to avoid buffer overflow. It is the user’s responsibility to ensure correct data collection for those two methods (e.g. using MPI Gather). Here are examples of collecting the local simulation results for qml.state() and qml.probs() methods:

The workflow for collecting local state vector (using the qml.state() method) to rank 0 is as follows:

from mpi4py import MPI
import pennylane as qml
comm = MPI.COMM_WORLD
rank = comm.Get_rank()

dev = qml.device('lightning.kokkos', wires=30, mpi=True)
@qml.qnode(dev)
def circuit_mpi():
    qml.PauliX(wires=[0])
    return qml.state()

local_state_vector = circuit_mpi()

#rank 0 will collect the local state vector
state_vector = comm.gather(local_state_vector, root=0)
if rank == 0:
    print(state_vector)

The workflow for collecting local probability (using the qml.probs() method) to rank 0 is as follows:

from mpi4py import MPI
import pennylane as qml
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
dev = qml.device('lightning.kokkos', wires=30, mpi=True)
prob_wires = [0, 1]

@qml.qnode(dev)
def mpi_circuit():
    qml.Hadamard(wires=1)
    return qml.probs(wires=prob_wires)

local_probs = mpi_circuit()

#For data collection across MPI processes.
recv_counts = comm.gather(len(local_probs),root=0)
if rank == 0:
    probs = np.zeros(2**len(prob_wires))
else:
    probs = None

comm.Gatherv(local_probs,[probs,recv_counts],root=0)
if rank == 0:
    print(probs)

Then the python script can be executed with the following command (for example on 4 MPI processes):

$ mpirun -np 4 python pennylane_quantum_script.py

lightning_kokkos/device

 

Download Python script

 

Download Notebook

 

View on GitHub