Cleanup log_softmax & softmax (#791)

Summary:
Pull Request resolved: #791

Refactors helper functions for checking arguments, to be reused by optimized `log_softmax`
ghstack-source-id: 203560636
exported-using-ghexport

Reviewed By: cbilgin

Differential Revision: D50130021

fbshipit-source-id: f964002a195563aef7d1b573b453da5423585647

Author: manuelcandales

kernels/portable/cpu/op_log_softmax.cpp

@@ -8,6 +8,7 @@
 
 #include <cmath>
 
+#include <executorch/kernels/portable/cpu/util/activation_ops_util.h>
 #include <executorch/kernels/portable/cpu/util/functional_util.h>
 #include <executorch/kernels/portable/cpu/util/reduce_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
@@ -18,37 +19,6 @@ namespace native {
 
 using Tensor = exec_aten::Tensor;
 
-namespace {
-
-void check_preconditions(
-    const Tensor& in,
-    int64_t dim,
-    bool half_to_float,
-    Tensor& out) {
-  // Ensure half_to_float is not true
-  ET_CHECK_MSG(
-      !half_to_float,
-      "log_softmax with half to float conversion is not supported on CPU");
-  // Check both in and out are of the same dtype
-  ET_CHECK_SAME_DTYPE2(in, out);
-  // Check both in and out have the same number of dimensions
-  ET_CHECK_MSG(
-      in.dim() == out.dim(),
-      "in.dim() %zd!= out.dim() %zd",
-      in.dim(),
-      out.dim());
-  // Ensure dim is valid
-  if (in.dim() == 0) {
-    ET_CHECK_MSG(dim == 0 || dim == -1, "dim must be 0 or -1 for 0-D tensor");
-  } else {
-    ET_CHECK_VALID_DIM(dim, in.dim());
-  }
-  ET_CHECK_DEFAULT_OR_CHANNELSLAST_DIMORDER(in);
-  ET_CHECK_DEFAULT_OR_CHANNELSLAST_DIMORDER(out);
-}
-
-} // namespace
-
 Tensor& log_softmax_out(
     RuntimeContext& ctx,
     const Tensor& in,
@@ -57,14 +27,13 @@ Tensor& log_softmax_out(
     Tensor& out) {
   (void)ctx;
 
-  check_preconditions(in, dim, half_to_float, out);
+  check_log_softmax_args(in, dim, half_to_float, out);
 
-  Error err = resize_tensor(out, in.sizes());
-  ET_CHECK_MSG(
-      err == Error::Ok, "Failed to resize out tensor in log_softmax_out");
+  ET_KERNEL_CHECK(
+      ctx, resize_tensor(out, in.sizes()) == Error::Ok, InvalidArgument, out);
 
   // Adjust for negative dim
-  dim = dim < 0 ? dim + in.dim() : dim;
+  dim = dim < 0 ? dim + nonzero_dim(in) : dim;
 
   ET_SWITCH_FLOAT_TYPES(in.scalar_type(), ctx, "log_softmax", CTYPE, [&]() {
     const CTYPE* const in_data = in.const_data_ptr<CTYPE>();

kernels/portable/cpu/op_softmax.cpp

@@ -8,6 +8,7 @@
 
 #include <cmath>
 
+#include <executorch/kernels/portable/cpu/util/activation_ops_util.h>
 #include <executorch/kernels/portable/cpu/util/functional_util.h>
 #include <executorch/kernels/portable/cpu/util/reduce_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
@@ -18,37 +19,6 @@ namespace native {
 
 using Tensor = exec_aten::Tensor;
 
-namespace {
-
-void check_preconditions(
-    const Tensor& in,
-    int64_t dim,
-    bool half_to_float,
-    Tensor& out) {
-  // Ensure half_to_float is not true
-  ET_CHECK_MSG(
-      !half_to_float,
-      "softmax with half to float conversion is not supported on CPU");
-  // Check both in and out are of the same dtype
-  ET_CHECK_SAME_DTYPE2(in, out);
-  // Check both in and out have the same number of dimensions
-  ET_CHECK_MSG(
-      in.dim() == out.dim(),
-      "in.dim() %zd!= out.dim() %zd",
-      in.dim(),
-      out.dim());
-  // Ensure dim is valid
-  if (in.dim() == 0) {
-    ET_CHECK(dim == 0 || dim == -1);
-  } else {
-    ET_CHECK_VALID_DIM(dim, in.dim());
-  }
-  ET_CHECK_DEFAULT_OR_CHANNELSLAST_DIMORDER(in);
-  ET_CHECK_DEFAULT_OR_CHANNELSLAST_DIMORDER(out);
-}
-
-} // namespace
-
 Tensor& softmax_out(
     RuntimeContext& ctx,
     const Tensor& in,
@@ -57,13 +27,13 @@ Tensor& softmax_out(
     Tensor& out) {
   (void)ctx;
 
-  check_preconditions(in, dim, half_to_float, out);
+  check_softmax_args(in, dim, half_to_float, out);
 
-  Error err = resize_tensor(out, in.sizes());
-  ET_CHECK_MSG(err == Error::Ok, "Failed to resize out tensor in softmax_out");
+  ET_KERNEL_CHECK(
+      ctx, resize_tensor(out, in.sizes()) == Error::Ok, InvalidArgument, out);
 
   // Adjust for negative dim
-  dim = dim < 0 ? dim + in.dim() : dim;
+  dim = dim < 0 ? dim + nonzero_dim(in) : dim;
 
   ET_SWITCH_FLOAT_TYPES(in.scalar_type(), ctx, "softmax", CTYPE, [&]() {
     const CTYPE* const in_data = in.const_data_ptr<CTYPE>();

kernels/portable/cpu/targets.bzl

@@ -434,6 +434,7 @@ _ATEN_OPS = (
         name = "op_log_softmax",
         deps = [
             ":vec_ops",
+            "//executorch/kernels/portable/cpu/util:activation_ops_util",
             "//executorch/kernels/portable/cpu/util:functional_util",
             "//executorch/kernels/portable/cpu/util:reduce_util",
         ],
@@ -693,6 +694,7 @@ _ATEN_OPS = (
         name = "op_softmax",
         deps = [
             ":vec_ops",
+            "//executorch/kernels/portable/cpu/util:activation_ops_util",
             "//executorch/kernels/portable/cpu/util:functional_util",
             "//executorch/kernels/portable/cpu/util:reduce_util",
         ],

kernels/portable/cpu/util/activation_ops_util.cpp

@@ -23,5 +23,27 @@ bool check_gelu_args(const Tensor& in, string_view approximate, Tensor& out) {
   return true;
 }
 
+bool check_log_softmax_args(
+    const Tensor& in,
+    int64_t dim,
+    bool half_to_float,
+    Tensor& out) {
+  ET_LOG_MSG_AND_RETURN_IF_FALSE(
+      !half_to_float, "half to float conversion is not supported on CPU");
+  ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out));
+  ET_LOG_AND_RETURN_IF_FALSE(tensor_has_dim(in, dim));
+  ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_or_channels_last_dim_order(in));
+  ET_LOG_AND_RETURN_IF_FALSE(tensor_is_default_or_channels_last_dim_order(out));
+  return true;
+}
+
+bool check_softmax_args(
+    const Tensor& in,
+    int64_t dim,
+    bool half_to_float,
+    Tensor& out) {
+  return check_log_softmax_args(in, dim, half_to_float, out);
+}
+
 } // namespace executor
 } // namespace torch

kernels/portable/cpu/util/activation_ops_util.h

@@ -15,5 +15,17 @@ namespace executor {
 
 bool check_gelu_args(const Tensor& in, string_view approximate, Tensor& out);
 
+bool check_log_softmax_args(
+    const Tensor& in,
+    int64_t dim,
+    bool half_to_float,
+    Tensor& out);
+
+bool check_softmax_args(
+    const Tensor& in,
+    int64_t dim,
+    bool half_to_float,
+    Tensor& out);
+
 } // namespace executor
 } // namespace torch