convert atom to task graph (#4230)

* Intial attempt to convert the Atom/RHI/FrameScheduler to use the new TaskGraph api

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* Avoid enqueuing work on the active task thread if the submitted task
graph is waitable

When submitting a task graph, supplying a wait event implies that
dependent jobs must occur on threads that do not wait on the event (in
the absence of work stealing). This change prevents this by adding a
notion of a task thread enable/disable state, and prohibiting dependent
jobs from being enqueued on waiting threads.

Signed-off-by: Jeremy Ong <jcong@amazon.com>

* Convert RPI/Scene to use TaskGraph pass 1, Culling jobs remain on the old system

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* RemoveTask Graph changes from the FrameScheduler::ExecuteGroups, use old job system instead

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* Per review, removing commented out code

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* Cleanup debug code, & build fix

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* Add a cvar & interface to query whether to use jobs or task graph

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* Make TaskGraph assert if you try to wait inside a job

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* Fix TaskTest SpawnSubgraph to account for the new TaskGraphEvent assert on wait in a running task

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* 3 minor cleanups. 1) Events always store a ptr to their executor 2) Fix clang compile error 3) remove an early out.

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* Fix double group end that was causing assert/crash plus misc minor diff's with development

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* Fix deallocation failure on deactivation of the TaskGraphSystemComponent. Also make the system component account for multiple creation in Unit Tests.

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* Update with PR feedback
1) Rename UseTaskGraph to IsTaskGraphActive & update related code
2) prefer TaskExecutor::SetInstance
3) add comments and remove commented out code

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* Fix incorrect RTTI name for TaskGraphActiveInterface

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

* Move TaskGraphSystemComponent CRC calculation to a shared variable

Signed-off-by: rgba16f <82187279+rgba16f@users.noreply.github.com>

Co-authored-by: Jeremy Ong <jcong@amazon.com>

Code/Framework/AzCore/AzCore/AzCoreModule.cpp

@@ -22,6 +22,7 @@
 #include <AzCore/Time/TimeSystemComponent.h>
 #include <AzCore/Console/LoggerSystemComponent.h>
 #include <AzCore/EBus/EventSchedulerSystemComponent.h>
+#include <AzCore/Task/TaskGraphSystemComponent.h>
 
 namespace AZ
 {
@@ -41,6 +42,7 @@ namespace AZ
             TimeSystemComponent::CreateDescriptor(),
             LoggerSystemComponent::CreateDescriptor(),
             EventSchedulerSystemComponent::CreateDescriptor(),
+            TaskGraphSystemComponent::CreateDescriptor(),
 
 #if !defined(AZCORE_EXCLUDE_LUA)
             ScriptSystemComponent::CreateDescriptor(),
@@ -55,6 +57,7 @@ namespace AZ
             azrtti_typeid<TimeSystemComponent>(),
             azrtti_typeid<LoggerSystemComponent>(),
             azrtti_typeid<EventSchedulerSystemComponent>(),
+            azrtti_typeid<TaskGraphSystemComponent>(),
         };
     }
 }

Code/Framework/AzCore/AzCore/Task/TaskExecutor.cpp

@@ -190,11 +190,13 @@ namespace AZ
         class TaskWorker
         {
         public:
-            void Spawn(::AZ::TaskExecutor& executor, size_t id, AZStd::semaphore& initSemaphore, bool affinitize)
+            static thread_local TaskWorker* t_worker;
+
+            void Spawn(::AZ::TaskExecutor& executor, uint32_t id, AZStd::semaphore& initSemaphore, bool affinitize)
             {
                 m_executor = &executor;
 
-                AZStd::string threadName = AZStd::string::format("TaskWorker %zu", id);
+                AZStd::string threadName = AZStd::string::format("TaskWorker %u", id);
                 AZStd::thread_desc desc = {};
                 desc.m_name = threadName.c_str();
                 if (affinitize)
@@ -205,12 +207,29 @@ namespace AZ
 
                 m_thread = AZStd::thread{ [this, &initSemaphore]
                                           {
+                                              t_worker = this;
                                               initSemaphore.release();
                                               Run();
                                           },
                                           &desc };
             }
 
+            // Threads that wait on a graph to complete are disqualified from receiving tasks until the wait finishes
+            void Disable()
+            {
+                m_enabled = false;
+            }
+
+            void Enable()
+            {
+                m_enabled = true;
+            }
+
+            bool Enabled() const
+            {
+                return m_enabled;
+            }
+
             void Join()
             {
                 m_active.store(false, AZStd::memory_order_release);
@@ -222,11 +241,7 @@ namespace AZ
             {
                 m_queue.Enqueue(task);
 
-                if (!m_busy.exchange(true))
-                {
-                    // The worker was idle prior to enqueueing the task, release the semaphore
-                    m_semaphore.release();
-                }
+                m_semaphore.release();
             }
 
         private:
@@ -234,7 +249,6 @@ namespace AZ
             {
                 while (m_active)
                 {
-                    m_busy = false;
                     m_semaphore.acquire();
 
                     if (!m_active)
@@ -242,8 +256,6 @@ namespace AZ
                         return;
                     }
 
-                    m_busy = true;
-
                     Task* task = m_queue.TryDequeue();
                     while (task)
                     {
@@ -271,12 +283,15 @@ namespace AZ
 
             AZStd::thread m_thread;
             AZStd::atomic<bool> m_active;
-            AZStd::atomic<bool> m_busy;
+            AZStd::atomic<bool> m_enabled = true;
             AZStd::binary_semaphore m_semaphore;
 
             ::AZ::TaskExecutor* m_executor;
             TaskQueue m_queue;
+            friend class ::AZ::TaskExecutor;
         };
+
+        thread_local TaskWorker* TaskWorker::t_worker = nullptr;
     } // namespace Internal
 
     static EnvironmentVariable<TaskExecutor*> s_executor;
@@ -291,13 +306,16 @@ namespace AZ
         return **s_executor;
     }
 
-    // TODO: Create the default executor as part of a component (as in TaskManagerComponent)
     void TaskExecutor::SetInstance(TaskExecutor* executor)
     {
-        AZ_Assert(!s_executor, "Attempting to set the global task executor more than once");
-
-        s_executor = AZ::Environment::CreateVariable<TaskExecutor*>("GlobalTaskExecutor");
-        s_executor.Set(executor);
+        if (!executor)
+        {
+            s_executor.Reset();
+        }
+        else if (!s_executor) // ignore any calls to set after the first (this happens in unit tests that create new system entities)
+        {
+            s_executor = AZ::Environment::CreateVariable<TaskExecutor*>(s_executorName, executor);
+        }
     }
 
     TaskExecutor::TaskExecutor(uint32_t threadCount)
@@ -307,14 +325,12 @@ namespace AZ
 
         m_workers = reinterpret_cast<Internal::TaskWorker*>(azmalloc(m_threadCount * sizeof(Internal::TaskWorker)));
 
-        bool affinitize = m_threadCount == AZStd::thread::hardware_concurrency();
-
         AZStd::semaphore initSemaphore;
 
-        for (size_t i = 0; i != m_threadCount; ++i)
+        for (uint32_t i = 0; i != m_threadCount; ++i)
         {
             new (m_workers + i) Internal::TaskWorker{};
-            m_workers[i].Spawn(*this, i, initSemaphore, affinitize);
+            m_workers[i].Spawn(*this, i, initSemaphore, false);
         }
 
         for (size_t i = 0; i != m_threadCount; ++i)
@@ -334,9 +350,21 @@ namespace AZ
         azfree(m_workers);
     }
 
