LWP100-1624 : non-preemptible scheduler locking

lib/include/rocRoller/CodeGen/Instruction.hpp

@@ -92,8 +92,9 @@ namespace rocRoller
         static Instruction Wait(WaitCount const& wait);
         static Instruction Wait(WaitCount&& wait);
 
-        static Instruction Lock(Scheduling::Dependency const& dependency, std::string comment);
-        static Instruction Unlock(std::string comment);
+        static Instruction Lock(Scheduling::Dependency dependency, std::string comment = "");
+        static Instruction Unlock(Scheduling::Dependency dependency, std::string comment = "");
+        static Instruction Unlock(std::string comment = "");
 
         /**
          * This instruction will be considered to write `reg` even though it
@@ -167,8 +168,8 @@ namespace rocRoller
         void        toStream(std::ostream&, LogLevel level) const;
         std::string toString(LogLevel level) const;
 
-        constexpr int          getLockValue() const;
-        Scheduling::Dependency getDependency() const;
+        constexpr Scheduling::LockOperation getLockValue() const;
+        constexpr Scheduling::Dependency    getDependency() const;
 
         std::string const&                                 getOpCode() const;
         std::array<std::string, Instruction::MaxModifiers> getModifiers() const;
@@ -178,8 +179,9 @@ namespace rocRoller
             return m_nopCount;
         }
 
-        Instruction lock(Scheduling::Dependency const& dependency, std::string comment);
-        Instruction unlock(std::string comment);
+        Instruction& lock(Scheduling::Dependency dependency, std::string comment = "");
+        Instruction& unlock(Scheduling::Dependency dependency, std::string comment = "");
+        Instruction& unlock(std::string comment = "");
 
         void                    addControlOp(int id);
         std::vector<int> const& controlOps() const;
@@ -273,7 +275,8 @@ namespace rocRoller
 
         std::string m_opcode;
 
-        Scheduling::Dependency m_dependency = Scheduling::Dependency::None;
+        Scheduling::LockOperation m_lockOp     = Scheduling::LockOperation::None;
+        Scheduling::Dependency    m_dependency = Scheduling::Dependency::None;
 
         std::array<Register::ValuePtr, MaxDstRegisters> m_inoutDsts;
         std::array<Register::ValuePtr, MaxDstRegisters> m_dst;

lib/include/rocRoller/CodeGen/Instruction_impl.hpp

@@ -249,24 +249,22 @@ namespace rocRoller
         return rv;
     }
 
-    inline Instruction Instruction::Lock(Scheduling::Dependency const& dependency,
-                                         std::string                   comment = "")
+    inline Instruction Instruction::Lock(Scheduling::Dependency dependency, std::string comment)
     {
-        AssertFatal(dependency != Scheduling::Dependency::Unlock
-                        && dependency != Scheduling::Dependency::Count,
-                    "Can not create lock instruction with Unlock or Count dependency");
-
         Instruction rv;
-        rv.m_dependency = dependency;
-        rv.addComment(comment);
+        rv.lock(dependency, std::move(comment));
         return rv;
     }
 
-    inline Instruction Instruction::Unlock(std::string comment = "")
+    inline Instruction Instruction::Unlock(std::string comment)
+    {
+        return Unlock(Scheduling::Dependency::None, std::move(comment));
+    }
+
+    inline Instruction Instruction::Unlock(Scheduling::Dependency dependency, std::string comment)
     {
         Instruction rv;
-        rv.m_dependency = Scheduling::Dependency::Unlock;
-        rv.addComment(comment);
+        rv.unlock(dependency, std::move(comment));
         return rv;
     }
 
@@ -355,7 +353,8 @@ namespace rocRoller
             && m_label.empty()
             && m_waitCount == WaitCount()
             && m_opcode.empty()
-            && m_dependency == Scheduling::Dependency::None;
+            && m_dependency == Scheduling::Dependency::None
+        && m_lockOp == Scheduling::LockOperation::None;
         // clang-format on
     }
 
@@ -608,42 +607,43 @@ namespace rocRoller
         throw std::runtime_error("Too many allocations!");
     }
 
-    inline Instruction Instruction::lock(Scheduling::Dependency const& dependency,
-                                         std::string                   comment = "")
+    inline Instruction& Instruction::lock(Scheduling::Dependency dependency, std::string comment)
     {
-        AssertFatal(dependency != Scheduling::Dependency::Unlock
+        AssertFatal(m_lockOp == Scheduling::LockOperation::None,
+                    "An instruction can only lock or unlock once.");
+
+        AssertFatal(dependency != Scheduling::Dependency::None
                         && dependency != Scheduling::Dependency::Count,
                     "Can not create lock instruction with Unlock or Count dependency");
 
+        m_lockOp     = Scheduling::LockOperation::Lock;
         m_dependency = dependency;
-        addComment(comment);
+        addComment(std::move(comment));
         return *this;
     }
 
-    inline Instruction Instruction::unlock(std::string comment = "")
+    inline Instruction& Instruction::unlock(std::string comment)
     {
-        m_dependency = Scheduling::Dependency::Unlock;
-        addComment(comment);
+        return unlock(Scheduling::Dependency::None, std::move(comment));
+    }
+
+    inline Instruction& Instruction::unlock(Scheduling::Dependency dependency, std::string comment)
+    {
+        AssertFatal(m_lockOp == Scheduling::LockOperation::None,
+                    "An instruction can only lock or unlock once.");
+
+        m_lockOp     = Scheduling::LockOperation::Unlock;
+        m_dependency = dependency;
+        addComment(std::move(comment));
         return *this;
     }
 
-    inline constexpr int Instruction::getLockValue() const
+    inline constexpr Scheduling::LockOperation Instruction::getLockValue() const
     {
-        if(m_dependency == Scheduling::Dependency::Unlock)
-        {
-            return -1;
-        }
-        else if(m_dependency == Scheduling::Dependency::None)
-        {
-            return 0;
-        }
-        else
-        {
-            return 1;
-        }
+        return m_lockOp;
     }
 
-    inline Scheduling::Dependency Instruction::getDependency() const
+    inline constexpr Scheduling::Dependency Instruction::getDependency() const
     {
         return m_dependency;
     }

lib/include/rocRoller/Scheduling/Costs/Cost.hpp

@@ -73,6 +73,8 @@ namespace rocRoller
              */
             Result operator()(std::vector<Generator<Instruction>::iterator>&) const;
 
