Fix reduction over dim list for empty input

Summary:
Allow empty input in `MapReduceOverDimListPlan`
Avoid integer division by zero in `parallel_for_each_reduce_over_dim_list_output_index`
Add empty input tests for ops: any, mean, sum & var

Differential Revision: D81383049

Author: manuelcandales

kernels/portable/cpu/util/reduce_util.h

@@ -531,7 +531,6 @@ class MapReduceOverDimListPlan {
       const executorch::aten::Tensor& in,
       const std::optional<executorch::aten::ArrayRef<int64_t>>& dim_list)
       : plan_(in, dim_list, 1, -1) {
-    ET_CHECK_MSG(in.numel() > 0, "Input tensor must be nonempty");
   }
 
   template <
@@ -834,10 +833,12 @@ template <typename Func>
     const Func& func) {
 #ifdef ET_USE_THREADPOOL
   const ssize_t reduction_size = get_reduced_dim_product(in, dim_list);
-  const auto grain_size = std::max(
-      static_cast<ssize_t>(1),
-      static_cast<ssize_t>(executorch::extension::internal::GRAIN_SIZE) /
-          reduction_size);
+  const auto grain_size = reduction_size == 0
+      ? 1
+      : std::max(
+            static_cast<ssize_t>(1),
+            static_cast<ssize_t>(executorch::extension::internal::GRAIN_SIZE) /
+                reduction_size);
 #else // ET_USE_THREADPOOL
   const auto grain_size = 1;
 #endif // ET_USE_THREADPOOL

kernels/test/op_any_test.cpp

@@ -148,3 +148,31 @@ TEST_F(OpAnyOutTest, SmokeTest) {
   op_any_out(self, dim, keepdim, out);
   EXPECT_TENSOR_CLOSE(out, out_expected);
 }
+
+TEST_F(OpAnyOutTest, EmptyInput) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Bool> tfBool;
+
+  Tensor x = tf.make({2, 0, 3}, {});
+  optional<ArrayRef<int64_t>> dim_list = ArrayRef<int64_t>{};
+  Tensor out = tfBool.make({2, 0, 3}, {});
+
+  op_any_dims_out(x, dim_list, /*keepdim=*/true, out);
+  EXPECT_TENSOR_CLOSE(out, tfBool.zeros({2, 0, 3}));
+
+  out = tfBool.ones({2, 0, 3});
+  op_any_dims_out(x, dim_list, /*keepdim=*/false, out);
+  EXPECT_TENSOR_CLOSE(out, tfBool.zeros({2, 0, 3}));
+
+  int64_t dims1[1] = {1};
+  dim_list = ArrayRef<int64_t>{dims1, 1};
+  out = tfBool.ones({2, 3});
+  op_any_dims_out(x, dim_list, /*keepdim=*/false, out);
+  EXPECT_TENSOR_CLOSE(out, tfBool.zeros({2, 3}));
+
+  int64_t dims2[1] = {2};
+  dim_list = ArrayRef<int64_t>{dims2, 1};
+  out = tfBool.make({2, 0, 1}, {});
+  op_any_dims_out(x, dim_list, /*keepdim=*/true, out);
+  EXPECT_TENSOR_CLOSE(out, tfBool.make({2, 0, 1}, {}));
+}

kernels/test/op_mean_test.cpp

@@ -551,3 +551,30 @@ TEST_F(OpMeanOutTest, DTypeOutFloatNAN) {
   Tensor ret = op_mean_dtype_out(x, ScalarType::Float, out);
   EXPECT_TENSOR_CLOSE(out, expected_result);
 }