-    void TaskExecutor::Submit(Internal::CompiledTaskGraph& graph)
+    Internal::TaskWorker* TaskExecutor::GetTaskWorker()
+    {
+        if (Internal::TaskWorker::t_worker && Internal::TaskWorker::t_worker->m_executor == this)
+        {
+            return Internal::TaskWorker::t_worker;
+        }
+        return nullptr;
+    }
+
+    void TaskExecutor::Submit(Internal::CompiledTaskGraph& graph, TaskGraphEvent* event)
     {
         ++m_graphsRemaining;
+
+        event->m_executor = this; // Used to validate event is not waited for inside a job
+
         // Submit all tasks that have no inbound edges
         for (Internal::Task& task : graph.Tasks())
         {
@@ -352,11 +380,24 @@ namespace AZ
         // TODO: Something more sophisticated is likely needed here.
         // First, we are completely ignoring affinity.
         // Second, some heuristics on core availability will help distribute work more effectively
-        m_workers[++m_lastSubmission % m_threadCount].Enqueue(&task);
+        uint32_t nextWorker = ++m_lastSubmission % m_threadCount;
+        while (!m_workers[nextWorker].Enabled())
+        {
+            // Graphs that are waiting for the completion of a task graph cannot enqueue tasks onto
+            // the thread issuing the wait.
+            nextWorker = ++m_lastSubmission % m_threadCount;
+        }
+
+        m_workers[nextWorker].Enqueue(&task);
     }
 
     void TaskExecutor::ReleaseGraph()
     {
         --m_graphsRemaining;
     }
+
+    void TaskExecutor::ReactivateTaskWorker()
+    {
+        GetTaskWorker()->Enable();
+    }
 } // namespace AZ

Code/Framework/AzCore/AzCore/Task/TaskExecutor.h

@@ -72,14 +72,19 @@ namespace AZ
         explicit TaskExecutor(uint32_t threadCount = 0);
         ~TaskExecutor();
 
-        void Submit(Internal::CompiledTaskGraph& graph);
+        // Submit a task graph for execution. Waitable task graphs cannot enqueue work on the task thread
+        // that is currently active
+        void Submit(Internal::CompiledTaskGraph& graph, TaskGraphEvent* event);
 
         void Submit(Internal::Task& task);
 
     private:
         friend class Internal::TaskWorker;
+        friend class TaskGraphEvent;
 
+        Internal::TaskWorker* GetTaskWorker();
         void ReleaseGraph();
+        void ReactivateTaskWorker();
 
         Internal::TaskWorker* m_workers;
         uint32_t m_threadCount = 0;

Code/Framework/AzCore/AzCore/Task/TaskGraph.cpp

@@ -14,6 +14,12 @@ namespace AZ
 {
     using Internal::CompiledTaskGraph;
 
+    void TaskGraphEvent::Wait()
+    {
+        AZ_Assert(m_executor->GetTaskWorker() == nullptr, "Waiting in a task is unsupported");
+        m_semaphore.acquire();
+    }
+
     void TaskToken::PrecedesInternal(TaskToken& comesAfter)
     {
         AZ_Assert(!m_parent.m_submitted, "Cannot mutate a TaskGraph that was previously submitted.");
@@ -71,7 +77,7 @@ namespace AZ
             m_compiledTaskGraph->m_tasks[i].Init();
         }
 
-        executor.Submit(*m_compiledTaskGraph);
+        executor.Submit(*m_compiledTaskGraph, waitEvent);
 
         if (m_retained)
         {

Code/Framework/AzCore/AzCore/Task/TaskGraph.h

@@ -22,10 +22,19 @@ namespace AZ
     namespace Internal
     {
         class CompiledTaskGraph;
+        class TaskWorker;
     }
     class TaskExecutor;
     class TaskGraph;
 
+    class TaskGraphActiveInterface
+    {
+    public:
+        AZ_RTTI(TaskGraphActiveInterface, "{08118074-B139-4EF9-B8FD-29F1D6DC9233}");
+
+        virtual bool IsTaskGraphActive() const = 0;
+    };
+
     // A TaskToken is returned each time a Task is added to the TaskGraph. TaskTokens are used to
     // express dependencies between tasks within the graph, and have no purpose after the graph
     // is submitted (simply let them go out of scope)
@@ -70,9 +79,12 @@ namespace AZ
     private:
         friend class ::AZ::Internal::CompiledTaskGraph;
         friend class TaskGraph;
+        friend class TaskExecutor;
+
         void Signal();
 
         AZStd::binary_semaphore m_semaphore;
+        TaskExecutor* m_executor = nullptr;
     };
 
     // The TaskGraph encapsulates a set of tasks and their interdependencies. After adding
@@ -89,6 +101,9 @@ namespace AZ
         // Reset the state of the task graph to begin recording tasks and edges again
         // NOTE: Graph must be in a "settled" state (cannot be in-flight)
         void Reset();
+        
+        // Returns false if 1 or more tasks have been added to the graph
+        bool IsEmpty();
 
         // Add a task to the graph, retrieiving a token that can be used to express dependencies
         // between tasks. The first argument specifies the TaskKind, used for tracking the task.

Code/Framework/AzCore/AzCore/Task/TaskGraph.inl

@@ -33,11 +33,6 @@ namespace AZ
         return m_semaphore.try_acquire_for(AZStd::chrono::milliseconds{ 0 });
     }
 