+            float operator()(Instruction const& inst) const;
+
             /**
              * @brief Gets the cost of one generator for the given iteration
              *

lib/include/rocRoller/Scheduling/Scheduler.hpp

@@ -26,6 +26,7 @@
 
 #pragma once
 
+#include <stack>
 #include <string>
 #include <vector>
 
@@ -40,31 +41,81 @@ namespace rocRoller
 {
     namespace Scheduling
     {
+        constexpr bool isNonPreemptibleDependency(Dependency dep);
+
+        /**
+         * Locking Rules
+         *
+         * A scheduler has a number of streams which each will yield a sequence of instructions.
+         * The job of the scheduler is to pick (i.e. schedule) the instruction from the beginning of one of the streams, and then repeat until there are no more streams with any instructions left.
+         *
+         * - If a scheduler schedules an exclusive lock, it must continue to
+         * select instructions from that same stream until that lock has been
+         * unlocked.
+         *   - That stream might include further lock/unlock instructions which
+         * must occur in a last-in, first-out order, those should be treated
+         * as a stack to track when the original lock has been unlocked.
+         *
+         * - If a stream yields any kind of lock, it cannot yield a lower-ranked
+         *   lock until it releases the higher-ranked lock.
+         * - If a scheduler schedules a non-exclusive lock, it cannot schedule
+         *   the same kind of lock from any other stream until that lock is
+         *   released.
+         *      Example:
+         *          1. Stream 0 locks M0.
+         *          2. Stream 1 locks VCC.
+         *          3. Stream 0 tries to lock VCC.  It must wait until Stream 1 unlocks VCC.
+         *          4. Stream 1 unlocks VCC.
+         *          5. Stream 0 locks VCC.
+         *          6. Stream 1 locks SCC.  Pull from Stream 1 until SCC is unlocked.
+         *          7. Stream 0 locks SCC.  Pull from Stream 0 until SCC is unlocked.
+         *          8. Stream 0 unlocks VCC.
+         *          9. Stream 0 unlocks M0.
+         *
+         * - If a scheduler schedules a non-exclusive lock, it cannot schedule a
+         *   lower-ranked exclusive lock from any stream until that lock is
+         *   released.
+         *      Examples:
+         *          - If stream 0 locks M0, and then we see stream 3 try to lock
+         *            Branch, we can't pull from stream 3 until stream 0
+         *            releases M0.
+         *          - If stream 2 locks VCC, stream 1 can lock SCC.  We will
+         *            then have to pull from stream 1 until SCC is released.
+         *
+         */
         class LockState
         {
         public:
             explicit LockState(ContextPtr ctx);
             LockState(ContextPtr ctx, Dependency dependency);
 
-            void add(Instruction const& instr);
-            bool isLocked() const;
-            void isValid(bool locked = false) const;
+            void add(Instruction const& instr, int streamId);
+            bool isNonPreemptibleStream(int streamId) const;
+            bool isSchedulable(Instruction const& instr, int streamId) const;
+            bool isLocked(Dependency dependency, int streamId) const;
 
             /**
              * @brief Extra checks to verify lock state integrity.
              *
              * Note: disabled in Release mode.
              *
              * @param instr The instruction to verify
+         * @param streamId The instruction's stream ID
              */
-            void lockCheck(Instruction const& instr);
+            void lockCheck(Instruction const& instr, int streamId) const;
 
-            Dependency getDependency() const;
-            int        getLockDepth() const;
+            Dependency getTopDependency(int streamId) const;
+            int        getLockDepth(int streamId) const;
 
         private:
-            int                               m_lockdepth;
-            Dependency                        m_dependency;
+            void lock(Dependency dep, int streamId);
+            void unlock(Dependency dep, int streamId);
+
+            std::map<int, std::stack<Dependency>> m_stack;
+            std::map<Dependency, int>             m_stream;
+            std::unordered_multiset<Dependency>   m_locks;
+            std::optional<int>                    m_nonPreemptibleStream;
+
             std::weak_ptr<rocRoller::Context> m_ctx;
         };
 
