@extern: add support for spir-v locations and descriptors

lib/std/builtin.zig

@@ -1000,6 +1000,17 @@ pub const ExternOptions = struct {
     is_thread_local: bool = false,
     is_dll_import: bool = false,
     relocation: Relocation = .any,
+    decoration: ?Decoration = null,
+
+    pub const Decoration = union(enum) {
+        location: u32,
+        descriptor: Descriptor,
+
+        pub const Descriptor = struct {
+            binding: u32,
+            set: u32,
+        };
+    };
 
     pub const Relocation = enum(u1) {
         /// Any type of relocation is allowed.

lib/std/gpu.zig

@@ -20,30 +20,6 @@ pub extern const global_invocation_id: @Vector(3, u32) addrspace(.input);
 pub extern const vertex_index: u32 addrspace(.input);
 pub extern const instance_index: u32 addrspace(.input);
 
-/// Forms the main linkage for `input` and `output` address spaces.
-/// `ptr` must be a reference to variable or struct field.
-pub fn location(comptime ptr: anytype, comptime loc: u32) void {
-    asm volatile (
-        \\OpDecorate %ptr Location $loc
-        :
-        : [ptr] "" (ptr),
-          [loc] "c" (loc),
-    );
-}
-
-/// Forms the main linkage for `input` and `output` address spaces.
-/// `ptr` must be a reference to variable or struct field.
-pub fn binding(comptime ptr: anytype, comptime set: u32, comptime bind: u32) void {
-    asm volatile (
-        \\OpDecorate %ptr DescriptorSet $set
-        \\OpDecorate %ptr Binding $bind
-        :
-        : [ptr] "" (ptr),
-          [set] "c" (set),
-          [bind] "c" (bind),
-    );
-}
-
 pub const ExecutionMode = union(Tag) {
     /// Sets origin of the framebuffer to the upper-left corner
     origin_upper_left,

src/InternPool.zig

@@ -2270,6 +2270,7 @@ pub const Key = union(enum) {
         is_threadlocal: bool,
         is_dll_import: bool,
         relocation: std.builtin.ExternOptions.Relocation,
+        decoration: ?std.builtin.ExternOptions.Decoration,
         is_const: bool,
         alignment: Alignment,
         @"addrspace": std.builtin.AddressSpace,
@@ -5985,6 +5986,8 @@ pub const Tag = enum(u8) {
         flags: Flags,
         owner_nav: Nav.Index,
         zir_index: TrackedInst.Index,
+        location_or_descriptor_set: u32,
+        descriptor_binding: u32,
 
         pub const Flags = packed struct(u32) {
             linkage: std.builtin.GlobalLinkage,
@@ -5993,10 +5996,22 @@ pub const Tag = enum(u8) {
             is_dll_import: bool,
             relocation: std.builtin.ExternOptions.Relocation,
             source: Source,
-            _: u24 = 0,
+            decoration_type: DecorationType,
+            _: u22 = 0,
 
             pub const Source = enum(u1) { builtin, syntax };
+            pub const DecorationType = enum(u2) { none, location, descriptor };
         };
+
+        pub fn decoration(self: Extern) ?std.builtin.ExternOptions.Decoration {
+            return switch (self.flags.decoration_type) {
+                .none => null,
+                .location => std.builtin.ExternOptions.Decoration{
+                    .location = self.location_or_descriptor_set,
+                },
+                .descriptor => std.builtin.ExternOptions.Decoration{ .descriptor = .{ .set = self.location_or_descriptor_set, .binding = self.descriptor_binding } },
+            };
+        }
     };
 
     /// Trailing:
@@ -7354,6 +7369,7 @@ pub fn indexToKey(ip: *const InternPool, index: Index) Key {
                 .is_threadlocal = extra.flags.is_threadlocal,
                 .is_dll_import = extra.flags.is_dll_import,
                 .relocation = extra.flags.relocation,
+                .decoration = extra.decoration(),
                 .is_const = nav.status.fully_resolved.is_const,
                 .alignment = nav.status.fully_resolved.alignment,
                 .@"addrspace" = nav.status.fully_resolved.@"addrspace",
@@ -9243,15 +9259,22 @@ pub fn getExtern(
         .@"linksection" = .none,
         .@"addrspace" = key.@"addrspace",
     }) catch unreachable; // capacity asserted above
+    const decoration_type, const location_or_descriptor_set, const descriptor_binding = if (key.decoration) |decoration| switch (decoration) {
+        .location => |location| .{ Tag.Extern.Flags.DecorationType.location, location, undefined },
+        .descriptor => |descriptor| .{ Tag.Extern.Flags.DecorationType.descriptor, descriptor.binding, descriptor.set },
+    } else .{ Tag.Extern.Flags.DecorationType.none, undefined, undefined };
     const extra_index = addExtraAssumeCapacity(extra, Tag.Extern{
         .ty = key.ty,
         .lib_name = key.lib_name,
+        .location_or_descriptor_set = location_or_descriptor_set,
+        .descriptor_binding = descriptor_binding,
         .flags = .{
             .linkage = key.linkage,
             .visibility = key.visibility,
             .is_threadlocal = key.is_threadlocal,
             .is_dll_import = key.is_dll_import,
             .relocation = key.relocation,
+            .decoration_type = decoration_type,
             .source = key.source,
         },
         .zir_index = key.zir_index,

src/Sema.zig

@@ -25640,6 +25640,7 @@ fn resolveExternOptions(
     is_thread_local: bool,
     is_dll_import: bool,
     relocation: std.builtin.ExternOptions.Relocation,
+    decoration: ?std.builtin.ExternOptions.Decoration,
 } {
     const pt = sema.pt;
     const zcu = pt.zcu;
@@ -25656,6 +25657,7 @@ fn resolveExternOptions(
     const thread_local_src = block.src(.{ .init_field_thread_local = src.offset.node_offset_builtin_call_arg.builtin_call_node });
     const dll_import_src = block.src(.{ .init_field_dll_import = src.offset.node_offset_builtin_call_arg.builtin_call_node });
     const relocation_src = block.src(.{ .init_field_relocation = src.offset.node_offset_builtin_call_arg.builtin_call_node });
+    const decoration_src = block.src(.{ .init_field_decoration = src.offset.node_offset_builtin_call_arg.builtin_call_node });
 
     const name_ref = try sema.fieldVal(block, src, options, try ip.getOrPutString(gpa, pt.tid, "name", .no_embedded_nulls), name_src);
     const name = try sema.toConstString(block, name_src, name_ref, .{ .simple = .extern_options });
@@ -25690,6 +25692,10 @@ fn resolveExternOptions(
     const relocation_val = try sema.resolveConstDefinedValue(block, relocation_src, relocation_ref, .{ .simple = .extern_options });
     const relocation = try sema.interpretBuiltinType(block, relocation_src, relocation_val, std.builtin.ExternOptions.Relocation);
 
+    const decoration_ref = try sema.fieldVal(block, src, options, try ip.getOrPutString(gpa, pt.tid, "decoration", .no_embedded_nulls), decoration_src);
+    const decoration_val = try sema.resolveConstDefinedValue(block, decoration_src, decoration_ref, .{ .simple = .extern_options });
+    const decoration = try sema.interpretBuiltinType(block, decoration_src, decoration_val, ?std.builtin.ExternOptions.Decoration);
+
     if (name.len == 0) {
         return sema.fail(block, name_src, "extern symbol name cannot be empty", .{});
     }
@@ -25706,6 +25712,7 @@ fn resolveExternOptions(
         .is_thread_local = is_thread_local_val.toBool(),
         .is_dll_import = is_dll_import_val.toBool(),
         .relocation = relocation,
+        .decoration = decoration,
     };
 }
 
@@ -25765,6 +25772,7 @@ fn zirBuiltinExtern(
         .is_threadlocal = options.is_thread_local,
         .is_dll_import = options.is_dll_import,
         .relocation = options.relocation,
+        .decoration = options.decoration,
         .is_const = ptr_info.flags.is_const,
         .alignment = ptr_info.flags.alignment,
         .@"addrspace" = ptr_info.flags.address_space,

src/Zcu.zig

@@ -2062,6 +2062,7 @@ pub const SrcLoc = struct {
             .init_field_thread_local,
             .init_field_dll_import,
             .init_field_relocation,
+            .init_field_decoration,
             => |builtin_call_node| {
                 const wanted = switch (src_loc.lazy) {
                     .init_field_name => "name",
@@ -2075,6 +2076,7 @@ pub const SrcLoc = struct {
                     .init_field_thread_local => "thread_local",
                     .init_field_dll_import => "dll_import",
                     .init_field_relocation => "relocation",
+                    .init_field_decoration => "decoration",
                     else => unreachable,
                 };
                 const tree = try src_loc.file_scope.getTree(zcu);
@@ -2543,6 +2545,7 @@ pub const LazySrcLoc = struct {
         init_field_thread_local: Ast.Node.Offset,
         init_field_dll_import: Ast.Node.Offset,
         init_field_relocation: Ast.Node.Offset,
+        init_field_decoration: Ast.Node.Offset,
         /// The source location points to the value of an item in a specific
         /// case of a `switch`.
         switch_case_item: SwitchItem,

src/Zcu/PerThread.zig

@@ -1254,6 +1254,7 @@ fn analyzeNavVal(pt: Zcu.PerThread, nav_id: InternPool.Nav.Index) Zcu.CompileErr
                 .visibility = .default,
                 .is_dll_import = false,
                 .relocation = .any,
+                .decoration = null,
                 .is_const = is_const,
                 .alignment = modifiers.alignment,
                 .@"addrspace" = modifiers.@"addrspace",
@@ -3449,6 +3450,7 @@ pub fn getCoerced(pt: Zcu.PerThread, val: Value, new_ty: Type) Allocator.Error!V
                 .visibility = e.visibility,
                 .is_dll_import = e.is_dll_import,
                 .relocation = e.relocation,
+                .decoration = e.decoration,
                 .alignment = e.alignment,
                 .@"addrspace" = e.@"addrspace",
                 .zir_index = e.zir_index,

src/codegen/spirv/CodeGen.zig

@@ -254,7 +254,7 @@ pub fn genNav(cg: *CodeGen, do_codegen: bool) Error!void {
             try cg.module.debugName(func_result_id, nav.fqn.toSlice(ip));
         },
         .global => {
-            assert(ip.indexToKey(val.toIntern()) == .@"extern");
+            const key = ip.indexToKey(val.toIntern()).@"extern";
 
             const storage_class = cg.module.storageClass(nav.getAddrspace());
             assert(storage_class != .generic); // These should be instance globals
@@ -277,14 +277,32 @@ pub fn genNav(cg: *CodeGen, do_codegen: bool) Error!void {
                         }
                     }
 
-                    switch (ip.indexToKey(ty.toIntern())) {
-                        .func_type, .opaque_type => {},
-                        else => {
-                            try cg.module.decorate(ptr_ty_id, .{
-                                .array_stride = .{ .array_stride = @intCast(ty.abiSize(zcu)) },
+                    try cg.module.decorate(ptr_ty_id, .{
+                        .array_stride = .{ .array_stride = @intCast(ty.abiSize(zcu)) },
+                    });
+
+                    if (key.decoration) |decoration| switch (decoration) {
+                        .location => |location| {
+                            if (storage_class != .output and storage_class != .input and storage_class != .uniform_constant) {
+                                return cg.fail("storage class must be one of (output, input, uniform_constant) but is {s}", .{@tagName(storage_class)});
+                            }
+                            try cg.module.decorate(result_id, .{
+                                .location = .{ .location = location },
                             });
                         },
-                    }
+                        .descriptor => |descriptor| {
+                            if (storage_class != .storage_buffer and storage_class != .uniform and storage_class != .uniform_constant) {
+                                return cg.fail("storage class must be one of (storage_buffer, uniform, uniform_constant) but is {s}", .{@tagName(storage_class)});
+                            }
+                            try cg.module.decorate(result_id, .{
+                                .binding = .{ .binding_point = descriptor.binding },
+                            });
+
+                            try cg.module.decorate(result_id, .{
+                                .descriptor_set = .{ .descriptor_set = descriptor.set },
+                            });
+                        },
+                    };
                 },
                 else => {},
             }