New OQD end-to-end pipeline is added to Catalyst

[rniczh]

Context:

Description of the Change:

This pipeline enables Catalyst to compile circuit for ARTIQ-based quantum device. When device_db configuration is provided, the compilation follows the ARTIQ route:

• ions-decomposition: Decompose quantum operations for trapped ion systems
• gates-to-pulses: Convert gate operations to pulse sequences
• convert-ion-to-rtio: Convert ion operations to RTIO dialect
• convert-rtio-event-to-artiq: Lower RTIO event to ARTIQ's primitives
• llvm-dialect-lowering-stage: Lower to LLVM IR
• emit-artiq-runtime: Generate ARTIQ runtime entry point as wrapper for ARITQ device to execute

The final stage compiles LLVM IR to an ELF binary targeting the ARTIQ device (ARM Cortex-A9). Due to current limitations where Catalyst's internal LLVM build does not include the corresponding ARM backend, the compilation will use llc to compile .ll to object file, and use ld.lld to compile to .elf

We added a compile_to_artiq helper function to the oqd module, so it can compile and link Catalyst-generated LLVM IR to ARTIQ's binary, keep OQD-specific logic out of Catalyst core.

Example:
Replace the artiq configs to your own env setting:

import os

import numpy as np
import pennylane as qml

from catalyst import qjit
from catalyst.third_party.oqd import OQDDevice, OQDDevicePipeline, compile_to_artiq

OQD_PIPELINES = OQDDevicePipeline(
    os.path.join("calibration_data", "device.toml"),
    os.path.join("calibration_data", "qubit.toml"),
    os.path.join("calibration_data", "gate.toml"),
    os.path.join("device_db", "device_db.json"),
)


def test_rx_gate():
    """Test RX gate with ARTIQ linking done in user code."""
    artiq_config = {
        "kernel_ld": "/path/to/kernel.ld",
        "llc_path": "/path/to/llc",
        "lld_path": "/path/to/ld.lld",
    }

    oqd_dev = OQDDevice(
        backend="default",
        shots=4,
        wires=1,
        artiq_config=artiq_config
    )
    qml.capture.enable()

    # Compile to LLVM IR only
    @qml.qnode(oqd_dev)
    def circuit():
        x = np.pi / 2
        qml.RX(x, wires=0)
        return qml.counts(wires=0)
        
    compiled_circuit = QJIT(circuit, CompileOptions(link=False, pipelines=OQD_PIPELINES))

    # Compile to ARTIQ ELF
    output_elf_path = compile_to_artiq(compiled_circuit, oqd_dev.artiq_config)
    print(f"ARTIQ ELF file generated: {output_elf_path}")

test_rx_gate()

The result will store to circuit.elf

[ARTIQ] Generated ELF: circuit.elf

And you can run it on artiq device:

artiq_run  --device-db device_db.py circuit.elf

Benefits:

Possible Drawbacks:
It causes the OQD specific compile stuffs steps into the original catalyst compiler driver's pipepline

Related GitHub Issues:
[sc-100853]

[codecov]

Codecov Report

❌ Patch coverage is 45.00000% with 55 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.96%. Comparing base (6a496da) to head (0751f89).
⚠️ Report is 6 commits behind head on main.

Files with missing lines	Patch %	Lines
frontend/catalyst/third_party/oqd/oqd_compile.py	16.92%	54 Missing ⚠️
frontend/catalyst/jit.py	93.75%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2299      +/-   ##
==========================================
- Coverage   97.35%   96.96%   -0.40%     
==========================================
  Files         107      108       +1     
  Lines       13136    13223      +87     
  Branches     1074     1089      +15     
==========================================
+ Hits        12789    12822      +33     
- Misses        286      339      +53     
- Partials       61       62       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[paul0403]

Thanks Hongsheng, looks really nice! I think this is the minimally intrusive way we can do it without altering too much of the base jit flow!

I left some comments but all are minor. Another thing is can you add your PR description's example as an end-to-end pytest in the frontend oqd tests? Obviously no need for execution, just check that the compiled ELF exist (maybe check a little bit of its content as well)? You may need to do things like subprocess.run("which llc", shell=True) in the test, but I think that's possible.

[mehrdad2m]

Thanks @rniczh, I did a first round of review with some minor comments. However, I would like to play around with a bit more and maybe more review on Monday.

[dime10]

Alright I'll approve the qjit updates (note I haven't reviewed the OQD portions) because the changes are kept fairly small, however there is change in contract / behaviour of QJIT class in the following ways:

• compile is no longer guaranteed to output a CompiledFunction (wrapping binary code) which it was guaranteed to do until now
• __call__ / jit_compile will no longer force compilation of the function down to binary if the target is set as "llvmir", as a result run will raise an error for this case
  • this is different than the other target levels ("jaxpr", "mlir") which only control aot compilation, while calling the function will force the full compilation pipeline so the program can be executed
  • once the "llvmir" target is set there is no way to compile the program for execution, except by mutating the compile options in place or defining a new qjit object

The inconsistency could be solved by limiting the change (e.g. the internal link option) to take effect in AOT mode only. Alternatively, one could expand the target option to affect both AOT and and JIT mode for all targets, but I still think this comes with downsides (see the last bullet point above).
However, the change in the contract structure of QJIT (main idea: QJIT is broken down into discrete stages, each of which produces well defined results that can feed into the next stage) might only be resolvable by breaking the compile step up such that we have separate compile and link stages.

If @mehrdad2m is happy with the current state we can merge this PR but I think these two points will be added as technical debt.

[mehrdad2m]

If @mehrdad2m is happy with the current state we can merge this PR but I think these two points will be added as technical debt.

Oh I thought we agreed on keeping things as before for this PR as there is no consensus on the ideal behaviour. This is a very user facing feature and not also very related to this PR, so even if we want to change anything, I think it's best to get in sync with product team first and do it in another PR. @rniczh If you don't mind, let's revert the change that raises error for target=llvmir in jit mode for now until we discuss it more later. Thanks

[rniczh]

@mehrdad2m Sure, I reverted it in this commit: 67dba09

[mehrdad2m]

Wouldn't this still crash the jit_compile path if we have target=llvmir? 🤔

[rniczh]

Yeah, should remove those setting as well. So the current goal is to remove all LLVM IR specializations in the pipeline. For now, we're reverting to only adding the link setting. Since end-to-end is not possible without disabling link, when using OQD's end-to-end, we need to change to actively construct using QJIT.

Changed: 9d80502

Therefore, in terms of usage, we're changing the approach to first create a qnode with @qnode, then use QJIT with the current compile options to achieve this.

oqd_pipelines = _get_oqd_pipelines()

@qml.set_shots(4)
@qml.qnode(oqd_dev)
def circuit():
    x = np.pi / 2
    qml.RX(x, wires=0)
    return qml.counts(all_outcomes=True)

# Get the LLVM IR
compiled_circuit = QJIT(circuit, CompileOptions(link=False, pipelines=oqd_pipelines))
llvm_ir = compiled_circuit.llvmir

[mehrdad2m]

Thanks @rniczh . I think this looks good for now. This could be improved if we divide compile() into compile() and link() like @dime10 suggested. But we don't need to address this in this PR.

[mehrdad2m]

Did we remove a test for compile_to_artiq? Codecov is not happy about it, and it seems to be not used anywhere except in documentation.