Federated Plan Computation & Optimization in Program-Level Fed Planner

[min-guk]

Overview

This PR enhances Program-Level Fed Planner by improving federated plan computation, cost modeling, visualization, and redundant TWrite/TRead handling.

---

1. Federated Plan Computation

• Computes federated plans for all Hop DAGs in DMLProgram.
• Implements rewiring logic for TWrite/TRead:
  • Uses TransTable (Map<String, Hop>) to track TWrite.
  • Retrieves from TransTable and connects at TRead.

---

2. Cost Model & Optimization

• Statement Block-Based Cost Weighting:
  • while block: ×10
  • if block: ×0.5
  • Example: if inside while → 1 × 10 × 0.5 = 5
• Cost Deduplication:
  • Cost *= weight / # of Parents

---

3. Federated Plan Visualization

• Converts FedPlan text output to a graph using: NetworkX, Matplotlib, Graphviz
  • Since only imports are used without modifying library code, there are no license conflicts.
• Helps in debugging.

---

4. Handling Redundant TWrite/TRead in If-Else Blocks

• Problem:
  • TWrite occurs in both if and else, but TRead needs both inputs.
  • Solution:
    • Immutable table (pre-if-else state).
    • If-table & else-table → merge into a unified table.
    • Supports nested if-else handling.

---

5. Hop DAG Disconnection Fixes

• Disconnected TWrite in Predicates:
  • Fix: Connect TWrite pred to root dummy node.
• Unreferenced Hops (e.g., print):
• Loop TWrite Not Used Outside:
  • No computing/forwarding cost.

---

6. TWrite/TRead Constraint Handling

• Forwarding Cost: Network cost when a node's federated output type differs from its parent.
• Self Computing Cost: Computation cost at the given node.
• Cumulative Cost = Self Computing Cost + Child’s Cumulative & Forwarding Cost.

Node Type	Federated Output Type	Computing Cost	Forwarding Cost
Child of TWrite	Identical	✅	❌ (Op → TWrite)
TWrite	Identical	❌	❌ (TWrite → TRead)
TRead	Identical	❌	✅ (TRead → Parent)
Parent of TRead	Can Change	✅	✅

[mboehm7]

LGTM - thanks for the patch @min-guk and sorry for the delay. During the merge I moved the group to staging, and fixed the formatting of java classes (tabs over spaces), missing license header in new python files, and a java doc issue.