Pattern: unary_ufunc_real

Reviewed By: kirklandsign

Differential Revision: D48380141

fbshipit-source-id: 4a47b9aa41e7010cf81fa90aa58a7ae4735f87ed

Author: manuelcandales

kernels/portable/cpu/op_floor.cpp

@@ -6,11 +6,9 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <cmath>
-
-#include <executorch/kernels/portable/cpu/util/functional_util.h>
+#include <executorch/kernels/portable/cpu/pattern/pattern.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
-#include <executorch/runtime/platform/assert.h>
+#include <cmath>
 
 namespace torch {
 namespace executor {
@@ -19,24 +17,7 @@ namespace native {
 using exec_aten::Tensor;
 
 Tensor& floor_out(RuntimeContext& ctx, const Tensor& in, Tensor& out) {
-  (void)ctx;
-
-  // Resize for dynamic shape
-  auto error = resize_tensor(out, in.sizes());
-  ET_CHECK_MSG(error == Error::Ok, "Failed to resize output tensor.");
-  ET_CHECK_SAME_SHAPE_AND_DTYPE2(in, out);
-
-  ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "floor", CTYPE, [&] {
-    apply_unary_map_fn(
-        [](const CTYPE val_in) {
-          return static_cast<CTYPE>(std::floor(val_in));
-        },
-        in.const_data_ptr<CTYPE>(),
-        out.mutable_data_ptr<CTYPE>(),
-        in.numel());
-  });
-
-  return out;
+  return internal::unary_ufunc_real(std::floor, ctx, in, out);
 }
 
 } // namespace native

kernels/portable/cpu/pattern/pattern.h

@@ -51,6 +51,18 @@ namespace executor {
 namespace native {
 namespace internal {
 
+/**
+ * Implements an op pattern for ops that take a single input tensor of any
+ * real dtye, no additional arguments, and outputs a tensor of the same size
+ * and dtype. The function fn specifies the math operation which is applied to
+ * the input tensor element-wise.
+ */
+Tensor& unary_ufunc_real(
+    FunctionRef<double(double)> fn,
+    RuntimeContext& ctx,
+    const Tensor& in,
+    Tensor& out);
+
 /**
  * Implements an op pattern for ops that take a single input tensor of any
  * realb dtye (real and boolean), no additional arguments, and outputs a

kernels/portable/cpu/pattern/targets.bzl

@@ -13,6 +13,7 @@ def define_common_targets():
             "binary_ufunc_realb_realb_to_realb_logical.cpp",
             "unary_ufunc_realb_to_bool.cpp",
             "unary_ufunc_realb_to_float.cpp",
+            "unary_ufunc_real.cpp",
         ],
         exported_headers = [
             "pattern.h",

kernels/portable/cpu/pattern/unary_ufunc_real.cpp

@@ -0,0 +1,45 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/kernels/portable/cpu/pattern/pattern.h>
+#include <executorch/kernels/portable/cpu/util/functional_util.h>
+#include <executorch/runtime/core/function_ref.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+
+namespace torch {
+namespace executor {
+namespace native {
+namespace internal {
+
+Tensor& unary_ufunc_real(
+    FunctionRef<double(double)> fn,
+    RuntimeContext& ctx,
+    const Tensor& in,
+    Tensor& out) {
+  (void)ctx;
+
+  // Resize for dynamic shape
+  auto error = resize_tensor(out, in.sizes());
+  ET_CHECK_MSG(error == Error::Ok, "Failed to resize output tensor.");
+  ET_CHECK_SAME_SHAPE_AND_DTYPE2(in, out);
+
+  ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, __func__, CTYPE, [&] {
+    apply_unary_map_fn(
+        [fn](const CTYPE val_in) { return static_cast<CTYPE>(fn(val_in)); },
+        in.const_data_ptr<CTYPE>(),
+        out.mutable_data_ptr<CTYPE>(),
+        in.numel());
+  });
+
+  return out;
+}
+
+} // namespace internal
+} // namespace native
+} // namespace executor
+} // namespace torch

kernels/portable/cpu/targets.bzl

@@ -287,7 +287,7 @@ _ATEN_OPS = (
     op_target(
         name = "op_floor",
         deps = [
-            "//executorch/kernels/portable/cpu/util:functional_util",
+            "//executorch/kernels/portable/cpu/pattern:pattern",
         ],
     ),
     op_target(