Fix & cleanup op ones (#699)

Summary:
Pull Request resolved: #699

Resize out tensor
ghstack-source-id: 203341575
exported-using-ghexport

Reviewed By: SS-JIA

Differential Revision: D49735857

fbshipit-source-id: 75c30806c541f097641a73cc502443834d071789

Author: manuelcandales

kernels/portable/cpu/op_ones.cpp

@@ -8,67 +8,24 @@
 
 #include <executorch/runtime/kernel/kernel_includes.h>
 
-#include <cstring>
-
 namespace torch {
 namespace executor {
 namespace native {
 
-using exec_aten::Scalar;
-using ScalarType = exec_aten::ScalarType;
-
-namespace {
-
-/**
- * Checks sizes passed to `ones_out` (`size_int64_t`) matches those within
- * `out` tensor (`size_int32_t`).
- */
-void size_check(
-    exec_aten::ArrayRef<int64_t> size_int64_t,
-    exec_aten::ArrayRef<int32_t> size_int32_t) {
-  ET_CHECK(size_int64_t.size() == size_int32_t.size());
-  for (int i = 0; i < size_int64_t.size(); i++) {
-    ET_CHECK(((int64_t)size_int32_t[i] == size_int64_t[i]));
-  }
-};
-
-/**
- * Fills the `out` tensor with value 1.
- */
-template <class CTYPE>
-void ones_kernel(Tensor& out) {
-  // Create pointer over `out` data with `CTYPE`.
-  auto data_out = out.mutable_data_ptr<CTYPE>();
-
-  // Set each element of the tensor to the "1" value for the type.
-  for (size_t i = 0; i < out.numel(); i++) {
-    data_out[i] = static_cast<CTYPE>(1);
-  }
-};
-
-} // namespace
-
-/**
- * `ones_out` implementation.
- */
 Tensor& ones_out(RuntimeContext& ctx, IntArrayRef size, Tensor& out) {
   (void)ctx;
-  size_check(size, out.sizes());
-
-#define ONES_OUT(ctype, dtype) \
-  case ScalarType::dtype:      \
-    ones_kernel<ctype>(out);   \
-    break;
 
-  switch (out.scalar_type()) {
-    ET_FORALL_REAL_TYPES_AND(Bool, ONES_OUT)
-    default:
-      ET_CHECK_MSG(
-          false,
-          "out tensor should be a real or bool dtype, but got %" PRId8,
-          static_cast<int8_t>(out.scalar_type()));
-  }
-#undef ONES_OUT
+  // Resize for dynamic shape
+  ET_KERNEL_CHECK(
+      ctx, resize_tensor(out, size) == Error::Ok, InvalidArgument, out);
+
+  ScalarType out_type = out.scalar_type();
+  ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, __func__, CTYPE, [&] {
+    auto out_data = out.mutable_data_ptr<CTYPE>();
+    for (size_t i = 0; i < out.numel(); i++) {
+      out_data[i] = static_cast<CTYPE>(1);
+    }
+  });
 
   return out;
 }