Modernize indexing kernels

Reviewed By: manuelcandales

Differential Revision: D48405501

fbshipit-source-id: 672ce23825a364200949e66355b4b9636ed38963

Author: SS-JIA

kernels/portable/cpu/op_index.cpp

@@ -22,29 +22,13 @@ using Tensor = exec_aten::Tensor;
 
 namespace {
 
-void check_index_args(
-    const Tensor& input,
-    exec_aten::ArrayRef<exec_aten::optional<Tensor>> indices,
-    Tensor& output) {
-  // size of indices must not exceed the number of dimensions
-  ET_CHECK_MSG(
-      indices.size() <= input.dim(),
-      "indices.size() %zd > input.dim() %zd",
-      ssize_t(indices.size()),
-      ssize_t(input.dim()));
-
-  check_indices(input, indices);
-
-  check_index_result_size(input, indices, output);
-}
-
 template <typename CTYPE_IN, typename CTYPE_OUT>
 void index_out_impl_mask(
-    const Tensor& input,
+    const Tensor& in,
     exec_aten::ArrayRef<exec_aten::optional<Tensor>> indices,
     Tensor& out) {
   // Data pointers
-  const CTYPE_IN* const in_data = input.const_data_ptr<CTYPE_IN>();
+  const CTYPE_IN* const in_data = in.const_data_ptr<CTYPE_IN>();
   CTYPE_OUT* const out_data = out.mutable_data_ptr<CTYPE_OUT>();
 
   const Tensor& mask = indices[0].value();
@@ -60,11 +44,11 @@ void index_out_impl_mask(
 
 template <typename CTYPE_IN, typename CTYPE_OUT>
 void index_out_impl_list(
-    const Tensor& input,
+    const Tensor& in,
     exec_aten::ArrayRef<exec_aten::optional<Tensor>> indices,
     Tensor& out) {
   // Data pointers
-  const CTYPE_IN* const in_data = input.const_data_ptr<CTYPE_IN>();
+  const CTYPE_IN* const in_data = in.const_data_ptr<CTYPE_IN>();
   CTYPE_OUT* dst = out.mutable_data_ptr<CTYPE_OUT>();
 
   size_t num_idx_queries = get_indices_broadcast_len(indices);
@@ -73,13 +57,13 @@ void index_out_impl_list(
 
     // For each index query, align the src and dst pointers to the position
     // described by the query.
-    size_t offset = get_index_query_pos_offset(idx, input, indices);
+    size_t offset = get_index_query_pos_offset(idx, in, indices);
     src += offset;
 
     // Calculate the region of data to copy for this query.
     // For example, a 2x4x3x5 tensor indexing at [1, 1, :, :] should copy 15
     // elements.
-    size_t copy_len = getTrailingDims(input, indices.size() - 1);
+    size_t copy_len = getTrailingDims(in, indices.size() - 1);
 
     for (size_t i = 0; i < copy_len; ++i) {
       dst[i] = static_cast<CTYPE_OUT>(src[i]);
@@ -88,107 +72,50 @@ void index_out_impl_list(
   }
 }
 
-template <typename CTYPE_IN, typename CTYPE_OUT>
-void index_out_impl(
-    const Tensor& input,
-    exec_aten::ArrayRef<exec_aten::optional<Tensor>> indices,
-    Tensor& out) {
-  if (is_index_mask(input, indices)) {
-    index_out_impl_mask<CTYPE_IN, CTYPE_OUT>(input, indices, out);
-  } else {
-    index_out_impl_list<CTYPE_IN, CTYPE_OUT>(input, indices, out);
-  }
-}
-
-template <typename CTYPE_IN>
-inline void index_out_switch_out(
-    const Tensor& input,
-    exec_aten::ArrayRef<exec_aten::optional<Tensor>> indices,
-    Tensor& out) {
-  auto out_type = out.scalar_type();
-#define INDEX_COPY_SWITCH_OUTPUT_CASE(ctype, dtype)       \
-  case ScalarType::dtype:                                 \
-    index_out_impl<CTYPE_IN, ctype>(input, indices, out); \
-    break;
-
-  switch (out_type) {
-    ET_FORALL_REAL_TYPES_AND(Bool, INDEX_COPY_SWITCH_OUTPUT_CASE);
-    default:
-      ET_CHECK_MSG(
-          false, "%hhd scalar type is not supported for output", out_type);
-  }
-
-#undef INDEX_COPY_SWITCH_OUTPUT_CASE
-}
-
-inline void index_out_switch_input(
-    const Tensor& input,
-    exec_aten::ArrayRef<exec_aten::optional<Tensor>> indices,
-    Tensor& out) {
-  auto input_type = input.scalar_type();
-#define INDEX_COPY_SWITCH_INPUT_CASE(ctype, dtype)    \
-  case ScalarType::dtype:                             \
-    index_out_switch_out<ctype>(input, indices, out); \
-    break;
-
-  switch (input_type) {
-    ET_FORALL_REAL_TYPES_AND(Bool, INDEX_COPY_SWITCH_INPUT_CASE);
-    default:
-      ET_CHECK_MSG(
-          false, "%hhd scalar type is not supported for input", input_type);
-  }
-
-#undef INDEX_COPY_SWITCH_INPUT_CASE
-}
-
-// expected output dim: 1 + (remaining dimension). Shape: [indices.size,
-// *remaining dimension shape]. E.g., 3x3x3x3 tensor, index at [(1, 2), (0,
-// 2), :, :] gives output shape [2, 3, 3].
-Error resize_out(
-    const Tensor& input,
-    Tensor& out,
-    ArrayRef<exec_aten::optional<Tensor>> indices) {
-  size_t out_ndim = 0;
-  Tensor::SizesType out_sizes[kTensorDimensionLimit];
-  get_index_result_size(input, indices, out_sizes, out_ndim);
-
-  ArrayRef<Tensor::SizesType> output_size{out_sizes, out_ndim};
-  auto error = resize_tensor(out, output_size);
-
-  return error;
-}
 } // namespace
 
-/// aten::index.Tensor_out(Tensor self, Tensor?[] indices, *, Tensor(a!) out) ->
-/// Tensor(a!)
 Tensor& index_Tensor_out(
     RuntimeContext& ctx,
-    const Tensor& input,
+    const Tensor& in,
     exec_aten::ArrayRef<exec_aten::optional<Tensor>> indices,
     Tensor& out) {
-  (void)ctx;
+  ET_KERNEL_CHECK(
+      ctx, check_index_args(in, indices, out), InvalidArgument, out);
 
   if (indices.empty()) {
-    auto error = resize_tensor(out, input.sizes());
-    ET_CHECK_MSG(error == Error::Ok, "Failed to resize output tensor.");
+    ET_KERNEL_CHECK(
+        ctx, resize_tensor(out, in.sizes()) == Error::Ok, InvalidArgument, out);
     memcpy(
-        out.mutable_data_ptr<char>(),
-        input.const_data_ptr<char>(),
-        input.nbytes());
+        out.mutable_data_ptr<char>(), in.const_data_ptr<char>(), in.nbytes());
     return out;
   }
 
-  // resize out tensor
-  auto error = resize_out(input, out, indices);
-  // TODO: Construct error message with requested output sizes.
-  ET_CHECK_MSG(error == Error::Ok, "Failed to resize output tensor.");
-  check_index_args(input, indices, out);
+  size_t expected_ndim = 0;
+  Tensor::SizesType expected_size[kTensorDimensionLimit];
+  get_index_out_target_size(in, indices, expected_size, &expected_ndim);
+  ET_KERNEL_CHECK(
+      ctx,
+      resize_tensor(out, {expected_size, expected_ndim}) == Error::Ok,
+      InvalidArgument,
+      out);
+
+  check_index_args(in, indices, out);
 
-  if (input.numel() == 0) {
+  if (in.numel() == 0) {
     return out;
   }
 
-  index_out_switch_input(input, indices, out);
+  ScalarType in_type = in.scalar_type();
+  ScalarType out_type = out.scalar_type();
+  ET_SWITCH_REAL_TYPES_AND(Bool, in_type, ctx, "index", CTYPE_IN, [&]() {
+    ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, "index", CTYPE_OUT, [&]() {
+      if (is_index_mask(in, indices)) {
+        index_out_impl_mask<CTYPE_IN, CTYPE_OUT>(in, indices, out);
+      } else {
+        index_out_impl_list<CTYPE_IN, CTYPE_OUT>(in, indices, out);
+      }
+    });
+  });
 
   return out;
 }