Skip to content

Use checked arithmetic for GATHER shape products - #10846

Closed
Thanatos9404 wants to merge 1 commit into
google-ai-edge:mainfrom
Thanatos9404:checked-arithmetic-gather
Closed

Thanatos9404 wants to merge 1 commit into
google-ai-edge:mainfrom
Thanatos9404:checked-arithmetic-gather

Conversation

@Thanatos9404

Copy link
Copy Markdown

Summary

reference_ops::Gather views the input as [batch, outer, axis, inner] and the positions as [batch, coord]. It computes its batch, outer, inner and flat sizes, and its copy offsets, with int. The GATHER kernel did not check before Eval that any of these products fit. A tensor with a zero dimension has no elements, so BytesRequired accepts it and nothing is allocated. Its other dimensions can still multiply past INT_MAX before the reference code reaches the zero. For example, an input of shape {0, 65536, 32768} gathered on axis 0 gives an inner size of 2^31.

GATHER is on the in-scope operator list in docs/instructions/security.md. This PR follows that document's guidance: use checked arithmetic for every product that feeds a shape, loop bound or offset, and do not let a zero dimension short-circuit that validation.

Changes

  • New ValidateGatherFlattenedShapes, called from Prepare after the existing axis and batch_dims checks. It checks the batch, outer and inner sizes of the input, the flat input size, and the size of positions past batch_dims. It uses kernel_utils::CheckedDimensionProduct, the existing helper that ARG_MIN and ARG_MAX use. The helper accumulates in CheckedInt<int> and fails as soon as a partial product overflows, so a later zero dimension does not hide the overflow.
  • After ResizeTensor, Prepare checks the flat output size the same way, because reference_ops::Gather also computes it with RuntimeShape::FlatSize(). The constant-input path (noop) and the regular path now share this code. The constant path evaluates the op during Prepare, and it now does so only after the check.

Valid models are not affected. The only newly rejected inputs are zero-element shapes whose other dimensions multiply past INT_MAX. ARG_MIN and ARG_MAX already reject the same kind of shape (see IntermediateInnerProductOverflowRejected in arg_min_max_test.cc).

Tests

New tests in gather_test.cc:

  • GatherPrepareTest.Rejects{Inner,Outer,Positions,Output}SizeOverflowWithZeroElements use zero-element shapes in which one of the products that Gather computes is 2^31, one more than INT_MAX. The shapes are chosen so that the product fits a 32-bit size_t, so the overflow is in Gather's own int arithmetic and not in the element count check. The positions and output cases use shapes where only that one product overflows. Each test expects AllocateTensors() to fail.
  • GatherPrepareTest.ZeroElementsWithLargeDimensionStillWork checks that {0, INT_MAX} is still accepted and invoked, and that the output shape is unchanged.

PrepareOnlyGatherOpModel now takes optional input and positions shapes. Existing call sites are unchanged.

Testing

CMake build with -DTFLITE_KERNEL_TEST=ON on x86_64 Linux (GCC 13.3, Release):

  • gather_test: 70 of 70 tests pass.
  • With gather.cc restored to the base version and the new tests kept, the 4 rejection tests fail because AllocateTensors() returns kTfLiteOk. The other 66 tests pass, including ZeroElementsWithLargeDimensionStillWork.

reference_ops::Gather views the input as [batch, outer, axis, inner]
and the positions as [batch, coord]. It computes its batch, outer,
inner and flat sizes and its copy offsets with int, and Prepare did not
check that any of these products fit. A tensor with a zero dimension has
no elements, so it passes the size_t element count check in
BytesRequired. The other dimensions can still multiply past INT_MAX
before the kernel reaches the zero.

Prepare now checks the batch, outer and inner sizes of the input, the
flat input size, and the size of positions past batch_dims with
kernel_utils::CheckedDimensionProduct, the same helper used for ARG_MIN
and ARG_MAX. The helper accumulates in CheckedInt<int> and fails as soon
as a partial product overflows, so a later zero dimension does not hide
the overflow. After the output is resized, Prepare checks the flat
output size the same way. The constant-input path, which evaluates the
op during Prepare, now runs after this check as well.

New tests cover zero-element shapes where the inner, outer, positions
or output product is 2^31, which is one more than INT_MAX, and a
zero-element shape with a large dimension that must keep working.
@snnn

snnn commented Oct 14, 2026

Copy link
Copy Markdown
Member

Per our security guidelines
: "If a sanitization check can be performed in the graph verifier (including shape inference), it must be done there instead of being duplicated across individual TFLite delegates or per-op kernels."

Because GATHER's sub-shapes (batch, outer, inner, coord) and output shape depend only on the input/positions tensor shapes and static GatherOptions (axis, batch_dims) rather than runtime tensor values, these 32-bit dimension-product checks can be enforced once at model ingress in litert::Verify (litert/core/model/verifier.cc) and InferGather (litert/core/model/ops/gather.h) rather than in tflite/kernels/gather.cc:

In litert/core/model/verifier.cc (ValidateFlatBufferTensors and ValidateLiteRtTensor), checking that the product of all positive dimensions (dim > 1) of every tensor fits in int32_t (using SafeMultiply) guarantees that every sub-range product (batch_size, outer_size, inner_size, coord_size, input_size) and the inferred output tensor product fit in int32_t, regardless of where a 0 dimension appears, while still allowing {0, INT_MAX}.
In litert/core/model/ops/gather.h, InferGather should be updated to validate batch_dims and slice indices_shape[batch_dims:] when inferring output_shapes[0].

Enforcing this in the verifier avoids adding extra checks to the TFLite kernel (preserving backwards compatibility for trusted-model callers) and protects all backends/delegates simultaneously. We are adding extra checks to the verifier to cover this. The PR are stuck because there are some build break in Tensorflow CI build pipelines. We are waiting TF repo to recover

@snnn snnn closed this Oct 14, 2026
@Thanatos9404

Copy link
Copy Markdown
Author

Thanks @snnn for the pointer to the security guidelines. That makes sense: checks that only depend on shapes and static options belong in the graph verifier and shape inference, not in each kernel.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants