fix(xla): bound diagnostic local-task threads

src/lib/mlxcel-xla/build.rs

@@ -67,9 +67,17 @@ fn compile_gemma3n_qmv_ptx() {
     }
 }
 
+fn add_iree_shim_sources(build: &mut cc::Build) {
+    build.file("csrc/xla_iree.c").file("csrc/xla_aux.c");
+    if env::var_os("CARGO_FEATURE_DIAGNOSTICS").is_some() {
+        build.file("csrc/xla_diagnostic_flags.c");
+    }
+}
+
 fn main() {
     println!("cargo:rerun-if-changed=csrc/xla_iree.c");
     println!("cargo:rerun-if-changed=csrc/xla_aux.c");
+    println!("cargo:rerun-if-changed=csrc/xla_diagnostic_flags.c");
     println!("cargo:rerun-if-changed=csrc/gemma3n_qmv.cu");
     println!("cargo:rerun-if-env-changed=NVCC");
     println!("cargo:rerun-if-env-changed=IREE_DIST");
@@ -103,9 +111,9 @@ fn main() {
             "IREE_CUDA_HOME={} is not an iree source+build tree (missing src/runtime/src/iree/runtime/api.h)",
             home.display()
         );
-        cc::Build::new()
-            .file("csrc/xla_iree.c")
-            .file("csrc/xla_aux.c")
+        let mut build = cc::Build::new();
+        add_iree_shim_sources(&mut build);
+        build
             .include(&src_inc)
             .include(&bld_inc)
             .define("XLA_GATE_CUDA", None)
@@ -138,9 +146,9 @@ fn main() {
              src/runtime/src/iree/runtime/api.h); run scripts/iree/setup-macos.sh",
             home.display()
         );
-        cc::Build::new()
-            .file("csrc/xla_iree.c")
-            .file("csrc/xla_aux.c")
+        let mut build = cc::Build::new();
+        add_iree_shim_sources(&mut build);
+        build
             .include(&src_inc)
             .include(&bld_inc)
             .compile("xla_iree");
@@ -164,11 +172,9 @@ fn main() {
         dist.display()
     );
 
-    cc::Build::new()
-        .file("csrc/xla_iree.c")
-        .file("csrc/xla_aux.c")
-        .include(&include)
-        .compile("xla_iree");
+    let mut build = cc::Build::new();
+    add_iree_shim_sources(&mut build);
+    build.include(&include).compile("xla_iree");
 
     // Bake the dist path so the session can find `bin/iree-compile` at runtime
     // (a runtime `IREE_DIST` env var still takes precedence). The vmfb must be

src/lib/mlxcel-xla/csrc/xla_diagnostic_flags.c

@@ -0,0 +1,44 @@
+// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Compiled only by mlxcel-xla's `diagnostics` feature. Production IREE startup
+// does not parse or override global flags.
+
+#include <stdio.h>
+
+#include "iree/base/tooling/flags.h"
+
+int xla_diagnostics_configure_local_task_threads(void) {
+  char program[] = "mlxcel-xla-diagnostics";
+  char topology_flag[] = "--task_topology_group_count=1";
+  char worker_stack_flag[] = "--task_worker_stack_size=0";
+  char* argv_storage[] = {program, topology_flag, worker_stack_flag, NULL};
+  char** argv = argv_storage;
+  int argc = 3;
+
+  iree_status_t status =
+      iree_flags_parse(IREE_FLAGS_PARSE_MODE_DEFAULT, &argc, &argv);
+  if (!iree_status_is_ok(status)) {
+    int status_code = (int)iree_status_code(status);
+    fputs("failed to configure diagnostics-only IREE task threads: ", stderr);
+    iree_status_fprint(stderr, status);
+    iree_status_ignore(status);
+    return status_code == 0 ? 1 : status_code;
+  }
+  if (argc != 1) {
+    fputs("diagnostics-only IREE task thread flags were not consumed\n", stderr);
+    return 1;
+  }
+  return 0;
+}

src/lib/mlxcel-xla/src/diagnostic_flags.rs

@@ -0,0 +1,76 @@
+// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#[cfg(feature = "diagnostics")]
+use std::ffi::c_int;
+#[cfg(feature = "diagnostics")]
+use std::sync::OnceLock;
+
+#[cfg(feature = "diagnostics")]
+static LOCAL_TASK_THREADS: OnceLock<Result<(), String>> = OnceLock::new();
+
+#[cfg(feature = "diagnostics")]
+unsafe extern "C" {
+    fn xla_diagnostics_configure_local_task_threads() -> c_int;
+}
+
+/// Configure bounded local-task threading before diagnostics create IREE.
+///
+/// This selects one topology group and lets pthreads choose the host-default
+/// worker stack instead of IREE's exact `PTHREAD_STACK_MIN` request. The IREE
+/// flag registry is process-global, so cache both success and failure and make
+/// repeated diagnostic calls observe one immutable parse result.
+#[cfg(feature = "diagnostics")]
+pub fn configure_diagnostic_local_task_threads() -> Result<(), String> {
+    LOCAL_TASK_THREADS
+        .get_or_init(|| {
+            // SAFETY: the diagnostics-only C helper takes no pointers, owns its
+            // synthetic argv for the duration of the call, and returns a status.
+            let status = unsafe { xla_diagnostics_configure_local_task_threads() };
+            if status == 0 {
+                Ok(())
+            } else {
+                Err(format!(
+                    "failed to configure diagnostics-only IREE local-task threads \
+                     (status {status})"
+                ))
+            }
+        })
+        .clone()
+}
+
+/// Whether the diagnostics-only threading override completed successfully.
+#[doc(hidden)]
+#[cfg(feature = "diagnostics")]
+pub fn diagnostic_local_task_threads_are_configured() -> bool {
+    LOCAL_TASK_THREADS.get().is_some_and(Result::is_ok)
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn native_helper_pins_bounded_host_thread_flags() {
+        let source = include_str!("../csrc/xla_diagnostic_flags.c");
+        assert!(source.contains("--task_topology_group_count=1"));
+        assert!(source.contains("--task_worker_stack_size=0"));
+        assert!(source.contains("IREE_FLAGS_PARSE_MODE_DEFAULT"));
+    }
+
+    #[test]
+    fn rust_wrapper_caches_the_process_global_parse_result() {
+        let source = include_str!("diagnostic_flags.rs");
+        assert!(source.contains("OnceLock<Result<(), String>>"));
+        assert!(source.contains(".get_or_init(||"));
+    }
+}

src/lib/mlxcel-xla/src/lib.rs

@@ -50,6 +50,8 @@ use std::path::{Path, PathBuf};
 use mlxcel_core::session::{InferenceSession, PreparedPrefill, SessionCapabilities};
 
 mod context;
+#[cfg(any(feature = "diagnostics", test))]
+mod diagnostic_flags;
 #[cfg(any(feature = "iree", test))]
 #[allow(dead_code)]
 mod numeric_dtype_contract;
@@ -150,6 +152,10 @@ pub use context::{
     CONTEXT_CAPACITY_ENV, ContextCapacityError, DEFAULT_CONTEXT_CAPACITY,
     context_capacity_from_env, validate_request_capacity,
 };
+#[cfg(feature = "diagnostics")]
+pub use diagnostic_flags::{
+    configure_diagnostic_local_task_threads, diagnostic_local_task_threads_are_configured,
+};
 #[cfg(all(feature = "diagnostics", xla_iree_cuda))]
 pub use emitter::{
     run_gemma3n_altup_correct_diagnostic_probe, run_gemma3n_altup_predict_diagnostic_probe,