MLIR extraction for custom circuit

Catalyst
1

Dear Community,
I am new to the Catalyst package and I am currently on learning curve. I have been going through the documentation to understand the Catalyst better. At the page Catalyst Compiler Passes — Catalyst 0.6.0 documentation, I see there is an example of Catalyst’s IR and then I got curious how can I extract IR for my circuit. I looked in to the code of jit.py file catalyst/frontend/catalyst/jit.py at main · PennyLaneAI/catalyst · GitHub. I can see that there is a function generate_ir() for the IR. Since this is part of qjit decorator I am not sure how to extract IR string.
Is there any other way where I can extract MLIR for my custom circuit?

Thanks

2

hi @quantum-girl! For sure, the main way we can extract IR from a compiled circuit is through the .mlir attribute:

@qjit(keep_intermediate=True)
def func(x: float):
    return x

func.mlir

Note that you will need to make sure that the qjit function has been compiled before calling the .mlir attribute.

If you are interested in more finegrained details, you can also checkout the print_compilation_stage() function:

@qjit(keep_intermediate=True)
def func(x: float):
    return x

catalyst.debug.print_compilation_stage(func, "HLOLoweringPass")
1 Like
3

@josh Thanks! that’s what I was looking for.
So that means, function has to be compiled. Either it’s called first time or argument type hints are provided for “ahead of the time” compilation.

May be a stupid question, Why it has to be compiled? is it because of the you need to first convert fun to JAXPR and then lower from JAXPR to the quantum MLIR representation?

4

May be a stupid question, Why it has to be compiled? is it because of the you need to first convert fun to JAXPR and then lower from JAXPR to the quantum MLIR representation?

Hi @quantum-girl, that’s exactly right! Technically the function doesn’t have to be compiled all the way down to binary code, but the compilation process has to proceed at least far enough for us to generate the MLIR. Hope that helps :slight_smile:

1 Like
5

Hello.

Since I have a requirement for extracting the MLIR before compilation, I was wondering if there’s any workaround that I can implement that helps me achieve this.

Thanks.

6

Hi @annkay, thank you for posting your question!

Could you clarify what you mean by “before compilation”?

The .mlir attribute that @josh mentioned returns the program before it is processed by the “core” Catatyst compiler, so that is the earliest form of MLIR we generate.
Technically, there is one compilation step that happens in Python prior to reaching the Catalyst compiler, which is gate decomposition according to the chosen device’s supported gates and features. This doesn’t happen in MLIR however, so you cannot generate an MLIR representation of the program before decomposition.

Let me know if this answers your question :slight_smile:

7

@David_Ittah Thank you for the swift response.

We want to extract the MLIR within the compilation step, without executing/simulating the circuit with Pennylane. We manage to trigger circuit compilation only by running the quantum circuit directly, which ends up simulating it, and this is what we are trying to avoid.

8

Oh I see, in that case it sounds like you’re looking for ahead-of-time compilation, which will compile your qjit function upon definition. Alternately, if you don’t want to provide the type hints, you can call this method manually, which will only compile but not run the function:

my_function.jit_compile((x, y, z, ...))
9

Ahead-of-time compliation was just what we were looking for. Thank you @David_Ittah.

1 Like