Accept function parameters and closure cells as arguments with compile=True

[rostan-t]

Category:

New feature (non-breaking change which adds functionality)

Description:

#6382 detects constant variables to accept them as argument for transparent pipelining, as long as those variables are defined locally in a function. This PR adds support for parameters and closure cells.

The following can now be accepted:

def transform(jpegs, device, crop, mean, std):
    images = ndd.decoders.image(images, device=device)
    images = ndd.crop_mirror_normalize(images, mean=mean, std=std)
    return images


reader = ndd.readers.File(file_root=DATA_DIR)
for jpegs, _ in reader.next_epoch(batch=BATCH_SIZE):
    images = transform(
        jpegs,
        device="gpu",
        crop=(224, 224),
        mean=(0.485 * 255, 0.456 * 255, 0.406 * 255),
        std=(0.229 * 255, 0.224 * 255, 0.225 * 255),
    )

Things that still cannot be reliably captured include variables of mutable types, e.g. lists.

Additional information:

Affected modules and functionalities:

Transparent pipelining.

Key points relevant for the review:

• Is there any low hanging fruit that we could capture additionally?
• Do the choices made for closure cells make sense?
• Is the test coverage correct?

Tests:

• Existing tests apply:
  • test_compile.py
• New tests added
  • Python tests
  • GTests
  • Benchmark
  • Other
• N/A

Checklist

Documentation

• Existing documentation applies
• Documentation updated
  • Docstring
  • Doxygen
  • RST
  • Jupyter
  • Other
• N/A

DALI team only

Requirements

• Implements new requirements
• Affects existing requirements
• N/A

REQ IDs: N/A

JIRA TASK: DALI-4740

[greptile-apps]

Greptile Summary

This PR extends transparent pipelining to accept function parameters and closure cells as invariant arguments under compile=True, building on the constant-variable detection from #6382. It introduces a Binding datatype, replaces ModuleInfo.local_rhs with the more general binding(), and adds a tracing chain (_is_binding_invariant → _is_param_invariant → _is_arg_invariant) that walks the live call stack and CST to verify that what was passed to a parameter was itself an invariant expression.

• _source_analysis.py: ModuleInfo.binding() now returns Binding(in_scope, rhs) for both local variables and parameters; the new _is_param_invariant/_is_arg_invariant methods resolve the caller's call expression and use inspect.Signature.bind (with CST nodes as argument values) to map the parameter back to its call-site expression, recursing through chained and closure cases. _matches_callee handles plain functions, bound methods, and functools.partial.
• _compile.py: _matches now correctly handles equality comparisons that return numpy arrays or np.bool_ (rather than Python bool) by falling back to np.all(result).item().
• test_compile_invariants.py: 165 new test cases covering parameters, closures, cross-file, methods, functools.partial, default values, and all expected-rejection paths.

Confidence Score: 5/5

Safe to merge; the changes are additive and conservative — misclassification falls back to eager mode rather than producing incorrect compiled output.

The tracing logic is carefully gated: _is_binding_invariant and _is_immutable_value form independent guards so no mutable value can be captured even if one gate is over-generous. The _compile.py fix correctly handles numpy bool equality. Test coverage is thorough across the common patterns. The one design note (O(depth) tracing for recursive functions) affects only an unusual usage pattern and only on the first-trace pass.

No files require special attention; _source_analysis.py is the complex core and warrants a careful read of _is_binding_invariant's three branches.

Important Files Changed

