[Graph] Support member ndarrays for qd.checkpoint and qd.graph_do_while

docs/source/user_guide/graph.md

@@ -87,7 +87,7 @@ solve(x, counter)
 # x is now incremented 10 times; counter is 0
 ```
 
-The argument to `qd.graph_do_while()` must be the name of a scalar `qd.i32` ndarray parameter. The loop body repeats while this value is non-zero.
+The argument to `qd.graph_do_while()` must reference a scalar `qd.i32` ndarray that the kernel can access — a bare kernel parameter (`qd.graph_do_while(counter)`), a `@qd.data_oriented` member ndarray (`qd.graph_do_while(self.counter)`), or a `@dataclasses.dataclass` parameter member (`qd.graph_do_while(params.counter)`). The loop body repeats while this value is non-zero.
 
 - On CUDA SM 9.0+ (Hopper), this uses CUDA conditional while nodes — the entire iteration runs on the GPU with no host involvement.
 - On older CUDA GPUs, AMDGPU, and non-GPU backends, it falls back to a host-side do-while loop (see the [backend support table](#backend-support)).
@@ -256,7 +256,7 @@ while status.yielded:
 
 - Must be used inside `@qd.kernel(graph=True, checkpoints=True)`. Without the flag, `qd.checkpoint(...)` raises `QuadrantsSyntaxError` at compile time with a fix-it pointing at `checkpoints=True`.
 - `cp_id` must be an int literal or an `IntEnum` value, and must be unique across the kernel.
-- `yield_on=` must be a kernel parameter that is a 0-d `qd.types.ndarray(qd.i32, ndim=0)`; expressions are not supported.
+- `yield_on=` must reference a 0-d `qd.types.ndarray(qd.i32, ndim=0)` — a bare kernel parameter (`yield_on=flag`), a `@qd.data_oriented` member ndarray (`yield_on=self.flag`), or a `@dataclasses.dataclass` parameter member (`yield_on=params.flag`). Arbitrary expressions are not supported.
 - Checkpoints cannot be nested inside other checkpoints. Checkpoints inside a `qd.graph_do_while` body are fine.
 - The body of a `with qd.checkpoint(...)` block cannot contain bare top-level statements (assignments, augmented assignments, or bare call/expression statements). Every top-level statement must be inside a `for`-loop (or other control-flow construct). A docstring as the first statement is allowed. Bare statements raise `QuadrantsSyntaxError` at compile time with a fix-it pointing at the explicit one-iteration `for`-wrap:
 

python/quadrants/lang/_fast_caching/src_hasher.py

@@ -18,8 +18,10 @@
 from .python_side_cache import PythonSideCache
 
 # Bumped whenever the persisted CacheValue schema changes (see create_cache_key). v2 replaced the single
-# graph_do_while_arg string with a nested level table.
-_CACHE_VALUE_SCHEMA_VERSION = "cachevalue-v2-gdw-levels"
+# graph_do_while_arg string with a nested level table. v3 added the AST-resolved flat C++ arg-ids for
+# qd.graph_do_while conditions and qd.checkpoint(yield_on=...) targets so the launch path can forward them
+# directly without per-launch name matching (necessary for @qd.data_oriented member ndarrays).
+_CACHE_VALUE_SCHEMA_VERSION = "cachevalue-v3-ast-resolved-ids"
 
 
 def create_cache_key(
@@ -69,17 +71,31 @@ class CacheValue(BaseModel):
     frontend_cache_key: str
     hashed_function_source_infos: list[HashedFunctionSourceInfo]
     used_py_dataclass_parameters: set[str]
-    # Nested graph_do_while level table as (cond_arg_name, parent_id) pairs, indexed by level id. None / empty for
-    # kernels without graph_do_while.
-    graph_do_while_levels: list[tuple[str, int]] | None = None
+    # Nested graph_do_while level table as (cond_arg_name, parent_id, cond_cpp_arg_id) triples, indexed by level
+    # id. None / empty for kernels without graph_do_while. ``cond_cpp_arg_id`` is the flat C++ arg-id resolved at
+    # AST-build time by ``ASTTransformer._resolve_ndarray_kernel_arg_id`` and is required by the launch path to
+    # support `@qd.data_oriented` member conditions (`qd.graph_do_while(self.counter)`) -- name-matching against
+    # ``arg_metas`` only resolves top-level parameters.
+    graph_do_while_levels: list[tuple[str, int, int]] | None = None
+    # AST-build-time-resolved checkpoint metadata, indexed by internal cp_id. Empty for kernels without any
+    # `with qd.checkpoint(...)` block. See `Kernel.checkpoint_yield_on_args` /
+    # `Kernel.checkpoint_yield_on_cpp_arg_ids` / `Kernel.checkpoint_user_labels_by_cp_id` for what each entry means.
+    # Restored alongside the C++-side cached kernel so the launch path can forward `yield_on=` arg-ids and
+    # translate `from_checkpoint=` labels without re-running the AST transformer.
+    checkpoint_yield_on_args: list[str | None] = []
+    checkpoint_yield_on_cpp_arg_ids: list[int] = []
+    checkpoint_user_labels_by_cp_id: list[int | None] = []
 
 
 def store(
     frontend_cache_key: str,
     fast_cache_key: str,
     function_source_infos: Iterable[FunctionSourceInfo],
     used_py_dataclass_parameters: set[str],
-    graph_do_while_levels: list[tuple[str, int]] | None = None,
+    graph_do_while_levels: list[tuple[str, int, int]] | None = None,
+    checkpoint_yield_on_args: list[str | None] | None = None,
+    checkpoint_yield_on_cpp_arg_ids: list[int] | None = None,
+    checkpoint_user_labels_by_cp_id: list[int | None] | None = None,
 ) -> None:
     """
     Note that unlike other caches, this cache is not going to store the actual value we want.
@@ -108,6 +124,9 @@ def store(
         hashed_function_source_infos=list(hashed_function_source_infos),
         used_py_dataclass_parameters=used_py_dataclass_parameters,
         graph_do_while_levels=graph_do_while_levels,
+        checkpoint_yield_on_args=checkpoint_yield_on_args or [],
+        checkpoint_yield_on_cpp_arg_ids=checkpoint_yield_on_cpp_arg_ids or [],
+        checkpoint_user_labels_by_cp_id=checkpoint_user_labels_by_cp_id or [],
     )
     cache.store(fast_cache_key, cache_value_obj.model_dump_json())
 
@@ -125,23 +144,19 @@ def _try_load(cache_key: str) -> CacheValue | None:
     return cache_value_obj
 
 
-def load(
-    cache_key: str,
-) -> tuple[set[str], str, list[tuple[str, int]] | None] | tuple[None, None, None]:
-    """
-    loads function source infos from cache, if available
-    checks the hashes against the current source code
+def load(cache_key: str) -> CacheValue | None:
+    """Load a validated ``CacheValue`` for *cache_key* if one exists and its source hashes still match, else None.
+
+    Returns the full ``CacheValue`` (rather than the historical 3-tuple) so callers can pick off the
+    AST-transformer-produced metadata (graph_do_while levels, checkpoint tables) without the loader having to grow
+    a new return slot every time we cache a new piece of AST output.
     """
     cache_value = _try_load(cache_key)
     if cache_value is None:
-        return None, None, None
+        return None
     if function_hasher.validate_hashed_function_infos(cache_value.hashed_function_source_infos):
-        return (
-            cache_value.used_py_dataclass_parameters,
-            cache_value.frontend_cache_key,
-            cache_value.graph_do_while_levels,
-        )
-    return None, None, None
+        return cache_value
+    return None
 
 
 def dump_stats() -> None:

python/quadrants/lang/ast/ast_transformer.py

@@ -1352,19 +1352,67 @@ def build_For(ctx: ASTTransformerFuncContext, node: ast.For) -> None:
                 return ASTTransformer.build_struct_for(ctx, node, is_grouped=False)
 
     @staticmethod
-    def _is_graph_do_while_call(node: ast.expr) -> str | None:
-        """If *node* is ``qd.graph_do_while(var)`` return the arg name, else None."""
+    def _is_graph_do_while_call(node: ast.expr) -> ast.expr | None:
+        """If *node* is ``qd.graph_do_while(arg)`` return the arg AST node, else None.
+
+        ``arg`` may be an ``ast.Name`` (a bare kernel parameter, e.g. ``counter``) or an ``ast.Attribute`` chain
+        (a ``@qd.data_oriented`` member ndarray such as ``self.counter`` or a ``@dataclasses.dataclass`` parameter
+        member such as ``params.counter``). The actual resolution to a kernel ndarray argument happens in
+        ``build_While`` via ``_resolve_ndarray_kernel_arg_id``.
+        """
         if not isinstance(node, ast.Call):
             return None
         func = node.func
-        if isinstance(func, ast.Attribute) and func.attr == "graph_do_while":
-            if len(node.args) == 1 and isinstance(node.args[0], ast.Name):
-                return node.args[0].id
-        if isinstance(func, ast.Name) and func.id == "graph_do_while":
-            if len(node.args) == 1 and isinstance(node.args[0], ast.Name):
-                return node.args[0].id
+        is_gdw = (isinstance(func, ast.Attribute) and func.attr == "graph_do_while") or (
+            isinstance(func, ast.Name) and func.id == "graph_do_while"
+        )
+        if not is_gdw:
+            return None
+        if len(node.args) == 1 and isinstance(node.args[0], (ast.Name, ast.Attribute)):
+            return node.args[0]
         return None
 
+    @staticmethod
+    def _resolve_ndarray_kernel_arg_id(
+        ctx: ASTTransformerFuncContext,
+        kernel,
+        node: ast.expr,
+        usage: str,
+    ) -> tuple[str, int]:
+        """Resolve an ndarray-referencing expression to ``(label, flat_cpp_arg_id)`` at AST-build time.
+
+        Shared between ``qd.checkpoint(yield_on=...)`` and ``qd.graph_do_while(...)`` to turn the control-flag
+        argument into the flat C++ arg-id the runtime matches against. ``node`` is an ``ast.Name`` (a bare kernel
+        parameter, e.g. ``flag``) or an ``ast.Attribute`` chain (e.g. ``self.flag`` for a ``@qd.data_oriented``
+        owner, or ``params.flag`` where ``params`` is a ``@dataclasses.dataclass`` kernel parameter). We build the
+        expression through the normal AST machinery and read the arg-id off the resulting external-tensor
+        expression -- this unifies the bare-param and member-ndarray cases, since both flatten to a real ndarray
+        kernel argument carrying its arg-id on the ``ExternalTensorExpression``.
+
+        ``usage`` is the call form (e.g. ``"qd.checkpoint(yield_on=...)"``) used in the error message. Raises
+        ``QuadrantsSyntaxError`` if the expression does not resolve to an ndarray kernel argument.
+        """
+        from quadrants.lang.any_array import AnyArray  # pylint: disable=C0415
+
+        label = ast.unparse(node)
+        bad = QuadrantsSyntaxError(
+            f"{usage} got {label!r} which does not resolve to an ndarray kernel parameter of "
+            f"{kernel.func.__name__!r}. The argument must reference an ndarray kernel parameter (e.g. "
+            f"`flag`) or a @qd.data_oriented member ndarray (e.g. `self.flag`); other expressions are not "
+            f"supported."
+        )
+        try:
+            built = build_stmt(ctx, node)
+        except Exception as e:  # noqa: BLE001 - any resolution failure is a user-facing misuse
+            raise bad from e
+        resolved_expr = built.ptr if isinstance(built, AnyArray) else built
+        if not (hasattr(resolved_expr, "is_external_tensor_expr") and resolved_expr.is_external_tensor_expr()):
+            raise bad
+        arg_id = _qd_core.get_external_tensor_arg_id(resolved_expr)
+        if not arg_id:
+            raise bad
+        return label, int(arg_id[0])
+
     @staticmethod
     def _is_checkpoint_call(node: ast.expr, global_vars: dict):
         """Thin forwarding wrapper around ``CheckpointTransformer.is_checkpoint_call``; the actual logic lives in module
@@ -1377,18 +1425,11 @@ def build_While(ctx: ASTTransformerFuncContext, node: ast.While) -> None:
         if node.orelse:
             raise QuadrantsSyntaxError("'else' clause for 'while' not supported in Quadrants kernels")
 
-        graph_do_while_arg = ASTTransformer._is_graph_do_while_call(node.test)
-        if graph_do_while_arg is not None:
+        graph_do_while_node = ASTTransformer._is_graph_do_while_call(node.test)
+        if graph_do_while_node is not None:
             from quadrants.lang.kernel import GraphDoWhileLevel  # pylint: disable=C0415
 
             kernel = ctx.global_context.current_kernel
-            arg_names = [m.name for m in kernel.arg_metas]
-            if graph_do_while_arg not in arg_names:
-                raise QuadrantsSyntaxError(
-                    f"qd.graph_do_while({graph_do_while_arg!r}) does not match any "
-                    f"parameter of kernel {kernel.func.__name__!r}. "
-                    f"Available parameters: {arg_names}"
-                )
             if not kernel.use_graph:
                 raise QuadrantsSyntaxError("qd.graph_do_while() requires @qd.kernel(graph=True)")
             # graph_do_while emits no loop IR; its body's for-loops must be top-level (offloaded) tasks. So it may only
@@ -1399,15 +1440,22 @@ def build_While(ctx: ASTTransformerFuncContext, node: ast.While) -> None:
                     "qd.graph_do_while() must be at the kernel top level or directly nested inside "
                     "another qd.graph_do_while(); it cannot appear inside a for-loop."
                 )
+            # Resolve the condition ndarray (bare parameter or @qd.data_oriented member) to its flat C++ arg-id at
+            # AST-build time -- the same id the runtime needs -- so the launch path forwards it directly with no
+            # per-launch name matching. ``cond_arg_name`` keeps the readable label (e.g. "counter" or "self.counter")
+            # for introspection and for the legacy ``graph_do_while_arg`` alias surfaced on Kernel.
+            cond_label, cond_cpp_arg_id = ASTTransformer._resolve_ndarray_kernel_arg_id(
+                ctx, kernel, graph_do_while_node, "qd.graph_do_while(...)"
+            )
             # Register this loop as a new nesting level (the body restriction is validated up-front in
             # FunctionDefTransformer). Outer loops get lower ids than the inner loops they contain.
             parent_id = kernel._graph_do_while_level_stack[-1] if kernel._graph_do_while_level_stack else -1
             level_id = len(kernel.graph_do_while_levels)
             kernel.graph_do_while_levels.append(
-                GraphDoWhileLevel(cond_arg_name=graph_do_while_arg, parent_id=parent_id)
+                GraphDoWhileLevel(cond_arg_name=cond_label, parent_id=parent_id, cond_cpp_arg_id=cond_cpp_arg_id)
             )
             if level_id == 0:
-                kernel.graph_do_while_arg = graph_do_while_arg
+                kernel.graph_do_while_arg = cond_label
             kernel._graph_do_while_level_stack.append(level_id)
             ctx.ast_builder.set_graph_do_while_level_id(level_id)
             try: