From 8c7e618d84fa354755b89cb74f5b49e9f3f91816 Mon Sep 17 00:00:00 2001 From: Hugh Perkins Date: Thu, 23 Jul 2026 06:57:46 -0700 Subject: [PATCH 1/6] [test] Add fastcache pruning tests for dead-static-branch and swapped-slot forwarding (#808) Adds failing repros for the two fastcache parameter-pruning bugs from #808: - dead-static-branch field forwarded through a caller path walked before the live instantiation populates the shared used-set (flat and 3-level nested) - same-flat-name forwarding into swapped callee slots, cross-caller and same-caller Also adds a reversed walk-order variant of the flat dead-static-branch case, which passes on main and guards against an order-dependent fix. --- tests/python/test_py_dataclass.py | 228 ++++++++++++++++++++++++++++++ 1 file changed, 228 insertions(+) diff --git a/tests/python/test_py_dataclass.py b/tests/python/test_py_dataclass.py index e58641fe1b..c3d700f915 100644 --- a/tests/python/test_py_dataclass.py +++ b/tests/python/test_py_dataclass.py @@ -1694,6 +1694,234 @@ 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_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_pruning_with_keyword_rename() -> None: @dataclasses.dataclass From ae35346c5ab4cef0de3036508911df065f6215d0 Mon Sep 17 00:00:00 2001 From: Hugh Perkins Date: Thu, 23 Jul 2026 08:58:03 -0700 Subject: [PATCH 2/6] [fastcache] Fix pruning for dead-static-branch and swapped-slot forwarding (#808) Record per-call-site call-graph edges during the discovery pass and propagate used-sets to a worklist fixpoint afterwards, instead of a single-shot copy keyed by callee. - Dead-static-branch (bug A): a callee's used-set (shared across template instantiations via its func id) can grow after an earlier caller copied it, so the enforcing pass pruned a parameter the callee still needed. The fixpoint makes propagation order-independent. - Swapped-slot forwarding (bug B): the caller-arg -> callee-param map was keyed by callee func id alone, so call sites forwarding the same flat name into different callee slots collided. Edges are now keyed per call site (source position). filter_call_args now takes the caller func id and looks up the per-call-site edge; kernel.materialize runs propagate_fixpoint() between discovery and the prefix-collapse. The fastcache-load path is unaffected (it skips discovery and only re-parses the kernel function def). --- python/quadrants/lang/_pruning.py | 175 +++++++++++------- .../ast/ast_transformers/call_transformer.py | 4 +- python/quadrants/lang/kernel.py | 3 + 3 files changed, 115 insertions(+), 67 deletions(-) diff --git a/python/quadrants/lang/_pruning.py b/python/quadrants/lang/_pruning.py index 3289365767..bdc5f98f72 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,35 @@ 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 (keyed by caller arg name) and drives ``filter_call_args``. + + 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, str] = {} + + class Pruning: """ We use the func id to uniquely identify each function. @@ -28,6 +56,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 +74,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 +88,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 +101,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[caller_arg_name] = 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 +180,11 @@ 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. 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 +194,11 @@ 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 {} + 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 +211,14 @@ 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 - ) + callee_param_name = positional_map.get(caller_arg_name) 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..fbff838aa9 100644 --- a/python/quadrants/lang/ast/ast_transformers/call_transformer.py +++ b/python/quadrants/lang/ast/ast_transformers/call_transformer.py @@ -393,7 +393,9 @@ 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: From 4725667cfafa1106541ce980d6bb357c6608de1c Mon Sep 17 00:00:00 2001 From: Hugh Perkins Date: Thu, 23 Jul 2026 09:30:38 -0700 Subject: [PATCH 3/6] [test] Add kwarg-forwarding variant of dead-static-branch fastcache pruning test (#808) Forwards the dataclass by keyword (md=md) at every call site, exercising the kwarg pruning path (_expand_Call_dataclass_kwargs gated on the callee used-set) rather than the positional filter_call_args, confirming the used-set fixpoint also closes over keyword-forwarded edges. use_deep stays positional to avoid passing the template value by keyword (an unrelated concern). --- tests/python/test_py_dataclass.py | 48 +++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/tests/python/test_py_dataclass.py b/tests/python/test_py_dataclass.py index c3d700f915..b2cbfb152a 100644 --- a/tests/python/test_py_dataclass.py +++ b/tests/python/test_py_dataclass.py @@ -1789,6 +1789,54 @@ def k1(md: MyDataclass) -> None: 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 From 9f73e7b40cae853f71477fe52885b8cc4336bcea Mon Sep 17 00:00:00 2001 From: Hugh Perkins Date: Thu, 23 Jul 2026 11:32:35 -0700 Subject: [PATCH 4/6] [fastcache] Collapse filter_call_args call to a single line (black) --- .../quadrants/lang/ast/ast_transformers/call_transformer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/quadrants/lang/ast/ast_transformers/call_transformer.py b/python/quadrants/lang/ast/ast_transformers/call_transformer.py index fbff838aa9..da92c9ae82 100644 --- a/python/quadrants/lang/ast/ast_transformers/call_transformer.py +++ b/python/quadrants/lang/ast/ast_transformers/call_transformer.py @@ -393,9 +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( - ctx.func.func_id, 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) From 4b2cf70e3cf93ac1e1fc32b475683f160ef9e669 Mon Sep 17 00:00:00 2001 From: Hugh Perkins Date: Thu, 23 Jul 2026 11:48:48 -0700 Subject: [PATCH 5/6] [test] Add same-name-into-two-slots fastcache pruning regression test (#808) --- tests/python/test_py_dataclass.py | 34 +++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/python/test_py_dataclass.py b/tests/python/test_py_dataclass.py index b2cbfb152a..f0c143fcef 100644 --- a/tests/python/test_py_dataclass.py +++ b/tests/python/test_py_dataclass.py @@ -1970,6 +1970,40 @@ def k1(p: MyDataclass, q: MyDataclass) -> None: 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 From 7654dab4c3bc5d2a9d91cbd4d3a69b1d3d3fad21 Mon Sep 17 00:00:00 2001 From: Hugh Perkins Date: Thu, 23 Jul 2026 11:48:48 -0700 Subject: [PATCH 6/6] [fastcache] Fix positional pruning when one flat name feeds multiple callee slots (#808) Key CallEdge.positional_map by caller arg name with an ordered list of callee params (one per occurrence) instead of a single value, and consume them in left-to-right order during the enforcing pass. A single name->name map let a later slot (e.g. the second md in inner(md, md)) overwrite an earlier one, so the enforcing pass could prune a positional arg the callee still needs, raising 'Missing argument'. Keying by name (not slot index) keeps the mapping robust to the positional shift that occurs when an upstream caller prunes fields out of a forwarded dataclass. --- python/quadrants/lang/_pruning.py | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/python/quadrants/lang/_pruning.py b/python/quadrants/lang/_pruning.py index bdc5f98f72..4c196ddf9d 100644 --- a/python/quadrants/lang/_pruning.py +++ b/python/quadrants/lang/_pruning.py @@ -27,7 +27,13 @@ class CallEdge: ``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 (keyed by caller arg name) and drives ``filter_call_args``. + ``__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. @@ -39,7 +45,7 @@ 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, str] = {} + self.positional_map: dict[str, list[str]] = {} class Pruning: @@ -128,7 +134,7 @@ def record_after_call( callee_param_name = callee_func.arg_metas_expanded[arg_id + self_offset].name # type: ignore edge.pairs.append((caller_arg_name, callee_param_name)) if caller_arg_name.startswith("__qd_"): - edge.positional_map[caller_arg_name] = callee_param_name + 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: @@ -183,8 +189,10 @@ def filter_call_args( used in build_Call, before making the call, in the enforcing pass (pass 1) 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. Note that this ONLY - handles args, not kwargs (kwargs are pruned in _expand_Call_dataclass_kwargs). + 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 ( @@ -197,6 +205,10 @@ def filter_call_args( 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 = [] for i, arg in enumerate(node_args): @@ -215,7 +227,10 @@ def filter_call_args( if type(arg) in {Name}: caller_arg_name = arg.id # type: ignore if caller_arg_name.startswith("__qd_"): - callee_param_name = positional_map.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 callee_used_args and callee_param_name.startswith("__qd_") ):