Validate SpaceDepth shape arithmetic

[danielsongmicrosoft]

Description

Validates blocksize and uses overflow-safe shape arithmetic shared by SpaceToDepth and DepthToSpace. Removes the obsolete warning suppression and adds boundary tests.

Motivation and Context

Invalid or extreme block sizes could cause divide-by-zero or signed overflow during output-shape calculation.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Hariharan Seshadri (hariharans29)]

Review — PR #32039: Validate SpaceDepth shape arithmetic

Bug analysis — what was actually broken pre-PR

Two related issues in onnxruntime/core/providers/cpu/tensor/space_depth_ops.h's InputValidationsAndOutputDimsCalc, both explaining the C4723 warning suppression the PR removes:

1. blocksize == 0 → modulo-by-zero UB. No pre-check on blocksize. SpaceToDepth executes input_height % blocksize, DepthToSpace executes input_depth % (blocksize * blocksize) — both are x % 0 if blocksize == 0. Signed modulo-by-zero is UB; MSVC's C4723 static-detected the DepthToSpace one along at least one folded path (which is why the pragma sat over the DepthToSpace branch specifically).
2. blocksize * blocksize overflow in the DepthToSpace divisibility check. With any blocksize >= 2^32, the product blocksize * blocksize overflows int64 — signed overflow is UB, and if it happens to fold to 0 the check becomes x % 0. The pragma also covered this.

The negative-blocksize case is a related but distinct correctness concern: input_depth / blocksize etc. don't produce sensible dimensions with a negative blocksize, so barring it upfront is the right call.

The fix — three components

1. Early guard: if (blocksize <= 0) return INVALID_ARGUMENT(...) placed before any use of blocksize. Kills both the zero and negative paths. ✓ Placed correctly before the batch/input_depth assignments, so no partially-populated out-params (matches pre-PR early-return behavior for the rank check just above).
2. Divisibility rewrite for DepthToSpace: replaces input_depth % (blocksize * blocksize) != 0 with input_depth % blocksize != 0 || (input_depth / blocksize) % blocksize != 0. Provably equivalent for blocksize > 0: blocksize² | d iff blocksize | d ∧ blocksize | (d/blocksize). The rewrite never forms the product, so no overflow — the actual UB source is removed rather than masked. ✓
3. Overflow guards on output dim multiplication: pre-multiply comparisons against int64_max / blocksize. Short-circuit ordering is correct — the second clause input_depth * blocksize > int64_max / blocksize only evaluates after the first has confirmed input_depth * blocksize doesn't overflow. ✓

Pragma removal is the natural consequence: with the actual UB sources eliminated, MSVC has nothing to warn about.

Small correctness point on the guards

For SpaceToDepth:

if (input_depth > int64_max / blocksize || input_depth * blocksize > int64_max / blocksize) {
  return INVALID_ARGUMENT("SpaceToDepth output depth exceeds int64_t limits");
}
output_depth = input_depth * blocksize * blocksize;

I walked the algebra: this correctly guards input_depth * blocksize * blocksize. Second clause reads slightly odd — it's checking that (input_depth * blocksize) * blocksize won't overflow by asking whether (input_depth * blocksize) > int64_max / blocksize, which is the classical pre-multiply overflow test recursively applied. A one-line comment noting "guard both multiplications" would spare the next reader the derivation, but the logic is right.

For DepthToSpace:

if (input_height > int64_max / blocksize || input_width > int64_max / blocksize) {
  return INVALID_ARGUMENT("DepthToSpace output dimensions exceed int64_t limits");
}
output_height = input_height * blocksize;
output_width = input_width * blocksize;

Straightforward, correct.

Style nits (all minor)

• const auto int64_max = std::numeric_limits<int64_t>::max(); — should be constexpr. Or hoist out of the branch entirely to file scope. The current shape duplicates the constant across the two branches. Non-blocking.
• Message "SpaceToDepth output depth exceeds int64_t limits" — good. "DepthToSpace output dimensions exceed int64_t limits" — could distinguish which dimension (height vs width); low-value polish.
• #include <limits> correctly added. ✓
• New file dependency in the test: #include "core/framework/allocator.h" for CPUAllocator — appropriate.