-    inline void TaskGraphEvent::Wait()
-    {
-        m_semaphore.acquire();
-    }
-
     inline void TaskGraphEvent::Signal()
     {
         m_semaphore.release();
@@ -59,6 +54,11 @@ namespace AZ
         return { AddTask(descriptor, AZStd::forward<Lambdas>(lambdas))... };
     }
 
+    inline bool TaskGraph::IsEmpty()
+    {
+        return m_tasks.empty();
+    }
+
     inline void TaskGraph::Detach()
     {
         m_retained = false;

Code/Framework/AzCore/AzCore/Task/TaskGraphSystemComponent.cpp

@@ -0,0 +1,88 @@
+/*
+ * Copyright (c) Contributors to the Open 3D Engine Project.
+ * For complete copyright and license terms please see the LICENSE at the root of this distribution.
+ *
+ * SPDX-License-Identifier: Apache-2.0 OR MIT
+ *
+ */
+
+#include <AzCore/Console/IConsole.h>
+#include <AzCore/Interface/Interface.h>
+#include <AzCore/Task/TaskGraphSystemComponent.h>
+#include <AzCore/Serialization/SerializeContext.h>
+#include <AzCore/Serialization/EditContext.h>
+
+// Create a cvar as a central location for experimentation with switching from the Job system to TaskGraph system.
+AZ_CVAR(bool, cl_activateTaskGraph, false, nullptr, AZ::ConsoleFunctorFlags::Null, "Flag clients of TaskGraph to switch between jobs/taskgraph (Note does not disable task graph system)");
+static constexpr uint32_t TaskExecutorServiceCrc = AZ_CRC_CE("TaskExecutorService");
+
+namespace AZ
+{
+    void TaskGraphSystemComponent::Activate()
+    {
+        AZ_Assert(m_taskExecutor == nullptr, "Error multiple activation of the TaskGraphSystemComponent");
+
+        if (Interface<TaskGraphActiveInterface>::Get() == nullptr)
+        {
+            Interface<TaskGraphActiveInterface>::Register(this);
+            m_taskExecutor = aznew TaskExecutor();
+            TaskExecutor::SetInstance(m_taskExecutor);
+        }
+    }
+
+    void TaskGraphSystemComponent::Deactivate()
+    {
+        if (&TaskExecutor::Instance() == m_taskExecutor) // check that our instance is the global instance (not always true in unit tests)
+        {
+            m_taskExecutor->SetInstance(nullptr);
+        }
+        if (m_taskExecutor)
+        {
+            azdestroy(m_taskExecutor);
+            m_taskExecutor = nullptr;
+        }
+        if (Interface<TaskGraphActiveInterface>::Get() == this)
+        {
+            Interface<TaskGraphActiveInterface>::Unregister(this);
+        }
+    }
+
+    void TaskGraphSystemComponent::GetProvidedServices(ComponentDescriptor::DependencyArrayType& provided)
+    {
+        provided.push_back(TaskExecutorServiceCrc);
+    }
+
+    void TaskGraphSystemComponent::GetIncompatibleServices(ComponentDescriptor::DependencyArrayType& incompatible)
+    {
+        incompatible.push_back(TaskExecutorServiceCrc);
+    }
+
+    void TaskGraphSystemComponent::GetDependentServices([[maybe_unused]] ComponentDescriptor::DependencyArrayType& dependent)
+    {
+    }
+
+    void TaskGraphSystemComponent::Reflect(ReflectContext* context)
+    {
+        if (SerializeContext* serializeContext = azrtti_cast<SerializeContext*>(context))
+        {
+            serializeContext->Class<TaskGraphSystemComponent, AZ::Component>()
+                ->Version(1)
+                ;
+
+            if (AZ::EditContext* ec = serializeContext->GetEditContext())
+            {
+                ec->Class<TaskGraphSystemComponent>
+                    ("TaskGraph", "System component to create the default executor")
+                    ->ClassElement(AZ::Edit::ClassElements::EditorData, "")
+                        ->Attribute(AZ::Edit::Attributes::Category, "Engine")
+                        ->Attribute(AZ::Edit::Attributes::AppearsInAddComponentMenu, AZ_CRC_CE("System"))
+                    ;
+            }
+        }
+    }
+
+    bool TaskGraphSystemComponent::IsTaskGraphActive() const
+    {
+        return cl_activateTaskGraph;
+    }
+} // namespace AZ

Code/Framework/AzCore/AzCore/Task/TaskGraphSystemComponent.h

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) Contributors to the Open 3D Engine Project.
+ * For complete copyright and license terms please see the LICENSE at the root of this distribution.
+ *
+ * SPDX-License-Identifier: Apache-2.0 OR MIT
+ *
+ */
+#pragma once
+
+#include <AzCore/Component/Component.h>
+#include <AzCore/Math/Crc.h>
+#include <AzCore/Task/TaskExecutor.h>
+#include <AzCore/Task/TaskGraph.h>
+
+namespace AZ
+{
+    class TaskGraphSystemComponent
+        : public Component
+        , public TaskGraphActiveInterface
+    {
+    public:
+        AZ_COMPONENT(AZ::TaskGraphSystemComponent, "{5D56B829-1FEB-43D5-A0BD-E33C0497EFE2}")
+
+        TaskGraphSystemComponent() = default;
+
+        // Implement TaskGraphActiveInterface
+        bool IsTaskGraphActive() const override;
+
+    private:
+        //////////////////////////////////////////////////////////////////////////
+        // Component base
+        void Activate() override;
+        void Deactivate() override;
+        //////////////////////////////////////////////////////////////////////////
+
+        /// \ref ComponentDescriptor::GetProvidedServices
+        static void GetProvidedServices(ComponentDescriptor::DependencyArrayType& provided);
+        /// \ref ComponentDescriptor::GetIncompatibleServices
+        static void GetIncompatibleServices(ComponentDescriptor::DependencyArrayType& incompatible);
+        /// \ref ComponentDescriptor::GetDependentServices
+        static void GetDependentServices(ComponentDescriptor::DependencyArrayType& dependent);
+        /// \red ComponentDescriptor::Reflect
+        static void Reflect(ReflectContext* reflection);
+
+        AZ::TaskExecutor*   m_taskExecutor = nullptr;
+    };
+}

