// Copyright Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

#include <rocRoller/CodeGen/WaitCount.hpp>

#include "rocRoller/GPUArchitecture/GPUArchitectureLibrary.hpp"
#include "rocRoller/GPUArchitecture/GPUArchitectureTarget.hpp"
#include <rocRoller/CodeGen/Instruction.hpp>
#include <rocRoller/Utilities/Settings.hpp>
#include <rocRoller/Utilities/String.hpp>

namespace rocRoller
{
    WaitCount::WaitCount(GPUArchitecture const& arch, std::string const& message)
    {
        m_isSplitCounter = arch.HasCapability(GPUCapability::HasSplitWaitCounters);
        m_hasVSCnt       = arch.HasCapability(GPUCapability::SeparateVscnt);
        m_hasEXPCnt      = arch.HasCapability(GPUCapability::HasExpcnt);
        m_hasTensorCnt   = arch.HasCapability(GPUCapability::HasTensorcnt);
        if(message.length() > 0)
        {
            m_comments = {message};
        }
    }

    WaitCount::WaitCount(GPUArchitecture const& arch,
                         int                    loadcnt,
                         int                    storecnt,
                         int                    vscnt,
                         int                    dscnt,
                         int                    kmcnt,
                         int                    expcnt,
                         int                    tensorcnt)
        : m_loadcnt(loadcnt)
        , m_storecnt(storecnt)
        , m_vscnt(vscnt)
        , m_dscnt(dscnt)
        , m_kmcnt(kmcnt)
        , m_expcnt(expcnt)
        , m_tensorcnt(tensorcnt)
        , m_isSplitCounter(arch.HasCapability(GPUCapability::HasSplitWaitCounters))
        , m_hasVSCnt(arch.HasCapability(GPUCapability::SeparateVscnt))
        , m_hasEXPCnt(arch.HasCapability(GPUCapability::HasExpcnt))
        , m_hasTensorCnt(arch.HasCapability(GPUCapability::HasTensorcnt))
    {
    }

    WaitCount::WaitCount(GPUArchitecture const& arch, GPUWaitQueue queue, int count)
        : m_loadcnt(-1)
        , m_storecnt(-1)
        , m_vscnt(-1)
        , m_dscnt(-1)
        , m_kmcnt(-1)
        , m_expcnt(-1)
        , m_tensorcnt(-1)
        , m_isSplitCounter(arch.HasCapability(GPUCapability::HasSplitWaitCounters))
        , m_hasVSCnt(arch.HasCapability(GPUCapability::SeparateVscnt))
        , m_hasEXPCnt(arch.HasCapability(GPUCapability::HasExpcnt))
        , m_hasTensorCnt(arch.HasCapability(GPUCapability::HasTensorcnt))
    {
        switch(queue)
        {
        case GPUWaitQueue::LoadQueue:
            m_loadcnt = count;
            break;
        case GPUWaitQueue::StoreQueue:
            m_storecnt = count;
            break;
        case GPUWaitQueue::DSQueue:
            m_dscnt = count;
            break;
        case GPUWaitQueue::KMQueue:
            m_kmcnt = count;
            break;
        case GPUWaitQueue::EXPQueue:
            m_expcnt = count;
            break;
        case GPUWaitQueue::VSQueue:
            m_vscnt = count;
            break;
        case GPUWaitQueue::TensorQueue:
            m_tensorcnt = count;
            break;
        case GPUWaitQueue::None:
        case GPUWaitQueue::Count:
            Throw<FatalError>("Invalid GPUWaitQueue!");
            break;
        }
    }

    WaitCount::WaitCount(GPUArchitecture const&       arch,
                         EnumBitset<GPUWaitQueueType> queuesToSync,
                         std::string const&           message)
        : m_queuesToSync(queuesToSync)
        , m_comments({message})
    {
    }

    WaitCount WaitCount::LoadCnt(GPUArchitecture const& arch, int value, std::string const& message)
    {
        WaitCount rv{arch, message};
        rv.m_loadcnt = value;

        return rv;
    }

    WaitCount
        WaitCount::StoreCnt(GPUArchitecture const& arch, int value, std::string const& message)
    {
        WaitCount rv{arch, message};
        rv.m_storecnt = value;

        return rv;
    }

    WaitCount WaitCount::VSCnt(GPUArchitecture const& arch, int value, std::string const& message)
    {
        WaitCount rv{arch, message};
        rv.m_vscnt = value;

        return rv;
    }

