schedule.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_SCHEDULE_H
#define PXR_EXEC_VDF_SCHEDULE_H


#include "pxr/pxr.h"

#include "pxr/exec/vdf/api.h"
#include "pxr/exec/vdf/countingIterator.h"
#include "pxr/exec/vdf/node.h"
#include "pxr/exec/vdf/request.h"
#include "pxr/exec/vdf/scheduleNode.h"
#include "pxr/exec/vdf/scheduleTasks.h"
#include "pxr/exec/vdf/types.h"

#include "pxr/base/tf/bits.h"

#include <vector>

PXR_NAMESPACE_OPEN_SCOPE

class VdfConnection;
class VdfNetwork;


class VdfSchedule
{
public:
    template <typename Iterator>
    class IteratorRange {
    public:
        template <typename IteratorConvertible>
        IteratorRange(IteratorConvertible begin, IteratorConvertible end) :
            _begin(begin), _end(end) {}

        Iterator begin() const { return _begin; }
        Iterator end() const { return _end; }
        bool empty() const { return _begin == _end; }
        size_t size() const { return std::distance(_begin, _end); }

    private:
        Iterator _begin;
        Iterator _end;
    };

    VdfSchedule(const VdfSchedule &) = delete;
    VdfSchedule &operator=(const VdfSchedule &) = delete;

    using ScheduleNodeVector = std::vector<VdfScheduleNode>;

    using TaskIdRange = IteratorRange<Vdf_CountingIterator<VdfScheduleTaskId>>;

    using InputDependencyRange =
        IteratorRange<std::vector<VdfScheduleInputDependency>::const_iterator>;

    using InputsRange =
        IteratorRange<std::vector<VdfScheduleInput>::const_iterator>;

    class OutputId {
    public:
        bool IsValid() const {
            return _scheduleNodeIndex >= 0 && _secondaryIndex >= 0;
        }

        OutputId &operator++() {
            _secondaryIndex++;
            return *this;
        }

        bool operator==(const OutputId &rhs) const {
            return _scheduleNodeIndex == rhs._scheduleNodeIndex &&
                   _secondaryIndex == rhs._secondaryIndex;
        }

        bool operator!=(const OutputId &rhs) const {
            return !(*this == rhs);
        }

    private:
        // Construct an OutputId with a specific schedule node index
        // and secondary index.
        //
        // The scheduleNodeIndex is expected to be one of the possible values
        // stored in VdfSchedule::_nodesToIndexMap (i.e. [0, _nodes.size()-1])
        // or a negative value to indicate an invalid id.
        //
        // And invalid id signifies that an output is not scheduled.
        //
        // The secondaryIndex is an index into the associated
        // VdfScheduleNode's VdfScheduleOutputs vector, which stores data
        // about the scheduled node explicitly.
        //
        OutputId(int scheduleNodeIndex, int secondaryIndex) :
            _scheduleNodeIndex(scheduleNodeIndex),
            _secondaryIndex(secondaryIndex)
        {}

        // Only VdfSchedule is allowed to construct instances of VdfOuputId.
        friend class VdfSchedule;

        // Data members
        int _scheduleNodeIndex;
        int _secondaryIndex;
    };

    VDF_API
    VdfSchedule();

    VDF_API
    ~VdfSchedule();

    VDF_API
    void Clear();

    bool IsValid() const {
        return _isValid;
    }

    const VdfNetwork *GetNetwork() const { return _network; }


    VDF_API
    bool IsScheduled(const VdfNode &node) const;

    VDF_API
    OutputId GetOutputId(const VdfOutput &output) const;

    VDF_API
    OutputId GetOrCreateOutputId(const VdfOutput &output);

    VDF_API
    void AddInput(const VdfConnection &connection, const VdfMask &mask);

    VDF_API
    void DeduplicateInputs();

    VDF_API
    const VdfNode *GetNode(const OutputId &outputId) const;

    VDF_API
    OutputId GetOutputIdsBegin(const VdfNode &node) const;

    VDF_API
    OutputId GetOutputIdsEnd(const VdfNode &node) const;

    VDF_API
    InputsRange GetInputs(const VdfNode &node) const;

    VDF_API
    bool IsAffective(const OutputId &outputId) const;



    VDF_API
    const VdfOutput *GetOutput(const OutputId &outputId) const;

    VDF_API
    const VdfOutput *GetOutputToClear(const VdfNode &node) const;

    VDF_API
    const VdfMask &GetRequestMask(const OutputId &outputId) const;

    const VdfMask &GetRequestMask(
        const VdfScheduleTaskIndex invocationIndex) const {
        TF_DEV_AXIOM(!VdfScheduleTaskIsInvalid(invocationIndex));
        return _nodeInvocations[invocationIndex].requestMask;
    }

