mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add flags to exclude from cache-key generation.

Browse Source 
Some flags do not affect the compilation output. These should
not be part of the cache key, otherwise changing them will change
the key causing an unnecessary cache miss.

Synchronize the exclusions between the command-line flags and
DebugOptions. Add if-this-then-that lint checks to keep them
in sync.

PiperOrigin-RevId: 564474189
This commit is contained in:
jax authors 2023-09-11 12:52:47 -07:00
parent 05d2432e9e
commit 6c3b42d33c
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
jax/_src/cache_key.py
View File

@ -197,6 +197,7 @@ def _hash_serialized_compile_options(hash_obj, compile_options_obj):
# strings to empty. The exact values used to clear are not relevant as long
# as the same values are used everytime for each field.
debug_options = compile_options_copy.executable_build_options.debug_options
# LINT.IfChange(debug_options)
debug_options.xla_force_host_platform_device_count = 0
debug_options.xla_dump_to = ""
debug_options.xla_dump_hlo_module_re = ""
@ -214,6 +215,7 @@ def _hash_serialized_compile_options(hash_obj, compile_options_obj):
debug_options.xla_dump_hlo_as_long_text = False
debug_options.xla_dump_disable_metadata = False
debug_options.xla_dump_hlo_pipeline_re = ""
# LINT.ThenChange(:xla_flags)
return hash_obj.update(compile_options_copy.SerializeAsString())
@ -315,6 +317,7 @@ def _hash_platform(hash_obj, backend):
def _hash_xla_flags(hash_obj, extra_flag_prefixes: list[str]):
# LINT.IfChange(xla_flags)
xla_flags_to_exclude_from_cache_key = [
"--xla_dump_compress_protos",
"--xla_dump_module_metadata",
@ -327,6 +330,9 @@ def _hash_xla_flags(hash_obj, extra_flag_prefixes: list[str]):
"--xla_dump_hlo_as_url",
"--xla_dump_hlo_as_proto",
"--xla_dump_hlo_as_text",
"--xla_dump_hlo_as_long_text",
"--xla_dump_hlo_as_html",
"--xla_dump_hlo_as_dot",
"--xla_dump_to",
"--xla_force_host_platform_device_count",
"--xla_dump_disable_metadata",
@ -334,6 +340,7 @@ def _hash_xla_flags(hash_obj, extra_flag_prefixes: list[str]):
"--xla_tpu_sdc_checker_streamz_metric",
"--xla_tpu_sdc_checker_enable_sdc_event_callbacks",
]
# LINT.ThenChange(:debug_options)
xla_flags = []