Multi-IR passmanager for Rust/C/Python

[Julien Gacon (Cryoris)]

This RFC drafts a new passmanager for Qiskit that is written in Rust and will be exposed to C and to Python. It has two main goals:

1. Expose a single compiler infrastructure to Python and C (and by an FFI to other languages, like Rust), and
2. to formalize the multi-IR pipelines in a typed language -- which will facilitate the planned FTQC compiler pipelines.

This is targeting Qiskit v2.6.

[Ian Hincks (ihincks)]

Thanks Julien Gacon (@Cryoris)!

I'd like to sketch an adapter that allows a nested compiler structure. By nested, I mean the situation where an outer IR owns components which themselves are complicated enough that you want to treat them via a separate pass manager. That is, you have a pass, the "adapter", acting as Outer->Outer and that owns a task Component->Component. The job of the adapter is to search for components and apply the inner task. For example, when a quantum circuit owns a conditional instruction, there should be a generic way to run a pass manager/task on each block of every conditional instruction, and recursively. But the example I'm really interested in is a QuantumProgram that contains a bunch of QuantumCircuits: we want to run standard pass managers on those as one step of compiling the QuantumProgram.

/// A way of decomposing an outer IR into owned component units.
/// 
/// Example: `impl UnrollQuantumCircuitForLoop for Nesting`
pub trait Nesting {
    type OuterIR;
    type ComponentIR;
    type Id: Clone + Eq + Hash + Debug + 'static;

    // Whether extracted components might share any scope (affects possible parallelism)
    const ISOLATED: bool;

    fn take_components(
        &self,
        outer: &mut Self::OuterIR,
    ) -> Vec<(Self::Id, Self::ComponentIR)>;

    fn put_component(
        &self,
        outer: &mut Self::OuterIR,
        id: Self::Id,
        component: Self::ComponentIR,
    );
}

/// Runs `inner` over the components of an outer IR.
pub struct NestedAdaptor<N: Nesting> {
    nesting: N,
    /// Task is `N::ComponentIR -> N::ComponentIR` checked at construction
    inner: Task
}

impl<N: Nesting> Pass for NestedAdaptor<N> {
    type InputIR = N::OuterIR;
    type OutputIR = N::OuterIR;

    fn run(&self, mut outer: N::OuterIR, context: &mut PassContext)
        -> Result<Self::OutputIR>
    { /* ... invoke take_components and put_component, possibly with parallelism */ }
}

This looks like it mostly works well. However, it seems like PassContext will step on its own toes a bit.
Taking the example Outer = QuantumProgram and Inner = QuantumCircuit, and assuming that we want to pass the PassContext
from the OuterIR down to each QuantumCircuit rather than start a fresh context each time, then all of the strings will
collide. Should we introduce namespacing or scoping into the context?

I have similar questions about how this will eventually interact with debug information. The user will care just as much about
the debug information inside of components as they do about debug information in the outer scope.

[Ian Hincks (ihincks)]

Pointing out something obvious in case it sparks discussion: making preserves() a static function, instead of an output of run(), forces it to be conservative. If a pass finds nothing to do on a particular input, it must declare that it did something here, despite the fact that it didn't.

[Julien Gacon (Cryoris)]

Oh that's a good point that we actually discussed previously too. It would be better to have it be an output of run, I agree. In that case we can make this decision dynamically!

[Ian Hincks (ihincks)]

A PassManager and a Task::Group are essentially the same thing. I wonder if we can unify them a bit by having

impl From<PassManager> for Task {
    fn from(pm: PassManager) -> Task { Task::Group(pm.tasks) }
}

This ties into my main comment, where I would want the ability to set inner to a PassManager.

[Julien Gacon (Cryoris)]

Conceptually, a PassManager and Task are different in this RFC and it's a distinction we'd like to make clear compared to the current situation: a PassManager is the sole executor, a Task only provides the description of the flow.

Further, there is only a single pass manager as executor, so "nesting PM1 inside PM2" would rather amount to "appending the tasks of PM1 to PM2" using a Task::Group. What you wrote above is one option, alternatively we could be more explicit about the distinction