Filename	Overview
dali/python/nvidia/dali/experimental/dynamic/_source_analysis.py	Core change: ModuleInfo.binding() replaces local_rhs(), exposing parameter and closure-cell bindings; new _is_binding_invariant, _is_param_invariant, _is_arg_invariant, _live_owner_frame, _split_call_args, _matches_callee methods implement the full tracing chain.
dali/python/nvidia/dali/experimental/dynamic/_compile.py	Bug fix in _matches: actual == expected can return a numpy array or numpy bool_ (not a Python bool) when compared against a numpy scalar; the np.all(result).item() fallback handles this correctly.
dali/test/python/experimental_mode/test_compile_invariants.py	165 new test cases covering parameters, closure cells, cross-file, method types, partial application, defaults, and all expected-rejection paths.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["classify(frame, inputs, kwargs)"] --> B["_Classifier.classify()"]
    B --> C{"call_at(frame)?"}
    C -- None --> D["return None (eager)"]
    C -- Call node --> E["_split_call_args(call)"]
    E -- has *args/**kwargs --> D
    E -- ok --> F["_capture_arg(node, value)"]
    F --> G{"is CompiledBatch?"}
    G -- yes --> H["return CompileRef"]
    G -- no --> I["is_invariant(node)"]
    I --> J{"node type?"}
    J -- literal/True/False/None --> K["True"]
    J -- Name --> L["_is_name_invariant"]
    J -- Attribute --> M["_is_dali_chain"]
    L --> O["ModuleInfo.binding"]
    O -- None --> P["False"]
    O -- Binding --> Q["_is_binding_invariant"]
    Q -- in_scope+rhs --> R["is_invariant(rhs)"]
    Q -- in_scope+param --> S["_is_param_invariant"]
    Q -- closure live owner --> T["_live_owner_frame"]
    T --> U["_Classifier(owner_frame)"]
    U --> S
    Q -- closure dead owner --> V["True (frozen cell)"]
    S --> W["resolve_callsite_frame"]
    W --> X["_is_arg_invariant"]
    X --> Y{"sig.bind"}
    Y -- omitted --> Z["default != empty"]
    Y -- provided --> AA["is_invariant(bound_arg)"]
    AA --> I
    L --> AB["_is_immutable_value(resolved)"]

Loading

_{Reviews (5): Last reviewed commit: "Add compile invariant tests for params a..." | Re-trigger Greptile}

[rostan-t]

!build

[jantonguirao]

[Minor / Design note] self.module_info is reused here as the CST analysis for the owner's frame. This is correct because Python closures are always lexically nested in the same source file, so the owner frame's variable bindings are covered by the same ModuleInfo. Worth a brief comment clarifying this assumption for readers who might wonder why we don't do _get_module_info(frame.f_code.co_filename) here like _is_param_invariant does.

[rostan-t]

What would you suggest?

[jantonguirao]

Maybe just:

          elif frame := self._live_owner_frame(name_node.value):
              # A live closure owner is a lexical parent, so it shares this module's CST metadata.
              classifier = _Classifier(self.module_info, frame)

The contrast with _is_param_invariant is that call arguments can come from a different caller module, while a live
closure owner is lexically enclosing this function, so self.module_info is the right one.

[jantonguirao]

[Minor] The comment # owner returned: frozen cell is accurate but doesn't tell the full story: returning True here only says the binding is stable (the cell can't be rebound once the owner has returned); the immutability of the value is still checked later in _is_name_invariant via _is_immutable_value(_static_eval(...)). This two-phase separation is the right design, but a one-liner like # cell is frozen; value immutability is checked by the caller would make the invariant clearer.

[rostan-t]

I don't think it's necessary to re-specify this for this specific case. Immutability is checked by the caller for all names, not just cells of closures with dead parents.

[jantonguirao]

[Test coverage] A few functools.partial and special parameter kinds are untested:

1. Positional partial: functools.partial(fn, first_arg_value) — the first positional argument is pre-filled, not a keyword. _matches_callee and inspect.signature should handle this, but it is worth confirming with a test.

2. Keyword-only parameters: def rotate(imgs, *, angle=60): — angle is KEYWORD_ONLY. The path through _is_arg_invariant should work (it doesn't special-case KEYWORD_ONLY), but there is no test.

3. Positional-only parameters: def rotate(imgs, angle, /): — sig.bind(*pos_nodes, **kw_nodes) handles these positionally, but again no test.

Even a quick expect_captured=True and one expect_captured=False for each would guard against regressions.

[rostan-t]

1. Partial functions with bound positional arguments are actually not supported because inspect.BoundArguments.apply_defaults doesn't restore positional arguments. However, I don't think that we should add a test for this case because it's a result of how BoundArguments handles arguments that we can accept because it's uncommon enough and there's an easy workaround.
2. and 3. fixed by 10c82ab.

[jantonguirao]

[Low-hanging fruit / design question] test_import_name is currently rejected because a from math import pi statement produces an ImportAssignment (not a plain Assignment) in LibCST's scope analysis, and ModuleInfo.binding filters these out with type(assignment) is not Assignment.

Local imports of immutable scalars seem like a natural case users would write:

def transform(images):
    from math import pi
    return ndd.rotate(images, angle=pi)

pi is a float (immutable), there is only one assignment, and it is function-scoped. Accepting ImportAssignment alongside Assignment (still guarded by _is_immutable_value + the single-assignment check) would cover this case with minimal risk. Is this intentionally deferred, or just not considered?

[rostan-t]

This is true but I would leave it as a follow-up because it's not in the scope for this PR.
I wouldn't call it a low-hanging fruit because it requires a whitelist of modules that we can accept. Accepting all imported variables of immutable types would not work as they can always be reassigned somewhere else.

[dali-automaton]

CI MESSAGE: [54075060]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [54075060]: BUILD PASSED

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [54208281]: BUILD STARTED

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [54208755]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [54208755]: BUILD PASSED

[rostan-t]

!build

[dali-automaton]

CI MESSAGE: [54548519]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [54548519]: BUILD PASSED