Set setstate/getstate methods to Config (#868)

Author: jansel

helion/runtime/config.py

@@ -111,6 +111,12 @@ def __eq__(self, other: object) -> bool:
     def __hash__(self) -> int:
         return hash(frozenset([(k, _list_to_tuple(v)) for k, v in self.config.items()]))
 
+    def __getstate__(self) -> dict[str, object]:
+        return dict(self.config)
+
+    def __setstate__(self, state: dict[str, object]) -> None:
+        self.config = dict(state)
+
     def to_json(self) -> str:
         """Convert the config to a JSON string."""
         return json.dumps(self.config, indent=2)

test/test_autotuner.py

@@ -4,6 +4,7 @@
 import math
 import os
 from pathlib import Path
+import pickle
 import random
 import tempfile
 from types import SimpleNamespace
@@ -135,6 +136,21 @@ def test_save_load_config(self):
             self.assertEqual(config, loaded_config)
         self.assertExpectedJournal(config.to_json())
 
+    def test_config_pickle_roundtrip(self):
+        config = helion.Config(
+            block_sizes=[64, 64, 32],
+            loop_orders=[[1, 0]],
+            num_warps=4,
+            num_stages=2,
+            indexing="tensor_descriptor",
+            extra_metadata={"nested": [1, 2, 3]},
+        )
+        restored = pickle.loads(pickle.dumps(config))
+        self.assertIsInstance(restored, helion.Config)
+        self.assertEqual(config, restored)
+        self.assertIsNot(config, restored)
+        self.assertIsNot(config.config, restored.config)
+
     def test_run_fixed_config(self):
         @helion.kernel(
             config=helion.Config(