adding derivatives support

src/GLSL_std_450/opcodes.zig

@@ -93,6 +93,7 @@ pub fn initRuntimeDispatcher() void {
     runtime_dispatcher[@intFromEnum(ext.GLSLOp.Log)]       = MathEngine(.Float, .Log).opSingleOperator;
     runtime_dispatcher[@intFromEnum(ext.GLSLOp.Log2)]      = MathEngine(.Float, .Log2).opSingleOperator;
     runtime_dispatcher[@intFromEnum(ext.GLSLOp.Normalize)] = opNormalize;
+    runtime_dispatcher[@intFromEnum(ext.GLSLOp.Pow)]       = MathEngine(.Float, .Pow).opDoubleOperators;
     runtime_dispatcher[@intFromEnum(ext.GLSLOp.Round)]     = MathEngine(.Float, .Round).opSingleOperator;
     runtime_dispatcher[@intFromEnum(ext.GLSLOp.SAbs)]      = MathEngine(.SInt,  .SAbs).opSingleOperator;
     runtime_dispatcher[@intFromEnum(ext.GLSLOp.SClamp)]    = MathEngine(.SInt, .SClamp).opTripleOperators;
@@ -102,6 +103,7 @@ pub fn initRuntimeDispatcher() void {
     runtime_dispatcher[@intFromEnum(ext.GLSLOp.Tan)]       = MathEngine(.Float, .Tan).opSingleOperator;
     runtime_dispatcher[@intFromEnum(ext.GLSLOp.Trunc)]     = MathEngine(.Float, .Trunc).opSingleOperator;
     runtime_dispatcher[@intFromEnum(ext.GLSLOp.UClamp)]    = MathEngine(.UInt, .UClamp).opTripleOperators;
+    runtime_dispatcher[@intFromEnum(ext.GLSLOp.Cross)]     = opCross;
     runtime_dispatcher[@intFromEnum(ext.GLSLOp.FindILsb)]  = IntBitEngine(.FindILsb).op;
     runtime_dispatcher[@intFromEnum(ext.GLSLOp.FindSMsb)]  = IntBitEngine(.FindSMsb).op;
     runtime_dispatcher[@intFromEnum(ext.GLSLOp.FindUMsb)]  = IntBitEngine(.FindUMsb).op;
@@ -205,6 +207,7 @@ fn MathEngine(comptime T: PrimitiveType, comptime Op: MathOp) type {
                     return switch (Op) {
                         .FMax => @max(l, r),
                         .FMin => @min(l, r),
+                        .Pow => if (comptime isFloatOrF32Vector(TT)) @exp(@log(l) * r) else RuntimeError.InvalidSpirV,
                         else => RuntimeError.InvalidSpirV,
                     };
                 }
@@ -397,6 +400,34 @@ fn IntBitEngine(comptime op_kind: IntBitOp) type {
     };
 }
 
+fn opCross(_: std.mem.Allocator, target_type_id: SpvWord, id: SpvWord, _: SpvWord, rt: *Runtime) RuntimeError!void {
+    const target_type = (try rt.results[target_type_id].getVariant()).Type;
+    const dst = try rt.results[id].getValue();
+    const lhs = try rt.results[try rt.it.next()].getValue();
+    const rhs = try rt.results[try rt.it.next()].getValue();
+
+    const lane_bits = try Result.resolveLaneBitWidth(target_type, rt);
+    if (try Result.resolveLaneCount(target_type) != 3)
+        return RuntimeError.InvalidSpirV;
+
+    switch (lane_bits) {
+        inline 16, 32, 64 => |bits| {
+            const FloatT = Value.getPrimitiveFieldType(.Float, bits);
+            const lx = try Value.readLane(.Float, bits, lhs, 0);
+            const ly = try Value.readLane(.Float, bits, lhs, 1);
+            const lz = try Value.readLane(.Float, bits, lhs, 2);
+            const rx = try Value.readLane(.Float, bits, rhs, 0);
+            const ry = try Value.readLane(.Float, bits, rhs, 1);
+            const rz = try Value.readLane(.Float, bits, rhs, 2);
+
+            try Value.writeLane(.Float, bits, dst, 0, @as(FloatT, ly * rz - lz * ry));
+            try Value.writeLane(.Float, bits, dst, 1, @as(FloatT, lz * rx - lx * rz));
+            try Value.writeLane(.Float, bits, dst, 2, @as(FloatT, lx * ry - ly * rx));
+        },
+        else => return RuntimeError.InvalidSpirV,
+    }
+}
+
 fn opLength(_: std.mem.Allocator, target_type_id: SpvWord, id: SpvWord, _: SpvWord, rt: *Runtime) RuntimeError!void {
     const target_type = (try rt.results[target_type_id].getVariant()).Type;
     const dst = try rt.results[id].getValue();

src/Runtime.zig

@@ -47,6 +47,23 @@ pub const SpecializationEntry = struct {
     size: usize,
 };
 
+pub const Derivative = struct {
+    dx: Value,
+    dy: Value,
+
+    pub fn dupe(self: *const @This(), allocator: std.mem.Allocator) RuntimeError!@This() {
+        return .{
+            .dx = try self.dx.dupe(allocator),
+            .dy = try self.dy.dupe(allocator),
+        };
+    }
+
+    pub fn deinit(self: *@This(), allocator: std.mem.Allocator) void {
+        self.dx.deinit(allocator);
+        self.dy.deinit(allocator);
+    }
+};
+
 pub const Function = struct {
     source_location: usize,
     result: *Result,
@@ -86,6 +103,7 @@ current_label: ?SpvWord,
 previous_label: ?SpvWord,
 
 specialization_constants: std.AutoHashMapUnmanaged(u32, []const u8),
+derivatives: std.AutoHashMapUnmanaged(SpvWord, Derivative),
 
 image_api: ImageAPI,
 
@@ -116,6 +134,7 @@ pub fn init(allocator: std.mem.Allocator, module: *Module, image_api: ImageAPI)
         .current_label = null,
         .previous_label = null,
         .specialization_constants = .empty,
+        .derivatives = .empty,
         .image_api = image_api,
     };
 }
@@ -131,6 +150,12 @@ pub fn deinit(self: *Self, allocator: std.mem.Allocator) void {
         allocator.free(entry.value_ptr.*);
     }
     self.specialization_constants.deinit(allocator);
+
+    var derivatives = self.derivatives.iterator();
+    while (derivatives.next()) |entry| {
+        entry.value_ptr.deinit(allocator);
+    }
+    self.derivatives.deinit(allocator);
 }
 
 pub fn addSpecializationInfo(self: *Self, allocator: std.mem.Allocator, entry: SpecializationEntry, data: []const u8) RuntimeError!void {
@@ -149,6 +174,62 @@ pub fn copySpecializationConstantsFrom(self: *Self, allocator: std.mem.Allocator
     }
 }
 