pm2.try_push_task(Task::Group(pm1.tasks));

[Ian Hincks (ihincks)]

"nesting PM1 inside PM2" would rather amount to "appending the tasks of PM1 to PM2" using a Task::Group. What you wrote above is one option, alternatively we could be more explicit about the distinction

I think there must be a fundamental misunderstanding somewhere between us. I don't think that concatenating list-like passes is equivalent to nesting passes.

[Julien Gacon (Cryoris)]

Right, I probably didn't write the above well enough: you can nest passes arbitrarily using the Task::Group. What I meant above is that nesting a PM inside a pass would mean to define a Task::Group(pm.tasks). I was assuming you then want to append this nested object onto another pass manager, hence the "append" above. But we can also have a chat offline!

[Ian Hincks (ihincks)]

Can you remind me where a Target lives? In here?

[Julien Gacon (Cryoris)]

Right now the Target is not a runtime argument, but the idea is to keep the existing model and have it be an input argument to the pass struct. So something like

struct UnitarySynthesis {
  target: Option<Target>,
  ...
}

impl UnitarySynthesis {
  fn new(target: Option<Target>, ...) -> Self { ... }
}

impl Pass for UnitarySynthesis {
  ...
}

[Ian Hincks (ihincks)]

the idea is to keep the existing model

Just for my own sake: is this for the continuity, or because it's a better idea? I suppose allowing every pass to have it's own definition of a target is more flexible.

[Julien Gacon (Cryoris)]

It sounds like the right separation to me: the pass context contains program-specific information gathered at runtime, whereas the Target (and other arguments like approximation_degree etc.) specify the pass' action independent of the program. I'll mention this in the doc!

[Ian Hincks (ihincks)]

Prompted by Jake Lishman (@jakelishman), this is a rework of my original comment where I extend the Task enum with a
new variant instead of trying to make a new Pass. It is cleaner to own the entire pipeline, including
actions on regions, as a single object instead of some adhoc Pass implementations that themselves own some tasks. Moreover, it is
important for this RFC to prove that we are not shooting ourselves in the foot by disallowing
extensibility: in general terms, we need to make sure that variants of the signature of what this RFC currently
calls Pass::run (as opposed to what some of the proposals below may or may not change about it) can
be added without breaking anything that uses the public interface.

The main deficiency I see with the proposed Pass::run (relative specifically to my needs for "nesting") is that
it gets the IR by value. This is awkward when the IR represents a region, because the outer IR that it is a part of
naturally wants to own it to; youneed to implement a work around, like replacing it with a place-holder while
running the pass, or something like that.

One of Jake Lishman (@jakelishman)'s main points is that what I'm considering here is just one of many examples (as yet unknown)
where we will find the need to have new signatures of Pass::run. So, this is not a specific proposal, but sort
of just some mild proof that we can do it if we need to.

Region walker

I'm going to write two alternative implementation sketches, but both of them use this RegionWalker trait.
A trait is used just for this instead of trying to tie it to a new IR trait or something like that because
I'd strongly prefer region discovery to live outside the IR so that a single IR can have multiple notions of
what a region is, depending on the needs of the pipeline.

/// Discovers regions and calls a function on each. Notably, provides mutable references.
pub trait RegionWalker {
    type OuterIR: 'static;
    type RegionIR: 'static;

    fn walk(
        &self,
        outer: &mut Self::OuterIR,
        body: &mut dyn FnMut(&mut Self::RegionIR) -> anyhow::Result<()>,
    ) -> anyhow::Result<()>;
}

trait AnyRegionWalker { /* type erased version of RegionWalker */ }
impl<W: RegionWalker> AnyRegionWalker for W { /* ... */ }

Alternative 1: multiple traits, each with one ::run().

Remove Pass as a trait, and instead have multiple traits. For every new run() signature
required in the future, a new trait is added.

pub trait Transform {
    type IR: 'static;

    fn run(&self, ir: &mut Self::IR, context: &mut PassContext) -> anyhow::Result<()>;
    fn preserves(&self) -> PreservedAnalysis;
}

