Skip to content

[codex] Add scalar pointer out-return support #344

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
isVoid wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

[codex] Add scalar pointer out-return support #344

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a445de3
Add scalar pointer out-return support
isVoid May 12, 2026
136bd00
Apply pre-commit formatting
isVoid May 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 14 additions & 2 deletions docs/source/argument_intents.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,10 @@ C++ source of truth
__device__ bool stats_update_and_get_zscore(
RunningStats &state, float x, float &zscore_out);

__device__ uint4 texture_footprint(
cudaTextureObject_t texture, float x, float y,
unsigned int *single_mip_level);

Example config
--------------

Expand All @@ -38,6 +42,8 @@ Example config
stats_update_and_get_zscore:
state: inout_ptr
zscore_out: out_return
texture_footprint:
single_mip_level: out_return

Intent semantics
----------------
Expand Down Expand Up @@ -69,6 +75,8 @@ Intent semantics

- Parameter is removed from the visible Python call arguments.
- Numbast allocates temporary storage, passes it to C++, then returns the value to Python.
- For scalar pointer outputs such as ``unsigned int*``, the returned value is
the pointee scalar, not ``CPointer(T)``.
- If C++ also returns a non-``void`` value, generated return type is packed as a tuple.

Generated Python signatures
Expand All @@ -94,11 +102,15 @@ Representative signatures for the example API:
float32,
)

# out_return on a scalar pointer output:
signature(types.Tuple((uint32x4, uint32)), cudaTextureObject_t, float32, float32)

Notes
-----

- ``inout_ptr``, ``out_ptr``, and ``out_return`` are only supported on C++
reference parameters (``T&`` / ``T&&``).
- ``inout_ptr`` and ``out_ptr`` are only supported on C++ reference parameters
(``T&`` / ``T&&``). ``out_return`` also supports scalar pointer output
parameters (``T*``).
- In ``Function Argument Intents``, parameter overrides can be keyed by
parameter name or 0-based parameter index.

Expand Down
46 changes: 40 additions & 6 deletions numbast/src/numbast/callconv.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,30 @@
from numbast.intent import IntentPlan

# NBST:BEGIN_CALLCONV
from typing import NamedTuple

from numba.cuda import types, cgutils

from llvmlite import ir


class _OutReturnPtr(NamedTuple):
numba_ty: types.Type
ptr: ir.Value


def _get_out_return_ptr_mask(plan):
mask = getattr(plan, "out_return_ptr_mask", ())
if not mask:
return (False,) * len(plan.intents)
if len(mask) != len(plan.intents):
raise ValueError(
"IntentPlan out_return_ptr_mask length does not match intents: "
f"{len(mask)} != {len(plan.intents)}"
)
return tuple(bool(v) for v in mask)


class BaseCallConv:
shim_function_template = "{mangled_name}_nbst"

Expand Down Expand Up @@ -146,7 +165,7 @@ def _lower_impl(self, builder, context, sig, args):
# - default: pass pointer-to-value to shim (alloca + store)
# - for C++ reference args mapped to CPointer(T): pass pointer value directly
ptrs = []
out_return_ptrs: list[tuple[types.Type, ir.Value]] = []
out_return_ptrs: list[_OutReturnPtr] = []
if self._intent_plan is None:
for argty, arg, passthrough in zip(sig.args, args, pass_ptr_mask):
vty = context.get_value_type(argty)
Expand All @@ -168,16 +187,28 @@ def _lower_impl(self, builder, context, sig, args):
orig_to_out = [None] * n_orig
for out_pos, orig_idx in enumerate(plan.out_return_indices):
orig_to_out[orig_idx] = out_pos
out_return_ptr_mask = _get_out_return_ptr_mask(plan)

for orig_idx in range(n_orig):
out_pos = orig_to_out[orig_idx]
if out_pos is not None:
out_nbty = self._out_return_types[out_pos]
vty = context.get_value_type(out_nbty)
ptr = cgutils.alloca_once(builder, vty)
ptrs.append(ptr)
arg_pointer_types.append(ir.PointerType(vty))
out_return_ptrs.append((out_nbty, ptr))
if out_return_ptr_mask[orig_idx]:
ptr_slot_ty = ir.PointerType(vty)
ptr_slot = cgutils.alloca_once(
builder, ptr_slot_ty, name="out_return_ptr"
)
builder.store(ptr, ptr_slot)
ptrs.append(ptr_slot)
arg_pointer_types.append(ir.PointerType(ptr_slot_ty))
else:
ptrs.append(ptr)
arg_pointer_types.append(ir.PointerType(vty))
out_return_ptrs.append(
_OutReturnPtr(numba_ty=out_nbty, ptr=ptr)
)
continue

