rigetti / pyquil

A Python library for quantum programming using Quil.
http://docs.rigetti.com
Apache License 2.0
1.39k stars 341 forks source link

feat: Bootstrap experimental, toggleable high-performance backend for Quil programs. #1755

Open MarquessV opened 3 months ago

MarquessV commented 3 months ago

tl;dr: Make pyQuil a hybrid Rust package that wraps quil-rs directly, get better performance and code quality.

I'll start with some leading questions to inform the review:

Description

When we built pyQuil v4, we started by building quil, a Python package with Rust bindings to quil-rs. However, quil-rs was re-built from first principles, with all the lessons of pyQuil, to serve solely as a reference to the Quil spec. Its approach and API is different, and in order to not break pyQuil, we added a compatibility layer as to not rock the boat.

This compatibility layer is implemented in Python and has a negligible performance impact when dealing with a small amount of instructions, however, it scales poorly as instruction count grows. Recently, we investigated the performance of iterating over a Program's instructions property, and saw less than ideal performance as instructions got into the thousands (see #1754).

Currently, the path from quil-rs, to pyQuil's public facing API looks something like this:

image

The compatibility layer, being in Python, is slow. Not only that, but quil is really just a middle-man between quil-rs and pyQuil. The changes in this PR illustrate a pattern that would look something like this:

image

By porting all of pyQuil's core Quil logic into Rust and wrapping quil-rs directly, we eliminate the need for quil entirely, simplifying the dependency chain. In addition, the transformation logic is all performed in the more performant Rust. This improves performance (receipts below), and I think, improves the quality of the code. Due to not wanting to sacrifice the quality of quil-py, we made certain compromises in pyQuil. For example, we use a custom metaclass to make all quil-rs types appear as AbstractInstruction types. This mostly works, but has some odd corner cases if you start to mix quil-rs types with pyQuil types. That problem stems from the fact that the quil Instruction class can't be used as a subclass because it's an enum (a limitation with pyo3/C ABI), which also means we have lots of duplicated code for every instruction to implement things like __copy__ or pickling. Support for other things like pickling are inconsistent because of this. Moving the compatibility layer into Rust solves all of these problems.

Performance improvements

For each backend, parse a program from a string containing 9001 lines, then iterate over every instruction 1000 times. Benchmark with hyperfine.

Command Mean [s] Min [s] Max [s] Relative
Current pyQuil backend 20.089 ± 0.213 19.736 20.322 8.70 ± 0.16
Proposed Rust core 2.310 ± 0.034 2.276 2.365 1.00

Data collected on a 2021 Macbook Pro M1 Max

Benchmark script

hyperfine 'python bench_quil_backends.py python' 'python bench_quil_backends.py rust' --warmup 2
import argparse

from pyquil._core import program as program_core
from pyquil.quil import Program

PROGRAM_BACKENDS = {
    "python": Program,
    "rust": program_core.Program,
}

with open("test/bench/fixtures/over-9000.quil", "r") as f:
    PROGRAM_STRING = f.read()

def iterate_over_instructions(program):
    for _ in range(1000):
        for _ in program:
            continue

def main():
    parser = argparse.ArgumentParser(description="Iterates over instructions using the  specified backend.")
    parser.add_argument(
        "backend", choices=["python", "rust"], help="The backend to run benchmarks for. Options: python, rust."
    )

    args = parser.parse_args()

    backend = PROGRAM_BACKENDS[args.backend]

    program = backend(PROGRAM_STRING)
    assert program.out() == PROGRAM_STRING, "Program should round trip correctly"
    iterate_over_instructions(program)

if __name__ == "__main__":
    main()

Tracking issue #1760

MarquessV commented 3 months ago

To be transparent, I am hesitant about the approach for these reasons:

* in our migration to Rust and pyQuil v4, we intended to leave pyquil "as-is" as much as possible. This is a clear departure from that and appears to be an unintended consequence of the move to Rust. That said: it should still serve the purposes of its users, so long as all of the helper functions and python internals that they rely on today are upheld by the new Rust-based classes.

* it's a lot of work to keep the API and functionality roughly the same

That said, the current performance limitations (within the python-side instruction processor) are not things we can leave unfixed, and this is a principled way of fixing them.

I think we share the same concerns. I would add that even though it is a departure from keeping pyQuil "as-is", in hindsight, the compatibility layer is a thin illusion that makes it seem like we kept things the same more than we actually did.

my only recommendation would be to expand the tests to cover all the pythonic ways that users might use or query instructions - all the dunder methods, etc.

Agreed, I think test_quilbase.py is the right approach, but I need to add test cases for the default dunder methods Python provides that I took for granted. Fortunately, the proposed backend will make this much easier, since we can actually implement common functionality on the Instruction base class, which isn't something we can do with quil-py today.

I'd also want to validate this version through our power users' notebooks before merge. I know that some are covered today, but from the conversation last week it sounds like others may be missing.

Same, and I think I can make this easy in a low risk way. I want to introduce a configuration option (e.g. a PYQUIL_EXPERIMENTAL_BACKEND environment variable) that, when enabled, would patch the existing instruction and program module with this proposed backend:

The one downside is that we have to maintain two backends, but changes are relatively infrequent, and I think that is the right tradeoff to make for the benefits above.

can you add the other function from the internal report, copy_everything_except_instructions, for a similar program as reported?

Will do!

* needs a linked and prefixed issue. Every MR needs one.

Right, created #1760

github-actions[bot] commented 1 month ago

☂️ Python Coverage

current status: ✅

Overall Coverage

Lines Covered Coverage Threshold Status
7221 6364 88% 87% 🟢

New Files

No new covered files...

Modified Files

File Coverage Status
pyquil/init.py 78% 🟢
pyquil/quilatom.py 83% 🟢
pyquil/quilbase.py 94% 🟢
TOTAL 85% 🟢

updated for commit: 19ab61f by action🐍