Fix op permute_copy (#726)

Summary:
Pull Request resolved: #726

Fix the check for duplicates in dimension list
ghstack-source-id: 203418805
exported-using-ghexport

Reviewed By: SS-JIA

Differential Revision: D50089231

fbshipit-source-id: e209bc8a1384992b0b9d88ae9ead29bd957450ea

Author: manuelcandales

kernels/portable/cpu/util/copy_ops_util.cpp

@@ -152,20 +152,25 @@ bool check_permute_copy_args(const Tensor& in, IntArrayRef dims, Tensor& out) {
   ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
 
   // Make sure no dimensions are duplicated and all in the range [-in.dim(),
-  // in.dim() - 1]. Use gaussian sum to check this.
-  size_t expected_sum = (dims.size() * (dims.size() + 1)) / 2;
-  size_t gauss_sum = 0;
+  // in.dim() - 1].
+  bool dim_exist[kTensorDimensionLimit];
+  memset(dim_exist, false, sizeof(dim_exist));
+
   for (int i = 0; i < dims.size(); i++) {
-    // Convert dimension to a non-negative number. dim_base is in the range
+    ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(in, dims[i]));
+    // Convert dimension to a non-negative number in the range
     // [0 .. in.dim() - 1].
-    size_t dim = dims[i] > -1 ? dims[i] : in.dim() + dims[i];
-    ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(in, dim));
-    gauss_sum += dim + 1;
-  }
+    size_t dim = dims[i] >= 0 ? dims[i] : in.dim() + dims[i];
 
-  ET_LOG_MSG_AND_RETURN_IF_FALSE(
-      gauss_sum == expected_sum,
-      "The dims passed to permute_copy must contain one of each dim!");
+    // Internal check, since we have already validated this
+    ET_CHECK(dim < kTensorDimensionLimit && dim >= 0);
+
+    // Check that the dimension hasn't been seen previously.
+    ET_LOG_MSG_AND_RETURN_IF_FALSE(
+        dim_exist[dim] == false, "duplicate dims are not allowed.");
+
+    dim_exist[dim] = true;
+  }
 
   return true;
 }

kernels/test/op_permute_copy_test.cpp

@@ -327,6 +327,20 @@ TEST(OpPermuteCopyKernelTest, DupeDimensionPos) {
       t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out));
 }
 
+TEST(OpPermuteCopyKernelTest, DupeDimensionPos2) {
+  TensorFactory<ScalarType::Int> tf;
+
+  const std::vector<int64_t> new_dim = {1, 1, 1};
+
+  const std::vector<int32_t> sizes = {1, 1, 1};
+  Tensor t_int = tf.make(sizes, {1});
+
+  Tensor out = tf.zeros(sizes);
+
+  ET_EXPECT_KERNEL_FAILURE(op_permute_copy_out(
+      t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out));
+}
+
 TEST(OpPermuteCopyKernelTest, DupeDimensionNeg) {
   TensorFactory<ScalarType::Int> tf;
 
@@ -341,6 +355,20 @@ TEST(OpPermuteCopyKernelTest, DupeDimensionNeg) {
       t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out));
 }
 
+TEST(OpPermuteCopyKernelTest, DupeDimensionNeg2) {
+  TensorFactory<ScalarType::Int> tf;
+
+  const std::vector<int64_t> new_dim = {0, 1, -5};
+
+  const std::vector<int32_t> sizes = {1, 1, 1};
+  Tensor t_int = tf.make(sizes, {1});
+
+  Tensor out = tf.zeros(sizes);
+
+  ET_EXPECT_KERNEL_FAILURE(op_permute_copy_out(
+      t_int, ArrayRef<int64_t>(new_dim.data(), new_dim.size()), out));
+}
+
 TEST(OpPermuteCopyKernelTest, MismatchDim) {
   TensorFactory<ScalarType::Int> tf;