parallelExecutorDataVector.h

//
// 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_EXECUTOR_DATA_VECTOR_H
#define PXR_EXEC_VDF_PARALLEL_EXECUTOR_DATA_VECTOR_H
 
#include "pxr/pxr.h"
 
#include "pxr/exec/vdf/api.h"
#include "pxr/exec/vdf/executorBufferData.h"
#include "pxr/exec/vdf/executorInvalidationData.h"
#include "pxr/exec/vdf/output.h"
#include "pxr/exec/vdf/types.h"
 
#include "pxr/base/work/zeroAllocator.h"
 
#include <tbb/concurrent_vector.h>
#include <tbb/spin_mutex.h>
 
#include <atomic>
#include <climits>
#include <cstdint>
#include <memory>
 
PXR_NAMESPACE_OPEN_SCOPE
 
class VdfNetwork;
 
class Vdf_ParallelExecutorDataVector
{
public:
 
    using DataHandle = size_t;
 
    static const size_t InvalidHandle = size_t(-1);
 
    Vdf_ParallelExecutorDataVector() :
        _numSegments(0),
        _numTouched(0)
    {}
 
    VDF_API
    ~Vdf_ParallelExecutorDataVector();
 
    VDF_API
    void Resize(const VdfNetwork &network);
 
    inline DataHandle GetOrCreateDataHandle(const VdfId outputId) const;
 
    inline DataHandle GetDataHandle(const VdfId outputId) const;
 
    VdfExecutorBufferData *GetPrivateBufferData(const DataHandle handle) const {
        const uint8_t idx = _outputData[handle].bufferIndices.GetPrivateIndex();
        return &_bufferData[idx][handle];
    }
 
    VdfExecutorBufferData *GetScratchBufferData(const DataHandle handle) const {
        const uint8_t idx = _outputData[handle].bufferIndices.GetScratchIndex();
        return &_bufferData[idx][handle];
    }
 
    VdfExecutorBufferData *GetPublicBufferData(const DataHandle handle) const {
        const uint8_t idx = _outputData[handle].bufferIndices.GetPublicIndex();
        return &_bufferData[idx][handle];
    }
 
    void PublishPrivateBufferData(const DataHandle handle) const {
        _outputData[handle].bufferIndices.PublishPrivateIndex();    
    }
 
    void PublishScratchBufferData(const DataHandle handle) const {
        _outputData[handle].bufferIndices.PublishScratchIndex();
    }
 
    VdfExecutorBufferData *GetTransferredBufferData(
        const DataHandle handle) const {
        return _inboxes[handle].Get();
    }
 
    bool TransferBufferData(
        const DataHandle handle,
        VdfVector *value,
        const VdfMask &mask) {
        return _inboxes[handle].Take(value, mask);
    }
 
    void ResetTransferredBufferData(const DataHandle handle) {
        _inboxes[handle].Reset();
    }
 
    VdfExecutorInvalidationData *GetInvalidationData(
        const DataHandle handle) const {
        return &_invalidationData[handle];
    }
 
    VdfInvalidationTimestamp GetInvalidationTimestamp(
        const DataHandle handle) const {
        return _outputData[handle].invalidationTimestamp;
    }
 
    void SetInvalidationTimestamp(
        const DataHandle handle,
        VdfInvalidationTimestamp timestamp) {
        _outputData[handle].invalidationTimestamp = timestamp;
    }
 
    bool IsTouched(const VdfId outputId) const {
        const VdfIndex outputIdx = VdfOutput::GetIndexFromId(outputId);
        const uint32_t idx = outputIdx / _TouchedWordBits;
        const uint64_t bit =
            UINT64_C(1) << (outputIdx & (_TouchedWordBits - 1));
        return
            idx < _numTouched.load(std::memory_order_acquire) &&
            (_touched[idx].load(std::memory_order_relaxed) & bit) != 0;
    }
 
    void Touch(const VdfId outputId) {
        // Lazily initialize the touched array, if necessary.
        if (ARCH_UNLIKELY(_numTouched.load(std::memory_order_acquire) == 0)) {
            _GrowTouched();
        }
 
        const VdfIndex outputIdx = VdfOutput::GetIndexFromId(outputId);
        const uint32_t idx = outputIdx / _TouchedWordBits;
        const uint64_t bit =
            UINT64_C(1) << (outputIdx & (_TouchedWordBits - 1));
        if ((_touched[idx].load(std::memory_order_relaxed) & bit) == 0) {
            _touched[idx].fetch_or(bit);
        }
    }
 
