Merge pull request #62 from AI-Hypercomputer/shangkun-fix-evaluation

Support nested outputs correctness check and use Xprof time instead of wall time by default

Author: shangkunwang01

MaxKernel/evaluation/benchmark.py

@@ -153,6 +153,9 @@ def benchmark(
     # Convert the new result to a dictionary
     res_dict = asdict(result)
     res_dict["speedup"] = result.speedup if result.speedup is not None else None
+    res_dict["speed_up_xprof"] = (
+      result.speed_up_xprof if result.speed_up_xprof is not None else None
+    )
 
     # Append to the main list and save immediately
     results.append(res_dict)

MaxKernel/evaluation/custom_types/evaluation_result.py

@@ -1,5 +1,5 @@
-from dataclasses import dataclass
-from typing import Optional
+from dataclasses import dataclass, field
+from typing import List, Optional
 
 
 @dataclass
@@ -14,9 +14,16 @@ class EvaluationResult:
   xprof_reference_time_ms: float = 0.0
   xprof_optimized_time_ms: float = 0.0
   error_trace: Optional[str] = None
+  logs: List[str] = field(default_factory=list)
 
   @property
   def speedup(self) -> Optional[float]:
     if self.optimized_time_ms == 0 or self.reference_time_ms == 0:
       return None
     return self.reference_time_ms / self.optimized_time_ms
+
+  @property
+  def speed_up_xprof(self) -> Optional[float]:
+    if self.xprof_optimized_time_ms == 0 or self.xprof_reference_time_ms == 0:
+      return None
+    return self.xprof_reference_time_ms / self.xprof_optimized_time_ms

MaxKernel/evaluation/evaluation_utils.py

@@ -90,18 +90,30 @@ def print_eval_result(result: EvaluationResult):
     print("-" * 40)
     print(f"Reference time:       {result.reference_time_ms:.3f} ms")
     print(f"Optimized time:       {result.optimized_time_ms:.3f} ms")
-    speedup = result.speedup
-    if speedup is not None:
-      print(f"Speedup:        {speedup:.2f}x")
+    wall_time_speedup = result.speedup
+    if wall_time_speedup is not None:
+      print(f"Wall time speedup:  {wall_time_speedup:.2f}x")
     else:
-      print("Speedup:        N/A")
+      print("Wall time speedup:  N/A")
+
+    xprof_speedup = result.speed_up_xprof
+    if xprof_speedup is not None:
+      print(f"XProf speedup:      {xprof_speedup:.2f}x")
+    else:
+      print("XProf speedup:      N/A")
+
+    if result.logs:
+      print("Harness Logs:")
+      for log_msg in result.logs:
+        print(f"  - {log_msg}")
   print("=" * 40 + "\n")
 
 
 def summarize_results(
   results: list,
   speedup_threshold: float,
   output_dir: Optional[str] = None,
+  use_xprof_speedup: bool = True,
 ) -> None:
   """
   Calculates and prints summary statistics for a list of evaluation results.
@@ -110,6 +122,7 @@ def summarize_results(
       results: A list of dictionaries, where each dictionary is an evaluation result.
       speedup_threshold: The minimum speedup factor to consider an improvement.
       output_dir: Optional directory path to save the summary report and stats.
+      use_xprof_speedup: Whether to use XProf speedup for the summary.
   """
   total_attempted = len(results)
   if not total_attempted:
@@ -123,9 +136,20 @@ def summarize_results(
   num_correct = len(correct_tasks)
 
   # Speedup calculations should only be on tasks that are numerically correct.
-  speedups = [
-    r["speedup"] for r in correct_tasks if r.get("speedup") is not None
-  ]
+  speedups = []
+  for r in correct_tasks:
+    s = None
+    if use_xprof_speedup:
+      if r.get("speed_up_xprof") is not None:
+        s = r["speed_up_xprof"]
+      elif r.get("speedup") is not None:
+        s = r["speedup"]
+    else:
+      if r.get("speedup") is not None:
+        s = r["speedup"]
+
+    if s is not None:
+      speedups.append(s)
 
   improvements = [s for s in speedups if s > speedup_threshold]
   num_improved = len(improvements)
@@ -214,7 +238,9 @@ def summarize_results(
     logger.info(f"Saved evaluation summary and stats to {output_dir}")
 
 
-def visualize_speed_up(results: list, output_dir: str) -> None:
+def visualize_speed_up(
+  results: list, output_dir: str, use_xprof_speedup: bool = True
+) -> None:
   """
   Visualizes the evaluation results.
 
@@ -223,6 +249,7 @@ def visualize_speed_up(results: list, output_dir: str) -> None:
       output_dir: Directory path to save the output PNG files.
                   Will generate speedup_distribution.png and
                   speedup_barplot.png in this directory.
+      use_xprof_speedup: Whether to use XProf speedup for the visualization.
   """
   os.makedirs(output_dir, exist_ok=True)
 
@@ -238,7 +265,16 @@ def set_log_ticks(ax, log_values):
   plot_data = []
   for r in results:
     is_valid = r.get("compiled_successfully") and r.get("numerically_correct")
-    s = r.get("speedup") if is_valid else None
+    s = None
+    if is_valid:
+      if use_xprof_speedup:
+        if r.get("speed_up_xprof") is not None:
+          s = r["speed_up_xprof"]
+        else:
+          s = r.get("speedup")
+      else:
+        s = r.get("speedup")
+
     log_s = math.log2(s) if s and s > 0 else -10.0
     plot_data.append((r["task_id"], log_s, not is_valid))
 

MaxKernel/evaluation/harness_code.py

@@ -149,6 +149,25 @@ def main():
     out_base_cpu = jax.device_get(out_base)
     del out_base
 
+    harness_logs = []
+
+    # Dirty all HBM memory leaves with NaN / Sentinel values to prevent cache reuse
+    try:
+      for leaf in jax.tree_util.tree_leaves(out_base_cpu):
+        if hasattr(leaf, "shape") and hasattr(leaf, "dtype"):
+          if jnp.issubdtype(leaf.dtype, jnp.floating) or jnp.issubdtype(leaf.dtype, jnp.complexfloating):
+            val = jnp.nan
+          elif jnp.issubdtype(leaf.dtype, jnp.bool_):
+            val = True
+          else:
+            val = 123  # Fits within int8/uint8 and all larger integer dtypes
+
+          dummy = jnp.full(leaf.shape, val, dtype=leaf.dtype)
+          dummy.block_until_ready()
+          del dummy
+    except Exception as e:
+      harness_logs.append(f"Failed to dirty HBM memory: {e}")
+
     try:
       jit_optimized = jax.jit(optimized_mod.computation, static_argnums=static_argnums)
       out_optimized = jax.block_until_ready(jit_optimized(*args))
@@ -160,16 +179,23 @@ def main():
           "error": str(e),
           "traceback": traceback.format_exc()
       }
+      if harness_logs:
+        result["logs"] = harness_logs
       with open("result.json", "w", encoding="utf-8") as f:
         json.dump(result, f)
       return
 
-    is_correct = jnp.allclose(out_base_cpu,
-                              out_optimized_cpu,
-                              atol={atol},
-                              rtol={rtol})
-    max_abs_diff = float(jnp.max(jnp.abs(out_base_cpu - out_optimized_cpu)))
-    max_rel_diff = float(jnp.max(jnp.abs((out_base_cpu - out_optimized_cpu) / out_base_cpu)))
+    out_base_flat = jax.tree_util.tree_leaves(out_base_cpu)
+    out_optimized_flat = jax.tree_util.tree_leaves(out_optimized_cpu)
+
+    is_correct = True
+    max_abs_diff = 0.0
+    max_rel_diff = 0.0
+
+    for b, o in zip(out_base_flat, out_optimized_flat):
+      is_correct = is_correct and bool(jnp.allclose(b, o, atol={atol}, rtol={rtol}))
+      max_abs_diff = max(max_abs_diff, float(jnp.max(jnp.abs(b - o))))
+      max_rel_diff = max(max_rel_diff, float(jnp.max(jnp.abs((b - o) / b))))
 
     if not is_correct:
       result = {
@@ -178,6 +204,8 @@ def main():
           "max_abs_diff": max_abs_diff,
           "max_rel_diff": max_rel_diff,
       }
+      if harness_logs:
+        result["logs"] = harness_logs
       with open("result.json", "w", encoding="utf-8") as f:
         json.dump(result, f)
       return
@@ -195,6 +223,8 @@ def main():
         "xprof_reference_time_ms": xprof_time_base,
         "xprof_optimized_time_ms": xprof_time_optimized,
     }
+    if harness_logs:
+      result["logs"] = harness_logs
     with open("result.json", "w", encoding="utf-8") as f:
       json.dump(result, f)
   except Exception as e:

MaxKernel/evaluation/xprof_utils.py

@@ -4,7 +4,7 @@
 
 
 def extract_xprof_time(
-  trace_dir: str, event_name: str, num_runs: int = 20
+  trace_dir: str, event_name: str, num_runs: int = 20, line_name: str = None
 ) -> float:
   """Extracts execution time for a named event from xprof trace.
 
@@ -80,12 +80,19 @@ def extract_xprof_time(
     except Exception as e:
       logging.warning(f"Failed to parse {file_path}: {e}")
 
-  if xla_op_events:
-    logging.info("Found kernel events starting with %. Using them.")
-    target_events = xla_op_events
+  if line_name == "XLA Ops":
+    if xla_op_events:
+      logging.info(f"Using XLA Ops events. Found {len(xla_op_events)} events.")
+      target_events = xla_op_events
+    else:
+      logging.info(
+        "XLA Ops events requested but not found. Falling back to XLA Modules events."
+      )
+      target_events = xla_module_events
   else:
+    # Default behavior or explicit "XLA Modules"
     logging.info(
-      "No xla_op_events events found. Falling back to xla_module_events matched events."
+      f"Using XLA Modules events. Found {len(xla_module_events)} events."
     )
     target_events = xla_module_events