SABRE qubit mapping and routing implementation in Catalyst

frontend/catalyst/api_extensions/callbacks.py

@@ -138,7 +138,7 @@ def f(x):
     return accelerate_impl(func, dev=dev)
 
 
-def pure_callback(callback_fn, result_type=None):
+def pure_callback(callback_fn=None, *, result_type=None):
     """Execute and return the results of a functionally pure Python
     function from within a qjit-compiled function.
 
@@ -261,6 +261,11 @@ def vjp(x, dy) -> (jax.ShapeDtypeStruct((2,), jnp.float64),):
         >>> f(jnp.array([0.1, 0.2]))
         Array([-0.01071923,  0.82698717], dtype=float64)
     """
+    # Just for convenience
+    if callback_fn is None:
+        kwargs = copy.copy(locals())
+        kwargs.pop("callback_fn")
+        return functools.partial(pure_callback, **kwargs)
 
     # Verify inputs
     if result_type is None:

frontend/catalyst/compiled_functions.py

@@ -92,36 +92,45 @@ def load_symbols(self):
             CFuncPtr: handle to the teardown function, which tears down the device
             CFuncPtr: handle to the memory transfer function for program results
         """
+        try:
 
-        setup = self.shared_object.setup
-        setup.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_char_p)]
-        setup.restypes = ctypes.c_int
+            setup = self.shared_object.setup
+            setup.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_char_p)]
+            setup.restypes = ctypes.c_int
 
-        teardown = self.shared_object.teardown
-        teardown.argtypes = None
-        teardown.restypes = None
+            teardown = self.shared_object.teardown
+            teardown.argtypes = None
+            teardown.restypes = None
 
-        # We are calling the c-interface
-        function = self.shared_object["_catalyst_pyface_" + self.func_name]
-        # Guaranteed from _mlir_ciface specification
-        function.restypes = None
-        # Not needed, computed from the arguments.
-        # function.argyptes
+        except:
+            setup = None
+            teardown = None
+        finally:
 
-        mem_transfer = self.shared_object["_mlir_memory_transfer"]
+            # We are calling the c-interface
+            function = self.shared_object["_catalyst_pyface_" + self.func_name]
+            # Guaranteed from _mlir_ciface specification
+            function.restypes = None
+            # Not needed, computed from the arguments.
+            # function.argyptes
 
-        return function, setup, teardown, mem_transfer
+            mem_transfer = self.shared_object["_mlir_memory_transfer"]
+
+            return function, setup, teardown, mem_transfer
 
     def __enter__(self):
-        params_to_setup = [b"jitted-function"]
-        argc = len(params_to_setup)
-        array_of_char_ptrs = (ctypes.c_char_p * len(params_to_setup))()
-        array_of_char_ptrs[:] = params_to_setup
-        self.setup(ctypes.c_int(argc), array_of_char_ptrs)
+        if self.setup:
+            params_to_setup = [b"jitted-function"]
+            argc = len(params_to_setup)
+            array_of_char_ptrs = (ctypes.c_char_p * len(params_to_setup))()
+            array_of_char_ptrs[:] = params_to_setup
+            self.setup(ctypes.c_int(argc), array_of_char_ptrs)
+            return self
         return self
 
     def __exit__(self, _type, _value, _traceback):
-        self.teardown()
+        if self.teardown:
+            self.teardown()
 
 
 class CompiledFunction:
@@ -329,10 +338,10 @@ def get_cmain(self, *args):
         return get_template(self.func_name, self.restype, *buffer)
 
     def __call__(self, *args, **kwargs):
-        static_argnums = self.compile_options.static_argnums
+        static_argnums = self.compile_options.static_argnums if self.compile_options else ()
         dynamic_args = filter_static_args(args, static_argnums)
 
-        if self.compile_options.abstracted_axes is not None:
+        if self.compile_options and self.compile_options.abstracted_axes is not None:
             abstracted_axes = self.compile_options.abstracted_axes
             dynamic_args = get_implicit_and_explicit_flat_args(
                 abstracted_axes, *dynamic_args, **kwargs

mlir/include/Quantum/Transforms/Passes.h

@@ -34,5 +34,6 @@ std::unique_ptr<mlir::Pass> createDisentangleCNOTPass();
 std::unique_ptr<mlir::Pass> createDisentangleSWAPPass();
 std::unique_ptr<mlir::Pass> createIonsDecompositionPass();
 std::unique_ptr<mlir::Pass> createLoopBoundaryOptimizationPass();
+std::unique_ptr<mlir::Pass> createRoutingPass();
 
 } // namespace catalyst

mlir/include/Quantum/Transforms/Passes.td

@@ -125,6 +125,22 @@ def LoopBoundaryOptimizationPass : Pass<"loop-boundary"> {
 
     let constructor = "catalyst::createLoopBoundaryOptimizationPass()";
 }
+
+def RoutingPass : Pass<"route-circuit"> {
+    let summary = "Perform mapping and routing of quantum circuit to a hardware.";
+
+    let constructor = "catalyst::createRoutingPass()";
+
+  let options = [
+    Option<
+      /*C++ variable name=*/"hardwareGraph",
+      /*CLI argument=*/"hardware-graph",
+      /*type=*/"std::string",
+      /*default=*/"",
+      /*description=*/"Hardware graph represented by a list of edges connecting physical qubits."
+    >,
+  ];
+}
 // ----- Quantum circuit transformation passes end ----- //
 
 #endif // QUANTUM_PASSES

mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp

@@ -71,4 +71,5 @@ void catalyst::registerAllCatalystPasses()
     mlir::registerPass(catalyst::createScatterLoweringPass);
     mlir::registerPass(catalyst::createSplitMultipleTapesPass);
     mlir::registerPass(catalyst::createTestPass);
+    mlir::registerPass(catalyst::createRoutingPass);
 }

mlir/lib/Quantum/Transforms/CMakeLists.txt

@@ -20,6 +20,7 @@ file(GLOB SRC
     IonsDecompositionPatterns.cpp
     loop_boundary_optimization.cpp
     LoopBoundaryOptimizationPatterns.cpp
+    route_circuit.cpp
 )
 
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)