Repository: jax-ml/jax
Pull request #22899 (2024-08-06 18:07:53)

When caching is enabled, also enable XLA caching features as well
This PR makes it easier to enable all of the caching features in JAX and XLA with a single option. Now, when the JAX persistent cache is enabled (`JAX_COMPILATION_CACHE_DIR`), some XLA caching features will also be enabled to subdirectories of the JAX cache dir. The XLA caching features that are used can be selected via `JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES`.

Currently, there is an issue related to kernel naming when both `xla_gpu_kernel_cache_file` and the JAX persistent cache are enabled together, so only the autotune cache is enabled by default now. Once this is fixed, the default value of `JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES` should be `all`.

~Requires https://github.com/openxla/xla/pull/15636~
~Requires https://github.com/openxla/xla/pull/18450~

trevor-m: > The change is fine, but there are CI failures (possibly stale).

@hawkinsp Thanks for reviewing! I've rebased which should fix the CI failures
hawkinsp: One more thing: please squash your commits.
dfm: @trevor-m — Thanks for your patience here! Can you rebase your PR onto the current `main` branch? We'll get this in ASAP after that. Thanks!
trevor-m: @dfm Thanks for looking at this. However, we may need to hold off merging this a bit longer. We think there will be issues when using this feature with multihost. To solve it, we can set `xla_gpu_experimental_autotune_cache_mode` to `update` for rank 0 only and set it to `read` for the other ranks. We will need to expose that flag in the xla python bindings first.

