Pattern: unary realb to float (exp & log)

Reviewed By: digantdesai

Differential Revision: D47118015

fbshipit-source-id: 03a889582cc669c5c6e2817db97e78e8a1a51775

Author: manuelcandales

kernels/optimized/cpu/op_exp.cpp

@@ -56,7 +56,7 @@ Tensor& opt_exp_out(RuntimeContext& ctx, const Tensor& in, Tensor& out) {
   auto error = resize_tensor(out, in.sizes());
   ET_CHECK_MSG(error == Error::Ok, "Failed to resize output tensor.");
 
-  ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "exp", CTYPE_IN, [&] {
+  ET_SWITCH_REAL_TYPES_AND(Bool, in.scalar_type(), ctx, "exp", CTYPE_IN, [&] {
     ET_SWITCH_FLOAT_TYPES(out.scalar_type(), ctx, "exp", CTYPE_OUT, [&] {
       exp_data<CTYPE_IN, CTYPE_OUT>(
           in.const_data_ptr<CTYPE_IN>(),

kernels/portable/cpu/op_exp.cpp

@@ -1,37 +1,15 @@
-// Copyright (c) Meta Platforms, Inc. and affiliates.
-#include <cmath>
+// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
 
 #include <executorch/kernels/kernel_includes.h>
-#include <executorch/kernels/portable/cpu/util/functional_util.h>
+#include <executorch/kernels/portable/cpu/pattern/pattern.h>
+#include <cmath>
 
 namespace torch {
 namespace executor {
 namespace native {
 
-using exec_aten::Tensor;
-
 Tensor& exp_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_SHAPE2(in, out);
-
-  ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "exp", CTYPE_IN, [&] {
-    ET_SWITCH_FLOAT_TYPES(out.scalar_type(), ctx, "exp", CTYPE_OUT, [&] {
-      apply_unary_map_fn(
-          [](const CTYPE_IN val_in) {
-            CTYPE_OUT xi = static_cast<CTYPE_OUT>(val_in);
-            return std::exp(xi);
-          },
-          in.const_data_ptr<CTYPE_IN>(),
-          out.mutable_data_ptr<CTYPE_OUT>(),
-          in.numel());
-    });
-  });
-
-  return out;
+  return internal::unary_ufunc_realb_to_float(std::exp, ctx, in, out);
 }
 
 } // namespace native

kernels/portable/cpu/op_log.cpp

@@ -1,39 +1,15 @@
-// Copyright (c) Meta Platforms, Inc. and affiliates.
-
-#include <cmath>
+// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
 
 #include <executorch/kernels/kernel_includes.h>
-#include <executorch/kernels/portable/cpu/util/functional_util.h>
+#include <executorch/kernels/portable/cpu/pattern/pattern.h>
+#include <cmath>
 
 namespace torch {
 namespace executor {
 namespace native {
 
-using exec_aten::Tensor;
-
 Tensor& log_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_SHAPE2(in, out);
-
-  ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "log", CTYPE_IN, [&] {
-    ET_SWITCH_FLOAT_TYPES(out.scalar_type(), ctx, "log", CTYPE_OUT, [&] {
-      apply_unary_map_fn(
-          [](const CTYPE_IN val_in) {
-            CTYPE_OUT xi = static_cast<CTYPE_OUT>(val_in);
-            ET_CHECK_MSG(xi > 0.0, "input must be greater than 0.");
-            return static_cast<CTYPE_OUT>(log(xi));
-          },
-          in.const_data_ptr<CTYPE_IN>(),
-          out.mutable_data_ptr<CTYPE_OUT>(),
-          in.numel());
-    });
-  });
-
-  return out;
+  return internal::unary_ufunc_realb_to_float(std::log, ctx, in, out);
 }
 
 } // namespace native

kernels/test/op_exp_test.cpp

@@ -72,24 +72,18 @@ TEST(OpExpOutKernelTest, AllFloatInputDoubleOutputSupport) {
 #undef TEST_ENTRY
 }
 
-TEST(OpExpOutKernelTest, UnhandledInputDtypeDies) {
-  if (SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle bool dtype";
-  }
-
-  // _exp_out() doesn't handle Bool as input.
+TEST(OpExpOutKernelTest, HandleBoolInput) {
+  // _exp_out() handles Bool as input.
   TensorFactory<ScalarType::Bool> tf_bool;
   TensorFactory<ScalarType::Float> tf_float;
 
-  const std::vector<int32_t> sizes = {2, 2};
-  Tensor a = tf_bool.make(sizes, /*data=*/{false, true, false, true});
+  const std::vector<int32_t> sizes = {1, 2};
 
-  // Destination for the log
+  Tensor a = tf_bool.make(sizes, /*data=*/{true, false});
   Tensor out = tf_float.zeros(sizes);
+  Tensor res = tf_float.make(sizes, /*data=*/{2.718282, 1});
 
-  // Boolean tensor should cause an assertion and kill the
-  // test process.
-  ET_EXPECT_KERNEL_FAILURE(_exp_out(a, out));
+  EXPECT_TENSOR_CLOSE(_exp_out(a, out), res);
 }
 
 // Mismatched shape tests.

kernels/test/op_log_test.cpp

@@ -29,18 +29,12 @@ void test__log_out() {
 
   const std::vector<int32_t> sizes = {2, 2};
 
-  // Invalid input zero should die
   Tensor out = tf_out.zeros(sizes);
-  if (SupportedFeatures::get()->is_aten) {
-    // However, ATen can handle when input is zero
-  } else {
-    ET_EXPECT_KERNEL_FAILURE(_log_out(tf.zeros(sizes), out));
-  }
 
   // Valid input should give the expected output
-  _log_out(tf.make(sizes, /*data=*/{1, 2, 4, 8}), out);
+  _log_out(tf.make(sizes, /*data=*/{0, 1, 2, 4}), out);
   EXPECT_TENSOR_CLOSE(
-      out, tf_out.make(sizes, /*data=*/{0, 0.693147, 1.386294, 2.079441}));
+      out, tf_out.make(sizes, /*data=*/{-INFINITY, 0, 0.693147, 1.386294}));
 }
 
 TEST(OpLogOutKernelTest, AllRealInputFloatOutputSupport) {
@@ -57,24 +51,18 @@ TEST(OpLogOutKernelTest, AllRealInputDoubleOutputSupport) {
 #undef TEST_ENTRY
 }
 
-TEST(OpLogOutKernelTest, UnhandledInputDtypeDies) {
-  if (SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle bool dtype";
-  }
-
-  // _log_out() doesn't handle Bool as input.
+TEST(OpLogOutKernelTest, HandleBoolInput) {
+  // _log_out() handles Bool as input.
   TensorFactory<ScalarType::Bool> tf_bool;
   TensorFactory<ScalarType::Float> tf_float;
 
-  const std::vector<int32_t> sizes = {2, 2};
-  Tensor a = tf_bool.make(sizes, /*data=*/{false, true, false, true});
+  const std::vector<int32_t> sizes = {1, 2};
 
-  // Destination for the log
+  Tensor a = tf_bool.make(sizes, /*data=*/{true, false});
   Tensor out = tf_float.zeros(sizes);
+  Tensor res = tf_float.make(sizes, /*data=*/{0, -INFINITY});
 
-  // Boolean tensor should cause an assertion and kill the
-  // test process.
-  ET_EXPECT_KERNEL_FAILURE(_log_out(a, out));
+  EXPECT_TENSOR_EQ(_log_out(a, out), res);
 }
 
 // Mismatched shape tests.