Add check for negative index in cpu kernel op_index

Fast path cannot be taken when the index is negative in op_index. It
results in a failure in `check_fast_path_args`. Add a check to
`check_fast_path_conditions` for negative index and return `false` if
one is encountered.

Change-Id: I6c3fb1793da9a0831ef41a47dc2ab4423aa38be3
Signed-off-by: Sebastian Larsson <[email protected]>

Author: Sebastian-Larsson

kernels/portable/cpu/op_index.cpp

@@ -1,6 +1,7 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
+ * Copyright 2025 Arm Limited and/or its affiliates.
  *
  * This source code is licensed under the BSD-style license found in the
  * LICENSE file in the root directory of this source tree.
@@ -47,6 +48,23 @@ bool check_fast_path_conditions(
       if (index.dim() != 1) {
         return false;
       }
+
+      // Fast path only supports non-negative indices.
+      if (ix_type == ScalarType::Int) {
+        const int32_t* const data = index.const_data_ptr<int32_t>();
+        for (const auto j : c10::irange(index.numel())) {
+          if (data[j] < 0) {
+            return false;
+          }
+        }
+      } else { // ScalarType::Long
+        const int64_t* const data = index.const_data_ptr<int64_t>();
+        for (const auto j : c10::irange(index.numel())) {
+          if (data[j] < 0) {
+            return false;
+          }
+        }
+      }
     }
   }