L4 completely add support for multi predictor, working on ver 5 mat para

2026-04-10 18:05:40 +08:00
parent b5bee5c442
commit 8a66c1697a
16 changed files with 602 additions and 191 deletions
@@ -1,6 +1,6 @@
 //-----------------------------------------------------------------------
 // -----------------------------------------------------------------------
-// file="ConcurrentList"
+// file="ConcurrentFlexibaleArray"
 // Project: StgSharp
 // AuthorGroup: Nitload
 // Copyright (c) Nitload. All rights reserved.
@@ -96,5 +96,23 @@ namespace StgSharp.Common.Collections
            }
        }

+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal Span<T> AsSpan()
+        {
+            return _array.AsSpan();
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal void ReleaseSpanOperation()
+        {
+            _rwLock.ExitReadLock();
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal void RequestSpanOperation()
+        {
+            _rwLock.EnterReadLock();
+        }
+
    }
 }
@@ -45,10 +45,10 @@ namespace StgSharp.HighPerformance.Memory
        {

            internal int _id;
+            internal int _predictor;
            internal int _refCount;
            internal nuint _address;
            internal nuint _origin;
-            internal uint _predictor;
            internal uint _profile;

        }
@@ -62,6 +62,10 @@ namespace StgSharp.HighPerformance.Memory
            int remain = count;
            Vector128<byte> three = Vector128.Create((byte)0x03);
            byte* p = (byte*)_eviction;
+            _predictors.RequestSpanOperation();
+            Span<IL4Predict> pSpan = _predictors.AsSpan();
+            try
+            {
                for (int i = 0; (i < CacheLineCount) && (remain > 0); i += 16)
                {
                    int evictMask = _intrinsic.ScanEvict(p + i, three);
@@ -75,7 +79,7 @@ namespace StgSharp.HighPerformance.Memory
                            // EvictLine(i + index);
                            _ = _map.TryRemove(head->_origin, out _);
                            int id = (int)head->_predictor;
-                        IL4Predict predict = _predictions[id] ??
+                            IL4Predict predict = pSpan[id] ??
                                throw new L4PredictUnregisteredException(id);
                            predict.WriteBack(head->_origin, head->_profile, new ReadOnlySpan<byte>((byte*)head->_address, Unsafe.SizeOf<CacheLineData>()));
                            RecycleLine(index);
@@ -86,6 +90,11 @@ namespace StgSharp.HighPerformance.Memory
                }
                return remain;
            }
+            finally
+            {
+                _predictors.ReleaseSpanOperation();
+            }
+        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private void RecycleLine(
@@ -28,6 +28,7 @@
 using StgSharp.Common.Collections;
 using StgSharp.Mathematics.Memory;
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 using System.Text;
@@ -40,14 +41,16 @@ namespace StgSharp.HighPerformance.Memory

        private ConcurrentBufferStack<int> _idReuse = new(32);

-        private ConcurrentFlexibleArray<IL4Predict> _predictions = new(32);
+        private ConcurrentFlexibleArray<IL4Predict> _predictors = new(32);

-        private int _largestId = 0;
+        private ConcurrentQueue<int> _predictionQueue = new();

-        private Queue<int> _predictionQueue = [];
+        private int _largestId ;

        private IL4Predict CurrentPredict { get; set; }

+        private int CurrentPredictIndex { get; set; }
+
        public int RegisterPredict(
                   IL4Predict predict
        )
@@ -58,7 +61,7 @@ namespace StgSharp.HighPerformance.Memory
            if (!_idReuse.TryPop(out int id)) {
                id = _largestId++;
            }
-            _predictions[id] = predict;
+            _predictors[id] = predict;
            _predictionQueue.Enqueue(id);
            return id;
        }
@@ -70,29 +73,30 @@ namespace StgSharp.HighPerformance.Memory
            if (id < 0) {
                return null;
            }
-            IL4Predict p = _predictions[id];
+            IL4Predict p = _predictors[id];
            if (p == null) {
                return p;
            }
-            _predictions[id] = null!;
-
-            // TODO: scan and write back
-
+            _predictors[id] = null!;
            _idReuse.Push(id);
            return p;
        }

-        private void UseNewPredict()
+        private bool TryGetNextPredict(
+        )
        {
-            if (_predictionQueue.Count == 0) {
-                return;
+            if (_predictionQueue.IsEmpty) {
+                return false;
            }
            if (_predictionQueue.TryDequeue(out int id))
            {
-                IL4Predict predict = _predictions[id] ??
+                IL4Predict predict = _predictors[id] ??
                    throw new L4PredictUnregisteredException(id);
                CurrentPredict = predict;
+                CurrentPredictIndex = id;
+                return true;
            }
+            return false;
        }

    }
@@ -45,6 +45,11 @@ namespace StgSharp.HighPerformance.Memory
    public interface IL4Predict
    {

+        bool IsSamePolicy(
+             uint policy,
+             IL4Predict predict
+        );
+
        PredictResult Predict(
                      Span<CacheLinePrediction> node,
                      out int count
@@ -65,7 +70,7 @@ namespace StgSharp.HighPerformance.Memory
        /// </param>
        void Prefetch(
             nuint origin,
-             ulong policy,
+             uint policy,
             Span<byte> cache
        );

@@ -42,9 +42,13 @@ namespace StgSharp.HighPerformance.Memory
            if (_activeCount > MaxActivateCacheLineCount) {
                throw new OverflowException("Too many activated cache line in external reference.");
            }
+            IL4Predict predict = CurrentPredict;
            Span<CacheLinePrediction> span = new(_prediction, PredictionCount);
-            PredictResult result = PrefetchPredict.Predict(span, out int actualCount);
-            if ((result == PredictResult.Exhausted) && (actualCount <= 0))
+            PredictResult result = predict.Predict(span, out int actualCount);
+            if ((result & PredictResult.Exhausted) != 0) {
+                _exhausted = !TryGetNextPredict();
+            }
+            if (actualCount <= 0)
            {
                _exhausted = true;
                return;
@@ -62,18 +66,43 @@ namespace StgSharp.HighPerformance.Memory
                    missIndex++;
                }
            }
-
+            _predictors.RequestSpanOperation();
+            Span<IL4Predict> predictorSpan = _predictors.AsSpan();
+            try
+            {
                // now prefetch the misses
                for (int i = 0; i < missIndex; i++)
                {
-                ref CacheLinePrediction handle = ref span[i];
-                nuint address = handle.Address;
-                if (TestHit(address))
+                    CacheLineHead* head;
+                    ref CacheLinePrediction pHandle = ref span[i];
+                    nuint address = pHandle.Address;
+                    int index;
+                    if (_map.TryGet(address, out nuint handle))
                    {
+                        head = (CacheLineHead*)handle;
+                        index = head->_id;
+                        (*_eviction)[index] = 0; // reset eviction counter
+
+                        #region update predictor belonging
+
+                        int predictId = head->_predictor;
+                        IL4Predict belong = predictorSpan[predictId];
+                        if (predict.IsSamePolicy(head->_profile, belong))
+                        {
+                            head->_predictor = predictId; // wrong line
+                        } else
+                        {
+                            belong.WriteBack(head->_origin, head->_profile, new ReadOnlySpan<byte>((byte*)head->_address, Unsafe.SizeOf<CacheLineData>()));
+                            head->_profile = pHandle.MapPolicy;
+                            predict.Prefetch(head->_origin, head->_profile, new Span<byte>((byte*)head->_address, Unsafe.SizeOf<CacheLineData>()));
+                        }
+
+                        #endregion
+
                        continue;
                    }

-                int retryCount = 0, index, remain = missIndex - i;
+                    int retryCount = 0, remain = missIndex - i;
                    const int maxRetry = 1024;
                    while (!TryGetEmptyLine(out index))
                    {
@@ -82,33 +111,27 @@ namespace StgSharp.HighPerformance.Memory
                        if (retryCount > maxRetry) {
                            throw new OverflowException($"Failed to evict cache line after {maxRetry} retries.");
                        }
+                        if (retryCount > maxRetry / 2) {
+                            Thread.Sleep(1);
                        }
-                CacheLineHead* head = &_head[index];
+                    }
+                    head = &_head[index];
                    head->_id = index;
                    head->_origin = address;
-                head->_profile = handle.MapPolicy;
+                    head->_profile = pHandle.MapPolicy;
                    head->_refCount = 0;
                    head->_address = (nuint)(_data + index);

                    // pre fetch them here
-                PrefetchPredict.Prefetch(head->_origin, head->_profile, new Span<byte>((byte*)head->_address, Unsafe.SizeOf<CacheLineData>()));
+                    predict.Prefetch(head->_origin, head->_profile, new Span<byte>((byte*)head->_address, Unsafe.SizeOf<CacheLineData>()));
                    _ = _map.TryAddOrSet(address, (nuint)head, out _);
                }
                _ = Interlocked.Add(ref _aheadCount, actualCount);
            }
-
-        private partial bool TestHit(
-                             nuint origin
-        )
+            finally
            {
-            if (_map.TryGet(origin, out nuint handle))
-            {
-                CacheLineHead* head = (CacheLineHead*)handle;
-                int index = head->_id;
-                (*_eviction)[index] = 0; // reset eviction counter
-                return true;
+                _predictors.ReleaseSpanOperation();
            }
-            return false;
        }

    }
@@ -26,6 +26,7 @@
 // -----------------------------------------------------------------------
 // -----------------------------------------------------------------------
 using StgSharp.Collections;
+using StgSharp.Common.HighPerformance.ProcessorAbstraction;
 using System.Data;
 using System.Runtime.InteropServices;

@@ -49,14 +50,20 @@ namespace StgSharp.HighPerformance.Memory
        private int _activeCount;

        private int _aheadCount;
+
+        private readonly int _bindedNuma;
        private int _maxId;
        private volatile int _prefetching;
        private readonly nuint _baseAddress;
        private readonly SwissTable _map;
        private readonly UnmanagedStack<int> _idMux;

-        public L4()
+        public L4(
+               int bindedNuma
+        )
        {
+            _bindedNuma = bindedNuma;
+
            // init order is fallow by align requirement of each component.
            nuint cacheSize = ((nuint)Unsafe.SizeOf<CacheLineData>()) * CacheLineCount;   // SIMDID, no less than 16 byte bit
            nuint entrySize = ((nuint)Unsafe.SizeOf<CacheLineHead>()) * CacheLineCount;   // 16 byte align
@@ -67,7 +74,7 @@ namespace StgSharp.HighPerformance.Memory

            // malloc all memory at once, and free at once in Dispose.
            // This is to avoid fragmentation and improve performance.
-            _baseAddress = (nuint)NativeMemory.AlignedAlloc(cacheSize + entrySize + stackSize + mapSize + predictionSize + rrpvSize, 16);
+            _baseAddress = (nuint)Numa.Alloc(bindedNuma, cacheSize + entrySize + stackSize + mapSize + predictionSize + rrpvSize);
            nuint ptr = _baseAddress;

            // init all components with the buffer, and move the buffer pointer forward.
@@ -86,6 +93,8 @@ namespace StgSharp.HighPerformance.Memory
            ptr += stackSize;
        }

+        public static nuint AllocSize { get; }
+
        // public IL4Predict PrefetchPredict { get; set; }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -114,7 +123,15 @@ namespace StgSharp.HighPerformance.Memory
            _disposed = true;
            _map.Dispose();
            _idMux.Dispose();
-            NativeMemory.AlignedFree((void*)_baseAddress);
+
+
+            nuint cacheSize = ((nuint)Unsafe.SizeOf<CacheLineData>()) * CacheLineCount;   // SIMDID, no less than 16 byte bit
+            nuint entrySize = ((nuint)Unsafe.SizeOf<CacheLineHead>()) * CacheLineCount;   // 16 byte align
+            nuint predictionSize = ((nuint)Unsafe.SizeOf<CacheLinePrediction>()) * PredictionCount;   //16 byte align
+            nuint mapSize = (nuint)SwissTable.RequestedNativeMemorySize;     // 16 byte align
+            nuint rrpvSize = (nuint)Unsafe.SizeOf<EvictionArray>();           // 16 byte align
+            nuint stackSize = ((nuint)sizeof(int)) * CacheLineCount;                      // int align
+            Numa.Free((void*)_baseAddress, _bindedNuma, cacheSize + entrySize + stackSize + mapSize + predictionSize + rrpvSize);
            GC.SuppressFinalize(this);
        }

@@ -164,10 +181,6 @@ namespace StgSharp.HighPerformance.Memory

        private partial void Prefetch();

-        private partial bool TestHit(
-                             nuint origin
-        );
-
        private struct CacheLineData
        {

@@ -25,8 +25,10 @@
 //     
 // -----------------------------------------------------------------------
 // -----------------------------------------------------------------------
+using StgSharp.HighPerformance.Memory;
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.Contracts;
 using System.Linq;
 using System.Runtime.InteropServices;
 using System.Text;
@@ -43,6 +45,35 @@ namespace StgSharp.Common.HighPerformance.ProcessorAbstraction

        public static ReadOnlySpan<NumaNode> AllNumaNode => _numaNodes;

+        public static unsafe void* Alloc(
+                                   int numaId,
+                                   nuint size
+        )
+        {
+            if (_numaNodes.Length == 0) {
+                throw new InvalidOperationException("Numa info has not been loaded.");
+            }
+            if ((numaId < 0) || (numaId >= _numaNodes.Length)) {
+                throw new ArgumentOutOfRangeException(nameof(numaId), "Invalid NUMA node ID.");
+            }
+            return NumaAllocNative(size, numaId);
+        }
+
+        public static unsafe void Free(
+                                  void* address,
+                                  int numaId,
+                                  nuint size
+        )
+        {
+            if (_numaNodes.Length == 0) {
+                throw new InvalidOperationException("Numa info has not been loaded.");
+            }
+            if ((numaId < 0) || (numaId >= _numaNodes.Length)) {
+                throw new ArgumentOutOfRangeException(nameof(numaId), "Invalid NUMA node ID.");
+            }
+            NumaFreeNative(address, size, numaId);
+        }
+
        [LibraryImport(Native.LibName, EntryPoint = "sn_get_numa_each")]
        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
        [DefaultDllImportSearchPaths(DllImportSearchPath.AssemblyDirectory)]
@@ -60,6 +91,20 @@ namespace StgSharp.Common.HighPerformance.ProcessorAbstraction
                                  Span<int> coreCountInNuma
        );

+        public static L4[] RequestBufferForEachNuma()
+        {
+            if (_numaNodes.Length == 0) {
+                throw new InvalidOperationException("Numa info has not been loaded.");
+            }
+            L4[] buffers = new L4[NumaNodeCount];
+            for (int i = 0; i < NumaNodeCount; i++)
+            {
+                L4 buffer = new L4(i);
+                buffers[i] = buffer;
+            }
+            return buffers;
+        }
+
        internal static void Initialize()
        {
            Span<int> coreCountInNuma = stackalloc int[(Environment.ProcessorCount / 32) + 1];
@@ -73,6 +118,13 @@ namespace StgSharp.Common.HighPerformance.ProcessorAbstraction
            }
        }

+        [LibraryImport(Native.LibName, EntryPoint = "sn_set_thread_affinity")]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
+        [DefaultDllImportSearchPaths(DllImportSearchPath.AssemblyDirectory)]
+        internal static partial void SetThreadAffinity(
+                                     int coreId
+        );
+
        [MethodImpl(MethodImplOptions.NoInlining)]
        private static NumaNode InitNumaNode(
                                uint numaId,
@@ -84,5 +136,22 @@ namespace StgSharp.Common.HighPerformance.ProcessorAbstraction
            return new(numaId, count, globalIndex);
        }

+        [LibraryImport(Native.LibName, EntryPoint = "sn_numa_alloc")]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
+        [DefaultDllImportSearchPaths(DllImportSearchPath.AssemblyDirectory)]
+        private static unsafe partial void* NumaAllocNative(
+                                            nuint size,
+                                            int numaId
+        );
+
+        [LibraryImport(Native.LibName, EntryPoint = "sn_numa_free")]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
+        [DefaultDllImportSearchPaths(DllImportSearchPath.AssemblyDirectory)]
+        private static unsafe partial void NumaFreeNative(
+                                           void* ptr,
+                                           nuint size,
+                                           int numa_id
+        );
+
    }
 }
@@ -0,0 +1,48 @@
+//-----------------------------------------------------------------------
+// -----------------------------------------------------------------------
+// file="MatrixParallel.Numa"
+// Project: StgSharp
+// AuthorGroup: Nitload
+// Copyright (c) Nitload. All rights reserved.
+//     
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+//     
+// -----------------------------------------------------------------------
+// -----------------------------------------------------------------------
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace StgSharp.Mathematics.Numeric
+{
+    public static unsafe partial class MatrixParallel
+    {
+
+        private static int[] NumaPayload;
+        private static MatrixParallelThread[,] NumaThreads;
+
+        private static void UpdateThreadPayload()
+        {
+            for (int i = 0; i < NumaPayload.Length; i++) { }
+        }
+
+    }
+}
@@ -69,12 +69,8 @@ namespace StgSharp.Mathematics.Numeric

        public ThreadPriority AfterWork { get; set; }

-        internal ref readonly IntrinsicContext Intrinsic => ref (*_context);
-
        internal bool IsSingle { get; set; }

-        internal MatrixParallelWrap BindedWrap { get; set; }
-
        internal MatrixThreadState State { get; private set; }

        internal SIMDID Simd { get; private set; }
@@ -163,7 +159,6 @@ namespace StgSharp.Mathematics.Numeric
                }
                State = MatrixThreadState.Idle;
                State = MatrixThreadState.Running;
-                AfterWork = BindedWrap.AfterWork;
                _managedThread.Priority = ThreadPriority.Highest;
                if (IsSingle)
                {
@@ -173,17 +168,6 @@ namespace StgSharp.Mathematics.Numeric
                    IsSingle = false;
                } else
                {
-                    PredictPolicy policy = BindedWrap.Policy;
-                    switch (policy)
-                    {
-                        case PredictPolicy.Transparent:
-                            break;
-                        case PredictPolicy.BufferSequential:
-                            break;
-                        default:
-                            break;
-                    }
-
                    // MatrixParallelQueue taskQueue = BindedWrap.TaskQueue;
                    /*
                    if (0 > (int)p ->ComputeHandle) {
@@ -198,7 +182,6 @@ namespace StgSharp.Mathematics.Numeric
                            break;
                    }
                    /**/
-                    BindedWrap.ReportThreadAccomplished();
                }
                _managedThread.Priority = AfterWork;
                State = MatrixThreadState.Idle;
@@ -208,19 +191,19 @@ namespace StgSharp.Mathematics.Numeric

 #pragma warning disable CA5393
        [LibraryImport(Native.LibName, EntryPoint = "sn_get_simd_level_local")]
-        [UnmanagedCallConv(CallConvs = [ typeof(CallConvCdecl) ])]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
        [DefaultDllImportSearchPaths(DllImportSearchPath.AssemblyDirectory)]
        internal static partial ulong GetSimdId();

        [LibraryImport(Native.LibName, EntryPoint = "sn_set_thread_affinity")]
-        [UnmanagedCallConv(CallConvs = [ typeof(CallConvCdecl) ])]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
        [DefaultDllImportSearchPaths(DllImportSearchPath.AssemblyDirectory)]
        internal static partial void SetThreadAffinity(
                                     int coreId
        );

        [LibraryImport(Native.LibName, EntryPoint = "load_intrinsic_function")]
-        [UnmanagedCallConv(CallConvs = [ typeof(CallConvCdecl) ])]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
        [DefaultDllImportSearchPaths(DllImportSearchPath.AssemblyDirectory)]
        private static unsafe partial void LoadIntrinsicFunction(
                                           IntrinsicContext* context,
@@ -1,73 +0,0 @@
-#ifdef _MSC_VER
-#pragma once
-#endif
-
-#ifndef SSC
-#define SSC
-
-#include "ssc_internal.h"
-
-#define STBIDEF SSCAPI
-#include "stb_image.h"
-
-#include "gl.h"
-#include "wgl.h"
-#include "glfw3.h"
-#include <immintrin.h>
-#include <assert.h>
-
-#ifndef SSCAPI_DEFINE
-
-#ifdef _WIN32
-#define SSCAPI __declspec(dllexport)
-#elif defined _LINUX_
-#define SSCAPI __attribute__((dllexport))
-#elif defined _APPLE_
-#define SSCAPI __attribute__((visibility("default")))
-#else
-#define SSCAPI
-#endif
-
-#endif // ! SSCAPI_DEFINE
-
-#if __cplusplus
-extern "C" {
-#endif
-
-#define byte char
-#define cs_char uint16_t
-#define ssc_null_assert(ptr) assert(ptr != NULL)
-
-#define ALIGN(simdVec) _mm_loadu_ps(&simdVec)
-#define ALIGN256(simdVec) _mm256_loadu_ps(&simdVec)
-#define ALIGN512(simdVec) _mm512_loadu_ps(&simdVec)
-
-typedef struct Image {
-        int width;
-        int height;
-        int channel;
-        char *pixelPtr;
-} Image;
-
-typedef char *(*imageLoader)(char const *filename, int *x, int *y, int *channels_in_file,
-                             int desired_channels);
-
-SSCAPI char infolog[512];
-SSCAPI int SSCDECL glCheckShaderStat(GladGLContext *context, uint64_t shaderHandle, int key,
-                                     char **logRef);
-SSCAPI void SSCDECL initGL(int majorVersion, int minorVersion);
-SSCAPI unsigned int SSCDECL linkShaderProgram(GladGLContext *context, GLuint shadertype);
-SSCAPI void SSCDECL loadImageData(char *location, Image *out, imageLoader loader);
-SSCAPI GLFWglproc SSCDECL loadGlfuncDefault(char *procName);
-SSCAPI void SSCDECL unloadImageData(Image *out);
-SSCAPI char *SSCDECL readLog(void);
-
-#pragma endregion
-
-SSCAPI int SSCDECL load_intrinsic_function(void *intrinsic_context);
-
-#ifdef __cplusplus
-}
-#endif
-
-#endif
@@ -55,7 +55,8 @@ SN_API SIMDID SN_DECL sn_get_simd_level_local();
 SN_API SIMDID SN_DECL sn_get_unite_simd(SIMDID left, SIMDID right);
 SN_API int SN_DECL sn_compare_simd(SIMDID left, SIMDID right);
 SN_API void SN_DECL sn_set_thread_affinity(int core_id);
-
+SN_API void* SN_DECL sn_numa_alloc(size_t size, int numa_id);
+SN_API void SN_DECL sn_numa_free(void *ptr, size_t size, int numa_id);
 SN_API int SN_DECL sn_get_numa_each(int numa_id, uint64_t *internal, uint32_t *global);
 SN_API int SN_DECL sn_get_numa_overall(int *out_numa_count, int *count_of_each_numa);
 #endif
@@ -1,25 +1,236 @@
+#define _GNU_SOURCE
+
 #include "sn_target.h"
-#include "StgSharpNative.h"

 #if SN_IS_OS(SN_OS_LINUX)
-#define _GNU_SOURCE
+#include <errno.h>
+#include <inttypes.h>
+#include <limits.h>
+#include <linux/mempolicy.h>
 #include <pthread.h>
 #include <sched.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/mman.h>
+#include <sys/syscall.h>
 #include <unistd.h>

+typedef int (*sn_linux_id_visitor)(int value, void *context);
+
+typedef struct sn_linux_counter_context {
+        int count;
+} sn_linux_counter_context;
+
+typedef struct sn_linux_max_context {
+        int max_value;
+} sn_linux_max_context;
+
+typedef struct sn_linux_fill_cpu_context {
+        uint64_t *internal;
+        uint32_t *global;
+        int index;
+} sn_linux_fill_cpu_context;
+
+typedef struct sn_linux_fill_node_count_context {
+        int *count_of_each_numa;
+} sn_linux_fill_node_count_context;
+
+static int sn_linux_read_first_line(char const *path, char *buffer, size_t buffer_size)
+{
+        FILE *file = fopen(path, "r");
+        if (file == NULL)
+                return 0;
+
+        char *line = fgets(buffer, (int)buffer_size, file);
+        fclose(file);
+        return line != NULL;
+}
+
+static int sn_linux_parse_nonnegative_int(char const **cursor, int *value)
+{
+        char *end = NULL;
+        uintmax_t parsed_value = 0;
+
+        if (**cursor == '-')
+                return 0;
+
+        errno = 0;
+        parsed_value = strtoumax(*cursor, &end, 10);
+        if (end == *cursor || errno == ERANGE || parsed_value > (uintmax_t)INT_MAX)
+                return 0;
+
+        *cursor = end;
+        *value = (int)parsed_value;
+        return 1;
+}
+
+static int sn_linux_visit_id_list(char const *text, sn_linux_id_visitor visitor, void *context)
+{
+        char const *cursor = text;
+
+        while (*cursor != '\0') {
+                while (*cursor == ' ' || *cursor == '\t' || *cursor == '\n' || *cursor == ',')
+                        cursor++;
+
+                if (*cursor == '\0')
+                        break;
+
+                int begin = 0;
+                if (!sn_linux_parse_nonnegative_int(&cursor, &begin))
+                        return 0;
+
+                int finish = begin;
+
+                if (*cursor == '-') {
+                        cursor++;
+                        if (!sn_linux_parse_nonnegative_int(&cursor, &finish))
+                                return 0;
+                }
+
+                if (finish < begin)
+                        return 0;
+
+                for (int value = begin;; value++) {
+                        if (!visitor(value, context))
+                                return 0;
+
+                        if (value == finish)
+                                break;
+                }
+        }
+
+        return 1;
+}
+
+static int sn_linux_count_id(int value, void *context)
+{
+        (void)value;
+        ((sn_linux_counter_context *)context)->count++;
+        return 1;
+}
+
+static int sn_linux_track_max_id(int value, void *context)
+{
+        sn_linux_max_context *max_context = (sn_linux_max_context *)context;
+        if (value > max_context->max_value)
+                max_context->max_value = value;
+        return 1;
+}
+
+static int sn_linux_fill_cpu_id(int value, void *context)
+{
+        sn_linux_fill_cpu_context *fill_context = (sn_linux_fill_cpu_context *)context;
+
+        if (fill_context->internal != NULL)
+                fill_context->internal[fill_context->index] = (uint64_t)fill_context->index;
+        if (fill_context->global != NULL)
+                fill_context->global[fill_context->index] = (uint32_t)value;
+
+        fill_context->index++;
+        return 1;
+}
+
+static int sn_linux_count_cpu_list(char const *path)
+{
+        char buffer[256];
+        if (!sn_linux_read_first_line(path, buffer, sizeof(buffer)))
+                return 0;
+
+        sn_linux_counter_context counter = {0};
+        if (!sn_linux_visit_id_list(buffer, sn_linux_count_id, &counter))
+                return 0;
+
+        return counter.count;
+}
+
+static int sn_linux_fill_node_cpu_count(int value, void *context)
+{
+        char path[128];
+        snprintf(path, sizeof(path), "/sys/devices/system/node/node%d/cpulist", value);
+
+        ((sn_linux_fill_node_count_context *)context)->count_of_each_numa[value] =
+                sn_linux_count_cpu_list(path);
+        return 1;
+}
+
+static int sn_linux_node_exists(int numa_id)
+{
+        char path[128];
+        snprintf(path, sizeof(path), "/sys/devices/system/node/node%d", numa_id);
+        return access(path, F_OK) == 0;
+}
+
+static size_t sn_linux_round_up_to_page_size(size_t size)
+{
+        int64_t page_size_value = (int64_t)sysconf(_SC_PAGESIZE);
+        size_t effective_page_size = (page_size_value > 0) ? (size_t)page_size_value : 4096u;
+
+        size_t rounded = size + effective_page_size - 1;
+        if (rounded < size)
+                return 0;
+
+        return rounded - (rounded % effective_page_size);
+}
+
+static int sn_linux_bind_memory_to_node(void *address, size_t size, int numa_id,
+                                        uint32_t flags)
+{
+#if defined(SYS_mbind)
+        size_t bits_per_word = sizeof(unsigned long) * 8u;
+        size_t word_count = ((size_t)numa_id / bits_per_word) + 1u;
+        unsigned long *node_mask = (unsigned long *)calloc(word_count, sizeof(unsigned long));
+        if (node_mask == NULL)
+                return 0;
+
+        node_mask[(size_t)numa_id / bits_per_word] = 1ul << ((size_t)numa_id % bits_per_word);
+
+        int syscall_result = (int)syscall(SYS_mbind, address, size,
+                                          MPOL_BIND | MPOL_F_STATIC_NODES, node_mask,
+                                          (unsigned long)numa_id + 1ul, flags);
+        free(node_mask);
+        return syscall_result == 0;
+#else
+        (void)address;
+        (void)size;
+        (void)numa_id;
+        (void)flags;
+        return 0;
+#endif
+}
+
+static int sn_linux_populate_writable_mapping(void *address, size_t size)
+{
+#if defined(MADV_POPULATE_WRITE)
+        if (madvise(address, size, MADV_POPULATE_WRITE) == 0)
+                return 1;
+
+        return errno == EINVAL;
+#else
+        (void)address;
+        (void)size;
+        return 1;
+#endif
+}
+
 /* Pin the calling thread to a single logical processor.
 * core_id >= 0 : bind to that logical processor (global index, across groups).
 * core_id <  0 : clear affinity constraint within the current group. */
 SN_API void SN_DECL sn_set_thread_affinity(int core_id)
 {
-        long ncpus = sysconf(_SC_NPROCESSORS_ONLN);
+        int cpu_count = (int)sysconf(_SC_NPROCESSORS_ONLN);
+        if (cpu_count <= 0)
+                return;
+
+        cpu_set_t *set = CPU_ALLOC(cpu_count);
+        size_t size = CPU_ALLOC_SIZE(cpu_count);
+        if (set == NULL)
+                return;

-        cpu_set_t *set = CPU_ALLOC(ncpus);
-        size_t size = CPU_ALLOC_SIZE(ncpus);
        CPU_ZERO_S(size, set);

        if (core_id < 0) {
-                for (long i = 0; i < ncpus; i++)
+                for (int i = 0; i < cpu_count; i++)
                        CPU_SET_S((int)i, size, set);
        } else {
                CPU_SET_S(core_id, size, set);
@@ -31,16 +242,89 @@ SN_API void SN_DECL sn_set_thread_affinity(int core_id)

 SN_API int SN_DECL sn_get_numa_overall(int *out_numa_count, int *count_of_each_numa)
 {
-        (void)out_numa_count;
-        (void)count_of_each_numa;
+        char buffer[256];
+        if (!sn_linux_read_first_line("/sys/devices/system/node/online", buffer, sizeof(buffer)))
                return 0;
+
+        sn_linux_max_context max_context = {-1};
+        if (!sn_linux_visit_id_list(buffer, sn_linux_track_max_id, &max_context) ||
+            max_context.max_value < 0)
+                return 0;
+
+        if (out_numa_count != NULL)
+                *out_numa_count = max_context.max_value + 1;
+
+        if (count_of_each_numa != NULL) {
+                for (int index = 0; index <= max_context.max_value; index++)
+                        count_of_each_numa[index] = 0;
+
+                sn_linux_fill_node_count_context fill_context = {count_of_each_numa};
+                if (!sn_linux_visit_id_list(buffer, sn_linux_fill_node_cpu_count, &fill_context))
+                        return 0;
+        }
+
+        return 1;
 }

-SN_API int SN_DECL sn_get_numa_each(int numa_id, uint64_t *in_numa_cpuid, uint64_t *global_cpuid)
+SN_API int SN_DECL sn_get_numa_each(int numa_id, uint64_t *internal, uint32_t *global)
+{
+        if (numa_id < 0)
+                return 0;
+
+        char path[128];
+        char buffer[256];
+        snprintf(path, sizeof(path), "/sys/devices/system/node/node%d/cpulist", numa_id);
+        if (!sn_linux_read_first_line(path, buffer, sizeof(buffer)))
+                return 0;
+
+        sn_linux_fill_cpu_context fill_context = {internal, global, 0};
+        return sn_linux_visit_id_list(buffer, sn_linux_fill_cpu_id, &fill_context);
+}
+
+SN_API void *SN_DECL sn_numa_alloc(size_t size, int numa_id)
+{
+        size_t allocation_size = sn_linux_round_up_to_page_size(size);
+        if (allocation_size == 0)
+                return NULL;
+
+        int mmap_flags = MAP_PRIVATE | MAP_ANONYMOUS;
+#if defined(MAP_POPULATE)
+        if (numa_id >= 0)
+                mmap_flags |= MAP_POPULATE;
+#endif
+
+        void *memory = mmap(NULL, allocation_size, PROT_READ | PROT_WRITE,
+                            mmap_flags, -1, 0);
+        if (memory == MAP_FAILED)
+                return NULL;
+
+        if (numa_id < 0)
+                return memory;
+
+        uint32_t bind_flags = 0u;
+#if defined(MAP_POPULATE)
+        bind_flags = MPOL_MF_MOVE | MPOL_MF_STRICT;
+#endif
+
+        if (!sn_linux_node_exists(numa_id) ||
+            !sn_linux_bind_memory_to_node(memory, allocation_size, numa_id, bind_flags) ||
+            !sn_linux_populate_writable_mapping(memory, allocation_size)) {
+                munmap(memory, allocation_size);
+                return NULL;
+        }
+
+        return memory;
+}
+
+SN_API void SN_DECL sn_numa_free(void *ptr, size_t size, int numa_id)
 {
        (void)numa_id;
-        (void)in_numa_cpuid;
-        (void)global_cpuid;
-        return 0;
+
+        if (ptr == NULL || size == 0)
+                return;
+
+        size_t allocation_size = sn_linux_round_up_to_page_size(size);
+        if (allocation_size != 0)
+                munmap(ptr, allocation_size);
 }
 #endif
@@ -56,14 +56,41 @@ SN_API int SN_DECL sn_get_numa_each(int numa_id, uint64_t *internal, uint32_t *g

        int intra = 0;
        while (mask) {
-                uint32_t idx;
+                DWORD idx;
                _BitScanForward64(&idx, mask);
-                int _ = (internal == NULL) ? intra : (internal[intra] = intra);
-                global[intra] = (int)idx;
+                if (internal != NULL)
+                        internal[intra] = (uint64_t)intra;
+                global[intra] = (uint32_t)idx;
                intra++;
                mask &= mask - 1;
        }
        return 1;
 }

+
+SN_API void *SN_DECL sn_numa_alloc(size_t size, int numa_id)
+{
+        if (size == 0)
+                return NULL;
+
+        if (numa_id < 0)
+                return VirtualAlloc(NULL, size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+
+        ULONG highest = 0;
+        if (!GetNumaHighestNodeNumber(&highest) || (ULONG)numa_id > highest)
+                return NULL;
+
+        return VirtualAllocExNuma(GetCurrentProcess(), NULL, size, MEM_RESERVE | MEM_COMMIT,
+                                  PAGE_READWRITE, (DWORD)numa_id);
+}
+
+SN_API void SN_DECL sn_numa_free(void *ptr, size_t size, int numa_id)
+{
+        (void)size;
+        (void)numa_id;
+
+        if (ptr != NULL)
+                VirtualFree(ptr, 0, MEM_RELEASE);
+}
+
 #endif