    bool Untouch(const VdfId outputId) {
        const VdfIndex outputIdx = VdfOutput::GetIndexFromId(outputId);
        const uint32_t idx = outputIdx / _TouchedWordBits;
        const uint64_t bit =
            UINT64_C(1) << (outputIdx & (_TouchedWordBits - 1));
        if (idx < _numTouched.load(std::memory_order_acquire) &&
            (_touched[idx].load(std::memory_order_relaxed) & bit) != 0) {
            return (_touched[idx].fetch_and(~bit) & bit) != 0;
        }
        return false;
    }
 
    size_t GetNumData() const {
        return _outputData.size();
    }
 
    VDF_API
    void Reset(const DataHandle handle, const VdfId outputId) const;
 
    VDF_API
    void Clear();
 
private:
 
    // An 8-bit field containing indices into the executor buffer data array.
    //
    class _BufferIndices
    {
    public:
 
        // Default constructor.
        //
        _BufferIndices() {
            Reset();
        }
 
        // Reset the buffer indices.
        //
        void Reset() {
            _indices = 
                (0 << _privateOffset) |
                (1 << _scratchOffset) |
                (2 << _publicOffset);
        }
 
        // Returns the private index.
        //
        const uint8_t GetPrivateIndex() const {
            return (_indices.load(std::memory_order_acquire) & _privateMask) >>
                _privateOffset;
        }
 
        // Returns the scratch index.
        //
        const uint8_t GetScratchIndex() const {
            return (_indices.load(std::memory_order_acquire) & _scratchMask) >>
                _scratchOffset;
        }
 
        // Returns the public index.
        //
        const uint8_t GetPublicIndex() const {
            return (_indices.load(std::memory_order_acquire) & _publicMask) >>
                _publicOffset;
        }
 
        // Swaps the private buffer with the public buffer index.
        //
        void PublishPrivateIndex() {
            uint8_t indices = _indices.load(std::memory_order_relaxed);
            uint8_t newIndices =
                // Scratch index stays.
                (indices & _scratchMask) |
                // Public index replaces the private index.
                ((indices >> _publicOffset) & _privateMask) |
                // Private index replaces the public index.
                ((indices & _privateMask) << _publicOffset);
            _indices.compare_exchange_strong(indices, newIndices);
        }
 
        // Swaps the scratch buffer with the public buffer index.
        //
        void PublishScratchIndex() {
            uint8_t indices = _indices.load(std::memory_order_relaxed);
            uint8_t newIndices =
                // Private index stays.
                (indices & _privateMask) |
                // Public index replaces the scratch index.
                ((indices >> _scratchOffset) & _scratchMask) |
                // Scratch index replaces the public index.
                ((indices & _scratchMask) << _scratchOffset);
            _indices.compare_exchange_strong(indices, newIndices);
        }
 
    private:
 
        // The bit offsets into the indices bitset.
        //
        constexpr static uint8_t _privateOffset = 0;
        constexpr static uint8_t _scratchOffset = 2;
        constexpr static uint8_t _publicOffset  = 4;
 
        // The bitmasks for each entry in the indices bitset. Two bits per
        // buffer index.
        //
        constexpr static uint8_t _privateMask   = 0x03;
        constexpr static uint8_t _scratchMask   = 0x0C;
        constexpr static uint8_t _publicMask    = 0x30;
 
        // The buffer indices packed into a bitset.
        //
        std::atomic<uint8_t> _indices;
    };
 
    // The generic output data stored for each entry in this container. Note,
    // the memory underpinning this structure must be zero-initialized. This
    // is required for the "constructed flags" to work properly.
    //
    struct _OutputData
    {
        // Noncopyable.
        //
        _OutputData(const _OutputData &) = delete;
        _OutputData &operator=(const _OutputData &) = delete;
 
        // Default constructor.
        //
        explicit _OutputData(const VdfId oid) {
            Reset(oid);
        }
 
        // Reset the output data to its defaul constructed state.
        //
        void Reset(const VdfId oid) {
            // Set the default buffer indices.
            bufferIndices.Reset();
 
            // Reset the invalidation timestamp.
            invalidationTimestamp =
                VdfExecutorInvalidationData::InitialInvalidationTimestamp;
 
            // Enforce release semantics on the outputId to synchronize the
            // non-atomic and dependent atomic write above.
            outputId.store(oid, std::memory_order_release);
        }
 
        // The output id.
        //
        std::atomic<VdfId> outputId;
 
        // The buffer indices.
        //
        _BufferIndices bufferIndices;
 
