parallelTaskSync.h

Go to the documentation of this file.

//
// Copyright 2025 Pixar
//
// Licensed under the terms set forth in the LICENSE.txt file available at
// https://openusd.org/license.
//
#ifndef PXR_EXEC_VDF_PARALLEL_TASK_SYNC_H
#define PXR_EXEC_VDF_PARALLEL_TASK_SYNC_H
 
 
#include "pxr/pxr.h"
 
#include "pxr/exec/vdf/api.h"
#include "pxr/exec/vdf/parallelTaskWaitlist.h"
 
#include "pxr/base/work/taskGraph.h"
 
#include <atomic>
#include <cstdint>
#include <memory>
 
PXR_NAMESPACE_OPEN_SCOPE
 
class VdfParallelTaskSync
{
public:
    VdfParallelTaskSync(const VdfParallelTaskSync &) = delete;
    VdfParallelTaskSync &operator=(const VdfParallelTaskSync &) = delete;
 
    VdfParallelTaskSync(WorkTaskGraph *taskGraph) 
        : _waitlists(1000)
        , _num(0)
        , _taskGraph(taskGraph)
    {}
 
    VDF_API
    void Reset(const size_t num);
 
    enum class State {
        Done,       
 
        Wait,       
 
        Claimed     
    };
 
    inline State Claim(const size_t idx, WorkTaskGraph::BaseTask *successor);
 
    inline void MarkDone(const size_t idx);
 
private:
 
    // The different states a task can be in. 
    enum _TaskState : uint8_t {
        _TaskStateUnclaimed,
        _TaskStateClaimed,
        _TaskStateDone
    };
 
    // A byte-array indicating the state of each task.
    std::unique_ptr<std::atomic<uint8_t>[]> _state;
 
    // A pointer to the waiting queue head for each task.
    std::unique_ptr<VdfParallelTaskWaitlist::HeadPtr[]> _waiting;
 
    // The waitlist instance for managing the queues.
    VdfParallelTaskWaitlist _waitlists;
 
    // The number of tasks in this graph.
    size_t _num;
 
    // The task graph for running pending tasks. 
    WorkTaskGraph *_taskGraph;
 
};
 
 
VdfParallelTaskSync::State
VdfParallelTaskSync::Claim(const size_t idx, WorkTaskGraph::BaseTask *successor)
{
    // Get the current task state.
    uint8_t state = _state[idx].load(std::memory_order_acquire);
 
    // If the task has completed, bail out.
    if (state == _TaskStateDone) {
        return State::Done;
    }
 
    // If the task has not been claimed, yet, attempt to atomically claim it
    // now. Return if this succeeds.
    else if (state == _TaskStateUnclaimed && 
        _state[idx].compare_exchange_strong(state, _TaskStateClaimed)) {
        return State::Claimed;
    }
    
    // If we have to wait, try to enqueue in the waiting list, but bail
    // out if the task completes while attempting to do so.
    return _waitlists.WaitOn(&_waiting[idx], successor)
        ? State::Wait
        : State::Done;
}
 
void
VdfParallelTaskSync::MarkDone(const size_t idx)
{
    // Mark the task done in the state array.
    _state[idx].store(_TaskStateDone, std::memory_order_release);
 
    // Close the corresponding wait list and notify all waiting tasks.
    _waitlists.CloseAndNotify(&_waiting[idx], _taskGraph);
}
 
PXR_NAMESPACE_CLOSE_SCOPE
 
#endif