    WaitCount WaitCount::DSCnt(GPUArchitecture const& arch, int value, std::string const& message)
    {
        WaitCount rv{arch, message};
        rv.m_dscnt = value;

        return rv;
    }

    WaitCount WaitCount::KMCnt(GPUArchitecture const& arch, int value, std::string const& message)
    {
        WaitCount rv{arch, message};
        rv.m_kmcnt = value;

        return rv;
    }

    WaitCount WaitCount::EXPCnt(GPUArchitecture const& arch, int value, std::string const& message)
    {
        WaitCount rv{arch, message};
        rv.m_expcnt = value;

        return rv;
    }

    WaitCount
        WaitCount::TensorCnt(GPUArchitecture const& arch, int value, std::string const& message)
    {
        WaitCount rv{arch, message};
        rv.m_tensorcnt = value;

        return rv;
    }

    WaitCount WaitCount::Zero(GPUArchitecture const& architecture, std::string const& message)
    {
        WaitCount rv{architecture, message};

        rv.m_loadcnt  = 0;
        rv.m_storecnt = 0;
        if(architecture.HasCapability(GPUCapability::SeparateVscnt))
        {
            rv.m_vscnt = 0;
        }
        rv.m_dscnt = 0;
        rv.m_kmcnt = 0;
        if(architecture.HasCapability(GPUCapability::HasExpcnt))
        {
            rv.m_expcnt = 0;
        }
        if(architecture.HasCapability(GPUCapability::HasTensorcnt))
        {
            rv.m_tensorcnt = 0;
        }

        return rv;
    }

    WaitCount WaitCount::Max(GPUArchitecture const& architecture, std::string const& message)
    {
        WaitCount rv{architecture, message};

        if(architecture.HasCapability(GPUCapability::MaxVmcnt))
        {
            rv.m_loadcnt  = architecture.GetCapability(GPUCapability::MaxVmcnt);
            rv.m_storecnt = architecture.GetCapability(GPUCapability::MaxVmcnt);
        }

        if(architecture.HasCapability(GPUCapability::MaxLgkmcnt))
        {
            rv.m_dscnt = architecture.GetCapability(GPUCapability::MaxLgkmcnt);
            rv.m_kmcnt = architecture.GetCapability(GPUCapability::MaxLgkmcnt);
        }

        if(architecture.HasCapability(GPUCapability::HasExpcnt)
           && architecture.HasCapability(GPUCapability::MaxExpcnt))
        {
            rv.m_expcnt = architecture.GetCapability(GPUCapability::MaxExpcnt);
        }

        if(architecture.HasCapability(GPUCapability::HasTensorcnt)
           && architecture.HasCapability(GPUCapability::MaxTensorcnt))
        {
            rv.m_tensorcnt = architecture.GetCapability(GPUCapability::MaxTensorcnt);
        }

        return rv;
    }

    WaitCount WaitCount::SyncQueue(GPUArchitecture const& arch,
                                   GPUWaitQueueType       queue,
                                   std::string const&     message)
    {
        return SyncQueues(arch, EnumBitset<GPUWaitQueueType>{queue}, message);
    }

    WaitCount WaitCount::SyncQueues(GPUArchitecture const&       arch,
                                    EnumBitset<GPUWaitQueueType> queues,
                                    std::string const&           message)
    {
        return WaitCount(arch, queues, message);
    }

    EnumBitset<GPUWaitQueueType> const& WaitCount::queuesToSync() const
    {
        return m_queuesToSync;
    }

    int WaitCount::CombineValues(int lhs, int rhs)
    {
        if(lhs < 0)
        {
            return rhs;
        }

        if(rhs < 0)
        {
            return lhs;
        }

        return std::min(lhs, rhs);
    }

    WaitCount& WaitCount::combine(WaitCount const& other)
    {
        m_loadcnt   = CombineValues(m_loadcnt, other.m_loadcnt);
        m_storecnt  = CombineValues(m_storecnt, other.m_storecnt);
        m_vscnt     = CombineValues(m_vscnt, other.m_vscnt);
        m_dscnt     = CombineValues(m_dscnt, other.m_dscnt);
        m_kmcnt     = CombineValues(m_kmcnt, other.m_kmcnt);
        m_expcnt    = CombineValues(m_expcnt, other.m_expcnt);
        m_tensorcnt = CombineValues(m_tensorcnt, other.m_tensorcnt);

        m_isSplitCounter = other.m_isSplitCounter;
        m_hasVSCnt       = other.m_hasVSCnt;
        m_hasEXPCnt      = other.m_hasEXPCnt;
        m_hasTensorCnt   = other.m_hasTensorCnt;

        m_comments.insert(m_comments.end(), other.m_comments.begin(), other.m_comments.end());

        m_queuesToSync |= other.m_queuesToSync;

        return *this;
    }

