You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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
Summary
Follow-up items identified during review of #1908 (Allow freezing of FunctionGraph for hashing).
Items
Remove mutable
fgraphfrom Scan, use onlyfrozen_fgraph— Scan currently stores both a mutablefgraphand 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 | DiscussionFrozenApply 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 | DiscussionAvoid double-cloning in Composite and ScalarLoop
make_node— When input types don't match, both Ops unfreeze then rebuild viarebuild_collect_shared, effectively cloning twice. Should be a single pass. Composite | ScalarLoop | DiscussionAllow 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 | DiscussionAdd
FrozenFunctionGraph.replacefor trivial replacements — For simple input substitutions (not full rewrites), aFrozenFunctionGraph.replacethat returns a new frozen graph would avoid the unfreeze-clone-refreeze round-trip. Code | DiscussionAdd dummy
OutputApply toFrozenFunctionGraph.clients— RegularFunctionGraphincludes a dummyOutputnode in clients of output variables.FrozenFunctionGraphcurrently omits this. Code | DiscussionInner 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. DiscussionComposite 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