Function cmd_plan

orchestrate.py:716–790

Show the full optimization plan with Amdahl's law analysis.

(state: dict)

def cmd_plan(state: dict) -> None:
    """Show the full optimization plan with Amdahl's law analysis."""
    plan = load_plan()
    if plan is None:
        print("ERROR: No optimization_plan.json found. Run extract.py first.")
        sys.exit(1)

    kernels_plan = plan.get("kernels_to_optimize", plan.get("kernels", []))
    kernels_state = state["kernels"]

    # Build a lookup from file -> state entry
    state_by_file: dict[str, dict] = {}
    for k in kernels_state:
        state_by_file[k["file"]] = k
        state_by_file[Path(k["file"]).name] = k

    print()
    print("=" * 65)
    print("  AutoKernel -- Optimization Plan")
    print("=" * 65)
    print()

    # Plan table
    total_gpu_time = plan.get("total_gpu_time_ms", 0)
    if total_gpu_time > 0:
        print(f"  Total profiled GPU time: {total_gpu_time:.1f} ms")
        print()

    print(f"  {'Rank':<5} {'Op Type':<20} {'Shape':<30} {'GPU Time (ms)':<15} {'% Total':<10} {'Status':<12} {'Speedup':<10}")
    print(f"  {'-'*5} {'-'*20} {'-'*30} {'-'*15} {'-'*10} {'-'*12} {'-'*10}")

    for kp in kernels_plan:
        rank = kp.get("rank", "?")
        op_type = kp.get("op_type", "unknown")
        shape = kp.get("shape", "")
        if isinstance(shape, dict):
            shape = ", ".join(f"{k}={v}" for k, v in shape.items())
        elif isinstance(shape, list):
            shape = str(shape)
        gpu_time = kp.get("gpu_time_ms", 0)
        pct_total = kp.get("pct_total", 0)

        # Match to state
        file_key = kp.get("file", "")
        sk = state_by_file.get(file_key) or state_by_file.get(Path(file_key).name) if file_key else None

        status = sk["status"].upper() if sk else "UNKNOWN"
        speedup_str = f"{sk['speedup']:.2f}x" if sk and sk.get("speedup") else "--"

        # Truncate shape for display
        shape_disp = shape[:28] + ".." if len(str(shape)) > 30 else str(shape)

        print(f"  {rank:<5} {op_type:<20} {shape_disp:<30} {gpu_time:<15.2f} {pct_total:<10.1f} {status:<12} {speedup_str:<10}")

    print()

    # Amdahl's law what-if analysis
    print("  Amdahl's Law What-If Analysis:")