+pub fn setDerivative(self: *Self, allocator: std.mem.Allocator, result: SpvWord, dx: *const Value, dy: *const Value) RuntimeError!void {
+    const derivative: Derivative = .{
+        .dx = try dx.dupe(allocator),
+        .dy = try dy.dupe(allocator),
+    };
+    errdefer {
+        var tmp = derivative;
+        tmp.deinit(allocator);
+    }
+
+    const gop = self.derivatives.getOrPut(allocator, result) catch return RuntimeError.OutOfMemory;
+    if (gop.found_existing) {
+        gop.value_ptr.deinit(allocator);
+    }
+    gop.value_ptr.* = derivative;
+}
+
+pub fn setDerivativeFromMemory(self: *Self, allocator: std.mem.Allocator, result: SpvWord, dx: []const u8, dy: []const u8) RuntimeError!void {
+    const target_type = try self.getResultTargetTypeWord(result);
+
+    var dx_value = try Value.init(allocator, self.results, target_type, false);
+    defer dx_value.deinit(allocator);
+    _ = try dx_value.write(dx);
+
+    var dy_value = try Value.init(allocator, self.results, target_type, false);
+    defer dy_value.deinit(allocator);
+    _ = try dy_value.write(dy);
+
+    try self.setDerivative(allocator, result, &dx_value, &dy_value);
+}
+
+fn getResultTargetTypeWord(self: *const Self, result: SpvWord) RuntimeError!SpvWord {
+    return switch ((try self.results[result].getConstVariant()).*) {
+        .Variable => |v| v.type_word,
+        .Constant => |c| c.type_word,
+        .FunctionParameter => |p| p.type_word,
+        .AccessChain => |a| a.target,
+        else => return RuntimeError.InvalidSpirV,
+    };
+}
+
+pub fn clearDerivative(self: *Self, allocator: std.mem.Allocator, result: SpvWord) void {
+    if (self.derivatives.fetchRemove(result)) |kv| {
+        var derivative = kv.value;
+        derivative.deinit(allocator);
+    }
+}
+
+pub fn copyDerivative(self: *Self, allocator: std.mem.Allocator, dst: SpvWord, src: SpvWord) RuntimeError!void {
+    if (self.derivatives.get(src)) |derivative| {
+        try self.setDerivative(allocator, dst, &derivative.dx, &derivative.dy);
+    } else {
+        self.clearDerivative(allocator, dst);
+    }
+}
+
 pub fn getEntryPointByName(self: *const Self, name: []const u8) RuntimeError!SpvWord {
     for (self.mod.entry_points.items, 0..) |entry_point, i| {
         if (blk: {
@@ -523,6 +604,12 @@ pub fn hasResultDecoration(self: *const Self, result: SpvWord, decoration: spv.S
 }
 
 pub fn resetInvocation(self: *Self, allocator: std.mem.Allocator) void {
+    var derivatives = self.derivatives.iterator();
+    while (derivatives.next()) |entry| {
+        entry.value_ptr.deinit(allocator);
+    }
+    self.derivatives.clearRetainingCapacity();
+
     for (self.results) |*result| {
         if (result.variant) |*variant| {
             switch (variant.*) {

src/Value.zig

@@ -749,9 +749,41 @@ pub const Value = union(Type) {
 
         return switch (v.*) {
             .Int => (try getPrimitiveField(T, bits, @constCast(v))).*,
+            .Float => (try getPrimitiveField(T, bits, @constCast(v))).*,
 
             .Vector => |lanes| (try getPrimitiveField(T, bits, &lanes[lane_index])).*,
 
+            .Vector2f32 => |*vec| switch (lane_index) {
+                inline 0...1 => |i| blk: {
+                    if (comptime T == .Float and bits == 32) {
+                        break :blk @as(TT, vec[i]);
+                    } else {
+                        return RuntimeError.InvalidSpirV;
+                    }
+                },
+                else => return RuntimeError.InvalidSpirV,
+            },
+            .Vector3f32 => |*vec| switch (lane_index) {
+                inline 0...2 => |i| blk: {
+                    if (comptime T == .Float and bits == 32) {
+                        break :blk @as(TT, vec[i]);
+                    } else {
+                        return RuntimeError.InvalidSpirV;
+                    }
+                },
+                else => return RuntimeError.InvalidSpirV,
+            },
+            .Vector4f32 => |*vec| switch (lane_index) {
+                inline 0...3 => |i| blk: {
+                    if (comptime T == .Float and bits == 32) {
+                        break :blk @as(TT, vec[i]);
+                    } else {
+                        return RuntimeError.InvalidSpirV;
+                    }
+                },
+                else => return RuntimeError.InvalidSpirV,
+            },
+
             .Vector2i32 => |*vec| switch (lane_index) {
                 inline 0...1 => |i| blk: {
                     if (bits == 32) {
@@ -821,9 +853,35 @@ pub const Value = union(Type) {
     pub inline fn writeLane(comptime T: PrimitiveType, comptime bits: u32, dst: *Value, lane_index: usize, value: getPrimitiveFieldType(T, bits)) RuntimeError!void {
         switch (dst.*) {
             .Int => (try getPrimitiveField(T, bits, dst)).* = value,
+            .Float => (try getPrimitiveField(T, bits, dst)).* = value,
 
             .Vector => |lanes| try setScalarLaneValue(T, bits, &lanes[lane_index], value),
 
+            .Vector2f32 => |*vec| switch (lane_index) {
+                inline 0...1 => |i| if (comptime T == .Float and bits == 32) {
+                    vec[i] = value;
+                } else {
+                    return RuntimeError.InvalidSpirV;
+                },
+                else => return RuntimeError.InvalidSpirV,
+            },
+            .Vector3f32 => |*vec| switch (lane_index) {
+                inline 0...2 => |i| if (comptime T == .Float and bits == 32) {
+                    vec[i] = value;
+                } else {
+                    return RuntimeError.InvalidSpirV;
+                },
+                else => return RuntimeError.InvalidSpirV,
+            },
+            .Vector4f32 => |*vec| switch (lane_index) {
+                inline 0...3 => |i| if (comptime T == .Float and bits == 32) {
+                    vec[i] = value;
+                } else {
+                    return RuntimeError.InvalidSpirV;
+                },
+                else => return RuntimeError.InvalidSpirV,
+            },
+
             .Vector2i32 => |*vec| switch (lane_index) {
                 inline 0...1 => |i| if (bits == 32) {
                     vec[i] = @bitCast(value);
@@ -881,7 +939,17 @@ pub const Value = union(Type) {
     fn setScalarLaneValue(comptime value_type: PrimitiveType, comptime bits: u32, dst: *Value, v: getPrimitiveFieldType(value_type, bits)) RuntimeError!void {
         switch (bits) {
             inline 8, 16, 32, 64 => {
-                dst.* = .{ .Int = .{
+                dst.* = switch (value_type) {
+                    .Float => .{ .Float = .{
+                        .bit_count = bits,
+                        .value = switch (bits) {
+                            16 => .{ .float16 = v },
+                            32 => .{ .float32 = v },
+                            64 => .{ .float64 = v },
+                            else => return RuntimeError.InvalidSpirV,
+                        },
+                    } },
+                    .SInt, .UInt => .{ .Int = .{
                         .bit_count = bits,
                         .is_signed = if (value_type == .SInt) true else false,
                         .value = switch (value_type) {
@@ -899,9 +967,11 @@ pub const Value = union(Type) {
                                 64 => .{ .uint64 = v },
                                 else => unreachable,
                             },
-                        else => return RuntimeError.InvalidSpirV,
+                            else => unreachable,
                         },
-                } };
+                    } },
+                    else => return RuntimeError.InvalidSpirV,
+                };
             },
             else => return RuntimeError.InvalidSpirV,
         }

src/opcodes.zig

@@ -141,6 +141,12 @@ pub const SetupDispatcher = block: {
         .Decorate = opDecorate,
         .DecorationGroup = opDecorationGroup,
         .Dot = autoSetupConstant,
+        .DPdx = autoSetupConstant,
+        .DPdxCoarse = autoSetupConstant,
+        .DPdxFine = autoSetupConstant,
+        .DPdy = autoSetupConstant,
+        .DPdyCoarse = autoSetupConstant,
+        .DPdyFine = autoSetupConstant,
         .EntryPoint = opEntryPoint,
         .ExecutionMode = opExecutionMode,
         .ExtInst = autoSetupConstant,
@@ -303,6 +309,12 @@ pub fn initRuntimeDispatcher() void {
     runtime_dispatcher[@intFromEnum(spv.SpvOp.ConvertSToF)]            = ConversionEngine(.SInt, .Float).op;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.ConvertUToF)]            = ConversionEngine(.UInt, .Float).op;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.CopyMemory)]             = opCopyMemory;
+    runtime_dispatcher[@intFromEnum(spv.SpvOp.DPdx)]                   = DerivativeEngine(.x).op;
+    runtime_dispatcher[@intFromEnum(spv.SpvOp.DPdxCoarse)]             = DerivativeEngine(.x).op;
+    runtime_dispatcher[@intFromEnum(spv.SpvOp.DPdxFine)]               = DerivativeEngine(.x).op;
+    runtime_dispatcher[@intFromEnum(spv.SpvOp.DPdy)]                   = DerivativeEngine(.y).op;
+    runtime_dispatcher[@intFromEnum(spv.SpvOp.DPdyCoarse)]             = DerivativeEngine(.y).op;
+    runtime_dispatcher[@intFromEnum(spv.SpvOp.DPdyFine)]               = DerivativeEngine(.y).op;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.Dot)]                    = opDot;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.ExtInst)]                = opExtInst;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.FAdd)]                   = MathEngine(.Float, .Add, false).op;
@@ -341,7 +353,6 @@ pub fn initRuntimeDispatcher() void {
     runtime_dispatcher[@intFromEnum(spv.SpvOp.ImageSampleImplicitLod)] = ImageEngine(.SampleImplicitLod).op;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.ImageTexelPointer)]      = opImageTexelPointer;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.ImageWrite)]             = ImageEngine(.Write).op;
-    runtime_dispatcher[@intFromEnum(spv.SpvOp.SampledImage)]           = opSampledImage;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.InBoundsAccessChain)]    = opAccessChain;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.IsFinite)]               = CondEngine(.Float, .IsNan).op;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.IsInf)]                  = CondEngine(.Float, .IsInf).op;
@@ -373,6 +384,7 @@ pub fn initRuntimeDispatcher() void {
     runtime_dispatcher[@intFromEnum(spv.SpvOp.SMulExtended)]           = opSMulExtended;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.SNegate)]                = MathEngine(.SInt, .Negate, false).opSingle;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.SRem)]                   = MathEngine(.SInt, .Rem, false).op;