Test coverage — where it lands, and where it doesn't

Test #1 RejectsNonPositiveBlocksizeBeforeShapeArithmetic: covers blocksize = 0 and blocksize = -1 across both branches (via the is_space_to_depth pairing). The blocksize <= 0 guard sits before any branch, so both parameter pairs hit the same code — the coverage is redundant across the two rows but harmless, and asserting the same substring for both is fine.

Test #2 RejectsBlocksizeWhoseSquareExceedsInt64: the test name doesn't match what it exercises. With blocksize = 1<<32, input_shape = {1, 1, 2, 2}, is_space_to_depth = false, it goes down the DepthToSpace branch:

• blocksize <= 0 → false, skipped.
• Divisibility rewrite: input_depth % blocksize = 1 % (1<<32) = 1 ≠ 0. First clause fires, returns INVALID_ARGUMENT("...multiple of (block_size * block_size)").
• The new overflow guards (input_height > int64_max / blocksize || input_width > int64_max / blocksize) are unreachable because the divisibility check errors out first.

So Test #2 actually covers the divisibility rewrite, not the overflow guards. That's still valuable — the rewrite is the piece that actually removes the pre-PR UB (integer overflow on blocksize * blocksize). But the name should reflect what it exercises, e.g., DepthToSpaceRejectsHugeBlocksizeViaDivisibilityCheck or similar.

Actual overflow guards are not exercised by any test. To hit them you'd need an input_depth (SpaceToDepth) or input_height/input_width (DepthToSpace) large enough that multiplication overflows int64, while simultaneously passing the divisibility checks. That requires constructing a Tensor with a shape dimension near int64_max, which the current test scaffolding can't do because it also allocates a data buffer.

Two reasonable ways to close this coverage gap — pick one:

• Refactor the shape-arithmetic helper to take individual dims (or a TensorShape without an owning Tensor), letting tests pass huge dims without allocation. Small API change but tightens testability of the overflow guards.
• Accept the gap and add a comment on the guard blocks noting they're validated by inspection (the arithmetic is small and the guards are the standard pre-multiply pattern).

Not blocking, but worth choosing consciously.

One additional gap — CRD vs DCR mode

ReadIsDCR and the two DepthToSpace layout modes aren't touched by the shape arithmetic — the divisibility and output-dim math is the same for both — so this PR is complete for the mode dimension. Just calling it out.

Broader observation (for follow-up, not this PR)

The blocksize² | d divisibility check pattern also shows up in the CUDA / DML kernels for the same ops. Worth grepping providers/cuda/tensor/space_to_depth* and DML equivalents for the same blocksize * blocksize pattern. If they use the same idiom without guards, they carry the same UB. Separate PR.

CI

1/1 check OK on a105028. The 1-check total is unusually low — probably means only a subset of pipelines ran (Copilot review pending, no /azp run yet). Would want a full CI run before merge just to be sure the pragma removal doesn't resurrect warnings on unusual toolchains (e.g., older MSVC on x86 release, which was the original trigger for the suppression).

Recommendation

Approve pending:

1. Rename Test Remove vsts test runner in cmake file #2 to match what it actually tests (the divisibility rewrite path), or add a comment inside the test noting that. Non-code change.
2. A full CI run against MSVC x86 release to confirm the pragma removal doesn't reintroduce C4723. (Original suppression was _MSC_VER-specific — the x86 release build was the historical trigger.)

Nice to have (fast follow-up, not blocking):

3. Consider refactoring the shape helper to accept individual dims so the overflow guards themselves can be exercised.
4. constexpr auto int64_max hoisted before the is_space_to_depth branch to avoid duplication.

The core change is small, precise, and eliminates real UB rather than papering over it with a warning suppression. Divisibility rewrite is correct-and-safe. Overflow guards are the standard idiom applied correctly. Test #1 is on the money.