[MLIR + Frontend] Quantum Module Abstraction #1144

erick-xanadu · 2024-09-19T15:44:47Z

Context: We create nested modules for qnodes that contain all the program for a specific qnode. This will allow us to perform transformations targetting specific qnodes.

Description of the Change:

Adds catalyst.launch_kernel that can call functions in nested modules.
Adds pass to keep previous behaviour by:
- determining the name of a function given a scope
- renaming all functions in nested modules to be unique
- creating a map from flat symbol names to symbol names with path
- inlining nested modules
- transforming catalyst.launch_kernel to func.call

Benefits:

Nothing is currently using this pass / abstraction. But now we have the scaffolding for qnode specific transformations and lowerings.

Possible Drawbacks: More compilation time. Nested nested qnodes semantics are still needed. It was opted to disable tests with nested qnodes as they are not fully supported. Only nested qnodes where no qnode execution overlap was supported previously, but this is not the general case and could easily lead to runtime errors.

Related GitHub Issues:

TODO (minor changes):

cleaning up the frontend
(optional) upstream changes to symbol table, to avoid patterns on gradient interface and zne.
use values from ctx.module_context instead of hard coded ones
can we avoid changing the order of functions as they appear? (This would help lit testing)
remove func_p?

[sc-67095]

codecov · 2024-09-19T20:30:13Z

Codecov Report

Attention: Patch coverage is 99.16667% with 1 line in your changes missing coverage. Please review.

Project coverage is 97.89%. Comparing base (3bea1b6) to head (8e0890e).
Report is 199 commits behind head on main.

Files with missing lines	Patch %	Lines
frontend/catalyst/jax_primitives.py	98.98%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1144      +/-   ##
==========================================
+ Coverage   97.86%   97.89%   +0.03%     
==========================================
  Files          76       76              
  Lines       10819    10888      +69     
  Branches     1283     1290       +7     
==========================================
+ Hits        10588    10659      +71     
+ Misses        179      178       -1     
+ Partials       52       51       -1

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

frontend/catalyst/jax_primitives.py

dime10

Nice work 💯 I left some suggestions regarding the frontend which relate to a lot of your Todos as well I think :)

doc/releases/changelog-dev.md

frontend/catalyst/compiler.py

frontend/catalyst/jax_primitives.py

mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp

mlir/test/Catalyst/NestedModule5.mlir

frontend/catalyst/jax_primitives.py

frontend/test/lit/test_peephole_optimizations.py

frontend/catalyst/jax_primitives.py

paul0403 · 2024-09-26T20:33:33Z

Thanks! 🥇

An unrelated comment: once this qnode module scope is in place I'd like to utilize it and convert how we schedule passes in transform_named_sequence. No need to do that in this PR.

mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp

Co-authored-by: David Ittah <[email protected]>

dime10

Nice work 💯 I'm happy to approve pending resolution of the open comments :)

frontend/catalyst/from_plxpr.py

frontend/catalyst/jax_primitives.py

mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp

dime10

🦝 🚀 🌔

**Context:** We create nested modules for qnodes that contain all the program for a specific qnode. This will allow us to perform transformations targetting specific qnodes. **Description of the Change:** * Adds `catalyst.launch_kernel` that can call functions in nested modules. * Adds pass to keep previous behaviour by: * determining the name of a function given a scope * renaming all functions in nested modules to be unique * creating a map from flat symbol names to symbol names with path * inlining nested modules * transforming `catalyst.launch_kernel` to `func.call` **Benefits:** Nothing is currently using this pass / abstraction. But now we have the scaffolding for qnode specific transformations and lowerings. **Possible Drawbacks:** More compilation time. Nested nested qnodes semantics are still needed. It was opted to disable tests with nested qnodes as they are not fully supported. Only nested qnodes where no qnode execution overlap was supported previously, but this is not the general case and could easily lead to runtime errors. **Related GitHub Issues:** [sc-67095] --------- Co-authored-by: David Ittah <[email protected]>

erick-xanadu marked this pull request as draft September 19, 2024 15:44

erick-xanadu force-pushed the eochoa/2024-09-16/nested-module branch 4 times, most recently from 3bc3bd4 to 6acdf4d Compare September 19, 2024 19:42

erick-xanadu marked this pull request as ready for review September 19, 2024 19:43

Add inline nested modules pass.

4d9d6cd

erick-xanadu force-pushed the eochoa/2024-09-16/nested-module branch 2 times, most recently from 009342d to 845e927 Compare September 19, 2024 19:49

erick-xanadu added 2 commits September 19, 2024 16:04

Frontend changes