Code/Framework/AzCore/AzCore/azcore_files.cmake

@@ -633,6 +633,8 @@ set(FILES
     Task/TaskGraph.cpp
     Task/TaskGraph.h
     Task/TaskGraph.inl
+    Task/TaskGraphSystemComponent.h
+    Task/TaskGraphSystemComponent.cpp
     Threading/ThreadSafeDeque.h
     Threading/ThreadSafeDeque.inl
     Threading/ThreadSafeObject.h

Code/Framework/AzCore/Tests/TaskTests.cpp

@@ -34,7 +34,7 @@ namespace UnitTest
             AZ::AllocatorInstance<AZ::PoolAllocator>::Create();
             AZ::AllocatorInstance<AZ::ThreadPoolAllocator>::Create();
 
-            m_executor = aznew TaskExecutor(4);
+            m_executor = aznew TaskExecutor();
         }
 
         void TearDown() override
@@ -236,6 +236,82 @@ namespace UnitTest
         EXPECT_EQ(x, 1);
     }
 
+    TEST_F(TaskGraphTestFixture, SingleTask)
+    {
+        AZStd::atomic_int32_t x = 0;
+
+        TaskGraph graph;
+        graph.AddTask(
+            defaultTD,
+            [&x]
+            {
+                x = 1;
+            });
+
+        TaskGraphEvent ev;
+        graph.SubmitOnExecutor(*m_executor, &ev);
+        ev.Wait();
+
+        EXPECT_EQ(1, x);
+    }
+
+
+    TEST_F(TaskGraphTestFixture, SingleTaskChain)
+    {
+        AZStd::atomic_int32_t x = 0;
+
+        TaskGraph graph;
+        auto a = graph.AddTask(
+            defaultTD,
+            [&x]
+            {
+                x += 1;
+            });
+        auto b = graph.AddTask(
+            defaultTD,
+            [&x]
+            {
+                x += 1;
+            });
+        b.Precedes(a);
+
+        TaskGraphEvent ev;
+        graph.SubmitOnExecutor(*m_executor, &ev);
+        ev.Wait();
+
+        EXPECT_EQ(2, x);
+    }
+
+    TEST_F(TaskGraphTestFixture, MultipleIndependentTaskChains)
+    {
+        AZStd::atomic_int32_t x = 0;
+        constexpr int numChains = 5;
+
+        TaskGraph graph;
+        for( int i = 0; i < numChains; ++i)
+        {
+            auto a = graph.AddTask(
+                defaultTD,
+                [&x]
+                {
+                    x += 1;
+                });
+            auto b = graph.AddTask(
+                defaultTD,
+                [&x]
+                {
+                    x += 1;
+                });
+            b.Precedes(a);
+        }
+
+        TaskGraphEvent ev;
+        graph.SubmitOnExecutor(*m_executor, &ev);
+        ev.Wait();
+
+        EXPECT_EQ(2*numChains, x);
+    }
+
     TEST_F(TaskGraphTestFixture, VariadicInterface)
     {
         int x = 0;
@@ -388,6 +464,7 @@ namespace UnitTest
         EXPECT_EQ(3, x);
     }
 
