Issue #1714: Add variable names to the IR #2054

rturrado · 2025-09-18T22:12:26Z

Context:
We can investigate the option of passing down source code information like variable names from the frontend down to the generated ir, using MLIR's location mechanism. The additional info will make generated code more readable.

Description of the Change:

User -> Frontend: @qjit(use_nameloc=True) to use name location information in the generated MLIR code.
Frontend -> Catalyst CLI: whenever use_nameloc is enabled:
- Pass stdin=self.mlir_module.operation.get_asm(enable_debug_info=True) instead of stdin=str(self.mlir_module).
- Pass --use-nameloc-as-prefix in options.
Catalyst CLI -> frontend: whenever --use-nameloc-as-prefix is enabled:
- Enable useNameLocAsPrefix printing flag.
- Print the MLIR module to stdout with printing flags, e.g., mlirModule->print(oss, opPrintingFlags) instead of oss << *mlirModule.

Benefits:

@qjit(use_nameloc=True) is an eay way for the user to specify that they want to use name location information.
The same applies to Catalyst CLI's --use-nameloc-as-prefix option.

Possible Drawbacks:

get_arg_names function may not be robust enough.
Name location information works nicely for function parameter names (e.g., %x instead of %arg0), but maybe not so nicely for results of equations (e.g., %jit28ok292Fjit28jit_f292Fmul instead of %0). An option would be to add name location information only for function parameters.

Related GitHub Issues:
#1714

Add test_alice_and_bob_parameters. Rename test_alice to test_alice_qubit. Rename test_alice_and_bob to test_alice_and_bob_qubits.

lower_jaxpr_to_mlir: add arg_names, and pass them on to jaxpr_to_mlir. jaxpr_to_mlir: change parameter order, add arg_names, and pass them on to custom_lower_jaxpr_to_module. custom_lower_jaxpr_to_module: add arg_names, and pass them on to lower_jaxpr_to_fun. This latter function builds arg_locs with the arg_names and appends them to the function's entry_block. qjit: add use_nameloc parameter. CompileOptions: add use_nameloc option. _options_to_cli_flags: add '-mlir-use-nameloc-as-prefix' if use_nameloc=True. test_debug.py: add test_option_use_nameloc.Add frontend/test/lit/test_option_use_nameloc.py. TODO: the name location information is added to the MLIR module, but still not consumed by quantum-opt. TODO: should we try adding name location information for other than funcion arguments?

… name location information in Catalyst Add QJIT.get_arg_names() to construct a list of argument names for a qjit function. Change QJIT.mlir and QJIT.mlir_opt to optionally use an MLIR string with debug info. Change canonicalize function to optionally pass --use-nameloc-as-prefix. Add --use-nameloc-as-prefix to Catalyst CLI options. Change QuantumDriverMain to optionally use useNameLocAsPrefix printing flag when printing the MLIR module. Add test_get_arg_names.py. Change test_option_use_nameloc.py to check that: 1) MLIR module contains location information. 2) MLIR code contains location information. 3) f.mlir uses the qjit function parameter names, %x and %y, instead of %arg0 and %arg1.

…mes_to_the_ir

mlxd · 2025-09-25T13:37:59Z

Thanks @rturrado for the submission.
We are currently finalizing our plans for Q4, so may take a little longer to get around to reviewing this.
Thanks in advance for your patience.

rturrado · 2025-09-25T13:49:06Z

Thanks @rturrado for the submission.
We are currently finalizing our plans for Q4, so may take a little longer to get around to reviewing this.
Thanks in advance for your patience.

Many thanks for the update, @mlxd. No rush at all!

codecov · 2025-10-08T13:50:25Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.50%. Comparing base (687c7e8) to head (9634bc9).
⚠️ Report is 7 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2054      +/-   ##
==========================================
+ Coverage   97.49%   97.50%   +0.01%     
==========================================
  Files          91       91              
  Lines       10604    10624      +20     
  Branches      995     1001       +6     
==========================================
+ Hits        10338    10359      +21     
+ Misses        211      210       -1     
  Partials       55       55

☔ 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.

dime10 · 2025-10-08T22:48:43Z

Thank you for patience Roberto, I should be able to come back to this PR in the next few days :)

Add print_mlir_opt to frontend/test/lit/utils.py. Invoke print_mlir_opt from frontend/test/lit/test_option_use_nameloc.py.

rturrado · 2025-10-09T16:12:27Z

Thank you for patience Roberto, I should be able to come back to this PR in the next few days :)

No rush at all, @dime10 !

dime10

Nice work, another excellent PR! 🥇

dime10 · 2025-10-14T21:43:24Z

doc/catalyst-cli/catalyst-cli.rst

+``--use-nameloc-as-prefix[=<true|false>]``
+""""""""""""""""""""""""""""""""""""""""""
+
+Print SSA IDs using their name location, if provided, as prefix. By default, name location information is not used.


Maybe we could say a bit more about what a "name location" is :)

dime10 · 2025-10-14T21:47:35Z

frontend/test/lit/test_option_use_nameloc.py

+assert str(f.mlir_module.body.operations[0].arguments[1].location) == 'loc("y")'
+
+print_mlir(f, 0.3, 0.4)
+print_mlir_opt(f, 0.3, 0.4)


Does the second line do anything here? Typically we the print the IR in the lit tests so that the CHECK statements can be matched against it, but if they already matched once, they aren't going to be matched against anything further down in the stream.

dime10 · 2025-10-14T21:49:26Z

mlir/test/Catalyst/PrintUseNameLocAsPrefix.mlir

