Add exec_time_optimization_effort and memory_fitting_effort flags.

These flags control the amount of effort the compiler spends minimizing execution time and memory usage, respectively. They can be set via the command line, e.g. . Valid values are between -1.0 and 1.0, default is 0.0.

Author: google-labs-jules[bot]

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

jax/_src/compiler.py

@@ -200,6 +200,9 @@ def get_compile_options(
         setattr(build_options, name, env_options_overrides.pop(name))
     compile_options.env_option_overrides = list(env_options_overrides.items())
 
+  build_options.exec_time_optimization_effort = config.exec_time_optimization_effort.value
+  build_options.memory_fitting_effort = config.memory_fitting_effort.value
+
   debug_options = compile_options.executable_build_options.debug_options
   if lib.cuda_path is not None:
     debug_options.xla_gpu_cuda_data_dir = lib.cuda_path

jax/_src/config.py

@@ -1993,3 +1993,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].'
+)