work: Provide malloc-tag specific task wrappers. Evidently just

having an extra unused pointer in the task impacts some internal perf tests

(Internal change: 2394417)

Author: gitamohr

pxr/base/work/dispatcher.h

@@ -61,9 +61,17 @@ class Work_Dispatcher
 
     template <class Callable>
     inline void Run(Callable &&c) {
-        _dispatcher.Run(
-            _InvokerTask<typename std::remove_reference<Callable>::type>(
-                std::forward<Callable>(c), &_errors));
+        if (TfMallocTag::IsInitialized()) {
+            _dispatcher.Run(
+                _MallocTagsInvokerTask<
+                typename std::remove_reference<Callable>::type>(
+                    std::forward<Callable>(c), &_errors));
+        }
+        else {
+            _dispatcher.Run(
+                _InvokerTask<typename std::remove_reference<Callable>::type>(
+                    std::forward<Callable>(c), &_errors));
+        }
     }
 
     template <class Callable, class A0, class ... Args>
@@ -103,20 +111,51 @@ class Work_Dispatcher
     template <class Fn>
     struct _InvokerTask {
         explicit _InvokerTask(Fn &&fn, _ErrorTransports *err) 
+            : _fn(std::move(fn))
+            , _errors(err) {}
+
+        explicit _InvokerTask(Fn const &fn, _ErrorTransports *err) 
+            : _fn(fn)
+            , _errors(err) {}
+
+        // Ensure only moves happen, no copies.
+        _InvokerTask(_InvokerTask &&other) = default;
+        _InvokerTask(const _InvokerTask &other) = delete;
+        _InvokerTask &operator=(const _InvokerTask &other) = delete;
+
+        void operator()() const {
+            TfErrorMark m;
+            _fn();
+            if (!m.IsClean())
+                Work_Dispatcher::_TransportErrors(m, _errors);
+        }
+    private:
+        Fn _fn;
+        _ErrorTransports *_errors;
+    };
+
+    // Function invoker helper that wraps the invocation with an ErrorMark so we
+    // can transmit errors that occur back to the thread that Wait() s for tasks
+    // to complete.  This version also duplicates the caller's malloc tag stack
+    // to the callee's thread.
+    template <class Fn>
+    struct _MallocTagsInvokerTask {
+        explicit _MallocTagsInvokerTask(Fn &&fn, _ErrorTransports *err) 
             : _fn(std::move(fn))
             , _errors(err)
             , _mallocTagStack(TfMallocTag::GetCurrentStackState())
             {}
 
-        explicit _InvokerTask(Fn const &fn, _ErrorTransports *err) 
+        explicit _MallocTagsInvokerTask(Fn const &fn, _ErrorTransports *err) 
             : _fn(fn)
             , _errors(err)
             , _mallocTagStack(TfMallocTag::GetCurrentStackState()) {}
 
         // Ensure only moves happen, no copies.
-        _InvokerTask(_InvokerTask &&other) = default;
-        _InvokerTask(const _InvokerTask &other) = delete;
-        _InvokerTask &operator=(const _InvokerTask &other) = delete;
+        _MallocTagsInvokerTask(_MallocTagsInvokerTask &&other) = default;
+        _MallocTagsInvokerTask(const _MallocTagsInvokerTask &other) = delete;
+        _MallocTagsInvokerTask &
+        operator=(const _MallocTagsInvokerTask &other) = delete;
 
         void operator()() const {
             TfErrorMark m;

pxr/base/work/loops.h

@@ -31,6 +31,31 @@ class Work_LoopsTaskWrapper
         Fn &&callback,
         Work_ErrorTransports *errors)
     : _callback(callback)
+    , _errors(errors) {}
+
+    template <typename ... Args>
+    void operator()(Args&&... args) const {
+        TfErrorMark m;
+        _callback(std::forward<Args>(args)...);
+        if (!m.IsClean()) {
+            TfErrorTransport transport = m.Transport();
+            _errors->grow_by(1)->swap(transport);
+        }
+    }
+
+private:
+    Fn & _callback;
+    Work_ErrorTransports *_errors;
+};
+
+template <class Fn>
+class Work_MallocTagsLoopsTaskWrapper
+{
+public:
+    Work_MallocTagsLoopsTaskWrapper(
+        Fn &&callback,
+        Work_ErrorTransports *errors)
+    : _callback(callback)
     , _errors(errors)
     , _mallocTagStack(TfMallocTag::GetCurrentStackState()) {}
 
@@ -59,6 +84,31 @@ class Work_LoopsForEachTaskWrapper
         Fn &&callback,
         Work_ErrorTransports *errors)
     : _callback(callback)
