Repository navigation
Use checked arithmetic for GATHER shape products - #10846
Thanatos9404 wants to merge 1 commit into
Conversation
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.
|
Per our security guidelines 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}. 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 |
|
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. |
Summary
reference_ops::Gatherviews 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, withint. The GATHER kernel did not check before Eval that any of these products fit. A tensor with a zero dimension has no elements, soBytesRequiredaccepts it and nothing is allocated. Its other dimensions can still multiply pastINT_MAXbefore 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
ValidateGatherFlattenedShapes, called fromPrepareafter the existing axis andbatch_dimschecks. It checks the batch, outer and inner sizes of the input, the flat input size, and the size of positions pastbatch_dims. It useskernel_utils::CheckedDimensionProduct, the existing helper that ARG_MIN and ARG_MAX use. The helper accumulates inCheckedInt<int>and fails as soon as a partial product overflows, so a later zero dimension does not hide the overflow.ResizeTensor,Preparechecks the flat output size the same way, becausereference_ops::Gatheralso computes it withRuntimeShape::FlatSize(). The constant-input path (noop) and the regular path now share this code. The constant path evaluates the op duringPrepare, 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 (seeIntermediateInnerProductOverflowRejectedinarg_min_max_test.cc).Tests
New tests in
gather_test.cc:GatherPrepareTest.Rejects{Inner,Outer,Positions,Output}SizeOverflowWithZeroElementsuse zero-element shapes in which one of the products that Gather computes is 2^31, one more thanINT_MAX. The shapes are chosen so that the product fits a 32-bitsize_t, so the overflow is in Gather's ownintarithmetic and not in the element count check. The positions and output cases use shapes where only that one product overflows. Each test expectsAllocateTensors()to fail.GatherPrepareTest.ZeroElementsWithLargeDimensionStillWorkchecks that{0, INT_MAX}is still accepted and invoked, and that the output shape is unchanged.PrepareOnlyGatherOpModelnow takes optional input and positions shapes. Existing call sites are unchanged.Testing
CMake build with
-DTFLITE_KERNEL_TEST=ONon x86_64 Linux (GCC 13.3, Release):gather_test: 70 of 70 tests pass.gather.ccrestored to the base version and the new tests kept, the 4 rejection tests fail becauseAllocateTensors()returnskTfLiteOk. The other 66 tests pass, includingZeroElementsWithLargeDimensionStillWork.