        // A zero-initialized checksum to synchronize the corresponding data
        // construction. The checksum will be incremented after construction of
        // each piece of data. Once the checksum has reached a specific value,
        // all data is guaranteed to be constructed.
        //
        // The correctness of the checksum implementation requires that this
        // member be initialized to zero prior to becoming visible to other
        // threads.  Due to the design of concurrent_vector, zero_allocator is
        // used to ensure that the storage is zeroed but the lifetime of the
        // _OutputData object is not guaranteed to begin before other threads
        // are able to observe its entry in the vector.  Aside from this
        // undefined behavior, there is another issue.  As of C++20,
        // std::atomic's default constructor does value-initialization.  This
        // zeroes the checksum again *after* other threads may have already
        // incremented its value.  Placing the checksum in a union and
        // omitting any explicit initialization in the _OutputData constructor
        // dodges this specific re-zeroing problem.  However, the workaround
        // cannot solve the fundamental issue of object lifetime outlined
        // above.
        //
        union {
            std::atomic<uint8_t> constructionChecksum;
        };
 
        // The invalidation timestamp. We store this information in the
        // generic data vector in order to make it available during evaluation
        // (mung buffer locking).
        //
        VdfInvalidationTimestamp invalidationTimestamp;
 
    };
 
    // A simple container that contains output values that had their ownership
    // transferred into this data manager. The operations provided on this
    // class are thread-safe.
    //
    class _Inbox
    {
    public:
 
        // Constructor.
        //
        _Inbox() : _buffer() {}
 
        // Destructor.
        //
        ~_Inbox() {
            VdfExecutorBufferData *buffer =
                _buffer.load(std::memory_order_acquire);
            if (buffer) {
                delete buffer;
            }
        }
 
        // Takes ownership of the \p value, and returns \c true if ownership
        // has successfully been transferred. Note that this instance will
        // assume responsibility of lifetime management over \p value if this
        // operation succeeds. Otherwise, the call site will maintain this
        // responsibility.
        //
        VDF_API
        bool Take(VdfVector *value, const VdfMask &mask);
 
        // Gets the current value.
        //
        VdfExecutorBufferData *Get() const {
            return _buffer.load(std::memory_order_acquire);
        }
 
        // Clears out the inbox.
        //
        void Reset() {
            if (_buffer.load(std::memory_order_relaxed)) {
                delete _buffer.exchange(nullptr);
            }
        }
 
    private:
 
        // The buffer. May be nullptr, if this inbox is empty.
        // 
        std::atomic<VdfExecutorBufferData *> _buffer;
 
    };
 
    // Type of each segment in the locations array. An array of atomic integers.
    using _LocationsSegment = std::atomic<int> *;
 
    // Expected to be 16 bytes in size.
    static_assert(sizeof(_OutputData) == 16, 
        "_OutputData expected to be 16 bytes in size.");
 
    // Resize the touched bitset to accommodate touching any output currently
    // in the network. It is safe to call this concurrently.
    //
    VDF_API
    void _GrowTouched();
 
    // Atomically create a new locations segment and return its pointer.
    //
    VDF_API
    _LocationsSegment _CreateSegment(size_t segmentIndex) const;
 
    // Create a new location index and its corresponding data for the output
    // with the given output id.
    //
    VDF_API
    int _CreateLocation(
        const VdfId outputId,
        int currentLocation,
        std::atomic<int> *newLocation) const;
 
    // Pushes a new data entry into the internal vectors for the output with
    // the given outputId.
    // 
    VDF_API
    int _CreateData(const VdfId outputId) const;
 
    // Reset the data stored at the location.
    //
    VDF_API
    void _ResetLocation(
        const VdfId outputId,
        int currentLocation,
        std::atomic<int> *newLocation) const;
 
    // Wait for the location entry to become available.
    // 
    VDF_API
    int _WaitForLocation(
        int currentLocation,
        std::atomic<int> *newLocation) const;
 
    // The number of output buffers (public, private, scratch).
    //
    constexpr static size_t _NumBuffers = 3;
 
    // The initial number of entries reserved in the data vectors.
    //
    constexpr static size_t _InitialDataNum = 1024;
 
    // Sentinels for invalid (not yet created) and pending (currently being
    // created) location indices.
    //
    constexpr static int _LocationInvalid = -1;
    constexpr static int _LocationPending = -2;
 
    // The number of bits in a word of the touched array.
    //
    constexpr static uint32_t _TouchedWordBits = sizeof(uint64_t) * CHAR_BIT;
 
    // The size of a segment in the segmented locations array. Must be a power
    // of two.
    //
    constexpr static size_t _SegmentSize = 4096;
 
    // The locations array, mapping from output index to output data,
    // evaluation data and invalidation data index. The array is segmented, and
    // segments will be lazily allocated.
    //
    size_t _numSegments;
    std::unique_ptr<std::atomic<_LocationsSegment>[]> _locations;
 
    // The touched bitset. This data needs to be stored outside of _OutputData
    // to allow us to set bits concurrently from speculation node evaluation.
    //
    std::atomic<size_t> _numTouched;
    std::unique_ptr<std::atomic<uint64_t>[]> _touched;
    tbb::spin_mutex _touchedMutex;
 
    // The output data.
    //
    using _OutputDataVector =
        tbb::concurrent_vector<_OutputData, WorkZeroAllocator<_OutputData>>;
    mutable _OutputDataVector _outputData;
 
    // The arrays of buffer data corresponding with the output data.
    //
    using _BufferDataVector = tbb::concurrent_vector<VdfExecutorBufferData>;
    mutable _BufferDataVector _bufferData[_NumBuffers];
 
    // The array of inboxes corresponding with the output data.
    //
    using _InboxVector = tbb::concurrent_vector<_Inbox>;
    mutable _InboxVector _inboxes;
 
    // The invalidation specific data corresponding with the output data.
    //
    using _InvalidationDataVector =
        tbb::concurrent_vector<VdfExecutorInvalidationData>;
    mutable _InvalidationDataVector _invalidationData;
 
};
 
 
Vdf_ParallelExecutorDataVector::DataHandle
Vdf_ParallelExecutorDataVector::GetOrCreateDataHandle(
    const VdfId outputId) const
{
    // Get the output index.
    const VdfIndex outputIndex = VdfOutput::GetIndexFromId(outputId);
 
    // Compute the index to the segment.
    const size_t segmentIndex = outputIndex / _SegmentSize;
    TF_DEV_AXIOM(segmentIndex < _numSegments);
 
    // Retrieve the segment.
    _LocationsSegment segment =
        _locations[segmentIndex].load(std::memory_order_acquire);
 
    // Allocate the segment, if required.
    if (!segment) {
        segment = _CreateSegment(segmentIndex);
    }
 
    // Using the output index, look up the location in the data vector.
    const size_t segmentOffset = outputIndex & (_SegmentSize - 1);
    std::atomic<int> *location = &segment[segmentOffset];
    const int currentLocation = location->load(std::memory_order_acquire);
 
    // Create a new entry, if the location is still uninitialized.
    if (currentLocation < 0) {
        return _CreateLocation(outputId, currentLocation, location);
    }
 
    // Make sure the output id matches at the location, and reset the entry
    // if there is a mismatch.
    else if (outputId != _outputData[currentLocation].outputId.load(
        std::memory_order_acquire)) {
        _ResetLocation(outputId, currentLocation, location);
    }
 
    // Return the existing location.
    return currentLocation;
}
 
Vdf_ParallelExecutorDataVector::DataHandle
Vdf_ParallelExecutorDataVector::GetDataHandle(const VdfId outputId) const
{
    // Get the output index.
    const VdfIndex outputIndex = VdfOutput::GetIndexFromId(outputId);
 
    // Compute the index to the segment.
    const size_t segmentIndex = outputIndex / _SegmentSize;
 
    // If the index is out of bounds, we can bail out right away.
    if (segmentIndex >= _numSegments) {
        return InvalidHandle;
    }
 
    // Retrieve the segment.
    _LocationsSegment segment =
        _locations[segmentIndex].load(std::memory_order_acquire);
 
    // If the segment is not allocated, we can bail out right away.
    if (!segment) {
        return InvalidHandle;
    }
 
    // If the location points to a valid entry in the data vector, and the
    // data at that index matches the output id, we can return the data.
    // Otherwise, the location may either be garbage, or the output version
    // may have changed.
    const size_t segmentOffset = outputIndex & (_SegmentSize - 1);
    std::atomic<int> * const location = &segment[segmentOffset];
    int currentLocation = location->load(std::memory_order_acquire);
 
    // Because of the ABA problems that can occur when we call
    // _ResetLocation, it's possible to observe _LocationPending
    // as the location for this output. In this case we have to wait
    if (ARCH_UNLIKELY(currentLocation == _LocationPending)) {
        currentLocation = _WaitForLocation(currentLocation, location);
    }
 
    return
        currentLocation >= 0 &&
        _outputData[currentLocation].outputId.load(
            std::memory_order_acquire) == outputId
            ? currentLocation
            : InvalidHandle;
}
 
PXR_NAMESPACE_CLOSE_SCOPE
 
#endif