Add squeeze_copy.dims (#655)

SS-JIA · facebook-github-bot · commit a5cb1e2f26d3 · 2023-10-06T09:58:05.000-07:00
Summary: Pull Request resolved: #655 Add `squeeze_copy.dims`, which `squeeze_copy.dim` will decompose to. Reviewed By: manuelcandales Differential Revision: D49988507 fbshipit-source-id: b273db0a6cb9286e583c996a6aec58c0eeec4ce1
diff --git a/exir/dialects/edge/op/sample_input.py b/exir/dialects/edge/op/sample_input.py
@@ -1109,6 +1109,15 @@
             Return(ArgType.Tensor, size=[2]),
         ],
     },
+    "squeeze_copy.dims": {  # (Tensor self, int[] dims) -> Tensor
+        "args": [
+            InArg(ArgType.Tensor, size=[1, 2, 1, 5]),
+            InArg(ArgType.Param, value=[0, 2]),
+        ],
+        "returns": [
+            Return(ArgType.Tensor, size=[2, 5]),
+        ],
+    },
     "stack.default": {  # (Tensor[] tensors, int dim=0) -> Tensor
         "args": [
             InArg(
diff --git a/kernels/aten/functions.yaml b/kernels/aten/functions.yaml
@@ -290,6 +290,8 @@
 
 - op: squeeze_copy.dim_out
 
+- op: squeeze_copy.dims_out
+
 - op: squeeze_copy.out
 
 - op: stack.out
diff --git a/kernels/portable/cpu/op_squeeze_copy.cpp b/kernels/portable/cpu/op_squeeze_copy.cpp
@@ -10,6 +10,7 @@
 #include <cstdint>
 #include <cstring>
 
+#include <executorch/kernels/portable/cpu/util/copy_ops_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 
 namespace torch {
@@ -18,113 +19,65 @@ namespace native {
 
 using Tensor = exec_aten::Tensor;
 
-namespace {
-
-void check_squeeze_copy_dim_out(
-    const Tensor input,
-    int64_t dim,
-    const Tensor out) {
-  if (input.dim() != 0 && input.size(dim) == 1) {
-    ET_CHECK(input.dim() == out.dim() + 1);
-
-    for (size_t d = 0; d < out.dim(); ++d) {
-      if (d < dim) {
-        // d < dim
-        ET_CHECK_MSG(
-            input.size(d) == out.size(d),
-            "input.size(%zu) %zd != out.size(%zu) %zd | dim = %" PRId64,
-            d,
-            input.size(d),
-            d,
-            out.size(d),
-            dim);
-      } else {
-        // d >= dim
-        ET_CHECK_MSG(
-            input.size(d + 1) == out.size(d),
-            "input.size(%zu) %zd != out.size(%zu) %zd | dim = %" PRId64,
-            d + 1,
-            input.size(d),
-            d,
-            out.size(d),
-            dim);
-      }
-    }
-  } else {
-    ET_CHECK(input.dim() == out.dim());
-
-    for (size_t d = 0; d < out.dim(); ++d) {
-      ET_CHECK_MSG(
-          input.size(d) == out.size(d),
-          "input.size(%zu) %zd != out.size(%zu) %zd | dim = %" PRId64,
-          d,
-          input.size(d),
-          d,
-          out.size(d),
-          dim);
-    }
-  }
-}
-} // namespace
-
-//
-// squeeze_copy.dim_out(Tensor self, int dim, Tensor(a!) out) -> Tensor(a!)
-//
 Tensor& squeeze_copy_dim_out(
     RuntimeContext& ctx,
-    const Tensor& self,
+    const Tensor& in,
     int64_t dim,
     Tensor& out) {
   (void)ctx;
-  Tensor::SizesType expected_output_size[kTensorDimensionLimit];
-
-  // The input and out shall share same dtype
-  ET_CHECK_SAME_DTYPE2(self, out);
-
-  // A valid dim must be in [-self.dim(), self.dim())
-  if (self.dim() == 0 && dim == -1) {
-    dim = 0;
-  }
-  ET_CHECK_MSG(
-      (self.dim() == 0 && dim == 0) || (dim >= -self.dim() && dim < self.dim()),
-      "dim %" PRId64 " out of range [-%zd,%zd)",
-      dim,
-      self.dim(),
-      self.dim());
 
+  // TODO(ssjia): use nonzero_dim() instead
   if (dim < 0) {
-    dim += self.dim();
+    dim += in.dim();
   }
 
-  size_t expected_out_dim = (self.dim() == 0 || self.size(dim) != 1)
-      ? self.dim()
-      : std::max<ssize_t>(self.dim() - 1, 0);
+  ET_KERNEL_CHECK(
+      ctx, check_squeeze_copy_dim_args(in, dim, out), InvalidArgument, out);
+
+  Tensor::SizesType expected_out_size[kTensorDimensionLimit];
+  size_t expected_out_dim = 0;
+  get_squeeze_copy_dim_out_target_size(
+      in, dim, expected_out_size, &expected_out_dim);
+  ET_KERNEL_CHECK(
+      ctx,
+      resize_tensor(out, {expected_out_size, expected_out_dim}) == Error::Ok,
+      InvalidArgument,
+      out);
 
-  if (dim == self.dim() || self.size(dim) != 1) {
-    for (size_t i = 0; i < expected_out_dim; ++i) {
-      expected_output_size[i] = self.size(i);
-    }
-  } else {
-    // 0 <= dim < self.dim() AND self.size(dim) == 1
-    for (size_t i = 0; i < expected_out_dim; ++i) {
-      if (i < dim) {
-        expected_output_size[i] = self.size(i);
-      } else {
-        // Squeeze the given dimension 'dim'
-        expected_output_size[i] = self.size(i + 1);
-      }
-    }
+  if (in.nbytes() > 0) {
+    // Note that this check is important. It's valid for a tensor with numel 0
+    // to have a null data pointer, but in some environments it's invalid to
+    // pass a null pointer to memcpy() even when the size is zero.
+    memcpy(out.mutable_data_ptr(), in.const_data_ptr(), in.nbytes());
   }
-  ET_CHECK_MSG(
-      Error::Ok == resize_tensor(out, {expected_output_size, expected_out_dim}),
-      "Failed to resize output tensor.");
-  check_squeeze_copy_dim_out(self, dim, out);
+  return out;
+}
+
+Tensor& squeeze_copy_dims_out(
+    RuntimeContext& ctx,
+    const Tensor& in,
+    exec_aten::ArrayRef<int64_t> dims,
+    Tensor& out) {
+  (void)ctx;
+
+  ET_KERNEL_CHECK(
+      ctx, check_squeeze_copy_dims_args(in, dims, out), InvalidArgument, out);
+
+  Tensor::SizesType expected_out_size[kTensorDimensionLimit];
+  size_t expected_out_dim = 0;
+  get_squeeze_copy_dims_out_target_size(
+      in, dims, expected_out_size, &expected_out_dim);
+  ET_KERNEL_CHECK(
+      ctx,
+      resize_tensor(out, {expected_out_size, expected_out_dim}) == Error::Ok,
+      InvalidArgument,
+      out);
 
-  if (self.nbytes() > 0) {
+  if (in.nbytes() > 0) {
     // Note that this check is important. It's valid for a tensor with numel 0
     // to have a null data pointer, but in some environments it's invalid to
     // pass a null pointer to memcpy() even when the size is zero.
-    memcpy(out.mutable_data_ptr(), self.const_data_ptr(), self.nbytes());
+    memcpy(out.mutable_data_ptr(), in.const_data_ptr(), in.nbytes());
   }
   return out;
 }
diff --git a/kernels/portable/cpu/targets.bzl b/kernels/portable/cpu/targets.bzl
@@ -705,6 +705,9 @@ _ATEN_OPS = (
     ),
     op_target(
         name = "op_squeeze_copy",
+        deps = [
+            "//executorch/kernels/portable/cpu/util:copy_ops_util",
+        ],
     ),
     op_target(
         name = "op_stack",
diff --git a/kernels/portable/cpu/util/copy_ops_util.cpp b/kernels/portable/cpu/util/copy_ops_util.cpp
@@ -204,6 +204,111 @@ void get_split_with_sizes_copy_out_target_size(
   out_sizes[dim] = split_size;
 }
 
+bool check_squeeze_copy_dim_args(
+    const Tensor in,
+    int64_t dim,
+    const Tensor out) {
+  ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
+  ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(in, dim));
+
+  return true;
+}
+
+void get_squeeze_copy_dim_out_target_size(
+    const Tensor in,
+    int64_t dim,
+    Tensor::SizesType* out_sizes,
+    size_t* out_ndim) {
+  // For 0 dim tensors, the output should also be 0 dim.
+  if (in.dim() == 0) {
+    *out_ndim = 0;
+    return;
+  }
+
+  // Specified dim is only removed if the size at the given dim is 1.
+  if (in.size(dim) == 1) {
+    *out_ndim = in.dim() - 1;
+  } else {
+    *out_ndim = in.dim();
+  }
+
+  size_t out_d = 0;
+  for (size_t in_d = 0; in_d < in.dim(); ++in_d) {
+    if (in_d != dim || in.size(in_d) > 1) {
+      out_sizes[out_d] = in.size(in_d);
+      ++out_d;
+    }
+  }
+}
+
+bool check_squeeze_copy_dims_args(
+    const Tensor in,
+    const exec_aten::ArrayRef<int64_t> dims,
+    const Tensor out) {
+  ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
+
+  const int64_t dim_adjust = in.dim() == 0 ? 1 : in.dim();
+  for (size_t i = 0; i < dims.size(); ++i) {
+    // TODO(ssjia): use nonzero_dim() instead
+    const int64_t dim = dims[i] < 0 ? dims[i] + dim_adjust : dims[i];
+    ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(in, dim));
+
+    // Check that a dim does not appear twice in dims
+    for (size_t j = 0; j < dims.size(); ++j) {
+      if (i != j) {
+        const int64_t dim_temp = dims[j] < 0 ? dims[j] + dim_adjust : dims[j];
+        ET_LOG_MSG_AND_RETURN_IF_FALSE(
+            dim != dim_temp,
+            "dim %" PRId64 " appears multiple times in dims!",
+            dim);
+      }
+    }
+  }
+
+  return true;
+}
+
+void get_squeeze_copy_dims_out_target_size(
+    const Tensor in,
+    const exec_aten::ArrayRef<int64_t> dims,
+    Tensor::SizesType* out_sizes,
+    size_t* out_ndim) {
+  // For 0 dim tensors, the output should also be 0 dim.
+  if (in.dim() == 0) {
+    *out_ndim = 0;
+    return;
+  }
+
+  int64_t dim_adjust = in.dim() == 0 ? 1 : in.dim();
+  // A dim is only removed if the size at the given dim is 1.
+  Tensor::SizesType dims_to_remove = 0;
+  for (size_t i = 0; i < dims.size(); ++i) {
+    // TODO(ssjia): use nonzero_dim() instead
+    int64_t dim = dims[i] < 0 ? dims[i] + dim_adjust : dims[i];
+    if (in.size(dim) == 1) {
+      ++dims_to_remove;
+    }
+  }
+  *out_ndim = in.dim() - dims_to_remove;
+
+  size_t out_d = 0;
+  for (size_t in_d = 0; in_d < in.dim(); ++in_d) {
+    bool in_d_in_dims = false;
+    for (size_t i = 0; i < dims.size(); ++i) {
+      // TODO(ssjia): use nonzero_dim() instead
+      int64_t dim = dims[i] < 0 ? dims[i] + dim_adjust : dims[i];
+      if (in_d == dim) {
+        in_d_in_dims = true;
+        break;
+      }
+    }
+    if (!in_d_in_dims || in.size(in_d) > 1) {
+      out_sizes[out_d] = in.size(in_d);
+      ++out_d;
+    }
+  }
+}
+
 bool check_stack_args(
     exec_aten::ArrayRef<Tensor> tensors,
     int64_t dim,
diff --git a/kernels/portable/cpu/util/copy_ops_util.h b/kernels/portable/cpu/util/copy_ops_util.h
@@ -56,6 +56,28 @@ void get_split_with_sizes_copy_out_target_size(
     Tensor::SizesType* out_sizes,
     size_t* out_ndim);
 
+bool check_squeeze_copy_dim_args(
+    const Tensor in,
+    int64_t dim,
+    const Tensor out);
+
+void get_squeeze_copy_dim_out_target_size(
+    const Tensor in,
+    int64_t dim,
+    Tensor::SizesType* out_sizes,
+    size_t* out_ndim);
+
+bool check_squeeze_copy_dims_args(
+    const Tensor in,
+    const exec_aten::ArrayRef<int64_t> dims,
+    const Tensor out);
+
+void get_squeeze_copy_dims_out_target_size(
+    const Tensor in,
+    const exec_aten::ArrayRef<int64_t> dims,
+    Tensor::SizesType* out_sizes,
+    size_t* out_ndim);
+
 bool check_stack_args(
     exec_aten::ArrayRef<Tensor> tensors,
     int64_t dim,
diff --git a/kernels/portable/functions.yaml b/kernels/portable/functions.yaml
@@ -647,6 +647,11 @@
     - arg_meta: null
       kernel_name: torch::executor::squeeze_copy_dim_out
 
+- op: squeeze_copy.dims_out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::squeeze_copy_dims_out
+
 - op: stack.out
   kernels:
     - arg_meta: null
diff --git a/kernels/test/op_squeeze_copy_test.cpp b/kernels/test/op_squeeze_copy_test.cpp