+    , _errors(errors) {}
+
+    template <typename Arg>
+    void operator()(Arg &&arg) const {
+        TfErrorMark m;
+        _callback(std::forward<Arg>(arg));
+        if (!m.IsClean()) {
+            TfErrorTransport transport = m.Transport();
+            _errors->grow_by(1)->swap(transport);
+        }
+    }
+
+private:
+    Fn & _callback;
+    Work_ErrorTransports *_errors;
+};
+
+template <class Fn>
+class Work_MallocTagsLoopsForEachTaskWrapper
+{
+public:
+    Work_MallocTagsLoopsForEachTaskWrapper(
+        Fn &&callback,
+        Work_ErrorTransports *errors)
+    : _callback(callback)
     , _errors(errors)
     , _mallocTagStack(TfMallocTag::GetCurrentStackState()) {}
 
@@ -124,9 +174,17 @@ WorkParallelForN(size_t n, Fn &&callback, size_t grainSize)
     if (WorkHasConcurrency()) {
         PXR_WORK_IMPL_NAMESPACE_USING_DIRECTIVE;
         Work_ErrorTransports errorTransports;
-        Work_LoopsTaskWrapper<Fn>
-            task(std::forward<Fn>(callback), &errorTransports);
-        WorkImpl_ParallelForN(n, task, grainSize);
+        if (TfMallocTag::IsInitialized()) {
+            Work_MallocTagsLoopsTaskWrapper<Fn>
+                task(std::forward<Fn>(callback), &errorTransports);
+            WorkImpl_ParallelForN(n, task, grainSize);
+        }
+        else {
+            Work_LoopsTaskWrapper<Fn>
+                task(std::forward<Fn>(callback), &errorTransports);
+            WorkImpl_ParallelForN(n, task, grainSize);
+        }
+
         for (auto &et: errorTransports) {
             et.Post();
         }
@@ -179,9 +237,16 @@ WorkParallelForTBBRange(const RangeType &range, Fn &&callback)
         // dispatcher.
 #if defined WORK_IMPL_HAS_PARALLEL_FOR_TBB_RANGE
         Work_ErrorTransports errorTransports;
-        Work_LoopsTaskWrapper<Fn>
-            task(std::forward<Fn>(callback), &errorTransports);
-        WorkImpl_ParallelForTBBRange(range, task);
+        if (TfMallocTag::IsInitialized()) {
+            Work_MallocTagsLoopsTaskWrapper<Fn>
+                task(std::forward<Fn>(callback), &errorTransports);
+            WorkImpl_ParallelForTBBRange(range, task);
+        }
+        else {
+            Work_LoopsTaskWrapper<Fn>
+                task(std::forward<Fn>(callback), &errorTransports);
+            WorkImpl_ParallelForTBBRange(range, task);
+        }
         for (auto &et: errorTransports) {
             et.Post();
         }
@@ -253,9 +318,16 @@ WorkParallelForEach(
     if (WorkHasConcurrency()) {
         PXR_WORK_IMPL_NAMESPACE_USING_DIRECTIVE;
         Work_ErrorTransports errorTransports;
-        Work_LoopsForEachTaskWrapper<Fn>
-            task(std::forward<Fn>(fn), &errorTransports);
-        WorkImpl_ParallelForEach(first, last, task);
+        if (TfMallocTag::IsInitialized()) {
+            Work_MallocTagsLoopsForEachTaskWrapper<Fn>
+                task(std::forward<Fn>(fn), &errorTransports);
+            WorkImpl_ParallelForEach(first, last, task);
+        }
+        else {
+            Work_LoopsForEachTaskWrapper<Fn>
+                task(std::forward<Fn>(fn), &errorTransports);
+            WorkImpl_ParallelForEach(first, last, task);
+        }
         for (auto &et: errorTransports) {
             et.Post();
         }