VertexColoring.hpp

#pragma once
 
#include <wmtk/threading/dynamic_parallel_for.hpp>
#include <wmtk/threading/task_group.hpp>
 
#include <algorithm>
#include <atomic>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <exception>
#include <memory>
#include <mutex>
#include <thread>
#include <vector>
 
namespace wmtk::utils {
 
template <typename NeighborFn>
std::vector<std::vector<size_t>> greedy_vertex_coloring(
    const std::vector<size_t>& vids,
    size_t vert_capacity,
    int num_threads,
    NeighborFn&& neighbors)
{
    constexpr uint32_t kNone = UINT32_MAX;
    std::vector<uint32_t> slot(vert_capacity, kNone);
    for (size_t i = 0; i < vids.size(); ++i) {
        // A repeated vertex would not see its other copy as a neighbour and could land twice in
        // one class, smoothed by two threads at once.
        assert(slot[vids[i]] == kNone && "greedy_vertex_coloring: a vertex is listed twice");
        slot[vids[i]] = uint32_t(i);
    }
 
    // Each entry is written by the one task that owns its index.
    std::vector<std::vector<uint32_t>> adjacency(vids.size());
    threading::dynamic_parallel_for(vids.size(), num_threads, 256, [&](size_t b, size_t e) {
        std::vector<size_t> nb;
        for (size_t i = b; i < e; ++i) {
            nb.clear();
            neighbors(vids[i], nb);
            for (const size_t u : nb) {
                if (u < slot.size() && slot[u] != kNone && slot[u] != i) {
                    adjacency[i].push_back(slot[u]);
                }
            }
        }
    });
 
    // Greedy in `vids` order: the smallest color no already-colored neighbor holds.
    std::vector<uint32_t> color(vids.size(), kNone);
    std::vector<size_t> seen_at; // seen_at[c] == i + 1: a neighbor of vertex i holds color c
    uint32_t n_colors = 0;
    for (size_t i = 0; i < vids.size(); ++i) {
        for (const uint32_t j : adjacency[i]) {
            const uint32_t c = color[j];
            if (c != kNone) {
                seen_at[c] = i + 1;
            }
        }
        uint32_t c = 0;
        while (c < n_colors && seen_at[c] == i + 1) {
            ++c;
        }
        if (c == n_colors) {
            ++n_colors;
            seen_at.push_back(0);
        }
        color[i] = c;
    }
 
    std::vector<std::vector<size_t>> classes(n_colors);
    for (size_t i = 0; i < vids.size(); ++i) {
        classes[color[i]].push_back(vids[i]);
    }
    return classes;
}
 
inline size_t color_class_chunk(size_t n, int num_threads, size_t max_chunk)
{
    const size_t per = n / (8 * size_t(num_threads > 0 ? num_threads : 1));
    return per < 1 ? 1 : (per > max_chunk ? max_chunk : per);
}
 
template <typename F>
size_t for_each_in_classes(
    const std::vector<std::vector<size_t>>& classes,
    int num_threads,
    size_t max_chunk,
    F&& fn)
{
    size_t n_total = 0;
    for (const auto& cls : classes) n_total += cls.size();
    if (n_total == 0) return 0;
    const size_t nt = size_t(num_threads > 1 ? num_threads : 1);
    if (nt == 1) {
        size_t successes = 0;
        for (const auto& cls : classes) {
            for (const size_t v : cls) {
                if (fn(v)) ++successes;
            }
        }
        return successes;
    }
 
    // next[c]: the first vertex of class c not yet handed out.
    auto next = std::make_unique<std::atomic<size_t>[]>(classes.size());
    for (size_t c = 0; c < classes.size(); ++c) next[c].store(0, std::memory_order_relaxed);
    // A reusable barrier: the last to arrive starts the next generation.
    std::atomic<size_t> arrived{0};
    std::atomic<size_t> generation{0};
    const auto barrier = [&] {
        const size_t gen = generation.load(std::memory_order_acquire);
        if (arrived.fetch_add(1, std::memory_order_acq_rel) + 1 == nt) {
            arrived.store(0, std::memory_order_relaxed);
            generation.fetch_add(1, std::memory_order_release);
        } else {
            while (generation.load(std::memory_order_acquire) == gen) {
                std::this_thread::yield();
            }
        }
    };
    std::atomic<size_t> successes{0};
    std::atomic<bool> failed{false};
    std::exception_ptr error;
    std::mutex error_mutex;
    // The start gate: kLaunching until every task is launched, then kGo -- or kAbort if a launch
    // failed. Declared before the task_group, which waits for the tasks when it is destroyed.
    enum : int { kLaunching = 0, kGo = 1, kAbort = 2 };
    std::atomic<int> gate{kLaunching};
 
    threading::task_group tg;
    const auto task = [&] {
        int state;
        while ((state = gate.load(std::memory_order_acquire)) == kLaunching) {
            std::this_thread::yield();
        }
        if (state == kAbort) return;
        size_t mine = 0;
        for (size_t c = 0; c < classes.size(); ++c) {
            const std::vector<size_t>& cls = classes[c];
            const size_t chunk = color_class_chunk(cls.size(), int(nt), max_chunk);
            while (!failed.load(std::memory_order_relaxed)) {
                const size_t b = next[c].fetch_add(chunk, std::memory_order_relaxed);
                if (b >= cls.size()) break;
                const size_t e = std::min(cls.size(), b + chunk);
                try {
                    for (size_t i = b; i < e; ++i) {
                        if (fn(cls[i])) ++mine;
                    }
                } catch (...) {
                    std::lock_guard<std::mutex> lock(error_mutex);
                    if (!error) error = std::current_exception();
                    failed.store(true, std::memory_order_relaxed);
                }
            }
            // Not after the last class: tg.wait() is the barrier there.
            if (c + 1 < classes.size()) barrier();
        }
        successes.fetch_add(mine, std::memory_order_relaxed);
    };
    try {
        for (size_t t = 0; t < nt; ++t) {
            tg.run(task);
        }
    } catch (...) {
        gate.store(kAbort, std::memory_order_release);
        throw;
    }
    gate.store(kGo, std::memory_order_release);
    tg.wait();
    if (error) std::rethrow_exception(error);
    return successes.load();
}
 
template <typename Mesh>
size_t
smooth_classes_checked(Mesh& m, const std::vector<std::vector<size_t>>& classes, size_t max_chunk)
{
    struct Restore
    {
        bool& flag;
        bool saved;
        ~Restore() { flag = saved; }
    } restore{m.m_check_smoothing_stays_in_star, m.m_check_smoothing_stays_in_star},
        restore_pass{m.m_in_operation_pass, m.m_in_operation_pass};
    m.m_check_smoothing_stays_in_star = true;
    // A pass is running (see TetOptimizerMesh::m_in_operation_pass); this sweep does not go
    // through run_pass().
    m.m_in_operation_pass = true;
    return for_each_in_classes(classes, m.NUM_THREADS, max_chunk, [&m](size_t v) {
        return m.smooth_vertex(m.tuple_from_vertex(v));
    });
}
 
} // namespace wmtk::utils