ONNX Runtime Shape Inference Safety Audit

Contrib-Op Shape-Inference Output-Index Safety

Reusable method for finding and fixing the bug class where an operator's TypeAndShapeInference function guards an output write with getNumOutputs() > N but then writes an output index greater than N. For a node that declares fewer outputs, the written index is past the end of the inference context's output vector.

Scope: schema-level shape inference in onnxruntime/core/graph/contrib_ops/*.cc and shape_inference_functions.cc. This runs once during Graph::Resolve (model-load time), EP-agnostic - there is no per-EP (CPU/CUDA/ROCm) kernel duplicate of this code to chase. Op kernels allocate outputs via the bounds-safe OpKernelContext::Output(index) and are a separate concern.

1. The pattern

// onnxruntime/core/graph/contrib_ops/bert_defs.cc  (before)
propagateElemTypeFromInputToOutput(ctx, 0, 0);
if (ctx.getNumOutputs() > 1) {                       // guard says "> 1"
  propagateElemTypeFromInputToOutput(ctx, 0, 1);
  propagateElemTypeFromInputToOutput(ctx, 0, 2);     // but writes index 2
}

The guard getNumOutputs() > 1 admits a node with exactly 2 outputs (indices 0, 1), yet the body writes index 2. The implication "> 1 ⇒ index 2 exists" is false: > 1 only guarantees indices 0 and 1.

Why a smaller output count is valid

Trailing outputs declared OpSchema::Optional lower min_output. ONNX derives min_output = number of required outputs, max_output = total declared. The model checker (checker::check_node) only enforces min_output <= N <= max_output.

So a node with output=['out','present_key'] is schema-valid, passes the checker, and then reaches the index-2 write. A passing checker is not a guarantee the index is in range.

2. The sink (why the write is not caught)

// onnxruntime/core/graph/graph.cc  -  InferenceContextImpl
const TypeProto* getInputType(size_t index) const override {
  return node_.InputDefs().at(index)->TypeAsProto();   // .at()  -> bounds-checked
}
TypeProto* getOutputType(size_t index) override {
  return &node_output_types_[index];                   // operator[]  -> NOT bounds-checked
}

• node_output_types_ is sized to node.OutputDefs().size() in the InferenceContextImpl ctor, so for a 2-output node it has 2 elements; getOutputType(2) returns one past the end.
• getInputType uses .at() (would throw on a bad index); getOutputType uses raw operator[] (no check) - the asymmetry is the root cause.
• The call runs at Graph::Resolve → InferAndVerifyTypeMatch → RunInferencing. The surrounding ORT_TRY/ORT_CATCH(const std::exception&) only catches thrown fail_shape_inference; a raw out-of-range operator[] does not throw, so the catch does not help.
• Because this is schema-level inference, it is EP-independent - no CUDA/ROCm copy.

3. Audit technique — always sweep siblings

Do not stop at the reported function. Grep every shape-inference guard and compare its threshold against the highest output index written before the next guard.

git grep -n 'getNumOutputs' -- \
  onnxruntime/core/graph/contrib_ops/*.cc \
  onnxruntime/core/graph/contrib_ops/shape_inference_functions.cc

For each if (ctx.getNumOutputs() > N) block, find the largest index passed to propagateElemTypeFromInputToOutput(ctx, _, index) / updateOutputShape(ctx, index, _) / getOutputType(index) inside it. Rule: the guard must require strictly more outputs than the highest index written (write index k ⇒ guard must ensure getNumOutputs() > k).

Correct exemplars already in the tree to copy:

Gotcha — conditional writes can hide a vacuous audit. A write may sit behind an inner condition (e.g. hasInputShape(past_key_index) before writing index 2). The site is still a bug, but you can only observe it when that inner condition is also satisfied. Keep this in mind both for the audit and for tests (§5).

4. Fix patterns

Point fix (required): raise the guard to cover the highest index written.

// before
if (ctx.getNumOutputs() > 1) { ... writes idx 2 ... }
// after
if (ctx.getNumOutputs() > 2) {  // both present_key (idx 1) AND present_value (idx 2)
  ...
}

Justify the threshold with the op's output semantics. For these attention ops the two trailing outputs - present_key (idx 1) and present_value (idx 2) for MultiHeadAttention, new_key_cache / new_value_cache for DecoderAttention (see the §1 table for each op's exact output names) - are a both-or-neither pair: there is no valid configuration that emits one without the other, so requiring all three outputs before populating indices 1 and 2 is behavior-preserving. (PagedAttention encodes the same invariant via its nested != 3 check.)

Defense-in-depth (recommended): bound the sink so a future author cannot reintroduce the class.

// onnxruntime/core/graph/graph.cc  -  InferenceContextImpl::getOutputType
TypeProto* getOutputType(size_t index) override {
  if (index >= node_output_types_.size()) {
    fail_type_inference("output index ", index, " is out of range; node has ",
                        node_output_types_.size(), " outputs");
  }
  return &node_output_types_[index];
}

This mirrors getInputType's .at() and the existing bounds checks in the sibling DataPropagationContextImpl. Placing it at the base layer transitively protects the NHWC and quantization wrapper contexts. After the point fix this branch is unreachable through a normal model (the guard already prevents the out-of-range index), so it is pure defense-in-depth. Its failure mode is build-dependent: with exceptions enabled, fail_type_inference raises InferenceError (a std::exception), caught by the existing ORT_CATCH(const std::exception&) around RunInferencing and surfaced as a clean load-time error; under ORT_NO_EXCEPTIONS it is not compiled out - ONNX's no-exceptions path prints the message to std::cerr and calls abort(), a deterministic fail-fast (consistent with getInputType's .at(), which likewise terminates under no-exceptions). Either way the result is a controlled failure rather than an out-of-range write.