ExecutionScheduler.hpp

#pragma once
 
#include "wmtk/TetMesh.h"
#include "wmtk/TriMesh.h"
#include "wmtk/utils/Logger.hpp"
 
// clang-format off
#include <functional>
#include <limits>
#include <wmtk/utils/DisableWarnings.hpp>
#include <tbb/concurrent_priority_queue.h>
#include <tbb/concurrent_queue.h>
#include <tbb/parallel_for.h>
#include <tbb/parallel_reduce.h>
#include <tbb/spin_mutex.h>
#include <tbb/task_arena.h>
#include <tbb/task_group.h>
#include <wmtk/utils/EnableWarnings.hpp>
// clang-format on
 
#include <atomic>
#include <cassert>
#include <cstddef>
#include <queue>
#include <stdexcept>
#include <type_traits>
 
namespace wmtk {
enum class ExecutionPolicy { kSeq, kUnSeq, kPartition, kColor, kMax };
 
using Op = std::string;
 
template <class AppMesh>
struct ExecutePass
{
    using Tuple = typename AppMesh::Tuple;
    std::map<
        Op, // strings
        std::function<std::optional<std::vector<Tuple>>(AppMesh&, const Tuple&)>>
        edit_operation_maps;
    std::function<double(const AppMesh&, Op op, const Tuple&)> priority =
        [](const AppMesh&, Op, const Tuple&) { return 0.; };
    std::function<bool(double)> should_renew = [](double) { return true; };
    std::function<std::vector<std::pair<Op, Tuple>>(const AppMesh&, Op, const std::vector<Tuple>&)>
        renew_neighbor_tuples =
            [](const AppMesh&, Op, const std::vector<Tuple>&) -> std::vector<std::pair<Op, Tuple>> {
        return {};
    };
    std::function<bool(AppMesh&, const Tuple&, int task_id)> lock_vertices =
        [](const AppMesh&, const Tuple&, int task_id) { return true; };
    std::function<bool(const AppMesh&)> stopping_criterion = [](const AppMesh&) {
        return false; // non-stop, process everything
    };
    size_t stopping_criterion_checking_frequency = std::numeric_limits<size_t>::max();
    std::function<bool(const AppMesh&, const std::tuple<double, Op, Tuple>& t)>
        is_weight_up_to_date = [](const AppMesh& m, const std::tuple<double, Op, Tuple>& t) {
            // always do.
            assert(std::get<2>(t).is_valid(m));
            return true;
        };
    std::function<void(const AppMesh&, Op, const Tuple& t)> on_fail =
        [](const AppMesh&, Op, const Tuple& t) {};
 
    ExecutionPolicy policy;
 
    int num_threads = 1;
 
    size_t max_retry_limit = 10;
    ExecutePass(const ExecutionPolicy& policy_ = ExecutionPolicy::kSeq)
        : policy(policy_)
    {
        if constexpr (std::is_base_of<TetMesh, AppMesh>::value) {
            edit_operation_maps = {
                {"edge_collapse",
                 [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     std::vector<Tuple> ret;
                     if (m.collapse_edge(t, ret))
                         return ret;
                     else
                         return {};
                 }},
                {"edge_swap",
                 [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     std::vector<Tuple> ret;
                     if (m.swap_edge(t, ret))
                         return ret;
                     else
                         return {};
                 }},
                {"edge_swap_44",
                 [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     std::vector<Tuple> ret;
                     if (m.swap_edge_44(t, ret))
                         return ret;
                     else
                         return {};
                 }},
                {"edge_swap_56",
                 [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     std::vector<Tuple> ret;
                     if (m.swap_edge_56(t, ret))
                         return ret;
                     else
                         return {};
                 }},
                {"edge_split",
                 [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     std::vector<Tuple> ret;
                     if (m.split_edge(t, ret))
                         return ret;
                     else
                         return {};
                 }},
                {"face_swap",
                 [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     std::vector<Tuple> ret;
                     if (m.swap_face(t, ret))
                         return ret;
                     else
                         return {};
                 }},
                {"vertex_smooth",
                 [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     if (m.smooth_vertex(t))
                         return std::vector<Tuple>{};
                     else
                         return {};
                 }},
                {"face_split",
                 [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     std::vector<Tuple> ret;
                     if (m.split_face(t, ret))
                         return ret;
                     else
                         return {};
                 }},
                {"tet_split", [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     std::vector<Tuple> ret;
                     if (m.split_tet(t, ret))
                         return ret;
                     else
                         return {};
                 }}};
        }
        if constexpr (std::is_base_of<TriMesh, AppMesh>::value) {
            edit_operation_maps = {
                {"edge_collapse",
                 [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     std::vector<Tuple> ret;
                     if (m.collapse_edge(t, ret))
                         return ret;
                     else
                         return {};
                 }},
                {"edge_swap",
                 [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     std::vector<Tuple> ret;
                     if (m.swap_edge(t, ret))
                         return ret;
                     else
                         return {};
                 }},
                {"edge_split",
                 [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     std::vector<Tuple> ret;
                     if (m.split_edge(t, ret))
                         return ret;
                     else
                         return {};
                 }},
                {"vertex_smooth",
                 [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     if (m.smooth_vertex(t))
                         return std::vector<Tuple>{};
                     else
                         return {};
                 }},
                {"face_split", [](AppMesh& m, const Tuple& t) -> std::optional<std::vector<Tuple>> {
                     std::vector<Tuple> ret;
                     if (m.split_face(t, ret))
                         return ret;
                     else
                         return {};
                 }}};
        }
    };
 
