Centralize Calling Convention and Simplify Binding Generation (#271)

This PR centralizes the C++ FFI convention to callconv.py and reuses the
callconv everywhere to reduce redundancy.

Additionally, CUDA13 header adds a few compiler native keywords / native
types that prevents clang from parsing the code without syntax error.
This PR adds in patches to prevent those error. Note that they should
not be considered as official support for the patched keyword /
datatypes.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Refactor**
* Unified call-convention framework for C++ function bindings,
constructors, conversion operators, and methods, simplifying shim
generation and lowering.

* **New Features**
* Added call-convention utilities into generated bindings and tooling to
centralize shim/call logic.
  * Introduced a utility to prepare IR argument types for CUDA targets.

* **Bug Fixes / Chores**
* Ensure shim include path is always resolved (now fails fast if
missing).
  * Include header shim files in distribution wheels.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Michael Wang <isVoid@users.noreply.github.com>

Author: isVoid

ast_canopy/CMakeLists.txt

@@ -32,3 +32,12 @@ target_link_libraries(pylibastcanopy
   PUBLIC astcanopy::astcanopy)
 
 install(TARGETS pylibastcanopy LIBRARY DESTINATION ast_canopy)
+
+# Ensure header shims are shipped inside the wheel under:
+#   site-packages/ast_canopy/shim_include/
+# scikit-build-core builds wheels from the CMake install tree, so non-built
+# sources must be installed explicitly.
+install(
+  DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/ast_canopy/shim_include
+  DESTINATION ast_canopy
+)

ast_canopy/ast_canopy/api.py

@@ -28,11 +28,18 @@
 logger = logging.getLogger(f"AST_Canopy.{__name__}")
 
 
-def _get_shim_include_dir() -> str | None:
-    """Return the absolute path to the local shim include directory, if present."""
+def _get_shim_include_dir() -> str:
+    """Return the absolute path to the local shim include directory"""
     here = os.path.dirname(__file__)
     shim_dir = os.path.join(here, "shim_include")
-    return shim_dir if os.path.isdir(shim_dir) else None
+
+    if not os.path.isdir(shim_dir):
+        raise RuntimeError(
+            f"Shim include directory not found at {shim_dir}. "
+            "This indicates a packaging issue. Please reinstall ast_canopy."
+        )
+
+    return shim_dir
 
 
 @dataclass
@@ -435,7 +442,7 @@ def parse_declarations_from_source(
         f"-std={cxx_standard}",
         f"-resource-dir={clang_resource_dir}",
         # Place shim include dir early so it can intercept vendor headers.
-        *([f"-I{_get_shim_include_dir()}"] if _get_shim_include_dir() else []),
+        f"-I{_get_shim_include_dir()}",
         # cuda_wrappers_dir precede libstdc++ search includes to shadow certain
         # libstdc++ headers
         f"-isystem{cuda_wrappers_dir}",

ast_canopy/pyproject.toml

@@ -13,6 +13,9 @@ readme = { file = "README.md", content-type = "text/markdown" }
 [tool.scikit-build]
 cmake.targets = ["pylibastcanopy"]
 wheel.license-files = ["../LICENSE"]
+# Be explicit about shipping header shims in sdists so wheels built from sdist
+# (e.g., on PyPI) always contain these files.
+sdist.include = ["ast_canopy/shim_include/**"]
 
 [tool.scikit-build.metadata.version]
 provider = "scikit_build_core.metadata.setuptools_scm"

numbast/src/numbast/args.py

@@ -0,0 +1,31 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+from numba.cuda.target import CUDATargetContext
+from llvmlite import ir
+
+
+def prepare_ir_types(
+    context: CUDATargetContext, argtys: list[ir.Type]
+) -> list[ir.Type]:
+    """
+    Prepare IR types for passing arguments via pointers in function calls.
+
+    This utility wraps each argument type in a PointerType to enable
+    the call convention used by FunctionCallConv, where arguments are
+    passed by reference.
+
+    Parameters
+    ----------
+    context : context object
+        The compilation context providing the get_value_type method.
+    argtys : list[ir.Type]
+        List of LLVM IR types representing function arguments.
+
+    Returns
+    -------
+    list[ir.Type]
+        List of pointer types wrapping the value types of each argument.
+    """
+    return [ir.PointerType(context.get_value_type(argty)) for argty in argtys]

numbast/src/numbast/callconv.py

@@ -0,0 +1,88 @@
+from numbast.args import prepare_ir_types
+from numba.cuda import types, cgutils
+
+from llvmlite import ir
+
+
+class BaseCallConv:
+    shim_function_template = "{mangled_name}_nbst"
+
+    def __init__(
+        self,
+        itanium_mangled_name: str,
+        shim_writer: object,
+        shim_code: str,
+    ):
+        self.shim_writer = shim_writer
+        self.itanium_mangled_name = itanium_mangled_name
+        self.shim_code = shim_code
+
+        self.shim_function_name = self.shim_function_template.format(
+            mangled_name=self.itanium_mangled_name
+        )
+
+    def _lazy_write_shim(self, shim_code: str):
+        self.shim_writer.write_to_shim(shim_code, self.shim_function_name)
+
+    def _lower(self, builder, context, sig, args):
+        self._lazy_write_shim(self.shim_code)
+        return self._lower_impl(builder, context, sig, args)
+
+    def _lower_impl(self, builder, context, sig, args):
+        raise NotImplementedError
+
+    def __call__(self, builder, context, sig, args):
+        return self._lower(builder, context, sig, args)
+
+
+class FunctionCallConv(BaseCallConv):
+    def __init__(
+        self,
+        itanium_mangled_name: str,
+        shim_writer: object,
+        shim_code: str,
+        return_type: types.Type,
+    ):
+        super().__init__(itanium_mangled_name, shim_writer, shim_code)
+        self.return_type = return_type
+
+    def _lower_impl(self, builder, context, sig, args):
+        # 1. Prepare return value pointer
+        if self.return_type == types.void:
+            # Void return type in C++ is shimmed as int& ignored
+            retval_ty = ir.IntType(32)
+            retval_ptr = builder.alloca(retval_ty, name="ignored")
+        else:
+            retval_ty = context.get_value_type(self.return_type)
+            retval_ptr = builder.alloca(retval_ty, name="retval")
+
+        # 2. Prepare arguments
+        arg_pointer_types = prepare_ir_types(context, sig.args)
+
+        # All arguments are passed by pointer
+        ptrs = [
+            cgutils.alloca_once(builder, context.get_value_type(argty))
+            for argty in sig.args
+        ]
+        for ptr, argty, arg in zip(ptrs, sig.args, args):
+            builder.store(arg, ptr, align=getattr(argty, "alignof_", None))
+
+        # 3. Declare shim
+        # Shim signature: int (retval_type*, arg0_type*, ...)
+        fnty = ir.FunctionType(
+            ir.IntType(32), [ir.PointerType(retval_ty)] + arg_pointer_types
+        )
+        fn = cgutils.get_or_insert_function(
+            builder.module, fnty, self.shim_function_name
+        )
+
+        # 4. Call shim
+        builder.call(fn, (retval_ptr, *ptrs))
+
+        # 5. Return
+        if self.return_type == types.void:
+            return None
+        else:
+            return builder.load(
+                retval_ptr, align=getattr(self.return_type, "alignof_", None)
+            )

numbast/src/numbast/class_template.py

@@ -21,7 +21,6 @@
     CallableTemplate,
 )
 from numba.core.datamodel.models import StructModel, OpaqueModel
-from numba.cuda import declare_device
 from numba.cuda.cudadecl import register_global, register, register_attr
 from numba.cuda.cudaimpl import lower
 from numba.cuda.core.imputils import numba_typeref_ctor
@@ -44,25 +43,14 @@
 )
 from numbast.utils import (
     deduplicate_overloads,
-    make_device_caller_with_nargs,
-    assemble_arglist_string,
-    assemble_dereferenced_params_string,
+    make_struct_ctor_shim,
     make_struct_regular_method_shim,
 )
