lib/pluck/src/wmc.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const sys = @import("sys");
3 const Allocator = std.mem.Allocator;
4
5 const bdd = @import("bdd.zig");
6 const Bdd = bdd.Bdd;
7 const Manager = bdd.Manager;
8 const WmcParams = bdd.WmcParams;
9
10 const weight_dd = @import("weight.zig");
11 const WeightDD = weight_dd.WeightDD;
12 const Weight = weight_dd.Weight;
13
14 const runtime = @import("runtime.zig");
15 const RuntimeValue = runtime.RuntimeValue;
16
17 pub const World = runtime.GuardedWorld;
18
19 pub const WeightedResult = struct {
20 value: *RuntimeValue,
21 probability: f64,
22 };
23
24 pub const DeferredCaches = struct {
25 deferred: std.AutoHashMap(u64, f64),
26 weighted: std.AutoHashMap(u64, f64),
27
28 pub fn init(allocator: Allocator) DeferredCaches {
29 return .{
30 .deferred = std.AutoHashMap(u64, f64).init(allocator),
31 .weighted = std.AutoHashMap(u64, f64).init(allocator),
32 };
33 }
34
35 pub fn deinit(self: *DeferredCaches) void {
36 self.weighted.deinit();
37 self.deferred.deinit();
38 }
39 };
40
41 const ParallelWmcContext = struct {
42 cache_allocator: Allocator,
43 worlds: []const World,
44 weighted_results: []WeightedResult,
45 wmc_params: *const WmcParams,
46 weight_dd: *const WeightDD,
47 weight_root: Weight,
48 };
49
50 fn parallelWmcWorker(ctx: *const ParallelWmcContext, start: usize, end: usize) void {
51 var cache = std.AutoHashMap(u64, f64).init(ctx.cache_allocator);
52 defer cache.deinit();
53
54 for (start..end) |i| {
55 ctx.weighted_results[i] = WeightedResult{
56 .value = ctx.worlds[i].value,
57 .probability = weight_dd.wmcWeightedWithCache(
58 ctx.weight_dd,
59 ctx.worlds[i].guard,
60 ctx.weight_root,
61 ctx.wmc_params,
62 &cache,
63 ),
64 };
65 }
66 }
67
68 pub fn computeWmcParallel(
69 result_allocator: Allocator,
70 cache_allocator: Allocator,
71 worlds: []const World,
72 wmc_params: *const WmcParams,
73 weight_dd_ctx: *const WeightDD,
74 weight_root: Weight,
75 num_threads: ?usize,
76 ) ![]WeightedResult {
77 if (!sys.thread.threadsSupported()) {
78 return computeWmcSequential(result_allocator, worlds, wmc_params, weight_dd_ctx, weight_root);
79 } else {
80 if (worlds.len == 0) {
81 return try result_allocator.alloc(WeightedResult, 0);
82 }
83
84 const weighted_results = try result_allocator.alloc(WeightedResult, worlds.len);
85 errdefer result_allocator.free(weighted_results);
86
87 const worker_count = @max(1, num_threads orelse 1);
88
89 const ctx = ParallelWmcContext{
90 .cache_allocator = cache_allocator,
91 .worlds = worlds,
92 .weighted_results = weighted_results,
93 .wmc_params = wmc_params,
94 .weight_dd = weight_dd_ctx,
95 .weight_root = weight_root,
96 };
97
98 if (worker_count == 1 or worlds.len < 2) {
99 parallelWmcWorker(&ctx, 0, worlds.len);
100 return weighted_results;
101 }
102
103 const total_threads = @min(worker_count, worlds.len);
104 const worker_threads = total_threads - 1;
105 const threads = try result_allocator.alloc(sys.thread.JoinHandle, worker_threads);
106 defer result_allocator.free(threads);
107 var spawned_threads: usize = 0;
108 errdefer {
109 for (threads[0..spawned_threads]) |thread| {
110 thread.join();
111 }
112 }
113
114 var start: usize = 0;
115 const base_chunk_size = worlds.len / total_threads;
116 const remainder = worlds.len % total_threads;
117
118 for (threads, 0..) |*thread, i| {
119 const extra: usize = if (i < remainder) 1 else 0;
120 const chunk_size = base_chunk_size + extra;
121 const end = start + chunk_size;
122 thread.* = try sys.thread.spawn(parallelWmcWorker, .{ &ctx, start, end });
123 spawned_threads += 1;
124 start = end;
125 }
126
127 if (start < worlds.len) {
128 parallelWmcWorker(&ctx, start, worlds.len);
129 }
130
131 for (threads) |thread| {
132 thread.join();
133 }
134
135 return weighted_results;
136 }
137 }
138
139 pub fn computeWmcSequential(
140 allocator: Allocator,
141 worlds: []const World,
142 wmc_params: *const WmcParams,
143 weight_dd_ctx: *const WeightDD,
144 weight_root: Weight,
145 ) ![]WeightedResult {
146 var weighted_results = try allocator.alloc(WeightedResult, worlds.len);
147 errdefer allocator.free(weighted_results);
148
149 var cache = std.AutoHashMap(u64, f64).init(allocator);
150 defer cache.deinit();
151
152 for (worlds, 0..) |world, i| {
153 weighted_results[i] = WeightedResult{
154 .value = world.value,
155 .probability = weight_dd.wmcWeightedWithCache(weight_dd_ctx, world.guard, weight_root, wmc_params, &cache),
156 };
157 }
158 return weighted_results;
159 }
160
161 test "DeferredCaches separates deferred and weighted key spaces" {
162 var caches = DeferredCaches.init(std.testing.allocator);
163 defer caches.deinit();
164
165 const key: u64 = 42;
166 try caches.deferred.put(key, 1.25);
167 try caches.weighted.put(key, 2.5);
168
169 try std.testing.expectApproxEqAbs(@as(f64, 1.25), caches.deferred.get(key).?, 1e-12);
170 try std.testing.expectApproxEqAbs(@as(f64, 2.5), caches.weighted.get(key).?, 1e-12);
171 }
172
173 test "parallel WMC produces same results as sequential" {
174 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
175 defer arena.deinit();
176 const allocator = arena.allocator();
177
178 var manager = try Manager.init(allocator);
179 defer manager.deinit();
180
181 var weight_ctx = try WeightDD.init(allocator, &manager);
182 defer weight_ctx.deinit();
183 const weight_root = try weight_ctx.leaf(1.0);
184
185 const num_worlds = 32;
186 const worlds = try allocator.alloc(World, num_worlds);
187
188 var vars: [5]Bdd = undefined;
189 for (0..5) |i| {
190 vars[i] = try manager.newVar(true);
191 }
192
193 for (0..num_worlds) |i| {
194 const val = try RuntimeValue.initNative(allocator, .{ .int = @intCast(i) });
195 var guard = Bdd.TRUE;
196 for (0..5) |j| {
197 const use_var = (i >> @intCast(j)) & 1 == 1;
198 if (use_var) {
199 guard = try manager.bddAnd(guard, vars[j]);
200 }
201 }
202 worlds[i] = World{ .value = val, .guard = guard };
203 }
204
205 var wmc_params = WmcParams.init(allocator);
206 defer wmc_params.deinit();
207 for (0..manager.numVars()) |v| {
208 try wmc_params.setWeight(@intCast(v), 0.5, 0.5);
209 }
210
211 const seq_results = try computeWmcSequential(allocator, worlds, &wmc_params, &weight_ctx, weight_root);
212 defer allocator.free(seq_results);
213
214 const par_results = try computeWmcParallel(allocator, std.testing.allocator, worlds, &wmc_params, &weight_ctx, weight_root, 2);
215 defer allocator.free(par_results);
216
217 try std.testing.expectEqual(seq_results.len, par_results.len);
218 for (seq_results, par_results) |seq, par| {
219 try std.testing.expectEqual(seq.value, par.value);
220 try std.testing.expectApproxEqAbs(seq.probability, par.probability, 1e-15);
221 }
222 }
223
224 test "parallel WMC handles exact divisibility by thread count" {
225 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
226 defer arena.deinit();
227 const allocator = arena.allocator();
228
229 var manager = try Manager.init(allocator);
230 defer manager.deinit();
231
232 var weight_ctx = try WeightDD.init(allocator, &manager);
233 defer weight_ctx.deinit();
234 const weight_root = try weight_ctx.leaf(1.0);
235
236 const num_worlds = 9;
237 const worlds = try allocator.alloc(World, num_worlds);
238
239 var vars: [4]Bdd = undefined;
240 for (0..4) |i| {
241 vars[i] = try manager.newVar(true);
242 }
243
244 for (0..num_worlds) |i| {
245 const val = try RuntimeValue.initNative(allocator, .{ .int = @intCast(i) });
246 var guard = Bdd.TRUE;
247 for (0..4) |j| {
248 const use_var = (i >> @intCast(j)) & 1 == 1;
249 if (use_var) {
250 guard = try manager.bddAnd(guard, vars[j]);
251 }
252 }
253 worlds[i] = World{ .value = val, .guard = guard };
254 }
255
256 var wmc_params = WmcParams.init(allocator);
257 defer wmc_params.deinit();
258 for (0..manager.numVars()) |v| {
259 try wmc_params.setWeight(@intCast(v), 0.5, 0.5);
260 }
261
262 const seq_results = try computeWmcSequential(allocator, worlds, &wmc_params, &weight_ctx, weight_root);
263 defer allocator.free(seq_results);
264
265 const par_results = try computeWmcParallel(allocator, std.testing.allocator, worlds, &wmc_params, &weight_ctx, weight_root, 2);
266 defer allocator.free(par_results);
267
268 try std.testing.expectEqual(seq_results.len, par_results.len);
269 for (seq_results, par_results, 0..) |seq, par, i| {
270 try std.testing.expectEqual(seq.value, par.value);
271 try std.testing.expectApproxEqAbs(seq.probability, par.probability, 1e-15);
272 try std.testing.expect(par.probability >= 0 and par.probability <= 1);
273 _ = i;
274 }
275 }