diff --git a/python/quadrants/lang/_pruning.py b/python/quadrants/lang/_pruning.py index 3289365767..4c196ddf9d 100644 --- a/python/quadrants/lang/_pruning.py +++ b/python/quadrants/lang/_pruning.py @@ -1,12 +1,11 @@ from ast import Name, Starred, expr, keyword -from collections import defaultdict +from collections import defaultdict, deque from typing import TYPE_CHECKING, Any from ._exceptions import raise_exception from ._quadrants_callable import BoundQuadrantsCallable, QuadrantsCallable from .exception import QuadrantsSyntaxError from .func import Func -from .kernel_arguments import ArgMetadata if TYPE_CHECKING: import ast @@ -14,6 +13,41 @@ from .ast.ast_transformer_utils import ASTTransformerFuncContext +# (caller_func_id, callee_func_id, call node lineno, call node col_offset). We key call edges by the +# call site's source position rather than by object identity because the AST is re-parsed on every +# pass, so the ``ast.Call`` node is a different object in the discovery and enforcing passes; the +# source position is stable across passes and unique within a caller. +CallSiteKey = tuple[int, int, int, int] + + +class CallEdge: + """ + One ``@qd.func`` / ``@qd.kernel`` -> ``@qd.func`` call site, recorded during the discovery pass. + + ``pairs`` holds every ``(caller_arg_flat_name, callee_param_flat_name)`` correspondence, for both + positional args and kwargs; it drives the used-set fixpoint (a caller needs every argument it + forwards into a callee parameter the callee needs). ``positional_map`` holds only the positional + ``__qd_`` Name args and drives ``filter_call_args``. It maps each caller arg flat name to the list of + callee param flat names it feeds, in call-site (left-to-right) order - a list rather than a single + value because the same flat name can be forwarded into several slots of one call (e.g. + ``inner(md, md)``); a plain name->name map would let a later slot overwrite an earlier one and prune + an argument the callee still needs. We key by name (not by slot index) so the mapping survives the + positional shift that happens when an upstream caller prunes fields out of a forwarded dataclass, + which shortens this call's expanded ``node_args`` between the discovery and enforcing passes. + + Edges are keyed per call site (source position), not per callee, so two call sites that forward the + same flat name into different callee slots (swapped-slot forwarding) stay independent. + """ + + __slots__ = ("caller_func_id", "callee_func_id", "pairs", "positional_map") + + def __init__(self, caller_func_id: int, callee_func_id: int) -> None: + self.caller_func_id = caller_func_id + self.callee_func_id = callee_func_id + self.pairs: list[tuple[str, str]] = [] + self.positional_map: dict[str, list[str]] = {} + + class Pruning: """ We use the func id to uniquely identify each function. @@ -28,6 +62,15 @@ class Pruning: Note that we unify handling of func and kernel by using func_id KERNEL_FUNC_ID to denote the kernel. + + Propagation of used parameters from callee to caller is done as a fixpoint over the recorded call + edges (``propagate_fixpoint``), run once after the discovery pass. This is required because a + callee's used-set (shared across template instantiations via its func id) keeps growing as later + call sites and instantiations are discovered - e.g. a field read only inside a + ``qd.static(True)`` branch is marked used only when that instantiation is walked. A single forward + copy at call-record time would miss such a field for any caller walked earlier, so the enforcing + pass would prune a parameter the callee still needs. Iterating to a fixpoint makes the result + independent of discovery order. """ KERNEL_FUNC_ID = 0 @@ -37,8 +80,9 @@ def __init__(self, kernel_used_parameters: set[str] | None) -> None: self.used_vars_by_func_id: dict[int, set[str]] = defaultdict(set) if kernel_used_parameters is not None: self.used_vars_by_func_id[Pruning.KERNEL_FUNC_ID].update(kernel_used_parameters) - # only needed for args, not kwargs - self.callee_param_by_caller_arg_name_by_func_id: dict[int, dict[str, str]] = defaultdict(dict) + # One entry per call site (source position), recorded during discovery. Consumed by + # propagate_fixpoint() (used-set propagation) and filter_call_args() (per-call-site arg pruning). + self.edges_by_call_site: dict[CallSiteKey, CallEdge] = {} def mark_used(self, func_id: int, parameter_flat_name: str) -> None: assert not self.enforcing @@ -50,6 +94,10 @@ def enforce(self) -> None: def is_used(self, func_id: int, var_flat_name: str) -> bool: return var_flat_name in self.used_vars_by_func_id[func_id] + @staticmethod + def _call_site_key(caller_func_id: int, callee_func_id: int, node: "ast.Call") -> CallSiteKey: + return (caller_func_id, callee_func_id, node.lineno, node.col_offset) + def record_after_call( self, ctx: "ASTTransformerFuncContext", @@ -59,67 +107,78 @@ def record_after_call( node_keywords: list[keyword], ) -> None: """ - called from build_Call, after making the call, in pass 0 + called from build_Call, after making the call, in the discovery pass (pass 0) - note that this handles both args and kwargs + Records the call-graph edge for this call site (handles both args and kwargs). Used-set + propagation is deferred to ``propagate_fixpoint``; nothing here mutates the used-sets. """ if type(func) not in {QuadrantsCallable, BoundQuadrantsCallable}: return - my_func_id = ctx.func.func_id + caller_func_id = ctx.func.func_id callee_func_id = func.wrapper.func_id # type: ignore - # Copy the used parameters from the child function into our own function. - callee_used_vars = self.used_vars_by_func_id[callee_func_id] - vars_to_unprune: set[str] = set() - arg_id = 0 - # node.args ordering will match that of the called function's metas_expanded, - # because of the way calling with sequential args works. - # We need to look at the child's declaration - via metas - in order to get the name they use. - # We can't tell their name just by looking at our own metas. + # node.args ordering will match that of the called function's arg_metas_expanded, because of + # the way calling with sequential args works. We read the callee's declared (flat) parameter + # name from its metas - we can't tell their name just by looking at our own metas. # - # One issue is when calling data-oriented methods, there will be a `self`. We'll detect this - # by seeing if the childs arg_metas_expanded is exactly 1 longer than len(node.args) + len(node.kwargs) + # One issue is when calling data-oriented methods, there will be a `self`, which occupies the + # first callee meta slot; we skip it with self_offset. callee_func: Func = node.func.ptr.wrapper # type: ignore has_self = type(func) is BoundQuadrantsCallable self_offset = 1 if has_self else 0 - for i, arg in enumerate(node_args): + + edge = CallEdge(caller_func_id, callee_func_id) + for arg_id, arg in enumerate(node_args): if type(arg) in {Name}: caller_arg_name = arg.id # type: ignore callee_param_name = callee_func.arg_metas_expanded[arg_id + self_offset].name # type: ignore - if callee_param_name in callee_used_vars: - vars_to_unprune.add(caller_arg_name) - arg_id += 1 - # Note that our own arg_metas ordering will in general NOT match that of the child's. That's - # because our ordering is based on the order in which we pass arguments to the function, but the - # child's ordering is based on the ordering of their declaration; and these orderings might not - # match. - # This is not an issue because, for keywords, we don't need to look at the child's metas. - # We can get the child's name directly from our own keyword node. + edge.pairs.append((caller_arg_name, callee_param_name)) + if caller_arg_name.startswith("__qd_"): + edge.positional_map.setdefault(caller_arg_name, []).append(callee_param_name) + # For keywords we don't need the callee metas (whose ordering need not match ours): the + # callee's parameter name is available directly from our own keyword node. for kwarg in node_keywords: if type(kwarg.value) in {Name}: caller_arg_name = kwarg.value.id # type: ignore callee_param_name = kwarg.arg - if callee_param_name in callee_used_vars: - vars_to_unprune.add(caller_arg_name) - arg_id += 1 - self.used_vars_by_func_id[my_func_id].update(vars_to_unprune) - - used_callee_vars = self.used_vars_by_func_id[callee_func_id] - child_arg_id = 0 - child_metas: list[ArgMetadata] = node.func.ptr.wrapper.arg_metas_expanded # type: ignore - callee_param_by_called_arg_name = self.callee_param_by_caller_arg_name_by_func_id[callee_func_id] - for i, arg in enumerate(node_args): - if type(arg) in {Name}: - caller_arg_name = arg.id # type: ignore - if caller_arg_name.startswith("__qd_"): - callee_param_name = child_metas[child_arg_id + self_offset].name - if callee_param_name in used_callee_vars or not callee_param_name.startswith("__qd_"): - callee_param_by_called_arg_name[caller_arg_name] = callee_param_name - child_arg_id += 1 - self.callee_param_by_caller_arg_name_by_func_id[callee_func_id] = callee_param_by_called_arg_name + edge.pairs.append((caller_arg_name, callee_param_name)) # type: ignore + + self.edges_by_call_site[self._call_site_key(caller_func_id, callee_func_id, node)] = edge + + def propagate_fixpoint(self) -> None: + """ + Propagate used-sets from callees up to callers along the recorded call edges, until they stop + growing. Run once after the discovery pass, before the enforcing pass. + + A caller needs every argument it forwards into a callee parameter that the callee needs. Used- + sets only grow and parameters are finite, so this terminates. See the class docstring for why a + fixpoint (rather than a single forward copy at record time) is required. + """ + assert not self.enforcing + edges_by_callee: dict[int, list[CallEdge]] = defaultdict(list) + for edge in self.edges_by_call_site.values(): + edges_by_callee[edge.callee_func_id].append(edge) + + worklist: deque[int] = deque(self.used_vars_by_func_id.keys()) + queued: set[int] = set(worklist) + while worklist: + callee_func_id = worklist.popleft() + queued.discard(callee_func_id) + callee_used = self.used_vars_by_func_id[callee_func_id] + for edge in edges_by_callee.get(callee_func_id, ()): + caller_used = self.used_vars_by_func_id[edge.caller_func_id] + grew = False + for caller_arg, callee_param in edge.pairs: + if callee_param in callee_used and caller_arg not in caller_used: + caller_used.add(caller_arg) + grew = True + if grew and edge.caller_func_id not in queued: + worklist.append(edge.caller_func_id) + queued.add(edge.caller_func_id) def filter_call_args( self, + caller_func_id: int, quadrants_callable: "QuadrantsCallable", node: "ast.Call", node_args: list[expr], @@ -127,9 +186,13 @@ def filter_call_args( py_args: list[Any], ) -> list[Any]: """ - used in build_Call, before making the call, in pass 1 + used in build_Call, before making the call, in the enforcing pass (pass 1) - note that this ONLY handles args, not kwargs + Prunes positional args the callee does not need. Keyed per call site (via caller_func_id + + the call node position) so swapped-slot forwarding stays independent. When the same flat name is + forwarded into several slots (e.g. ``inner(md, md)``), the recorded callee params are consumed in + left-to-right order (one per occurrence) so each slot is decided against its own callee param. + Note that this ONLY handles args, not kwargs (kwargs are pruned in _expand_Call_dataclass_kwargs). """ # We can be called with callables other than qd.func, so filter those out: if ( @@ -139,18 +202,15 @@ def filter_call_args( return py_args func: Func = quadrants_callable.wrapper # type: ignore callee_func_id = func.func_id - caller_used_args = self.used_vars_by_func_id[callee_func_id] + callee_used_args = self.used_vars_by_func_id[callee_func_id] + edge = self.edges_by_call_site.get(self._call_site_key(caller_func_id, callee_func_id, node)) + positional_map = edge.positional_map if edge is not None else {} + # Per-name cursor: the k-th occurrence of a caller arg name consumes the k-th recorded callee + # param for that name. node_args is walked in the same left-to-right order as when the edge was + # recorded, so occurrences line up even if an upstream prune dropped some fields in between. + occurrence_by_name: dict[str, int] = {} + new_args = [] - callee_param_id = 0 - callee_metas: list[ArgMetadata] = node.func.ptr.wrapper.arg_metas_expanded # type: ignore - callee_metas_pruned = [] - for _callee_meta in callee_metas: - if _callee_meta.name.startswith("__qd_"): - if _callee_meta.name in caller_used_args: - callee_metas_pruned.append(_callee_meta) - else: - callee_metas_pruned.append(_callee_meta) - callee_metas = callee_metas_pruned for i, arg in enumerate(node_args): is_starred = type(arg) is Starred if is_starred: @@ -163,19 +223,17 @@ def filter_call_args( # we'll just dump the rest of the py_args in: new_args.extend(py_args[i:]) - callee_param_id += len(py_args[i:]) break if type(arg) in {Name}: caller_arg_name = arg.id # type: ignore if caller_arg_name.startswith("__qd_"): - callee_param_name = self.callee_param_by_caller_arg_name_by_func_id[callee_func_id].get( - caller_arg_name - ) + mapped = positional_map.get(caller_arg_name) + occurrence = occurrence_by_name.get(caller_arg_name, 0) + occurrence_by_name[caller_arg_name] = occurrence + 1 + callee_param_name = mapped[occurrence] if mapped is not None and occurrence < len(mapped) else None if callee_param_name is None or ( - callee_param_name not in caller_used_args and callee_param_name.startswith("__qd_") + callee_param_name not in callee_used_args and callee_param_name.startswith("__qd_") ): continue new_args.append(py_args[i]) - callee_param_id += 1 - py_args = new_args - return py_args + return new_args diff --git a/python/quadrants/lang/ast/ast_transformers/call_transformer.py b/python/quadrants/lang/ast/ast_transformers/call_transformer.py index 3530f7964d..da92c9ae82 100644 --- a/python/quadrants/lang/ast/ast_transformers/call_transformer.py +++ b/python/quadrants/lang/ast/ast_transformers/call_transformer.py @@ -393,7 +393,7 @@ def build_Call(ctx: ASTTransformerFuncContext, node: ast.Call, build_stmt, build try: pruning = ctx.global_context.pruning if pruning.enforcing: - py_args = pruning.filter_call_args(func, node, node_args, node_keywords, py_args) + py_args = pruning.filter_call_args(ctx.func.func_id, func, node, node_args, node_keywords, py_args) node.ptr = func(*py_args, **py_kwargs) diff --git a/python/quadrants/lang/kernel.py b/python/quadrants/lang/kernel.py index c56ccb3deb..00c731579b 100644 --- a/python/quadrants/lang/kernel.py +++ b/python/quadrants/lang/kernel.py @@ -554,6 +554,9 @@ def materialize(self, key: "CompiledKernelKeyType | None", py_args: tuple[Any, . if struct_primitive_launch_info: self._struct_primitive_launch_info_by_key[key] = struct_primitive_launch_info else: + # Propagate used-sets across call edges to a fixpoint (order-independent) before + # collapsing flat names to their used prefixes below. + pruning.propagate_fixpoint() for used_parameters in pruning.used_vars_by_func_id.values(): new_used_parameters = set() for param in used_parameters: diff --git a/tests/python/test_py_dataclass.py b/tests/python/test_py_dataclass.py index e58641fe1b..f0c143fcef 100644 --- a/tests/python/test_py_dataclass.py +++ b/tests/python/test_py_dataclass.py @@ -1694,6 +1694,316 @@ def k1(envs_idx: qd.types.NDArray[qd.i32, 1], md1: MyDataclass1, md2: MyDataclas k1(envs_idx, md1, md2=md2) +@test_utils.test() +def test_prune_used_parameters_fastcache_dead_static_branch(tmp_path: Path): + # inner() reads md.deep only inside a dead qd.static branch, so md.deep is marked used under inner's + # (instantiation-shared) func id via the qd.static(True) instantiation alone. Two separate caller paths + # forward md to inner: path_dead (qd.static(False)) is walked before path_live (qd.static(True)), so the + # single-shot used-set copy at path_dead's call misses md.deep. Without cross-call fixpoint propagation + # the enforcing pass still forwards md.deep from path_dead, which never bound it -> compile failure. + arch_name = qd.lang.impl.current_cfg().arch.name + for _it in range(3): + qd.init(arch=getattr(qd, arch_name), offline_cache_file_path=str(tmp_path), offline_cache=True) + + @dataclasses.dataclass + class MyDataclass: + base: qd.types.NDArray[qd.i32, 1] + deep: qd.types.NDArray[qd.i32, 1] + not_used: qd.types.NDArray[qd.i32, 1] + + @qd.func + def inner(md: MyDataclass, use_deep: qd.template()) -> None: + md.base[0] = 1 + if qd.static(use_deep): + md.deep[0] = 99 + + @qd.func + def path_dead(md: MyDataclass) -> None: + inner(md, False) + + @qd.func + def path_live(md: MyDataclass) -> None: + inner(md, True) + + @qd.kernel(fastcache=True) + def k1(md: MyDataclass) -> None: + path_dead(md) + path_live(md) + + base = qd.ndarray(qd.i32, (4,)) + deep = qd.ndarray(qd.i32, (4,)) + not_used = qd.ndarray(qd.i32, (4,)) + md = MyDataclass(base=base, deep=deep, not_used=not_used) + + k1(md) + assert base[0] == 1 + assert deep[0] == 99 + kernel_args_count_by_type = k1._primal.launch_stats.kernel_args_count_by_type + assert kernel_args_count_by_type[KernelBatchedArgType.QD_ARRAY] == 2 + + +@test_utils.test() +def test_prune_used_parameters_fastcache_dead_static_branch_reversed_order(tmp_path: Path): + # Same setup as the flat dead-static-branch case, but the kernel walks path_live (qd.static(True)) before + # path_dead (qd.static(False)). In this order inner's used-set already contains md.deep by the time + # path_dead's call is recorded, so this case compiles even without the fix. It is included as a regression + # guard: a correct fix must stay order-independent, so both walk orders must keep passing. + arch_name = qd.lang.impl.current_cfg().arch.name + for _it in range(3): + qd.init(arch=getattr(qd, arch_name), offline_cache_file_path=str(tmp_path), offline_cache=True) + + @dataclasses.dataclass + class MyDataclass: + base: qd.types.NDArray[qd.i32, 1] + deep: qd.types.NDArray[qd.i32, 1] + not_used: qd.types.NDArray[qd.i32, 1] + + @qd.func + def inner(md: MyDataclass, use_deep: qd.template()) -> None: + md.base[0] = 1 + if qd.static(use_deep): + md.deep[0] = 99 + + @qd.func + def path_dead(md: MyDataclass) -> None: + inner(md, False) + + @qd.func + def path_live(md: MyDataclass) -> None: + inner(md, True) + + @qd.kernel(fastcache=True) + def k1(md: MyDataclass) -> None: + path_live(md) + path_dead(md) + + base = qd.ndarray(qd.i32, (4,)) + deep = qd.ndarray(qd.i32, (4,)) + not_used = qd.ndarray(qd.i32, (4,)) + md = MyDataclass(base=base, deep=deep, not_used=not_used) + + k1(md) + assert base[0] == 1 + assert deep[0] == 99 + kernel_args_count_by_type = k1._primal.launch_stats.kernel_args_count_by_type + assert kernel_args_count_by_type[KernelBatchedArgType.QD_ARRAY] == 2 + + +@test_utils.test() +def test_prune_used_parameters_fastcache_dead_static_branch_kwargs(tmp_path: Path): + # Same dead-static-branch bug as the flat case, but md is forwarded by keyword (md=md) at every call site + # instead of positionally. Keyword-forwarded dataclasses are pruned in _expand_Call_dataclass_kwargs (gated + # on the callee used-set), a different path from the positional filter_call_args, so this confirms the used- + # set fixpoint also closes over keyword-forwarded edges. use_deep stays positional so the template value is + # not passed by keyword (an unrelated concern). + arch_name = qd.lang.impl.current_cfg().arch.name + for _it in range(3): + qd.init(arch=getattr(qd, arch_name), offline_cache_file_path=str(tmp_path), offline_cache=True) + + @dataclasses.dataclass + class MyDataclass: + base: qd.types.NDArray[qd.i32, 1] + deep: qd.types.NDArray[qd.i32, 1] + not_used: qd.types.NDArray[qd.i32, 1] + + @qd.func + def inner(use_deep: qd.template(), md: MyDataclass) -> None: + md.base[0] = 1 + if qd.static(use_deep): + md.deep[0] = 99 + + @qd.func + def path_dead(md: MyDataclass) -> None: + inner(False, md=md) + + @qd.func + def path_live(md: MyDataclass) -> None: + inner(True, md=md) + + @qd.kernel(fastcache=True) + def k1(md: MyDataclass) -> None: + path_dead(md=md) + path_live(md=md) + + base = qd.ndarray(qd.i32, (4,)) + deep = qd.ndarray(qd.i32, (4,)) + not_used = qd.ndarray(qd.i32, (4,)) + md = MyDataclass(base=base, deep=deep, not_used=not_used) + + k1(md) + assert base[0] == 1 + assert deep[0] == 99 + kernel_args_count_by_type = k1._primal.launch_stats.kernel_args_count_by_type + assert kernel_args_count_by_type[KernelBatchedArgType.QD_ARRAY] == 2 + + +@test_utils.test() +def test_prune_used_parameters_fastcache_dead_static_branch_nested(tmp_path: Path): + # Same dead-static-branch forwarding bug as the flat case, but the forwarded field lives three dataclass + # levels deep (top.mid.leaf.deep), confirming the fix closes the used set across arbitrary nesting depth. + arch_name = qd.lang.impl.current_cfg().arch.name + for _it in range(3): + qd.init(arch=getattr(qd, arch_name), offline_cache_file_path=str(tmp_path), offline_cache=True) + + @dataclasses.dataclass + class Leaf: + deep: qd.types.NDArray[qd.i32, 1] + not_used: qd.types.NDArray[qd.i32, 1] + + @dataclasses.dataclass + class Mid: + not_used: qd.types.NDArray[qd.i32, 1] + leaf: Leaf + + @dataclasses.dataclass + class Top: + base: qd.types.NDArray[qd.i32, 1] + mid: Mid + + @qd.func + def inner(top: Top, use_deep: qd.template()) -> None: + top.base[0] = 1 + if qd.static(use_deep): + top.mid.leaf.deep[0] = 99 + + @qd.func + def path_dead(top: Top) -> None: + inner(top, False) + + @qd.func + def path_live(top: Top) -> None: + inner(top, True) + + @qd.kernel(fastcache=True) + def k1(top: Top) -> None: + path_dead(top) + path_live(top) + + base = qd.ndarray(qd.i32, (4,)) + deep = qd.ndarray(qd.i32, (4,)) + mid_not_used = qd.ndarray(qd.i32, (4,)) + leaf_not_used = qd.ndarray(qd.i32, (4,)) + top = Top(base=base, mid=Mid(not_used=mid_not_used, leaf=Leaf(deep=deep, not_used=leaf_not_used))) + + k1(top) + assert base[0] == 1 + assert deep[0] == 99 + kernel_args_count_by_type = k1._primal.launch_stats.kernel_args_count_by_type + assert kernel_args_count_by_type[KernelBatchedArgType.QD_ARRAY] == 2 + + +@test_utils.test() +def test_prune_used_parameters_fastcache_forward_same_name_swapped_slots(tmp_path: Path): + # inner() reads only its first struct (a); b is entirely unused. caller1 and caller2 declare identically + # named parameters and forward them into swapped inner slots, so the flat name __qd_md__qd_x binds inner's + # used slot a in one caller and its unused slot b in the other. Keyed by callee alone, the mapping from a + # caller argument name to its callee parameter lets the second caller overwrite the first, so the enforcing + # pass prunes the field the first caller needs -> a Missing argument failure and a write to the wrong struct. + # Keying the mapping by (caller, callee) keeps the two call sites independent. + arch_name = qd.lang.impl.current_cfg().arch.name + for _it in range(3): + qd.init(arch=getattr(qd, arch_name), offline_cache_file_path=str(tmp_path), offline_cache=True) + + @dataclasses.dataclass + class MyDataclass: + x: qd.types.NDArray[qd.i32, 1] + + @qd.func + def inner(a: MyDataclass, b: MyDataclass) -> None: + a.x[0] = 42 + + @qd.func + def caller1(md: MyDataclass, other: MyDataclass) -> None: + inner(md, other) + + @qd.func + def caller2(md: MyDataclass, other: MyDataclass) -> None: + inner(other, md) + + @qd.kernel(fastcache=True) + def k1(p: MyDataclass, q: MyDataclass) -> None: + caller1(p, q) + caller2(p, q) + + p_x = qd.ndarray(qd.i32, (4,)) + q_x = qd.ndarray(qd.i32, (4,)) + k1(MyDataclass(x=p_x), MyDataclass(x=q_x)) + assert p_x[0] == 42 + assert q_x[0] == 42 + kernel_args_count_by_type = k1._primal.launch_stats.kernel_args_count_by_type + assert kernel_args_count_by_type[KernelBatchedArgType.QD_ARRAY] == 2 + + +@test_utils.test() +def test_prune_used_parameters_fastcache_forward_same_name_swapped_slots_same_caller(tmp_path: Path): + # Same swapped-slot forwarding as the cross-caller case, but both call sites live in a single caller: md + # binds inner's used slot a on the first line and its unused slot b on the second. The forwarding map is + # keyed per call site (source position), so the two calls stay independent; a map shared for the whole + # (caller, callee) pair would let the second line overwrite the first and prune the field the first needs. + arch_name = qd.lang.impl.current_cfg().arch.name + for _it in range(3): + qd.init(arch=getattr(qd, arch_name), offline_cache_file_path=str(tmp_path), offline_cache=True) + + @dataclasses.dataclass + class MyDataclass: + x: qd.types.NDArray[qd.i32, 1] + + @qd.func + def inner(a: MyDataclass, b: MyDataclass) -> None: + a.x[0] = 42 + + @qd.func + def caller(md: MyDataclass, other: MyDataclass) -> None: + inner(md, other) + inner(other, md) + + @qd.kernel(fastcache=True) + def k1(p: MyDataclass, q: MyDataclass) -> None: + caller(p, q) + + p_x = qd.ndarray(qd.i32, (4,)) + q_x = qd.ndarray(qd.i32, (4,)) + k1(MyDataclass(x=p_x), MyDataclass(x=q_x)) + assert p_x[0] == 42 + assert q_x[0] == 42 + kernel_args_count_by_type = k1._primal.launch_stats.kernel_args_count_by_type + assert kernel_args_count_by_type[KernelBatchedArgType.QD_ARRAY] == 2 + + +@test_utils.test() +def test_prune_used_parameters_fastcache_forward_same_name_same_call_two_slots(tmp_path: Path): + # A single call site forwards the SAME dataclass into two positional slots: inner(md, md). inner reads + # only its first struct (a); b is entirely unused, so the expanded call carries the same flat name in both + # slots. A positional map keyed by the caller flat name lets slot b overwrite slot a, so the enforcing pass + # prunes the field slot a needs -> a Missing argument failure. Keying the map by slot index keeps the two + # occurrences independent. + arch_name = qd.lang.impl.current_cfg().arch.name + for _it in range(3): + qd.init(arch=getattr(qd, arch_name), offline_cache_file_path=str(tmp_path), offline_cache=True) + + @dataclasses.dataclass + class MyDataclass: + x: qd.types.NDArray[qd.i32, 1] + + @qd.func + def inner(a: MyDataclass, b: MyDataclass) -> None: + a.x[0] = 42 + + @qd.func + def caller(md: MyDataclass) -> None: + inner(md, md) + + @qd.kernel(fastcache=True) + def k1(p: MyDataclass) -> None: + caller(p) + + p_x = qd.ndarray(qd.i32, (4,)) + k1(MyDataclass(x=p_x)) + assert p_x[0] == 42 + kernel_args_count_by_type = k1._primal.launch_stats.kernel_args_count_by_type + assert kernel_args_count_by_type[KernelBatchedArgType.QD_ARRAY] == 1 + + @test_utils.test() def test_pruning_with_keyword_rename() -> None: @dataclasses.dataclass