proof of concept 2

Author: erick-xanadu

frontend/catalyst/from_plxpr.py

@@ -183,6 +183,7 @@ def handle_qnode(
     non_const_args = args[n_consts:]
 
     f = partial(QFuncPlxprInterpreter(device, shots).eval, qfunc_jaxpr, consts)
+    f = jax.jit(f)
 
     return quantum_kernel_p.bind(
         wrap_init(f, debug_info=qfunc_jaxpr.debug_info),

frontend/catalyst/pipelines.py

@@ -173,6 +173,8 @@ def get_enforce_runtime_invariants_stage(_options: CompileOptions) -> List[str]:
         # keep inlining modules targeting the Catalyst runtime.
         # But qnodes targeting other backends may choose to lower
         # this into something else.
+        "builtin.module(inline)",
+        "split-multiple-tapes",
         "inline-nested-module",
     ]
     return enforce_runtime_invariants
@@ -217,7 +219,6 @@ def get_bufferization_stage(_options: CompileOptions) -> List[str]:
     """Returns the list of passes that performs bufferization"""
     bufferization = [
         "one-shot-bufferize{dialect-filter=memref}",
-        "inline",
         "gradient-preprocess",
         "gradient-bufferize",
         "scf-bufferize",

mlir/lib/Catalyst/IR/CatalystDialect.cpp

@@ -28,6 +28,17 @@ using namespace catalyst;
 //===----------------------------------------------------------------------===//
 // Catalyst dialect.
 //===----------------------------------------------------------------------===//
+namespace {
+struct CatalystInlinerInterface : public DialectInlinerInterface {
+    using DialectInlinerInterface::DialectInlinerInterface;
+
+    /// Operations in Gradient dialect are always legal to inline.
+    bool isLegalToInline(Operation *op, Region *, bool, IRMapping &valueMapping) const final
+    {
+        return isa<CallbackCallOp>(op);
+    }
+};
+}
 
 void CatalystDialect::initialize()
 {
@@ -40,6 +51,7 @@ void CatalystDialect::initialize()
 #define GET_OP_LIST
 #include "Catalyst/IR/CatalystOps.cpp.inc"
         >();
+    addInterface<CatalystInlinerInterface>();
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Quantum/IR/QuantumDialect.cpp

@@ -14,6 +14,7 @@
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/IR/DialectImplementation.h" // needed for generated type parser
+#include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/TypeSwitch.h"           // needed for generated type parser
 
 #include "Quantum/IR/QuantumDialect.h"
@@ -28,6 +29,20 @@ using namespace catalyst::quantum;
 
 #include "Quantum/IR/QuantumOpsDialect.cpp.inc"
 
+namespace {
+struct QuantumInlinerInterface : public DialectInlinerInterface {
+    using DialectInlinerInterface::DialectInlinerInterface;
+
+    /// Operations in Gradient dialect are always legal to inline.
+    bool isLegalToInline(Operation *, Region *, bool, IRMapping &valueMapping) const final
+    {
+        return true;
+    }
+};
+} // namespace
+
+
+
 void QuantumDialect::initialize()
 {
     addTypes<
@@ -48,6 +63,7 @@ void QuantumDialect::initialize()
     declarePromisedInterfaces<bufferization::BufferizableOpInterface, QubitUnitaryOp, HermitianOp,
                               HamiltonianOp, SampleOp, CountsOp, ProbsOp, StateOp, SetStateOp,
                               SetBasisStateOp>();
+    addInterface<QuantumInlinerInterface>();
 }
 
 //===----------------------------------------------------------------------===//