[ET][Portable] Fix & cleanup op as_strided_copy

kernels/portable/cpu/op_as_strided_copy.cpp

@@ -6,8 +6,8 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/kernels/portable/cpu/util/copy_ops_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
-#include <cstring>
 
 namespace torch {
 namespace executor {
@@ -17,40 +17,6 @@ using Tensor = exec_aten::Tensor;
 using ScalarType = exec_aten::ScalarType;
 
 namespace {
-size_t compute_storage_nbytes(
-    IntArrayRef sizes,
-    IntArrayRef strides,
-    size_t itemsize_bytes) {
-  // size of the underlying storage is 1 bigger than the offset
-  // of the last element according to stride
-  size_t size = 1;
-  for (size_t i = 0; i < sizes.size(); ++i) {
-    if (sizes[i] == 0) {
-      return 0;
-    }
-    size += strides[i] * (sizes[i] - 1);
-  }
-  return size * itemsize_bytes;
-}
-
-void check_inbounds_for_storage(
-    const Tensor& self,
-    ArrayRef<int64_t> size,
-    ArrayRef<int64_t> stride,
-    int64_t storage_offset) {
-  size_t storage_size_bytes =
-      compute_storage_nbytes(size, stride, self.element_size());
-  size_t storage_offset_bytes = storage_offset * self.element_size();
-  if (storage_size_bytes == 0) {
-    return;
-  }
-  size_t new_storage_size_bytes = self.nbytes();
-  ET_CHECK_MSG(
-      storage_size_bytes + storage_offset_bytes <= new_storage_size_bytes,
-      "Requiring a storage size of %zd are out of bounds for storage of size %zd",
-      storage_size_bytes + storage_offset_bytes,
-      new_storage_size_bytes);
-}
 
 /**
  * Copy input_data to output_data according to the stride and shape recursively
@@ -81,39 +47,8 @@ void as_strided_copy(
   }
 }
 
-void check_preconditions(
-    const Tensor& self,
-    ArrayRef<int64_t> size,
-    ArrayRef<int64_t> stride,
-    optional<int64_t> storage_offset,
-    Tensor& out) {
-  ET_CHECK_SAME_DTYPE2(self, out);
-  ET_CHECK_MSG(
-      size.size() == stride.size(), "mismatch in length of strides and shape");
-  for (const auto& val : stride) {
-    ET_CHECK_MSG(
-        val >= 0,
-        "as_strided: Negative strides are not supported at the moment");
-  }
-  ET_CHECK_MSG(
-      out.sizes().size() == size.size(),
-      "output tensor should have same shape as size");
-  for (size_t i = 0; i < out.sizes().size(); ++i) {
-    ET_CHECK_MSG(
-        out.sizes().at(i) == size.at(i),
-        "output tensor should have same shape as size");
-  }
-  int64_t offset = storage_offset.has_value() ? storage_offset.value() : 0;
-  ET_CHECK_MSG(offset >= 0, "Negative storage offset");
-  check_inbounds_for_storage(self, size, stride, offset);
-}
-
 } // namespace
 
-/**
- * Copy the tener `self` to `out`, assume `self` and `out` have same type and
- * shape
- */
 Tensor& as_strided_copy_out(
     RuntimeContext& ctx,
     const Tensor& self,
@@ -123,34 +58,35 @@ Tensor& as_strided_copy_out(
     Tensor& out) {
   (void)ctx;
 
-  torch::executor::Error err = resize_tensor(out, size);
-  ET_CHECK_MSG(
-      err == torch::executor::Error::Ok,
-      "Failed to resize out Tensor in as_strided_copy_out");
+  ET_KERNEL_CHECK(
+      ctx,
+      check_as_strided_copy_args(self, size, stride, storage_offset, out),
+      InvalidArgument,
+      out);
+
+  ET_KERNEL_CHECK(
+      ctx,
+      resize_tensor(out, size) == torch::executor::Error::Ok,
+      InvalidArgument,
+      out);
+
+  if (self.numel() == 0) {
+    return out;
+  }
 
-  check_preconditions(self, size, stride, storage_offset, out);
   size_t offset = storage_offset.has_value() ? storage_offset.value() : 0;
 
-#define AS_STRIDED_COPY_TENSOR(ctype, dtype)                    \
-  case ScalarType::dtype:                                       \
-    as_strided_copy<ctype>(                                     \
-        /*input_data=*/self.mutable_data_ptr<ctype>() + offset, \
-        /*output_data=*/out.mutable_data_ptr<ctype>(),          \
-        out,                                                    \
-        size,                                                   \
-        stride,                                                 \
-        /*dim=*/0);                                             \
-    break;
+  ET_SWITCH_ALL_TYPES(self.scalar_type(), ctx, __func__, CTYPE, [&] {
+    CTYPE* self_data = self.mutable_data_ptr<CTYPE>() + offset;
+    CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
+
+    if (size.empty()) {
+      out_data[0] = self_data[0];
+    } else {
+      as_strided_copy<CTYPE>(self_data, out_data, out, size, stride, 0);
+    }
+  });
 
-  switch (self.scalar_type()) {
-    ET_FORALL_SCALAR_TYPES(AS_STRIDED_COPY_TENSOR)
-    default:
-      ET_CHECK_MSG(
-          false,
-          "Unhandled dtype %" PRId8,
-          static_cast<int8_t>(self.scalar_type()));
-  }
-#undef AS_STRIDED_COPY_TENSOR
   return out;
 }
 

kernels/portable/cpu/targets.bzl

@@ -92,7 +92,7 @@ _ATEN_OPS = (
     op_target(
         name = "op_as_strided_copy",
         deps = [
-            ":scalar_utils",
+            "//executorch/kernels/portable/cpu/util:copy_ops_util",
         ],
     ),
     op_target(

kernels/portable/cpu/util/copy_ops_util.cpp

@@ -15,6 +15,60 @@ namespace executor {
 
 using Tensor = exec_aten::Tensor;
 
+namespace {
+
+size_t as_strided_copy_compute_storage_nbytes(
+    IntArrayRef sizes,
+    IntArrayRef strides,
+    size_t itemsize_bytes) {
+  // size of the underlying storage is 1 bigger than the offset
+  // of the last element according to stride
+  size_t size = 1;
+  for (size_t i = 0; i < sizes.size(); ++i) {
+    if (sizes[i] == 0) {
+      return 0;
+    }
+    size += strides[i] * (sizes[i] - 1);
+  }
+  return size * itemsize_bytes;
+}
+
+} // namespace
+
+bool check_as_strided_copy_args(
+    const Tensor& in,
+    ArrayRef<int64_t> size,
+    ArrayRef<int64_t> stride,
+    optional<int64_t> storage_offset,
+    Tensor& out) {
+  ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      size.size() == stride.size(), "mismatch in length of strides and shape");
+  for (const auto& val : stride) {
+    ET_LOG_MSG_AND_RETURN_IF_FALSE(
+        val >= 0,
+        "as_strided: Negative strides are not supported at the moment");
+  }
+
+  int64_t offset = storage_offset.has_value() ? storage_offset.value() : 0;
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(offset >= 0, "Negative storage offset");
+
+  // Check that the requested storage is within bounds of input storage
+  size_t storage_size_bytes =
+      as_strided_copy_compute_storage_nbytes(size, stride, in.element_size());
+  size_t storage_offset_bytes = offset * in.element_size();
+  if (storage_size_bytes == 0) {
+    return true;
+  }
+  size_t new_storage_size_bytes = in.nbytes();
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      storage_size_bytes + storage_offset_bytes <= new_storage_size_bytes,
+      "Requiring a storage size of %zd are out of bounds for storage of size %zd",
+      storage_size_bytes + storage_offset_bytes,
+      new_storage_size_bytes);
+  return true;
+}
+
 bool check_cat_args(
     exec_aten::ArrayRef<Tensor> tensors,
     int64_t dim,

kernels/portable/cpu/util/copy_ops_util.h

@@ -13,6 +13,13 @@
 namespace torch {
 namespace executor {
 
+bool check_as_strided_copy_args(
+    const Tensor& in,
+    ArrayRef<int64_t> size,
+    ArrayRef<int64_t> stride,
+    optional<int64_t> storage_offset,
+    Tensor& out);
+
 bool check_cat_args(
     exec_aten::ArrayRef<Tensor> tensors,
     int64_t dim,