Skip to content

Commit b429dff

Browse files
Copilottianleiwu
andauthored
Fix Inverse kernel rank underflow before indexing trailing dims (microsoft#28400)
### Description `Inverse::Compute` (CPU and CUDA) read `dims[num_dim - 2]` and `dims[num_dim - 1]` before validating `num_dim >= 2`. Because `num_dim` is `size_t`, scalar or 1D inputs silently underflow the subtraction and produce out-of-bounds indices. - **`contrib_ops/cpu/inverse.cc`** — add `ORT_RETURN_IF_NOT(num_dim >= 2, ...)` immediately after shape retrieval, before any dimension indexing. - **`contrib_ops/cuda/inverse.cc`** — same guard in `ComputeInternal`. - **`test/contrib_ops/inverse_test.cc`** — add `scalar_input_fails` and `one_dim_input_fails` test cases. ### Motivation and Context Without the guard, passing a 0-D or 1-D tensor to the Inverse op causes `size_t` underflow on `num_dim - 2`, reading garbage memory before any validation occurs. The fix is a defense-in-depth kernel-level check that mirrors the existing ONNX shape-inference rejection. Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com> Co-authored-by: Tianlei Wu <tlwu@microsoft.com>
1 parent eb63cc6 commit b429dff

3 files changed

Lines changed: 19 additions & 0 deletions

File tree

onnxruntime/contrib_ops/cpu/inverse.cc

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,8 @@ Status Inverse::Compute(OpKernelContext* ctx) const {
7070
const auto num_dim = input_shape.NumDimensions();
7171
auto* output = ctx->Output(0, input_shape);
7272

73+
ORT_RETURN_IF_NOT(num_dim >= 2, "Input tensor rank must be >= 2, got: ", num_dim);
74+
7375
int64_t num_batches = 1;
7476
const int64_t rows = input_shape.GetDims()[num_dim - 2];
7577
const int64_t cols = input_shape.GetDims()[num_dim - 1];

onnxruntime/contrib_ops/cuda/inverse.cc

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,8 @@ Status Inverse::ComputeInternal(OpKernelContext* ctx) const {
142142
const auto num_dim = input_shape.NumDimensions();
143143
auto* output = ctx->Output(0, input_shape);
144144

145+
ORT_RETURN_IF_NOT(num_dim >= 2, "Input tensor rank must be >= 2, got: ", num_dim);
146+
145147
size_t num_batches = 1;
146148
const size_t rows = static_cast<size_t>(input_shape.GetDims()[num_dim - 2]);
147149
const size_t cols = static_cast<size_t>(input_shape.GetDims()[num_dim - 1]);

onnxruntime/test/contrib_ops/inverse_test.cc

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,5 +86,20 @@ TEST(InverseContribOpTest, four_by_four_batches_float) {
8686
test.AddOutput<float>("Output", {3, 4, 4, 4}, output);
8787
test.Run();
8888
}
89+
90+
TEST(InverseContribOpTest, scalar_input_fails) {
91+
OpTester test("Inverse", 1, kMSDomain);
92+
test.AddInput<float>("X", {}, {4.f});
93+
test.AddOutput<float>("Y", {}, {0.25f});
94+
test.Run(OpTester::ExpectResult::kExpectFailure, "rank must be >= 2");
95+
}
96+
97+
TEST(InverseContribOpTest, one_dim_input_fails) {
98+
OpTester test("Inverse", 1, kMSDomain);
99+
test.AddInput<float>("X", {4}, {4.f, 7.f, 2.f, 6.f});
100+
test.AddOutput<float>("Y", {4}, {0.f, 0.f, 0.f, 0.f});
101+
test.Run(OpTester::ExpectResult::kExpectFailure, "rank must be >= 2");
102+
}
103+
89104
} // namespace test
90105
} // namespace onnxruntime

0 commit comments

Comments
 (0)