Re-enable persistent cache on cpu.

CPU cache key now includes machine attributes, so there should no longer
be a problem with incompatible CPUs accessing the same cache entry.

PiperOrigin-RevId: 621341638

Author: jax authors

jax/_src/compiler.py

@@ -33,6 +33,7 @@
 from jax._src import traceback_util
 from jax._src.interpreters import mlir
 from jax._src.lib import xla_client as xc
+from jax._src.lib import xla_extension_version
 from jax._src.lib.mlir import ir
 from jax._src.xla_bridge import process_count
 import numpy as np
@@ -253,8 +254,8 @@ def compile_or_get_cached(
   # that supports serialization of executables.
   # TODO(skye): add warning when initializing cache on unsupported default platform
   supported_platforms = ["tpu", "gpu"]
-  # TODO(b/323256224): Add back support for CPU together with extra fields in a
-  # cache key with underlying hardware features (xla_extension_version >= 230).
+  if xla_extension_version >= 253:
+    supported_platforms.append("cpu")
   use_compilation_cache = (
       config.enable_compilation_cache.value
       and getattr(backend, "supports_executable_serialization", True)

tests/compilation_cache_test.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections import Counter
 from functools import partial
 import math
 import os
 import platform
 import tempfile
-from collections import Counter
 import unittest
 from unittest import mock
 from unittest import SkipTest
@@ -33,11 +33,12 @@
 from jax._src import compiler
 from jax._src import config
 from jax._src import distributed
-from jax._src.maps import xmap
 from jax._src import monitoring
 from jax._src import test_util as jtu
 from jax._src import xla_bridge
 from jax._src.lib import xla_client
+from jax._src.lib import xla_extension_version
+from jax._src.maps import xmap
 from jax.experimental.pjit import pjit
 from jax.sharding import PartitionSpec as P
 import numpy as np
@@ -71,9 +72,9 @@ class CompilationCacheTest(jtu.JaxTestCase):
 
   def setUp(self):
     super().setUp()
-    # TODO(b/323256224): Add back support for CPU together with extra fields in
-    # a cache key with underlying hardware features.
     supported_platforms = ["tpu", "gpu"]
+    if xla_extension_version >= 253:
+      supported_platforms.append("cpu")
 
     if not jtu.test_device_matches(supported_platforms):
       raise SkipTest(