Context:
Several places in the Catalyst codebase assume that the first non-declaration func.func in the module is the program entry point:
frontend/catalyst/jit.py : accesses self.mlir_module.body.operations[0]
frontend/catalyst/utils/gen_mlir.py : adds llvm.emit_c_interface to the first op
mlir/lib/Catalyst/Analysis/ResourceAnalysis.cpp : counts qubit arguments only for the first function
- . . .
This is not guaranteed by MLIR and relies on JAX's tracing placing the entry function first. If JAX ever reorders module contents, or if future passes move/clone functions, this assumption silently breaks.
Proposed Solution:
Introduce an explicit attribute (e.g., catalyst.entry_point or catalyst.main) that is attached to the entry function during tracing. All downstream consumers (resource analysis, codegen, etc.) should detect the entry point via this attribute rather than positional assumption.
The existing llvm.emit_c_interface attribute or the jit_ name prefix could serve as detection mechanisms, but a dedicated Catalyst attribute is cleaner and self-documenting.
Raised in PR #2782 discussion.
Context:
Several places in the Catalyst codebase assume that the first non-declaration func.func in the module is the program entry point:
frontend/catalyst/jit.py: accessesself.mlir_module.body.operations[0]frontend/catalyst/utils/gen_mlir.py: addsllvm.emit_c_interfaceto the first opmlir/lib/Catalyst/Analysis/ResourceAnalysis.cpp: counts qubit arguments only for the first functionThis is not guaranteed by MLIR and relies on JAX's tracing placing the entry function first. If JAX ever reorders module contents, or if future passes move/clone functions, this assumption silently breaks.
Proposed Solution:
Introduce an explicit attribute (e.g., catalyst.entry_point or catalyst.main) that is attached to the entry function during tracing. All downstream consumers (resource analysis, codegen, etc.) should detect the entry point via this attribute rather than positional assumption.
The existing
llvm.emit_c_interfaceattribute or thejit_name prefix could serve as detection mechanisms, but a dedicated Catalyst attribute is cleaner and self-documenting.Raised in PR #2782 discussion.