+
+TEST_F(OpMeanOutTest, EmptyInput) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor x = tf.make({2, 0, 3}, {});
+  optional<ScalarType> dtype = ScalarType::Float;
+  optional<ArrayRef<int64_t>> dim_list = ArrayRef<int64_t>{};
+  Tensor out = tf.zeros({1, 1, 1});
+  op_mean_out(x, dim_list, /*keepdim=*/true, dtype, out);
+  EXPECT_TENSOR_CLOSE(out, tf.make({1, 1, 1}, {NAN}));
+
+  out = tf.zeros({});
+  op_mean_out(x, dim_list, /*keepdim=*/false, dtype, out);
+  EXPECT_TENSOR_CLOSE(out, tf.make({}, {NAN}));
+
+  int64_t dims1[1] = {1};
+  dim_list = ArrayRef<int64_t>{dims1, 1};
+  out = tf.zeros({2, 3});
+  op_mean_out(x, dim_list, /*keepdim=*/false, dtype, out);
+  EXPECT_TENSOR_CLOSE(out, tf.make({2, 3}, {NAN, NAN, NAN, NAN, NAN, NAN}));
+
+  int64_t dims2[1] = {2};
+  dim_list = ArrayRef<int64_t>{dims2, 1};
+  out = tf.make({2, 0, 1}, {});
+  op_mean_out(x, dim_list, /*keepdim=*/true, dtype, out);
+  EXPECT_TENSOR_CLOSE(out, tf.make({2, 0, 1}, {}));
+}

kernels/test/op_sum_test.cpp

@@ -490,3 +490,30 @@ TEST_F(OpSumOutTest, InfinityAndNANTest) {
     }));
   // clang-format on
 }
+
+TEST_F(OpSumOutTest, EmptyInput) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor x = tf.make({2, 0, 3}, {});
+  optional<ScalarType> dtype = ScalarType::Float;
+  optional<ArrayRef<int64_t>> dim_list = ArrayRef<int64_t>{};
+  Tensor out = tf.ones({1, 1, 1});
+  op_sum_intlist_out(x, dim_list, /*keepdim=*/true, dtype, out);
+  EXPECT_TENSOR_CLOSE(out, tf.zeros({1, 1, 1}));
+
+  out = tf.ones({});
+  op_sum_intlist_out(x, dim_list, /*keepdim=*/false, dtype, out);
+  EXPECT_TENSOR_CLOSE(out, tf.zeros({}));
+
+  int64_t dims1[1] = {1};
+  dim_list = ArrayRef<int64_t>{dims1, 1};
+  out = tf.ones({2, 3});
+  op_sum_intlist_out(x, dim_list, /*keepdim=*/false, dtype, out);
+  EXPECT_TENSOR_CLOSE(out, tf.zeros({2, 3}));
+
+  int64_t dims2[1] = {2};
+  dim_list = ArrayRef<int64_t>{dims2, 1};
+  out = tf.make({2, 0, 1}, {});
+  op_sum_intlist_out(x, dim_list, /*keepdim=*/true, dtype, out);
+  EXPECT_TENSOR_CLOSE(out, tf.make({2, 0, 1}, {}));
+}

kernels/test/op_var_test.cpp

@@ -468,3 +468,30 @@ TEST_F(OpVarCorrectionOutTest, SmokeTest) {
   ET_FORALL_FLOATHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
+
+TEST_F(OpVarOutTest, EmptyInput) {
+  TensorFactory<ScalarType::Float> tf;
+
+  Tensor x = tf.make({2, 0, 3}, {});
+  bool unbiased = true;
+  optional<ArrayRef<int64_t>> dim_list = ArrayRef<int64_t>{};
+  Tensor out = tf.zeros({1, 1, 1});
+  op_var_out(x, dim_list, unbiased, /*keepdim=*/true, out);
+  EXPECT_TENSOR_CLOSE(out, tf.make({1, 1, 1}, {NAN}));
+
+  out = tf.zeros({});
+  op_var_out(x, dim_list, unbiased, /*keepdim=*/false, out);
+  EXPECT_TENSOR_CLOSE(out, tf.make({}, {NAN}));
+
+  int64_t dims1[1] = {1};
+  dim_list = ArrayRef<int64_t>{dims1, 1};
+  out = tf.zeros({2, 3});
+  op_var_out(x, dim_list, unbiased, /*keepdim=*/false, out);
+  EXPECT_TENSOR_CLOSE(out, tf.make({2, 3}, {NAN, NAN, NAN, NAN, NAN, NAN}));
+
+  int64_t dims2[1] = {2};
+  dim_list = ArrayRef<int64_t>{dims2, 1};
+  out = tf.make({2, 0, 1}, {});
+  op_var_out(x, dim_list, unbiased, /*keepdim=*/true, out);
+  EXPECT_TENSOR_CLOSE(out, tf.make({2, 0, 1}, {}));
+}