cpu_serial_kernel Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/Loops.h lines 357–370

template <typename func_t>
void cpu_serial_kernel(TensorIteratorBase& iter, func_t&& op, const Range& range) {
  using traits = function_traits<func_t>;
  constexpr bool result_void = std::is_void_v<typename traits::result_type>;
  TORCH_INTERNAL_ASSERT(iter.ninputs() == traits::arity &&
                        ((result_void && iter.noutputs() == 0) || (!result_void && iter.noutputs() == 1)));
  // dynamic casting not currently supported on CPU
  TORCH_INTERNAL_ASSERT(!needs_dynamic_casting<func_t>::check(iter));

  iter.serial_for_each([&](char** data, const int64_t* strides, int64_t n) {
    basic_loop(data, strides, 0, n, op);
  }, range);
  iter.cast_outputs();
}