Add check for negative index in cpu kernel op_index (#15366)

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.

cc @freddan80 @per @zingo @oscarandersson8218 @digantdesai

Signed-off-by: Sebastian Larsson <[email protected]>
(cherry picked from commit d6de876)

Author: Sebastian-Larsson

kernels/portable/cpu/op_index.cpp

@@ -1,11 +1,13 @@
 /*
  * 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.
  */
 
+#include <algorithm>
 #include <cinttypes>
 #include <cstdint>
 #include <cstring>
@@ -47,6 +49,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>();
+        if (std::any_of(data, data + index.numel(), [](const auto x) {
+              return x < 0;
+            })) {
+          return false;
+        }
+      } else { // ScalarType::Long
+        const int64_t* const data = index.const_data_ptr<int64_t>();
+        if (std::any_of(data, data + index.numel(), [](const auto x) {
+              return x < 0;
+            })) {
+          return false;
+        }
+      }
     }
   }
 

kernels/test/op_index_test.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.
@@ -480,6 +481,36 @@ TEST_F(OpIndexTensorOutTest, AllDtypesSupportedForIndex) {
   test_dtype<ScalarType::Double, ScalarType::Int, ScalarType::Double>();
 }
 
+TEST_F(OpIndexTensorOutTest, NegativeIndexSupportedForLong) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+
+  Tensor x = tf.make({3}, {1., 2., 3.});
+  Tensor out = tf.zeros({1});
+  Tensor expected = tf.make({1}, {3.});
+
+  std::array<optional<Tensor>, 1> indices = {
+      optional<Tensor>(tfl.make({1}, {-1}))};
+
+  Tensor ret = op_index_tensor_out(x, indices, out);
+  EXPECT_TENSOR_EQ(ret, expected);
+}
+
+TEST_F(OpIndexTensorOutTest, NegativeIndexSupportedForInt) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Int> tfi;
+
+  Tensor x = tf.make({3}, {1., 2., 3.});
+  Tensor out = tf.zeros({1});
+  Tensor expected = tf.make({1}, {3.});
+
+  std::array<optional<Tensor>, 1> indices = {
+      optional<Tensor>(tfi.make({1}, {-1}))};
+
+  Tensor ret = op_index_tensor_out(x, indices, out);
+  EXPECT_TENSOR_EQ(ret, expected);
+}
+
 //
 // Death Tests
 //