AuroraRuntime/Source/Threading/Primitives/SMTYield.hpp

/***
    Copyright (C) 2023 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: SMTYield.hpp
    Date: 2023-3-12
    Author: Reece
***/
#pragma once

namespace Aurora::Threading
{
    inline AuUInt32 gHasThreadLocalTimeout {};
    inline thread_local AuUInt32 tlsSpinCountLocal {};
}

#define SPIN_FOUR 1

namespace Aurora::Threading::Primitives
{
    inline AuUInt32 gSpinAdaptiveThreshold {};
    inline AuUInt32 gSpinAdaptiveCurrentCount {};

    inline AuUInt32 gUseFutexRWLock {};

    void InitAdaptiveThreshold();
    void InitAdaptiveThresholdFirstTime();

    static auline void SMPPause()
    {
    #if (defined(AURORA_ARCH_X64) || defined(AURORA_ARCH_X86))
        _mm_pause();
    #elif defined(AURORA_ARCH_ARM)
    #if defined(AURORA_COMPILER_GCC)
        asm volatile("yield");
    #else
        __yield();
    #endif
    #else
        // TODO: your platform here
        AuThreading::ContextYield();
    #endif
    }

    template <typename T>
    bool auline YieldToSharedCore(long spin, T callback)
    {
        if (callback())
        {
            return true;
        }

        if (gSpinAdaptiveThreshold)
        {
            auto uNow = AuAtomicAdd(&gSpinAdaptiveCurrentCount, 1u);

            if (uNow <= gSpinAdaptiveThreshold)
            {
                auto uCount = spin;
            #if defined(AURORA_ARCH_X64) || defined(AURORA_ARCH_X86)
                auto perfCounter = __rdtsc() + uCount;
                while (__rdtsc() < perfCounter)
            #else
                while (uCount > 0)
            #endif
                {
                    if (callback())
                    {
                        AuAtomicSub(&gSpinAdaptiveCurrentCount, 1u);
                        return true;
                    }
                    else
                    {
                    #if defined(SPIN_FOUR) && SPIN_FOUR == 1
                        SMPPause();
                        SMPPause();
                        SMPPause();
                        SMPPause();
                        uCount -= 4;
                    #else
                        SMPPause();
                        uCount -= 1;
                    #endif
                    }
                }

                if (gHasThreadLocalTimeout)
                {
                    auto uCount = tlsSpinCountLocal;
                #if defined(AURORA_ARCH_X64) || defined(AURORA_ARCH_X86)
                    auto perfCounter = __rdtsc() + uCount;
                    while (__rdtsc() < perfCounter)
                #else
                    while (uCount > 0)
                #endif
                    {
                        if (callback())
                        {
                            AuAtomicSub(&gSpinAdaptiveCurrentCount, 1u);
                            return true;
                        }
                        else
                        {
                            SMPPause();
                            uCount--;
                        }
                    }
                }

                AuAtomicSub(&gSpinAdaptiveCurrentCount, 1u);
            }
            else
            {
                auto uCount = (spin) / 3;
            #if defined(AURORA_ARCH_X64) || defined(AURORA_ARCH_X86)
                auto perfCounter = __rdtsc() + uCount;
                while (__rdtsc() < perfCounter)
            #else
                while (uCount > 0)
            #endif
                {
                    if (callback())
                    {
                        AuAtomicSub(&gSpinAdaptiveCurrentCount, 1u);
                        return true;
                    }
                    else
                    {
                        SMPPause();
                        uCount--;
                    }
                }
            }

            AuAtomicSub(&gSpinAdaptiveCurrentCount, 1u);
        }
        else
        {
            auto uCount = spin;
        #if defined(AURORA_ARCH_X64) || defined(AURORA_ARCH_X86)
            auto perfCounter = __rdtsc() + uCount;
            while (__rdtsc() < perfCounter)
        #else
            while (uCount > 0)
        #endif            
            {
                if (callback())
                {
                    return true;
                }
                else
                {
                #if defined(SPIN_FOUR) && SPIN_FOUR == 1
                    SMPPause();
                    SMPPause();
                    SMPPause();
                    SMPPause();
                    uCount -= 4;
                #else
                    SMPPause();
                    uCount -= 1;
                #endif
                }
            }

            if (gHasThreadLocalTimeout)
            {
                auto uCount = tlsSpinCountLocal;
            #if defined(AURORA_ARCH_X64) || defined(AURORA_ARCH_X86)
                auto perfCounter = __rdtsc() + uCount;
                while (__rdtsc() < perfCounter)
            #else
                while (uCount > 0)
            #endif
                {
                    if (callback())
                    {
                        return true;
                    }
                    else
                    {
                        SMPPause();
                        uCount--;
                    }
                }
            }
        }

        return callback();
    }

    template <typename T>
    bool auline DoTryIf(T callback)
    {
        if (gRuntimeConfig.threadingConfig.bPlatformIsSMPProcessorOptimized)
        {
            return YieldToSharedCore(gRuntimeConfig.threadingConfig.uSpinLoopPowerA, callback);
        }
        else
        {
            return callback();
        }
    }
}