CaptureKernelCall Class — pytorch Architecture

Source Code

aten/src/ATen/core/dispatch/Dispatcher.h lines 646–674

template <typename ReturnType>
struct CaptureKernelCall {
  template <typename F, typename... Args>
  CaptureKernelCall(
      const F& kernel,
      const TypedOperatorHandle<ReturnType(Args...)>& op,
      const DispatchKeySet& dispatchKeySet,
      Args&&... args)
      // Calls the kernel and capture the result in output_.
      : output_{kernel.template call<ReturnType, Args...>(
            op,
            dispatchKeySet,
            std::forward<Args>(args)...)} {}
  // Wraps the return values in a Stack.
  Stack getOutputs() {
    Stack stack;
    impl::push_outputs<ReturnType, false>::copy(output_, &stack);
    return stack;
  }
  // Since we are returning the output_, we don't expect the output_ to be used
  // afterward. Copy elision and RVO do not apply to class data members. Using
  // move semantic to avoid copies when possible.
  ReturnType release() && {
    return std::move(output_);
  }

 private:
  ReturnType output_;
};