loops.h

Go to the documentation of this file.

//
// Copyright 2016 Pixar
//
// Licensed under the terms set forth in the LICENSE.txt file available at
// https://openusd.org/license.
//
#ifndef PXR_BASE_WORK_LOOPS_H
#define PXR_BASE_WORK_LOOPS_H
 
#include "pxr/pxr.h"
#include "pxr/base/work/api.h"
#include "pxr/base/work/dispatcher.h"
#include "pxr/base/work/impl.h"
#include "pxr/base/work/threadLimits.h"
 
#include "pxr/base/tf/diagnosticTrap.h"
#include "pxr/base/tf/diagnosticTransport.h"
#include "pxr/base/tf/mallocTag.h"
 
#include <algorithm>
 
PXR_NAMESPACE_OPEN_SCOPE
 
using Work_DiagnosticTransports = tbb::concurrent_vector<TfDiagnosticTransport>;
 
template <class Fn>
class Work_LoopsTaskWrapper
{
public:
    Work_LoopsTaskWrapper(
        Fn &&callback,
        Work_DiagnosticTransports *diagnostics)
    : _callback(callback)
    , _diagnostics(diagnostics) {}
 
    template <typename ... Args>
    void operator()(Args&&... args) const {
        TfDiagnosticTrap trap;
        _callback(std::forward<Args>(args)...);
        if (!trap.IsClean()) {
            *_diagnostics->grow_by(1) = trap.Transport();
        }
    }
 
private:
    Fn & _callback;
    Work_DiagnosticTransports *_diagnostics;
};
 
template <class Fn>
class Work_MallocTagsLoopsTaskWrapper
{
public:
    Work_MallocTagsLoopsTaskWrapper(
        Fn &&callback,
        Work_DiagnosticTransports *diagnostics)
    : _callback(callback)
    , _diagnostics(diagnostics)
    , _mallocTagStack(TfMallocTag::GetCurrentStackState()) {}
 
    template <typename ... Args>
    void operator()(Args&&... args) const {
        TfDiagnosticTrap trap;
        TfMallocTag::StackOverride ovr(_mallocTagStack);
        _callback(std::forward<Args>(args)...);
        if (!trap.IsClean()) {
            *_diagnostics->grow_by(1) = trap.Transport();
        }
    }
 
private:
    Fn & _callback;
    Work_DiagnosticTransports *_diagnostics;
    TfMallocTag::StackState _mallocTagStack;
};
 
template <class Fn>
class Work_LoopsForEachTaskWrapper
{
public:
    Work_LoopsForEachTaskWrapper(
        Fn &&callback,
        Work_DiagnosticTransports *diagnostics)
    : _callback(callback)
    , _diagnostics(diagnostics) {}
 
    template <typename Arg>
    void operator()(Arg &&arg) const {
        TfDiagnosticTrap trap;
        _callback(std::forward<Arg>(arg));
        if (!trap.IsClean()) {
            *_diagnostics->grow_by(1) = trap.Transport();
        }
    }
 
private:
    Fn & _callback;
    Work_DiagnosticTransports *_diagnostics;
};
 
template <class Fn>
class Work_MallocTagsLoopsForEachTaskWrapper
{
public:
    Work_MallocTagsLoopsForEachTaskWrapper(
        Fn &&callback,
        Work_DiagnosticTransports *diagnostics)
    : _callback(callback)
    , _diagnostics(diagnostics)
    , _mallocTagStack(TfMallocTag::GetCurrentStackState()) {}
 
    template <typename Arg>
    void operator()(Arg &&arg) const {
        TfDiagnosticTrap trap;
        TfMallocTag::StackOverride ovr(_mallocTagStack);
        _callback(std::forward<Arg>(arg));
        if (!trap.IsClean()) {
            *_diagnostics->grow_by(1) = trap.Transport();
        }
    }
 
private:
    Fn & _callback;
    Work_DiagnosticTransports *_diagnostics;
    TfMallocTag::StackState _mallocTagStack;
};
 
template<typename Fn>
void
WorkSerialForN(size_t n, Fn &&fn)
{
    std::forward<Fn>(fn)(0, n);
}
 
