Merge pull request #341 from eclipse-qrisp/catalyst_device

Support for different Catalyst backends

Author: positr0nium

src/qrisp/jasp/evaluation_tools/catalyst_interface.py

@@ -99,14 +99,16 @@ def jaspr_to_catalyst_jaxpr(jaspr):
     return make_jaxpr(eval_jaxpr(jaspr, eqn_evaluator=catalyst_eqn_evaluator))(*args)
 
 
-def jaspr_to_catalyst_function(jaspr):
+def jaspr_to_catalyst_function(jaspr, device=None):
 
     # This function takes a jaspr and returns a function that performs a sequence
     # of .bind calls of Catalyst primitives, such that the function (when compiled)
     # by Catalyst reproduces the semantics of jaspr
 
     # Initiate Catalyst backend info
-    device = qml.device("lightning.qubit", wires=0)
+    if device==None:
+        device = qml.device("lightning.qubit", wires=0)
+
     backend_info = catalyst.device.extract_backend_info(device)
 
     def catalyst_function(*args):
@@ -138,10 +140,10 @@ def catalyst_function(*args):
 
 
 @lru_cache(int(1e5))
-def jaspr_to_catalyst_qjit(jaspr, function_name="jaspr_function"):
+def jaspr_to_catalyst_qjit(jaspr, function_name="jaspr_function", device=None):
     # This function takes a jaspr and turns it into a Catalyst QJIT object.
     # Perform the code specified by the Catalyst developers
-    catalyst_function = jaspr_to_catalyst_function(jaspr)
+    catalyst_function = jaspr_to_catalyst_function(jaspr, device=device)
     catalyst_function.__name__ = function_name
     jit_object = catalyst.QJIT(catalyst_function, catalyst.CompileOptions())
     jit_object.jaxpr = make_jaxpr(catalyst_function)(

src/qrisp/jasp/evaluation_tools/catalyst_qjit.py

@@ -20,7 +20,7 @@
 from qrisp.jasp.jasp_expression import make_jaspr
 
 
-def qjit(function):
+def qjit(function=None, device=None):
     """
     Decorator to leverage the jasp + Catalyst infrastructure to compile the given
     function to QIR and run it on the Catalyst QIR runtime.
@@ -29,12 +29,40 @@ def qjit(function):
     ----------
     function : callable
         A function performing Qrisp code.
+    device : object
+        The `PennyLane device <https://docs.pennylane.ai/projects/catalyst/en/stable/dev/devices.html>`_ to execute the function. 
+        The default device is `"lightning.qubit" <https://docs.pennylane.ai/projects/lightning/en/stable/lightning_qubit/device.html>`_, 
+        a fast state-vector qubit simulator.
 
     Returns
     -------
     callable
         A function executing the compiled code.
 
+    Notes
+    -----
+
+    Lightning-GPU is compatible with systems featuring NVIDIA Volta (SM 7.0) GPUs or newer. 
+    It is specifically optimized for Linux environments on X86-64 or ARM64 architectures running CUDA-12.
+
+    To install Lightning-GPU with NVIDIA CUDA support, the following packages need to be installed
+     
+    ::
+
+        pip install custatevec_cu12
+        pip install pennylane-lightning-gpu
+
+        
+    Pre-built wheels for Lightning-AMDGPU are available for AMD MI300 series GPUs and systems running ROCm 7.0 or newer.
+
+    ::
+
+        pip install pennylane-lightning-amdgpu
+
+    If the setup uses an older version of ROCm or a different AMD GPU series, Lightning-AMDGPU must be built manually from source.
+
+    Installation instructions for different platforms are available at `pennylane.ai/install <https://pennylane.ai/install#high-performance-computing-and-gpus>`_.
+
     Examples
     --------
 
@@ -65,8 +93,31 @@ def test_fun(i):
     >>> test_fun(5)
     [array(7.25, dtype=float64)]
 
+
+    For executing on "lightning.gpu" we specify the device:
+
+    ::
+
+        import pennylane as qml
+        from qrisp import *
+        from qrisp.jasp import qjit
+
+        dev = qml.device("lightning.gpu", wires=0)
+
+        @qjit(device=dev)
+        def test_fun(i):
+            qv = QuantumFloat(i, -2)
+            with invert():
+                cx(qv[0], qv[qv.size-1])
+                h(qv[0])
+            meas_res = measure(qv)
+            return meas_res + 3
+
     """
 
+    if function is None:
+        return lambda x: qjit(x, device=device)
+
     def jitted_function(*args):
 
         if not hasattr(function, "jaspr_dict"):
@@ -79,7 +130,7 @@ def jitted_function(*args):
             function.jaspr_dict[signature] = make_jaspr(function)(*args)
 
         return function.jaspr_dict[signature].qjit(
-            *args, function_name=function.__name__
+            *args, function_name=function.__name__, device=device
         )
 
     return jitted_function

src/qrisp/jasp/jasp_expression/centerclass.py

@@ -670,7 +670,7 @@ def embedd(self, *args, name=None, inline=False):
         qs.abs_qc = new_abs_qc
         return res
 
-    def qjit(self, *args, function_name="jaspr_function"):
+    def qjit(self, *args, function_name="jaspr_function", device=None):
         """
         Leverages the Catalyst pipeline to compile a QIR representation of
         this function and executes that function using the Catalyst QIR runtime.
@@ -679,6 +679,10 @@ def qjit(self, *args, function_name="jaspr_function"):
         ----------
         *args : iterable
             The arguments to call the function with.
+        device : object
+            The `PennyLane device <https://docs.pennylane.ai/projects/catalyst/en/stable/dev/devices.html>`_ to execute the function. 
+            The default device is `"lightning.qubit" <https://docs.pennylane.ai/projects/lightning/en/stable/lightning_qubit/device.html>`_, 
+            a fast state-vector qubit simulator.
 
         Returns
         -------
@@ -691,7 +695,7 @@ def qjit(self, *args, function_name="jaspr_function"):
             jaspr_to_catalyst_qjit,
         )
 
-        qjit_obj = jaspr_to_catalyst_qjit(flattened_jaspr, function_name=function_name)
+        qjit_obj = jaspr_to_catalyst_qjit(flattened_jaspr, function_name=function_name, device=device)
         res = qjit_obj.compiled_function(*args)
         if not isinstance(res, (tuple, list)):
             return res