Repository: jax-ml/jax
Issue category: Feature Request

Issue:
Allow specifying `exec_time_optimization_effort` and `memory_fitting_effort` via CLI.
New compiler options (see title) have recently been [exposed](https://github.com/jax-ml/jax/issues/24625) in JAX.

To use them users have to change their Python code like this:

```
jax.jit(f, compiler_options={
  "exec_time_optimization_effort": 0.5, # change here
 })
``` 

It would be very convenient to not modify the code but rely on CLI instead. For example:

```
JAX_EXEC_TIME_OPTIMIZATION_EFFORT=0.5 python3 jax_program.py
```


Hints:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ce8b040439c0..b5758d107077 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -79,6 +79,7 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
   * {func}`jax.lax.linalg.eig` and the related `jax.numpy` functions
     ({func}`jax.numpy.linalg.eig` and {func}`jax.numpy.linalg.eigvals`) are now
     supported on GPU. See {jax-issue}`#24663` for more details.
+  * Added two new configuration flags, `jax_exec_time_optimization_effort` and `jax_memory_fitting_effort`, to control the amount of effort the compiler spends minimizing execution time and memory usage, respectively.  Valid values are between -1.0 and 1.0, default is 0.0.
 
 * Bug fixes
   * Fixed a bug where the GPU implementations of LU and QR decomposition would
diff --git a/jax/_src/compiler.py b/jax/_src/compiler.py
index d40ded556a5b..2d032069eb02 100644
--- a/jax/_src/compiler.py
+++ b/jax/_src/compiler.py
@@ -36,6 +36,7 @@
 from jax._src.interpreters import mlir
 from jax._src.lib import version as jaxlib_version
 from jax._src.lib import xla_client as xc
+from jax._src.lib import xla_extension_version
 from jax._src.lib.mlir import ir
 import numpy as np
 
@@ -191,6 +192,10 @@ def get_compile_options(
     assert device_assignment.computation_count() == num_partitions
     compile_options.device_assignment = device_assignment
 
+  if xla_extension_version >= 294:
+    build_options.exec_time_optimization_effort = config.exec_time_optimization_effort.value
+    build_options.memory_fitting_effort = config.memory_fitting_effort.value
+
   if env_options_overrides is not None:
     # Some overrides are passed directly on build_options.
     overrides_on_build_options = [
diff --git a/jax/_src/config.py b/jax/_src/config.py
index eea686fa0f3c..afb7cb498581 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -2000,3 +2000,15 @@ def _update_garbage_collection_guard(state, key, val):
         'to use this feature.'
     ),
 )
+
+exec_time_optimization_effort = float_state(
+    name='jax_exec_time_optimization_effort',
+    default=0.0,
+    help='Effort for minimizing execution time (higher means more effort), valid range [-1.0, 1.0].'
+)
+
+memory_fitting_effort = float_state(
+    name='jax_memory_fitting_effort',
+    default=0.0,
+    help='Effort for minimizing memory usage (higher means more effort), valid range [-1.0, 1.0].'
+)


Test patch: