refactor: update kernel planning prompt to handle revision and new plan better, propagate autotuning summaries in pipeline agent and search classes

Author: shangkunwang01

MaxKernel/auto_agent/subagents/kernel_writing/prompts/kernel_planning_prompt.py

@@ -14,12 +14,13 @@
 Identify whether you are creating a **NEW plan** or performing a **REVISION**.
 
 *   **NEW Plan:**
-    *   No existing plan file is mentioned.
-    *   User provides a kernel filename or pastes code to optimize.
+    *   The plan path (`{kernel_plan_path?}`) is either not provided, or the file at that path is empty/missing.
+    *   Your primary input is the base kernel code that needs optimization.
 *   **REVISION:**
-    *   An existing plan path is provided: `{kernel_plan_path?}`.
-    *   You receive results from other subagents (compilation status, test results, profiling summary) indicating issues.
-
+    *   The plan file at `{kernel_plan_path?}` already contains an existing plan.
+    *   The `{optimized_kernel_path?}` already contains an optimized kernel implementation.
+    *   You receive execution results (compilation status, test results, profiling summary) from a previous attempt that need to be addressed.
+ 
 ### Step 2: Gather Context (Conditional)
 
 **For NEW Plans:**
@@ -32,7 +33,8 @@
 
 **For REVISIONS:**
 1.  **Read current plan:** Use `filesystem_tool` to read the existing plan at `{kernel_plan_path?}`.
-2.  **Review execution results:** Analyze the following to identify what needs improvement:
+2.  **Read optimized kernel:** Use `filesystem_tool` to read the `{optimized_kernel_path?}`.
+3.  **Review execution results:** Analyze the following to identify what needs improvement:
     *   Compilation Status: `{kernel_compilation_status?}`
     *   Test Results: `{test_results?}`
     *.  Autotune Summary: `{autotuning_summary?}`
@@ -48,7 +50,8 @@
 
 ## 1. Current Kernel Analysis
 - Brief description of what the kernel does
-- Current implementation approach
+- Current implementation approach 
+    - Analyze the `{base_kernel_path}` for NEW plans, or the `{optimized_kernel_path?}` for REVISIONS
 - Identified performance bottlenecks or issues
 
 ## 2. Optimization Strategy

MaxKernel/auto_agent/subagents/pipeline_agent.py

@@ -251,6 +251,7 @@ def _record_history_snapshot(self, ctx: InvocationContext, iteration: int):
       ),
       "test_status": ctx.session.state.get("test_results", {}),
       "latency_ms": latency,
+      "autotuning_summary": ctx.session.state.get("autotuning_summary", ""),
       "profiling_summary": ctx.session.state.get("profiling_summary", ""),
     }
 

MaxKernel/auto_search/graph.py

@@ -13,6 +13,7 @@ class EvaluationResult(BaseModel):
   correct: bool = False
   latency_ms: Optional[float] = None
   profiling_summary: Optional[str] = None
+  autotuning_summary: Optional[str] = None
   compilation_error: Optional[str] = None
   test_error: Optional[str] = None
 

MaxKernel/auto_search/worker.py

@@ -177,6 +177,7 @@ def _process_results(
     test_error = test_status.get("output") if not correct else None
 
     latency_ms = best_run.get("latency_ms")
+    autotuning_summary = best_run.get("autotuning_summary")
     profiling_summary = best_run.get("profiling_summary")
 
     # Read optimized code if it exists
@@ -206,6 +207,7 @@ def _process_results(
         compiled=compiled,
         correct=correct,
         latency_ms=latency_ms,
+        autotuning_summary=autotuning_summary,
         profiling_summary=profiling_summary,
         compilation_error=compilation_error,
         test_error=test_error,