Capture Enum Underlying Type in AST_CANOPY, documenting the C++/Python Enum lowering (#267)

This PR adds feature to capture the c++ enum underlying integral type in
ast_canopy. As well as capturing the handling of binding to function
with enum argument in docs.

This PR also tests whether `cuda.bindings.runtime.cudaRoundMode` works
well with existing bindings.

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

* **New Features**
* Enum underlying integer types are exposed to Python bindings (includes
cudaRoundMode); enum handling and registration improved.

* **Refactor**
* Call/return handling internals simplified (no user-visible API
change); several type-resolution flows consolidated.
* Public typing and metadata structures extended to include enum
underlying-type information.

* **Tests**
* Added and updated tests exercising enum underlying types, device-side
enum helpers, and streamlined binding-generation fixtures.

* **Documentation**
  * FAQ expanded with notes on enum binding and lowering behavior.

<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/ast_canopy/pylibastcanopy.cpp

@@ -49,16 +49,24 @@ PYBIND11_MODULE(pylibastcanopy, m) {
       .def_readwrite("name", &Enum::name)
       .def_readwrite("enumerators", &Enum::enumerators)
       .def_readwrite("enumerator_values", &Enum::enumerator_values)
+      .def_readwrite("underlying_type", &Enum::underlying_type)
       .def(py::pickle(
           [](const Enum &e) {
-            return py::make_tuple(e.name, e.enumerators, e.enumerator_values);
+            return py::make_tuple(e.name, e.enumerators, e.enumerator_values,
+                                  e.underlying_type);
           },
           [](py::tuple t) {
-            if (t.size() != 3)
+            // Backward compat: older pickles only stored (name, enumerators,
+            // enumerator_values).
+            if (t.size() != 3 && t.size() != 4)
               throw std::runtime_error("Invalid enum state during unpickle!");
-            return Enum{t[0].cast<std::string>(),
-                        t[1].cast<std::vector<std::string>>(),
-                        t[2].cast<std::vector<std::string>>()};
+            Enum e{t[0].cast<std::string>(),
+                   t[1].cast<std::vector<std::string>>(),
+                   t[2].cast<std::vector<std::string>>()};
+            if (t.size() == 4) {
+              e.underlying_type = t[3].cast<Type>();
+            }
+            return e;
           }));
 
   py::class_<Type>(m, "Type")

ast_canopy/ast_canopy/pylibastcanopy.pyi

@@ -1,3 +1,5 @@
+import collections.abc
+import typing
 from _typeshed import Incomplete
 from typing import ClassVar, overload
 
@@ -6,13 +8,19 @@ class ClassTemplate(Template):
     record: Incomplete
     def __init__(self, *args, **kwargs) -> None: ...
 
+class ClassTemplateSpecialization(Record):
+    actual_template_arguments: list[str]
+    class_template: ClassTemplate
+    def __init__(self, *args, **kwargs) -> None: ...
+
 class ConstExprVar:
     name: str
     type_: Type
     value: str
     def __init__(self) -> None: ...
 
 class Declarations:
+    class_template_specializations: list[ClassTemplateSpecialization]
     class_templates: list[ClassTemplate]
     enums: list[Enum]
     function_templates: list[FunctionTemplate]
@@ -25,6 +33,7 @@ class Enum:
     enumerator_values: list[str]
     enumerators: list[str]
     name: str
+    underlying_type: Incomplete
     def __init__(self, arg0) -> None: ...
 
 class Field:
@@ -34,6 +43,7 @@ class Field:
     def __init__(self, *args, **kwargs) -> None: ...
 
 class Function:
+    attributes: set[str]
     exec_space: execution_space
     is_constexpr: bool
     mangled_name: str
@@ -73,7 +83,11 @@ class Record:
 class Template:
     num_min_required_args: int
     template_parameters: list[TemplateParam]
-    def __init__(self, arg0: list[TemplateParam], arg1: int) -> None: ...
+    def __init__(
+        self,
+        arg0: collections.abc.Sequence[TemplateParam],
+        arg1: typing.SupportsInt,
+    ) -> None: ...
 
 class TemplateParam:
     kind: template_param_kind
@@ -104,7 +118,7 @@ class access_kind:
     private_: ClassVar[access_kind] = ...
     protected_: ClassVar[access_kind] = ...
     public_: ClassVar[access_kind] = ...
-    def __init__(self, value: int) -> None: ...
+    def __init__(self, value: typing.SupportsInt) -> None: ...
     def __eq__(self, other: object) -> bool: ...
     def __hash__(self) -> int: ...
     def __index__(self) -> int: ...
@@ -123,7 +137,7 @@ class execution_space:
     host: ClassVar[execution_space] = ...
     host_device: ClassVar[execution_space] = ...
     undefined: ClassVar[execution_space] = ...
-    def __init__(self, value: int) -> None: ...
+    def __init__(self, value: typing.SupportsInt) -> None: ...
     def __eq__(self, other: object) -> bool: ...
     def __hash__(self) -> int: ...
     def __index__(self) -> int: ...
@@ -145,7 +159,7 @@ class method_kind:
     move_constructor: ClassVar[method_kind] = ...
     other: ClassVar[method_kind] = ...
     other_constructor: ClassVar[method_kind] = ...
-    def __init__(self, value: int) -> None: ...
+    def __init__(self, value: typing.SupportsInt) -> None: ...
     def __eq__(self, other: object) -> bool: ...
     def __hash__(self) -> int: ...
     def __index__(self) -> int: ...
@@ -162,7 +176,7 @@ class template_param_kind:
     non_type: ClassVar[template_param_kind] = ...
     template_: ClassVar[template_param_kind] = ...
     type_: ClassVar[template_param_kind] = ...
-    def __init__(self, value: int) -> None: ...
+    def __init__(self, value: typing.SupportsInt) -> None: ...
     def __eq__(self, other: object) -> bool: ...
     def __hash__(self) -> int: ...
     def __index__(self) -> int: ...
@@ -174,8 +188,10 @@ class template_param_kind:
     def value(self) -> int: ...
 
 def parse_declarations_from_command_line(
-    arg0: list[str], arg1: list[str]
+    arg0: collections.abc.Sequence[str],
+    arg1: collections.abc.Sequence[str],
+    arg2: bool,
 ) -> Declarations: ...
 def value_from_constexpr_vardecl(
-    arg0: list[str], arg1: str
+    arg0: collections.abc.Sequence[str], arg1: str
 ) -> ConstExprVar | None: ...

ast_canopy/cpp/include/ast_canopy/ast_canopy.hpp

@@ -40,18 +40,6 @@ enum class template_param_kind { type, non_type, template_ };
 
 enum class access_kind { public_, protected_, private_ };
 
-struct Enum {
-  Enum(const std::string &name, const std::vector<std::string> &enumerators,
-       const std::vector<std::string> &enumerator_values)
-      : name(name), enumerators(enumerators),
-        enumerator_values(enumerator_values) {}
-  Enum(const clang::EnumDecl *);
-
-  std::string name;
-  std::vector<std::string> enumerators;
-  std::vector<std::string> enumerator_values;
-};
-
 struct Type {
   Type() = default;
   Type(std::string name, std::string unqualified_non_ref_type_name,
@@ -69,6 +57,19 @@ struct Type {
   bool _is_left_reference;
 };
 
+struct Enum {
+  Enum(const std::string &name, const std::vector<std::string> &enumerators,
+       const std::vector<std::string> &enumerator_values)
+      : name(name), enumerators(enumerators),
+        enumerator_values(enumerator_values) {}
+  Enum(const clang::EnumDecl *);
+
+  std::string name;
+  std::vector<std::string> enumerators;
+  std::vector<std::string> enumerator_values;
+  Type underlying_type;
+};
+
 struct ConstExprVar {
   ConstExprVar() = default;
   ConstExprVar(const clang::VarDecl *VD);

ast_canopy/cpp/src/enum.cpp

@@ -7,7 +7,9 @@
 
 namespace ast_canopy {
 
-Enum::Enum(const clang::EnumDecl *ED) : name(ED->getNameAsString()) {
+Enum::Enum(const clang::EnumDecl *ED)
+    : name(ED->getNameAsString()),
+      underlying_type(ED->getIntegerType(), ED->getASTContext()) {
   for (const auto *enumerator : ED->enumerators()) {
     enumerators.push_back(enumerator->getNameAsString());
 

ast_canopy/tests/data/sample_enum.cu

@@ -7,3 +7,7 @@ enum Foo { A = 1, B = 2, C = 3 };
 enum NoDefaultFoo { D, E, F };
 
 enum class Bar { A = 1 };
+
+enum Fruit : uint64_t { Apple, Banana, Orange };
+enum Car : int16_t { Sedan, SUV, Pickup, Hatchback };
+enum Plane : char { Boeing, Airbus, Embraer };

ast_canopy/tests/test_parse_from_source.py

@@ -467,7 +467,7 @@ def test_load_enum(sample_enum_source, test_pickle):
         pickled = [pickle.dumps(e) for e in enums]
         enums = [pickle.loads(p) for p in pickled]
 
-    assert len(enums) == 3
+    assert len(enums) == 6
     assert enums[0].name == "Foo"
     assert len(enums[0].enumerators) == 3
     assert enums[0].enumerators[0] == "A"
@@ -495,6 +495,24 @@ def test_load_enum(sample_enum_source, test_pickle):
     assert enums[2].enumerator_values[0] == "1"
 
 
+def test_load_enum_underlying_type(sample_enum_source, test_pickle):
+    decls = parse_declarations_from_source(
+        sample_enum_source, [sample_enum_source], "sm_80"
+    )
+
+    enums = decls.enums
+    if test_pickle:
+        pickled = [pickle.dumps(e) for e in enums]
+        enums = [pickle.loads(p) for p in pickled]
+
+    assert enums[3].name == "Fruit"
+    assert enums[3].underlying_type.name == "uint64_t"
+    assert enums[4].name == "Car"
+    assert enums[4].underlying_type.name == "int16_t"
+    assert enums[5].name == "Plane"
+    assert enums[5].underlying_type.name == "char"
+
+
 def test_load_struct_function_execution_space(
     sample_execution_space_source, test_pickle
 ):

docs/source/faq.rst

@@ -167,3 +167,26 @@ bindings generated with a specific version of Numbast are tested against a speci
 .. note::
    These version restrictions may be relaxed or removed once ``numba-cuda`` releases a stable 1.0 version with
    stabilized public APIs. Until then, bindings are tested against specific version ranges to ensure compatibility.
+
+
+C++ Enum Binding Generation Notes
+---------------------------------
+
+**Why do Numbast bindings treat C++ enums as ``int64`` in Numba?**
+
+Numba represents Python ``IntEnum`` values using ``IntEnumMember(..., int64)`` (i.e., enum values are lowered as
+64-bit integers). Numbast follows this convention in both dynamic and static binding generation so that Python-side
+typing and lowering are consistent.
+
+**But C++ enums can have different underlying integer types. Why don't we track and truncate to that type in lowering?**
+
+Numbast does not keep a per-enum “underlying integer type” registry and does not perform explicit truncation during
+lowering because the device-side shim is compiled by NVRTC, and the shim call site is where C++ type checking happens.
+Even though the Python/Numba side lowers enum values as 64-bit integers, NVRTC can resolve the target enum type and emit
+the appropriate conversion when the shim calls the original function that takes the C++ enum parameter. This means we
+don't need to track per-enum underlying integer types in Python or add special-case truncation/casting logic in Numbast
+lowering.
+
+If you are binding code that depends on unusual enum representations or non-standard ABIs, you may need a custom
+adapter. For typical CUDA device code, this approach keeps the implementation simpler and avoids maintaining extra
+metadata for every enum type.

numbast/src/numbast/args.py

@@ -1,8 +1,8 @@
 # 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
 
 

numbast/src/numbast/callconv.py

@@ -36,24 +36,15 @@ def __call__(self, 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):
+        return_type = sig.return_type
         # 1. Prepare return value pointer
-        if self.return_type == types.void:
+        if 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_ty = context.get_value_type(return_type)
             retval_ptr = builder.alloca(retval_ty, name="retval")
 
         # 2. Prepare arguments
@@ -80,9 +71,9 @@ def _lower_impl(self, builder, context, sig, args):
         builder.call(fn, (retval_ptr, *ptrs))
 
         # 5. Return
-        if self.return_type == types.void:
+        if return_type == types.void:
             return None
         else:
             return builder.load(
-                retval_ptr, align=getattr(self.return_type, "alignof_", None)
+                retval_ptr, align=getattr(return_type, "alignof_", None)
             )

numbast/src/numbast/class_template.py

@@ -110,9 +110,7 @@ def bind_cxx_struct_ctor(
         shim_name=shim_func_name, struct_name=struct_name, params=ctor.params
     )
 
-    ctor_cc = FunctionCallConv(
-        mangled_name, shim_writer, shim, s_type_ref.instance_type
-    )
+    ctor_cc = FunctionCallConv(mangled_name, shim_writer, shim)
 
     @lower(numba_typeref_ctor, s_type_ref, *param_types)
     def ctor_impl(context, builder, sig, args):
@@ -197,7 +195,7 @@ def bind_cxx_struct_regular_method(
         params=method_decl.params,
     )
 
-    method_cc = FunctionCallConv(mangled_name, shim_writer, shim, return_type)
+    method_cc = FunctionCallConv(mangled_name, shim_writer, shim)
 
     qualname = f"{s_type}.{method_decl.name}"