    VDF_API
    void GetRequestAndAffectsMask(
        const OutputId &outputId,
        const VdfMask **requestMask,
        const VdfMask **affectsMask) const;

    void GetRequestAndAffectsMask(
        const VdfScheduleTaskIndex invocationIndex,
        const VdfMask **requestMask,
        const VdfMask **affectsMask) const {
        TF_DEV_AXIOM(!VdfScheduleTaskIsInvalid(invocationIndex));
        *requestMask = &_nodeInvocations[invocationIndex].requestMask;
        *affectsMask = &_nodeInvocations[invocationIndex].affectsMask;
    }

    VDF_API
    const VdfMask &GetAffectsMask(const OutputId &outputId) const;

    VDF_API
    const VdfMask &GetKeepMask(const OutputId &outputId) const;

    const VdfMask &GetKeepMask(
        const VdfScheduleTaskIndex invocationIndex) const {
        TF_DEV_AXIOM(!VdfScheduleTaskIsInvalid(invocationIndex));
        return _nodeInvocations[invocationIndex].keepMask;
    }

    VDF_API
    const VdfOutput *GetPassToOutput(const OutputId &outputId) const;

    VDF_API
    const VdfOutput *GetFromBufferOutput(const OutputId &outputId) const;

    bool HasSMBL() const {
        return _hasSMBL;
    }


    VDF_API
    void ForEachScheduledOutput(
        const VdfNode &node,
        const VdfScheduledOutputCallback &callback) const;

    size_t GetNumUniqueInputDependencies() const {
        return _numUniqueInputDeps;
    }

    size_t GetNumComputeTasks() const {
        return _computeTasks.size();
    }

    size_t GetNumInputsTasks() const {
        return _inputsTasks.size();
    }

    size_t GetNumPrepTasks() const {
        return _numPrepTasks;
    }

    size_t GetNumKeepTasks() const {
        return _numKeepTasks;
    }

    const TaskIdRange GetComputeTaskIds(const VdfNode &node) const {
        int scheduleNodeIndex = _GetScheduleNodeIndex(node);
        TF_DEV_AXIOM(scheduleNodeIndex >= 0);
        return TaskIdRange(
            _nodesToComputeTasks[scheduleNodeIndex].taskId,
            _nodesToComputeTasks[scheduleNodeIndex].taskId +
            _nodesToComputeTasks[scheduleNodeIndex].taskNum);
    }

    TaskIdRange GetComputeTaskIds(
        const VdfScheduleInputDependency &input) const {
        return TaskIdRange(
            input.computeOrKeepTaskId,
            input.computeOrKeepTaskId + input.computeTaskNum);
    }

    const VdfScheduleTaskIndex GetKeepTaskIndex(const VdfNode &node) const {
        int scheduleNodeIndex = _GetScheduleNodeIndex(node);
        return scheduleNodeIndex >= 0
            ? _nodesToKeepTasks[scheduleNodeIndex]
            : VdfScheduleTaskInvalid;
    }

    const VdfScheduleComputeTask &GetComputeTask(
        const VdfScheduleTaskIndex index) const {
        TF_DEV_AXIOM(index < _computeTasks.size());
        return _computeTasks[index];
    }

    const VdfScheduleInputsTask &GetInputsTask(
        const VdfScheduleTaskIndex index) const {
        TF_DEV_AXIOM(index < _inputsTasks.size());
        return _inputsTasks[index];
    }

    InputDependencyRange GetPrereqInputDependencies(
        const VdfScheduleInputsTask &task) const {
        std::vector<VdfScheduleInputDependency>::const_iterator begin =
            _inputDeps.begin() + task.inputDepIndex;
        return InputDependencyRange(begin, begin + task.prereqsNum);
    }

    InputDependencyRange GetOptionalInputDependencies(
        const VdfScheduleInputsTask &task) const {
        std::vector<VdfScheduleInputDependency>::const_iterator begin =
            _inputDeps.begin() + task.inputDepIndex + task.prereqsNum;
        return InputDependencyRange(begin, begin + task.optionalsNum);
    }

    InputDependencyRange GetRequiredInputDependencies(
        const VdfScheduleComputeTask &task) const {
        std::vector<VdfScheduleInputDependency>::const_iterator begin =
            _inputDeps.begin() + task.requiredsIndex;
        return InputDependencyRange(begin, begin + task.requiredsNum);
    }

    VDF_API
    VdfScheduleInputDependencyUniqueIndex GetUniqueIndex(
        const OutputId outputId) const;


    bool IsSmallSchedule() const { return _isSmallSchedule; }

    VDF_API
    void SetRequest(const VdfRequest &request);

    const VdfRequest &GetRequest() const { return _request; }

