Build the optimization_plan.json data structure.
(
extracted: List[Dict[str, Any]],
)
| 428 | # --------------------------------------------------------------------------- |
| 429 | |
| 430 | def generate_optimization_plan( |
| 431 | extracted: List[Dict[str, Any]], |
| 432 | ) -> Dict[str, Any]: |
| 433 | """Build the optimization_plan.json data structure.""" |
| 434 | kernels_to_optimize = [] |
| 435 | total_pct = 0.0 |
| 436 | |
| 437 | for entry in extracted: |
| 438 | total_pct += entry["pct_total"] |
| 439 | kernels_to_optimize.append({ |
| 440 | "rank": entry["rank"], |
| 441 | "file": entry["output_file"], |
| 442 | "op_type": entry["op_type"], |
| 443 | "model_shape": entry["model_shape"], |
| 444 | "gpu_time_ms": entry["gpu_time_ms"], |
| 445 | "pct_total": entry["pct_total"], |
| 446 | "estimated_speedup_potential": SPEEDUP_ESTIMATES.get( |
| 447 | entry["op_type"], "1.5-2x" |
| 448 | ), |
| 449 | }) |
| 450 | |
| 451 | return { |
| 452 | "kernels_to_optimize": kernels_to_optimize, |
| 453 | "total_optimization_targets": len(kernels_to_optimize), |
| 454 | "covered_gpu_time_pct": round(total_pct, 1), |
| 455 | } |
| 456 | |
| 457 | |
| 458 | # --------------------------------------------------------------------------- |