Enable multi-threading in Jax Context with shared thread pool

Author: vfdev-5

jax/_src/interpreters/mlir.py

@@ -593,6 +593,15 @@ def module_to_bytecode(module: ir.Module) -> bytes:
 
 # Translation rules
 
+# Create one global thread pool that can be shared between multiple ir.Contexts
+# and enabling multi-threading
+# TODO: remove this check after jaxlib 0.5.4
+if hasattr(ir, "ThreadPool"):
+  global_thread_pool = ir.ThreadPool()
+else:
+  global_thread_pool = None
+
+
 class JaxIrContext(ir.Context):
   def __init__(self, *args, **kwargs):
     # Note: we're very intentionally *not* calling the __init__() of our
@@ -607,12 +616,16 @@ def make_ir_context() -> ir.Context:
   context.append_dialect_registry(upstream_dialects)
   context.load_all_available_dialects()
 
-  # 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)
+  # TODO: remove this check after v0.5.4 jaxlib
+  if global_thread_pool is not None:
+    context.set_thread_pool(global_thread_pool)
+  else:
+    # 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)
   # TODO(bartchr): Once JAX is released with SDY, remove the if.
   if dialects.sdy:
     dialects.sdy.register_dialect(context)