template <typename Fn>
void
WorkParallelForN(size_t n, Fn &&callback, size_t grainSize)
{
    if (n == 0)
        return;
 
    // Don't bother with parallel_for, if concurrency is limited to 1.
    if (WorkHasConcurrency()) {
        PXR_WORK_IMPL_NAMESPACE_USING_DIRECTIVE;
        Work_DiagnosticTransports diagnosticTransports;
        if (TfMallocTag::IsInitialized()) {
            Work_MallocTagsLoopsTaskWrapper<Fn>
                task(std::forward<Fn>(callback), &diagnosticTransports);
            WorkImpl_ParallelForN(n, task, grainSize);
        }
        else {
            Work_LoopsTaskWrapper<Fn>
                task(std::forward<Fn>(callback), &diagnosticTransports);
            WorkImpl_ParallelForN(n, task, grainSize);
        }
 
        for (auto &dt: diagnosticTransports) {
            dt.Post();
        }
    } else {
        // If concurrency is limited to 1, execute serially.
        WorkSerialForN(n, std::forward<Fn>(callback));
    }
}
 
template <typename Fn>
void
WorkParallelForN(size_t n, Fn &&callback)
{
    WorkParallelForN(n, std::forward<Fn>(callback), 1);
}
 
template <typename RangeType, typename Fn>
void
WorkParallelForTBBRange(const RangeType &range, Fn &&callback)
{
    // Don't bother with parallel_for, if concurrency is limited to 1.
    if (WorkHasConcurrency()) {
        PXR_WORK_IMPL_NAMESPACE_USING_DIRECTIVE;
        // Use the work backend's ParallelForTBBRange if one exists
        // otherwise use the default implementation below that builds off of the 
        // dispatcher.
#if defined WORK_IMPL_HAS_PARALLEL_FOR_TBB_RANGE
        Work_DiagnosticTransports diagnosticTransports;
        if (TfMallocTag::IsInitialized()) {
            Work_MallocTagsLoopsTaskWrapper<Fn>
                task(std::forward<Fn>(callback), &diagnosticTransports);
            WorkImpl_ParallelForTBBRange(range, task);
        }
        else {
            Work_LoopsTaskWrapper<Fn>
                task(std::forward<Fn>(callback), &diagnosticTransports);
            WorkImpl_ParallelForTBBRange(range, task);
        }
        for (auto &dt: diagnosticTransports) {
            dt.Post();
        }
#else
        // The parallel task responsible for recursively sub-dividing the range
        // and invoking the callback on the sub-ranges.
        class _RangeTask
        {
        public:
            _RangeTask(
                WorkDispatcher &dispatcher,
                RangeType &&range,
                const Fn &callback)
            : _dispatcher(dispatcher)
            , _range(std::move(range))
            , _callback(callback) {}
 
            void operator()() const {
                // Subdivide the given range until it is no longer divisible, and
                // recursively spawn _RangeTasks for the right side of the split.
                RangeType &leftRange = _range;
                while (leftRange.is_divisible()) {
                    RangeType rightRange(leftRange, tbb::split());
                    _dispatcher.Run(_RangeTask(
                        _dispatcher, std::move(rightRange), _callback));
                }
 
                // If there are any more entries remaining in the left-most side
                // of the given range, invoke the callback on the left-most range.
                if (!leftRange.empty()) {
                    std::invoke(_callback, leftRange);
                }
            }
 
        private:
            WorkDispatcher &_dispatcher;
            mutable RangeType _range;
            const Fn &_callback;
        };
 
        WorkDispatcher dispatcher;
        RangeType range = range;
        dispatcher.Run(_RangeTask(
            dispatcher, range, std::forward<Fn>(callback)));
#endif
    } else {
        // If concurrency is limited to 1, execute serially.
        std::forward<Fn>(callback)(range);
    }
}
 
template <typename InputIterator, typename Fn>
inline void
WorkParallelForEach(
    InputIterator first, InputIterator last, Fn &&fn)
{
    if (WorkHasConcurrency()) {
        PXR_WORK_IMPL_NAMESPACE_USING_DIRECTIVE;
        Work_DiagnosticTransports diagnosticTransports;
        if (TfMallocTag::IsInitialized()) {
            Work_MallocTagsLoopsForEachTaskWrapper<Fn>
                task(std::forward<Fn>(fn), &diagnosticTransports);
            WorkImpl_ParallelForEach(first, last, task);
        }
        else {
            Work_LoopsForEachTaskWrapper<Fn>
                task(std::forward<Fn>(fn), &diagnosticTransports);
            WorkImpl_ParallelForEach(first, last, task);
        }
        for (auto &dt: diagnosticTransports) {
            dt.Post();
        }
    } else {
        std::for_each(first, last, std::forward<Fn>(fn));
    }
}
 
PXR_NAMESPACE_CLOSE_SCOPE
 
#endif // PXR_BASE_WORK_LOOPS_H