spirv: add support for workItemId, workGroupId, workGroupSize

1 files changed, 55 insertions, 0 deletions

diff --git a/src/codegen/spirv.zig b/src/codegen/spirv.zig
index 09185211ef..42a42b10e7 100644
--- a/src/codegen/spirv.zig
+++ b/src/codegen/spirv.zig
@@ -3376,6 +3376,11 @@ const DeclGen = struct {
             .call_always_tail  => try self.airCall(inst, .always_tail),
             .call_never_tail   => try self.airCall(inst, .never_tail),
             .call_never_inline => try self.airCall(inst, .never_inline),
+
+            .work_item_id => try self.airWorkItemId(inst),
+            .work_group_size => try self.airWorkGroupSize(inst),
+            .work_group_id => try self.airWorkGroupId(inst),
+
             // zig fmt: on
 
             else => |tag| return self.todo("implement AIR tag {s}", .{@tagName(tag)}),
@@ -6533,6 +6538,56 @@ const DeclGen = struct {
         return result_id;
     }
 
+    fn builtin3D(self: *DeclGen, result_ty: Type, builtin: spec.BuiltIn, dimension: u32, out_of_range_value: anytype) !IdRef {
+        const mod = self.module;
+        if (dimension >= 3) {
+            return try self.constInt(result_ty, out_of_range_value, .direct);
+        }
+        const vec_ty = try mod.vectorType(.{
+            .len = 3,
+            .child = result_ty.toIntern(),
+        });
+        const ptr_ty_id = try self.ptrType(vec_ty, .Input);
+        const spv_decl_index = try self.spv.builtin(ptr_ty_id, builtin);
+        try self.func.decl_deps.put(self.spv.gpa, spv_decl_index, {});
+        const ptr = self.spv.declPtr(spv_decl_index).result_id;
+        const vec = try self.load(vec_ty, ptr, .{});
+        return try self.extractVectorComponent(result_ty, vec, dimension);
+    }
+
+    fn airWorkItemId(self: *DeclGen, inst: Air.Inst.Index) !?IdRef {
+        if (self.liveness.isUnused(inst)) return null;
+        const pl_op = self.air.instructions.items(.data)[@intFromEnum(inst)].pl_op;
+        const dimension = pl_op.payload;
+        // TODO: Should we make these builtins return usize?
+        const result_id = try self.builtin3D(Type.u64, .LocalInvocationId, dimension, 0);
+        const tmp = Temporary.init(Type.u64, result_id);
+        const result = try self.buildIntConvert(Type.u32, tmp);
+        return try result.materialize(self);
+    }
+
+    fn airWorkGroupSize(self: *DeclGen, inst: Air.Inst.Index) !?IdRef {
+        if (self.liveness.isUnused(inst)) return null;
+        const pl_op = self.air.instructions.items(.data)[@intFromEnum(inst)].pl_op;
+        const dimension = pl_op.payload;
+        // TODO: Should we make these builtins return usize?
+        const result_id = try self.builtin3D(Type.u64, .WorkgroupSize, dimension, 0);
+        const tmp = Temporary.init(Type.u64, result_id);
+        const result = try self.buildIntConvert(Type.u32, tmp);
+        return try result.materialize(self);
+    }
+
+    fn airWorkGroupId(self: *DeclGen, inst: Air.Inst.Index) !?IdRef {
+        if (self.liveness.isUnused(inst)) return null;
+        const pl_op = self.air.instructions.items(.data)[@intFromEnum(inst)].pl_op;
+        const dimension = pl_op.payload;
+        // TODO: Should we make these builtins return usize?
+        const result_id = try self.builtin3D(Type.u64, .WorkgroupId, dimension, 0);
+        const tmp = Temporary.init(Type.u64, result_id);
+        const result = try self.buildIntConvert(Type.u32, tmp);
+        return try result.materialize(self);
+    }
+
     fn typeOf(self: *DeclGen, inst: Air.Inst.Ref) Type {
         const mod = self.module;
         return self.air.typeOf(inst, &mod.intern_pool);