NEML2 2.0.0
All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
SimpleScheduler.h
1// Copyright 2024, UChicago Argonne, LLC
2// All Rights Reserved
3// Software Name: NEML2 -- the New Engineering material Model Library, version 2
4// By: Argonne National Laboratory
5// OPEN SOURCE LICENSE (MIT)
6//
7// Permission is hereby granted, free of charge, to any person obtaining a copy
8// of this software and associated documentation files (the "Software"), to deal
9// in the Software without restriction, including without limitation the rights
10// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11// copies of the Software, and to permit persons to whom the Software is
12// furnished to do so, subject to the following conditions:
13//
14// The above copyright notice and this permission notice shall be included in
15// all copies or substantial portions of the Software.
16//
17// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
23// THE SOFTWARE.
24
25#pragma once
26
27#include "neml2/dispatchers/WorkScheduler.h"
28#include "neml2/base/Registry.h"
29#include "neml2/base/NEML2Object.h"
30#include "neml2/base/Factory.h"
31
32namespace neml2
33{
44{
45public:
48
54 SimpleScheduler(const OptionSet & options);
55
56 std::vector<Device> devices() const override { return {_device}; }
57
58protected:
59 bool schedule_work_impl(Device &, std::size_t &) const override;
60 void dispatched_work_impl(Device, std::size_t) override;
61 void completed_work_impl(Device, std::size_t) override;
62 bool all_work_completed() const override;
63
64private:
66 Device _device;
67
69 std::size_t _batch_size;
70
72 std::size_t _capacity;
73
75 std::size_t _load = 0;
76};
77
78} // namespace neml2
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