[python] Move kernel signature inference out of MLIR compilation (#3930)

I've broken up this PR into two commits that can be reviewed
independently:

The first commit contains "clean-ups". It doesn't introduce any new
functionality nor should it change how anything is currently working. It
just removes function arguments that aren't used, renames a couple of
things and should make the code more readable. I'm including this here
because it made my actual work much easier. These breaking changes are
acceptable as they are within the `python/cudaq/kernel` folder, which is
not publicly exposed (discussed with Bettina).

The second commit is the real contribution:
- The main goal is to parse the kernel signature _before_ running the
MLIR compilation. This is required so that we can delay MLIR compilation
until the first kernel invocation (in a future PR).
- To this end, an AST visitor in `python/cudaq/kernel/analysis.py` is
extended to parse the function annotations. This can be run before MLIR
compilation
- A type `KernelSignature` is introduced to consolidate all kernel type
information that was scattered across many attributes in one place.
- `KernelSignature` now only stores the kernel types as MLIR types
(before, both Python and MLIR types were stored). This makes it easier
to handle as there is a single source of truth (and Python types can be
recovered using `mlirTypeToPyType`)

EDIT: Added a third commit that fixes issue #2895

I've noted that a lot of this Python code uses CamelCase, even though
our style guide ([Google style
guide](#2895)) recommends
snake_case. I'm using snake_case where new names/variables are
introduced, but limited the number of existing names that had to be
changed.

closes #2895

---------

Signed-off-by: Luca Mondada <luca@mondada.net>

Author: lmondada

python/cudaq/kernel/analysis.py

@@ -8,27 +8,41 @@
 
 import ast
 import inspect
-import importlib
 import textwrap
+from typing import Optional, Type
 
-from cudaq.mlir._mlir_libs._quakeDialects import cudaq_runtime
-from cudaq.mlir.dialects import cc
-from .utils import mlirTypeFromAnnotation
 
-
-class HasReturnNodeVisitor(ast.NodeVisitor):
+class FunctionDefVisitor(ast.NodeVisitor):
     """
-    This visitor will visit the function definition and report 
-    true if that function has a return statement.
+    This visitor will visit the function definition of `kernel_name` and report 
+    type annotations and whether the function has a return statement.
     """
 
-    def __init__(self):
-        self.hasReturnNode = False
+    arg_annotations: list[(str, Type)]
+    return_annotation: Optional[Type] = None
+    has_return_statement: bool = False
+    found: bool = False
+
+    def __init__(self, kernel_name: str):
+        self.kernel_name: str = kernel_name
+        self.arg_annotations = []
 
     def visit_FunctionDef(self, node):
-        for n in node.body:
-            if isinstance(n, ast.Return) and n.value != None:
-                self.hasReturnNode = True
+        if node.name == self.kernel_name:
+            self.found = True
+            self.arg_annotations = [
+                (arg.arg, arg.annotation) for arg in node.args.args
+            ]
+            self.return_annotation = node.returns
+            self.has_return_statement = any(
+                isinstance(n, ast.Return) and n.value != None
+                for n in node.body)
+
+    def generic_visit(self, node):
+        if self.found:
+            # skip traversing the rest of the AST once found
+            return
+        super().generic_visit(node)
 
 
 class FindDepFuncsVisitor(ast.NodeVisitor):

python/cudaq/kernel/ast_bridge.py

@@ -1360,39 +1360,36 @@ def resolve_callable_arg(self, insPt, target):
         closure here.
         Returns a `CreateLambdaOp` closure.
         """
+        # Add the target kernel to the current module.
         cudaq_runtime.updateModule(self.uniqName, self.module, target.qkeModule)
+        fulluniq = nvqppPrefix + target.uniqName
+        fn = recover_func_op(self.module, fulluniq)
+
         # build the closure to capture the lifted `args`
         thisPyMod = recover_calling_module()
         if target.defModule != thisPyMod:
             m = target.defModule
         else:
             m = None
-        fulluniq = nvqppPrefix + target.uniqName
-        fn = recover_func_op(self.module, fulluniq)
-        funcTy = fn.type
-        if target.firstLiftedPos:
-            moduloInTys = funcTy.inputs[:target.firstLiftedPos]
-        else:
-            moduloInTys = funcTy.inputs
-        callableTy = cc.CallableType.get(self.ctx, moduloInTys, funcTy.results)
+        funcTy = target.signature.get_lifted_type()
+        callableTy = target.signature.get_callable_type()
         with insPt, self.loc:
             lamb = cc.CreateLambdaOp(callableTy, loc=self.loc)
             lamb.attributes.__setitem__('function_type', TypeAttr.get(funcTy))
             initRegion = lamb.initRegion
-            initBlock = Block.create_at_start(initRegion, moduloInTys)
+            initBlock = Block.create_at_start(initRegion, target.arg_types())
             inner = InsertionPoint(initBlock)
             with inner:
                 vs = []
                 for ba in initBlock.arguments:
                     vs.append(ba)
-                for i, a in enumerate(target.liftedArgs):
-                    v = recover_value_of(a, m)
+                for var in target.captured_variables():
+                    v = recover_value_of(var.name, m)
                     if isa_kernel_decorator(v):
                         # The recursive step
                         v = self.resolve_callable_arg(inner, v)
                     else:
-                        argTy = funcTy.inputs[target.firstLiftedPos + i]
-                        v = self.__getMLIRValueFromPythonArg(v, argTy)
+                        v = self.__getMLIRValueFromPythonArg(v, var.type)
                     vs.append(v)
                 if funcTy.results:
                     call = func.CallOp(fn, vs).result