    int WaitCount::loadcnt() const
    {
        return m_loadcnt;
    }
    int WaitCount::storecnt() const
    {
        return m_storecnt;
    }
    int WaitCount::vmcnt() const
    {
        return WaitCount::CombineValues(m_loadcnt, m_storecnt);
    }
    int WaitCount::vscnt() const
    {
        return m_vscnt;
    }
    int WaitCount::dscnt() const
    {
        return m_dscnt;
    }
    int WaitCount::kmcnt() const
    {
        return m_kmcnt;
    }
    int WaitCount::expcnt() const
    {
        return m_expcnt;
    }
    int WaitCount::tensorcnt() const
    {
        return m_tensorcnt;
    }

    int WaitCount::getCount(GPUWaitQueue queue) const
    {
        switch(queue)
        {
        case GPUWaitQueue::LoadQueue:
            return m_loadcnt;
        case GPUWaitQueue::StoreQueue:
            return m_storecnt;
        case GPUWaitQueue::DSQueue:
            return m_dscnt;
        case GPUWaitQueue::KMQueue:
            return m_kmcnt;
        case GPUWaitQueue::EXPQueue:
            return m_expcnt;
        case GPUWaitQueue::VSQueue:
            return m_vscnt;
        case GPUWaitQueue::TensorQueue:
            return m_tensorcnt;
        default:
            return -1;
        }
    }

    void WaitCount::setLoadcnt(int value)
    {
        m_loadcnt = value;
    }
    void WaitCount::setStorecnt(int value)
    {
        m_storecnt = value;
    }
    void WaitCount::setVscnt(int value)
    {
        m_vscnt = value;
    }
    void WaitCount::setDScnt(int value)
    {
        m_dscnt = value;
    }
    void WaitCount::setKMcnt(int value)
    {
        m_kmcnt = value;
    }
    void WaitCount::setExpcnt(int value)
    {
        m_expcnt = value;
    }
    void WaitCount::setTensorcnt(int value)
    {
        m_tensorcnt = value;
    }

    WaitCount& WaitCount::combineLoadcnt(int value)
    {
        m_loadcnt = CombineValues(m_loadcnt, value);
        return *this;
    }

    WaitCount& WaitCount::combineStorecnt(int value)
    {
        m_storecnt = CombineValues(m_storecnt, value);
        return *this;
    }

    WaitCount& WaitCount::combineVscnt(int value)
    {
        m_vscnt = CombineValues(m_vscnt, value);
        return *this;
    }

    WaitCount& WaitCount::combineDScnt(int value)
    {
        m_dscnt = CombineValues(m_dscnt, value);
        return *this;
    }

    WaitCount& WaitCount::combineKMcnt(int value)
    {
        m_kmcnt = CombineValues(m_kmcnt, value);
        return *this;
    }

    WaitCount& WaitCount::combineExpcnt(int value)
    {
        m_expcnt = CombineValues(m_expcnt, value);
        return *this;
    }

    WaitCount& WaitCount::combineTensorcnt(int value)
    {
        m_tensorcnt = CombineValues(m_tensorcnt, value);
        return *this;
    }

    std::vector<std::string> const& WaitCount::comments() const
    {
        return m_comments;
    }

    void WaitCount::addComment(std::string const& comment)
    {
        m_comments.emplace_back(comment);
    }

    void WaitCount::addComment(std::string&& comment)
    {
        m_comments.emplace_back(std::move(comment));
    }

    WaitCount WaitCount::getAsSaturatedWaitCount(GPUArchitecture const& arch) const
    {
        int loadcnt   = m_loadcnt;
        int storecnt  = m_storecnt;
        int vscnt     = m_vscnt;
        int dscnt     = m_dscnt;
        int kmcnt     = m_kmcnt;
        int expcnt    = m_expcnt;
        int tensorcnt = m_tensorcnt;

        if(arch.HasCapability(GPUCapability::MaxVmcnt))
        {
            loadcnt  = std::min(loadcnt, arch.GetCapability(GPUCapability::MaxVmcnt));
            storecnt = std::min(storecnt, arch.GetCapability(GPUCapability::MaxVmcnt));
        }

        if(arch.HasCapability(GPUCapability::MaxLgkmcnt))
        {
            kmcnt = std::min(kmcnt, arch.GetCapability(GPUCapability::MaxLgkmcnt));
            dscnt = std::min(dscnt, arch.GetCapability(GPUCapability::MaxLgkmcnt));
        }

        if(arch.HasCapability(GPUCapability::MaxExpcnt))
        {
            expcnt = std::min(expcnt, arch.GetCapability(GPUCapability::MaxExpcnt));
        }

        if(arch.HasCapability(GPUCapability::MaxTensorcnt))
        {
            tensorcnt = std::min(tensorcnt, arch.GetCapability(GPUCapability::MaxTensorcnt));
        }

        return WaitCount(arch, loadcnt, storecnt, vscnt, dscnt, kmcnt, expcnt, tensorcnt);
    }

    void WaitCount::toStream(std::ostream& os, LogLevel level) const
    {
        auto commentIter = level > LogLevel::Terse ? m_comments.begin() : m_comments.end();

        if(m_loadcnt >= 0 || m_storecnt >= 0 || m_kmcnt >= 0 || m_dscnt >= 0 || m_expcnt >= 0
           || m_tensorcnt >= 0)
        {
            if(m_isSplitCounter)
            {
                if(m_loadcnt >= 0)
                {
                    os << "s_wait_loadcnt " << m_loadcnt << std::endl;
                }

                if(m_storecnt >= 0)
                {
                    os << "s_wait_storecnt " << m_storecnt << std::endl;
                }

                if(m_kmcnt >= 0)
                {
                    os << "s_wait_kmcnt " << m_kmcnt << std::endl;
                }

                if(m_dscnt >= 0)
                {
                    os << "s_wait_dscnt " << m_dscnt << std::endl;
                }

                if(m_expcnt >= 0)
                {
                    AssertFatal(m_hasEXPCnt,
                                "EXPCnt is not a valid counter in target architecture");
                    os << "s_wait_expcnt " << m_expcnt << std::endl;
                }

                if(m_tensorcnt >= 0)
                {
                    AssertFatal(m_hasTensorCnt,
                                "TensorCnt is not a valid counter in target architecture");
                    os << "s_wait_tensorcnt " << m_tensorcnt << std::endl;
                }
            }
            else
            {
                os << "s_waitcnt";

                if(vmcnt() >= 0)
                {
                    os << " vmcnt(" << vmcnt() << ")";
                }

                if(m_kmcnt >= 0 || m_dscnt >= 0)
                {
                    auto lgkmcnt = WaitCount::CombineValues(m_kmcnt, m_dscnt);
                    os << " lgkmcnt(" << lgkmcnt << ")";
                }

                if(m_expcnt >= 0)
                {
                    AssertFatal(m_hasEXPCnt,
                                "EXPCnt is not a valid counter in target architecture");
                    os << " expcnt(" << m_expcnt << ")";
                }
            }

            if(commentIter != m_comments.end())
            {
                for(auto const& line : EscapeComment(*commentIter))
                    os << line;
                commentIter++;
            }

            os << "\n";
        }

        if(m_vscnt >= 0)
        {
            AssertFatal(m_hasVSCnt, "VSCnt is not a valid counter in target architecture");

            os << "s_waitcnt_vscnt " << m_vscnt;

            if(commentIter != m_comments.end())
            {
                for(auto const& line : EscapeComment(*commentIter))
                    os << line;
                commentIter++;
            }

            os << "\n";
        }

        if(level > LogLevel::Debug)
        {
            auto fieldComment = fmt::format(
                "m_loadcnt({}) m_storecnt({}) m_vscnt({}) m_dscnt({}) m_kmcnt({}) m_expcnt({})",
                m_loadcnt,
                m_storecnt,
                m_vscnt,
                m_dscnt,
                m_kmcnt,
                m_expcnt);

            if(m_queuesToSync.any())
            {
                fieldComment
                    += fmt::format(" m_queuesToSync({})", rocRoller::toString(m_queuesToSync));
            }

            for(auto const& line : EscapeComment(fieldComment))
                os << line;
            os << "\n";
        }

        if(commentIter != m_comments.end())
        {
            for(; commentIter != m_comments.end(); commentIter++)
            {
                for(auto const& line : EscapeComment(*commentIter))
                    os << line;
            }
            os << "\n";
        }
    }

    std::string WaitCount::toString(LogLevel level) const
    {
        std::ostringstream oss;
        toStream(oss, level);
        return oss.str();
    }

    std::ostream& operator<<(std::ostream& stream, WaitCount const& wait)
    {
        auto logLevel = Settings::getInstance()->get(Settings::LogLvl);

        wait.toStream(stream, logLevel);
        return stream;
    }
}