From db069c0a6d304be5c58293774d05faf65d191d64 Mon Sep 17 00:00:00 2001
From: Akshay Sonawane <asonawane@microsoft.com>
Date: Tue, 11 Aug 2026 23:00:58 -0700
Subject: [PATCH] Handle extreme diagonal values in CPU Trilu

---
 .../core/providers/cpu/tensor/trilu.cc        | 16 +++++++++++
 .../providers/cpu/tensor/trilu_op_test.cc     | 27 ++++++++++++++++---
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/onnxruntime/core/providers/cpu/tensor/trilu.cc b/onnxruntime/core/providers/cpu/tensor/trilu.cc
index dad84f3f92f5a..1d2bb39be2efe 100644
--- a/onnxruntime/core/providers/cpu/tensor/trilu.cc
+++ b/onnxruntime/core/providers/cpu/tensor/trilu.cc
@@ -62,6 +62,14 @@ static Status TriluImpl(const Tensor* X, Tensor* Y, int64_t k_val, bool up) {
     }
 
     if (up) {
+      if (k_val <= 1 - matrix_h) {
+        continue;
+      }
+      if (k_val >= matrix_w) {
+        output_mat.setZero();
+        continue;
+      }
+
       int64_t start_i = k_val > 0 ? 0 : 1 - k_val;
       for (int64_t i = start_i; i < matrix_h; i++) {
         for (int64_t j = 0; j < i + k_val && j < matrix_w; j++) {
@@ -69,6 +77,14 @@ static Status TriluImpl(const Tensor* X, Tensor* Y, int64_t k_val, bool up) {
         }
       }
     } else {
+      if (k_val >= matrix_w - 1) {
+        continue;
+      }
+      if (k_val <= -matrix_h) {
+        output_mat.setZero();
+        continue;
+      }
+
       int64_t end_i = std::min(matrix_h, matrix_w - k_val);
       for (int64_t i = 0; i < end_i; i++) {
         for (int64_t j = std::max(static_cast<int64_t>(0), i + k_val + 1); j < matrix_w; j++) {
diff --git a/onnxruntime/test/providers/cpu/tensor/trilu_op_test.cc b/onnxruntime/test/providers/cpu/tensor/trilu_op_test.cc
index 81dd2aedf5592..179b1d9275536 100644
--- a/onnxruntime/test/providers/cpu/tensor/trilu_op_test.cc
+++ b/onnxruntime/test/providers/cpu/tensor/trilu_op_test.cc
@@ -1,13 +1,12 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <limits>
+
 #include "gtest/gtest.h"
 #include "test/providers/provider_test_utils.h"
 #include "core/util/math.h"
-
-#ifdef USE_WEBGPU
 #include "test/util/include/default_providers.h"
-#endif
 
 namespace onnxruntime {
 namespace test {
@@ -30,6 +29,28 @@ TEST(TriluOpTest, two_by_two_float_lower) {
   test.Run();
 }
 
+static void TestExtremeDiagonal(bool upper, int64_t k, const std::vector<float>& expected) {
+  OpTester test("Trilu", 14, kOnnxDomain);
+  test.AddAttribute<int64_t>("upper", upper ? 1 : 0);
+  test.AddInput<float>("X", {2, 2}, {4.f, 7.f, 2.f, 6.f});
+  test.AddInput<int64_t>("k", {}, {k});
+  test.AddOutput<float>("Y", {2, 2}, expected);
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCpuExecutionProvider());
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+}
+
+TEST(TriluOpTest, ExtremeNegativeDiagonal) {
+  TestExtremeDiagonal(true, std::numeric_limits<int64_t>::min(), {4.f, 7.f, 2.f, 6.f});
+  TestExtremeDiagonal(false, std::numeric_limits<int64_t>::min(), {0.f, 0.f, 0.f, 0.f});
+}
+
+TEST(TriluOpTest, ExtremePositiveDiagonal) {
+  TestExtremeDiagonal(true, std::numeric_limits<int64_t>::max(), {0.f, 0.f, 0.f, 0.f});
+  TestExtremeDiagonal(false, std::numeric_limits<int64_t>::max(), {4.f, 7.f, 2.f, 6.f});
+}
+
 TEST(TriluOpTest, two_by_two_double_upper) {
   OpTester test("Trilu", 14, kOnnxDomain);
   test.AddInput<double>("X", {2, 2}, {4, 7, 2, 6});