diff --git a/python/quadrants/lang/_pruning.py b/python/quadrants/lang/_pruning.py index 3289365767..3ec6f7aab6 100644 --- a/python/quadrants/lang/_pruning.py +++ b/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 ( diff --git a/python/quadrants/lang/ast/ast_transformers/call_transformer.py b/python/quadrants/lang/ast/ast_transformers/call_transformer.py index 3530f7964d..f3f6e44d7f 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, 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..9aaf121deb 100644 --- a/python/quadrants/lang/kernel.py +++ b/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: diff --git a/tests/python/test_py_dataclass.py b/tests/python/test_py_dataclass.py index e58641fe1b..9b5de77a51 100644 --- a/tests/python/test_py_dataclass.py +++ b/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