From a445de3404e3eb9598708fb107eac22b07734092 Mon Sep 17 00:00:00 2001 From: Michael Wang Date: Tue, 12 May 2026 15:28:37 -0700 Subject: [PATCH 1/2] Add scalar pointer out-return support --- docs/source/argument_intents.rst | 16 ++- numbast/src/numbast/callconv.py | 46 ++++++- numbast/src/numbast/class_template.py | 85 +++++-------- numbast/src/numbast/function.py | 24 +--- numbast/src/numbast/function_template.py | 38 +++--- numbast/src/numbast/intent.py | 51 +++++++- numbast/src/numbast/intent_defs.py | 6 +- numbast/src/numbast/intent_utils.py | 55 +++++++++ numbast/src/numbast/static/function.py | 109 +++++++++-------- numbast/src/numbast/static/struct.py | 114 +++++++++++------- .../static/tests/data/function_out.cuh | 2 + .../static/tests/data/src/function_out.cu | 22 ++++ .../tests/test_function_static_bindings.py | 41 ++++++- numbast/src/numbast/struct.py | 39 ++---- numbast/tests/data/sample_function_out.cuh | 9 ++ numbast/tests/test_function.py | 18 ++- numbast/tests/test_intent.py | 48 ++++++++ 17 files changed, 490 insertions(+), 233 deletions(-) create mode 100644 numbast/src/numbast/intent_utils.py create mode 100644 numbast/tests/test_intent.py diff --git a/docs/source/argument_intents.rst b/docs/source/argument_intents.rst index 6fa8ce4e..f2919393 100644 --- a/docs/source/argument_intents.rst +++ b/docs/source/argument_intents.rst @@ -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 -------------- @@ -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 ---------------- @@ -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 @@ -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. diff --git a/numbast/src/numbast/callconv.py b/numbast/src/numbast/callconv.py index 5e10eadd..fd14a3cf 100644 --- a/numbast/src/numbast/callconv.py +++ b/numbast/src/numbast/callconv.py @@ -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" @@ -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) @@ -168,6 +187,7 @@ 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] @@ -175,9 +195,20 @@ def _lower_impl(self, builder, context, sig, args): 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] @@ -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. diff --git a/numbast/src/numbast/class_template.py b/numbast/src/numbast/class_template.py index abd1dd1e..50134840 100644 --- a/numbast/src/numbast/class_template.py +++ b/numbast/src/numbast/class_template.py @@ -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, @@ -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: @@ -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 @@ -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: @@ -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. @@ -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: diff --git a/numbast/src/numbast/function.py b/numbast/src/numbast/function.py index 23049b14..7a460dcf 100644 --- a/numbast/src/numbast/function.py +++ b/numbast/src/numbast/function.py @@ -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, @@ -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. diff --git a/numbast/src/numbast/function_template.py b/numbast/src/numbast/function_template.py index 945aa099..e2712b4b 100644 --- a/numbast/src/numbast/function_template.py +++ b/numbast/src/numbast/function_template.py @@ -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 @@ -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( @@ -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 = [] @@ -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: diff --git a/numbast/src/numbast/intent.py b/numbast/src/numbast/intent.py index c3db6bf2..c4130fd0 100644 --- a/numbast/src/numbast/intent.py +++ b/numbast/src/numbast/intent.py @@ -60,6 +60,39 @@ def _is_ref_type(ast_type: Any) -> bool: return bool(is_ref) +def _type_name(ast_type: Any) -> str: + return str(getattr(ast_type, "unqualified_non_ref_type_name", "")).strip() + + +def _is_pointer_type(ast_type: Any) -> bool: + return _type_name(ast_type).endswith("*") + + +def pointee_type_name(type_name: str) -> str: + """ + Return the type string produced by dereferencing one C/C++ pointer layer. + """ + normalized = type_name.strip() + if not normalized.endswith("*"): + raise ValueError(f"expected pointer type, got {type_name!r}") + return normalized[:-1].rstrip() + + +def get_out_return_ptr_mask(plan: IntentPlan) -> tuple[bool, ...]: + """ + Return the per-original-parameter mask for scalar pointer out-returns. + """ + 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) + + def compute_intent_plan( *, params: list[Any], @@ -70,7 +103,7 @@ def compute_intent_plan( """ Compute a per-parameter intent plan for a function call. - This determines an ArgIntent for each parameter (defaulting to `ArgIntent.in_`), applies optional overrides (index-based and name-based; name-based overrides take precedence), validates intents against parameter types (non-`in_` intents require reference-like types), and produces an IntentPlan describing which parameters are visible, which are returned via out_return, and which should be passed as pointer-like arguments. + This determines an ArgIntent for each parameter (defaulting to `ArgIntent.in_`), applies optional overrides (index-based and name-based; name-based overrides take precedence), validates intents against parameter types (non-`in_` intents require reference-like types, except `out_return` also supports scalar pointer output parameters), and produces an IntentPlan describing which parameters are visible, which are returned via out_return, and which should be passed as pointer-like arguments. Parameters: params (list[Any]): Parameter-like objects (must have a `.name` attribute when name-based overrides are used). @@ -84,9 +117,10 @@ def compute_intent_plan( - visible_param_indices: tuple[int] indices of parameters that remain visible (not out_return) - out_return_indices: tuple[int] indices of parameters specified as `out_return` - pass_ptr_mask: tuple[bool] parallel to visible_param_indices indicating whether the parameter should be passed as a pointer (True for `inout_ptr`/`out_ptr`) + - out_return_ptr_mask: tuple[bool] parallel to original parameters indicating which out_return parameters are scalar pointer outputs Raises: - ValueError: If `params` and `param_types` lengths differ, an index override is out of range, a named override refers to an unknown parameter, a non-`in_` intent is applied to a non-reference type, or `out_return` is disallowed. + ValueError: If `params` and `param_types` lengths differ, an index override is out of range, a named override refers to an unknown parameter, a non-`in_` intent is applied to an unsupported type, or `out_return` is disallowed. TypeError: If override keys are not `int` or `str`, or override values are not `str` or `ArgIntent`. """ if len(params) != len(param_types): @@ -139,11 +173,21 @@ def compute_intent_plan( visible_param_indices: list[int] = [] out_return_indices: list[int] = [] pass_ptr_mask: list[bool] = [] + out_return_ptr_mask: list[bool] = [False] * len(params) for i, (intent, ty) in enumerate(zip(normalized, param_types)): is_ref = _is_ref_type(ty) + is_pointer = _is_pointer_type(ty) if intent != ArgIntent.in_: - if not is_ref: + if intent == ArgIntent.out_return: + if not (is_ref or is_pointer): + raise ValueError( + f"arg_intent[{i}]='out_return' is only supported for " + "reference parameters (T&/T&&) or pointer output " + "parameters (T*)" + ) + out_return_ptr_mask[i] = is_pointer and not is_ref + elif not is_ref: raise ValueError( f"arg_intent[{i}]={intent.value!r} is only supported for reference parameters (T&/T&&)" ) @@ -164,4 +208,5 @@ def compute_intent_plan( visible_param_indices=tuple(visible_param_indices), out_return_indices=tuple(out_return_indices), pass_ptr_mask=tuple(pass_ptr_mask), + out_return_ptr_mask=tuple(out_return_ptr_mask), ) diff --git a/numbast/src/numbast/intent_defs.py b/numbast/src/numbast/intent_defs.py index e68ef661..06b73fb8 100644 --- a/numbast/src/numbast/intent_defs.py +++ b/numbast/src/numbast/intent_defs.py @@ -18,8 +18,9 @@ class ArgIntent(str, Enum): as a value type on the Numba side. - `inout_ptr` / `out_ptr`: C++ reference parameter is exposed as a pointer (CPointer(T)) on the Numba side and passed through to the shim. - - `out_return`: C++ reference parameter is *not* exposed as an argument; a - temporary is allocated, passed to the shim, and then returned to the caller. + - `out_return`: C++ reference parameter or scalar pointer output parameter + is *not* exposed as an argument; a temporary is allocated, passed to the + shim, and then returned to the caller. """ in_ = "in" @@ -38,6 +39,7 @@ class IntentPlan: visible_param_indices: tuple[int, ...] # subset of [0..N) out_return_indices: tuple[int, ...] # subset of [0..N) pass_ptr_mask: tuple[bool, ...] # aligned with visible params only + out_return_ptr_mask: tuple[bool, ...] = () # aligned with original params # NBST:END_INTENT_DEFS diff --git a/numbast/src/numbast/intent_utils.py b/numbast/src/numbast/intent_utils.py new file mode 100644 index 00000000..27a5e06f --- /dev/null +++ b/numbast/src/numbast/intent_utils.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from numba import types as nbtypes + +from numbast.intent import get_out_return_ptr_mask, pointee_type_name +from numbast.intent_defs import ArgIntent, IntentPlan +from numbast.types import to_numba_type + + +def out_return_type_for_param(param_type, *, pointer_out: bool): + """ + Return the Numba-visible out-return type for one C++ parameter. + """ + type_name = param_type.unqualified_non_ref_type_name + if pointer_out: + type_name = pointee_type_name(type_name) + return to_numba_type(type_name) + + +def out_return_types_for_plan(param_types, plan: IntentPlan): + ptr_mask = get_out_return_ptr_mask(plan) + return [ + out_return_type_for_param(param_types[i], pointer_out=ptr_mask[i]) + for i in plan.out_return_indices + ] + + +def compose_return_type(cxx_return_type, out_return_types): + if not out_return_types: + return cxx_return_type + if cxx_return_type == nbtypes.void: + if len(out_return_types) == 1: + return out_return_types[0] + return nbtypes.Tuple(tuple(out_return_types)) + return nbtypes.Tuple(tuple([cxx_return_type, *out_return_types])) + + +def prepend_receiver_to_intent_plan(method_plan: IntentPlan) -> IntentPlan: + return 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, + out_return_ptr_mask=(False,) + get_out_return_ptr_mask(method_plan), + ) + + +def shim_arg_type_for_out_return(out_return_type, *, pointer_out: bool): + if pointer_out: + return nbtypes.CPointer(out_return_type) + return out_return_type diff --git a/numbast/src/numbast/static/function.py b/numbast/src/numbast/static/function.py index 3627bb99..26e96cbc 100644 --- a/numbast/src/numbast/static/function.py +++ b/numbast/src/numbast/static/function.py @@ -17,7 +17,12 @@ get_callconv_utils, ) from numbast.static.types import to_numba_type_str, to_numba_arg_type_str -from numbast.intent import ArgIntent, compute_intent_plan +from numbast.intent import ( + ArgIntent, + compute_intent_plan, + get_out_return_ptr_mask, + pointee_type_name, +) from numbast.utils import make_function_shim, _apply_prefix_removal from numbast.errors import TypeNotFoundError, MangledFunctionNameConflictError @@ -37,6 +42,54 @@ """A set of created function API names.""" +def _tuple_literal(items: list[str]) -> str: + """ + Builds a Python tuple literal from a list of string expressions. + """ + if not items: + return "()" + if len(items) == 1: + return f"({items[0]},)" + return f"({', '.join(items)})" + + +def _out_return_type_str(param_type, *, pointer_out: bool) -> str: + type_name = param_type.unqualified_non_ref_type_name + if pointer_out: + type_name = pointee_type_name(type_name) + return to_numba_type_str(type_name) + + +def _compose_return_type_str(cxx_return_type_str: str, out_return_types: list[str]): + if not out_return_types: + return cxx_return_type_str + if cxx_return_type_str == "void": + if len(out_return_types) == 1: + return out_return_types[0] + outs = ", ".join(out_return_types) + return f"types.Tuple(({outs},))" + outs = ", ".join([cxx_return_type_str, *out_return_types]) + return f"types.Tuple(({outs},))" + + +def _render_intent_plan(plan) -> str: + intents_str = _tuple_literal( + [ + f"ArgIntent.{i.value if i != ArgIntent.in_ else 'in_'}" + for i in plan.intents + ] + ) + return ( + "IntentPlan(" + f"intents={intents_str}, " + f"visible_param_indices={repr(plan.visible_param_indices)}, " + f"out_return_indices={repr(plan.out_return_indices)}, " + f"pass_ptr_mask={repr(plan.pass_ptr_mask)}, " + f"out_return_ptr_mask={repr(get_out_return_ptr_mask(plan))}" + ")" + ) + + def _matches_any_regex_pattern(name: str, patterns: list[str]) -> bool: """Check if a function name matches any of the provided regex patterns. @@ -219,61 +272,23 @@ def __init__( self._argument_numba_types ) + out_return_ptr_mask = get_out_return_ptr_mask(plan) out_return_types = [ - to_numba_type_str( - self._decl.param_types[i].unqualified_non_ref_type_name + _out_return_type_str( + self._decl.param_types[i], + pointer_out=out_return_ptr_mask[i], ) for i in plan.out_return_indices ] if out_return_types: self.Imports.add("from numba import types") - if self._cxx_return_type_str == "void": - if len(out_return_types) == 1: - self._return_numba_type_str = out_return_types[0] - else: - outs = ", ".join(out_return_types) - self._return_numba_type_str = f"types.Tuple(({outs},))" - else: - outs = ", ".join( - [self._cxx_return_type_str, *out_return_types] - ) - self._return_numba_type_str = f"types.Tuple(({outs},))" + self._return_numba_type_str = _compose_return_type_str( + self._cxx_return_type_str, out_return_types + ) else: self._return_numba_type_str = self._cxx_return_type_str - def _tuple_literal(items: list[str]) -> str: - """ - Builds a Python tuple literal from a list of string expressions. - - Parameters: - items (list[str]): String representations of tuple elements. - - Returns: - tuple_literal (str): A Python tuple literal. For an empty list returns "()"; for a single element returns "(element,)" (includes the trailing comma); otherwise returns "(elem1, elem2, ...)". - """ - if not items: - return "()" - if len(items) == 1: - return f"({items[0]},)" - return f"({', '.join(items)})" - - intents_str = _tuple_literal( - [ - f"ArgIntent.{i.value if i != ArgIntent.in_ else 'in_'}" - for i in plan.intents - ] - ) - visible_str = repr(plan.visible_param_indices) - out_str = repr(plan.out_return_indices) - mask_str = repr(plan.pass_ptr_mask) - self._intent_plan_rendered = ( - "IntentPlan(" - f"intents={intents_str}, " - f"visible_param_indices={visible_str}, " - f"out_return_indices={out_str}, " - f"pass_ptr_mask={mask_str}" - ")" - ) + self._intent_plan_rendered = _render_intent_plan(plan) if out_return_types: self._out_return_types_rendered = ( "[" + ", ".join(out_return_types) + "]" diff --git a/numbast/src/numbast/static/struct.py b/numbast/src/numbast/static/struct.py index d2dec854..58623117 100644 --- a/numbast/src/numbast/static/struct.py +++ b/numbast/src/numbast/static/struct.py @@ -24,7 +24,13 @@ to_numba_arg_type_str, CTYPE_TO_NBTYPE_STR, ) -from numbast.intent import ArgIntent, IntentPlan, compute_intent_plan +from numbast.intent import ( + ArgIntent, + IntentPlan, + compute_intent_plan, + get_out_return_ptr_mask, + pointee_type_name, +) from numbast.utils import ( make_struct_ctor_shim, make_struct_conversion_operator_shim, @@ -39,6 +45,62 @@ file_logger.addHandler(FileHandler(logger_path)) +def _tuple_literal(items: list[str]) -> str: + if not items: + return "()" + if len(items) == 1: + return f"({items[0]},)" + return f"({', '.join(items)})" + + +def _out_return_type_str(param_type, *, pointer_out: bool) -> str: + type_name = param_type.unqualified_non_ref_type_name + if pointer_out: + type_name = pointee_type_name(type_name) + return to_numba_type_str(type_name) + + +def _compose_return_type_str(cxx_return_type_str: str, out_return_types: list[str]): + if not out_return_types: + return cxx_return_type_str + if cxx_return_type_str == "void": + if len(out_return_types) == 1: + return out_return_types[0] + outs = ", ".join(out_return_types) + return f"types.Tuple(({outs},))" + outs = ", ".join([cxx_return_type_str, *out_return_types]) + return f"types.Tuple(({outs},))" + + +def _prepend_receiver_to_intent_plan(method_plan: IntentPlan) -> IntentPlan: + return 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, + out_return_ptr_mask=(False,) + get_out_return_ptr_mask(method_plan), + ) + + +def _render_intent_plan(plan: IntentPlan) -> str: + intents_str = _tuple_literal( + [ + f"ArgIntent.{i.value if i != ArgIntent.in_ else 'in_'}" + for i in plan.intents + ] + ) + return ( + "IntentPlan(" + f"intents={intents_str}, " + f"visible_param_indices={repr(plan.visible_param_indices)}, " + f"out_return_indices={repr(plan.out_return_indices)}, " + f"pass_ptr_mask={repr(plan.pass_ptr_mask)}, " + f"out_return_ptr_mask={repr(get_out_return_ptr_mask(plan))}" + ")" + ) + + class StaticStructMethodRenderer(BaseRenderer): """Base class for all struct methods TODO: merge all common code paths @@ -660,15 +722,7 @@ def __init__( 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) self._arg_is_ref = None self._nb_param_types = [] @@ -690,47 +744,23 @@ def __init__( ", ".join(map(str, self._nb_param_types)) or "" ) + out_return_ptr_mask = get_out_return_ptr_mask(method_plan) out_return_types = [ - to_numba_type_str( - self._method_decl.param_types[ - i - ].unqualified_non_ref_type_name + _out_return_type_str( + self._method_decl.param_types[i], + pointer_out=out_return_ptr_mask[i], ) for i in method_plan.out_return_indices ] if out_return_types: self.Imports.add("from numba import types") - if self._cxx_return_type_str == "void": - if len(out_return_types) == 1: - self._nb_return_type_str = out_return_types[0] - else: - outs = ", ".join(out_return_types) - self._nb_return_type_str = f"types.Tuple(({outs},))" - else: - outs = ", ".join( - [self._cxx_return_type_str, *out_return_types] - ) - self._nb_return_type_str = f"types.Tuple(({outs},))" + self._nb_return_type_str = _compose_return_type_str( + self._cxx_return_type_str, out_return_types + ) else: self._nb_return_type_str = self._cxx_return_type_str - intents_str = ( - "(" - + ", ".join( - f"ArgIntent.{i.value if i != ArgIntent.in_ else 'in_'}" - for i in intent_plan.intents - ) - + ("," if len(intent_plan.intents) == 1 else "") - + ")" - ) - self._intent_plan_rendered = ( - "IntentPlan(" - f"intents={intents_str}, " - f"visible_param_indices={repr(intent_plan.visible_param_indices)}, " - f"out_return_indices={repr(intent_plan.out_return_indices)}, " - f"pass_ptr_mask={repr(intent_plan.pass_ptr_mask)}" - ")" - ) + self._intent_plan_rendered = _render_intent_plan(intent_plan) if out_return_types: self._out_return_types_rendered = ( "[" + ", ".join(out_return_types) + "]" diff --git a/numbast/src/numbast/static/tests/data/function_out.cuh b/numbast/src/numbast/static/tests/data/function_out.cuh index 0ba25f4c..fb611b4b 100644 --- a/numbast/src/numbast/static/tests/data/function_out.cuh +++ b/numbast/src/numbast/static/tests/data/function_out.cuh @@ -10,3 +10,5 @@ void __device__ add_out(int &out, int x); int __device__ add_out_ret(int &out, int x); int __device__ add_in_ref(int &x); void __device__ add_inout_ref(int &x, int delta); +void __device__ add_ptr_out(unsigned int *out, unsigned int x); +int __device__ add_ptr_out_ret(int x, unsigned int *out); diff --git a/numbast/src/numbast/static/tests/data/src/function_out.cu b/numbast/src/numbast/static/tests/data/src/function_out.cu index 037f85f3..d0110bac 100644 --- a/numbast/src/numbast/static/tests/data/src/function_out.cu +++ b/numbast/src/numbast/static/tests/data/src/function_out.cu @@ -37,3 +37,25 @@ int __device__ add_in_ref(int &x) { return x + 5; } * @param delta Amount to add to `x`. */ void __device__ add_inout_ref(int &x, int delta) { x += delta; } + +/** + * @brief Writes x + 1 through a scalar pointer output parameter. + * + * @param out Pointer to the scalar output slot. + * @param x Input value whose increment is written to `out`. + */ +void __device__ add_ptr_out(unsigned int *out, unsigned int x) { + *out = x + 1; +} + +/** + * @brief Stores x + 2 through a scalar pointer and returns x + 3. + * + * @param x Input value used to compute both results. + * @param out Pointer to the scalar output slot. + * @return int The value x + 3. + */ +int __device__ add_ptr_out_ret(int x, unsigned int *out) { + *out = (unsigned int)(x + 2); + return x + 3; +} diff --git a/numbast/src/numbast/static/tests/test_function_static_bindings.py b/numbast/src/numbast/static/tests/test_function_static_bindings.py index 93fae793..fe914cc2 100644 --- a/numbast/src/numbast/static/tests/test_function_static_bindings.py +++ b/numbast/src/numbast/static/tests/test_function_static_bindings.py @@ -6,7 +6,7 @@ import numpy as np import cffi -from numba.cuda.types import int32, float32 +from numba.cuda.types import int32, float32, uint32 from numba import cuda from numba.cuda import device_array @@ -51,11 +51,18 @@ def decl_out(make_binding): intents = { "add_out": {"out": "out_return"}, "add_out_ret": {"out": "out_return"}, + "add_ptr_out": {"out": "out_return"}, + "add_ptr_out_ret": {"out": "out_return"}, } res = make_binding("function_out.cuh", {}, {}, "sm_50", intents) bindings = res["bindings"] - public_apis = ["add_out", "add_out_ret"] + public_apis = [ + "add_out", + "add_out_ret", + "add_ptr_out", + "add_ptr_out_ret", + ] assert all(public_api in bindings for public_api in public_apis) return bindings @@ -165,9 +172,11 @@ def kernel(arr): def test_out_return_function_bindings(decl_out, impl_out): add_out = decl_out["add_out"] add_out_ret = decl_out["add_out_ret"] + add_ptr_out = decl_out["add_ptr_out"] + add_ptr_out_ret = decl_out["add_ptr_out_ret"] @cuda.jit(link=[impl_out]) - def kernel(out_single, out_pair): + def kernel(out_single, out_pair, out_ptr_single, out_ptr_pair): """ CUDA kernel that invokes `add_out` and `add_out_ret` to populate provided output buffers. @@ -178,15 +187,39 @@ def kernel(out_single, out_pair): ret, out = add_out_ret(7) out_pair[0] = ret out_pair[1] = out + out_ptr_single[0] = add_ptr_out(uint32(10)) + ptr_ret, ptr_out = add_ptr_out_ret(int32(7)) + out_ptr_pair[0] = ptr_ret + out_ptr_pair[1] = ptr_out out_single = device_array((1,), "int32") out_pair = device_array((2,), "int32") - kernel[1, 1](out_single, out_pair) + out_ptr_single = device_array((1,), "uint32") + out_ptr_pair = device_array((2,), "uint32") + kernel[1, 1](out_single, out_pair, out_ptr_single, out_ptr_pair) assert out_single.copy_to_host()[0] == 11 host_pair = out_pair.copy_to_host() assert host_pair[0] == 10 assert host_pair[1] == 9 + assert out_ptr_single.copy_to_host()[0] == 11 + host_ptr_pair = out_ptr_pair.copy_to_host() + assert host_ptr_pair[0] == 10 + assert host_ptr_pair[1] == 9 + + +def test_scalar_pointer_out_return_static_rendering(make_binding): + intents = { + "add_ptr_out": {"out": "out_return"}, + "add_ptr_out_ret": {"out": "out_return"}, + } + res = make_binding("function_out.cuh", {}, {}, "sm_50", intents) + src = res["src"] + + assert "signature(uint32, uint32)" in src + assert "signature(types.Tuple((int32, uint32,)), int32)" in src + assert "out_return_ptr_mask=(True, False)" in src + assert "out_return_ptr_mask=(False, True)" in src def test_out_ptr_in_inout_function_bindings(decl_out_ptr, impl_out): diff --git a/numbast/src/numbast/struct.py b/numbast/src/numbast/struct.py index 682f21b4..b95bf23b 100644 --- a/numbast/src/numbast/struct.py +++ b/numbast/src/numbast/struct.py @@ -22,7 +22,12 @@ from ast_canopy.decl import Struct, StructMethod from numbast.types import CTYPE_MAPS as C2N, to_numba_type, to_numba_arg_type -from numbast.intent import ArgIntent, IntentPlan, compute_intent_plan +from numbast.intent import compute_intent_plan +from numbast.intent_utils import ( + compose_return_type, + out_return_types_for_plan, + prepend_receiver_to_intent_plan, +) from numbast.utils import ( deduplicate_overloads, make_struct_regular_method_shim, @@ -289,15 +294,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) # Visible param types for @lower exclude receiver param_types = [] @@ -310,24 +307,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 diff --git a/numbast/tests/data/sample_function_out.cuh b/numbast/tests/data/sample_function_out.cuh index ef56004b..3f04702d 100644 --- a/numbast/tests/data/sample_function_out.cuh +++ b/numbast/tests/data/sample_function_out.cuh @@ -13,3 +13,12 @@ __device__ int add_out_ret(int &out, int x) { } __device__ int add_in_ref(int &x) { return x + 5; } + +__device__ void add_ptr_out(unsigned int *out, unsigned int x) { + *out = x + 1; +} + +__device__ int add_ptr_out_ret(int x, unsigned int *out) { + *out = (unsigned int)(x + 2); + return x + 3; +} diff --git a/numbast/tests/test_function.py b/numbast/tests/test_function.py index 83b23017..8337ac15 100644 --- a/numbast/tests/test_function.py +++ b/numbast/tests/test_function.py @@ -13,6 +13,7 @@ from numbast import bind_cxx_functions, MemoryShimWriter import pytest +from numba.cuda.types import int32, uint32 @pytest.fixture @@ -146,6 +147,8 @@ def _sample_out_functions(): arg_intent={ "add_out": {"out": "out_return"}, "add_out_ret": {"out": "out_return"}, + "add_ptr_out": {"out": "out_return"}, + "add_ptr_out_ret": {"out": "out_return"}, }, ) @@ -173,9 +176,11 @@ def test_out_return_device_function_results(_sample_out_functions): func_bindings, shim_writer = _sample_out_functions add_out = find_binding(func_bindings, "add_out") add_out_ret = find_binding(func_bindings, "add_out_ret") + add_ptr_out = find_binding(func_bindings, "add_ptr_out") + add_ptr_out_ret = find_binding(func_bindings, "add_ptr_out_ret") @cuda.jit(link=shim_writer.links()) - def kernel(out_single, out_pair): + def kernel(out_single, out_pair, out_ptr_single, out_ptr_pair): """ Populate provided output buffers with results produced by bound device functions. @@ -187,13 +192,22 @@ def kernel(out_single, out_pair): ret, out = add_out_ret(7) out_pair[0] = ret out_pair[1] = out + out_ptr_single[0] = add_ptr_out(uint32(10)) + ptr_ret, ptr_out = add_ptr_out_ret(int32(7)) + out_ptr_pair[0] = ptr_ret + out_ptr_pair[1] = ptr_out out_single = np.array([0], dtype=np.int32) out_pair = np.array([0, 0], dtype=np.int32) - kernel[1, 1](out_single, out_pair) + out_ptr_single = np.array([0], dtype=np.uint32) + out_ptr_pair = np.array([0, 0], dtype=np.uint32) + kernel[1, 1](out_single, out_pair, out_ptr_single, out_ptr_pair) assert out_single[0] == 11 assert out_pair[0] == 10 assert out_pair[1] == 9 + assert out_ptr_single[0] == 11 + assert out_ptr_pair[0] == 10 + assert out_ptr_pair[1] == 9 DATA_FOLDER = os.path.join(os.path.dirname(__file__), "data") p = os.path.join(DATA_FOLDER, "sample_function_out.cuh") diff --git a/numbast/tests/test_intent.py b/numbast/tests/test_intent.py new file mode 100644 index 00000000..0b99ec5a --- /dev/null +++ b/numbast/tests/test_intent.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from ast_canopy import pylibastcanopy +from numba import types as nbtypes + +from numbast.intent import compute_intent_plan, get_out_return_ptr_mask +from numbast.intent_utils import out_return_types_for_plan + + +def _type(name: str, *, lref: bool = False, rref: bool = False): + return pylibastcanopy.Type(name, name, rref, lref) + + +def _param(name: str, type_name: str): + return pylibastcanopy.ParamVar(name, _type(type_name)) + + +def test_out_return_accepts_scalar_pointer_outputs(): + params = [_param("x", "int"), _param("singleMipLevel", "unsigned int *")] + param_types = [p.type_ for p in params] + + plan = compute_intent_plan( + params=params, + param_types=param_types, + overrides={"singleMipLevel": "out_return"}, + ) + + assert plan.visible_param_indices == (0,) + assert plan.out_return_indices == (1,) + assert get_out_return_ptr_mask(plan) == (False, True) + assert out_return_types_for_plan(param_types, plan) == [nbtypes.uint32] + + +def test_out_return_reference_outputs_do_not_use_pointer_slot(): + params = [_param("out", "int"), _param("x", "int")] + param_types = [_type("int", lref=True), params[1].type_] + + plan = compute_intent_plan( + params=params, + param_types=param_types, + overrides={"out": "out_return"}, + ) + + assert plan.visible_param_indices == (1,) + assert plan.out_return_indices == (0,) + assert get_out_return_ptr_mask(plan) == (False, False) + assert out_return_types_for_plan(param_types, plan) == [nbtypes.int32] From 136bd0092a57b259d2863efed1682c6367391653 Mon Sep 17 00:00:00 2001 From: Michael Wang Date: Tue, 12 May 2026 15:33:44 -0700 Subject: [PATCH 2/2] Apply pre-commit formatting --- numbast/src/numbast/static/function.py | 4 +++- numbast/src/numbast/static/struct.py | 4 +++- numbast/tests/data/sample_function_out.cuh | 4 +--- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/numbast/src/numbast/static/function.py b/numbast/src/numbast/static/function.py index 26e96cbc..f5314ccd 100644 --- a/numbast/src/numbast/static/function.py +++ b/numbast/src/numbast/static/function.py @@ -60,7 +60,9 @@ def _out_return_type_str(param_type, *, pointer_out: bool) -> str: return to_numba_type_str(type_name) -def _compose_return_type_str(cxx_return_type_str: str, out_return_types: list[str]): +def _compose_return_type_str( + cxx_return_type_str: str, out_return_types: list[str] +): if not out_return_types: return cxx_return_type_str if cxx_return_type_str == "void": diff --git a/numbast/src/numbast/static/struct.py b/numbast/src/numbast/static/struct.py index 58623117..db3afd1f 100644 --- a/numbast/src/numbast/static/struct.py +++ b/numbast/src/numbast/static/struct.py @@ -60,7 +60,9 @@ def _out_return_type_str(param_type, *, pointer_out: bool) -> str: return to_numba_type_str(type_name) -def _compose_return_type_str(cxx_return_type_str: str, out_return_types: list[str]): +def _compose_return_type_str( + cxx_return_type_str: str, out_return_types: list[str] +): if not out_return_types: return cxx_return_type_str if cxx_return_type_str == "void": diff --git a/numbast/tests/data/sample_function_out.cuh b/numbast/tests/data/sample_function_out.cuh index 3f04702d..1db68729 100644 --- a/numbast/tests/data/sample_function_out.cuh +++ b/numbast/tests/data/sample_function_out.cuh @@ -14,9 +14,7 @@ __device__ int add_out_ret(int &out, int x) { __device__ int add_in_ref(int &x) { return x + 5; } -__device__ void add_ptr_out(unsigned int *out, unsigned int x) { - *out = x + 1; -} +__device__ void add_ptr_out(unsigned int *out, unsigned int x) { *out = x + 1; } __device__ int add_ptr_out_ret(int x, unsigned int *out) { *out = (unsigned int)(x + 2);