@@ -118,7 +169,8 @@ namespace rocRoller
              * - At least one instruction
              * - If that first instruction locks the stream, yields until the stream is unlocked.
              */
-            Generator<Instruction> yieldFromStream(Generator<Instruction>::iterator& iter);
+            Generator<Instruction> yieldFromStream(Generator<Instruction>::iterator& iter,
+                                                   int                               streamId);
 
             /**
              * @brief Handles new nodes being added to the instruction streams being scheduled.
@@ -132,8 +184,9 @@ namespace rocRoller
                                std::vector<Generator<Instruction>::iterator>& iterators);
         };
 
-        std::ostream& operator<<(std::ostream&, SchedulerProcedure const&);
-        std::ostream& operator<<(std::ostream&, Dependency const&);
+        std::ostream& operator<<(std::ostream&, SchedulerProcedure proc);
+        std::ostream& operator<<(std::ostream&, Dependency dep);
+        std::ostream& operator<<(std::ostream& stream, LockOperation lockOp);
     }
 }
 

lib/include/rocRoller/Scheduling/Scheduler_fwd.hpp

@@ -41,21 +41,30 @@ namespace rocRoller
             Count
         };
 
-        enum class Dependency
+        enum class Dependency : int
+        {
+            None = 0, //< Temporary. Should only be used for unlocking.
+            Branch, //< Exclusive: Loops and ConditionalOp
+            M0, //< Non-exclusive: The M0 special-purpose register
+            VCC, //< Non-exclusive: The VCC special-purpose register
+            SCC, //< Exclusive: The SCC special-purpose register, which is
+            //  implicitly written by many instructions.
+            Count
+        };
+
+        enum class LockOperation : int
         {
             None = 0,
-            SCC,
-            VCC,
-            Branch,
+            Lock,
             Unlock,
-            M0,
             Count
         };
 
         class Scheduler;
         class LockState;
 
-        std::string toString(SchedulerProcedure const&);
-        std::string toString(Dependency const&);
+        std::string toString(SchedulerProcedure);
+        std::string toString(Dependency);
+        std::string toString(LockOperation);
     }
 }

lib/source/CodeGen/LowerFromKernelGraph.cpp

@@ -173,33 +173,30 @@ namespace rocRoller
 
                     // Generate code for all the nodes we found.
 
+                    // vector of instruction streams
                     std::vector<Generator<Instruction>> generators;
                     for(auto tag : nodes)
                     {
                         auto op = m_graph->control.getNode(tag);
                         generators.push_back(call(tag, op, coords));
                     }
 
-                    if(generators.size() == 1)
-                    {
-                        co_yield std::move(generators[0]);
-                    }
-                    else
                     {
                         co_yield Instruction::Comment(
                             concatenate("BEGIN Scheduler for operations ", nodes));
                         auto proc = Settings::getInstance()->get(Settings::Scheduler);
                         auto cost = Settings::getInstance()->get(Settings::SchedulerCost);
                         auto scheduler
                             = Component::GetNew<Scheduling::Scheduler>(proc, cost, m_context);
+                        auto generator = (*scheduler)(generators);
 
-                        if(!scheduler->supportsAddingStreams())
+                        if(generators.size() == 1 || !scheduler->supportsAddingStreams())
                         {
-                            co_yield (*scheduler)(generators);
+                            for(auto gen : generator)
+                                co_yield gen;
                         }
                         else
                         {
-                            auto generator         = (*scheduler)(generators);
                             auto numCompletedNodes = m_completedControlNodes.size();
 
                             for(auto iter = generator.begin(); iter != generator.end(); ++iter)
@@ -1241,7 +1238,14 @@ namespace rocRoller
 
             auto visitor = CodeGeneratorVisitor(graphPtr, kernel);
 
-            co_yield visitor.generate();
+            for(auto what : visitor.generate())
+            {
+                //if(!(what.getOpCode().empty()))
+                //    std::cout << what.getOpCode() << std::endl;
+                if(what.getLockValue() != Scheduling::LockOperation::None)
+                    std::cout << what.getLockValue() << " " << what.getDependency() << std::endl;
+                co_yield what;
+            }
         }
     }
 }

lib/source/Costs/Cost.cpp

@@ -90,11 +90,16 @@ namespace rocRoller
             return retval;
         }
 
-        float Cost::operator()(Generator<Instruction>::iterator& iter) const
+        float Cost::operator()(Instruction const& inst) const
         {
-            auto const& inst   = *iter;
-            auto        status = m_ctx.lock()->peek(inst);
+            auto status = m_ctx.lock()->peek(inst);
             return cost(inst, status);
         }
+
+        float Cost::operator()(Generator<Instruction>::iterator& iter) const
+        {
+            auto const& inst = *iter;
+            return (*this)(inst);
+        }
     }
 }