[Bug] Forward fastcache dataclass fields read only in a dead qd.static branch through every caller.

python/quadrants/lang/_pruning.py

@@ -37,8 +37,21 @@ 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)
+        # By-name arg forwarding recorded per call site, mapping caller argument names to the callee parameter each
+        # binds. The key is (caller_func_id, callee_func_id, node.lineno, node.col_offset), i.e. keyed per call site
+        # rather than per callee or per caller, so repeated calls to one callee that bind the same flat name to
+        # different slots - inner(md, other) then inner(other, md) - do not clobber each other. The source position
+        # is stable across the two passes because each parses the same function source. Only args, not kwargs.
+        self.callee_param_by_caller_arg_name_by_call_site: dict[tuple[int, int, int, int], dict[str, str]] = (
+            defaultdict(dict)
+        )
+        # Every caller -> callee forwarding recorded in pass 0, as (caller_func_id, callee_func_id, name_pairs)
+        # where each name_pair is (caller_arg_name, callee_param_name) for a by-name forwarded argument. The
+        # single-shot copy in ``record_after_call`` only reflects what the callee had marked used at the instant of
+        # that call; a callee's used set keeps growing as its later call sites (and its other qd.static template
+        # instantiations, which share the func id) are discovered. ``propagate_used_to_fixpoint`` replays these
+        # edges against the finalized callee sets so a caller's used set is a superset of everything it will forward.
+        self.call_edges: list[tuple[int, int, list[tuple[str, str]]]] = []
 
     def mark_used(self, func_id: int, parameter_flat_name: str) -> None:
         assert not self.enforcing
@@ -82,10 +95,12 @@ def record_after_call(
         callee_func: Func = node.func.ptr.wrapper  # type: ignore
         has_self = type(func) is BoundQuadrantsCallable
         self_offset = 1 if has_self else 0
+        name_pairs: list[tuple[str, str]] = []
         for i, 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
+                name_pairs.append((caller_arg_name, callee_param_name))
                 if callee_param_name in callee_used_vars:
                     vars_to_unprune.add(caller_arg_name)
             arg_id += 1
@@ -96,30 +111,63 @@ def record_after_call(
         # 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.
         for kwarg in node_keywords:
-            if type(kwarg.value) in {Name}:
+            # kwarg.arg is None for ``**d`` unpacking, which names no single callee parameter; only named
+            # keywords contribute a concrete caller -> callee edge.
+            if type(kwarg.value) in {Name} and kwarg.arg is not None:
                 caller_arg_name = kwarg.value.id  # type: ignore
                 callee_param_name = kwarg.arg
+                name_pairs.append((caller_arg_name, callee_param_name))
                 if callee_param_name in callee_used_vars:
                     vars_to_unprune.add(caller_arg_name)
             arg_id += 1
+        # propagate_used_to_fixpoint re-derives these forwarding additions against the finalized callee sets, so
+        # this per-call union is a subset of the final used set; it is kept as the direct in-pass signal and to
+        # keep the discovery pass self-contained.
         self.used_vars_by_func_id[my_func_id].update(vars_to_unprune)
+        self.call_edges.append((my_func_id, callee_func_id, name_pairs))
 
-        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]
+        call_site_key = (my_func_id, callee_func_id, node.lineno, node.col_offset)
+        callee_param_by_called_arg_name = self.callee_param_by_caller_arg_name_by_call_site[call_site_key]
         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_"):
+                    # Record the caller-arg -> callee-param mapping unconditionally: whether the callee actually
+                    # needs the field is decided in filter_call_args against the finalized used set, so gating the
+                    # mapping on the callee's used set here (which is not yet complete when the callee's live-branch
+                    # instantiation is discovered later) would drop a forwarded field the enforcing pass still emits.
                     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
+                    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
+        self.callee_param_by_caller_arg_name_by_call_site[call_site_key] = callee_param_by_called_arg_name
+
+    def propagate_used_to_fixpoint(self) -> None:
+        """Close the recorded caller -> callee forwarding edges so that every caller's used set contains each
+        by-name argument the enforcing pass will forward to a callee. Must run once after pass 0 (discovery) and
+        before the parent-prefix reduction, so the reduction sees the completed leaf sets.
+
+        A field a callee reads only inside a compile-time-dead ``qd.static`` branch of one template instantiation is
+        still marked used under the callee's (instantiation-shared) func id via its live-branch instantiation. The
+        enforcing pass forwards that field from every caller regardless of which instantiation a given call reaches,
+        so every such caller must bind it - which requires it in the caller's own used set. The single-shot copy at
+        call time misses it whenever the live-branch instantiation is discovered only after that caller was walked;
+        replaying to a fixpoint against the finalized callee sets repairs the omission along the whole call chain."""
+        changed = True
+        while changed:
+            changed = False
+            for caller_func_id, callee_func_id, name_pairs in self.call_edges:
+                callee_used = self.used_vars_by_func_id[callee_func_id]
+                caller_used = self.used_vars_by_func_id[caller_func_id]
+                for caller_arg_name, callee_param_name in name_pairs:
+                    if callee_param_name in callee_used and caller_arg_name not in caller_used:
+                        caller_used.add(caller_arg_name)
+                        changed = True
 
     def filter_call_args(
         self,
+        ctx: "ASTTransformerFuncContext",
         quadrants_callable: "QuadrantsCallable",
         node: "ast.Call",
         node_args: list[expr],
@@ -139,6 +187,8 @@ def filter_call_args(
             return py_args
         func: Func = quadrants_callable.wrapper  # type: ignore
         callee_func_id = func.func_id
+        caller_func_id = ctx.func.func_id
+        call_site_key = (caller_func_id, callee_func_id, node.lineno, node.col_offset)
         caller_used_args = self.used_vars_by_func_id[callee_func_id]
         new_args = []
         callee_param_id = 0
@@ -168,7 +218,7 @@ 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 = self.callee_param_by_caller_arg_name_by_func_id[callee_func_id].get(
+                    callee_param_name = self.callee_param_by_caller_arg_name_by_call_site[call_site_key].get(
                         caller_arg_name
                     )
                     if callee_param_name is None or (

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, node, node_args, node_keywords, py_args)
 
                 node.ptr = func(*py_args, **py_kwargs)
 

python/quadrants/lang/kernel.py

@@ -554,6 +554,11 @@ 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:
+                    # Complete the cross-call used-var propagation before collapsing to parent prefixes: a callee's
+                    # used set (shared across its qd.static template instantiations) is only final now, so every
+                    # caller's set must be re-closed over it, or a field forwarded from a dead-branch call site would
+                    # be unbound in the enforcing pass. See ``Pruning.propagate_used_to_fixpoint``.
+                    pruning.propagate_used_to_fixpoint()
                     for used_parameters in pruning.used_vars_by_func_id.values():
                         new_used_parameters = set()
                         for param in used_parameters:

tests/python/test_py_dataclass.py

@@ -1694,6 +1694,187 @@ 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_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