Solves Issue #1316: Implement an option to print out instructions in the null device

doc/releases/changelog-dev.md

@@ -2,6 +2,9 @@
 
 <h3>New features since last release</h3>
 
+* The PennyLane plugin now has the option to print out instructions in the null device
+  [(#1316)](https://github.com/PennyLaneAI/catalyst/pull/1346)
+
 <h3>Improvements 🛠</h3>
 
 * Replace pybind11 with nanobind for C++/Python bindings in the frontend and in the runtime.
@@ -113,4 +116,5 @@ Mehrdad Malekmohammadi,
 William Maxwell
 Romain Moyard,
 Raul Torres,
-Paul Haochen Wang.
+Paul Haochen Wang,
+Stefanie Muroya Lei.

frontend/catalyst/device/qjit_device.py

@@ -144,7 +144,7 @@
 # pylint: disable=too-many-branches
 @debug_logger
 def extract_backend_info(
     device: qml.devices.QubitDevice, capabilities: DeviceCapabilities
 ) -> BackendInfo:
     """Extract the backend info from a quantum device. The device is expected to carry a reference
     to a valid TOML config file."""
@@ -303,7 +303,7 @@
         return extract_backend_info(device, capabilities)
 
     @debug_logger_init
-    def __init__(self, original_device):
+    def __init__(self, original_device, print_instructions=False):
         self.original_device = original_device
 
         for key, value in original_device.__dict__.items():
@@ -331,6 +331,11 @@
         self.backend_name = backend.c_interface_name
         self.backend_lib = backend.lpath
         self.backend_kwargs = backend.kwargs
+
+        if original_device.name == "null.qubit":
+            # include 'print_instructions' as a keyword argument for the device constructor.
+            self.backend_kwargs["print_instructions"] = print_instructions
+
         self.capabilities = get_qjit_device_capabilities(device_capabilities)
 
     @debug_logger

frontend/catalyst/jit.py

@@ -87,6 +87,7 @@ def qjit(
     seed=None,
     experimental_capture=False,
     circuit_transform_pipeline=None,
+    print_instructions=False
 ):  # pylint: disable=too-many-arguments,unused-argument
     """A just-in-time decorator for PennyLane and JAX programs using Catalyst.
 
@@ -153,6 +154,8 @@ def qjit(
             dictionaries of valid keyword arguments and values for the specific pass.
             The order of keys in this dictionary will determine the pass pipeline.
             If not specified, the default pass pipeline will be applied.
+        print_instructions (Optional[bool]):
+            If set to True, instructions are printed when executing in the null device.
 
     Returns:
         QJIT object.
@@ -430,11 +433,12 @@ def sum_abstracted(arr):
     """
     kwargs = copy.copy(locals())
     kwargs.pop("fn")
+    kwargs.pop("print_instructions")
 
     if fn is None:
         return functools.partial(qjit, **kwargs)
 
-    return QJIT(fn, CompileOptions(**kwargs))
+    return QJIT(fn, CompileOptions(**kwargs), print_instructions=print_instructions)
 
 
 ## IMPL ##
@@ -452,6 +456,8 @@ class QJIT(CatalystCallable):
     Args:
         fn (Callable): the quantum or classical function to compile
         compile_options (CompileOptions): compilation options to use
+        print_instructions (Optional[bool]): If True, prints instructions
+            when executing in a null device.
 
     :ivar original_function: This attribute stores `fn`, the quantum or classical function
                              object to compile, as is, without any modifications
@@ -462,7 +468,10 @@ class QJIT(CatalystCallable):
     """
 
     @debug_logger_init
-    def __init__(self, fn, compile_options):
+    def __init__(self, fn, compile_options, print_instructions=False):
+        # flag for printing instructions in the null device
+        self.print_instructions = print_instructions
+
         functools.update_wrapper(self, fn)
         self.original_function = fn
         self.compile_options = compile_options
@@ -635,6 +644,10 @@ def closure(qnode, *args, **kwargs):
             params = {}
             params["static_argnums"] = kwargs.pop("static_argnums", static_argnums)
             params["_out_tree_expected"] = []
+
+            if qnode.device.name == "null.qubit":
+                kwargs["print_instructions"] = self.print_instructions
+
             return QFunc.__call__(
                 qnode,
                 pass_pipeline=self.compile_options.circuit_transform_pipeline,

frontend/catalyst/qfunc.py

@@ -122,7 +122,10 @@ def __call__(self, *args, **kwargs):
                 mcm_config.postselect_mode = mcm_config.postselect_mode or "hw-like"
                 return Function(dynamic_one_shot(self, mcm_config=mcm_config))(*args, **kwargs)
 
-        qjit_device = QJITDevice(self.device)
+        # retrieve the flag to print instructions, used for executing
+        # pre-compiled programs in a null device.
+        print_instructions = kwargs.pop("print_instructions", False)
+        qjit_device = QJITDevice(self.device, print_instructions)
 
         static_argnums = kwargs.pop("static_argnums", ())
         out_tree_expected = kwargs.pop("_out_tree_expected", [])