fixing phi management

src/Runtime.zig

@@ -388,18 +388,18 @@ fn clearPhiValues(self: *Self, allocator: std.mem.Allocator) void {
     self.phi_values.clearRetainingCapacity();
 }
 
-pub fn snapshotPhiValues(self: *Self, allocator: std.mem.Allocator) RuntimeError!void {
-    self.clearPhiValues(allocator);
+fn snapshotPhiValue(self: *Self, allocator: std.mem.Allocator, result_id: SpvWord) RuntimeError!void {
+    if (result_id >= self.results.len) return RuntimeError.InvalidSpirV;
 
-    for (self.results, 0..) |*result, result_id| {
-        const value = switch (result.variant orelse continue) {
+    const value = switch (self.results[result_id].variant orelse return) {
         .Constant => |*constant| &constant.value,
-            .FunctionParameter => |*parameter| parameter.value_ptr orelse continue,
-            else => continue,
+        .FunctionParameter => |*parameter| parameter.value_ptr orelse return,
+        else => return,
     };
-        if (std.meta.activeTag(value.*) == .Pointer) continue;
+    if (std.meta.activeTag(value.*) == .Pointer) return;
+
     const snapshot = try value.dupe(allocator);
-        const gop = self.phi_values.getOrPut(allocator, @intCast(result_id)) catch {
+    const gop = self.phi_values.getOrPut(allocator, result_id) catch {
         var tmp = snapshot;
         tmp.deinit(allocator);
         return RuntimeError.OutOfMemory;
@@ -407,6 +407,56 @@ pub fn snapshotPhiValues(self: *Self, allocator: std.mem.Allocator) RuntimeError
     if (gop.found_existing) gop.value_ptr.deinit(allocator);
     gop.value_ptr.* = snapshot;
 }
+
+pub fn snapshotPhiValuesForBranch(self: *Self, allocator: std.mem.Allocator, target_source_location: usize) RuntimeError!void {
+    self.clearPhiValues(allocator);
+
+    const predecessor = self.current_label orelse return;
+
+    var index = target_source_location;
+    if (index >= self.it.buffer.len) return RuntimeError.InvalidSpirV;
+
+    const label_opcode_data = self.it.buffer[index];
+    const label_word_count_with_header = (label_opcode_data & (~spv.SpvOpCodeMask)) >> spv.SpvWordCountShift;
+    if (label_word_count_with_header == 0) return RuntimeError.InvalidSpirV;
+    if ((label_opcode_data & spv.SpvOpCodeMask) != @intFromEnum(spv.SpvOp.Label))
+        return RuntimeError.InvalidSpirV;
+    index += label_word_count_with_header;
+
+    const phi_start = index;
+    var phi_end = phi_start;
+    while (index < self.it.buffer.len) {
+        const opcode_data = self.it.buffer[index];
+        const word_count_with_header = (opcode_data & (~spv.SpvOpCodeMask)) >> spv.SpvWordCountShift;
+        if (word_count_with_header == 0) return RuntimeError.InvalidSpirV;
+        const opcode = opcode_data & spv.SpvOpCodeMask;
+        if (opcode != @intFromEnum(spv.SpvOp.Phi)) break;
+        if (index + word_count_with_header > self.it.buffer.len) return RuntimeError.InvalidSpirV;
+
+        const word_count = word_count_with_header - 1;
+        if (word_count < 2 or ((word_count - 2) % 2) != 0) return RuntimeError.InvalidSpirV;
+
+        index += word_count_with_header;
+        phi_end = index;
+    }
+
+    index = phi_start;
+    while (index < phi_end) {
+        const opcode_data = self.it.buffer[index];
+        const word_count_with_header = (opcode_data & (~spv.SpvOpCodeMask)) >> spv.SpvWordCountShift;
+        var operand_index = index + 3;
+        const operand_end = index + word_count_with_header;
+        while (operand_index < operand_end) : (operand_index += 2) {
+            const value_id = self.it.buffer[operand_index];
+            const parent_label_id = self.it.buffer[operand_index + 1];
+            if (parent_label_id == predecessor) {
+                try self.snapshotPhiValue(allocator, value_id);
+                break;
+            }
+        }
+
+        index += word_count_with_header;
+    }
 }
 
 pub fn getPhiValueSnapshot(self: *Self, id: SpvWord) ?*const Value {
@@ -470,6 +520,8 @@ pub fn clearDerivative(self: *Self, allocator: std.mem.Allocator, result: SpvWor
 }
 
 pub fn copyDerivative(self: *Self, allocator: std.mem.Allocator, dst: SpvWord, src: SpvWord) RuntimeError!void {
+    if (self.derivatives.count() == 0) return;
+
     if (self.derivatives.get(src)) |derivative| {
         try self.setDerivative(allocator, dst, &derivative.dx, &derivative.dy);
     } else {

src/opcodes.zig

@@ -3802,7 +3802,9 @@ fn opAccessChain(allocator: std.mem.Allocator, word_count: SpvWord, rt: *Runtime
                     }
 
                     fn robustF32Pointer(gpa: std.mem.Allocator, ptr: ?*f32, window: ?[]u8, descriptor_backed: bool, backing: ?*Value, owns_backing: bool) RuntimeError!F32Pointer {
-                        if (window != null or !descriptor_backed) return .{ .ptr = ptr orelse return RuntimeError.InvalidSpirV, .backing = backing, .owns_backing = owns_backing };
+                        _ = descriptor_backed;
+                        if (ptr) |p| return .{ .ptr = p, .backing = backing, .owns_backing = owns_backing };
+                        if (window != null) return RuntimeError.OutOfBounds;
 
                         destroyBacking(gpa, backing, owns_backing);
                         const value = gpa.create(Value) catch return RuntimeError.OutOfMemory;
@@ -3811,7 +3813,9 @@ fn opAccessChain(allocator: std.mem.Allocator, word_count: SpvWord, rt: *Runtime
                     }
 
                     fn robustI32Pointer(gpa: std.mem.Allocator, ptr: ?*i32, window: ?[]u8, descriptor_backed: bool, backing: ?*Value, owns_backing: bool) RuntimeError!I32Pointer {
-                        if (window != null or !descriptor_backed) return .{ .ptr = ptr orelse return RuntimeError.InvalidSpirV, .backing = backing, .owns_backing = owns_backing };
+                        _ = descriptor_backed;
+                        if (ptr) |p| return .{ .ptr = p, .backing = backing, .owns_backing = owns_backing };
+                        if (window != null) return RuntimeError.OutOfBounds;
 
                         destroyBacking(gpa, backing, owns_backing);
                         const value = gpa.create(Value) catch return RuntimeError.OutOfMemory;
@@ -3820,7 +3824,9 @@ fn opAccessChain(allocator: std.mem.Allocator, word_count: SpvWord, rt: *Runtime
                     }
 
                     fn robustU32Pointer(gpa: std.mem.Allocator, ptr: ?*u32, window: ?[]u8, descriptor_backed: bool, backing: ?*Value, owns_backing: bool) RuntimeError!U32Pointer {
-                        if (window != null or !descriptor_backed) return .{ .ptr = ptr orelse return RuntimeError.InvalidSpirV, .backing = backing, .owns_backing = owns_backing };
+                        _ = descriptor_backed;
+                        if (ptr) |p| return .{ .ptr = p, .backing = backing, .owns_backing = owns_backing };
+                        if (window != null) return RuntimeError.OutOfBounds;
 
                         destroyBacking(gpa, backing, owns_backing);
                         const value = gpa.create(Value) catch return RuntimeError.OutOfMemory;
@@ -4258,12 +4264,13 @@ fn opBitcast(allocator: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeErro
 
 fn opBranch(allocator: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeError!void {
     const id = try rt.it.next();
-    try rt.snapshotPhiValues(allocator);
-    rt.previous_label = rt.current_label;
-    _ = rt.it.jumpToSourceLocation(switch ((try rt.results[id].getVariant()).*) {
+    const target = switch ((try rt.results[id].getVariant()).*) {
         .Label => |l| l.source_location,
         else => return RuntimeError.InvalidSpirV,
-    });
+    };
+    try rt.snapshotPhiValuesForBranch(allocator, target);
+    rt.previous_label = rt.current_label;
+    _ = rt.it.jumpToSourceLocation(target);
 }
 
 fn opBranchConditional(allocator: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeError!void {
@@ -4276,11 +4283,12 @@ fn opBranchConditional(allocator: std.mem.Allocator, _: SpvWord, rt: *Runtime) R
         .Label => |l| l.source_location,
         else => return RuntimeError.InvalidSpirV,
     };
-    try rt.snapshotPhiValues(allocator);
     rt.previous_label = rt.current_label;
     if (cond_value.Bool) {
+        try rt.snapshotPhiValuesForBranch(allocator, true_branch);
         _ = rt.it.jumpToSourceLocation(true_branch);
     } else {
+        try rt.snapshotPhiValuesForBranch(allocator, false_branch);
         _ = rt.it.jumpToSourceLocation(false_branch);
     }
 }
@@ -5019,7 +5027,7 @@ fn opSpecConstantFalse(allocator: std.mem.Allocator, _: SpvWord, rt: *Runtime) R
     }
 }
 
-fn opSwitch(_: std.mem.Allocator, word_count: SpvWord, rt: *Runtime) RuntimeError!void {
+fn opSwitch(allocator: std.mem.Allocator, word_count: SpvWord, rt: *Runtime) RuntimeError!void {
     if (word_count < 2)
         return RuntimeError.InvalidSpirV;
 
@@ -5062,11 +5070,13 @@ fn opSwitch(_: std.mem.Allocator, word_count: SpvWord, rt: *Runtime) RuntimeErro
         remaining -= literal_width + 1;
     }
 
-    rt.previous_label = rt.current_label;
-    _ = rt.it.jumpToSourceLocation(switch ((try rt.results[target].getVariant()).*) {
+    const target_source_location = switch ((try rt.results[target].getVariant()).*) {
         .Label => |l| l.source_location,
         else => return RuntimeError.InvalidSpirV,
-    });
+    };
+    try rt.snapshotPhiValuesForBranch(allocator, target_source_location);
+    rt.previous_label = rt.current_label;
+    _ = rt.it.jumpToSourceLocation(target_source_location);
 }
 
 fn opSpecConstantOp(allocator: std.mem.Allocator, word_count: SpvWord, rt: *Runtime) RuntimeError!void {
@@ -5557,7 +5567,7 @@ fn opCopyMemory(allocator: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeE
     try rt.copyDerivative(allocator, target, source);
 }
 
-fn opCopyObject(_: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeError!void {
+fn opCopyObject(allocator: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeError!void {
     _ = try rt.it.next(); // result type
     const target = try rt.it.next();
     const source = try rt.it.next();
@@ -5566,9 +5576,11 @@ fn opCopyObject(_: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeError!voi
     if (source_value.* == .Pointer) {
         target_value.* = source_value.*;
         target_value.Pointer.owns_uniform_backing_value = false;
+        try rt.copyDerivative(allocator, target, source);
         return;
     }
     try copyValue(target_value, source_value);
+    try rt.copyDerivative(allocator, target, source);
 }
 
 fn opDecorate(allocator: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeError!void {

test/maths.zig

@@ -78,6 +78,9 @@ test "Maths primitives" {
                 .expected_outputs = &.{
                     std.mem.asBytes(&[_]T{ expected, expected, expected, expected }),
                 },
+                .expected_output_tolerances = if (@typeInfo(T) == .float) &.{
+                    case.FloatTolerance.default(T),
+                } else &.{},
             });
         }
     }
@@ -150,6 +153,9 @@ test "Maths vectors" {
                     .expected_outputs = &.{
                         std.mem.asBytes(&@as([L]T, expected)),
                     },
+                    .expected_output_tolerances = if (@typeInfo(T) == .float) &.{
+                        case.FloatTolerance.default(T),
+                    } else &.{},
                 });
             }
         }
@@ -249,6 +255,9 @@ test "Maths vectors with scalars" {
                     .expected_outputs = &.{
                         std.mem.asBytes(&@as([L]T, expected)),
                     },
+                    .expected_output_tolerances = if (@typeInfo(T) == .float) &.{
+                        case.FloatTolerance.default(T),
+                    } else &.{},
                 });
             }
         }

test/root.zig

@@ -24,10 +24,37 @@ pub const case = struct {
         source: []const u32,
         inputs: []const []const u8 = &.{},
         expected_outputs: []const []const u8 = &.{},
+        expected_output_tolerances: []const ?FloatTolerance = &.{},
         descriptor_sets: []const []const []u8 = &.{},
         expected_descriptor_sets: []const []const []const u8 = &.{},
     };
 
+    const FloatKind = enum { f32, f64 };
+
+    pub const FloatTolerance = struct {
+        kind: FloatKind,
+        absolute: f64,
+
+        pub fn of(comptime T: type, absolute: f64) FloatTolerance {
+            return .{
+                .kind = switch (T) {
+                    f32 => .f32,
+                    f64 => .f64,
+                    else => @compileError("FloatTolerance only supports f32 and f64"),
+                },
+                .absolute = absolute,
+            };
+        }
+
+        pub fn default(comptime T: type) FloatTolerance {
+            return of(T, switch (T) {
+                f32 => 1e-6,
+                f64 => 1e-12,
+                else => @compileError("FloatTolerance only supports f32 and f64"),
+            });
+        }
+    };
+
     pub fn expect(config: Config) !void {
         const allocator = std.testing.allocator;
 
@@ -66,6 +93,12 @@ pub const case = struct {
                 defer allocator.free(output);
 
                 try rt.readOutput(output[0..], module.output_locations[n][0]);
+                if (n < config.expected_output_tolerances.len) {
+                    if (config.expected_output_tolerances[n]) |tolerance| {
+                        try expectApproxBytes(expected, output, tolerance);
+                        continue;
+                    }
+                }
                 try std.testing.expectEqualSlices(u8, expected, output);
             }
 
@@ -77,6 +110,30 @@ pub const case = struct {
         }
     }
 
+    fn expectApproxBytes(expected: []const u8, actual: []const u8, tolerance: FloatTolerance) !void {
+        try std.testing.expectEqual(expected.len, actual.len);
+
+        const stride: usize = switch (tolerance.kind) {
+            .f32 => @sizeOf(f32),
+            .f64 => @sizeOf(f64),
+        };
+        try std.testing.expectEqual(@as(usize, 0), expected.len % stride);
+
+        var i: usize = 0;
+        while (i < expected.len) : (i += stride) {
+            const expected_value = readFloat(expected[i .. i + stride], tolerance.kind);
+            const actual_value = readFloat(actual[i .. i + stride], tolerance.kind);
+            try std.testing.expectApproxEqAbs(expected_value, actual_value, tolerance.absolute);
+        }
+    }
+
+    fn readFloat(bytes: []const u8, kind: FloatKind) f64 {
+        return switch (kind) {
+            .f32 => @floatCast(@as(f32, @bitCast(std.mem.readInt(u32, bytes[0..@sizeOf(u32)], .little)))),
+            .f64 => @bitCast(std.mem.readInt(u64, bytes[0..@sizeOf(u64)], .little)),
+        };
+    }
+
     pub fn random(comptime T: type) T {
         var prng: std.Random.DefaultPrng = .init(@intCast(std.Io.Timestamp.now(std.testing.io, .real).toNanoseconds()));
         const rand = prng.random();