multiple_outputs_loop Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/Loops.h lines 169–192

template <typename func_t>
inline void
multiple_outputs_loop(char* C10_RESTRICT data[], const int64_t* strides_, int64_t i, int64_t n, func_t&& op) {
  using traits = function_traits<func_t>;

  using result_type = typename traits::result_type;
  constexpr int num_outputs = std::tuple_size_v<result_type>;
  constexpr int ntensors = traits::arity + num_outputs;

  // Copying strides to temporary array helps auto vectorization in older GCC
  // versions.
  int64_t strides[ntensors];
  for (const auto arg : c10::irange(ntensors)) {
    strides[arg] = strides_[arg];
  }

  for (; i < n; i++) {
    auto output = std::apply(op, dereference<traits>(
      &data[num_outputs],
      &strides[num_outputs],
      i));
    handle_tuple_outputs(data, strides, i, output);
  }
}