Handle extreme diagonal values in CPU Trilu

onnxruntime/core/providers/cpu/tensor/trilu.cc

@@ -62,13 +62,29 @@ 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++) {
           output_mat(onnxruntime::narrow<std::ptrdiff_t>(i), onnxruntime::narrow<std::ptrdiff_t>(j)) = 0;
         }
       }
     } 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++) {

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});