6f7df3c

changelog

1f41812

erick-xanadu force-pushed the eochoa/2024-09-16/nested-module branch from 845e927 to 1f41812 Compare September 19, 2024 20:09

erick-xanadu marked this pull request as draft September 19, 2024 21:26

erick-xanadu force-pushed the eochoa/2024-09-16/nested-module branch from 2979c40 to a9e0c6b Compare September 20, 2024 19:07

Generalize gradients to support SymbolRefAttr

97f4771

erick-xanadu force-pushed the eochoa/2024-09-16/nested-module branch from a9e0c6b to 97f4771 Compare September 20, 2024 19:09

erick-xanadu added 2 commits September 20, 2024 15:56

Generalize mitigation dialect

780cbd2

Mitigation changes frontendf

b6b5832

erick-xanadu marked this pull request as ready for review September 20, 2024 19:56

f

d2a9311

erick-xanadu commented Sep 20, 2024

View reviewed changes

frontend/catalyst/jax_primitives.py Outdated Show resolved Hide resolved

erick-xanadu commented Sep 20, 2024

View reviewed changes

frontend/catalyst/jax_primitives.py Outdated Show resolved Hide resolved

erick-xanadu commented Sep 20, 2024

View reviewed changes

frontend/catalyst/jax_primitives.py Outdated Show resolved Hide resolved

Apply suggestions from code review

fdaf78d

dime10 reviewed Sep 25, 2024

View reviewed changes

paul0403 reviewed Sep 26, 2024

View reviewed changes

mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp Outdated Show resolved Hide resolved

erick-xanadu and others added 5 commits October 1, 2024 15:27

Changelog

7b9199a

Codefactor

d07c923

Apply suggestions from code review

9aa6186

Co-authored-by: David Ittah <[email protected]>

isort

9d8747e

Use matchAndRewrite instead of match and rewrite

ec77fe9

erick-xanadu requested a review from dime10 October 1, 2024 20:47

erick-xanadu added 3 commits October 2, 2024 11:36

lit test

e695478

rename _module_to_mlir to _module_lowering

744ff53

pylint

7922121

dime10 reviewed Oct 2, 2024

View reviewed changes

erick-xanadu added 14 commits October 3, 2024 13:17

Merge branch 'main' into eochoa/2024-09-16/nested-module

fee4540

Change call_nested_module to launch_kernel

769fe0f

Change name of pass in python pipeline

9037dcd

Add comment for primitive lowering.

f830146

Rename module_p to quantum_kernel_p

de89c46

f

096ed3b

Rename _callable to callable_

15aa359

Rename fn to qfunc

5f57ee8

Remove redundant code

11523c8

Rename get to compute

057a9b9

Merge branch 'main' into eochoa/2024-09-16/nested-module

a3a54c7

Limit launch_kernel to symbols with public visibility

f6a901d

Fix lit tests

248f409

Add test for code coverage

8e0890e

dime10 approved these changes Oct 4, 2024

View reviewed changes

erick-xanadu merged commit aeb1e43 into main Oct 7, 2024
43 checks passed

erick-xanadu deleted the eochoa/2024-09-16/nested-module branch October 7, 2024 13:19

paul0403 mentioned this pull request Oct 7, 2024

Integrate transform dialect scheduling into quantum module abstraction #1185

Open

dime10 mentioned this pull request Oct 8, 2024

ZNE works with dynamic one shot #1169

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MLIR + Frontend] Quantum Module Abstraction #1144

[MLIR + Frontend] Quantum Module Abstraction #1144

erick-xanadu commented Sep 19, 2024 •

edited

Loading

codecov bot commented Sep 19, 2024 •

edited

Loading

dime10 left a comment

paul0403 commented Sep 26, 2024 •

edited

Loading

dime10 left a comment •

edited

Loading

dime10 left a comment

[MLIR + Frontend] Quantum Module Abstraction #1144

[MLIR + Frontend] Quantum Module Abstraction #1144

Conversation

erick-xanadu commented Sep 19, 2024 • edited Loading

codecov bot commented Sep 19, 2024 • edited Loading

Codecov Report

dime10 left a comment

Choose a reason for hiding this comment

paul0403 commented Sep 26, 2024 • edited Loading

dime10 left a comment • edited Loading

Choose a reason for hiding this comment

dime10 left a comment

Choose a reason for hiding this comment

erick-xanadu commented Sep 19, 2024 •

edited

Loading

codecov bot commented Sep 19, 2024 •

edited

Loading

paul0403 commented Sep 26, 2024 •

edited

Loading

dime10 left a comment •

edited

Loading