We will need to do something similar for the kernel cache.
hawkinsp: This is blocked on https://github.com/openxla/xla/pull/18450, right?
nouiz: > This is blocked on [openxla/xla#18450](https://github.com/openxla/xla/pull/18450), right?

Right, it got approved twice. I'm not sure why it isn't merged.
CI complain about clang-format issue on not changed code:
https://github.com/openxla/xla/actions/runs/11390477508/job/31692190615?pr=18450

trevor-m: @hawkinsp The prerequisite MR is now merged
dfm: It looks like all the new tests need to be conditioned on the jaxlib version. The issue here is that most of the CI jobs run with the released version of jaxlib, which doesn't include the `AutotuneCacheMode` enum yet. Something like:

```python
from jax._lib import version as jaxlib_version
```

and then, in the test:

```python
if jaxlib_version <= (0, 4, 35):
  self.skipTest("...")  # <- add a description here
```

should do the trick!
dfm: It also looks like you need to add an explicit dependency on `:path` here:

https://github.com/jax-ml/jax/blob/4b4fb9dae9eb7e2740d70de5b4a610f979530382/jax/BUILD#L425-L438
trevor-m: Thanks @dfm, I've added those changes.
dfm: @trevor-m — This should be good to go now, but can you rebase onto the current main branch so that we can rerun the import? Quite a few internal google tests were still failing because https://github.com/openxla/xla/pull/19101 hadn't been merged yet, but now that that's in, I expect everything should be green. Thanks for your patience with this!

Patch:
diff --git a/docs/persistent_compilation_cache.md b/docs/persistent_compilation_cache.md
index 246d3a6cb..37afa2f59 100644
--- a/docs/persistent_compilation_cache.md
+++ b/docs/persistent_compilation_cache.md
@@ -24,6 +24,7 @@ import jax.numpy as jnp
 jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
 jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
 jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
+jax.config.update("jax_persistent_cache_enable_xla_caches", "xla_gpu_per_fusion_autotune_cache_dir")
 
 @jax.jit
 def f(x):
@@ -87,6 +88,23 @@ cc.set_cache_dir("/tmp/jax_cache")
 
 Note that both criteria need to be satisfied for a function to be cached.
 
+### Additional caching
+
+XLA supports additional caching mechanism which can be enabled alongside JAX's
+persistent compilation cache to further improve recompilation time.
+
+* `jax_persistent_cache_enable_xla_caches`: Possible values:
+
+   * `all`: enable all XLA caching features
+
+   * `none`: don't enable any extra XLA caching features
+
+   * `xla_gpu_kernel_cache_file`: only enable the kernel cache
+
+   * `xla_gpu_per_fusion_autotune_cache_dir`: (default value) only enable the
+      autotuning cache
+
+
 ### Google Cloud
 
 When running on Google Cloud, the compilation cache can be placed on a Google
diff --git a/jax/BUILD b/jax/BUILD
index e8e817ece..0da99677d 100644
--- a/jax/BUILD
+++ b/jax/BUILD
@@ -430,6 +430,7 @@ pytype_strict_library(
         ":config",
         ":mlir",
         ":monitoring",
+        ":path",
         ":profiler",
         ":traceback_util",
         ":xla_bridge",
diff --git a/jax/_src/compiler.py b/jax/_src/compiler.py
index 113f7507c..ebb1a2b54 100644
--- a/jax/_src/compiler.py
+++ b/jax/_src/compiler.py
@@ -29,10 +29,12 @@ from jax._src import config as config
 from jax._src import distributed
 from jax._src import lib
 from jax._src import monitoring
+from jax._src import path as pathlib
 from jax._src import profiler
 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 version as jaxlib_version
 from jax._src.lib.mlir import ir
 import numpy as np
 
@@ -241,6 +243,31 @@ def get_compile_options(
 
   debug_options.xla_detailed_logging = detailed_logging
 
+  # If persistent cache is enabled, also enable additional XLA caching features.
+  if compilation_cache.is_persistent_cache_enabled() and jaxlib_version > (0, 4, 35):
+    # compilation_cache_dir can't be None here, but the type checker is a bit
+    # strict.
+    path = pathlib.Path(config.compilation_cache_dir.value or "")
+    enabled_flags = config.persistent_cache_enable_xla_caches.value or ""
+
+    if enabled_flags == "all" or "xla_gpu_kernel_cache_file" in enabled_flags:
+      kernel_cache_path = path / "xla_gpu_kernel_cache_file"
+      debug_options.xla_gpu_kernel_cache_file = str(kernel_cache_path)
+      # This option is required to use the kernel cache.
+      debug_options.xla_gpu_enable_llvm_module_compilation_parallelism = True
+      logger.debug("Enabling XLA kernel cache at '%s'", kernel_cache_path)
+
+    if enabled_flags == "all" or "xla_gpu_per_fusion_autotune_cache_dir" in enabled_flags:
+      autotune_cache_path = path / "xla_gpu_per_fusion_autotune_cache_dir"
+      debug_options.xla_gpu_per_fusion_autotune_cache_dir = str(autotune_cache_path)
+      logger.debug("Enabling XLA autotuning cache at '%s'", autotune_cache_path)
+
+      # Set caching mode so that only process 0 can write to the cache.
+      if distributed.global_state.process_id == 0:
+        debug_options.xla_gpu_experimental_autotune_cache_mode = xc.AutotuneCacheMode.UPDATE
+      else:
+        debug_options.xla_gpu_experimental_autotune_cache_mode = xc.AutotuneCacheMode.READ
+
   return compile_options
 
 
diff --git a/jax/_src/config.py b/jax/_src/config.py
index f3edde699..a44b0125a 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -1369,6 +1369,15 @@ persistent_cache_min_entry_size_bytes = int_state(
           '  filesystem being used for the cache. '
           '* > 0: the actual minimum size desired; no overrides.'))
 
+# TODO: Change default to all
+persistent_cache_enable_xla_caches = optional_string_state(
+    name='jax_persistent_cache_enable_xla_caches',
+    default='xla_gpu_per_fusion_autotune_cache_dir',
+    help=('When the persistent cache is enabled, additional XLA caching will '
+          'also be enabled automatically. This option can be used to configure'
+          'which XLA caching methods will be enabled.'),
+)
+
 compilation_cache_include_metadata_in_key = bool_state(
     name='jax_compilation_cache_include_metadata_in_key',
     default=False,


Test patch:
diff --git a/tests/compilation_cache_test.py b/tests/compilation_cache_test.py
index 40c2181a9..d10558afb 100644
--- a/tests/compilation_cache_test.py
+++ b/tests/compilation_cache_test.py
@@ -40,6 +40,8 @@ from jax._src import path as pathlib
 from jax._src import test_util as jtu
 from jax._src import xla_bridge
 from jax._src.compilation_cache_interface import CacheInterface
+from jax._src.lib import xla_client as xc
+from jax._src.lib import version as jaxlib_version
 from jax.experimental.pjit import pjit
 from jax.sharding import PartitionSpec as P
 import numpy as np
@@ -535,6 +537,42 @@ class CompilationCacheTest(CompilationCacheTestCase):
     self.assertEqual(
         executable.fingerprint, deserialized_executable.fingerprint)
 
+  def test_persistent_cache_enable_xla_caches(self):
+    if jaxlib_version <= (0, 4, 35):
+      self.skipTest("Test requires AutotuneCacheMode bindings")
+    with config.compilation_cache_dir("jax-cache"):
+      with config.persistent_cache_enable_xla_caches("none"):
+        compile_options = compiler.get_compile_options(
+          num_replicas=1, num_partitions=1
+        )
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_kernel_cache_file, "")
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_enable_llvm_module_compilation_parallelism, False)
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_per_fusion_autotune_cache_dir, "")
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_experimental_autotune_cache_mode, xc.AutotuneCacheMode.UPDATE)
+      with config.persistent_cache_enable_xla_caches("all"):
+        compile_options = compiler.get_compile_options(
+          num_replicas=1, num_partitions=1
+        )
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_kernel_cache_file, "jax-cache/xla_gpu_kernel_cache_file")
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_enable_llvm_module_compilation_parallelism, True)
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_per_fusion_autotune_cache_dir, "jax-cache/xla_gpu_per_fusion_autotune_cache_dir")
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_experimental_autotune_cache_mode, xc.AutotuneCacheMode.UPDATE)
+      with config.persistent_cache_enable_xla_caches("xla_gpu_kernel_cache_file"):
+        compile_options = compiler.get_compile_options(
+          num_replicas=1, num_partitions=1
+        )
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_kernel_cache_file, "jax-cache/xla_gpu_kernel_cache_file")
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_enable_llvm_module_compilation_parallelism, True)
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_per_fusion_autotune_cache_dir, "")
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_experimental_autotune_cache_mode, xc.AutotuneCacheMode.UPDATE)
+      with config.persistent_cache_enable_xla_caches("xla_gpu_per_fusion_autotune_cache_dir"):
+        compile_options = compiler.get_compile_options(
+          num_replicas=1, num_partitions=1
+        )
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_kernel_cache_file, "")
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_enable_llvm_module_compilation_parallelism, False)
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_per_fusion_autotune_cache_dir, "jax-cache/xla_gpu_per_fusion_autotune_cache_dir")
+        self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_experimental_autotune_cache_mode, xc.AutotuneCacheMode.UPDATE)
 
 @jtu.with_config(
     jax_enable_compilation_cache=False,
@@ -570,5 +608,17 @@ class CompilationCacheDisabledTest(CompilationCacheTestCase):
           "/jax/compilation_cache/task_disabled_cache"]
       self.assertEqual(count_after_second_use, count_after_first_use)
 
+  def test_persistent_cache_enable_xla_caches_disabled(self):
+    if jaxlib_version <= (0, 4, 35):
+      self.skipTest("Test requires AutotuneCacheMode bindings")
+    with config.enable_compilation_cache(False):
+      compile_options = compiler.get_compile_options(
+        num_replicas=1, num_partitions=1
+      )
+      self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_kernel_cache_file, "")
+      self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_enable_llvm_module_compilation_parallelism, False)
+      self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_per_fusion_autotune_cache_dir, "")
+      self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_experimental_autotune_cache_mode, xc.AutotuneCacheMode.UPDATE)
+
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())