27#include "neml2/dispatchers/WorkScheduler.h"
28#include "neml2/base/Registry.h"
29#include "neml2/base/NEML2Object.h"
30#include "neml2/base/Factory.h"
56 std::vector<Device>
devices()
const override {
return {_device}; }
69 std::size_t _batch_size;
72 std::size_t _capacity;
75 std::size_t _load = 0;
A custom map-like data structure. The keys are strings, and the values can be nonhomogeneously typed.
Definition OptionSet.h:52
bool schedule_work_impl(Device &, std::size_t &) const override
Implementation of the work scheduling.
Definition SimpleScheduler.cxx:62
void completed_work_impl(Device, std::size_t) override
Update the scheduler with the completion of the last batch.
Definition SimpleScheduler.cxx:79
SimpleScheduler(const OptionSet &options)
Construct a new WorkScheduler object.
Definition SimpleScheduler.cxx:52
void dispatched_work_impl(Device, std::size_t) override
Update the scheduler with the dispatch of the last batch.
Definition SimpleScheduler.cxx:73
std::vector< Device > devices() const override
Device options.
Definition SimpleScheduler.h:56
static OptionSet expected_options()
Options for the scheduler.
Definition SimpleScheduler.cxx:34
bool all_work_completed() const override
Check if all work has been completed.
Definition SimpleScheduler.cxx:86
WorkScheduler(const OptionSet &options)
Construct a new WorkScheduler object.
Definition WorkScheduler.cxx:39
Definition DiagnosticsInterface.cxx:30
c10::Device Device
Definition types.h:66