Reapply pytorch#11294 and pytorch#11295 (improve GLU test and implement using internal views to avoid copying) (pytorch#11539)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#11509 by
@swolchok
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/swolchok/451/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/swolchok/451/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/swolchok/451/orig
@diff-train-skip-merge

Co-authored-by: Scott Wolchok <[email protected]>

Author: pytorchbot

kernels/portable/cpu/op_glu.cpp

@@ -8,6 +8,7 @@
 
 #include <c10/util/irange.h>
 #include <executorch/kernels/portable/cpu/util/activation_ops_util.h>
+#include <executorch/kernels/portable/cpu/util/elementwise_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 #include <executorch/runtime/platform/assert.h>
 #include <cinttypes>
@@ -23,92 +24,46 @@ using ScalarType = executorch::aten::ScalarType;
 
 namespace {
 
-double exp_overload(double d) {
-  return exp(d);
-}
-
-float exp_overload(float f) {
-  return expf(f);
-}
-
-/**
- * In-place element-wise sigmoid function , i.e., f(x) = 1 / (1 + e^{-x})
- */
-// TODO: T146333648, refactor this as a common helper function
-template <typename CTYPE_OUT>
-void sigmoid_tensor(Tensor& out) {
-  CTYPE_OUT* out_data = out.mutable_data_ptr<CTYPE_OUT>();
-  for (const auto i : c10::irange(out.numel())) {
-    out_data[i] = 1.0 / (1.0 + exp_overload(-out_data[i]));
-  }
-}
-
-/**
- * Element-wise multiplication of the first half of `in` along the specified
- * dimension and `out`, overwriting `out`.
- */
-template <typename CTYPE_IN, typename CTYPE_OUT>
-void mul_tensors(const Tensor& in, int64_t dim, Tensor& out) {
-  size_t num_values = static_cast<size_t>(in.size(dim)) / 2;
-  size_t dim_length_in = static_cast<size_t>(in.size(dim));
-  size_t dim_length_out = static_cast<size_t>(out.size(dim));
-  size_t leading_dims = getLeadingDims(in, dim);
-  size_t trailing_dims = getTrailingDims(in, dim);
-
-  const CTYPE_IN* input_data_base = in.const_data_ptr<CTYPE_IN>();
-  CTYPE_OUT* output_data_base = out.mutable_data_ptr<CTYPE_OUT>();
-
-  for (const auto i : c10::irange(leading_dims)) {
-    const CTYPE_IN* input_data =
-        input_data_base + i * dim_length_in * trailing_dims;
-    CTYPE_OUT* output_data =
-        output_data_base + i * dim_length_out * trailing_dims;
-    for ([[maybe_unused]] const auto j : c10::irange(num_values)) {
-      for (const auto k : c10::irange(trailing_dims)) {
-        output_data[k] = static_cast<CTYPE_OUT>(input_data[k]) * output_data[k];
-      }
-      input_data += trailing_dims;
-      output_data += trailing_dims;
-    }
+struct SplitGLUInputTensor {
+  explicit SplitGLUInputTensor(const Tensor& self, int64_t dim);
+  using SizesArray =
+      std::array<executorch::aten::SizesType, kTensorDimensionLimit>;
+  SizesArray half_sizes;
+  TensorImpl first_half_impl;
+  TensorImpl second_half_impl;
+  Tensor first_half;
+  Tensor second_half;
+
+ private:
+  static SizesArray get_half_sizes(const Tensor& self, int64_t dim) {
+    SizesArray half_sizes;
+    std::copy(self.sizes().begin(), self.sizes().end(), half_sizes.begin());
+    half_sizes[dim] /= 2;
+    return half_sizes;
   }
-}
-
-/**
- * Slice the tensor in the given dim, from start to end, assume tensor in and
- * out have same shape and dtype, the dim is a non-negative number and start,
- * end are valid non-negative number
- */
-template <typename CTYPE_IN, typename CTYPE_OUT>
-void slice_tensor(
-    const Tensor& in,
-    int64_t dim,
-    int64_t start,
-    int64_t end,
-    Tensor& out) {
-  size_t num_values = static_cast<size_t>(end - start);
-  size_t dim_length_in = static_cast<size_t>(in.size(dim));
-  size_t dim_length_out = static_cast<size_t>(out.size(dim));
-  size_t non_negative_start = static_cast<size_t>(start);
-  size_t leading_dims = getLeadingDims(in, dim);
-  size_t trailing_dims = getTrailingDims(in, dim);
-
-  const CTYPE_IN* input_data_base = in.const_data_ptr<CTYPE_IN>();
-  CTYPE_OUT* output_data_base = out.mutable_data_ptr<CTYPE_OUT>();
-
-  for (const auto i : c10::irange(leading_dims)) {
-    const CTYPE_IN* input_data = input_data_base +
-        (i * dim_length_in + non_negative_start) * trailing_dims;
-    CTYPE_OUT* output_data =
-        output_data_base + i * dim_length_out * trailing_dims;
-    for ([[maybe_unused]] const auto j : c10::irange(num_values)) {
-      for (const auto k : c10::irange(trailing_dims)) {
-        output_data[k] = static_cast<CTYPE_OUT>(input_data[k]);
-      }
-      input_data += trailing_dims;
-      output_data += trailing_dims;
-    }
-  }
-}
+};
+
+SplitGLUInputTensor::SplitGLUInputTensor(const Tensor& self, int64_t dim)
+    : half_sizes(get_half_sizes(self, dim)),
+      first_half_impl(
+          self.scalar_type(),
+          self.dim(),
+          half_sizes.data(),
+          self.mutable_data_ptr(),
+          const_cast<executorch::aten::DimOrderType*>(self.dim_order().data()),
+          const_cast<executorch::aten::StridesType*>(self.strides().data()),
+          self.shape_dynamism()),
+      second_half_impl(
+          self.scalar_type(),
+          self.dim(),
+          half_sizes.data(),
+          reinterpret_cast<char*>(self.mutable_data_ptr()) +
+              self.strides()[dim] * self.size(dim) / 2 * self.element_size(),
+          const_cast<executorch::aten::DimOrderType*>(self.dim_order().data()),
+          const_cast<executorch::aten::StridesType*>(self.strides().data()),
+          self.shape_dynamism()),
+      first_half(&first_half_impl),
+      second_half(&second_half_impl) {}
 
 /**
  * Applies the gated linear unit function
@@ -120,11 +75,43 @@ void slice_tensor(
  *  2. The output shall be in float types (Float, Double)
  */
 template <typename CTYPE_IN, typename CTYPE_OUT>
-Tensor& glu_out_tensor(const Tensor& self, int64_t dim, Tensor& out) {
-  const auto self_size = self.size(dim);
-  slice_tensor<CTYPE_IN, CTYPE_OUT>(self, dim, self_size / 2, self_size, out);
-  sigmoid_tensor<CTYPE_OUT>(out);
-  mul_tensors<CTYPE_IN, CTYPE_OUT>(self, dim, out);
+Tensor& glu_out_tensor(
+    KernelRuntimeContext& ctx,
+    const Tensor& self,
+    int64_t dim,
+    Tensor& out) {
+  ET_KERNEL_CHECK(
+      ctx,
+      self.dim() <= static_cast<ssize_t>(kTensorDimensionLimit),
+      InvalidArgument,
+      out);
+  SplitGLUInputTensor split_input(self, dim);
+  ScalarType compute_type =
+      executorch::runtime::isFloatingType(self.scalar_type())
+      ? self.scalar_type()
+      : ScalarType::Float;
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] = "glu.out";
+  ET_SWITCH_FLOATHBF16_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() {
+    utils::apply_bitensor_elementwise_fn<
+        CTYPE_COMPUTE,
+        op_name,
+        utils::SupportedTensorDtypes::FLOATHBF16>(
+        [](const auto val_a, const auto val_b) -> CTYPE_COMPUTE {
+          // TODO: rewrite this to be vectorization-capable? the
+          // tensors might not be contiguous; need to have
+          // apply_bitensor_elementwise_fn check that.
+          const auto one = static_cast<decltype(val_a)>(1.0);
+          return val_a * (one / (one + std::exp(-val_b)));
+        },
+        ctx,
+        split_input.first_half,
+        utils::SupportedTensorDtypes::FLOATHBF16,
+        split_input.second_half,
+        utils::SupportedTensorDtypes::FLOATHBF16,
+        out,
+        utils::internal::SupportNoncontiguousTensors());
+  });
   return out;
 }
 } // namespace
