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

Disable threading in MLIR contexts.

Browse Source 
If threading is enabled, each MLIR context owns a threadpool of typically 3 threads. We might keep thousands of MLIR module objects alive in our caches, and we don't want to keep thousands of idle threads around. Disable the MLIR threading support. We don't do heavy compilation work from the Python-generated context anyway.

Fixes #16272
This commit is contained in:
Peter Hawkins 2023-06-06 19:46:30 -04:00
parent eea03ced0a
commit 08184cec86
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
jax/_src/interpreters/mlir.py
View File

@ -360,6 +360,14 @@ def _source_info_to_location(
def make_ir_context() -> ir.Context:
"""Creates an MLIR context suitable for JAX IR."""
context = ir.Context()
# If threading is enabled, each MLIR context will keep alive a thread pool.
# Since we cache MLIR modules (and hence contexts), this means we might keep
# several threads alive for each cache entry. This is a terrible idea. However
# we don't do any heavy computation on MLIR modules from Python anyway, so we
# just disable threading.
context.enable_multithreading(False)
dialects.mhlo.register_mhlo_dialect(context)
dialects.chlo.register_dialect(context)
dialects.stablehlo.register_dialect(context)