+    // Waiting inside a task is disallowed , test that it fails correctly
     TEST_F(TaskGraphTestFixture, SpawnSubgraph)
     {
         AZStd::atomic<int> x = 0;
@@ -434,7 +511,10 @@ namespace UnitTest
                 f.Precedes(g);
                 TaskGraphEvent ev;
                 subgraph.SubmitOnExecutor(*m_executor, &ev);
+                // TaskGraphEvent::Wait asserts if called on a worker thread, suppress & validate assert
+                AZ_TEST_START_TRACE_SUPPRESSION;
                 ev.Wait();
+                AZ_TEST_STOP_TRACE_SUPPRESSION(1);
             });
         auto d = graph.AddTask(
             defaultTD,
@@ -464,8 +544,6 @@ namespace UnitTest
         TaskGraphEvent ev;
         graph.SubmitOnExecutor(*m_executor, &ev);
         ev.Wait();
-
-        EXPECT_EQ(3 | 0b100000, x);
     }
 
     TEST_F(TaskGraphTestFixture, RetainedGraph)

Code/Framework/AzFramework/AzFramework/Application/Application.cpp

@@ -28,6 +28,7 @@
 #include <AzCore/NativeUI/NativeUISystemComponent.h>
 #include <AzCore/Module/ModuleManagerBus.h>
 #include <AzCore/Interface/Interface.h>
+#include <AzCore/Task/TaskGraphSystemComponent.h>
 
 #include <AzFramework/Asset/SimpleAsset.h>
 #include <AzFramework/Asset/AssetBundleManifest.h>
@@ -295,6 +296,7 @@ namespace AzFramework
             azrtti_typeid<AZ::ScriptSystemComponent>(),
             azrtti_typeid<AZ::JobManagerComponent>(),
             azrtti_typeid<AZ::SliceSystemComponent>(),
+            azrtti_typeid<AZ::TaskGraphSystemComponent>(),
 
             azrtti_typeid<AzFramework::AssetCatalogComponent>(),
             azrtti_typeid<AzFramework::CustomAssetTypeComponent>(),