Skip to content

Conversation

@dayo09
Copy link
Contributor

@dayo09 dayo09 commented Sep 10, 2025

For #219

@@ -0,0 +1,8 @@

def test():
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Redundant file

@jinevening
Copy link
Contributor

Summarized by cline.

Summary of the cond_support change

The diff titled cond_support introduces full support for the higher‑order PyTorch operator torch.cond (conditional execution) throughout the TICO code‑base. The main points are:

Area What was added / changed
New operator CircleIf (custom op circle_custom::if_) is registered in tico/utils/register_custom_op.py and a corresponding visitor CircleIfVisitor is added in tico/serialize/operators/op_circle_if.py.
IR mapping MapSubgraph pass (tico/experimental/controlflow/passes/map_subgraph.py) walks the exported program, extracts the true/false sub‑graphs of a torch.ops.higher_order.cond node, evaluates them with concrete inputs, and creates a Circle if node (torch.ops.circle_custom.if_).
Serialization circle_serializer.build_circle is extended to handle sub‑graphs: it now iterates over all graph modules (root + sub‑graphs) and exports tensors/operators for each. Helper functions _export_tensors and _export_tensors_for_subgraph were added.
Utility structs CondArgs and CircleIfArgs dataclasses were added to tico/utils/validate_args_kwargs.py to give a typed view of the arguments of torch.cond and the Circle if op.
Sub‑graph handling New tico/utils/subgraph.py provides freeze_subgraphs, get_frozen_subgraphs, and get_all_graph_modules utilities to manage and retrieve sub‑graphs during conversion.
Pass infrastructure PassBase.call signature was changed to accept a graph_module argument, and all passes were updated accordingly. The PassManager.run method now forwards the current graph module.
Shape inference freeze_subgraphs is called before the main conversion loop so that FakeTensor shape inference works for sub‑graphs.
Other adjustments Numerous files were updated to use the new graph_module argument, to import the new utilities, and to handle the new if operator (e.g., circle_serializer, convert.py, trace_decorators.py).
Tests / examples Example scripts (call.py, if.py) and a test module (test/modules/op/cond.py) were added to demonstrate conditional models.

Overall, the change adds a complete pipeline:

  1. Detect torch.ops.higher_order.cond in the exported graph.
  2. Extract its true/false sub‑graphs, evaluate them with concrete inputs to obtain meta‑values.
  3. Freeze sub‑graphs for shape inference.
  4. Create a Circle if operator that references the frozen sub‑graphs.
  5. Serialize the whole model (root + sub‑graphs) into a Circle binary.

@dayo09
Copy link
Contributor Author

dayo09 commented Oct 2, 2025

@jinevening FYI, I refactored the code and removed freeze_subgraph stages just now. It's got simpler. :-D

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants