lib/accy/src/eval/evaluator.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3 const choir = @import("choir");
4
5 const ir = choir.ir;
6 const dialects = choir.dialects;
7 const GpuDialect = choir.dialects.gpu.GpuDialect;
8 const Dimension = choir.dialects.gpu.Dimension;
9 const ShuffleMode = choir.dialects.gpu.ShuffleMode;
10 const WarpOpKind = choir.dialects.gpu.WarpOpKind;
11 const Bf16 = choir_abi.Bf16;
12
13 pub const EvalError = error{
14 InvalidFunction,
15 InputCountMismatch,
16 InputSizeMismatch,
17 UnsupportedOperation,
18 UnsupportedType,
19 MissingValue,
20 MissingAttribute,
21 InvalidAttribute,
22 InvalidCondition,
23 InvalidIndex,
24 InvalidMemref,
25 InvalidResultCount,
26 MissingTerminator,
27 DivisionByZero,
28 Overflow,
29 WhileIterationLimit,
30 OutOfMemory,
31 };
32
33 pub const max_while_iterations: usize = 1 << 22;
34
35 pub const Diagnostic = choir.backends.contract.VerificationDiagnostic;
36
37 pub fn executeKernelFunction(
38 allocator: std.mem.Allocator,
39 func_op: *ir.Operation,
40 args: []const KernelArgument,
41 launch: anytype,
42 ) EvalError!void {
43 return executeKernelFunctionWithDiagnostic(allocator, func_op, args, launch, null);
44 }
45
46 pub fn executeKernelFunctionWithDiagnostic(
47 allocator: std.mem.Allocator,
48 func_op: *ir.Operation,
49 args: []const KernelArgument,
50 launch: anytype,
51 diagnostic: ?*Diagnostic,
52 ) EvalError!void {
53 var machine = Machine.init(allocator);
54 defer machine.deinit();
55 return machine.runWithDiagnostic(func_op, args, launch, diagnostic);
56 }
57
58 pub const Machine = struct {
59 allocator: std.mem.Allocator,
60 evaluator: Evaluator,
61 bound: std.ArrayListUnmanaged(RuntimeValue) = .empty,
62 lanes: std.ArrayListUnmanaged(LaneState) = .empty,
63 lane_refs: std.ArrayListUnmanaged(*LaneState) = .empty,
64
65 pub fn init(allocator: std.mem.Allocator) Machine {
66 return .{
67 .allocator = allocator,
68 .evaluator = Evaluator.init(allocator),
69 };
70 }
71
72 pub fn deinit(self: *Machine) void {
73 for (self.lanes.items) |*lane| lane.deinit();
74 self.lanes.deinit(self.allocator);
75 self.lane_refs.deinit(self.allocator);
76 self.bound.deinit(self.allocator);
77 self.evaluator.deinit();
78 self.* = undefined;
79 }
80
81 pub fn run(
82 self: *Machine,
83 func_op: *ir.Operation,
84 args: []const KernelArgument,
85 launch: anytype,
86 ) EvalError!void {
87 return self.runWithDiagnostic(func_op, args, launch, null);
88 }
89
90 pub fn runWithDiagnostic(
91 self: *Machine,
92 func_op: *ir.Operation,
93 args: []const KernelArgument,
94 launch: anytype,
95 diagnostic: ?*Diagnostic,
96 ) EvalError!void {
97 if (diagnostic) |captured| captured.len = 0;
98 return self.runChecking(func_op, args, launch, diagnostic) catch |err| {
99 captureFunctionDiagnostic(diagnostic, "accy/kernel/run", func_op, err);
100 return err;
101 };
102 }
103
104 fn runChecking(
105 self: *Machine,
106 func_op: *ir.Operation,
107 args: []const KernelArgument,
108 launch: anytype,
109 diagnostic: ?*Diagnostic,
110 ) EvalError!void {
111 if (!std.mem.eql(u8, func_op.name.name, dialects.FuncDialect.FuncOp.operation_name)) {
112 return error.InvalidFunction;
113 }
114 const region = func_op.getRegion(0) orelse return error.InvalidFunction;
115 const entry = region.getEntryBlock() orelse return error.InvalidFunction;
116 if (entry.arguments.items.len != args.len) return error.InputCountMismatch;
117
118 self.evaluator.resetRun();
119 self.evaluator.setDiagnostic("accy/kernel/run", diagnostic);
120
121 try self.bound.resize(self.allocator, args.len);
122 for (args, 0..) |arg, i| {
123 const block_arg = entry.getArgument(i) orelse return error.InvalidFunction;
124 self.bound.items[i] = switch (arg) {
125 .memref => |bytes| .{ .memref = try self.evaluator.addBorrowedMemref(block_arg.type, bytes) },
126 .scalar => |value| .{ .scalar = try castScalar(try scalarKindOfType(block_arg.type), value) },
127 };
128 }
129
130 const lane_count = try launchBlockThreadCount(launch);
131 try self.ensureLanes(lane_count);
132 const lane_refs = self.lane_refs.items[0..lane_count];
133
134 for (0..@as(usize, @intCast(launch.grid[2]))) |block_z| {
135 for (0..@as(usize, @intCast(launch.grid[1]))) |block_y| {
136 for (0..@as(usize, @intCast(launch.grid[0]))) |block_x| {
137 self.evaluator.refillSharedSentinel();
138 var lane_index: usize = 0;
139 for (0..@as(usize, @intCast(launch.block[2]))) |thread_z| {
140 for (0..@as(usize, @intCast(launch.block[1]))) |thread_y| {
141 for (0..@as(usize, @intCast(launch.block[0]))) |thread_x| {
142 const lane = lane_refs[lane_index];
143 lane_index += 1;
144 lane.reset(launchIndex(launch, block_x, block_y, block_z, thread_x, thread_y, thread_z));
145 for (self.bound.items, 0..) |value, i| {
146 const block_arg = entry.getArgument(i) orelse return error.InvalidFunction;
147 try lane.setValue(block_arg, value);
148 }
149 }
150 }
151 }
152 std.debug.assert(lane_index == lane_count);
153 try self.evaluator.evalKernelBlock(entry, lane_refs);
154 }
155 }
156 }
157 }
158
159 fn ensureLanes(self: *Machine, lane_count: usize) EvalError!void {
160 while (self.lanes.items.len < lane_count) {
161 try self.lanes.append(self.allocator, LaneState.init(self.allocator, .{}));
162 }
163 try self.lane_refs.resize(self.allocator, self.lanes.items.len);
164 for (self.lanes.items, 0..) |*lane, index| {
165 self.lane_refs.items[index] = lane;
166 }
167 std.debug.assert(self.lane_refs.items.len >= lane_count);
168 }
169 };
170
171 fn captureFunctionDiagnostic(diagnostic: ?*Diagnostic, stage: []const u8, func_op: *ir.Operation, err: EvalError) void {
172 const captured = diagnostic orelse return;
173 if (captured.hasText()) return;
174 if (functionName(func_op)) |name| {
175 setFormattedDiagnostic(captured, "{s}: {s} in {s}", .{ stage, @errorName(err), name });
176 } else {
177 setFormattedDiagnostic(captured, "{s}: {s}", .{ stage, @errorName(err) });
178 }
179 }
180
181 fn captureOpDiagnostic(diagnostic: ?*Diagnostic, stage: []const u8, op: *ir.Operation, err: EvalError) void {
182 const captured = diagnostic orelse return;
183 if (captured.hasText()) return;
184 setFormattedDiagnostic(captured, "{s}: {s} at {s}", .{ stage, @errorName(err), op.name.name });
185 }
186
187 fn setFormattedDiagnostic(diagnostic: *Diagnostic, comptime fmt: []const u8, args: anytype) void {
188 const message = std.fmt.bufPrint(&diagnostic.buffer, fmt, args) catch {
189 diagnostic.set("accy/kernel/run: diagnostic truncated");
190 return;
191 };
192 diagnostic.len = message.len;
193 }
194
195 fn functionName(func_op: *ir.Operation) ?[]const u8 {
196 const func = dialects.FuncDialect.FuncOp{ .op = func_op };
197 return func.getName();
198 }
199
200 pub const ScalarKind = enum {
201 bool,
202 index,
203 i8,
204 i16,
205 i32,
206 u8,
207 u16,
208 u32,
209 i64,
210 u64,
211 f16,
212 bf16,
213 f32,
214 f64,
215 };
216
217 pub const Scalar = union(ScalarKind) {
218 bool: bool,
219 index: i64,
220 i8: i8,
221 i16: i16,
222 i32: i32,
223 u8: u8,
224 u16: u16,
225 u32: u32,
226 i64: i64,
227 u64: u64,
228 f16: f16,
229 bf16: Bf16,
230 f32: f32,
231 f64: f64,
232 };
233
234 pub const KernelArgument = union(enum) {
235 memref: []u8,
236 scalar: Scalar,
237 };
238
239 const LaunchIndex = struct {
240 global: [3]i64 = .{ 0, 0, 0 },
241 thread: [3]i64 = .{ 0, 0, 0 },
242 block: [3]i64 = .{ 0, 0, 0 },
243 grid_dim: [3]i64 = .{ 1, 1, 1 },
244 block_dim: [3]i64 = .{ 1, 1, 1 },
245 };
246
247 const LaneState = struct {
248 values: std.AutoHashMap(*ir.Value, RuntimeValue),
249 launch: LaunchIndex,
250
251 fn init(allocator: std.mem.Allocator, index: LaunchIndex) LaneState {
252 return .{
253 .values = std.AutoHashMap(*ir.Value, RuntimeValue).init(allocator),
254 .launch = index,
255 };
256 }
257
258 fn deinit(self: *LaneState) void {
259 self.values.deinit();
260 }
261
262 fn reset(self: *LaneState, index: LaunchIndex) void {
263 self.values.clearRetainingCapacity();
264 self.launch = index;
265 }
266
267 fn setValue(self: *LaneState, value: *ir.Value, runtime_value: RuntimeValue) EvalError!void {
268 try self.values.put(value, runtime_value);
269 }
270 };
271
272 const warp_size: usize = 32;
273
274 pub const shared_alloc_sentinel: u8 = 0x7F;
275
276 fn launchBlockThreadCount(launch: anytype) EvalError!usize {
277 const xy = std.math.mul(u32, launch.block[0], launch.block[1]) catch return error.Overflow;
278 const xyz = std.math.mul(u32, xy, launch.block[2]) catch return error.Overflow;
279 return std.math.cast(usize, xyz) orelse return error.Overflow;
280 }
281
282 fn launchIndex(
283 launch: anytype,
284 block_x: usize,
285 block_y: usize,
286 block_z: usize,
287 thread_x: usize,
288 thread_y: usize,
289 thread_z: usize,
290 ) LaunchIndex {
291 const block_i: [3]i64 = .{ @intCast(block_x), @intCast(block_y), @intCast(block_z) };
292 const thread_i: [3]i64 = .{ @intCast(thread_x), @intCast(thread_y), @intCast(thread_z) };
293 const block_dim_i: [3]i64 = .{ @intCast(launch.block[0]), @intCast(launch.block[1]), @intCast(launch.block[2]) };
294 return .{
295 .block = block_i,
296 .thread = thread_i,
297 .block_dim = block_dim_i,
298 .grid_dim = .{ @intCast(launch.grid[0]), @intCast(launch.grid[1]), @intCast(launch.grid[2]) },
299 .global = .{
300 block_i[0] * block_dim_i[0] + thread_i[0],
301 block_i[1] * block_dim_i[1] + thread_i[1],
302 block_i[2] * block_dim_i[2] + thread_i[2],
303 },
304 };
305 }
306
307 fn linearThreadIndex(index: LaunchIndex) EvalError!usize {
308 const x = std.math.cast(usize, index.thread[0]) orelse return error.Overflow;
309 const y = std.math.cast(usize, index.thread[1]) orelse return error.Overflow;
310 const z = std.math.cast(usize, index.thread[2]) orelse return error.Overflow;
311 const dim_x = std.math.cast(usize, index.block_dim[0]) orelse return error.Overflow;
312 const dim_y = std.math.cast(usize, index.block_dim[1]) orelse return error.Overflow;
313 const y_base = std.math.mul(usize, y, dim_x) catch return error.Overflow;
314 const plane = std.math.mul(usize, dim_x, dim_y) catch return error.Overflow;
315 const z_base = std.math.mul(usize, z, plane) catch return error.Overflow;
316 return std.math.add(usize, std.math.add(usize, x, y_base) catch return error.Overflow, z_base) catch return error.Overflow;
317 }
318
319 fn blockActiveMask(index: LaunchIndex) EvalError!u32 {
320 const current_warp = try warpId(index);
321 const dim_x = std.math.cast(usize, index.block_dim[0]) orelse return error.Overflow;
322 const dim_y = std.math.cast(usize, index.block_dim[1]) orelse return error.Overflow;
323 const dim_z = std.math.cast(usize, index.block_dim[2]) orelse return error.Overflow;
324 const xy = std.math.mul(usize, dim_x, dim_y) catch return error.Overflow;
325 const count = std.math.mul(usize, xy, dim_z) catch return error.Overflow;
326 var mask: u32 = 0;
327 var local: usize = 0;
328 while (local < count) : (local += 1) {
329 if (local / warp_size != current_warp) continue;
330 mask |= @as(u32, 1) << @intCast(local % warp_size);
331 }
332 return mask;
333 }
334
335 fn laneId(index: LaunchIndex) EvalError!usize {
336 return (try linearThreadIndex(index)) % warp_size;
337 }
338
339 fn warpId(index: LaunchIndex) EvalError!usize {
340 return (try linearThreadIndex(index)) / warp_size;
341 }
342
343 fn activeMaskForLanes(index: LaunchIndex, lanes: []*LaneState) EvalError!u32 {
344 const current_warp = try warpId(index);
345 var mask: u32 = 0;
346 for (lanes) |lane| {
347 if (try warpId(lane.launch) != current_warp) continue;
348 mask |= @as(u32, 1) << @intCast(try laneId(lane.launch));
349 }
350 return mask;
351 }
352
353 const RuntimeValue = union(enum) {
354 scalar: Scalar,
355 f32x4: [4]f32,
356 memref: usize,
357 };
358
359 const Memref = struct {
360 element: ScalarKind,
361 bytes: []u8,
362 owned: bool = true,
363 };
364
365 const LoopLane = struct {
366 lane: *LaneState,
367 iv: i64,
368 upper: i64,
369 step: i64,
370 };
371
372 const ControlScratch = struct {
373 frames: std.ArrayListUnmanaged(*Frame) = .empty,
374 depth: usize = 0,
375
376 const Frame = struct {
377 lanes_a: std.ArrayListUnmanaged(*LaneState) = .empty,
378 lanes_b: std.ArrayListUnmanaged(*LaneState) = .empty,
379 indices: std.ArrayListUnmanaged(usize) = .empty,
380 loops: std.ArrayListUnmanaged(LoopLane) = .empty,
381 carried: std.ArrayListUnmanaged(RuntimeValue) = .empty,
382
383 fn reset(self: *Frame) void {
384 self.lanes_a.clearRetainingCapacity();
385 self.lanes_b.clearRetainingCapacity();
386 self.indices.clearRetainingCapacity();
387 self.loops.clearRetainingCapacity();
388 self.carried.clearRetainingCapacity();
389 }
390
391 fn deinit(self: *Frame, allocator: std.mem.Allocator) void {
392 self.lanes_a.deinit(allocator);
393 self.lanes_b.deinit(allocator);
394 self.indices.deinit(allocator);
395 self.loops.deinit(allocator);
396 self.carried.deinit(allocator);
397 }
398 };
399
400 fn acquire(self: *ControlScratch, allocator: std.mem.Allocator) EvalError!*Frame {
401 std.debug.assert(self.depth <= self.frames.items.len);
402 if (self.depth == self.frames.items.len) {
403 const frame = try allocator.create(Frame);
404 errdefer allocator.destroy(frame);
405 frame.* = .{};
406 try self.frames.append(allocator, frame);
407 }
408 const frame = self.frames.items[self.depth];
409 self.depth += 1;
410 frame.reset();
411 return frame;
412 }
413
414 fn release(self: *ControlScratch) void {
415 std.debug.assert(self.depth > 0);
416 self.depth -= 1;
417 }
418
419 fn deinit(self: *ControlScratch, allocator: std.mem.Allocator) void {
420 for (self.frames.items) |frame| {
421 frame.deinit(allocator);
422 allocator.destroy(frame);
423 }
424 self.frames.deinit(allocator);
425 }
426 };
427
428 const Evaluator = struct {
429 allocator: std.mem.Allocator,
430 values: std.AutoHashMap(*ir.Value, RuntimeValue),
431 memrefs: std.ArrayList(Memref),
432 shared_allocs: std.AutoHashMap(*ir.Operation, usize),
433 shared_bytes: std.AutoHashMap(*ir.Operation, []u8),
434 scratch: ControlScratch = .{},
435 warps: std.AutoHashMap(usize, void),
436 launch: LaunchIndex = .{},
437 active_lanes: ?[]*LaneState = null,
438 diagnostic_stage: []const u8 = "accy/eval",
439 diagnostic: ?*Diagnostic = null,
440
441 const BinaryOp = enum { add, sub, mul, div, rem, min, max };
442
443 const UnaryOp = enum { neg, abs, sqrt, exp, log, tanh, sin, cos, tan, floor, round, trunc, tf32_round };
444
445 const LogicalOp = enum { and_, or_, xor };
446
447 const ShiftOp = enum { shl, shr, ushr };
448
449 const GpuIdKind = enum { global, thread, block, block_dim, grid_dim };
450
451 const WarpVoteKind = enum { all, any, ballot };
452
453 fn init(allocator: std.mem.Allocator) Evaluator {
454 return .{
455 .allocator = allocator,
456 .values = std.AutoHashMap(*ir.Value, RuntimeValue).init(allocator),
457 .memrefs = .empty,
458 .shared_allocs = std.AutoHashMap(*ir.Operation, usize).init(allocator),
459 .shared_bytes = std.AutoHashMap(*ir.Operation, []u8).init(allocator),
460 .warps = std.AutoHashMap(usize, void).init(allocator),
461 };
462 }
463
464 fn deinit(self: *Evaluator) void {
465 self.releaseOwnedMemrefs();
466 self.memrefs.deinit(self.allocator);
467 self.shared_allocs.deinit();
468 var shared_iter = self.shared_bytes.valueIterator();
469 while (shared_iter.next()) |bytes| {
470 self.allocator.free(bytes.*);
471 }
472 self.shared_bytes.deinit();
473 self.warps.deinit();
474 self.scratch.deinit(self.allocator);
475 self.values.deinit();
476 }
477
478 fn resetRun(self: *Evaluator) void {
479 std.debug.assert(self.scratch.depth == 0);
480 self.releaseOwnedMemrefs();
481 self.memrefs.clearRetainingCapacity();
482 self.shared_allocs.clearRetainingCapacity();
483 self.values.clearRetainingCapacity();
484 self.launch = .{};
485 self.active_lanes = null;
486 }
487
488 fn releaseOwnedMemrefs(self: *Evaluator) void {
489 for (self.memrefs.items) |memref| {
490 if (memref.owned) self.allocator.free(memref.bytes);
491 }
492 }
493
494 fn setDiagnostic(self: *Evaluator, stage: []const u8, diagnostic: ?*Diagnostic) void {
495 self.diagnostic_stage = stage;
496 self.diagnostic = diagnostic;
497 }
498
499 fn failOp(self: *Evaluator, op: *ir.Operation, err: EvalError) EvalError {
500 captureOpDiagnostic(self.diagnostic, self.diagnostic_stage, op, err);
501 return err;
502 }
503
504 fn refillSharedSentinel(self: *Evaluator) void {
505 var iter = self.shared_allocs.valueIterator();
506 while (iter.next()) |memref_index| {
507 std.debug.assert(memref_index.* < self.memrefs.items.len);
508 const memref = &self.memrefs.items[memref_index.*];
509 @memset(memref.bytes, shared_alloc_sentinel);
510 }
511 }
512
513 fn addBorrowedMemref(self: *Evaluator, typ: ir.Type, bytes: []u8) EvalError!usize {
514 const desc = try describeMemrefType(typ);
515 if (desc.byte_size) |byte_size| {
516 if (bytes.len < byte_size) return error.InputSizeMismatch;
517 }
518 return self.appendMemref(.{ .element = desc.element, .bytes = bytes, .owned = false });
519 }
520
521 fn appendMemref(self: *Evaluator, memref: Memref) EvalError!usize {
522 errdefer if (memref.owned) self.allocator.free(memref.bytes);
523 try self.memrefs.append(self.allocator, memref);
524 return self.memrefs.items.len - 1;
525 }
526
527 fn setValue(self: *Evaluator, value: *ir.Value, runtime_value: RuntimeValue) EvalError!void {
528 try self.values.put(value, runtime_value);
529 }
530
531 fn resolve(self: *Evaluator, value: *ir.Value) EvalError!RuntimeValue {
532 return self.values.get(value) orelse error.MissingValue;
533 }
534
535 fn evalBlock(self: *Evaluator, block: *ir.Block) EvalError![]RuntimeValue {
536 var iter = block.operations.head;
537 while (iter) |op_ptr| {
538 const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
539 if (std.mem.eql(u8, op.name.name, dialects.FuncDialect.ReturnOp.operation_name) or
540 std.mem.eql(u8, op.name.name, dialects.ScfDialect.YieldOp.operation_name))
541 {
542 return self.collectOperands(op) catch |err| return self.failOp(op, err);
543 }
544
545 try self.evalOpWithDiagnostic(op);
546 iter = op.next_op;
547 }
548 return error.MissingTerminator;
549 }
550
551 fn evalKernelBlock(self: *Evaluator, block: *ir.Block, lanes: []*LaneState) EvalError!void {
552 var iter = block.operations.head;
553 while (iter) |op_ptr| {
554 const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
555 if (std.mem.eql(u8, op.name.name, dialects.FuncDialect.ReturnOp.operation_name)) {
556 if (op.operands.items.len != 0) return self.failOp(op, error.UnsupportedOperation);
557 return;
558 }
559 if (std.mem.eql(u8, op.name.name, dialects.ScfDialect.YieldOp.operation_name)) {
560 return self.failOp(op, error.UnsupportedOperation);
561 }
562
563 const handled_collective = self.evalKernelCollective(op, lanes) catch |err| return self.failOp(op, err);
564 if (!handled_collective) {
565 if (std.mem.eql(u8, op.name.name, dialects.ScfDialect.IfOp.operation_name)) {
566 self.evalKernelIf(op, lanes) catch |err| return self.failOp(op, err);
567 } else if (std.mem.eql(u8, op.name.name, dialects.ScfDialect.ForOp.operation_name)) {
568 self.evalKernelFor(op, lanes) catch |err| return self.failOp(op, err);
569 } else if (std.mem.eql(u8, op.name.name, dialects.ScfDialect.WhileOp.operation_name)) {
570 self.evalKernelWhile(op, lanes) catch |err| return self.failOp(op, err);
571 } else {
572 for (lanes) |lane| self.evalLaneOp(lane, op, lanes) catch |err| return self.failOp(op, err);
573 }
574 }
575 iter = op.next_op;
576 }
577 return error.MissingTerminator;
578 }
579
580 fn evalKernelYieldBlock(self: *Evaluator, block: *ir.Block, lanes: []*LaneState, parent: *ir.Operation) EvalError!void {
581 var iter = block.operations.head;
582 while (iter) |op_ptr| {
583 const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
584 if (std.mem.eql(u8, op.name.name, dialects.ScfDialect.YieldOp.operation_name)) {
585 if (op.operands.items.len != parent.getNumResults()) return self.failOp(op, error.InvalidResultCount);
586 for (lanes) |lane| {
587 for (op.operands.items, 0..) |operand, result_index| {
588 const result = parent.getResult(result_index) orelse return error.InvalidResultCount;
589 lane.setValue(result, try self.resolveLane(lane, operand.value)) catch |err| return self.failOp(op, err);
590 }
591 }
592 return;
593 }
594 if (std.mem.eql(u8, op.name.name, dialects.FuncDialect.ReturnOp.operation_name)) {
595 return self.failOp(op, error.UnsupportedOperation);
596 }
597
598 const handled_collective = self.evalKernelCollective(op, lanes) catch |err| return self.failOp(op, err);
599 if (!handled_collective) {
600 if (std.mem.eql(u8, op.name.name, dialects.ScfDialect.IfOp.operation_name)) {
601 self.evalKernelIf(op, lanes) catch |err| return self.failOp(op, err);
602 } else if (std.mem.eql(u8, op.name.name, dialects.ScfDialect.ForOp.operation_name)) {
603 self.evalKernelFor(op, lanes) catch |err| return self.failOp(op, err);
604 } else if (std.mem.eql(u8, op.name.name, dialects.ScfDialect.WhileOp.operation_name)) {
605 self.evalKernelWhile(op, lanes) catch |err| return self.failOp(op, err);
606 } else {
607 for (lanes) |lane| self.evalLaneOp(lane, op, lanes) catch |err| return self.failOp(op, err);
608 }
609 }
610 iter = op.next_op;
611 }
612 return error.MissingTerminator;
613 }
614
615 fn evalKernelWhile(self: *Evaluator, op: *ir.Operation, lanes: []*LaneState) EvalError!void {
616 const while_op = dialects.ScfDialect.WhileOp{ .op = op };
617 const carry_count = op.operands.items.len;
618 if (op.getNumResults() != carry_count) return error.InvalidResultCount;
619 const before = while_op.getBeforeBlock();
620 const after = while_op.getAfterBlock();
621 if (before.arguments.items.len != carry_count) return error.UnsupportedOperation;
622 if (after.arguments.items.len != carry_count) return error.UnsupportedOperation;
623
624 for (lanes) |lane| {
625 for (0..carry_count) |i| {
626 const carried = try self.resolveLane(lane, op.operands.items[i].value);
627 try lane.setValue(before.getArgument(i).?, carried);
628 }
629 }
630
631 const frame = try self.scratch.acquire(self.allocator);
632 defer self.scratch.release();
633 const active = &frame.lanes_a;
634 const continuing = &frame.lanes_b;
635 try active.ensureTotalCapacity(self.allocator, lanes.len);
636 try continuing.ensureTotalCapacity(self.allocator, lanes.len);
637 active.appendSliceAssumeCapacity(lanes);
638
639 var iterations: usize = 0;
640 while (active.items.len != 0) {
641 if (iterations >= max_while_iterations) return error.WhileIterationLimit;
642 iterations += 1;
643
644 const terminator = try self.evalKernelBlockUntil(before, active.items, dialects.ScfDialect.ConditionOp.operation_name);
645 const condition_op = dialects.ScfDialect.ConditionOp{ .op = terminator };
646 const cond_value = condition_op.getCondition();
647 const args = condition_op.getArgs();
648 if (args.len != carry_count) return error.InvalidResultCount;
649
650 continuing.clearRetainingCapacity();
651 for (active.items) |lane| {
652 if (try self.boolOfLane(lane, cond_value)) {
653 for (args, 0..) |arg, i| {
654 try lane.setValue(after.getArgument(i).?, try self.resolveLane(lane, arg));
655 }
656 try continuing.append(self.allocator, lane);
657 } else {
658 for (args, 0..) |arg, i| {
659 const result = op.getResult(i) orelse return error.InvalidResultCount;
660 try lane.setValue(result, try self.resolveLane(lane, arg));
661 }
662 }
663 }
664 active.clearRetainingCapacity();
665 try active.appendSlice(self.allocator, continuing.items);
666 if (active.items.len == 0) break;
667
668 const yield_op = try self.evalKernelBlockUntil(after, active.items, dialects.ScfDialect.YieldOp.operation_name);
669 if (yield_op.operands.items.len != carry_count) return error.InvalidResultCount;
670 for (active.items) |lane| {
671 for (yield_op.operands.items, 0..) |operand, i| {
672 try lane.setValue(before.getArgument(i).?, try self.resolveLane(lane, operand.value));
673 }
674 }
675 }
676 }
677
678 fn evalKernelBlockUntil(
679 self: *Evaluator,
680 block: *ir.Block,
681 lanes: []*LaneState,
682 terminator_name: []const u8,
683 ) EvalError!*ir.Operation {
684 var iter = block.operations.head;
685 while (iter) |op_ptr| {
686 const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
687 if (std.mem.eql(u8, op.name.name, terminator_name)) return op;
688 if (std.mem.eql(u8, op.name.name, dialects.FuncDialect.ReturnOp.operation_name)) {
689 return self.failOp(op, error.UnsupportedOperation);
690 }
691
692 const handled_collective = self.evalKernelCollective(op, lanes) catch |err| return self.failOp(op, err);
693 if (!handled_collective) {
694 if (std.mem.eql(u8, op.name.name, dialects.ScfDialect.IfOp.operation_name)) {
695 self.evalKernelIf(op, lanes) catch |err| return self.failOp(op, err);
696 } else if (std.mem.eql(u8, op.name.name, dialects.ScfDialect.ForOp.operation_name)) {
697 self.evalKernelFor(op, lanes) catch |err| return self.failOp(op, err);
698 } else if (std.mem.eql(u8, op.name.name, dialects.ScfDialect.WhileOp.operation_name)) {
699 self.evalKernelWhile(op, lanes) catch |err| return self.failOp(op, err);
700 } else {
701 for (lanes) |lane| self.evalLaneOp(lane, op, lanes) catch |err| return self.failOp(op, err);
702 }
703 }
704 iter = op.next_op;
705 }
706 return error.MissingTerminator;
707 }
708
709 fn evalKernelIf(self: *Evaluator, op: *ir.Operation, lanes: []*LaneState) EvalError!void {
710 if (op.operands.items.len != 1) return error.UnsupportedOperation;
711 const frame = try self.scratch.acquire(self.allocator);
712 defer self.scratch.release();
713 const then_lanes = &frame.lanes_a;
714 const else_lanes = &frame.lanes_b;
715 try then_lanes.ensureTotalCapacity(self.allocator, lanes.len);
716 try else_lanes.ensureTotalCapacity(self.allocator, lanes.len);
717
718 for (lanes) |lane| {
719 if (try self.boolOfLane(lane, op.operands.items[0].value)) {
720 then_lanes.appendAssumeCapacity(lane);
721 } else {
722 else_lanes.appendAssumeCapacity(lane);
723 }
724 }
725
726 if (then_lanes.items.len != 0) {
727 const region = op.getRegion(0) orelse return error.UnsupportedOperation;
728 const block = region.getEntryBlock() orelse return error.UnsupportedOperation;
729 try self.evalKernelYieldBlock(block, then_lanes.items, op);
730 }
731 if (else_lanes.items.len != 0) {
732 const region = op.getRegion(1) orelse {
733 if (op.getNumResults() == 0) return;
734 return error.UnsupportedOperation;
735 };
736 const block = region.getEntryBlock() orelse return error.UnsupportedOperation;
737 try self.evalKernelYieldBlock(block, else_lanes.items, op);
738 }
739 }
740
741 fn evalKernelFor(self: *Evaluator, op: *ir.Operation, lanes: []*LaneState) EvalError!void {
742 if (op.operands.items.len < 3) return error.UnsupportedOperation;
743 const iter_count = op.operands.items.len - 3;
744 if (op.getNumResults() != iter_count) return error.InvalidResultCount;
745 const region = op.getRegion(0) orelse return error.UnsupportedOperation;
746 const body = region.getEntryBlock() orelse return error.UnsupportedOperation;
747 if (body.arguments.items.len != iter_count + 1) return error.UnsupportedOperation;
748
749 const frame = try self.scratch.acquire(self.allocator);
750 defer self.scratch.release();
751 try frame.loops.resize(self.allocator, lanes.len);
752 const loop_lanes = frame.loops.items;
753
754 const iter_value_count = std.math.mul(usize, lanes.len, iter_count) catch return error.Overflow;
755 try frame.carried.resize(self.allocator, iter_value_count);
756 const iter_values = frame.carried.items;
757
758 for (lanes, 0..) |lane, lane_index| {
759 const lower = try self.indexOfLane(lane, op.operands.items[0].value);
760 const upper = try self.indexOfLane(lane, op.operands.items[1].value);
761 const step = try self.indexOfLane(lane, op.operands.items[2].value);
762 if (step <= 0) return error.UnsupportedOperation;
763 loop_lanes[lane_index] = .{
764 .lane = lane,
765 .iv = lower,
766 .upper = upper,
767 .step = step,
768 };
769 for (0..iter_count) |iter_index| {
770 iter_values[lane_index * iter_count + iter_index] = try self.resolveLane(lane, op.operands.items[3 + iter_index].value);
771 }
772 }
773
774 const active_lanes = &frame.lanes_a;
775 const active_indices = &frame.indices;
776 try active_lanes.ensureTotalCapacity(self.allocator, lanes.len);
777 try active_indices.ensureTotalCapacity(self.allocator, lanes.len);
778
779 while (true) {
780 active_lanes.clearRetainingCapacity();
781 active_indices.clearRetainingCapacity();
782
783 for (loop_lanes, 0..) |loop_lane, lane_index| {
784 if (loop_lane.iv >= loop_lane.upper) continue;
785 const lane = loop_lane.lane;
786 try lane.setValue(body.getArgument(0) orelse return error.UnsupportedOperation, .{ .scalar = .{ .index = loop_lane.iv } });
787 for (0..iter_count) |iter_index| {
788 try lane.setValue(body.getArgument(iter_index + 1) orelse return error.UnsupportedOperation, iter_values[lane_index * iter_count + iter_index]);
789 }
790 active_lanes.appendAssumeCapacity(lane);
791 active_indices.appendAssumeCapacity(lane_index);
792 }
793
794 if (active_lanes.items.len == 0) break;
795
796 try self.evalKernelYieldBlock(body, active_lanes.items, op);
797
798 for (active_indices.items) |lane_index| {
799 const lane = loop_lanes[lane_index].lane;
800 for (0..iter_count) |iter_index| {
801 const result = op.getResult(iter_index) orelse return error.InvalidResultCount;
802 iter_values[lane_index * iter_count + iter_index] = try self.resolveLane(lane, result);
803 }
804 loop_lanes[lane_index].iv = std.math.add(i64, loop_lanes[lane_index].iv, loop_lanes[lane_index].step) catch return error.Overflow;
805 }
806 }
807
808 for (lanes, 0..) |lane, lane_index| {
809 for (0..iter_count) |iter_index| {
810 const result = op.getResult(iter_index) orelse return error.InvalidResultCount;
811 try lane.setValue(result, iter_values[lane_index * iter_count + iter_index]);
812 }
813 }
814 }
815
816 fn evalLaneOp(self: *Evaluator, lane: *LaneState, op: *ir.Operation, active_lanes: []*LaneState) EvalError!void {
817 const previous_values = self.values;
818 const previous_launch = self.launch;
819 const previous_active_lanes = self.active_lanes;
820 self.values = lane.values;
821 self.launch = lane.launch;
822 self.active_lanes = active_lanes;
823 defer {
824 lane.values = self.values;
825 self.values = previous_values;
826 self.launch = previous_launch;
827 self.active_lanes = previous_active_lanes;
828 }
829 try self.evalOpWithDiagnostic(op);
830 }
831
832 fn collectOperands(self: *Evaluator, op: *ir.Operation) EvalError![]RuntimeValue {
833 const values = try self.allocator.alloc(RuntimeValue, op.operands.items.len);
834 errdefer self.allocator.free(values);
835 for (op.operands.items, 0..) |operand, i| {
836 values[i] = try self.resolve(operand.value);
837 }
838 return values;
839 }
840
841 fn evalOp(self: *Evaluator, op: *ir.Operation) EvalError!void {
842 const name = op.name.name;
843 if (std.mem.eql(u8, name, dialects.ArithDialect.ConstantOp.operation_name)) return self.evalConstant(op);
844 if (std.mem.eql(u8, name, dialects.ArithDialect.AddOp.operation_name)) return self.evalBinary(op, .add);
845 if (std.mem.eql(u8, name, dialects.ArithDialect.SubOp.operation_name)) return self.evalBinary(op, .sub);
846 if (std.mem.eql(u8, name, dialects.ArithDialect.MulOp.operation_name)) return self.evalBinary(op, .mul);
847 if (std.mem.eql(u8, name, dialects.ArithDialect.UmulhiOp.operation_name)) return self.evalUmulhi(op);
848 if (std.mem.eql(u8, name, dialects.ArithDialect.DivOp.operation_name)) return self.evalBinary(op, .div);
849 if (std.mem.eql(u8, name, dialects.ArithDialect.RemOp.operation_name)) return self.evalBinary(op, .rem);
850 if (std.mem.eql(u8, name, dialects.ArithDialect.MinOp.operation_name)) return self.evalBinary(op, .min);
851 if (std.mem.eql(u8, name, dialects.ArithDialect.MaxOp.operation_name)) return self.evalBinary(op, .max);
852 if (std.mem.eql(u8, name, dialects.ArithDialect.NegOp.operation_name)) return self.evalUnary(op, .neg);
853 if (std.mem.eql(u8, name, dialects.ArithDialect.AbsOp.operation_name)) return self.evalUnary(op, .abs);
854 if (std.mem.eql(u8, name, dialects.ArithDialect.SqrtOp.operation_name)) return self.evalUnary(op, .sqrt);
855 if (std.mem.eql(u8, name, dialects.ArithDialect.ExpOp.operation_name)) return self.evalUnary(op, .exp);
856 if (std.mem.eql(u8, name, dialects.ArithDialect.LogOp.operation_name)) return self.evalUnary(op, .log);
857 if (std.mem.eql(u8, name, dialects.ArithDialect.TanhOp.operation_name)) return self.evalUnary(op, .tanh);
858 if (std.mem.eql(u8, name, dialects.ArithDialect.SinOp.operation_name)) return self.evalUnary(op, .sin);
859 if (std.mem.eql(u8, name, dialects.ArithDialect.CosOp.operation_name)) return self.evalUnary(op, .cos);
860 if (std.mem.eql(u8, name, dialects.ArithDialect.TanOp.operation_name)) return self.evalUnary(op, .tan);
861 if (std.mem.eql(u8, name, dialects.ArithDialect.FloorOp.operation_name)) return self.evalUnary(op, .floor);
862 if (std.mem.eql(u8, name, dialects.ArithDialect.RoundOp.operation_name)) return self.evalUnary(op, .round);
863 if (std.mem.eql(u8, name, dialects.ArithDialect.TruncOp.operation_name)) return self.evalUnary(op, .trunc);
864 if (std.mem.eql(u8, name, dialects.ArithDialect.Tf32RoundOp.operation_name)) return self.evalUnary(op, .tf32_round);
865 if (std.mem.eql(u8, name, dialects.ArithDialect.PowOp.operation_name)) return self.evalPow(op);
866 if (std.mem.eql(u8, name, dialects.ArithDialect.Atan2Op.operation_name)) return self.evalAtan2(op);
867 if (std.mem.eql(u8, name, dialects.ArithDialect.FmaOp.operation_name)) return self.evalFma(op);
868 if (std.mem.eql(u8, name, dialects.ArithDialect.CmpOp.operation_name)) return self.evalCmp(op);
869 if (std.mem.eql(u8, name, dialects.ArithDialect.CastOp.operation_name)) return self.evalCast(op);
870 if (std.mem.eql(u8, name, dialects.ArithDialect.SelectOp.operation_name)) return self.evalSelect(op);
871 if (std.mem.eql(u8, name, dialects.ArithDialect.AndOp.operation_name)) return self.evalLogical(op, .and_);
872 if (std.mem.eql(u8, name, dialects.ArithDialect.OrOp.operation_name)) return self.evalLogical(op, .or_);
873 if (std.mem.eql(u8, name, dialects.ArithDialect.XorOp.operation_name)) return self.evalLogical(op, .xor);
874 if (std.mem.eql(u8, name, dialects.ArithDialect.NotOp.operation_name)) return self.evalNot(op);
875 if (std.mem.eql(u8, name, dialects.ArithDialect.PopCountOp.operation_name)) return self.evalPopCount(op);
876 if (std.mem.eql(u8, name, dialects.ArithDialect.ShlOp.operation_name)) return self.evalShift(op, .shl);
877 if (std.mem.eql(u8, name, dialects.ArithDialect.ShrOp.operation_name)) return self.evalShift(op, .shr);
878 if (std.mem.eql(u8, name, dialects.ArithDialect.UshrOp.operation_name)) return self.evalShift(op, .ushr);
879 if (std.mem.eql(u8, name, dialects.ArithDialect.BitcastOp.operation_name)) return self.evalBitcast(op);
880 if (std.mem.eql(u8, name, dialects.MemrefDialect.AllocaOp.operation_name)) return self.evalAlloca(op);
881 if (std.mem.eql(u8, name, dialects.MemrefDialect.AllocOp.operation_name)) return self.evalAlloca(op);
882 if (std.mem.eql(u8, name, dialects.ArithDialect.ExtractOp.operation_name)) return self.evalVecExtract(op);
883 if (std.mem.eql(u8, name, dialects.ArithDialect.SplatOp.operation_name)) return self.evalVecSplat(op);
884 if (std.mem.eql(u8, name, dialects.ArithDialect.InsertOp.operation_name)) return self.evalVecInsert(op);
885 if (std.mem.eql(u8, name, dialects.MemrefDialect.LoadOp.operation_name)) return self.evalLoad(op);
886 if (std.mem.eql(u8, name, dialects.MemrefDialect.StoreOp.operation_name)) return self.evalStore(op);
887 if (std.mem.eql(u8, name, dialects.MemrefDialect.AtomicRmwOp.operation_name)) return self.evalAtomicRmw(op);
888 if (std.mem.eql(u8, name, dialects.MemrefDialect.AtomicCasOp.operation_name)) return self.evalAtomicCas(op);
889 if (std.mem.eql(u8, name, dialects.ScfDialect.ForOp.operation_name)) return self.evalFor(op);
890 if (std.mem.eql(u8, name, dialects.ScfDialect.WhileOp.operation_name)) return self.evalWhile(op);
891 if (std.mem.eql(u8, name, dialects.ScfDialect.IfOp.operation_name)) return self.evalIf(op);
892 if (std.mem.eql(u8, name, GpuDialect.GlobalIdxOp.operation_name)) return self.evalGpuId(op, .global);
893 if (std.mem.eql(u8, name, GpuDialect.ThreadIdxOp.operation_name)) return self.evalGpuId(op, .thread);
894 if (std.mem.eql(u8, name, GpuDialect.BlockIdxOp.operation_name)) return self.evalGpuId(op, .block);
895 if (std.mem.eql(u8, name, GpuDialect.BlockDimOp.operation_name)) return self.evalGpuId(op, .block_dim);
896 if (std.mem.eql(u8, name, GpuDialect.GridDimOp.operation_name)) return self.evalGpuId(op, .grid_dim);
897 if (std.mem.eql(u8, name, GpuDialect.LaneIdOp.operation_name)) return self.evalLaneId(op);
898 if (std.mem.eql(u8, name, GpuDialect.WarpIdOp.operation_name)) return self.evalWarpId(op);
899 if (std.mem.eql(u8, name, GpuDialect.ActiveMaskOp.operation_name)) return self.evalActiveMask(op);
900 if (std.mem.eql(u8, name, GpuDialect.SyncWarpOp.operation_name)) return;
901 if (std.mem.eql(u8, name, GpuDialect.WarpReduceOp.operation_name)) return error.UnsupportedOperation;
902 if (std.mem.eql(u8, name, GpuDialect.BarrierOp.operation_name)) return;
903 if (std.mem.eql(u8, name, GpuDialect.FenceOp.operation_name)) return;
904 if (std.mem.eql(u8, name, GpuDialect.CpAsyncSharedOp.operation_name)) return self.evalCpAsyncShared(op);
905 if (std.mem.eql(u8, name, GpuDialect.CpAsyncCommitOp.operation_name)) return;
906 if (std.mem.eql(u8, name, GpuDialect.CpAsyncWaitOp.operation_name)) return;
907 return error.UnsupportedOperation;
908 }
909
910 fn evalOpWithDiagnostic(self: *Evaluator, op: *ir.Operation) EvalError!void {
911 self.evalOp(op) catch |err| return self.failOp(op, err);
912 }
913
914 fn evalConstant(self: *Evaluator, op: *ir.Operation) EvalError!void {
915 const result = op.getResult(0) orelse return error.InvalidResultCount;
916 const kind = try scalarKindOfType(result.type);
917 const attr = op.getAttr("value") orelse return error.MissingAttribute;
918 try self.setOnlyResult(op, .{ .scalar = try scalarFromAttribute(kind, attr) });
919 }
920
921 fn evalBinary(self: *Evaluator, op: *ir.Operation, binary_op: BinaryOp) EvalError!void {
922 if (op.operands.items.len != 2) return error.UnsupportedOperation;
923 const result = op.getResult(0) orelse return error.InvalidResultCount;
924 const kind = try scalarKindOfType(result.type);
925 const lhs = try self.scalarOf(op.operands.items[0].value);
926 const rhs = try self.scalarOf(op.operands.items[1].value);
927 try self.setOnlyResult(op, .{ .scalar = try evalScalarBinary(kind, binary_op, lhs, rhs) });
928 }
929
930 fn evalUnary(self: *Evaluator, op: *ir.Operation, unary_op: UnaryOp) EvalError!void {
931 if (op.operands.items.len != 1) return error.UnsupportedOperation;
932 const result = op.getResult(0) orelse return error.InvalidResultCount;
933 const kind = try scalarKindOfType(result.type);
934 const input = try self.scalarOf(op.operands.items[0].value);
935 try self.setOnlyResult(op, .{ .scalar = try evalScalarUnary(kind, unary_op, input) });
936 }
937
938 fn evalPow(self: *Evaluator, op: *ir.Operation) EvalError!void {
939 if (op.operands.items.len != 2) return error.UnsupportedOperation;
940 const result = op.getResult(0) orelse return error.InvalidResultCount;
941 const kind = try scalarKindOfType(result.type);
942 const base = try self.scalarOf(op.operands.items[0].value);
943 const exponent = try self.scalarOf(op.operands.items[1].value);
944 try self.setOnlyResult(op, .{ .scalar = try evalScalarPow(kind, base, exponent) });
945 }
946
947 fn evalAtan2(self: *Evaluator, op: *ir.Operation) EvalError!void {
948 if (op.operands.items.len != 2) return error.UnsupportedOperation;
949 const result = op.getResult(0) orelse return error.InvalidResultCount;
950 const kind = try scalarKindOfType(result.type);
951 const y = try self.scalarOf(op.operands.items[0].value);
952 const x = try self.scalarOf(op.operands.items[1].value);
953 try self.setOnlyResult(op, .{ .scalar = try evalScalarAtan2(kind, y, x) });
954 }
955
956 fn evalFma(self: *Evaluator, op: *ir.Operation) EvalError!void {
957 if (op.operands.items.len != 3) return error.UnsupportedOperation;
958 const result = op.getResult(0) orelse return error.InvalidResultCount;
959 const kind = try scalarKindOfType(result.type);
960 const a = try self.scalarOf(op.operands.items[0].value);
961 const b = try self.scalarOf(op.operands.items[1].value);
962 const c = try self.scalarOf(op.operands.items[2].value);
963 try self.setOnlyResult(op, .{ .scalar = try evalScalarFma(kind, a, b, c) });
964 }
965
966 fn evalCmp(self: *Evaluator, op: *ir.Operation) EvalError!void {
967 if (op.operands.items.len != 2) return error.UnsupportedOperation;
968 const predicate = try cmpPredicate(op);
969 const lhs = try self.scalarOf(op.operands.items[0].value);
970 const rhs = try self.scalarOf(op.operands.items[1].value);
971 try self.setOnlyResult(op, .{ .scalar = .{ .bool = try evalCmpPredicate(predicate, lhs, rhs) } });
972 }
973
974 fn evalCast(self: *Evaluator, op: *ir.Operation) EvalError!void {
975 if (op.operands.items.len != 1) return error.UnsupportedOperation;
976 const result = op.getResult(0) orelse return error.InvalidResultCount;
977 const target = try scalarKindOfType(result.type);
978 const input = try self.scalarOf(op.operands.items[0].value);
979 try self.setOnlyResult(op, .{ .scalar = try castScalar(target, input) });
980 }
981
982 fn evalSelect(self: *Evaluator, op: *ir.Operation) EvalError!void {
983 if (op.operands.items.len != 3) return error.UnsupportedOperation;
984 const cond = try self.boolOf(op.operands.items[0].value);
985 const chosen = if (cond)
986 try self.resolve(op.operands.items[1].value)
987 else
988 try self.resolve(op.operands.items[2].value);
989 try self.setOnlyResult(op, chosen);
990 }
991
992 fn evalLogical(self: *Evaluator, op: *ir.Operation, logical_op: LogicalOp) EvalError!void {
993 if (op.operands.items.len != 2) return error.UnsupportedOperation;
994 const result = op.getResult(0) orelse return error.InvalidResultCount;
995 const kind = try scalarKindOfType(result.type);
996 const lhs = try self.scalarOf(op.operands.items[0].value);
997 const rhs = try self.scalarOf(op.operands.items[1].value);
998 try self.setOnlyResult(op, .{ .scalar = try evalScalarLogical(kind, logical_op, lhs, rhs) });
999 }
1000
1001 fn evalNot(self: *Evaluator, op: *ir.Operation) EvalError!void {
1002 if (op.operands.items.len != 1) return error.UnsupportedOperation;
1003 const result = op.getResult(0) orelse return error.InvalidResultCount;
1004 const kind = try scalarKindOfType(result.type);
1005 const input = try self.scalarOf(op.operands.items[0].value);
1006 try self.setOnlyResult(op, .{ .scalar = try evalScalarNot(kind, input) });
1007 }
1008
1009 fn evalPopCount(self: *Evaluator, op: *ir.Operation) EvalError!void {
1010 if (op.operands.items.len != 1) return error.UnsupportedOperation;
1011 const result = op.getResult(0) orelse return error.InvalidResultCount;
1012 const kind = try scalarKindOfType(result.type);
1013 const input = try self.scalarOf(op.operands.items[0].value);
1014 try self.setOnlyResult(op, .{ .scalar = try evalScalarPopCount(kind, input) });
1015 }
1016
1017 fn evalShift(self: *Evaluator, op: *ir.Operation, shift_op: ShiftOp) EvalError!void {
1018 if (op.operands.items.len != 2) return error.UnsupportedOperation;
1019 const result = op.getResult(0) orelse return error.InvalidResultCount;
1020 const kind = try scalarKindOfType(result.type);
1021 const value = try self.scalarOf(op.operands.items[0].value);
1022 const shift = try self.scalarOf(op.operands.items[1].value);
1023 try self.setOnlyResult(op, .{ .scalar = try evalScalarShift(kind, shift_op, value, shift) });
1024 }
1025
1026 fn evalBitcast(self: *Evaluator, op: *ir.Operation) EvalError!void {
1027 if (op.operands.items.len != 1) return error.UnsupportedOperation;
1028 const result = op.getResult(0) orelse return error.InvalidResultCount;
1029 const target = try scalarKindOfType(result.type);
1030 const input = try self.scalarOf(op.operands.items[0].value);
1031 try self.setOnlyResult(op, .{ .scalar = try evalScalarBitcast(target, input) });
1032 }
1033
1034 fn evalUmulhi(self: *Evaluator, op: *ir.Operation) EvalError!void {
1035 if (op.operands.items.len != 2) return error.UnsupportedOperation;
1036 const result = op.getResult(0) orelse return error.InvalidResultCount;
1037 const kind = try scalarKindOfType(result.type);
1038 const lhs = try self.scalarOf(op.operands.items[0].value);
1039 const rhs = try self.scalarOf(op.operands.items[1].value);
1040 try self.setOnlyResult(op, .{ .scalar = try evalScalarUmulhi(kind, lhs, rhs) });
1041 }
1042
1043 fn evalAlloca(self: *Evaluator, op: *ir.Operation) EvalError!void {
1044 if (op.operands.items.len > 1) return error.UnsupportedOperation;
1045 const result = op.getResult(0) orelse return error.InvalidResultCount;
1046 const desc = try describeMemrefType(result.type);
1047 const shared_alloc = desc.addr_space == .shared and std.mem.eql(u8, op.name.name, dialects.MemrefDialect.AllocOp.operation_name);
1048 if (shared_alloc) {
1049 if (self.shared_allocs.get(op)) |memref_index| {
1050 try self.setOnlyResult(op, .{ .memref = memref_index });
1051 return;
1052 }
1053 }
1054 const byte_size = desc.byte_size orelse blk: {
1055 if (op.operands.items.len != 1) return error.UnsupportedType;
1056 const elements = try self.indexOf(op.operands.items[0].value);
1057 if (elements < 0) return error.InvalidIndex;
1058 const element_count = std.math.cast(u64, elements) orelse return error.Overflow;
1059 const byte_size_u64 = std.math.mul(u64, element_count, scalarByteSize(desc.element)) catch return error.Overflow;
1060 break :blk std.math.cast(usize, byte_size_u64) orelse return error.Overflow;
1061 };
1062 if (shared_alloc) {
1063 const bytes = try self.retainedSharedBytes(op, byte_size);
1064 @memset(bytes, shared_alloc_sentinel);
1065 const memref_index = try self.appendMemref(.{ .element = desc.element, .bytes = bytes, .owned = false });
1066 try self.shared_allocs.put(op, memref_index);
1067 try self.setOnlyResult(op, .{ .memref = memref_index });
1068 return;
1069 }
1070 const bytes = try self.allocator.alloc(u8, byte_size);
1071 errdefer self.allocator.free(bytes);
1072 @memset(bytes, 0);
1073 const memref_index = try self.appendMemref(.{ .element = desc.element, .bytes = bytes });
1074 try self.setOnlyResult(op, .{ .memref = memref_index });
1075 }
1076
1077 fn retainedSharedBytes(self: *Evaluator, op: *ir.Operation, byte_size: usize) EvalError![]u8 {
1078 const entry = try self.shared_bytes.getOrPut(op);
1079 if (entry.found_existing and entry.value_ptr.len == byte_size) return entry.value_ptr.*;
1080 const bytes = self.allocator.alloc(u8, byte_size) catch |err| {
1081 if (!entry.found_existing) self.shared_bytes.removeByPtr(entry.key_ptr);
1082 return err;
1083 };
1084 if (entry.found_existing) self.allocator.free(entry.value_ptr.*);
1085 entry.value_ptr.* = bytes;
1086 return bytes;
1087 }
1088
1089 fn evalLoad(self: *Evaluator, op: *ir.Operation) EvalError!void {
1090 if (op.operands.items.len != 2) return error.UnsupportedOperation;
1091 const memref_index = try self.memrefIndexOf(op.operands.items[0].value);
1092 const index = try self.indexOf(op.operands.items[1].value);
1093 const memref = try self.memrefAt(memref_index);
1094 const result = op.getResult(0) orelse return error.InvalidResultCount;
1095 if (isVec4F32Type(result.type)) {
1096 if (memref.element != .f32) return error.UnsupportedType;
1097 var lanes: [4]f32 = undefined;
1098 for (&lanes, 0..) |*lane, offset| {
1099 lane.* = switch (try loadScalar(memref.*, index + @as(i64, @intCast(offset)))) {
1100 .f32 => |v| v,
1101 else => return error.UnsupportedType,
1102 };
1103 }
1104 try self.setOnlyResult(op, .{ .f32x4 = lanes });
1105 return;
1106 }
1107 const scalar = try loadScalar(memref.*, index);
1108 try self.setOnlyResult(op, .{ .scalar = scalar });
1109 }
1110
1111 fn evalStore(self: *Evaluator, op: *ir.Operation) EvalError!void {
1112 if (op.operands.items.len != 3) return error.UnsupportedOperation;
1113 const stored = try self.resolve(op.operands.items[0].value);
1114 const memref_index = try self.memrefIndexOf(op.operands.items[1].value);
1115 const index = try self.indexOf(op.operands.items[2].value);
1116 const memref = try self.memrefAt(memref_index);
1117 switch (stored) {
1118 .f32x4 => |lanes| {
1119 if (memref.element != .f32) return error.UnsupportedType;
1120 for (lanes, 0..) |lane, offset| {
1121 try storeScalar(memref, index + @as(i64, @intCast(offset)), .{ .f32 = lane });
1122 }
1123 },
1124 .scalar => |scalar| try storeScalar(memref, index, scalar),
1125 .memref => return error.UnsupportedType,
1126 }
1127 }
1128
1129 fn evalCpAsyncShared(self: *Evaluator, op: *ir.Operation) EvalError!void {
1130 if (op.operands.items.len != 4) return error.UnsupportedOperation;
1131 const copy = GpuDialect.CpAsyncSharedOp{ .op = op };
1132 const bytes = copy.getBytes() orelse return error.MissingAttribute;
1133 const dst_ref_index = try self.memrefIndexOf(copy.getDst());
1134 const dst_index = try self.indexOf(copy.getDstIndex());
1135 const src_ref_index = try self.memrefIndexOf(copy.getSrc());
1136 const src_index = try self.indexOf(copy.getSrcIndex());
1137 const src = try self.memrefAt(src_ref_index);
1138 const element_bytes: u32 = switch (src.element) {
1139 .f32, .i32, .u32 => 4,
1140 else => return error.UnsupportedType,
1141 };
1142 if (bytes % element_bytes != 0) return error.UnsupportedOperation;
1143 const count: i64 = @intCast(bytes / element_bytes);
1144 var offset: i64 = 0;
1145 while (offset < count) : (offset += 1) {
1146 const scalar = try loadScalar(src.*, src_index + offset);
1147 const dst = try self.memrefAt(dst_ref_index);
1148 if (dst.element != src.element) return error.UnsupportedType;
1149 try storeScalar(dst, dst_index + offset, scalar);
1150 }
1151 }
1152
1153 fn evalVecExtract(self: *Evaluator, op: *ir.Operation) EvalError!void {
1154 if (op.operands.items.len != 1) return error.UnsupportedOperation;
1155 const extract = dialects.ArithDialect.ExtractOp{ .op = op };
1156 const lane_index = extract.getIndex() orelse return error.MissingAttribute;
1157 if (lane_index < 0 or lane_index > 3) return error.InvalidAttribute;
1158 const lanes = switch (try self.resolve(op.operands.items[0].value)) {
1159 .f32x4 => |lanes| lanes,
1160 else => return error.UnsupportedType,
1161 };
1162 try self.setOnlyResult(op, .{ .scalar = .{ .f32 = lanes[@intCast(lane_index)] } });
1163 }
1164
1165 fn evalVecSplat(self: *Evaluator, op: *ir.Operation) EvalError!void {
1166 if (op.operands.items.len != 1) return error.UnsupportedOperation;
1167 const result = op.getResult(0) orelse return error.InvalidResultCount;
1168 if (!isVec4F32Type(result.type)) return error.UnsupportedType;
1169 const value = switch (try self.scalarOf(op.operands.items[0].value)) {
1170 .f32 => |v| v,
1171 else => return error.UnsupportedType,
1172 };
1173 try self.setOnlyResult(op, .{ .f32x4 = .{ value, value, value, value } });
1174 }
1175
1176 fn evalVecInsert(self: *Evaluator, op: *ir.Operation) EvalError!void {
1177 if (op.operands.items.len != 2) return error.UnsupportedOperation;
1178 const insert = dialects.ArithDialect.InsertOp{ .op = op };
1179 const lane_index = insert.getIndex() orelse return error.MissingAttribute;
1180 if (lane_index < 0 or lane_index > 3) return error.InvalidAttribute;
1181 var lanes = switch (try self.resolve(op.operands.items[0].value)) {
1182 .f32x4 => |resolved| resolved,
1183 else => return error.UnsupportedType,
1184 };
1185 lanes[@intCast(lane_index)] = switch (try self.scalarOf(op.operands.items[1].value)) {
1186 .f32 => |v| v,
1187 else => return error.UnsupportedType,
1188 };
1189 try self.setOnlyResult(op, .{ .f32x4 = lanes });
1190 }
1191
1192 fn evalAtomicRmw(self: *Evaluator, op: *ir.Operation) EvalError!void {
1193 if (op.operands.items.len != 3) return error.UnsupportedOperation;
1194 const kind = (dialects.MemrefDialect.AtomicRmwOp{ .op = op }).getKind() orelse return error.MissingAttribute;
1195 const operand = try self.scalarOf(op.operands.items[0].value);
1196 const memref_index = try self.memrefIndexOf(op.operands.items[1].value);
1197 const index = try self.indexOf(op.operands.items[2].value);
1198 const memref = try self.memrefAt(memref_index);
1199 const old = try loadScalar(memref.*, index);
1200 const updated = try atomicApply(kind, old, operand);
1201 try storeScalar(memref, index, updated);
1202 try self.setOnlyResult(op, .{ .scalar = old });
1203 }
1204
1205 fn evalAtomicCas(self: *Evaluator, op: *ir.Operation) EvalError!void {
1206 if (op.operands.items.len != 4) return error.UnsupportedOperation;
1207 const memref_index = try self.memrefIndexOf(op.operands.items[2].value);
1208 const index = try self.indexOf(op.operands.items[3].value);
1209 const memref = try self.memrefAt(memref_index);
1210 const expected = try castScalar(memref.element, try self.scalarOf(op.operands.items[0].value));
1211 const desired = try castScalar(memref.element, try self.scalarOf(op.operands.items[1].value));
1212 const old = try loadScalar(memref.*, index);
1213 if (try evalCmpPredicate(.eq, old, expected)) {
1214 try storeScalar(memref, index, desired);
1215 }
1216 try self.setOnlyResult(op, .{ .scalar = old });
1217 }
1218
1219 fn evalGpuId(self: *Evaluator, op: *ir.Operation, kind: GpuIdKind) EvalError!void {
1220 const dim = gpuDimension(op) orelse return error.MissingAttribute;
1221 const index = dimensionIndex(dim);
1222 const value = switch (kind) {
1223 .global => self.launch.global[index],
1224 .thread => self.launch.thread[index],
1225 .block => self.launch.block[index],
1226 .block_dim => self.launch.block_dim[index],
1227 .grid_dim => self.launch.grid_dim[index],
1228 };
1229 try self.setOnlyResult(op, .{ .scalar = .{ .index = value } });
1230 }
1231
1232 fn evalLaneId(self: *Evaluator, op: *ir.Operation) EvalError!void {
1233 try self.setOnlyResult(op, .{ .scalar = .{ .index = @intCast(try laneId(self.launch)) } });
1234 }
1235
1236 fn evalWarpId(self: *Evaluator, op: *ir.Operation) EvalError!void {
1237 try self.setOnlyResult(op, .{ .scalar = .{ .index = @intCast(try warpId(self.launch)) } });
1238 }
1239
1240 fn evalActiveMask(self: *Evaluator, op: *ir.Operation) EvalError!void {
1241 const mask = if (self.active_lanes) |lanes| try activeMaskForLanes(self.launch, lanes) else try blockActiveMask(self.launch);
1242 try self.setOnlyResult(op, .{ .scalar = .{ .i32 = @bitCast(mask) } });
1243 }
1244
1245 fn evalKernelCollective(self: *Evaluator, op: *ir.Operation, lanes: []*LaneState) EvalError!bool {
1246 if (std.mem.eql(u8, op.name.name, GpuDialect.WarpReduceOp.operation_name)) {
1247 try self.evalWarpReduce(op, lanes);
1248 return true;
1249 }
1250 if (std.mem.eql(u8, op.name.name, GpuDialect.WarpScanOp.operation_name)) {
1251 try self.evalWarpScan(op, lanes);
1252 return true;
1253 }
1254 if (std.mem.eql(u8, op.name.name, GpuDialect.ShflSyncOp.operation_name)) {
1255 try self.evalWarpShuffle(op, lanes);
1256 return true;
1257 }
1258 if (std.mem.eql(u8, op.name.name, GpuDialect.AllSyncOp.operation_name)) {
1259 try self.evalWarpVote(op, lanes, .all);
1260 return true;
1261 }
1262 if (std.mem.eql(u8, op.name.name, GpuDialect.AnySyncOp.operation_name)) {
1263 try self.evalWarpVote(op, lanes, .any);
1264 return true;
1265 }
1266 if (std.mem.eql(u8, op.name.name, GpuDialect.BallotSyncOp.operation_name)) {
1267 try self.evalWarpVote(op, lanes, .ballot);
1268 return true;
1269 }
1270 if (std.mem.eql(u8, op.name.name, GpuDialect.MmaSyncOp.operation_name)) {
1271 try self.evalMmaSync(op, lanes);
1272 return true;
1273 }
1274 return false;
1275 }
1276
1277 fn f32OfLane(self: *Evaluator, lane: *LaneState, value: *ir.Value) EvalError!f32 {
1278 return scalarToF32(try self.scalarOfLane(lane, value));
1279 }
1280
1281 fn evalMmaSync(self: *Evaluator, op: *ir.Operation, lanes: []*LaneState) EvalError!void {
1282 if (op.operands.items.len != 10) return error.UnsupportedOperation;
1283 if (op.getNumResults() != 4) return error.InvalidResultCount;
1284 const mma = GpuDialect.MmaSyncOp{ .op = op };
1285 const shape = mma.getShape() orelse return error.MissingAttribute;
1286 if (shape.m != 16 or shape.n != 8 or shape.k != 8) return error.UnsupportedOperation;
1287
1288 self.warps.clearRetainingCapacity();
1289 for (lanes) |seed| {
1290 const warp = try warpId(seed.launch);
1291 if (self.warps.contains(warp)) continue;
1292 try self.warps.put(warp, {});
1293
1294 var members: [warp_size]?*LaneState = @splat(null);
1295 for (lanes) |candidate| {
1296 if (try warpId(candidate.launch) != warp) continue;
1297 members[try laneId(candidate.launch)] = candidate;
1298 }
1299
1300 var a: [16][8]f32 = undefined;
1301 var b: [8][8]f32 = undefined;
1302 var c: [16][8]f32 = undefined;
1303 for (0..warp_size) |lane_id| {
1304 const member = members[lane_id] orelse return error.UnsupportedOperation;
1305 const group = lane_id >> 2;
1306 const tid = lane_id & 3;
1307 a[group][tid] = try self.f32OfLane(member, mma.getA(0));
1308 a[group + 8][tid] = try self.f32OfLane(member, mma.getA(1));
1309 a[group][tid + 4] = try self.f32OfLane(member, mma.getA(2));
1310 a[group + 8][tid + 4] = try self.f32OfLane(member, mma.getA(3));
1311 b[tid][group] = try self.f32OfLane(member, mma.getB(0));
1312 b[tid + 4][group] = try self.f32OfLane(member, mma.getB(1));
1313 c[group][tid * 2] = try self.f32OfLane(member, mma.getC(0));
1314 c[group][tid * 2 + 1] = try self.f32OfLane(member, mma.getC(1));
1315 c[group + 8][tid * 2] = try self.f32OfLane(member, mma.getC(2));
1316 c[group + 8][tid * 2 + 1] = try self.f32OfLane(member, mma.getC(3));
1317 }
1318
1319 var d: [16][8]f32 = undefined;
1320 for (0..16) |row| {
1321 for (0..8) |col| {
1322 var acc = c[row][col];
1323 for (0..8) |k| {
1324 acc = @mulAdd(f32, a[row][k], b[k][col], acc);
1325 }
1326 d[row][col] = acc;
1327 }
1328 }
1329
1330 for (0..warp_size) |lane_id| {
1331 const member = members[lane_id].?;
1332 const group = lane_id >> 2;
1333 const tid = lane_id & 3;
1334 try member.setValue(op.getResult(0).?, .{ .scalar = .{ .f32 = d[group][tid * 2] } });
1335 try member.setValue(op.getResult(1).?, .{ .scalar = .{ .f32 = d[group][tid * 2 + 1] } });
1336 try member.setValue(op.getResult(2).?, .{ .scalar = .{ .f32 = d[group + 8][tid * 2] } });
1337 try member.setValue(op.getResult(3).?, .{ .scalar = .{ .f32 = d[group + 8][tid * 2 + 1] } });
1338 }
1339 }
1340 }
1341
1342 fn evalWarpReduce(self: *Evaluator, op: *ir.Operation, lanes: []*LaneState) EvalError!void {
1343 if (op.operands.items.len != 2) return error.UnsupportedOperation;
1344 if (op.getNumResults() != 1) return error.InvalidResultCount;
1345 const reduce = GpuDialect.WarpReduceOp{ .op = op };
1346 const op_kind = reduce.getOpKind() orelse return error.MissingAttribute;
1347 const result = reduce.getResult();
1348 const kind = try scalarKindOfType(result.type);
1349 const mask_value = reduce.getMask();
1350 const input_value = reduce.getValue();
1351
1352 for (lanes) |lane| {
1353 const target_warp = try warpId(lane.launch);
1354 const mask = try scalarToMask(try self.scalarOfLane(lane, mask_value));
1355 var reduced: ?Scalar = null;
1356 for (lanes) |source_lane| {
1357 if (try warpId(source_lane.launch) != target_warp) continue;
1358 const source_lane_id = try laneId(source_lane.launch);
1359 if ((mask & (@as(u32, 1) << @intCast(source_lane_id))) == 0) continue;
1360 const source = try self.scalarOfLane(source_lane, input_value);
1361 reduced = if (reduced) |current| try evalWarpScalar(kind, op_kind, current, source) else try castScalar(kind, source);
1362 }
1363 try lane.setValue(result, .{ .scalar = reduced orelse return error.UnsupportedOperation });
1364 }
1365 }
1366
1367 fn evalWarpScan(self: *Evaluator, op: *ir.Operation, lanes: []*LaneState) EvalError!void {
1368 if (op.operands.items.len != 2) return error.UnsupportedOperation;
1369 if (op.getNumResults() != 1) return error.InvalidResultCount;
1370 const scan = GpuDialect.WarpScanOp{ .op = op };
1371 const op_kind = scan.getOpKind() orelse return error.MissingAttribute;
1372 const inclusive = scan.isInclusive();
1373 const result = scan.getResult();
1374 const kind = try scalarKindOfType(result.type);
1375 const mask_value = scan.getMask();
1376 const input_value = scan.getValue();
1377
1378 for (lanes) |lane| {
1379 const target_warp = try warpId(lane.launch);
1380 const target_lane_id = try laneId(lane.launch);
1381 const mask = try scalarToMask(try self.scalarOfLane(lane, mask_value));
1382 var scanned: ?Scalar = null;
1383 var source_lane_id: usize = 0;
1384 while (source_lane_id < warp_size) : (source_lane_id += 1) {
1385 if (source_lane_id > target_lane_id) break;
1386 if (!inclusive and source_lane_id == target_lane_id) break;
1387 const source_mask = @as(u32, 1) << @intCast(source_lane_id);
1388 if ((mask & source_mask) == 0) continue;
1389 for (lanes) |source_lane| {
1390 if (try warpId(source_lane.launch) != target_warp) continue;
1391 if (try laneId(source_lane.launch) != source_lane_id) continue;
1392 const source = try self.scalarOfLane(source_lane, input_value);
1393 scanned = if (scanned) |current| try evalWarpScalar(kind, op_kind, current, source) else try castScalar(kind, source);
1394 break;
1395 }
1396 }
1397 try lane.setValue(result, .{ .scalar = scanned orelse try warpIdentity(kind, op_kind) });
1398 }
1399 }
1400
1401 fn evalWarpShuffle(self: *Evaluator, op: *ir.Operation, lanes: []*LaneState) EvalError!void {
1402 if (op.operands.items.len != 3) return error.UnsupportedOperation;
1403 if (op.getNumResults() != 1) return error.InvalidResultCount;
1404 const shuffle = GpuDialect.ShflSyncOp{ .op = op };
1405 const mode = shuffle.getMode() orelse return error.MissingAttribute;
1406 const result = shuffle.getResult();
1407 const kind = try scalarKindOfType(result.type);
1408 const mask_value = shuffle.getMask();
1409 const input_value = shuffle.getSrc();
1410 const lane_or_delta_value = shuffle.getLaneOrDelta();
1411
1412 for (lanes) |lane| {
1413 const target_warp = try warpId(lane.launch);
1414 const target_lane_id = try laneId(lane.launch);
1415 const lane_or_delta = try self.indexOfLane(lane, lane_or_delta_value);
1416 const source_lane_id = try shuffleSourceLane(mode, target_lane_id, lane_or_delta);
1417 const source_mask = @as(u32, 1) << @intCast(source_lane_id);
1418 const mask = try scalarToMask(try self.scalarOfLane(lane, mask_value));
1419 if ((mask & source_mask) == 0) return error.UnsupportedOperation;
1420
1421 var found_source = false;
1422 for (lanes) |source_lane| {
1423 if (try warpId(source_lane.launch) != target_warp) continue;
1424 if (try laneId(source_lane.launch) != source_lane_id) continue;
1425 const source = try self.scalarOfLane(source_lane, input_value);
1426 try lane.setValue(result, .{ .scalar = try castScalar(kind, source) });
1427 found_source = true;
1428 break;
1429 }
1430 if (!found_source) return error.UnsupportedOperation;
1431 }
1432 }
1433
1434 fn evalWarpVote(self: *Evaluator, op: *ir.Operation, lanes: []*LaneState, vote: WarpVoteKind) EvalError!void {
1435 if (op.operands.items.len != 2) return error.UnsupportedOperation;
1436 if (op.getNumResults() != 1) return error.InvalidResultCount;
1437 const result = op.getResult(0) orelse return error.InvalidResultCount;
1438 const mask_value = op.operands.items[0].value;
1439 const predicate_value = op.operands.items[1].value;
1440
1441 for (lanes) |lane| {
1442 const target_warp = try warpId(lane.launch);
1443 const mask = try scalarToMask(try self.scalarOfLane(lane, mask_value));
1444 var saw_lane = false;
1445 var all_value = true;
1446 var any_value = false;
1447 var ballot_value: u32 = 0;
1448 for (lanes) |source_lane| {
1449 if (try warpId(source_lane.launch) != target_warp) continue;
1450 const source_lane_id = try laneId(source_lane.launch);
1451 const source_mask = @as(u32, 1) << @intCast(source_lane_id);
1452 if ((mask & source_mask) == 0) continue;
1453 saw_lane = true;
1454 const predicate = try self.boolOfLane(source_lane, predicate_value);
1455 if (predicate) {
1456 any_value = true;
1457 ballot_value |= source_mask;
1458 } else {
1459 all_value = false;
1460 }
1461 }
1462
1463 const scalar: Scalar = switch (vote) {
1464 .all => .{ .bool = saw_lane and all_value },
1465 .any => .{ .bool = any_value },
1466 .ballot => .{ .i32 = @bitCast(ballot_value) },
1467 };
1468 try lane.setValue(result, .{ .scalar = scalar });
1469 }
1470 }
1471
1472 fn scalarOfLane(self: *Evaluator, lane: *LaneState, value: *ir.Value) EvalError!Scalar {
1473 _ = self;
1474 return switch (lane.values.get(value) orelse return error.MissingValue) {
1475 .scalar => |scalar| scalar,
1476 .f32x4, .memref => error.UnsupportedType,
1477 };
1478 }
1479
1480 fn resolveLane(self: *Evaluator, lane: *LaneState, value: *ir.Value) EvalError!RuntimeValue {
1481 _ = self;
1482 return lane.values.get(value) orelse error.MissingValue;
1483 }
1484
1485 fn boolOfLane(self: *Evaluator, lane: *LaneState, value: *ir.Value) EvalError!bool {
1486 return switch (try self.scalarOfLane(lane, value)) {
1487 .bool => |v| v,
1488 else => error.InvalidCondition,
1489 };
1490 }
1491
1492 fn indexOfLane(self: *Evaluator, lane: *LaneState, value: *ir.Value) EvalError!i64 {
1493 return scalarToI64(try self.scalarOfLane(lane, value));
1494 }
1495
1496 fn evalFor(self: *Evaluator, op: *ir.Operation) EvalError!void {
1497 if (op.operands.items.len < 3) return error.UnsupportedOperation;
1498 const lower = try self.indexOf(op.operands.items[0].value);
1499 const upper = try self.indexOf(op.operands.items[1].value);
1500 const step = try self.indexOf(op.operands.items[2].value);
1501 if (step <= 0) return error.UnsupportedOperation;
1502
1503 const iter_count = op.operands.items.len - 3;
1504 if (op.getNumResults() != iter_count) return error.InvalidResultCount;
1505 const region = op.getRegion(0) orelse return error.UnsupportedOperation;
1506 const body = region.getEntryBlock() orelse return error.UnsupportedOperation;
1507 if (body.arguments.items.len != iter_count + 1) return error.UnsupportedOperation;
1508
1509 const iter_values = try self.allocator.alloc(RuntimeValue, iter_count);
1510 defer self.allocator.free(iter_values);
1511 for (iter_values, 0..) |*value, i| {
1512 value.* = try self.resolve(op.operands.items[3 + i].value);
1513 }
1514
1515 var iv = lower;
1516 while (iv < upper) {
1517 try self.setValue(body.getArgument(0).?, .{ .scalar = .{ .index = iv } });
1518 for (iter_values, 0..) |value, i| {
1519 try self.setValue(body.getArgument(i + 1).?, value);
1520 }
1521
1522 const yielded = try self.evalBlock(body);
1523 if (yielded.len != iter_count) {
1524 self.allocator.free(yielded);
1525 return error.InvalidResultCount;
1526 }
1527 @memcpy(iter_values, yielded);
1528 self.allocator.free(yielded);
1529
1530 iv = std.math.add(i64, iv, step) catch return error.Overflow;
1531 }
1532
1533 for (iter_values, 0..) |value, i| {
1534 const result = op.getResult(i) orelse return error.InvalidResultCount;
1535 try self.setValue(result, value);
1536 }
1537 }
1538
1539 fn evalWhile(self: *Evaluator, op: *ir.Operation) EvalError!void {
1540 const while_op = dialects.ScfDialect.WhileOp{ .op = op };
1541 const carry_count = op.operands.items.len;
1542 if (op.getNumResults() != carry_count) return error.InvalidResultCount;
1543 const before = while_op.getBeforeBlock();
1544 const after = while_op.getAfterBlock();
1545 if (before.arguments.items.len != carry_count) return error.UnsupportedOperation;
1546 if (after.arguments.items.len != carry_count) return error.UnsupportedOperation;
1547
1548 const carries = try self.allocator.alloc(RuntimeValue, carry_count);
1549 defer self.allocator.free(carries);
1550 for (carries, 0..) |*value, i| {
1551 value.* = try self.resolve(op.operands.items[i].value);
1552 }
1553
1554 var iterations: usize = 0;
1555 while (true) {
1556 if (iterations >= max_while_iterations) return error.WhileIterationLimit;
1557 iterations += 1;
1558
1559 for (carries, 0..) |value, i| {
1560 try self.setValue(before.getArgument(i).?, value);
1561 }
1562 const terminator = try self.evalBlockUntil(before, dialects.ScfDialect.ConditionOp.operation_name);
1563 const condition_op = dialects.ScfDialect.ConditionOp{ .op = terminator };
1564 const args = condition_op.getArgs();
1565 if (args.len != carry_count) return error.InvalidResultCount;
1566
1567 if (!(try self.boolOf(condition_op.getCondition()))) {
1568 for (args, 0..) |arg, i| {
1569 const result = op.getResult(i) orelse return error.InvalidResultCount;
1570 try self.setValue(result, try self.resolve(arg));
1571 }
1572 return;
1573 }
1574
1575 for (args, 0..) |arg, i| {
1576 try self.setValue(after.getArgument(i).?, try self.resolve(arg));
1577 }
1578 const yielded = try self.evalBlock(after);
1579 if (yielded.len != carry_count) {
1580 self.allocator.free(yielded);
1581 return error.InvalidResultCount;
1582 }
1583 @memcpy(carries, yielded);
1584 self.allocator.free(yielded);
1585 }
1586 }
1587
1588 fn evalBlockUntil(self: *Evaluator, block: *ir.Block, terminator_name: []const u8) EvalError!*ir.Operation {
1589 var iter = block.operations.head;
1590 while (iter) |op_ptr| {
1591 const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
1592 if (std.mem.eql(u8, op.name.name, terminator_name)) return op;
1593 try self.evalOpWithDiagnostic(op);
1594 iter = op.next_op;
1595 }
1596 return error.MissingTerminator;
1597 }
1598
1599 fn evalIf(self: *Evaluator, op: *ir.Operation) EvalError!void {
1600 if (op.operands.items.len != 1) return error.UnsupportedOperation;
1601 const cond = try self.boolOf(op.operands.items[0].value);
1602 const region = op.getRegion(if (cond) 0 else 1) orelse return error.UnsupportedOperation;
1603 const block = region.getEntryBlock() orelse return error.UnsupportedOperation;
1604 const yielded = try self.evalBlock(block);
1605 defer self.allocator.free(yielded);
1606 if (yielded.len != op.getNumResults()) return error.InvalidResultCount;
1607 for (yielded, 0..) |value, i| {
1608 const result = op.getResult(i) orelse return error.InvalidResultCount;
1609 try self.setValue(result, value);
1610 }
1611 }
1612
1613 fn setOnlyResult(self: *Evaluator, op: *ir.Operation, runtime_value: RuntimeValue) EvalError!void {
1614 if (op.getNumResults() != 1) return error.InvalidResultCount;
1615 const result = op.getResult(0) orelse return error.InvalidResultCount;
1616 try self.setValue(result, runtime_value);
1617 }
1618
1619 fn scalarOf(self: *Evaluator, value: *ir.Value) EvalError!Scalar {
1620 return switch (try self.resolve(value)) {
1621 .scalar => |scalar| scalar,
1622 .f32x4, .memref => error.UnsupportedType,
1623 };
1624 }
1625
1626 fn boolOf(self: *Evaluator, value: *ir.Value) EvalError!bool {
1627 return switch (try self.scalarOf(value)) {
1628 .bool => |v| v,
1629 else => error.InvalidCondition,
1630 };
1631 }
1632
1633 fn indexOf(self: *Evaluator, value: *ir.Value) EvalError!i64 {
1634 return scalarToI64(try self.scalarOf(value));
1635 }
1636
1637 fn memrefIndexOf(self: *Evaluator, value: *ir.Value) EvalError!usize {
1638 return switch (try self.resolve(value)) {
1639 .memref => |idx| idx,
1640 .scalar, .f32x4 => error.InvalidMemref,
1641 };
1642 }
1643
1644 fn memrefAt(self: *Evaluator, index: usize) EvalError!*Memref {
1645 if (index >= self.memrefs.items.len) return error.InvalidMemref;
1646 return &self.memrefs.items[index];
1647 }
1648 };
1649
1650 const MemrefDesc = struct {
1651 element: ScalarKind,
1652 byte_size: ?usize,
1653 addr_space: dialects.AddressSpace,
1654 };
1655
1656 fn describeMemrefType(typ: ir.Type) EvalError!MemrefDesc {
1657 const storage = typ.getDialectStorage() orelse return error.UnsupportedType;
1658 if (!std.mem.eql(u8, storage.name, dialects.MemrefDialect.name)) return error.UnsupportedType;
1659 const params = dialects.MemrefDialect.parseMemrefParams(storage.param_key) orelse return error.UnsupportedType;
1660 const element = scalarKindFromTypeName(params.element_type_name) orelse return error.UnsupportedType;
1661 const byte_size = if (params.size) |elements| blk: {
1662 const byte_size_u64 = std.math.mul(u64, elements, scalarByteSize(element)) catch return error.Overflow;
1663 break :blk std.math.cast(usize, byte_size_u64) orelse return error.Overflow;
1664 } else null;
1665 return .{ .element = element, .byte_size = byte_size, .addr_space = params.addr_space };
1666 }
1667
1668 fn scalarKindOfType(typ: ir.Type) EvalError!ScalarKind {
1669 const storage = typ.getDialectStorage() orelse return error.UnsupportedType;
1670 return scalarKindFromTypeName(storage.name) orelse error.UnsupportedType;
1671 }
1672
1673 fn scalarKindFromTypeName(name: []const u8) ?ScalarKind {
1674 const kind = dialects.arith.scalarKindFromTypeName(name) orelse return null;
1675 return switch (kind) {
1676 .bool => .bool,
1677 .index => .index,
1678 .i8 => .i8,
1679 .i16 => .i16,
1680 .i32 => .i32,
1681 .u8 => .u8,
1682 .u16 => .u16,
1683 .u32 => .u32,
1684 .i64 => .i64,
1685 .u64 => .u64,
1686 .f16 => .f16,
1687 .bf16 => .bf16,
1688 .f32 => .f32,
1689 .f64 => .f64,
1690 };
1691 }
1692
1693 fn gpuDimension(op: *ir.Operation) ?Dimension {
1694 const attr = op.getAttr("dim") orelse return null;
1695 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return null;
1696 return Dimension.fromString(dialect_attr.payload);
1697 }
1698
1699 fn dimensionIndex(dim: Dimension) usize {
1700 return switch (dim) {
1701 .x => 0,
1702 .y => 1,
1703 .z => 2,
1704 };
1705 }
1706
1707 fn scalarByteSize(kind: ScalarKind) u64 {
1708 return switch (kind) {
1709 .bool => 1,
1710 .i8, .u8 => 1,
1711 .i16, .u16, .f16, .bf16 => 2,
1712 .index, .i64, .u64, .f64 => 8,
1713 .i32, .u32, .f32 => 4,
1714 };
1715 }
1716
1717 fn scalarFromAttribute(kind: ScalarKind, attr: ir.Attribute) EvalError!Scalar {
1718 return switch (kind) {
1719 .bool => .{ .bool = dialects.ArithDialect.getBoolValue(attr) orelse return error.InvalidAttribute },
1720 .index => .{ .index = dialects.ArithDialect.getIntValue(attr) orelse return error.InvalidAttribute },
1721 .i8 => .{ .i8 = std.math.cast(i8, dialects.ArithDialect.getIntValue(attr) orelse return error.InvalidAttribute) orelse return error.Overflow },
1722 .i16 => .{ .i16 = std.math.cast(i16, dialects.ArithDialect.getIntValue(attr) orelse return error.InvalidAttribute) orelse return error.Overflow },
1723 .i32 => .{ .i32 = std.math.cast(i32, dialects.ArithDialect.getIntValue(attr) orelse return error.InvalidAttribute) orelse return error.Overflow },
1724 .u8 => .{ .u8 = std.math.cast(u8, dialects.ArithDialect.getIntValue(attr) orelse return error.InvalidAttribute) orelse return error.Overflow },
1725 .u16 => .{ .u16 = std.math.cast(u16, dialects.ArithDialect.getIntValue(attr) orelse return error.InvalidAttribute) orelse return error.Overflow },
1726 .u32 => .{ .u32 = std.math.cast(u32, dialects.ArithDialect.getIntValue(attr) orelse return error.InvalidAttribute) orelse return error.Overflow },
1727 .i64 => .{ .i64 = dialects.ArithDialect.getIntValue(attr) orelse return error.InvalidAttribute },
1728 .u64 => .{ .u64 = std.math.cast(u64, dialects.ArithDialect.getIntValue(attr) orelse return error.InvalidAttribute) orelse return error.Overflow },
1729 .f16 => .{ .f16 = @floatCast(dialects.ArithDialect.getFloatValue(attr) orelse return error.InvalidAttribute) },
1730 .bf16 => .{ .bf16 = Bf16.fromF32(@floatCast(dialects.ArithDialect.getFloatValue(attr) orelse return error.InvalidAttribute)) },
1731 .f32 => .{ .f32 = @floatCast(dialects.ArithDialect.getFloatValue(attr) orelse return error.InvalidAttribute) },
1732 .f64 => .{ .f64 = dialects.ArithDialect.getFloatValue(attr) orelse return error.InvalidAttribute },
1733 };
1734 }
1735
1736 fn isVec4F32Type(typ: ir.Type) bool {
1737 const name = typ.getDialectTypeName() orelse return false;
1738 return std.mem.eql(u8, name, dialects.arith.type_names.vec4xf32);
1739 }
1740
1741 fn loadScalar(memref: Memref, index_i64: i64) EvalError!Scalar {
1742 const index = try checkedElementIndex(memref, index_i64);
1743 const offset = index * @as(usize, @intCast(scalarByteSize(memref.element)));
1744 return switch (memref.element) {
1745 .bool => .{ .bool = memref.bytes[offset] != 0 },
1746 .i8 => .{ .i8 = @bitCast(memref.bytes[offset]) },
1747 .u8 => .{ .u8 = memref.bytes[offset] },
1748 .i16 => .{ .i16 = std.mem.readInt(i16, memref.bytes[offset..][0..2], .little) },
1749 .u16 => .{ .u16 = std.mem.readInt(u16, memref.bytes[offset..][0..2], .little) },
1750 .i32 => .{ .i32 = std.mem.readInt(i32, memref.bytes[offset..][0..4], .little) },
1751 .u32 => .{ .u32 = std.mem.readInt(u32, memref.bytes[offset..][0..4], .little) },
1752 .i64 => .{ .i64 = std.mem.readInt(i64, memref.bytes[offset..][0..8], .little) },
1753 .u64 => .{ .u64 = std.mem.readInt(u64, memref.bytes[offset..][0..8], .little) },
1754 .index => .{ .index = std.mem.readInt(i64, memref.bytes[offset..][0..8], .little) },
1755 .f16 => .{ .f16 = @bitCast(std.mem.readInt(u16, memref.bytes[offset..][0..2], .little)) },
1756 .bf16 => .{ .bf16 = .{ .bits = std.mem.readInt(u16, memref.bytes[offset..][0..2], .little) } },
1757 .f32 => .{ .f32 = @bitCast(std.mem.readInt(u32, memref.bytes[offset..][0..4], .little)) },
1758 .f64 => .{ .f64 = @bitCast(std.mem.readInt(u64, memref.bytes[offset..][0..8], .little)) },
1759 };
1760 }
1761
1762 fn storeScalar(memref: *Memref, index_i64: i64, scalar: Scalar) EvalError!void {
1763 const index = try checkedElementIndex(memref.*, index_i64);
1764 const offset = index * @as(usize, @intCast(scalarByteSize(memref.element)));
1765 const value = try castScalar(memref.element, scalar);
1766 switch (value) {
1767 .bool => |v| memref.bytes[offset] = if (v) 1 else 0,
1768 .i8 => |v| memref.bytes[offset] = @bitCast(v),
1769 .u8 => |v| memref.bytes[offset] = v,
1770 .i16 => |v| std.mem.writeInt(i16, memref.bytes[offset..][0..2], v, .little),
1771 .u16 => |v| std.mem.writeInt(u16, memref.bytes[offset..][0..2], v, .little),
1772 .i32 => |v| std.mem.writeInt(i32, memref.bytes[offset..][0..4], v, .little),
1773 .u32 => |v| std.mem.writeInt(u32, memref.bytes[offset..][0..4], v, .little),
1774 .i64 => |v| std.mem.writeInt(i64, memref.bytes[offset..][0..8], v, .little),
1775 .u64 => |v| std.mem.writeInt(u64, memref.bytes[offset..][0..8], v, .little),
1776 .index => |v| std.mem.writeInt(i64, memref.bytes[offset..][0..8], v, .little),
1777 .f16 => |v| std.mem.writeInt(u16, memref.bytes[offset..][0..2], @bitCast(v), .little),
1778 .bf16 => |v| std.mem.writeInt(u16, memref.bytes[offset..][0..2], v.bits, .little),
1779 .f32 => |v| std.mem.writeInt(u32, memref.bytes[offset..][0..4], @bitCast(v), .little),
1780 .f64 => |v| std.mem.writeInt(u64, memref.bytes[offset..][0..8], @bitCast(v), .little),
1781 }
1782 }
1783
1784 fn atomicApply(kind: dialects.AtomicRmwKind, old: Scalar, operand: Scalar) EvalError!Scalar {
1785 if (std.meta.activeTag(old) != std.meta.activeTag(operand)) return error.UnsupportedType;
1786 return switch (kind) {
1787 .add => switch (old) {
1788 .i8 => |v| .{ .i8 = v +% operand.i8 },
1789 .i16 => |v| .{ .i16 = v +% operand.i16 },
1790 .i32 => |v| .{ .i32 = v +% operand.i32 },
1791 .u8 => |v| .{ .u8 = v +% operand.u8 },
1792 .u16 => |v| .{ .u16 = v +% operand.u16 },
1793 .u32 => |v| .{ .u32 = v +% operand.u32 },
1794 .i64 => |v| .{ .i64 = v +% operand.i64 },
1795 .u64 => |v| .{ .u64 = v +% operand.u64 },
1796 .f32 => |v| .{ .f32 = v + operand.f32 },
1797 .f64 => |v| .{ .f64 = v + operand.f64 },
1798 .bool, .index, .f16, .bf16 => error.UnsupportedType,
1799 },
1800 .min => switch (old) {
1801 .i8 => |v| .{ .i8 = @min(v, operand.i8) },
1802 .i16 => |v| .{ .i16 = @min(v, operand.i16) },
1803 .i32 => |v| .{ .i32 = @min(v, operand.i32) },
1804 .u8 => |v| .{ .u8 = @min(v, operand.u8) },
1805 .u16 => |v| .{ .u16 = @min(v, operand.u16) },
1806 .u32 => |v| .{ .u32 = @min(v, operand.u32) },
1807 .i64 => |v| .{ .i64 = @min(v, operand.i64) },
1808 .u64 => |v| .{ .u64 = @min(v, operand.u64) },
1809 else => error.UnsupportedType,
1810 },
1811 .max => switch (old) {
1812 .i8 => |v| .{ .i8 = @max(v, operand.i8) },
1813 .i16 => |v| .{ .i16 = @max(v, operand.i16) },
1814 .i32 => |v| .{ .i32 = @max(v, operand.i32) },
1815 .u8 => |v| .{ .u8 = @max(v, operand.u8) },
1816 .u16 => |v| .{ .u16 = @max(v, operand.u16) },
1817 .u32 => |v| .{ .u32 = @max(v, operand.u32) },
1818 .i64 => |v| .{ .i64 = @max(v, operand.i64) },
1819 .u64 => |v| .{ .u64 = @max(v, operand.u64) },
1820 else => error.UnsupportedType,
1821 },
1822 .bit_and => switch (old) {
1823 .i8 => |v| .{ .i8 = v & operand.i8 },
1824 .i16 => |v| .{ .i16 = v & operand.i16 },
1825 .i32 => |v| .{ .i32 = v & operand.i32 },
1826 .u8 => |v| .{ .u8 = v & operand.u8 },
1827 .u16 => |v| .{ .u16 = v & operand.u16 },
1828 .u32 => |v| .{ .u32 = v & operand.u32 },
1829 .i64 => |v| .{ .i64 = v & operand.i64 },
1830 .u64 => |v| .{ .u64 = v & operand.u64 },
1831 else => error.UnsupportedType,
1832 },
1833 .bit_or => switch (old) {
1834 .i8 => |v| .{ .i8 = v | operand.i8 },
1835 .i16 => |v| .{ .i16 = v | operand.i16 },
1836 .i32 => |v| .{ .i32 = v | operand.i32 },
1837 .u8 => |v| .{ .u8 = v | operand.u8 },
1838 .u16 => |v| .{ .u16 = v | operand.u16 },
1839 .u32 => |v| .{ .u32 = v | operand.u32 },
1840 .i64 => |v| .{ .i64 = v | operand.i64 },
1841 .u64 => |v| .{ .u64 = v | operand.u64 },
1842 else => error.UnsupportedType,
1843 },
1844 .bit_xor => switch (old) {
1845 .i8 => |v| .{ .i8 = v ^ operand.i8 },
1846 .i16 => |v| .{ .i16 = v ^ operand.i16 },
1847 .i32 => |v| .{ .i32 = v ^ operand.i32 },
1848 .u8 => |v| .{ .u8 = v ^ operand.u8 },
1849 .u16 => |v| .{ .u16 = v ^ operand.u16 },
1850 .u32 => |v| .{ .u32 = v ^ operand.u32 },
1851 .i64 => |v| .{ .i64 = v ^ operand.i64 },
1852 .u64 => |v| .{ .u64 = v ^ operand.u64 },
1853 else => error.UnsupportedType,
1854 },
1855 .exchange => switch (old) {
1856 .bool => error.UnsupportedType,
1857 else => operand,
1858 },
1859 };
1860 }
1861
1862 fn checkedElementIndex(memref: Memref, index_i64: i64) EvalError!usize {
1863 if (index_i64 < 0) return error.InvalidIndex;
1864 const index = std.math.cast(usize, index_i64) orelse return error.InvalidIndex;
1865 const element_size = @as(usize, @intCast(scalarByteSize(memref.element)));
1866 if (element_size == 0 or memref.bytes.len % element_size != 0) return error.InvalidMemref;
1867 if (index >= memref.bytes.len / element_size) return error.InvalidIndex;
1868 return index;
1869 }
1870
1871 fn evalScalarBinary(kind: ScalarKind, op: Evaluator.BinaryOp, lhs: Scalar, rhs: Scalar) EvalError!Scalar {
1872 return switch (kind) {
1873 .bool => error.UnsupportedType,
1874 .index => .{ .index = try evalIndexBinary(op, try scalarToI64(lhs), try scalarToI64(rhs)) },
1875 .i8 => .{ .i8 = try evalIntBinary(i8, op, try scalarToI8(lhs), try scalarToI8(rhs)) },
1876 .i16 => .{ .i16 = try evalIntBinary(i16, op, try scalarToI16(lhs), try scalarToI16(rhs)) },
1877 .i32 => .{ .i32 = try evalIntBinary(i32, op, try scalarToI32(lhs), try scalarToI32(rhs)) },
1878 .u8 => .{ .u8 = try evalIntBinary(u8, op, try scalarToU8(lhs), try scalarToU8(rhs)) },
1879 .u16 => .{ .u16 = try evalIntBinary(u16, op, try scalarToU16(lhs), try scalarToU16(rhs)) },
1880 .u32 => .{ .u32 = try evalIntBinary(u32, op, try scalarToU32(lhs), try scalarToU32(rhs)) },
1881 .i64 => .{ .i64 = try evalIntBinary(i64, op, try scalarToI64(lhs), try scalarToI64(rhs)) },
1882 .u64 => .{ .u64 = try evalIntBinary(u64, op, try scalarToU64(lhs), try scalarToU64(rhs)) },
1883 .f16 => .{ .f16 = @floatCast(try evalFloatBinary(f32, op, try scalarToF32(lhs), try scalarToF32(rhs))) },
1884 .bf16 => .{ .bf16 = Bf16.fromF32(try evalFloatBinary(f32, op, try scalarToF32(lhs), try scalarToF32(rhs))) },
1885 .f32 => .{ .f32 = try evalFloatBinary(f32, op, try scalarToF32(lhs), try scalarToF32(rhs)) },
1886 .f64 => .{ .f64 = try evalFloatBinary(f64, op, try scalarToF64(lhs), try scalarToF64(rhs)) },
1887 };
1888 }
1889
1890 fn evalIntBinary(comptime T: type, op: Evaluator.BinaryOp, lhs: T, rhs: T) EvalError!T {
1891 return switch (op) {
1892 .add => lhs +% rhs,
1893 .sub => lhs -% rhs,
1894 .mul => lhs *% rhs,
1895 .div => if (rhs == 0) error.DivisionByZero else @divTrunc(lhs, rhs),
1896 .rem => if (rhs == 0) error.DivisionByZero else @rem(lhs, rhs),
1897 .min => @min(lhs, rhs),
1898 .max => @max(lhs, rhs),
1899 };
1900 }
1901
1902 fn evalIndexBinary(op: Evaluator.BinaryOp, lhs: i64, rhs: i64) EvalError!i64 {
1903 return switch (op) {
1904 .add => std.math.add(i64, lhs, rhs) catch error.Overflow,
1905 .sub => std.math.sub(i64, lhs, rhs) catch error.Overflow,
1906 .mul => std.math.mul(i64, lhs, rhs) catch error.Overflow,
1907 .div => if (rhs == 0) error.DivisionByZero else @divTrunc(lhs, rhs),
1908 .rem => if (rhs == 0) error.DivisionByZero else @rem(lhs, rhs),
1909 .min => @min(lhs, rhs),
1910 .max => @max(lhs, rhs),
1911 };
1912 }
1913
1914 fn evalFloatBinary(comptime T: type, op: Evaluator.BinaryOp, lhs: T, rhs: T) EvalError!T {
1915 return switch (op) {
1916 .add => lhs + rhs,
1917 .sub => lhs - rhs,
1918 .mul => lhs * rhs,
1919 .div => lhs / rhs,
1920 .rem => error.UnsupportedOperation,
1921 .min => @min(lhs, rhs),
1922 .max => @max(lhs, rhs),
1923 };
1924 }
1925
1926 fn evalScalarUnary(kind: ScalarKind, op: Evaluator.UnaryOp, input: Scalar) EvalError!Scalar {
1927 return switch (kind) {
1928 .bool => error.UnsupportedType,
1929 .index => .{ .index = switch (op) {
1930 .neg => try checkedNegSigned(i64, try scalarToI64(input)),
1931 else => return error.UnsupportedType,
1932 } },
1933 .i8 => .{ .i8 = switch (op) {
1934 .neg => try checkedNegSigned(i8, try scalarToI8(input)),
1935 .abs => try checkedAbsSigned(i8, try scalarToI8(input)),
1936 else => return error.UnsupportedType,
1937 } },
1938 .i16 => .{ .i16 = switch (op) {
1939 .neg => try checkedNegSigned(i16, try scalarToI16(input)),
1940 .abs => try checkedAbsSigned(i16, try scalarToI16(input)),
1941 else => return error.UnsupportedType,
1942 } },
1943 .i32 => .{ .i32 = switch (op) {
1944 .neg => try checkedNegSigned(i32, try scalarToI32(input)),
1945 .abs => try checkedAbsSigned(i32, try scalarToI32(input)),
1946 else => return error.UnsupportedType,
1947 } },
1948 .u8 => .{ .u8 = switch (op) {
1949 .neg => 0 -% try scalarToU8(input),
1950 .abs => try scalarToU8(input),
1951 else => return error.UnsupportedType,
1952 } },
1953 .u16 => .{ .u16 = switch (op) {
1954 .neg => 0 -% try scalarToU16(input),
1955 .abs => try scalarToU16(input),
1956 else => return error.UnsupportedType,
1957 } },
1958 .u32 => .{ .u32 = switch (op) {
1959 .neg => 0 -% try scalarToU32(input),
1960 .abs => try scalarToU32(input),
1961 else => return error.UnsupportedType,
1962 } },
1963 .i64 => .{ .i64 = switch (op) {
1964 .neg => try checkedNegSigned(i64, try scalarToI64(input)),
1965 .abs => try checkedAbsSigned(i64, try scalarToI64(input)),
1966 else => return error.UnsupportedType,
1967 } },
1968 .u64 => .{ .u64 = switch (op) {
1969 .neg => 0 -% try scalarToU64(input),
1970 .abs => try scalarToU64(input),
1971 else => return error.UnsupportedType,
1972 } },
1973 .f16 => .{ .f16 = @floatCast(try evalFloatUnary(f32, op, try scalarToF32(input))) },
1974 .bf16 => .{ .bf16 = Bf16.fromF32(try evalFloatUnary(f32, op, try scalarToF32(input))) },
1975 .f32 => .{ .f32 = try evalFloatUnary(f32, op, try scalarToF32(input)) },
1976 .f64 => .{ .f64 = try evalFloatUnary(f64, op, try scalarToF64(input)) },
1977 };
1978 }
1979
1980 fn checkedNegSigned(comptime T: type, value: T) EvalError!T {
1981 if (value == std.math.minInt(T)) return error.Overflow;
1982 return -value;
1983 }
1984
1985 fn checkedAbsSigned(comptime T: type, value: T) EvalError!T {
1986 if (value == std.math.minInt(T)) return error.Overflow;
1987 return if (value < 0) -value else value;
1988 }
1989
1990 fn evalFloatUnary(comptime T: type, op: Evaluator.UnaryOp, input: T) EvalError!T {
1991 return switch (op) {
1992 .neg => -input,
1993 .abs => @abs(input),
1994 .sqrt => @sqrt(input),
1995 .exp => @exp(input),
1996 .log => @log(input),
1997 .tanh => std.math.tanh(input),
1998 .sin => @sin(input),
1999 .cos => @cos(input),
2000 .tan => std.math.tan(input),
2001 .floor => @floor(input),
2002 .round => @round(input),
2003 .trunc => @trunc(input),
2004 .tf32_round => if (T == f32) tf32RoundF32(input) else error.UnsupportedType,
2005 };
2006 }
2007
2008 pub fn tf32RoundF32(value: f32) f32 {
2009 const bits: u32 = @bitCast(value);
2010 if (bits & 0x7F80_0000 == 0x7F80_0000) return value;
2011 const rounded = (bits +% 0x1000) & 0xFFFF_E000;
2012 return @bitCast(rounded);
2013 }
2014
2015 test "tf32RoundF32 truncates mantissa to tf32 grid and preserves specials" {
2016 try std.testing.expectEqual(@as(f32, 1.0), tf32RoundF32(1.0));
2017 try std.testing.expectEqual(@as(f32, 0.0), tf32RoundF32(0.0));
2018 try std.testing.expectEqual(@as(f32, -2.5), tf32RoundF32(-2.5));
2019
2020 const rounded = tf32RoundF32(1.00006103515625 + 0.00001);
2021 const rounded_bits: u32 = @bitCast(rounded);
2022 try std.testing.expectEqual(@as(u32, 0), rounded_bits & 0x1FFF);
2023
2024 const inf = std.math.inf(f32);
2025 try std.testing.expectEqual(inf, tf32RoundF32(inf));
2026 try std.testing.expect(std.math.isNan(tf32RoundF32(std.math.nan(f32))));
2027
2028 const third = tf32RoundF32(1.0 / 3.0);
2029 try std.testing.expectApproxEqAbs(@as(f32, 1.0 / 3.0), third, 0.0002);
2030 const third_bits: u32 = @bitCast(third);
2031 try std.testing.expectEqual(@as(u32, 0), third_bits & 0x1FFF);
2032 }
2033
2034 fn evalScalarPow(kind: ScalarKind, base: Scalar, exponent: Scalar) EvalError!Scalar {
2035 return switch (kind) {
2036 .f16 => .{ .f16 = @floatCast(std.math.pow(f32, try scalarToF32(base), try scalarToF32(exponent))) },
2037 .bf16 => .{ .bf16 = Bf16.fromF32(std.math.pow(f32, try scalarToF32(base), try scalarToF32(exponent))) },
2038 .f32 => .{ .f32 = std.math.pow(f32, try scalarToF32(base), try scalarToF32(exponent)) },
2039 .f64 => .{ .f64 = std.math.pow(f64, try scalarToF64(base), try scalarToF64(exponent)) },
2040 else => error.UnsupportedType,
2041 };
2042 }
2043
2044 fn evalScalarAtan2(kind: ScalarKind, y: Scalar, x: Scalar) EvalError!Scalar {
2045 return switch (kind) {
2046 .f16 => .{ .f16 = @floatCast(std.math.atan2(try scalarToF32(y), try scalarToF32(x))) },
2047 .bf16 => .{ .bf16 = Bf16.fromF32(std.math.atan2(try scalarToF32(y), try scalarToF32(x))) },
2048 .f32 => .{ .f32 = std.math.atan2(try scalarToF32(y), try scalarToF32(x)) },
2049 .f64 => .{ .f64 = std.math.atan2(try scalarToF64(y), try scalarToF64(x)) },
2050 else => error.UnsupportedType,
2051 };
2052 }
2053
2054 fn evalScalarFma(kind: ScalarKind, a: Scalar, b: Scalar, c: Scalar) EvalError!Scalar {
2055 return switch (kind) {
2056 .f16 => .{ .f16 = @floatCast(@mulAdd(f32, try scalarToF32(a), try scalarToF32(b), try scalarToF32(c))) },
2057 .bf16 => .{ .bf16 = Bf16.fromF32(@mulAdd(f32, try scalarToF32(a), try scalarToF32(b), try scalarToF32(c))) },
2058 .f32 => .{ .f32 = @mulAdd(f32, try scalarToF32(a), try scalarToF32(b), try scalarToF32(c)) },
2059 .f64 => .{ .f64 = @mulAdd(f64, try scalarToF64(a), try scalarToF64(b), try scalarToF64(c)) },
2060 else => error.UnsupportedType,
2061 };
2062 }
2063
2064 const Predicate = enum { eq, ne, lt, le, gt, ge, slt, sle, sgt, sge, ult, ule, ugt, uge };
2065
2066 fn cmpPredicate(op: *ir.Operation) EvalError!Predicate {
2067 const attr = op.getAttr("predicate") orelse return error.MissingAttribute;
2068 if (!std.mem.eql(u8, attr.abstract.name, "arith.predicate")) return error.InvalidAttribute;
2069 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return error.InvalidAttribute;
2070 inline for (
2071 @typeInfo(Predicate).@"enum".field_names,
2072 @typeInfo(Predicate).@"enum".field_values,
2073 ) |field_name, field_name_value| {
2074 const field = .{ .name = field_name, .value = field_name_value };
2075 if (std.mem.eql(u8, dialect_attr.payload, field.name)) return @fromBackingInt(@intCast(field.value));
2076 }
2077 return error.InvalidAttribute;
2078 }
2079
2080 fn evalCmpPredicate(predicate: Predicate, lhs: Scalar, rhs: Scalar) EvalError!bool {
2081 return switch (lhs) {
2082 .f16, .bf16, .f32, .f64 => evalFloatCmp(predicate, try scalarToF64(lhs), try scalarToF64(rhs)),
2083 .u8 => evalUnsignedCmp(u8, predicate, lhs.u8, try scalarToU8(rhs)),
2084 .u16 => evalUnsignedCmp(u16, predicate, lhs.u16, try scalarToU16(rhs)),
2085 .u32 => evalUnsignedCmp(u32, predicate, lhs.u32, try scalarToU32(rhs)),
2086 .u64 => evalUnsignedCmp(u64, predicate, lhs.u64, try scalarToU64(rhs)),
2087 else => evalIntCmp(predicate, try scalarToI64(lhs), try scalarToI64(rhs)),
2088 };
2089 }
2090
2091 fn evalIntCmp(predicate: Predicate, lhs: i64, rhs: i64) bool {
2092 return switch (predicate) {
2093 .eq => lhs == rhs,
2094 .ne => lhs != rhs,
2095 .lt, .slt => lhs < rhs,
2096 .le, .sle => lhs <= rhs,
2097 .gt, .sgt => lhs > rhs,
2098 .ge, .sge => lhs >= rhs,
2099 .ult => @as(u64, @bitCast(lhs)) < @as(u64, @bitCast(rhs)),
2100 .ule => @as(u64, @bitCast(lhs)) <= @as(u64, @bitCast(rhs)),
2101 .ugt => @as(u64, @bitCast(lhs)) > @as(u64, @bitCast(rhs)),
2102 .uge => @as(u64, @bitCast(lhs)) >= @as(u64, @bitCast(rhs)),
2103 };
2104 }
2105
2106 fn evalUnsignedCmp(comptime T: type, predicate: Predicate, lhs: T, rhs: T) bool {
2107 return switch (predicate) {
2108 .eq => lhs == rhs,
2109 .ne => lhs != rhs,
2110 .lt, .slt, .ult => lhs < rhs,
2111 .le, .sle, .ule => lhs <= rhs,
2112 .gt, .sgt, .ugt => lhs > rhs,
2113 .ge, .sge, .uge => lhs >= rhs,
2114 };
2115 }
2116
2117 fn evalFloatCmp(predicate: Predicate, lhs: f64, rhs: f64) EvalError!bool {
2118 return switch (predicate) {
2119 .eq => lhs == rhs,
2120 .ne => lhs != rhs,
2121 .lt, .slt, .ult => lhs < rhs,
2122 .le, .sle, .ule => lhs <= rhs,
2123 .gt, .sgt, .ugt => lhs > rhs,
2124 .ge, .sge, .uge => lhs >= rhs,
2125 };
2126 }
2127
2128 fn evalScalarLogical(kind: ScalarKind, op: Evaluator.LogicalOp, lhs: Scalar, rhs: Scalar) EvalError!Scalar {
2129 return switch (kind) {
2130 .bool => .{ .bool = switch (op) {
2131 .and_ => (try scalarToBool(lhs)) and (try scalarToBool(rhs)),
2132 .or_ => (try scalarToBool(lhs)) or (try scalarToBool(rhs)),
2133 .xor => (try scalarToBool(lhs)) != (try scalarToBool(rhs)),
2134 } },
2135 .index => .{ .index = switch (op) {
2136 .and_ => (try scalarToI64(lhs)) & (try scalarToI64(rhs)),
2137 .or_ => (try scalarToI64(lhs)) | (try scalarToI64(rhs)),
2138 .xor => (try scalarToI64(lhs)) ^ (try scalarToI64(rhs)),
2139 } },
2140 .i8 => .{ .i8 = switch (op) {
2141 .and_ => (try scalarToI8(lhs)) & (try scalarToI8(rhs)),
2142 .or_ => (try scalarToI8(lhs)) | (try scalarToI8(rhs)),
2143 .xor => (try scalarToI8(lhs)) ^ (try scalarToI8(rhs)),
2144 } },
2145 .i16 => .{ .i16 = switch (op) {
2146 .and_ => (try scalarToI16(lhs)) & (try scalarToI16(rhs)),
2147 .or_ => (try scalarToI16(lhs)) | (try scalarToI16(rhs)),
2148 .xor => (try scalarToI16(lhs)) ^ (try scalarToI16(rhs)),
2149 } },
2150 .i32 => .{ .i32 = switch (op) {
2151 .and_ => (try scalarToI32(lhs)) & (try scalarToI32(rhs)),
2152 .or_ => (try scalarToI32(lhs)) | (try scalarToI32(rhs)),
2153 .xor => (try scalarToI32(lhs)) ^ (try scalarToI32(rhs)),
2154 } },
2155 .u8 => .{ .u8 = switch (op) {
2156 .and_ => (try scalarToU8(lhs)) & (try scalarToU8(rhs)),
2157 .or_ => (try scalarToU8(lhs)) | (try scalarToU8(rhs)),
2158 .xor => (try scalarToU8(lhs)) ^ (try scalarToU8(rhs)),
2159 } },
2160 .u16 => .{ .u16 = switch (op) {
2161 .and_ => (try scalarToU16(lhs)) & (try scalarToU16(rhs)),
2162 .or_ => (try scalarToU16(lhs)) | (try scalarToU16(rhs)),
2163 .xor => (try scalarToU16(lhs)) ^ (try scalarToU16(rhs)),
2164 } },
2165 .u32 => .{ .u32 = switch (op) {
2166 .and_ => (try scalarToU32(lhs)) & (try scalarToU32(rhs)),
2167 .or_ => (try scalarToU32(lhs)) | (try scalarToU32(rhs)),
2168 .xor => (try scalarToU32(lhs)) ^ (try scalarToU32(rhs)),
2169 } },
2170 .i64 => .{ .i64 = switch (op) {
2171 .and_ => (try scalarToI64(lhs)) & (try scalarToI64(rhs)),
2172 .or_ => (try scalarToI64(lhs)) | (try scalarToI64(rhs)),
2173 .xor => (try scalarToI64(lhs)) ^ (try scalarToI64(rhs)),
2174 } },
2175 .u64 => .{ .u64 = switch (op) {
2176 .and_ => (try scalarToU64(lhs)) & (try scalarToU64(rhs)),
2177 .or_ => (try scalarToU64(lhs)) | (try scalarToU64(rhs)),
2178 .xor => (try scalarToU64(lhs)) ^ (try scalarToU64(rhs)),
2179 } },
2180 .f16, .bf16, .f32, .f64 => error.UnsupportedType,
2181 };
2182 }
2183
2184 fn evalWarpScalar(kind: ScalarKind, op_kind: WarpOpKind, lhs: Scalar, rhs: Scalar) EvalError!Scalar {
2185 return switch (op_kind) {
2186 .add => evalScalarBinary(kind, .add, lhs, rhs),
2187 .max => evalScalarBinary(kind, .max, lhs, rhs),
2188 .min => evalScalarBinary(kind, .min, lhs, rhs),
2189 .and_ => evalScalarLogical(kind, .and_, lhs, rhs),
2190 .or_ => evalScalarLogical(kind, .or_, lhs, rhs),
2191 .xor => evalScalarLogical(kind, .xor, lhs, rhs),
2192 };
2193 }
2194
2195 fn warpIdentity(kind: ScalarKind, op_kind: WarpOpKind) EvalError!Scalar {
2196 return switch (op_kind) {
2197 .add, .or_, .xor => switch (kind) {
2198 .bool => if (op_kind == .add) error.UnsupportedType else .{ .bool = false },
2199 .index => .{ .index = 0 },
2200 .i8 => .{ .i8 = 0 },
2201 .i16 => .{ .i16 = 0 },
2202 .i32 => .{ .i32 = 0 },
2203 .u8 => .{ .u8 = 0 },
2204 .u16 => .{ .u16 = 0 },
2205 .u32 => .{ .u32 = 0 },
2206 .i64 => .{ .i64 = 0 },
2207 .u64 => .{ .u64 = 0 },
2208 .f16 => if (op_kind == .add) .{ .f16 = 0 } else error.UnsupportedType,
2209 .bf16 => if (op_kind == .add) .{ .bf16 = Bf16.fromF32(0) } else error.UnsupportedType,
2210 .f32 => if (op_kind == .add) .{ .f32 = 0 } else error.UnsupportedType,
2211 .f64 => if (op_kind == .add) .{ .f64 = 0 } else error.UnsupportedType,
2212 },
2213 .and_ => switch (kind) {
2214 .bool => .{ .bool = true },
2215 .index => .{ .index = -1 },
2216 .i8 => .{ .i8 = -1 },
2217 .i16 => .{ .i16 = -1 },
2218 .i32 => .{ .i32 = -1 },
2219 .u8 => .{ .u8 = std.math.maxInt(u8) },
2220 .u16 => .{ .u16 = std.math.maxInt(u16) },
2221 .u32 => .{ .u32 = std.math.maxInt(u32) },
2222 .i64 => .{ .i64 = -1 },
2223 .u64 => .{ .u64 = std.math.maxInt(u64) },
2224 .f16, .bf16, .f32, .f64 => error.UnsupportedType,
2225 },
2226 .min => switch (kind) {
2227 .bool => error.UnsupportedType,
2228 .index => .{ .index = std.math.maxInt(i64) },
2229 .i8 => .{ .i8 = std.math.maxInt(i8) },
2230 .i16 => .{ .i16 = std.math.maxInt(i16) },
2231 .i32 => .{ .i32 = std.math.maxInt(i32) },
2232 .u8 => .{ .u8 = std.math.maxInt(u8) },
2233 .u16 => .{ .u16 = std.math.maxInt(u16) },
2234 .u32 => .{ .u32 = std.math.maxInt(u32) },
2235 .i64 => .{ .i64 = std.math.maxInt(i64) },
2236 .u64 => .{ .u64 = std.math.maxInt(u64) },
2237 .f16 => .{ .f16 = std.math.inf(f16) },
2238 .bf16 => .{ .bf16 = Bf16.fromF32(std.math.inf(f32)) },
2239 .f32 => .{ .f32 = std.math.inf(f32) },
2240 .f64 => .{ .f64 = std.math.inf(f64) },
2241 },
2242 .max => switch (kind) {
2243 .bool => error.UnsupportedType,
2244 .index => .{ .index = std.math.minInt(i64) },
2245 .i8 => .{ .i8 = std.math.minInt(i8) },
2246 .i16 => .{ .i16 = std.math.minInt(i16) },
2247 .i32 => .{ .i32 = std.math.minInt(i32) },
2248 .u8 => .{ .u8 = 0 },
2249 .u16 => .{ .u16 = 0 },
2250 .u32 => .{ .u32 = 0 },
2251 .i64 => .{ .i64 = std.math.minInt(i64) },
2252 .u64 => .{ .u64 = 0 },
2253 .f16 => .{ .f16 = -std.math.inf(f16) },
2254 .bf16 => .{ .bf16 = Bf16.fromF32(-std.math.inf(f32)) },
2255 .f32 => .{ .f32 = -std.math.inf(f32) },
2256 .f64 => .{ .f64 = -std.math.inf(f64) },
2257 },
2258 };
2259 }
2260
2261 fn shuffleSourceLane(mode: ShuffleMode, target_lane_id: usize, lane_or_delta: i64) EvalError!usize {
2262 if (lane_or_delta < 0) return error.InvalidIndex;
2263 const selector = std.math.cast(usize, lane_or_delta) orelse return error.InvalidIndex;
2264 const source = switch (mode) {
2265 .sync => selector,
2266 .down => std.math.add(usize, target_lane_id, selector) catch return error.Overflow,
2267 .up => if (target_lane_id >= selector) target_lane_id - selector else return error.InvalidIndex,
2268 .xor => target_lane_id ^ selector,
2269 };
2270 if (source >= warp_size) return error.InvalidIndex;
2271 return source;
2272 }
2273
2274 fn evalScalarNot(kind: ScalarKind, input: Scalar) EvalError!Scalar {
2275 return switch (kind) {
2276 .bool => .{ .bool = !(try scalarToBool(input)) },
2277 .index => .{ .index = ~(try scalarToI64(input)) },
2278 .i8 => .{ .i8 = ~(try scalarToI8(input)) },
2279 .i16 => .{ .i16 = ~(try scalarToI16(input)) },
2280 .i32 => .{ .i32 = ~(try scalarToI32(input)) },
2281 .u8 => .{ .u8 = ~(try scalarToU8(input)) },
2282 .u16 => .{ .u16 = ~(try scalarToU16(input)) },
2283 .u32 => .{ .u32 = ~(try scalarToU32(input)) },
2284 .i64 => .{ .i64 = ~(try scalarToI64(input)) },
2285 .u64 => .{ .u64 = ~(try scalarToU64(input)) },
2286 .f16, .bf16, .f32, .f64 => error.UnsupportedType,
2287 };
2288 }
2289
2290 fn evalScalarPopCount(kind: ScalarKind, input: Scalar) EvalError!Scalar {
2291 return switch (kind) {
2292 .i8 => .{ .i8 = @intCast(@popCount(@as(u8, @bitCast(try scalarToI8(input))))) },
2293 .u8 => .{ .u8 = @intCast(@popCount(try scalarToU8(input))) },
2294 .i16 => .{ .i16 = @intCast(@popCount(@as(u16, @bitCast(try scalarToI16(input))))) },
2295 .u16 => .{ .u16 = @intCast(@popCount(try scalarToU16(input))) },
2296 .i32 => .{ .i32 = @popCount(@as(u32, @bitCast(try scalarToI32(input)))) },
2297 .u32 => .{ .u32 = @popCount(try scalarToU32(input)) },
2298 .i64 => .{ .i64 = @popCount(@as(u64, @bitCast(try scalarToI64(input)))) },
2299 .u64 => .{ .u64 = @popCount(try scalarToU64(input)) },
2300 .index => .{ .index = @popCount(@as(u64, @bitCast(try scalarToI64(input)))) },
2301 .bool, .f16, .bf16, .f32, .f64 => error.UnsupportedType,
2302 };
2303 }
2304
2305 fn evalScalarUmulhi(kind: ScalarKind, lhs: Scalar, rhs: Scalar) EvalError!Scalar {
2306 return switch (kind) {
2307 .i8 => .{ .i8 = umulhiInt(i8, try scalarToI8(lhs), try scalarToI8(rhs)) },
2308 .u8 => .{ .u8 = umulhiInt(u8, try scalarToU8(lhs), try scalarToU8(rhs)) },
2309 .i16 => .{ .i16 = umulhiInt(i16, try scalarToI16(lhs), try scalarToI16(rhs)) },
2310 .u16 => .{ .u16 = umulhiInt(u16, try scalarToU16(lhs), try scalarToU16(rhs)) },
2311 .i32 => .{ .i32 = umulhiInt(i32, try scalarToI32(lhs), try scalarToI32(rhs)) },
2312 .u32 => .{ .u32 = umulhiInt(u32, try scalarToU32(lhs), try scalarToU32(rhs)) },
2313 .i64 => .{ .i64 = umulhiInt(i64, try scalarToI64(lhs), try scalarToI64(rhs)) },
2314 .u64 => .{ .u64 = umulhiInt(u64, try scalarToU64(lhs), try scalarToU64(rhs)) },
2315 .bool, .index, .f16, .bf16, .f32, .f64 => error.UnsupportedType,
2316 };
2317 }
2318
2319 fn umulhiInt(comptime T: type, lhs: T, rhs: T) T {
2320 const bit_count = @bitSizeOf(T);
2321 const U = @Int(.unsigned, bit_count);
2322 const W = @Int(.unsigned, bit_count * 2);
2323 const product = @as(W, @as(U, @bitCast(lhs))) * @as(W, @as(U, @bitCast(rhs)));
2324 return @bitCast(@as(U, @truncate(product >> bit_count)));
2325 }
2326
2327 fn evalScalarShift(kind: ScalarKind, op: Evaluator.ShiftOp, value: Scalar, shift: Scalar) EvalError!Scalar {
2328 return switch (kind) {
2329 .i8 => .{ .i8 = try evalSignedShift(i8, op, try scalarToI8(value), try scalarToI8(shift)) },
2330 .u8 => .{ .u8 = try evalUnsignedShift(u8, op, try scalarToU8(value), try scalarToU8(shift)) },
2331 .i16 => .{ .i16 = try evalSignedShift(i16, op, try scalarToI16(value), try scalarToI16(shift)) },
2332 .u16 => .{ .u16 = try evalUnsignedShift(u16, op, try scalarToU16(value), try scalarToU16(shift)) },
2333 .i32 => .{ .i32 = try evalSignedShift(i32, op, try scalarToI32(value), try scalarToI32(shift)) },
2334 .u32 => .{ .u32 = try evalUnsignedShift(u32, op, try scalarToU32(value), try scalarToU32(shift)) },
2335 .i64 => .{ .i64 = try evalSignedShift(i64, op, try scalarToI64(value), try scalarToI64(shift)) },
2336 .u64 => .{ .u64 = try evalUnsignedShift(u64, op, try scalarToU64(value), try scalarToU64(shift)) },
2337 .bool, .index, .f16, .bf16, .f32, .f64 => error.UnsupportedType,
2338 };
2339 }
2340
2341 fn evalSignedShift(comptime T: type, op: Evaluator.ShiftOp, value: T, shift: T) EvalError!T {
2342 if (shift < 0) return error.UnsupportedOperation;
2343 const bit_count = @bitSizeOf(T);
2344 if (@as(u64, @intCast(shift)) >= bit_count) return error.UnsupportedOperation;
2345 const shift_count: std.math.Log2Int(T) = @intCast(shift);
2346 const U = @Int(.unsigned, bit_count);
2347 return switch (op) {
2348 .shl => @bitCast(@as(U, @bitCast(value)) << shift_count),
2349 .shr => value >> shift_count,
2350 .ushr => @bitCast(@as(U, @bitCast(value)) >> shift_count),
2351 };
2352 }
2353
2354 fn evalUnsignedShift(comptime T: type, op: Evaluator.ShiftOp, value: T, shift: T) EvalError!T {
2355 const bit_count = @bitSizeOf(T);
2356 if (@as(u64, @intCast(shift)) >= bit_count) return error.UnsupportedOperation;
2357 const shift_count: std.math.Log2Int(T) = @intCast(shift);
2358 return switch (op) {
2359 .shl => value << shift_count,
2360 .shr, .ushr => value >> shift_count,
2361 };
2362 }
2363
2364 fn evalScalarBitcast(target: ScalarKind, input: Scalar) EvalError!Scalar {
2365 return switch (target) {
2366 .i8 => switch (input) {
2367 .i8 => input,
2368 .u8 => |v| .{ .i8 = @bitCast(v) },
2369 else => error.UnsupportedType,
2370 },
2371 .u8 => switch (input) {
2372 .u8 => input,
2373 .i8 => |v| .{ .u8 = @bitCast(v) },
2374 else => error.UnsupportedType,
2375 },
2376 .i16 => switch (input) {
2377 .i16 => input,
2378 .u16 => |v| .{ .i16 = @bitCast(v) },
2379 else => error.UnsupportedType,
2380 },
2381 .u16 => switch (input) {
2382 .u16 => input,
2383 .i16 => |v| .{ .u16 = @bitCast(v) },
2384 else => error.UnsupportedType,
2385 },
2386 .i32 => switch (input) {
2387 .i32 => input,
2388 .u32 => |v| .{ .i32 = @bitCast(v) },
2389 .f32 => |v| .{ .i32 = @bitCast(v) },
2390 else => error.UnsupportedType,
2391 },
2392 .u32 => switch (input) {
2393 .u32 => input,
2394 .i32 => |v| .{ .u32 = @bitCast(v) },
2395 .f32 => |v| .{ .u32 = @bitCast(v) },
2396 else => error.UnsupportedType,
2397 },
2398 .i64 => switch (input) {
2399 .i64 => input,
2400 .u64 => |v| .{ .i64 = @bitCast(v) },
2401 .f64 => |v| .{ .i64 = @bitCast(v) },
2402 else => error.UnsupportedType,
2403 },
2404 .u64 => switch (input) {
2405 .u64 => input,
2406 .i64 => |v| .{ .u64 = @bitCast(v) },
2407 .f64 => |v| .{ .u64 = @bitCast(v) },
2408 else => error.UnsupportedType,
2409 },
2410 .f32 => switch (input) {
2411 .f32 => input,
2412 .i32 => |v| .{ .f32 = @bitCast(v) },
2413 .u32 => |v| .{ .f32 = @bitCast(v) },
2414 else => error.UnsupportedType,
2415 },
2416 .f64 => switch (input) {
2417 .f64 => input,
2418 .i64 => |v| .{ .f64 = @bitCast(v) },
2419 .u64 => |v| .{ .f64 = @bitCast(v) },
2420 else => error.UnsupportedType,
2421 },
2422 .f16 => switch (input) {
2423 .f16 => input,
2424 else => error.UnsupportedType,
2425 },
2426 .bf16 => switch (input) {
2427 .bf16 => input,
2428 else => error.UnsupportedType,
2429 },
2430 .bool, .index => error.UnsupportedType,
2431 };
2432 }
2433
2434 fn castScalar(kind: ScalarKind, input: Scalar) EvalError!Scalar {
2435 return switch (kind) {
2436 .bool => .{ .bool = try scalarToBool(input) },
2437 .index => .{ .index = try scalarToI64(input) },
2438 .i8 => .{ .i8 = try scalarToI8(input) },
2439 .i16 => .{ .i16 = try scalarToI16(input) },
2440 .i32 => .{ .i32 = switch (input) {
2441 .u32 => |wide| @bitCast(wide),
2442 .i64 => |wide| @truncate(wide),
2443 .index => |wide| @truncate(wide),
2444 else => try scalarToI32(input),
2445 } },
2446 .u8 => .{ .u8 = try scalarToU8(input) },
2447 .u16 => .{ .u16 = try scalarToU16(input) },
2448 .u32 => .{ .u32 = switch (input) {
2449 .i32 => |wide| @bitCast(wide),
2450 .i64 => |wide| @truncate(@as(u64, @bitCast(wide))),
2451 .index => |wide| @truncate(@as(u64, @bitCast(wide))),
2452 else => try scalarToU32(input),
2453 } },
2454 .i64 => .{ .i64 = switch (input) {
2455 .u64 => |wide| @bitCast(wide),
2456 else => try scalarToI64(input),
2457 } },
2458 .u64 => .{ .u64 = switch (input) {
2459 .i32 => |wide| @bitCast(@as(i64, wide)),
2460 .i64 => |wide| @bitCast(wide),
2461 .index => |wide| @bitCast(wide),
2462 else => try scalarToU64(input),
2463 } },
2464 .f16 => .{ .f16 = try scalarToF16(input) },
2465 .bf16 => .{ .bf16 = Bf16.fromF32(try scalarToF32(input)) },
2466 .f32 => .{ .f32 = try scalarToF32(input) },
2467 .f64 => .{ .f64 = try scalarToF64(input) },
2468 };
2469 }
2470
2471 fn scalarToMask(input: Scalar) EvalError!u32 {
2472 return switch (input) {
2473 .i8 => |v| @bitCast(@as(i32, v)),
2474 .u8 => |v| v,
2475 .i16 => |v| @bitCast(@as(i32, v)),
2476 .u16 => |v| v,
2477 .i32 => |v| @bitCast(v),
2478 .u32 => |v| v,
2479 .u64 => |v| std.math.cast(u32, v) orelse return error.Overflow,
2480 .i64 => |v| std.math.cast(u32, v) orelse return error.Overflow,
2481 .index => |v| std.math.cast(u32, v) orelse return error.Overflow,
2482 else => error.UnsupportedType,
2483 };
2484 }
2485
2486 fn scalarToBool(input: Scalar) EvalError!bool {
2487 return switch (input) {
2488 .bool => |v| v,
2489 else => error.UnsupportedType,
2490 };
2491 }
2492
2493 fn scalarToI8(input: Scalar) EvalError!i8 {
2494 return switch (input) {
2495 .i8 => |v| v,
2496 .u8 => |v| @bitCast(v),
2497 .i16 => |v| @truncate(v),
2498 .u16 => |v| @bitCast(@as(u8, @truncate(v))),
2499 .i32 => |v| @truncate(v),
2500 .u32 => |v| @bitCast(@as(u8, @truncate(v))),
2501 .i64 => |v| @truncate(v),
2502 .u64 => |v| @bitCast(@as(u8, @truncate(v))),
2503 .index => |v| @truncate(v),
2504 .bool => |v| if (v) 1 else 0,
2505 .f16 => |v| @intFromFloat(v),
2506 .bf16 => |v| @intFromFloat(v.toF32()),
2507 .f32 => |v| @intFromFloat(v),
2508 .f64 => |v| @intFromFloat(v),
2509 };
2510 }
2511
2512 fn scalarToI16(input: Scalar) EvalError!i16 {
2513 return switch (input) {
2514 .i8 => |v| v,
2515 .u8 => |v| v,
2516 .i16 => |v| v,
2517 .u16 => |v| @bitCast(v),
2518 .i32 => |v| @truncate(v),
2519 .u32 => |v| @bitCast(@as(u16, @truncate(v))),
2520 .i64 => |v| @truncate(v),
2521 .u64 => |v| @bitCast(@as(u16, @truncate(v))),
2522 .index => |v| @truncate(v),
2523 .bool => |v| if (v) 1 else 0,
2524 .f16 => |v| @intFromFloat(v),
2525 .bf16 => |v| @intFromFloat(v.toF32()),
2526 .f32 => |v| @intFromFloat(v),
2527 .f64 => |v| @intFromFloat(v),
2528 };
2529 }
2530
2531 fn scalarToI32(input: Scalar) EvalError!i32 {
2532 return std.math.cast(i32, try scalarToI64(input)) orelse error.Overflow;
2533 }
2534
2535 fn scalarToU8(input: Scalar) EvalError!u8 {
2536 return switch (input) {
2537 .i8 => |v| @bitCast(v),
2538 .u8 => |v| v,
2539 .i16 => |v| @truncate(@as(u16, @bitCast(v))),
2540 .u16 => |v| @truncate(v),
2541 .i32 => |v| @truncate(@as(u32, @bitCast(v))),
2542 .u32 => |v| @truncate(v),
2543 .i64 => |v| @truncate(@as(u64, @bitCast(v))),
2544 .u64 => |v| @truncate(v),
2545 .index => |v| @truncate(@as(u64, @bitCast(v))),
2546 .bool => |v| if (v) 1 else 0,
2547 .f16 => |v| @intFromFloat(v),
2548 .bf16 => |v| @intFromFloat(v.toF32()),
2549 .f32 => |v| @intFromFloat(v),
2550 .f64 => |v| @intFromFloat(v),
2551 };
2552 }
2553
2554 fn scalarToU16(input: Scalar) EvalError!u16 {
2555 return switch (input) {
2556 .i8 => |v| @bitCast(@as(i16, v)),
2557 .u8 => |v| v,
2558 .i16 => |v| @bitCast(v),
2559 .u16 => |v| v,
2560 .i32 => |v| @truncate(@as(u32, @bitCast(v))),
2561 .u32 => |v| @truncate(v),
2562 .i64 => |v| @truncate(@as(u64, @bitCast(v))),
2563 .u64 => |v| @truncate(v),
2564 .index => |v| @truncate(@as(u64, @bitCast(v))),
2565 .bool => |v| if (v) 1 else 0,
2566 .f16 => |v| @intFromFloat(v),
2567 .bf16 => |v| @intFromFloat(v.toF32()),
2568 .f32 => |v| @intFromFloat(v),
2569 .f64 => |v| @intFromFloat(v),
2570 };
2571 }
2572
2573 fn scalarToU32(input: Scalar) EvalError!u32 {
2574 return switch (input) {
2575 .i8 => |v| @bitCast(@as(i32, v)),
2576 .u8 => |v| v,
2577 .i16 => |v| @bitCast(@as(i32, v)),
2578 .u16 => |v| v,
2579 .u32 => |v| v,
2580 .i32 => |v| @bitCast(v),
2581 .u64 => |v| std.math.cast(u32, v) orelse error.Overflow,
2582 else => std.math.cast(u32, try scalarToI64(input)) orelse error.Overflow,
2583 };
2584 }
2585
2586 fn scalarToI64(input: Scalar) EvalError!i64 {
2587 return switch (input) {
2588 .bool => |v| if (v) 1 else 0,
2589 .index => |v| v,
2590 .i8 => |v| v,
2591 .i16 => |v| v,
2592 .i32 => |v| v,
2593 .u8 => |v| v,
2594 .u16 => |v| v,
2595 .u32 => |v| v,
2596 .i64 => |v| v,
2597 .u64 => |v| std.math.cast(i64, v) orelse error.Overflow,
2598 .f16 => |v| @intFromFloat(v),
2599 .bf16 => |v| @intFromFloat(v.toF32()),
2600 .f32 => |v| @intFromFloat(v),
2601 .f64 => |v| @intFromFloat(v),
2602 };
2603 }
2604
2605 fn scalarToU64(input: Scalar) EvalError!u64 {
2606 return switch (input) {
2607 .bool => |v| if (v) 1 else 0,
2608 .index => |v| @bitCast(v),
2609 .i8 => |v| @bitCast(@as(i64, v)),
2610 .i16 => |v| @bitCast(@as(i64, v)),
2611 .i32 => |v| @bitCast(@as(i64, v)),
2612 .u8 => |v| v,
2613 .u16 => |v| v,
2614 .u32 => |v| v,
2615 .i64 => |v| @bitCast(v),
2616 .u64 => |v| v,
2617 .f16 => |v| @intFromFloat(v),
2618 .bf16 => |v| @intFromFloat(v.toF32()),
2619 .f32 => |v| @intFromFloat(v),
2620 .f64 => |v| @intFromFloat(v),
2621 };
2622 }
2623
2624 fn scalarToF16(input: Scalar) EvalError!f16 {
2625 return @floatCast(try scalarToF32(input));
2626 }
2627
2628 fn scalarToF32(input: Scalar) EvalError!f32 {
2629 return switch (input) {
2630 .bool => |v| if (v) 1.0 else 0.0,
2631 .index => |v| @floatFromInt(v),
2632 .i8 => |v| @floatFromInt(v),
2633 .i16 => |v| @floatFromInt(v),
2634 .i32 => |v| @floatFromInt(v),
2635 .u8 => |v| @floatFromInt(v),
2636 .u16 => |v| @floatFromInt(v),
2637 .u32 => |v| @floatFromInt(v),
2638 .i64 => |v| @floatFromInt(v),
2639 .u64 => |v| @floatFromInt(v),
2640 .f16 => |v| v,
2641 .bf16 => |v| v.toF32(),
2642 .f32 => |v| v,
2643 .f64 => |v| @floatCast(v),
2644 };
2645 }
2646
2647 fn scalarToF64(input: Scalar) EvalError!f64 {
2648 return switch (input) {
2649 .bool => |v| if (v) 1.0 else 0.0,
2650 .index => |v| @floatFromInt(v),
2651 .i8 => |v| @floatFromInt(v),
2652 .i16 => |v| @floatFromInt(v),
2653 .i32 => |v| @floatFromInt(v),
2654 .u8 => |v| @floatFromInt(v),
2655 .u16 => |v| @floatFromInt(v),
2656 .u32 => |v| @floatFromInt(v),
2657 .i64 => |v| @floatFromInt(v),
2658 .u64 => |v| @floatFromInt(v),
2659 .f16 => |v| v,
2660 .bf16 => |v| v.toF32(),
2661 .f32 => |v| v,
2662 .f64 => |v| v,
2663 };
2664 }