    ExecutePass(ExecutePass&) = delete;
 
private:
    void operation_cleanup(AppMesh& m)
    { //
        // class ResourceManger
        // what about RAII mesh edit locking?
        // release mutex, but this should be implemented in TetMesh class.
        if (policy == ExecutionPolicy::kSeq)
            return;
        else {
            m.release_vertex_mutex_in_stack();
        }
    }
 
    size_t get_partition_id(const AppMesh& m, const Tuple& e)
    {
        if (policy == ExecutionPolicy::kSeq) {
            return 0;
        }
        return m.get_partition_id(e);
    }
 
public:
    bool operator()(AppMesh& m, const std::vector<std::pair<Op, Tuple>>& operation_tuples)
    {
        using Elem = std::tuple<double, Op, Tuple, size_t>; // priority, operation, tuple, #retries
        using Queue = tbb::concurrent_priority_queue<Elem>;
 
        std::atomic<bool> stop(false);
        cnt_success = 0;
        cnt_fail = 0;
        cnt_update = 0;
 
        std::vector<Queue> queues(num_threads);
        Queue final_queue;
 
        auto run_single_queue = [&](Queue& Q, int task_id) {
            Elem ele_in_queue;
            while ([&]() { return Q.try_pop(ele_in_queue); }()) {
                auto& [weight, op, tup, retry] = ele_in_queue;
                if (!tup.is_valid(m)) {
                    continue;
                }
 
                std::vector<Elem> renewed_elements;
                {
                    auto locked_vid = lock_vertices(
                        m,
                        tup,
                        task_id); // Note that returning `Tuples` would be invalid.
                    if (!locked_vid) {
                        retry++;
                        if (retry < max_retry_limit) {
                            Q.emplace(ele_in_queue);
                        } else {
                            retry = 0;
                            final_queue.emplace(ele_in_queue);
                        }
                        continue;
                    }
                    if (tup.is_valid(m)) {
                        if (!is_weight_up_to_date(
                                m,
                                std::tuple<double, Op, Tuple>(weight, op, tup))) {
                            operation_cleanup(m);
                            continue;
                        } // this can encode, in qslim, recompute(energy) == weight.
                        auto newtup = edit_operation_maps[op](m, tup);
                        std::vector<std::pair<Op, Tuple>> renewed_tuples;
                        if (newtup) {
                            renewed_tuples = renew_neighbor_tuples(m, op, newtup.value());
                            cnt_success++;
                            cnt_update++;
                        } else {
                            on_fail(m, op, tup);
                            cnt_fail++;
                        }
                        for (const auto& [o, e] : renewed_tuples) {
                            auto val = priority(m, o, e);
                            if (should_renew(val)) {
                                renewed_elements.emplace_back(val, o, e, 0);
                            }
                        }
                    }
                    operation_cleanup(m); // Maybe use RAII
                }
                for (auto& e : renewed_elements) {
                    Q.emplace(e);
                }
 
                if (stop.load(std::memory_order_acquire)) {
                    return;
                }
                if (cnt_success > stopping_criterion_checking_frequency) {
                    if (stopping_criterion(m)) {
                        stop.store(true);
                        return;
                    }
                    cnt_update.store(0, std::memory_order_release);
                }
            }
        };
 
        if (policy == ExecutionPolicy::kSeq) {
            for (const auto& [op, e] : operation_tuples) {
                if (!e.is_valid(m)) {
                    continue;
                }
                final_queue.emplace(priority(m, op, e), op, e, 0);
            }
            run_single_queue(final_queue, 0);
        } else {
            for (const auto& [op, e] : operation_tuples) {
                if (!e.is_valid(m)) {
                    continue;
                }
                queues[get_partition_id(m, e)].emplace(priority(m, op, e), op, e, 0);
            }
            // Comment out parallel: work on serial first.
            tbb::task_arena arena(num_threads);
            tbb::task_group tg;
            arena.execute([&queues, &run_single_queue, &tg]() {
                for (int task_id = 0; task_id < queues.size(); task_id++) {
                    tg.run([&run_single_queue, &queues, task_id] {
                        run_single_queue(queues[task_id], task_id);
                    });
                }
                tg.wait();
            });
            logger().debug("Parallel Complete, remains element {}", final_queue.size());
            run_single_queue(final_queue, 0);
        }
 
        logger().info(
            "executed: {} | success / fail: {} / {}",
            (int)cnt_success + (int)cnt_fail,
            (int)cnt_success,
            (int)cnt_fail);
        return true;
    }
 
    int get_cnt_success() const { return cnt_success; }
    int get_cnt_fail() const { return cnt_fail; }
 
private:
    std::atomic_int cnt_update = 0;
    std::atomic_int cnt_success = 0;
    std::atomic_int cnt_fail = 0;
};
} // namespace wmtk