lib/simd/src/topology/native.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const builtin = @import("builtin");
  3 const topology = @import("root.zig");
  4 const linux = if (builtin.target.os.tag == .linux)
  5     @import("linux/root.zig")
  6 else
  7     struct {};
  8 const apple = if (builtin.target.os.tag.isDarwin())
  9     @import("apple/root.zig")
 10 else
 11     struct {};
 12 const windows = if (builtin.target.os.tag == .windows)
 13     @import("windows/root.zig")
 14 else
 15     struct {};
 16 const freebsd = if (builtin.target.os.tag == .freebsd)
 17     @import("freebsd/root.zig")
 18 else
 19     struct {};
 20 const wasm = if (builtin.target.os.tag == .emscripten)
 21     @import("wasm/root.zig")
 22 else
 23     struct {};
 24 
 25 var cache_guard: std.atomic.Mutex = .unlocked;
 26 var cache_state: std.atomic.Value(u8) = std.atomic.Value(u8).init(0);
 27 var cache_storage: [4]topology.Cache align(64) = @splat(.{});
 28 
 29 pub fn haveThreadingSupport() bool {
 30     if (comptime builtin.target.os.tag == .emscripten) {
 31         return wasm.haveThreadingSupport();
 32     }
 33     return if (comptime builtin.target.cpu.arch.isWasm())
 34         !builtin.single_threaded
 35     else
 36         true;
 37 }
 38 
 39 pub fn totalLogicalProcessors() usize {
 40     if (comptime builtin.target.os.tag == .linux) {
 41         return linux.totalLogicalProcessors();
 42     }
 43     if (comptime builtin.target.os.tag.isDarwin()) {
 44         return apple.totalLogicalProcessors();
 45     }
 46     if (comptime builtin.target.os.tag == .windows) {
 47         return windows.totalLogicalProcessors();
 48     }
 49     if (comptime builtin.target.os.tag == .emscripten) {
 50         return wasm.totalLogicalProcessors();
 51     }
 52     if (comptime builtin.target.os.tag == .freestanding) return 1;
 53     const count = std.Thread.getCpuCount() catch 1;
 54     return std.math.clamp(count, 1, topology.max_logical_processors);
 55 }
 56 
 57 pub fn getThreadAffinity(out: *topology.LogicalProcessorSet) bool {
 58     if (comptime builtin.target.os.tag == .linux) {
 59         return linux.getThreadAffinity(out);
 60     } else if (comptime builtin.target.os.tag == .windows) {
 61         return windows.getThreadAffinity(out);
 62     } else if (comptime builtin.target.os.tag == .freebsd) {
 63         return freebsd.getThreadAffinity(out);
 64     } else {
 65         return false;
 66     }
 67 }
 68 
 69 pub fn setThreadAffinity(lps: *const topology.LogicalProcessorSet) bool {
 70     if (comptime builtin.target.os.tag == .linux) {
 71         return linux.setThreadAffinity(lps);
 72     } else if (comptime builtin.target.os.tag == .windows) {
 73         return windows.setThreadAffinity(lps);
 74     } else if (comptime builtin.target.os.tag == .freebsd) {
 75         return freebsd.setThreadAffinity(lps);
 76     } else {
 77         return false;
 78     }
 79 }
 80 
 81 pub fn pinThreadToLogicalProcessor(lp: usize) bool {
 82     if (lp >= topology.max_logical_processors) return false;
 83     var lps = topology.LogicalProcessorSet{};
 84     lps.set(lp);
 85     return setThreadAffinity(&lps);
 86 }
 87 
 88 pub fn init(allocator: std.mem.Allocator) std.mem.Allocator.Error!topology.Topology {
 89     if (comptime builtin.target.os.tag == .linux) {
 90         var records: [topology.max_logical_processors]topology.ProcessorRecord =
 91             undefined;
 92         const count = linux.fillRecords(&records);
 93         if (count == 0) return .{};
 94         var detected = topology.Topology.initFromRecords(
 95             allocator,
 96             records[0..count],
 97         ) catch |err| switch (err) {
 98             error.OutOfMemory => return error.OutOfMemory,
 99             else => return .{},
100         };
101         linux.fillClusterCacheSizes(&detected);
102         return detected;
103     }
104     if (comptime builtin.target.os.tag.isDarwin()) {
105         var records: [topology.max_logical_processors]topology.ProcessorRecord =
106             undefined;
107         const count = apple.fillRecords(&records);
108         if (count == 0) return .{};
109         return topology.Topology.initFromRecords(
110             allocator,
111             records[0..count],
112         ) catch |err| switch (err) {
113             error.OutOfMemory => return error.OutOfMemory,
114             else => return .{},
115         };
116     }
117     if (comptime builtin.target.os.tag == .windows) {
118         var records: [topology.max_logical_processors]topology.ProcessorRecord =
119             undefined;
120         const count = try windows.fillRecords(allocator, &records);
121         if (count == 0) return .{};
122         return topology.Topology.initFromRecords(
123             allocator,
124             records[0..count],
125         ) catch |err| switch (err) {
126             error.OutOfMemory => return error.OutOfMemory,
127             else => return .{},
128         };
129     }
130     return .{};
131 }
132 
133 pub fn dataCaches() ?*const [4]topology.Cache {
134     const state = cache_state.load(.acquire);
135     if (state != 0) return if (state == 1) &cache_storage else null;
136     while (!cache_guard.tryLock()) std.atomic.spinLoopHint();
137     defer cache_guard.unlock();
138     const locked_state = cache_state.load(.monotonic);
139     if (locked_state == 0) {
140         const detected = if (comptime builtin.target.os.tag == .linux)
141             linux.dataCaches()
142         else if (comptime builtin.target.os.tag.isDarwin())
143             apple.dataCaches()
144         else if (comptime builtin.target.os.tag == .windows)
145             windows.dataCaches()
146         else
147             null;
148         if (detected) |caches| {
149             cache_storage = caches;
150             cache_state.store(1, .release);
151         } else {
152             cache_state.store(2, .release);
153         }
154     }
155     return if (cache_state.load(.acquire) == 1) &cache_storage else null;
156 }