Pre-check ConstantOfShape output size against input initializer before constant folding

Merged
CO
5mo ago
+168-2

A 152-byte ONNX model with a ConstantOfShape whose shape initializer encodes huge dims causes InferenceSession::Initialize() to materialize the full output tensor (287 MB in the PoC, up to PB-scale with fuzz-mutated dims) via the ConstantFolding optimizer. The existing pre-execution size cap in EstimateNodeOutputSizeInBytes relies on shape inference having populated output_def->Shape(), which is not guaranteed.

Description

  • onnxruntime/core/optimizer/constant_folding.cc

    • New EstimateConstantOfShapeOutputSizeInBytes(node, graph): looks up the shape input via Graph::GetConstantInitializer (it is constant by the time we reach this node), multiplies its int64 values with SafeInt<int64_t> (rejects negative dims, lets overflow propagate as an exception caught upstream), and multiplies by the element size derived from the value attribute's tensor type (defaulting to float per ONNX spec).
    • EstimateNodeOutputSizeInBytes now takes const Graph& and dispatches to the new estimator for ConstantOfShape, falling back to the generic shape-based path if the initializer can't be resolved.

  • onnxruntime/test/optimizer/graph_transform_test.cc

    • ConstantFoldingConstantOfShapeUsesInputInitializerForSizeCheck: shape [100M] with int64 value ⇒ 800 MB derived size; with kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes=256MB the node must remain unfolded, proving the size check fires from the initializer alone.
    • ConstantFoldingConstantOfShapeBlockedWhenOutputShapeMissing: same model, but the pre_graph_checker (which runs after Graph::Resolve() and before the transformer) calls ClearShape() on the ConstantOfShape output NodeArg to simulate the documented attack where shape inference has not propagated the output shape. With the inferred shape stripped, the generic shape-based estimator returns -1, so only the new EstimateConstantOfShapeOutputSizeInBytes path can derive the 800 MB size and block folding — isolating regressions in the new estimator from the pre-existing shape-inference path.

Motivation and Context

The byte cap added in #28055 only triggers when shape inference has propagated the output shape; a crafted model can bypass it and force unbounded allocation during Initialize(). Deriving the size from the (necessarily constant) shape input makes the cap effective for the documented attack vector and tightens the same code path used by the configurable kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes setting — no new knob, no behavior change for legitimate models within the existing 1 GB default.