Fix GEMM tuner config restoration

tests/test_tune_gemm_config.py

@@ -0,0 +1,95 @@
+from unittest.mock import patch
+
+from utils.tune_gemm_config import GemmTuner
+
+
+ORIGINAL_CONFIG = "#define ORIGINAL 1\n"
+CANDIDATE = {
+    "act_parallel": True,
+    "row_block_size": 8,
+    "col_block_size": 64,
+    "parallel_size": 4,
+}
+
+
+def make_tuner(tmp_path):
+    config_path = tmp_path / "gemm-config.h"
+    config_path.write_text(ORIGINAL_CONFIG)
+    return GemmTuner(config_path, "model.gguf"), config_path
+
+
+def record_candidate(tuner, *, with_result):
+    def test_configuration(**configuration):
+        tuner.generate_config(**configuration)
+        if not with_result:
+            return None
+
+        result = {
+            **configuration,
+            "config_name": "candidate",
+            "pp_throughput": 1.0,
+            "pp_std_dev": 0.0,
+        }
+        tuner.results.append(result)
+        return result
+
+    return test_configuration
+
+
+def assert_backup_removed(tuner):
+    assert not tuner.backup_path.exists()
+
+
+def test_restores_config_when_no_candidate_succeeds(tmp_path):
+    tuner, config_path = make_tuner(tmp_path)
+
+    with patch.object(
+        tuner,
+        "test_configuration",
+        side_effect=record_candidate(tuner, with_result=False),
+    ):
+        tuner.run_tuning([CANDIDATE], output_csv=tmp_path / "results.csv")
+
+    assert config_path.read_text() == ORIGINAL_CONFIG
+    assert_backup_removed(tuner)
+
+
+def test_restores_config_when_final_rebuild_fails(tmp_path):
+    tuner, config_path = make_tuner(tmp_path)
+
+    with (
+        patch.object(
+            tuner,
+            "test_configuration",
+            side_effect=record_candidate(tuner, with_result=True),
+        ),
+        patch.object(tuner, "rebuild_project", return_value=False),
+        patch("builtins.input", return_value="y"),
+    ):
+        tuner.run_tuning([CANDIDATE], output_csv=tmp_path / "results.csv")
+
+    assert config_path.read_text() == ORIGINAL_CONFIG
+    assert_backup_removed(tuner)
+
+
+def test_keeps_config_when_final_rebuild_succeeds(tmp_path):
+    tuner, config_path = make_tuner(tmp_path)
+
+    with (
+        patch.object(
+            tuner,
+            "test_configuration",
+            side_effect=record_candidate(tuner, with_result=True),
+        ),
+        patch.object(tuner, "rebuild_project", return_value=True),
+        patch("builtins.input", return_value="y"),
+    ):
+        tuner.run_tuning([CANDIDATE], output_csv=tmp_path / "results.csv")
+
+    assert config_path.read_text() == (
+        "#define ACT_PARALLEL\n"
+        "#define ROW_BLOCK_SIZE 8\n"
+        "#define COL_BLOCK_SIZE 64\n"
+        "#define PARALLEL_SIZE 4\n"
+    )
+    assert_backup_removed(tuner)

utils/tune_gemm_config.py

@@ -180,6 +180,7 @@ def run_tuning(self, configurations, output_csv=None):
 
         # Backup configuration
         self.backup_config()
+        keep_tuned_config = False
 
         try:
             # Test all configurations
@@ -231,32 +232,25 @@ def run_tuning(self, configurations, output_csv=None):
                         best['col_block_size'],
                         best['parallel_size']
                     )
-                    self.rebuild_project()
-                    print("✅ Best configuration applied and project rebuilt!")
+                    if self.rebuild_project():
+                        keep_tuned_config = True
+                        print("✅ Best configuration applied and project rebuilt!")
+                    else:
+                        print("⚠️  Failed to rebuild with the best configuration; restoring the original configuration")
                 else:
-                    self.restore_config()
-                    print("✅ Original configuration restored")
-
-                # Clean up backup file
-                if self.backup_path.exists():
-                    self.backup_path.unlink()
-                    print(f"🗑️  Removed backup file: {self.backup_path}")
+                    print("✅ Original configuration will be restored")
 
         except KeyboardInterrupt:
             print("\n⚠️  Tuning interrupted by user")
-            self.restore_config()
-            # Clean up backup file
-            if self.backup_path.exists():
-                self.backup_path.unlink()
-                print(f"🗑️  Removed backup file: {self.backup_path}")
         except Exception as e:
             print(f"\n❌ Error during tuning: {e}")
-            self.restore_config()
-            # Clean up backup file
+            raise
+        finally:
+            if not keep_tuned_config:
+                self.restore_config()
             if self.backup_path.exists():
                 self.backup_path.unlink()
                 print(f"🗑️  Removed backup file: {self.backup_path}")
-            raise
 
 
 def generate_configurations():