pub trait Lower {
    type InputIR: 'static;
    type OutputIR: 'static;

    fn run(&self, ir: Self::InputIR, context: &mut PassContext) -> anyhow::Result<Self::OutputIR>;
    fn preserves(&self) -> PreservedAnalysis;
}

trait AnyTransform { /* type erasure version of Transform */ } 
impl<T: Transform> AnyTransform for T { /* ... */ }

trait AnyLower { /* type erasure version of Lower */ } 
impl<T: Lower> AnyLower for T { /* ... */ }

All such traits are grouped into a non-exhaustive enum called Pass, and run methods are exposed.
Note that Pass::run() could be implemented for both Transform and Lower; there need not
be a 1-1 mapping between these methods and the traits, in general.

#[non_exhaustive]
pub enum Pass {
    Transform(Box<dyn AnyTransform>),
    Lower(Box<dyn AnyLower>),
}

impl Pass {
    fn io_types(&self) -> (TypeId, TypeId);

    /// `Some(ir)` for `Transform`, `None` for `Lower`.
    fn in_place_ir(&self) -> Option<TypeId>;

    fn run(&self, ir: Box<dyn Any>, context: &mut PassContext)
        -> Result<Box<dyn Any>, PassManagerError>;

    /// panics if not possible; pipeline construction checks should prevent an attempt at that
    fn run_in_place(&self, ir: &mut dyn Any, context: &mut PassContext)
        -> Result<(), PassManagerError>;
}

A task then gains a ForEachRegion variant, and the execution functions are taught how to use it.

#[non_exhaustive]
pub enum Task {
    Pass(Pass),
    Group(Vec<Task>),
    Stages(Vec<(String, Task)>),
    Switch { switch: fn(&dyn Any, &PassContext) -> usize, cases: Vec<Task> },
    Loop { condition: fn(&dyn Any, &PassContext) -> bool, body: Box<Task> },
    ForEachRegion { walker: Box<dyn AnyRegionWalker>, body: Box<Task> },
}

impl Task {
    /// `Some(ir)` if every leaf can run in place on the same ir
    fn in_place_ir(&self) -> Option<TypeId>;

    fn io_types(&self) -> Result<(TypeId, TypeId), PassManagerError> {
        match self {
            Task::ForEachRegion { walker, body } => {
                if body.in_place_ir() != Some(walker.region_ir()) {
                    return Err(PassManagerError::IncompatibleTypes);
                }
                let outer = walker.outer_ir();
                Ok((outer, outer))
            },
            // ...
        }
    }
}

fn execute(
    task: &Task,
    mut ir: Box<dyn Any>,
    context: &mut PassContext,
    callbacks: Option<&CallbackRegistry>,
) -> Result<Box<dyn Any>, PassManagerError> {
    match task {
        Task::Pass(pass) => pass.run(ir, context),
        Task::ForEachRegion { .. } => {
            execute_in_place(task, ir.as_mut(), context, callbacks)?;
            Ok(ir)
        },
        // ...
    }
}

fn execute_in_place(
    task: &Task,
    ir: &mut dyn Any,
    context: &mut PassContext,
    callbacks: Option<&CallbackRegistry>,
) -> Result<(), PassManagerError> {
    match task {
        Task::Pass(pass) => pass.run_in_place(ir, context),
        Task::ForEachRegion { walker, body } => walker
            .walk(ir, &mut |region| {
                execute_in_place(body, region, context, callbacks).map_err(anyhow::Error::from)
            })
            .map_err(PassManagerError::PassError),
        // ...
    }
}

Alternative 2: add optional methods to Pass

In this alternative, we simply start adding new methods to the existing trait, but we also introduce
a mechanism to query which ones are actually implemented. They must all have default implementations
to not trip up external code.

pub struct Signatures {
    pub by_value: bool,
    pub in_place: bool,
}

pub trait Pass {
    type InputIR: 'static;
    type OutputIR: 'static;

    fn signatures(&self) -> Signatures;
    fn preserves(&self) -> PreservedAnalysis;

    fn run(&self, ir: Self::InputIR, context: &mut PassContext) -> anyhow::Result<Self::OutputIR> {
        Err(PassError::NotImplemented.into())
    }

