WorkDispatcher.h

// Copyright 2024, UChicago Argonne, LLC
// All Rights Reserved
// Software Name: NEML2 -- the New Engineering material Model Library, version 2
// By: Argonne National Laboratory
// OPEN SOURCE LICENSE (MIT)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
 
#pragma once
 
#include <functional>
#include <thread>
#include <queue>
#include <mutex>
#include <condition_variable>
 
#include <ATen/Parallel.h>
 
#include "neml2/dispatchers/WorkGenerator.h"
#include "neml2/dispatchers/WorkScheduler.h"
#include "neml2/misc/assertions.h"
#include "neml2/misc/types.h"
#include "neml2/base/TracingInterface.h"
 
#include "neml2/tensors/R2.h"
 
// Pre-C++20 workaround for std::type_identity
// https://en.cppreference.com/w/cpp/types/type_identity
template <class T>
struct type_identity
{
  using type = T;
};
 
namespace neml2
{
template <typename I,
          typename O,
          typename Of = typename std::vector<O>,
          typename Ip = typename type_identity<I>::type,
          typename Op = typename type_identity<O>::type>
class WorkDispatcher : public TracingInterface
{
public:
  WorkDispatcher(WorkScheduler & scheduler, bool async, std::function<O(I &&, Device)> && do_work)
    : TracingInterface(scheduler),
      _scheduler(scheduler),
      _devices(scheduler.devices()),
      _async(async),
      _do_work(std::move(do_work))
  {
    init_thread_pool();
  }
 
  WorkDispatcher(WorkScheduler & scheduler,
                 bool async,
                 std::function<O(I &&, Device)> && do_work,
                 std::function<O(std::vector<O> &&)> && reduce)
    : TracingInterface(scheduler),
      _scheduler(scheduler),
      _devices(scheduler.devices()),
      _async(async),
      _do_work(std::move(do_work)),
      _reduce(std::move(reduce))
  {
    init_thread_pool();
  }
 
  WorkDispatcher(WorkScheduler & scheduler,
                 bool async,
                 std::function<O(I &&, Device)> && do_work,
                 std::function<Of(std::vector<Op> &&)> && reduce,
                 std::function<I(Ip &&, Device)> && preprocess,
                 std::function<Op(O &&)> && postprocess)
    : TracingInterface(scheduler),
      _scheduler(scheduler),
      _devices(scheduler.devices()),
      _async(async),
      _do_work(std::move(do_work)),
      _reduce(std::move(reduce)),
      _preprocess(std::move(preprocess)),
      _postprocess(std::move(postprocess))
  {
    init_thread_pool();
  }
 
  WorkDispatcher(WorkScheduler & scheduler,
                 bool async,
                 std::function<O(I &&, Device)> && do_work,
                 std::function<Of(std::vector<Op> &&)> && reduce,
                 std::function<I(Ip &&, Device)> && preprocess,
                 std::function<Op(O &&)> && postprocess,
                 std::function<void(Device)> && thread_init)
    : TracingInterface(scheduler),
      _scheduler(scheduler),
      _devices(scheduler.devices()),
      _async(async),
      _do_work(std::move(do_work)),
      _reduce(std::move(reduce)),
      _preprocess(std::move(preprocess)),
      _postprocess(std::move(postprocess)),
      _thread_init(std::move(thread_init))
  {
    if (_thread_init)
      neml_assert_dbg(
          _async, "Custom thread initialization functor is only supported in asynchronous mode");
 
    init_thread_pool();
  }
 
  WorkDispatcher() = delete;
  WorkDispatcher(WorkDispatcher &&) = delete;
  WorkDispatcher(const WorkDispatcher &) = delete;
  WorkDispatcher & operator=(WorkDispatcher &&) = delete;
  WorkDispatcher & operator=(const WorkDispatcher &) = delete;
  ~WorkDispatcher() override { stop_thread_pool(); }

  Of run(WorkGenerator<Ip> &);
 
protected:
  void init_thread_pool();

  void thread_pool_main(const Device &);

  bool should_unlock_thread();

  void stop_thread_pool();

  void validate() const;

  Of run_sync(WorkGenerator<Ip> &);

  Of run_async(WorkGenerator<Ip> &);

  WorkScheduler & _scheduler;

  const std::vector<Device> _devices;

  const bool _async;

  std::function<O(I &&, Device)> _do_work;

  std::function<Of(std::vector<Op> &&)> _reduce;

  std::function<I(Ip &&, Device)> _preprocess;

  std::function<Op(O &&)> _postprocess;