+    runtime_dispatcher[@intFromEnum(spv.SpvOp.SampledImage)]           = opSampledImage;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.Select)]                 = opSelect;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.ShiftLeftLogical)]       = BitEngine(.UInt, .ShiftLeft).op;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.ShiftRightArithmetic)]   = BitEngine(.SInt, .ShiftRightArithmetic).op;
@@ -384,6 +396,7 @@ pub fn initRuntimeDispatcher() void {
     runtime_dispatcher[@intFromEnum(spv.SpvOp.SpecConstantTrue)]       = opSpecConstantTrue;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.Store)]                  = opStore;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.Switch)]                 = opSwitch;
+    runtime_dispatcher[@intFromEnum(spv.SpvOp.TerminateInvocation)]    = opKill;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.UConvert)]               = ConversionEngine(.UInt, .UInt).op;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.UDiv)]                   = MathEngine(.UInt, .Div, false).op;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.UGreaterThan)]           = CondEngine(.UInt, .Greater).op;
@@ -392,6 +405,7 @@ pub fn initRuntimeDispatcher() void {
     runtime_dispatcher[@intFromEnum(spv.SpvOp.ULessThanEqual)]         = CondEngine(.UInt, .LessEqual).op;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.UMod)]                   = MathEngine(.UInt, .Mod, false).op;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.UMulExtended)]           = opUMulExtended;
+    runtime_dispatcher[@intFromEnum(spv.SpvOp.Unreachable)]            = opUnreachable;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.VectorShuffle)]          = opVectorShuffle;
     runtime_dispatcher[@intFromEnum(spv.SpvOp.VectorTimesMatrix)]      = MathEngine(.Float, .VectorTimesMatrix, false).op; // TODO
     runtime_dispatcher[@intFromEnum(spv.SpvOp.VectorTimesScalar)]      = MathEngine(.Float, .VectorTimesScalar, false).op;
@@ -3197,6 +3211,10 @@ fn opSMulExtended(_: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeError!v
     try opMulExtended(true, rt);
 }
 
+fn opUnreachable(_: std.mem.Allocator, _: SpvWord, _: *Runtime) RuntimeError!void {
+    return RuntimeError.Unreachable;
+}
+
 fn opSpecConstant(allocator: std.mem.Allocator, word_count: SpvWord, rt: *Runtime) RuntimeError!void {
     const location = rt.it.emitSourceLocation();
     _ = rt.it.skip();
@@ -3432,6 +3450,24 @@ fn opDot(_: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeError!void {
     }
 }
 
+fn DerivativeEngine(comptime axis: enum { x, y }) type {
+    return struct {
+        fn op(allocator: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeError!void {
+            _ = try rt.it.next(); // result type
+            const id = try rt.it.next();
+            const operand = try rt.it.next();
+
+            const derivative = rt.derivatives.get(operand) orelse return RuntimeError.UnsupportedSpirV;
+            const src = switch (axis) {
+                .x => &derivative.dx,
+                .y => &derivative.dy,
+            };
+            try copyValue(try rt.results[id].getValue(), src);
+            try rt.copyDerivative(allocator, id, operand);
+        }
+    };
+}
+
 fn opEntryPoint(allocator: std.mem.Allocator, word_count: SpvWord, rt: *Runtime) RuntimeError!void {
     const entry = rt.mod.entry_points.addOne(allocator) catch return RuntimeError.OutOfMemory;
     entry.exec_model = try rt.it.nextAs(spv.SpvExecutionModel);
@@ -3601,11 +3637,12 @@ fn opKill(_: std.mem.Allocator, _: SpvWord, _: *Runtime) RuntimeError!void {
     return RuntimeError.Killed;
 }
 
-fn opLoad(_: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeError!void {
+fn opLoad(allocator: std.mem.Allocator, _: SpvWord, rt: *Runtime) RuntimeError!void {
     _ = rt.it.skip();
     const id = try rt.it.next();
     const ptr_id = try rt.it.next();
     try copyValue(try rt.results[id].getValue(), try rt.results[ptr_id].getValue());
+    try rt.copyDerivative(allocator, id, ptr_id);
 }
 
 fn opMemberName(allocator: std.mem.Allocator, word_count: SpvWord, rt: *Runtime) RuntimeError!void {