Implement setThreadAffinity for Linux.

This commit is contained in:
Natalya McKay
2025-05-25 00:31:44 -04:00
parent c283f6eac4
commit 4660ab927c
2 changed files with 37 additions and 2 deletions

View File

@@ -5,6 +5,7 @@ import org.lwjgl.system.MemoryStack;
import org.lwjgl.system.MemoryUtil; import org.lwjgl.system.MemoryUtil;
import org.lwjgl.system.Platform; import org.lwjgl.system.Platform;
import org.lwjgl.system.windows.Kernel32; import org.lwjgl.system.windows.Kernel32;
import org.lwjgl.system.APIUtil;
//Platform specific code to assist in thread utilities //Platform specific code to assist in thread utilities
public class ThreadUtils { public class ThreadUtils {
@@ -15,6 +16,7 @@ public class ThreadUtils {
private static final boolean isWindows = Platform.get() == Platform.WINDOWS; private static final boolean isWindows = Platform.get() == Platform.WINDOWS;
private static final long SetThreadPriority; private static final long SetThreadPriority;
private static final long SetThreadSelectedCpuSetMasks; private static final long SetThreadSelectedCpuSetMasks;
private static final long schedSetaffinity;
static { static {
if (isWindows) { if (isWindows) {
SetThreadPriority = Kernel32.getLibrary().getFunctionAddress("SetThreadPriority"); SetThreadPriority = Kernel32.getLibrary().getFunctionAddress("SetThreadPriority");
@@ -23,6 +25,13 @@ public class ThreadUtils {
SetThreadPriority = 0; SetThreadPriority = 0;
SetThreadSelectedCpuSetMasks = 0; SetThreadSelectedCpuSetMasks = 0;
} }
if (Platform.get() == Platform.LINUX) {
var libc = APIUtil.apiCreateLibrary("libc.so.6");
schedSetaffinity = APIUtil.apiGetFunctionAddress(libc, "sched_setaffinity");
} else {
schedSetaffinity = 0;
}
} }
public static boolean SetThreadSelectedCpuSetMasksWin32(long mask) { public static boolean SetThreadSelectedCpuSetMasksWin32(long mask) {
@@ -70,4 +79,22 @@ public class ThreadUtils {
} }
return true; return true;
} }
public static boolean schedSetaffinityLinux(long masks[]) {
if (schedSetaffinity == 0 || isWindows) {
return false;
}
try (var stack = MemoryStack.stackPush()) {
long ptr = stack.ncalloc(8, masks.length, 8);
for (int i=0; i<masks.length; i++) {
MemoryUtil.memPutLong(ptr+i*8L, masks[i]);
}
int retVal = JNI.invokePPI(0, (long)masks.length*8, ptr, schedSetaffinity);
if (retVal != 0) {
throw new IllegalStateException();
}
return true;
}
}
} }

View File

@@ -25,7 +25,8 @@ public class CpuLayout {
} }
public static void setThreadAffinity(Affinity... affinities) { public static void setThreadAffinity(Affinity... affinities) {
if (Platform.get() == Platform.WINDOWS) { var platform = Platform.get();
if (platform == Platform.WINDOWS) {
long[] msks = new long[affinities.length]; long[] msks = new long[affinities.length];
short[] groups = new short[affinities.length];Arrays.fill(groups, (short) -1); short[] groups = new short[affinities.length];Arrays.fill(groups, (short) -1);
int i = 0; int i = 0;
@@ -36,8 +37,15 @@ public class CpuLayout {
msks[idx] |= a.msk; msks[idx] |= a.msk;
} }
ThreadUtils.SetThreadSelectedCpuSetMasksWin32(Arrays.copyOf(msks, i), Arrays.copyOf(groups, i)); ThreadUtils.SetThreadSelectedCpuSetMasksWin32(Arrays.copyOf(msks, i), Arrays.copyOf(groups, i));
} else if (platform == Platform.LINUX) {
Arrays.sort(affinities, (a, b) -> a.group - b.group);
long[] msks = new long[affinities.length];
for (int i=0; i<affinities.length; i++) {
msks[i] = affinities[i].msk;
}
ThreadUtils.schedSetaffinityLinux(msks);
} else { } else {
Logger.error("TODO: THIS"); Logger.error("Don't know how to set thread affinity on this platform.");
} }
} }