@@ -0,0 +1,31 @@
+// RUN: quantum-opt %s -mlir-use-nameloc-as-prefix -split-input-file | FileCheck %s


This is not our pass is it? If not, we're just testing upstream code right 🤔

dime10 · 2025-10-14T21:58:24Z

frontend/catalyst/jit.py

+        if self.compile_options.use_nameloc:
+            stdin = self.mlir_module.operation.get_asm(enable_debug_info=True)
+        else:
+            stdin = str(self.mlir_module)


Is this different from using this method unconditionally?

Suggested change

if self.compile_options.use_nameloc:

stdin = self.mlir_module.operation.get_asm(enable_debug_info=True)

else:

stdin = str(self.mlir_module)

stdin = self.mlir_module.operation.get_asm(enable_debug_info=self.compile_options.use_nameloc)

dime10 · 2025-10-14T22:02:25Z

frontend/catalyst/jit.py


        return jaxpr, out_type, treedef, dynamic_sig

+    def get_arg_names(self):


I'm hesitant to expose another method in the QJIT class as we want to keep it relatively stable and avoid class bloat. Could this be implemented as a free-standing function in the tracing submodule (or other jax-related modules if more appropriate, but I think type_signatures.py looks pretty relevant)?

rturrado and others added 11 commits August 26, 2025 13:27

Add catalyst/mlir/test/PrintUseNameLocAsPrefix.mlir file.

9972987

Update PrintUseNameLocAsPrefix.mlir.

55a84c2

Add test_alice_and_bob_parameters. Rename test_alice to test_alice_qubit. Rename test_alice_and_bob to test_alice_and_bob_qubits.

[no ci] bump nightly version

c0f02a5

[no ci] bump nightly version

5e58c62

Merge branch 'main' into issue_1714_add_variable_names_to_the_ir

f20430c

[no ci] bump nightly version

9e0a996

[no ci] bump nightly version

03d0a2b

[no ci] bump nightly version

f7f7a7e

Fix test_option_use_nameloc.py.

17b1cfd

github-actions bot added the external PRs where the author is not a part of PennyLane Org (or part of external contributors team) label Sep 18, 2025

rturrado and others added 4 commits September 19, 2025 00:25

Fix CodeFactor checks.

726f094

Merge branch 'PennyLaneAI:main' into main

0cbe469

Merge remote-tracking branch 'origin' into issue_1714_add_variable_na…

6ee021e

…mes_to_the_ir

[no ci] bump nightly version

daf2c99

rturrado changed the title ~~Issue 1714 add variable names to the ir~~ Issue #1714: add variable names to the IR Sep 19, 2025

rturrado changed the title ~~Issue #1714: add variable names to the IR~~ Issue #1714: Add variable names to the IR Sep 19, 2025

ringo-but-quantum added 4 commits September 20, 2025 03:47

[no ci] bump nightly version

9df3bf6

[no ci] bump nightly version

a7ebdd0

[no ci] bump nightly version

e139daf

[no ci] bump nightly version

e6b3d4f

ringo-but-quantum and others added 6 commits September 26, 2025 03:50

[no ci] bump nightly version

126e626

[no ci] bump nightly version

9c72207

[no ci] bump nightly version

8e3cc1a

[no ci] bump nightly version

a58843d

[no ci] bump nightly version

6da9eb1

Merge remote-tracking branch 'upstream/main'

58b79b0

Update to origin/main.

e81eb00

Run black.

67df639

rturrado added 2 commits October 9, 2025 18:07

Merge branch 'main' into issue_1714_add_variable_names_to_the_ir

7748331

Fix code coverage warning about uncovered code in mlir_opt

9634bc9

Add print_mlir_opt to frontend/test/lit/utils.py. Invoke print_mlir_opt from frontend/test/lit/test_option_use_nameloc.py.

dime10 reviewed Oct 14, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Issue #1714: Add variable names to the IR #2054

Issue #1714: Add variable names to the IR #2054

Uh oh!

rturrado commented Sep 18, 2025

Uh oh!

mlxd commented Sep 25, 2025

Uh oh!

rturrado commented Sep 25, 2025

Uh oh!

codecov bot commented Oct 8, 2025 •

edited

Loading

Uh oh!

dime10 commented Oct 8, 2025

Uh oh!

rturrado commented Oct 9, 2025

Uh oh!

dime10 left a comment

Uh oh!

dime10 Oct 14, 2025

Uh oh!

dime10 Oct 14, 2025

Uh oh!

dime10 Oct 14, 2025

Uh oh!

dime10 Oct 14, 2025

Uh oh!

dime10 Oct 14, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		@@ -0,0 +1,31 @@
		// RUN: quantum-opt %s -mlir-use-nameloc-as-prefix -split-input-file \| FileCheck %s


		return jaxpr, out_type, treedef, dynamic_sig

		def get_arg_names(self):

Issue #1714: Add variable names to the IR #2054

Are you sure you want to change the base?

Issue #1714: Add variable names to the IR #2054

Uh oh!

Conversation

rturrado commented Sep 18, 2025

Uh oh!

mlxd commented Sep 25, 2025

Uh oh!

rturrado commented Sep 25, 2025

Uh oh!

codecov bot commented Oct 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

dime10 commented Oct 8, 2025

Uh oh!

rturrado commented Oct 9, 2025

Uh oh!

dime10 left a comment

Choose a reason for hiding this comment

Uh oh!

dime10 Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

dime10 Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

dime10 Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

dime10 Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

dime10 Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov bot commented Oct 8, 2025 •

edited

Loading