Add op: gather.out

Differential Revision: D61822105

Pull Request resolved: #4939

Author: manuelcandales

kernels/aten/functions.yaml

@@ -173,6 +173,8 @@
 
 - op: full.out
 
+- op: gather.out
+
 - op: ge.Scalar_out
 
 - op: ge.Tensor_out

kernels/portable/cpu/op_gather.cpp

@@ -0,0 +1,98 @@
+/*
+ * 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 <cinttypes>
+#include <cstdint>
+#include <cstring>
+
+#include <executorch/kernels/portable/cpu/util/index_util.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+
+namespace torch {
+namespace executor {
+namespace native {
+
+using Tensor = exec_aten::Tensor;
+using ScalarType = exec_aten::ScalarType;
+
+namespace {
+
+template <typename CTYPE>
+void gather_helper(
+    const Tensor& in,
+    const Tensor& index,
+    Tensor& out,
+    int64_t dim) {
+  const CTYPE* in_data = in.const_data_ptr<CTYPE>();
+  const long* index_data = index.const_data_ptr<long>();
+  CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
+
+  if (index.dim() == 0) {
+    out_data[0] = in_data[index_data[0]];
+    return;
+  }
+
+  for (size_t ix = 0; ix < index.numel(); ++ix) {
+    size_t ix_coord[kTensorDimensionLimit];
+    indexToCoordinate(index, ix, ix_coord);
+
+    size_t in_coord[kTensorDimensionLimit];
+    for (size_t i = 0; i < out.dim(); ++i) {
+      if (i == dim) {
+        in_coord[i] = index_data[ix];
+      } else {
+        in_coord[i] = ix_coord[i];
+      }
+    }
+
+    size_t in_ix = coordinateToIndex(in, in_coord);
+    size_t out_ix = coordinateToIndex(out, ix_coord);
+
+    out_data[out_ix] = in_data[in_ix];
+  }
+}
+
+} // namespace
+
+Tensor& gather_out(
+    RuntimeContext& ctx,
+    const Tensor& in,
+    int64_t dim,
+    const Tensor& index,
+    bool sparse_grad,
+    Tensor& out) {
+  (void)ctx;
+
+  ET_KERNEL_CHECK(
+      ctx,
+      check_gather_args(in, dim, index, sparse_grad, out),
+      InvalidArgument,
+      out);
+
+  if (dim < 0) {
+    dim += nonzero_dim(in);
+  }
+
+  ET_KERNEL_CHECK(
+      ctx,
+      resize_tensor(out, index.sizes()) == Error::Ok,
+      InvalidArgument,
+      out);
+
+  constexpr auto name = "gather.out";
+
+  ET_SWITCH_REALHB_TYPES(in.scalar_type(), ctx, name, CTYPE, [&]() {
+    gather_helper<CTYPE>(in, index, out, dim);
+  });
+
+  return out;
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch

kernels/portable/cpu/util/index_util.cpp

@@ -12,6 +12,51 @@
 namespace torch {
 namespace executor {
 
+bool check_gather_args(
+    const Tensor& in,
+    int64_t dim,
+    const Tensor& index,
+    bool sparse_grad,
+    Tensor& out) {
+  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_MSG_AND_RETURN_IF_FALSE(
+      index.scalar_type() == ScalarType::Long,
+      "Expected dypte int64 for index");
+  if (index.numel() != 0) {
+    ET_LOG_MSG_AND_RETURN_IF_FALSE(
+        nonzero_dim(in) == nonzero_dim(index),
+        "self and index should have the same dimensionality when index is not empty "
+        "except for the case when one has dimension 0 and the other has dimension 1");
+  }
+
+  // Normalize dim to non-negative value
+  if (dim < 0) {
+    dim += nonzero_dim(in);
+  }
+
+  for (size_t d = 0; d < nonzero_dim(in); ++d) {
+    if (d != dim) {
+      ET_LOG_MSG_AND_RETURN_IF_FALSE(
+          nonempty_size(index, d) <= nonempty_size(in, d),
+          "size of dimension %zd of index should be smaller than the size of that dimension of input if dimension %zd != dim %zd",
+          d,
+          d,
+          (size_t)dim);
+    }
+  }
+  const long* index_data = index.const_data_ptr<long>();
+  for (size_t i = 0; i < index.numel(); ++i) {
+    ET_LOG_MSG_AND_RETURN_IF_FALSE(
+        index_data[i] >= 0 && index_data[i] < nonempty_size(in, dim),
+        "Index is out of bounds for dimension %zd with size %zd",
+        (size_t)dim,
+        nonempty_size(index, dim));
+  }
+
+  return true;
+}
+
 bool check_index_select_args(
     const Tensor& in,
     int64_t dim,

kernels/portable/cpu/util/index_util.h

@@ -14,6 +14,13 @@
 namespace torch {
 namespace executor {
 
+bool check_gather_args(
+    const Tensor& in,
+    int64_t dim,
+    const Tensor& index,
+    bool sparse_grad,
+    Tensor& output);
+
 bool check_index_select_args(
     const Tensor& in,
     int64_t dim,

kernels/portable/functions.yaml

@@ -392,6 +392,11 @@
     - arg_meta: null
       kernel_name: torch::executor::full_like_out
 
+- op: gather.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::gather_out
+
 - op: ge.Scalar_out
   kernels:
     - arg_meta: null