    // In-place execution is only sound when `InputIR` and `OutputIR` are the same type, which the trait
    // cannot state as a bound. The pipeline builder checks it instead.
    fn run_in_place(&self, ir: &mut Self::InputIR, context: &mut PassContext) -> anyhow::Result<()> {
        Err(PassError::NotImplemented.into())
    }
}

trait AnyPass { /* type erased version */ }
impl<P: Pass> AnyPass for P { /* ... */}

As in the first alternative, we grow the task enum and teach execution about it:

#[non_exhaustive]
pub enum Task {
    Pass(Box<dyn AnyPass>),
    Group(Vec<Task>),
    Stages(Vec<(String, Task)>),
    Switch { switch: fn(&dyn Any, &PassContext) -> usize, cases: Vec<Task> },
    Loop { condition: fn(&dyn Any, &PassContext) -> bool, body: Box<Task> },
    ForEachRegion { walker: Box<dyn AnyRegionWalker>, body: Box<Task> },
}

impl Task {
    /// `Some(ir)` if every leaf provides `run_in_place` and maps `ir` to itself.
    fn in_place_ir(&self) -> Option<TypeId> { /* ... */ }
    /// Validation lives in `io_types`, which `try_push_task` already calls and which already recurses for
    /// `Group` and `Switch`, so a nested `ForEachRegion` is checked by the same pass. A pass claiming
    /// `by_value: false` must also be rejected outside a `ForEachRegion`, so the `Task::Pass` arm of
    /// `io_types` consults `signatures` too.
    fn io_types(&self) -> Result<(TypeId, TypeId), PassManagerError> { /* ... */ }
}

fn execute(
    task: &Task,
    mut ir: Box<dyn Any>,
    context: &mut PassContext,
    callbacks: Option<&CallbackRegistry>,
) -> Result<Box<dyn Any>, PassManagerError> {
    match task {
        Task::Pass(pass) => pass.run(ir, context).map_err(PassManagerError::PassError),
        Task::ForEachRegion { .. } => {
            execute_in_place(task, ir.as_mut(), context, callbacks)?;
            Ok(ir)
        },
        // ...
    }
}

fn execute_in_place(
    task: &Task,
    ir: &mut dyn Any,
    context: &mut PassContext,
    callbacks: Option<&CallbackRegistry>,
) -> Result<(), PassManagerError> {
    match task {
        Task::Pass(pass) => pass.run_in_place(ir, context).map_err(PassManagerError::PassError),
        Task::ForEachRegion { walker, body } => walker
            .walk(ir, &mut |region| {
                execute_in_place(body, region, context, callbacks).map_err(anyhow::Error::from)
            })
            .map_err(PassManagerError::PassError),
        // ...
    }
}

[Ian Hincks (ihincks)]

I guess my previous comments omits the most important thing: what needs to change in the RFC to accommodate something like this. As far as I can tell it's this:

• #[non_exhaustive] on any public enums, plus maybe constructor functions.
• QkPass must be opaque handle rather than a #[repr(C)] struct: don't want to promise the size won't change

Alternative 1 additionally needs

• The pass leaf must be Task::Pass(Pass) where Pass is a #[non_exhaustive] enum holding today's single erased trait.

[Julien Gacon (Cryoris)]

Thanks for working this through, Ian! Firstly:

• I did not know #[non_exhaustive] that's super useful. I'll add that to the Task, it's something that we might want to extend anyways in the future. I see that there's some constraints with the #[non_exhaustive] attribute, but they seem reasonable as of now.
• QkPass as opaque pointer is something we wanted to do anyways, I just didn't get around to updating the doc yet -- so that sounds good to me too.

Then a question on your snippet: I would imagine that the pass manager would decide pass-by-pass which run method to use, so that you could have passes with different execution modes in the same compiler flow. Is that what you also had in mind?

[Ian Hincks (ihincks)]

the pass manager would decide pass-by-pass which run method to use,

I suppose so. In the current example, it's pretty straight forward, because there are only two, and built into the execute function.