ICM fixes (5/n)

[hariharans29]

Description

Fixes 3 ICM fixes:

https://portal.microsofticm.com/imp/v5/incidents/details/31000000572208/summary
https://portal.microsofticm.com/imp/v5/incidents/details/31000000573313/summary
https://portal.microsofticm.com/imp/v5/incidents/details/31000000575583/summary

Motivation and Context

Fix ICM issues

[Copilot]

Pull request overview

This PR addresses multiple ICM-reported robustness issues by adding missing input/attribute validation to CPU ML/math kernels and strengthening overflow handling in RotaryEmbedding, with new negative tests to prevent regressions.

Changes:

• Add validation for LinearRegressor input rank and coefficient vector sizing, plus a regression test for undersized coefficients.
• Tighten CumSum axis tensor validation to reject empty/multi-element axis inputs, plus a regression test.
• Add overflow checks in RotaryEmbedding (ONNX + contrib) and add tests for overflow/invalid sizing scenarios.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
onnxruntime/core/providers/cpu/ml/linearregressor.cc	Adds extra input/attribute validation and coefficient-size verification.
onnxruntime/test/providers/cpu/ml/linearregressor_test.cc	Adds failure test for undersized coefficients.
onnxruntime/core/providers/cpu/math/cumsum.h	Rejects axis tensors whose element-count is not exactly 1.
onnxruntime/test/providers/cpu/math/cumsum_test.cc	Adds failure test for empty axis tensor.
onnxruntime/core/providers/cpu/llm/rotary_embedding.cc	Adds ptrdiff overflow-checked loop bound calculations and safer iteration types.
onnxruntime/core/providers/cpu/llm/rotary_embedding_helper.h	Adds int32 narrowing/checked-multiply utilities and uses them during input validation.
onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc	Adds regression tests for overflow/shape-validation behavior.
onnxruntime/contrib_ops/cpu/bert/rotary_embedding.cc	Same overflow-checked loop bound updates for contrib RotaryEmbedding.
onnxruntime/contrib_ops/cpu/bert/rotary_embedding_helper.h	Same int32 narrowing/checked-multiply utilities for contrib input validation.
onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc	Adds contrib regression tests for overflow/shape-validation behavior.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 7 comments.

Comments suppressed due to low confidence (2)

onnxruntime/core/providers/cpu/llm/rotary_embedding.cc:120

• cost is computed using int arithmetic (head_size * sizeof(T) * 2 + rotary_emb_dim * 32) before being cast to double, which can overflow for large head_size/rotary_emb_dim and potentially yield a negative/incorrect cost passed to TryParallelFor. Consider promoting operands to double/int64_t before multiplying/adding so the cost calculation itself can’t overflow.

  // The cost is calculated as:
  //   - head_size * sizeof(T) for reading input
  //   - head_size * sizeof(T) for writing output
  //   - rotary_emb_dim * 32 for the rotary embedding operations (32 is an approximation of the number of CPU cycles)
  const double cost = static_cast<double>(head_size * sizeof(T) * 2 + rotary_emb_dim * 32);

onnxruntime/contrib_ops/cpu/bert/rotary_embedding.cc:142

• cost is computed with int arithmetic (head_size * sizeof(T) * 2 + rotary_emb_dim * 32) before converting to double, so it can overflow for large shapes and pass an incorrect/negative cost into TryParallelFor. Consider promoting operands to double/int64_t before the arithmetic to avoid overflow in the cost computation itself.

  // The cost is calculated as:
  //   - head_size * sizeof(T) for reading input
  //   - head_size * sizeof(T) for writing output
  //   - rotary_emb_dim * 32 for the rotary embedding operations (32 is an approximation of the number of CPU cycles)
  const double cost = static_cast<double>(head_size * sizeof(T) * 2 + rotary_emb_dim * 32);

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 6 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 4 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[tianleiwu]

Review Summary

Solid hardening PR — the checked arithmetic approach (NarrowNonNegativeToInt32, CheckedMulToInt32, CheckedMulToPtrdiff) is thorough and well-structured. CumSum empty-axis fix and LinearRegressor validations are clean. Test coverage is good across all three fix areas.

One high-priority concern: the contrib RotaryEmbedding helper has a potential division-by-zero that the llm variant correctly guards against (see inline comment).

[tianleiwu]

⚠️ Division-by-zero when num_heads == 0: This hidden_size % num_heads can be undefined behavior if num_heads is 0 (its default attribute value). The equivalent code in the llm variant (core/providers/cpu/llm/rotary_embedding_helper.h) correctly adds a num_heads <= 0 guard before any division:

if (num_heads <= 0) {
  return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                         "RotaryEmbedding: num_heads must be greater than 0 for rank-3 input");
}

The same guard should be added here before the modulo check, for rank-3 input when rotary_embedding_dim > 0.

[tianleiwu]

How about using SafeInt directly

std::ptrdiff_t position_count = SafeInt<std::ptrdiff_t>(batch_size) * sequence_length;

SafeInt will call ORT_THROW if there is overflow.