  std::function<void(Device)> _thread_init;

  std::vector<Op> _results;

  std::mutex _qmutex;
  std::condition_variable _thread_condition;
  bool _stop = false;
  // As it turns out it's undefined behavior to initialize this before the mutex and condition
  // variable
  std::vector<std::thread> _thread_pool;
  std::unordered_map<Device, std::queue<std::function<void()>>> _tasks;
};

// Implementation
template <typename I, typename O, typename Of, typename Ip, typename Op>
void
WorkDispatcher<I, O, Of, Ip, Op>::init_thread_pool()
{
  if (!_async)
    return;
 
#ifdef NEML2_HAS_JSON
  if (event_tracing_enabled())
    event_trace_writer().trace_duration_begin("thread pool", "WorkDispatcher");
#endif
 
  // Setup the task queue
  for (const auto & device : _devices)
    _tasks[device] = std::queue<std::function<void()>>();
 
  auto nthread = _devices.size();
  _thread_pool.reserve(nthread);
  for (std::size_t i = 0; i < nthread; ++i)
  {
    // This is necessary to initialize the torch linear algebra library prior to threaded calls
    // See: https://github.com/pytorch/pytorch/issues/90613
    auto res = R2::identity().to(_devices[i]).inverse();
    _thread_pool.emplace_back([this, i] { thread_pool_main(_devices[i]); });
  }
 
  // Initialize the thread
  if (_thread_init)
  {
    for (std::size_t i = 0; i < nthread; ++i)
    {
      auto device = _devices[i];
      auto task = [this, device = device]() mutable
      {
#ifdef NEML2_HAS_JSON
        if (event_tracing_enabled())
          event_trace_writer().trace_duration_begin(
              "thread init", "WorkDispatcher", {{"device", utils::stringify(device)}});
#endif
 
        _thread_init(device);
        _scheduler.completed_work(device, 1);
 
#ifdef NEML2_HAS_JSON
        if (event_tracing_enabled())
          event_trace_writer().trace_duration_end("thread init", "WorkDispatcher");
#endif
      };
      _scheduler.dispatched_work(device, 1);
      {
        std::lock_guard<std::mutex> lock(_qmutex);
        _tasks.at(device).push(task);
      }
      _thread_condition.notify_all();
    }
    _scheduler.wait_for_completion();
  }
}
 
template <typename I, typename O, typename Of, typename Ip, typename Op>
void
WorkDispatcher<I, O, Of, Ip, Op>::thread_pool_main(const Device & device)
{
#ifdef NEML2_HAS_JSON
  if (event_tracing_enabled())
    event_trace_writer().trace_duration_begin(
        "thread main", "WorkDispatcher", {{"device", utils::stringify(device)}});
#endif
 
  while (true)
  {
    std::function<void()> task;
    {
      std::unique_lock<std::mutex> lock(_qmutex);
      _thread_condition.wait(lock, [this, &device] { return _stop || !_tasks.at(device).empty(); });
      if (_stop && _tasks.at(device).empty())
        break;
      task = std::move(_tasks.at(device).front());
      _tasks.at(device).pop();
    }
    task();
  }
 
#ifdef NEML2_HAS_JSON
  if (event_tracing_enabled())
    event_trace_writer().trace_duration_end("thread main", "WorkDispatcher");
#endif
}
 
template <typename I, typename O, typename Of, typename Ip, typename Op>
void
WorkDispatcher<I, O, Of, Ip, Op>::stop_thread_pool()
{
  if (!_async)
    return;
  {
    std::unique_lock<std::mutex> lock(_qmutex);
    _stop = true;
  }
  _thread_condition.notify_all();
  for (auto & thread : _thread_pool)
    thread.join();
 
#ifdef NEML2_HAS_JSON
  if (event_tracing_enabled())
    event_trace_writer().trace_duration_end("thread pool", "WorkDispatcher");
#endif
}
 
template <typename I, typename O, typename Of, typename Ip, typename Op>
Of
WorkDispatcher<I, O, Of, Ip, Op>::run(WorkGenerator<Ip> & generator)
{
#ifdef NEML2_HAS_JSON
  if (event_tracing_enabled())
    event_trace_writer().trace_duration_begin("run", "WorkDispatcher");
#endif
 
  if (_async)
  {
    auto result = run_async(generator);
 
#ifdef NEML2_HAS_JSON
    if (event_tracing_enabled())
      event_trace_writer().trace_duration_end("run", "WorkDispatcher");
#endif
 
    return result;
  }
 
  auto result = run_sync(generator);
 
#ifdef NEML2_HAS_JSON
  if (event_tracing_enabled())
    event_trace_writer().trace_duration_end("run", "WorkDispatcher");
#endif
 
  return result;
}
 
template <typename I, typename O, typename Of, typename Ip, typename Op>
void
WorkDispatcher<I, O, Of, Ip, Op>::validate() const
{
  if (!_do_work)
    throw NEMLException("Do-work function is not set");
 
  if constexpr (!std::is_same_v<I, Ip>)
    if (!_preprocess)
      throw NEMLException("Preprocess function is not set");
 
  if constexpr (!std::is_same_v<O, Op>)
    if (!_postprocess)
      throw NEMLException("Postprocess function is not set");
 
  if constexpr (!std::is_same_v<Of, std::vector<Op>>)
    if (!_reduce)
      throw NEMLException("Reduce function is not set");
}
 
template <typename I, typename O, typename Of, typename Ip, typename Op>
Of
WorkDispatcher<I, O, Of, Ip, Op>::run_sync(WorkGenerator<Ip> & generator)
{
  validate();
 
  Device device = kCPU;
  std::size_t n = 0;
  _results.clear();
  while (generator.has_more())
  {
    _scheduler.schedule_work(device, n);
    if (n <= 0)
      throw NEMLException("Scheduler returned a batch size of " + std::to_string(n));
    // Generate work
    auto && [m, work] = generator.next(n);
    // Preprocess
    if (_preprocess)
      work = _preprocess(std::move(work), device);
    // Do work. Since there is no asynchronous execution, we do not notify the scheduler (this also
    // avoids potential parallel communication inccured by the scheduler)
    auto result = _do_work(std::move(work), device);
    // Postprocess
    if (_postprocess)
      result = _postprocess(std::move(result));
    _results.push_back(result);
  }
 
  if (_reduce)
    return _reduce(std::move(_results));
 
  if constexpr (std::is_same<Of, std::vector<Op>>::value)
    return _results;
 
  throw NEMLException("Internal error: unreachable code");
}
 
template <typename I, typename O, typename Of, typename Ip, typename Op>
Of
WorkDispatcher<I, O, Of, Ip, Op>::run_async(WorkGenerator<Ip> & generator)
{
  validate();
 
  Device device = kCPU;
  std::size_t n = 0;
  _results.clear();
 
  // Keep asking the scheduler for an available device
  // - If the generator has no more work, we break out of the loop
  // - If the scheduler schedules work, we dispatch the work and continue with the dispatching loop
  while (generator.has_more())
  {
    _scheduler.schedule_work(device, n);
    if (n <= 0)
      throw NEMLException("Scheduler returned a batch size of " + std::to_string(n));
    // Generate work
    auto && [m, work] = generator.next(n);
    // Reserve space for the result
    _results.resize(_results.size() + 1);
    auto i = _results.size() - 1;
    // Create the task
    auto task = [this, work = std::move(work), device = device, m = m, i = i]() mutable
    {
#ifdef NEML2_HAS_JSON
      if (event_tracing_enabled())
        event_trace_writer().trace_duration_begin(
            "task", "WorkDispatcher", {{"device", utils::stringify(device)}, {"batch size", m}});
#endif
      // Preprocess
      if (_preprocess)
        work = _preprocess(std::move(work), device);
      // Do work
      auto result = _do_work(std::move(work), device);
      // Postprocess
      if (_postprocess)
        result = _postprocess(std::move(result));
      // Collect result
      _results[i] = std::move(result);
      // Tell the scheduler that we have completed m batches
      _scheduler.completed_work(device, m);
 
#ifdef NEML2_HAS_JSON
      if (event_tracing_enabled())
        event_trace_writer().trace_duration_end("task", "WorkDispatcher");
#endif
    };
    // Tell the scheduler that we have dispatched m batches
    _scheduler.dispatched_work(device, m);
    // Enqueue the task
    {
      std::lock_guard<std::mutex> lock(_qmutex);
      _tasks.at(device).push(task);
    }
    // Notify the thread pool
    // Note: We notify_all instead of notify_one because we want the thread that's bind to the
    // target device to pick up the task
    _thread_condition.notify_all();
  }
 
  // Wait for all tasks to complete
  _scheduler.wait_for_completion();
 
  if (_reduce)
    return _reduce(std::move(_results));
 
  if constexpr (std::is_same<Of, std::vector<Op>>::value)
    return _results;
 
  throw NEMLException("Internal error: unreachable code");
}
} // namespace neml2