vis_pos = orig_to_vis[orig_idx]
Expand Down Expand Up @@ -231,9 +262,12 @@ def _lower_impl(self, builder, context, sig, args):
retval_ptr, align=getattr(cxx_return_type, "alignof_", None)
)
)
for out_ty, out_ptr in out_return_ptrs:
for out_return in out_return_ptrs:
ret_vals.append(
builder.load(out_ptr, align=getattr(out_ty, "alignof_", None))
builder.load(
out_return.ptr,
align=getattr(out_return.numba_ty, "alignof_", None),
)
)

# If Numba-visible return is a tuple, use context.make_tuple.
Expand Down
85 changes: 28 additions & 57 deletions numbast/src/numbast/class_template.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,14 @@
to_c_type_str,
to_numba_arg_type,
)
from numbast.intent import ArgIntent, IntentPlan, compute_intent_plan
from numbast.intent import ArgIntent, compute_intent_plan
from numbast.intent_utils import (
compose_return_type,
get_out_return_ptr_mask,
out_return_types_for_plan,
prepend_receiver_to_intent_plan,
shim_arg_type_for_out_return,
)
from numbast.utils import (
deduplicate_overloads,
make_struct_ctor_shim,
Expand Down Expand Up @@ -334,15 +341,7 @@ def bind_cxx_struct_regular_method(
overrides=overrides,
allow_out_return=True,
)
intent_plan = IntentPlan(
intents=(ArgIntent.in_,) + method_plan.intents,
visible_param_indices=(0,)
+ tuple(i + 1 for i in method_plan.visible_param_indices),
out_return_indices=tuple(
i + 1 for i in method_plan.out_return_indices
),
pass_ptr_mask=(False,) + method_plan.pass_ptr_mask,
)
intent_plan = prepend_receiver_to_intent_plan(method_plan)

param_types = []
for orig_idx in method_plan.visible_param_indices:
Expand All @@ -354,24 +353,10 @@ def bind_cxx_struct_regular_method(
else:
param_types.append(base)

out_return_types = [
to_numba_type(
method_decl.param_types[i].unqualified_non_ref_type_name
)
for i in method_plan.out_return_indices
]
if out_return_types:
if cxx_return_type == nbtypes.void:
if len(out_return_types) == 1:
return_type = out_return_types[0]
else:
return_type = nbtypes.Tuple(tuple(out_return_types))
else:
return_type = nbtypes.Tuple(
tuple([cxx_return_type, *out_return_types])
)
else:
return_type = cxx_return_type
out_return_types = out_return_types_for_plan(
method_decl.param_types, method_plan
)
return_type = compose_return_type(cxx_return_type, out_return_types)
arg_is_ref = None

# Lowering
Expand Down Expand Up @@ -519,35 +504,13 @@ def generic(
overrides=overrides,
allow_out_return=True,
)
intent_plan = IntentPlan(
intents=(ArgIntent.in_,) + method_plan.intents,
visible_param_indices=(0,)
+ tuple(i + 1 for i in method_plan.visible_param_indices),
out_return_indices=tuple(
i + 1 for i in method_plan.out_return_indices
),
pass_ptr_mask=(False,) + method_plan.pass_ptr_mask,
intent_plan = prepend_receiver_to_intent_plan(method_plan)
out_return_types = out_return_types_for_plan(
templated_method.function.param_types, method_plan
)
return_type = compose_return_type(
cxx_return_type, out_return_types
)
out_return_types = [
to_numba_type(
templated_method.function.param_types[
i
].unqualified_non_ref_type_name
)
for i in method_plan.out_return_indices
]
if out_return_types:
if cxx_return_type == nbtypes.void:
if len(out_return_types) == 1:
return_type = out_return_types[0]
else:
return_type = nbtypes.Tuple(tuple(out_return_types))
else:
return_type = nbtypes.Tuple(
tuple([cxx_return_type, *out_return_types])
)
else:
return_type = cxx_return_type

lowering_key = (qualname, recvr, param_types)
if lowering_key not in _TEMPLATED_METHOD_LOWERING_CACHE:
Expand Down Expand Up @@ -592,6 +555,9 @@ def _impl(
method_plan.out_return_indices
)
}
out_return_ptr_mask = get_out_return_ptr_mask(
method_plan
)
# Reconstruct full C++ param order by merging visible
# params with out_return slots, keeping a shim-aligned
# pass_ptr_mask.
Expand All @@ -602,7 +568,12 @@ def _impl(
out_pos = out_return_map.get(orig_idx)
if out_pos is not None:
param_types_for_shim_list.append(
out_return_types[out_pos]
shim_arg_type_for_out_return(
out_return_types[out_pos],
pointer_out=out_return_ptr_mask[
orig_idx
],
)
)
pass_ptr_mask_for_shim_list.append(False)
else:
Expand Down
24 changes: 5 additions & 19 deletions numbast/src/numbast/function.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

from numbast.types import to_numba_type, to_numba_arg_type
from numbast.intent import compute_intent_plan
from numbast.intent_utils import compose_return_type, out_return_types_for_plan
from numbast.utils import (
deduplicate_overloads,
make_function_shim,
Expand Down Expand Up @@ -196,25 +197,10 @@ def bind_cxx_non_operator_function(
else:
param_types.append(base)

out_return_types = [
to_numba_type(
func_decl.param_types[i].unqualified_non_ref_type_name
)
for i in intent_plan.out_return_indices
]

if out_return_types:
if cxx_return_type == nbtypes.void:
if len(out_return_types) == 1:
return_type = out_return_types[0]
else:
return_type = nbtypes.Tuple(tuple(out_return_types))
else:
return_type = nbtypes.Tuple(
tuple([cxx_return_type, *out_return_types])
)
else:
return_type = cxx_return_type
out_return_types = out_return_types_for_plan(
func_decl.param_types, intent_plan
)
return_type = compose_return_type(cxx_return_type, out_return_types)

# In intentful mode, pass-through pointers are controlled by intent_plan,
# not by whether the C++ parameter is a reference.
Expand Down
38 changes: 17 additions & 21 deletions numbast/src/numbast/function_template.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,12 @@
from numbast.callconv import FunctionCallConv
from numbast.deduction import deduce_templated_overloads
from numbast.intent import ArgIntent, compute_intent_plan
from numbast.intent_utils import (
compose_return_type,
get_out_return_ptr_mask,
out_return_types_for_plan,
shim_arg_type_for_out_return,
)
from numbast.types import to_c_type_str, to_numba_type
from numbast.utils import deduplicate_overloads, get_return_type_strings
from numbast.shim_writer import ShimWriterBase
Expand Down Expand Up @@ -265,26 +271,12 @@ def generic(self, args, kwds, overloads=overloads, overrides=overrides):
allow_out_return=True,
)
intent_plan = func_plan
out_return_types = [
to_numba_type(
templated_func.function.param_types[
i
].unqualified_non_ref_type_name
)
for i in func_plan.out_return_indices
]
if out_return_types:
if cxx_return_type == nbtypes.void:
if len(out_return_types) == 1:
return_type = out_return_types[0]
else:
return_type = nbtypes.Tuple(tuple(out_return_types))
else:
return_type = nbtypes.Tuple(
tuple([cxx_return_type, *out_return_types])
)
else:
return_type = cxx_return_type
out_return_types = out_return_types_for_plan(
templated_func.function.param_types, func_plan
)
return_type = compose_return_type(
cxx_return_type, out_return_types
)

@lower(func, *param_types)
def _impl(
Expand Down Expand Up @@ -321,6 +313,7 @@ def _impl(
func_plan.out_return_indices
)
}
out_return_ptr_mask = get_out_return_ptr_mask(func_plan)
visible_iter = iter(param_types_inner)
visible_mask_iter = iter(func_plan.pass_ptr_mask)
param_types_for_shim_list = []
Expand All @@ -329,7 +322,10 @@ def _impl(
out_pos = out_return_map.get(orig_idx)
if out_pos is not None:
param_types_for_shim_list.append(
out_return_types[out_pos]
shim_arg_type_for_out_return(
out_return_types[out_pos],
pointer_out=out_return_ptr_mask[orig_idx],
)
)
pass_ptr_mask_for_shim_list.append(False)
else:
Expand Down
Loading
Loading