    ScheduleNodeVector &GetScheduleNodeVector() { 
        return _nodes;
    }

    const ScheduleNodeVector &GetScheduleNodeVector() const { 
        return _nodes;
    }

    int GetScheduleNodeIndex(const OutputId &outputId) const {
        return outputId._scheduleNodeIndex;
    }

    const TfBits &GetScheduledNodeBits() const { return _scheduledNodes; }

    VDF_API
    void SetRequestMask(const OutputId &outputId, const VdfMask &mask);

    VDF_API
    void SetAffectsMask(const OutputId &outputId, const VdfMask &mask);

    VDF_API
    void SetKeepMask(const OutputId &outputId, const VdfMask &mask);    

    VDF_API
    void SetPassToOutput(const OutputId &outputId, const VdfOutput *output);

    VDF_API
    void SetFromBufferOutput(const OutputId &outputId, const VdfOutput *output);

    VDF_API
    void SetOutputToClear(const VdfNode &node, const VdfOutput *outputToClear);

    VDF_API
    void InitializeFromNetwork(const VdfNetwork &network);

    void SetHasSMBL(bool enable) {
        _hasSMBL = enable;
    }


private:

    // The VdfScheduler and its derived classes are the only objects allowed
    // to set a scheduler as valid.
    friend class VdfScheduler;

    // The VdfScheduler calls this method to make sure that this schedule
    // is marked as valid and registered with a particular network.
    //
    void _SetIsValidForNetwork(const VdfNetwork *network);

    // Returns the index into _nodes that corresponds to the given VdfNode.
    // If the node is not scheduled and thus has no corresponding _nodes entry,
    // this method returns a value less than 0.
    VDF_API
    int _GetScheduleNodeIndex(const VdfNode &node) const;

    // Ensures that \p node is in the schedule and returns its scheduleNode
    // index.
    //
    int _EnsureNodeInSchedule(const VdfNode &node);

    // Data Members

    // The total list of nodes that we have to execute.  This is where the
    // schedule nodes are owned.
    ScheduleNodeVector _nodes;

    // The request for this schedule
    VdfRequest _request;

    // This is a vector that maps VdfNodes to VdfScheduleNode index in _nodes.
    std::vector<int> _nodesToIndexMap;

    // The network that we are registered with.  All of our scheduled nodes 
    // belong to this network.
    const VdfNetwork *_network;

    // Bits are set for each schedule node's index.
    TfBits _scheduledNodes;

    // Flag as to whether or not the schedule is valid.
    bool _isValid;

    // A flag that determines whether this schedule's query methods will
    // use the small schedule optimization, which is to assume there is no
    // _nodesToIndexMap and instead find schedule nodes by searching the
    // _nodes array directly.
    bool _isSmallSchedule;

    // This flag indicates whether this schedule participates in sparse mung
    // buffer locking.
    bool _hasSMBL;

    // The number of unique input dependencies created for this schedule. Each
    // unique input dependencies refers to the same output and mask combination.
    size_t _numUniqueInputDeps;

    // The scheduled tasks for parallel evaluation.
    Vdf_DefaultInitVector<VdfScheduleComputeTask> _computeTasks;
    Vdf_DefaultInitVector<VdfScheduleInputsTask> _inputsTasks;
    size_t _numKeepTasks;
    size_t _numPrepTasks;

    // Scheduled node invocations for nodes with multiple invocations.
    Vdf_DefaultInitVector<VdfScheduleNodeInvocation> _nodeInvocations;

    // The array of input dependencies used to orchestrate task synchronization.
    std::vector<VdfScheduleInputDependency> _inputDeps;

    // Arrays that map from the scheduled node index to the scheduled tasks
    // corresponding to that node.
    std::vector<VdfScheduleNodeTasks> _nodesToComputeTasks;
    std::vector<VdfScheduleTaskIndex> _nodesToKeepTasks;
};


// Example usage:
// void MyFunction(const VdfSchedule &schedule, const VdfNode &node) {
//     VDF_FOR_EACH_SCHEDULED_OUTPUT_ID(outputId, schedule, node) {
//        DoThingsWithARequestMask(schedule->GetRequestMask(outputId));
//     }
// }
//
#define VDF_FOR_EACH_SCHEDULED_OUTPUT_ID(OUTPUT_ID_NAME,VDF_SCHEDULE,VDF_NODE)  \
    for (VdfSchedule::OutputId __endId =                                        \
            (VDF_SCHEDULE).GetOutputIdsEnd(VDF_NODE),                           \
         OUTPUT_ID_NAME = (VDF_SCHEDULE).GetOutputIdsBegin(VDF_NODE) ;          \
         OUTPUT_ID_NAME != __endId; ++OUTPUT_ID_NAME)



PXR_NAMESPACE_CLOSE_SCOPE

#endif