pyquil icon indicating copy to clipboard operation
pyquil copied to clipboard

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

Open MarquessV opened this issue 1 year ago • 2 comments

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

  • This PR bootstraps a new backend that uses said strategy by opting in the new Program via an environment variable. This allows us to make these changes in atomic pieces, without breaking the existing compatibility layer. When the backend is more mature, users can choose to opt-in to the new default in a relatively friction-free way.
  • As this new backend is incomplete, this feature is undocumented and only available if building from source. Future PRs will document this feature and distribute wheels that include the Rust backend when it is mature enough for users to experiment with.

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

  • Are the patterns illustrated by this PR acceptable?
  • Is there any other validation or pattern you'd want to see before pushing forward with this?
  • Any specific benchmarks you want to see?
  • Any other concerns, or reasons we shouldn't do this?

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 avatar Mar 25 '24 00:03 MarquessV

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:

  • That makes it easy to test the new backend without adding the friction of swapping between pyQuil releases, or frustrating users if things don't work.
  • It makes it easy for us to develop and fix issues in the new backend without having an RC in limbo, slowing down our ability to deliver features unrelated to the new backend.
  • We can transition to the new backend without getting rid of the old one, giving users a workaround if issues come up after we make it the default. When issues stop coming up, we can remove the old 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

MarquessV avatar Mar 29 '24 22:03 MarquessV

☂️ 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🐍

github-actions[bot] avatar May 23 '24 22:05 github-actions[bot]