Skip to content

Follow-up improvements to FrozenFunctionGraph infrastructure #2033

@ricardoV94

Description

@ricardoV94

Summary

Follow-up items identified during review of #1908 (Allow freezing of FunctionGraph for hashing).

Items

  • Remove mutable fgraph from Scan, use only frozen_fgraph — Scan currently stores both a mutable fgraph and a _frozen_fgraph. The mutable copy should be removed, with rewrites unfreezing as needed. Related: Scan Op is hashable but mutable in a way that affects its equality #1601. Code | Discussion

  • FrozenApply should not have a tag — Tags are runtime scratch space, which frozen nodes shouldn't need. Removing it requires unraveling downstream consumers that try to attach tags. Code | Discussion

  • Avoid double-cloning in Composite and ScalarLoop make_node — When input types don't match, both Ops unfreeze then rebuild via rebuild_collect_shared, effectively cloning twice. Should be a single pass. Composite | ScalarLoop | Discussion

  • Allow creating already-frozen FunctionGraphs directly — During fusion, we build a mutable FunctionGraph (via Composite(...)) only to immediately freeze it. We should support constructing a frozen graph in one step. Code | Discussion

  • Add FrozenFunctionGraph.replace for trivial replacements — For simple input substitutions (not full rewrites), a FrozenFunctionGraph.replace that returns a new frozen graph would avoid the unfreeze-clone-refreeze round-trip. Code | Discussion

  • Add dummy Output Apply to FrozenFunctionGraph.clients — Regular FunctionGraph includes a dummy Output node in clients of output variables. FrozenFunctionGraph currently omits this. Code | Discussion

  • Inner graph optimization should be an explicit rewrite — With this PR, we stop seeing the optimized inner graph in the compiled function. Optimizing inner graphs should be a dedicated rewrite pass, not something that happens implicitly at make_node/dispatch time. Discussion

  • Composite flattening (nested Composite) should be its own rewrite — Currently, nested Composites are eagerly flattened during construction. This should be a separate rewrite, since ScalarLoop can handle nested Composite/ScalarLoop inside. Discussion

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions