VirtualAlloc and affinity settings

microsoft · Sep 5, 2024 · e68aadd · e68aadd
1 parent a5a8853
commit e68aadd
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 0 deletions.
diff --git a/inc/cxplat_winkernel.h b/inc/cxplat_winkernel.h
@@ -119,6 +119,9 @@ CxPlatLogAssert(
 #define CXPLAT_ALLOC_PAGED(Size, Tag) ExAllocatePool2(POOL_FLAG_PAGED | POOL_FLAG_UNINITIALIZED, Size, Tag)
 #define CXPLAT_ALLOC_NONPAGED(Size, Tag) ExAllocatePool2(POOL_FLAG_NON_PAGED | POOL_FLAG_UNINITIALIZED, Size, Tag)
 #define CXPLAT_FREE(Mem, Tag) ExFreePoolWithTag((void*)Mem, Tag)
+#define CXPLAT_VIRTUAL_ALLOC(Size, TypeFlags, ProtectionFlags, Tag) \
+    CXPLAT_ALLOC_NONPAGED(Size, Tag)
+#define CXPLAT_VIRTUAL_FREE(Mem, Size, Type, Tag) CXPLAT_FREE(Mem, Tag);
 
 #define CxPlatZeroMemory RtlZeroMemory
 #define CxPlatCopyMemory RtlCopyMemory
@@ -603,6 +606,38 @@ CxPlatThreadCreate(
 typedef ULONG_PTR CXPLAT_THREAD_ID;
 #define CxPlatCurThreadID() ((CXPLAT_THREAD_ID)PsGetCurrentThreadId())
 
+inline
+BOOLEAN
+CxPlatSetThreadNodeAffinity(
+    LONG NodeAffinity
+    )
+{
+    GROUP_AFFINITY group;
+
+    KeQueryNodeActiveAffinity((USHORT)NodeAffinity, &group, NULL);
+    KeSetSystemGroupAffinityThread(&group, NULL);
+
+    return TRUE;
+}
+
+inline
+BOOLEAN
+CxPlatSetThreadGroupAffinity(
+    LONG GroupNumber,
+    DWORD_PTR CpuAffinity
+    )
+{
+    GROUP_AFFINITY group = {0};
+
+    group.Mask = CpuAffinity;
+    group.Group = (USHORT)GroupNumber;
+
+    KeSetSystemGroupAffinityThread(&group, NULL);
+
+    return TRUE;
+}
+
+
 //
 // Crypto Interfaces
 //

diff --git a/inc/cxplat_winuser.h b/inc/cxplat_winuser.h
@@ -146,6 +146,9 @@ CxPlatFree(
 #define CXPLAT_ALLOC_PAGED(Size, Tag) CxPlatAlloc(Size, Tag)
 #define CXPLAT_ALLOC_NONPAGED(Size, Tag) CxPlatAlloc(Size, Tag)
 #define CXPLAT_FREE(Mem, Tag) CxPlatFree((void*)Mem, Tag)
+#define CXPLAT_VIRTUAL_ALLOC(Size, TypeFlags, ProtectionFlags, Tag) \
+    VirtualAlloc(NULL, Size, TypeFlags, ProtectionFlags)
+#define CXPLAT_VIRTUAL_FREE(Mem, Size, Type, Tag) VirtualFree(Mem, Size, Type);
 
 #define CxPlatZeroMemory RtlZeroMemory
 #define CxPlatCopyMemory RtlCopyMemory
@@ -597,6 +600,46 @@ CxPlatThreadWaitWithTimeout(
 typedef uint32_t CXPLAT_THREAD_ID;
 #define CxPlatCurThreadID() GetCurrentThreadId()
 
+inline
+BOOLEAN
+CxPlatSetThreadNodeAffinity(
+    LONG NodeAffinity
+    )
+{
+    GROUP_AFFINITY group;
+
+    if (!GetNumaNodeProcessorMaskEx((USHORT)NodeAffinity, &group)) {
+        CXPLAT_DBG_ASSERT(FALSE);
+        return FALSE;
+    }
+    if (!SetThreadGroupAffinity(GetCurrentThread(), &group, NULL)) {
+        CXPLAT_DBG_ASSERT(FALSE);
+        return FALSE;
+    }
+
+    return TRUE;
+}
+
+inline
+BOOLEAN
+CxPlatSetThreadGroupAffinity(
+    LONG GroupNumber,
+    DWORD_PTR CpuAffinity
+    )
+{
+    GROUP_AFFINITY group = {0};
+
+    group.Mask = CpuAffinity;
+    group.Group = (WORD)GroupNumber;
+    if (!SetThreadGroupAffinity(GetCurrentThread(), &group, NULL)) {
+        CXPLAT_DBG_ASSERT(FALSE);
+        return FALSE;
+    }
+
+    return TRUE;
+}
+
+
 //
 // Crypto Interfaces
 //