From 7d03319be2ee7a8142a05d229452cd52216b238f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 18 Aug 2026 03:56:46 +0000 Subject: [PATCH] fix(runtime): infer Compress output shapes for large conditions A dynamic Compress output shape is the number of selected entries, so the condition tensor has to be read to compute it. The dynamic shape reader admitted only the small tensors that normally carry shape data, so a vision graph whose condition spans a full image failed shape inference with "no inferred shape ... produced by op Compress" and the request 500'd. Read a Compress condition through its own bounded path instead, so the operand it needs is materialized while an unrelated large tensor still is not. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0190e2eb-abe4-451f-b36d-44a035a99b7e --- .../src/executor/dispatch.rs | 34 ++++++++++++++++++- .../src/executor/geometry.rs | 8 +++++ .../src/executor/tests.rs | 18 ++++++++++ 3 files changed, 59 insertions(+), 1 deletion(-) diff --git a/crates/onnx-runtime-session/src/executor/dispatch.rs b/crates/onnx-runtime-session/src/executor/dispatch.rs index c89a99f8f9..be5dcef893 100644 --- a/crates/onnx-runtime-session/src/executor/dispatch.rs +++ b/crates/onnx-runtime-session/src/executor/dispatch.rs @@ -694,7 +694,13 @@ impl Executor { .iter() .enumerate() .map(|(i, v)| { - v.and_then(|vid| self.shape_input_i64(vid, &input_shapes[i], input_dtypes[i])) + v.and_then(|vid| { + if node.is_default_domain() && node.op_type == "Compress" && i == 1 { + self.compress_condition_i64(vid, &input_shapes[i], input_dtypes[i]) + } else { + self.shape_input_i64(vid, &input_shapes[i], input_dtypes[i]) + } + }) }) .collect(); // Only materialize a *float* input value for the specific inputs an @@ -1045,6 +1051,32 @@ impl Executor { self.input_i64(vid, shape, dtype) } + fn compress_condition_i64( + &self, + vid: ValueId, + shape: &[usize], + dtype: DataType, + ) -> Option> { + if !bounded_compress_condition(dtype, shape) { + return None; + } + let max_bytes = shape[0].checked_mul(dtype.byte_size())?; + if let Some(view) = self.views.get(&vid) { + let source = self.buffers.get(&view.source)?; + if source.len() > max_bytes { + return None; + } + } + if self + .seq_elem_values + .get(&vid) + .is_some_and(|elem| elem.root_len() > max_bytes) + { + return None; + } + self.input_i64(vid, shape, dtype) + } + pub(super) fn shape_input_f64( &self, vid: ValueId, diff --git a/crates/onnx-runtime-session/src/executor/geometry.rs b/crates/onnx-runtime-session/src/executor/geometry.rs index 7bf12cda8b..122434a1bc 100644 --- a/crates/onnx-runtime-session/src/executor/geometry.rs +++ b/crates/onnx-runtime-session/src/executor/geometry.rs @@ -385,6 +385,14 @@ pub(super) fn bounded_shape_input(dtype: DataType, shape: &[usize]) -> bool { .is_some_and(|count| count <= MAX_SHAPE_DATA_ELEMS) } +const MAX_COMPRESS_CONDITION_ELEMS: usize = 1 << 20; + +/// `Compress` is data-dependent on its full boolean condition, which can be +/// larger than ordinary scalar/vector shape metadata for image models. +pub(super) fn bounded_compress_condition(dtype: DataType, shape: &[usize]) -> bool { + dtype == DataType::Bool && shape.len() == 1 && shape[0] <= MAX_COMPRESS_CONDITION_ELEMS +} + /// Whether a node needs a float runtime input to resolve a data-dependent /// output extent. The list is deliberately explicit, so shape propagation never /// copies unrelated tensor data to host. diff --git a/crates/onnx-runtime-session/src/executor/tests.rs b/crates/onnx-runtime-session/src/executor/tests.rs index 466b96ab23..0b206274fc 100644 --- a/crates/onnx-runtime-session/src/executor/tests.rs +++ b/crates/onnx-runtime-session/src/executor/tests.rs @@ -3352,6 +3352,24 @@ fn dynamic_output_shapes_compress_counts_selected_values() { ); } +#[test] +fn compress_condition_allows_image_sized_boolean_vectors() { + let image_condition = MAX_SHAPE_DATA_ELEMS + 1; + assert!(bounded_compress_condition( + DataType::Bool, + &[image_condition] + )); + assert!(!bounded_shape_input(DataType::Bool, &[image_condition])); + assert!(!bounded_compress_condition( + DataType::Int64, + &[image_condition] + )); + assert!(!bounded_compress_condition( + DataType::Bool, + &[(1 << 20) + 1] + )); +} + /// The data-dependent shape sizer must return exactly one shape per output /// so the run loop's `out_shapes[oi]` indexing can never misindex. Slice is /// single-output, so it returns a 1-element Vec; the run loop additionally