+from numbast.callconv import FunctionCallConv
 from numbast.shim_writer import ShimWriterBase
 
 ConcreteTypeCache: dict[str, nbtypes.Type] = {}
 
-ConcreteTypeCache2: dict[str, nbtypes.Type] = {}
-
-struct_ctor_shim_layer_template = """
-extern "C" __device__ int
-{func_name}(int &ignore, {name} *self {arglist}) {{
-    new (self) {name}({args});
-    return 0;
-}}
-"""
-
 
 class MetaType(nbtypes.Type):
     def __init__(self, template_name):
@@ -111,65 +99,27 @@ def bind_cxx_struct_ctor(
     # Lowering
     # Note that libclang always consider the return type of a constructor
     # is void. So we need to manually specify the return type here.
-    func_name = deduplicate_overloads(f"{ctor.mangled_name}_nbst")
-
-    # FIXME: temporary solution for mismatching function prototype against definition.
-    # If params are passed by value, at prototype the signature of __nv_bfloat16 is set
-    # to `b32` type, but to `b64` at definition, causing a linker error. A temporary solution
-    # is to pass all params by pointer and dereference them in shim. See dereferencing at the
-    # shim generation below.
-    ctor_shim_decl = declare_device(
-        func_name,
-        nbtypes.int32(
-            nbtypes.CPointer(s_type_ref.instance_type),
-            *map(nbtypes.CPointer, param_types),
-        ),
-    )
-
-    ctor_shim_call = make_device_caller_with_nargs(
-        func_name + "_shim",
-        1 + len(param_types),  # the extra argument for placement new pointer
-        ctor_shim_decl,
-    )
+    mangled_name = deduplicate_overloads(ctor.mangled_name)
+    shim_func_name = f"{mangled_name}_nbst"
 
     # Dynamically generate the shim layer:
     # FIXME: All params are passed by pointers, then dereferenced in shim.
     # temporary solution for mismatching function prototype against definition.
     # See above lowering for details.
-    arglist = assemble_arglist_string(ctor.params)
+    shim = make_struct_ctor_shim(
+        shim_name=shim_func_name, struct_name=struct_name, params=ctor.params
+    )
 
-    shim = struct_ctor_shim_layer_template.format(
-        func_name=func_name,
-        name=struct_name,
-        arglist=arglist,
-        args=assemble_dereferenced_params_string(ctor.params),
+    ctor_cc = FunctionCallConv(
+        mangled_name, shim_writer, shim, s_type_ref.instance_type
     )
 
     @lower(numba_typeref_ctor, s_type_ref, *param_types)
     def ctor_impl(context, builder, sig, args):
-        s_type = s_type_ref.instance_type
-        # Delay writing the shim function at lowering time. This avoids writing
-        # shim functions from the parsed header that's unused in kernels.
-        shim_writer.write_to_shim(shim, func_name)
-
-        selfptr = builder.alloca(context.get_value_type(s_type), name="selfptr")
-        argptrs = [
-            builder.alloca(context.get_value_type(arg)) for arg in sig.args[1:]
-        ]
-        for ptr, ty, arg in zip(argptrs, sig.args[1:], args[1:]):
-            builder.store(arg, ptr, align=getattr(ty, "alignof_", None))
-
-        context.compile_internal(
-            builder,
-            ctor_shim_call,
-            nb_signature(
-                nbtypes.int32,
-                nbtypes.CPointer(s_type),
-                *map(nbtypes.CPointer, param_types),
-            ),
-            (selfptr, *argptrs),
-        )
-        return builder.load(selfptr, align=getattr(s_type, "alignof_", None))
+        # `numba_typeref_ctor` includes the typeref as the first argument; the
+        # generated shim expects only the actual constructor params.
+        ctor_sig = nb_signature(s_type_ref.instance_type, *param_types)
+        return ctor_cc(builder, context, ctor_sig, args[1:])
 
     return param_types
 
@@ -236,44 +186,24 @@ def bind_cxx_struct_regular_method(
     )
 
     # Lowering
-    func_name = deduplicate_overloads(f"__{method_decl.mangled_name}_nbst")
-
-    c_sig = return_type(
-        nbtypes.CPointer(s_type), *map(nbtypes.CPointer, param_types)
-    )
-
-    shim_decl = declare_device(func_name, c_sig)
-
-    shim_call = make_device_caller_with_nargs(
-        func_name + "_shim", 1 + len(param_types), shim_decl
-    )
+    mangled_name = deduplicate_overloads(f"__{method_decl.mangled_name}")
+    shim_func_name = f"{mangled_name}_nbst"
 
     shim = make_struct_regular_method_shim(
-        shim_name=func_name,
+        shim_name=shim_func_name,
         struct_name=struct_decl.name,
         method_name=method_decl.name,
         return_type=method_decl.return_type.unqualified_non_ref_type_name,
         params=method_decl.params,
     )
 
+    method_cc = FunctionCallConv(mangled_name, shim_writer, shim, return_type)
+
     qualname = f"{s_type}.{method_decl.name}"
 
     @lower(qualname, s_type, *param_types)
     def _method_impl(context, builder, sig, args):
-        shim_writer.write_to_shim(shim, func_name)
-
-        argptrs = [
-            builder.alloca(context.get_value_type(arg)) for arg in sig.args
-        ]
-        for ptr, ty, arg in zip(argptrs, sig.args, args):
-            builder.store(arg, ptr, align=getattr(ty, "alignof_", None))
-
-        return context.compile_internal(
-            builder,
-            shim_call,
-            c_sig,
-            argptrs,
-        )
+        return method_cc(builder, context, sig, args)
 
     return nb_signature(return_type, *param_types, recvr=s_type)