@@ -158,7 +145,7 @@ Tensor& glu_out(
 
   ET_SWITCH_FLOATHBF16_TYPES(in_dtype, ctx, "glu", CTYPE_IN, [&]() {
     ET_SWITCH_FLOATHBF16_TYPES(out.scalar_type(), ctx, "glu", CTYPE_OUT, [&]() {
-      glu_out_tensor<CTYPE_IN, CTYPE_OUT>(self, non_negative_dim, out);
+      glu_out_tensor<CTYPE_IN, CTYPE_OUT>(ctx, self, non_negative_dim, out);
     });
   });
 

kernels/portable/cpu/util/broadcast_indexes_range.h

@@ -43,7 +43,7 @@ inline bool sizes_match_ignoring_leading_1s(
       std::equal(lhs_begin, lhs_end, rhs_begin);
 }
 
-template <std::size_t kNumInputs>
+template <std::size_t kNumInputs, bool support_noncontiguous_tensors = false>
 class BroadcastIndexesIterator {
  public:
   using difference_type = ssize_t;
@@ -57,16 +57,20 @@ class BroadcastIndexesIterator {
   template <typename... Args>
   explicit BroadcastIndexesIterator(const Tensor& output, const Args&... args)
       : output_dim_or_zero_if_no_broadcasting_(
-            (sizes_match_ignoring_leading_1s(args.sizes(), output.sizes()) &&
-             ...)
+            !support_noncontiguous_tensors &&
+                    (sizes_match_ignoring_leading_1s(
+                         args.sizes(),
+                         output.sizes()) &&
+                     ...)
                 ? 0
                 : output.dim()),
         output_shape_(output.sizes()) {
     static_assert(
         sizeof...(args) == kNumInputs && (std::is_same_v<Args, Tensor> && ...),
         "BroadcastIndexesIterator constructor requires kNumInputs input tensor"
         "arguments!");
-    if (output_dim_or_zero_if_no_broadcasting_ != 0) {
+    if (support_noncontiguous_tensors ||
+        output_dim_or_zero_if_no_broadcasting_ != 0) {
       effective_input_broadcast_strides_ = {
           effective_input_broadcast_stride(output, args)...};
     }
@@ -249,11 +253,17 @@ class BroadcastIndexesIterator {
  * Unlike looping using delinearize_index() and
  * linearize_access_indexes(), BroadcastIndexesRange avoids expensive
  * division and modulo operations on each iteration.
+ *
+ * The support_noncontiguous_tensors argument disables an optimization
+ * that causes the iterators not to respect strides in some
+ * cases. This optimization is normally safe because ExecuTorch
+ * tensors are contiguous.
  */
-template <std::size_t kNumInputs>
+template <std::size_t kNumInputs, bool support_noncontiguous_tensors = false>
 class BroadcastIndexesRange {
  public:
-  using iterator = internal::BroadcastIndexesIterator<kNumInputs>;
+  using iterator = internal::
+      BroadcastIndexesIterator<kNumInputs, support_noncontiguous_tensors>;
 
   template <typename... Args>
   BroadcastIndexesRange(const Tensor& output, const Args&... args)