From 20a222fc746b3c8de885ebd88da59258c89ef339 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 27 Jul 2026 11:20:17 +0900 Subject: [PATCH] fix(xla): bound diagnostic local-task threads --- src/lib/mlxcel-xla/build.rs | 28 ++++--- .../mlxcel-xla/csrc/xla_diagnostic_flags.c | 44 +++++++++++ src/lib/mlxcel-xla/src/diagnostic_flags.rs | 76 +++++++++++++++++++ src/lib/mlxcel-xla/src/lib.rs | 6 ++ 4 files changed, 143 insertions(+), 11 deletions(-) create mode 100644 src/lib/mlxcel-xla/csrc/xla_diagnostic_flags.c create mode 100644 src/lib/mlxcel-xla/src/diagnostic_flags.rs diff --git a/src/lib/mlxcel-xla/build.rs b/src/lib/mlxcel-xla/build.rs index e64260b8f..b54a39a78 100644 --- a/src/lib/mlxcel-xla/build.rs +++ b/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 diff --git a/src/lib/mlxcel-xla/csrc/xla_diagnostic_flags.c b/src/lib/mlxcel-xla/csrc/xla_diagnostic_flags.c new file mode 100644 index 000000000..74931f489 --- /dev/null +++ b/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 + +#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; +} diff --git a/src/lib/mlxcel-xla/src/diagnostic_flags.rs b/src/lib/mlxcel-xla/src/diagnostic_flags.rs new file mode 100644 index 000000000..2f753dff4 --- /dev/null +++ b/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> = 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>")); + assert!(source.contains(".get_or_init(||")); + } +} diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 5848eddf5..2b1328267 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/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,