Merge pull request #21937 from Snektron/spirv-vulkan-ptrs

spirv: miscellaneous vulkan + zig stuff

2 files changed, 199 insertions, 53 deletions

diff --git a/src/codegen/spirv.zig b/src/codegen/spirv.zig
index 34497d8de8..16b4a6dfbd 100644
--- a/src/codegen/spirv.zig
+++ b/src/codegen/spirv.zig
@@ -1396,10 +1396,6 @@ const NavGen = struct {
 
         const child_ty_id = try self.resolveType(child_ty, child_repr);
 
-        if (storage_class == .Uniform or storage_class == .PushConstant) {
-            try self.spv.decorate(child_ty_id, .Block);
-        }
-
         try self.spv.sections.types_globals_constants.emit(self.spv.gpa, .OpTypePointer, .{
             .id_result = result_id,
             .storage_class = storage_class,
@@ -1643,7 +1639,11 @@ const NavGen = struct {
                     return try self.arrayType(1, elem_ty_id);
                 } else {
                     const result_id = try self.arrayType(total_len, elem_ty_id);
-                    try self.spv.decorate(result_id, .{ .ArrayStride = .{ .array_stride = @intCast(elem_ty.abiSize(zcu)) } });
+                    if (target.os.tag == .vulkan) {
+                        try self.spv.decorate(result_id, .{ .ArrayStride = .{
+                            .array_stride = @intCast(elem_ty.abiSize(zcu)),
+                        } });
+                    }
                     return result_id;
                 }
             },
@@ -1662,7 +1662,7 @@ const NavGen = struct {
                         else => unreachable,
                     }
 
-                    // Guaranteed by callConvSupportsVarArgs, there are nog SPIR-V CCs which support
+                    // Guaranteed by callConvSupportsVarArgs, there are no SPIR-V CCs which support
                     // varargs.
                     assert(!fn_info.is_var_args);
 
@@ -1698,8 +1698,15 @@ const NavGen = struct {
             .pointer => {
                 const ptr_info = ty.ptrInfo(zcu);
 
+                const child_ty = Type.fromInterned(ptr_info.child);
                 const storage_class = self.spvStorageClass(ptr_info.flags.address_space);
-                const ptr_ty_id = try self.ptrType(Type.fromInterned(ptr_info.child), storage_class);
+                const ptr_ty_id = try self.ptrType(child_ty, storage_class);
+
+                if (target.os.tag == .vulkan and ptr_info.flags.size == .Many) {
+                    try self.spv.decorate(ptr_ty_id, .{ .ArrayStride = .{
+                        .array_stride = @intCast(child_ty.abiSize(zcu)),
+                    } });
+                }
 
                 if (ptr_info.flags.size != .Slice) {
                     return ptr_ty_id;
@@ -1746,6 +1753,10 @@ const NavGen = struct {
                         defer self.gpa.free(type_name);
                         try self.spv.debugName(result_id, type_name);
 
+                        if (target.os.tag == .vulkan) {
+                            try self.spv.decorate(result_id, .Block); // Decorate all structs as block for now...
+                        }
+
                         return result_id;
                     },
                     .struct_type => ip.loadStructType(ty.toIntern()),
@@ -1791,6 +1802,10 @@ const NavGen = struct {
                 defer self.gpa.free(type_name);
                 try self.spv.debugName(result_id, type_name);
 
+                if (target.os.tag == .vulkan) {
+                    try self.spv.decorate(result_id, .Block); // Decorate all structs as block for now...
+                }
+
                 return result_id;
             },
             .optional => {
@@ -1882,7 +1897,7 @@ const NavGen = struct {
                 else => unreachable,
             },
             .shared => .Workgroup,
-            .local => .Private,
+            .local => .Function,
             .global => switch (target.os.tag) {
                 .opencl => .CrossWorkgroup,
                 .vulkan => .PhysicalStorageBuffer,
@@ -1893,6 +1908,7 @@ const NavGen = struct {
             .input => .Input,
             .output => .Output,
             .uniform => .Uniform,
+            .storage_buffer => .StorageBuffer,
             .gs,
             .fs,
             .ss,
@@ -4354,13 +4370,24 @@ const NavGen = struct {
         defer self.gpa.free(ids);
 
         const result_id = self.spv.allocId();
-        try self.func.body.emit(self.spv.gpa, .OpInBoundsPtrAccessChain, .{
-            .id_result_type = result_ty_id,
-            .id_result = result_id,
-            .base = base,
-            .element = element,
-            .indexes = ids,
-        });
+        const target = self.getTarget();
+        switch (target.os.tag) {
+            .opencl => try self.func.body.emit(self.spv.gpa, .OpInBoundsPtrAccessChain, .{
+                .id_result_type = result_ty_id,
+                .id_result = result_id,
+                .base = base,
+                .element = element,
+                .indexes = ids,
+            }),
+            .vulkan => try self.func.body.emit(self.spv.gpa, .OpPtrAccessChain, .{
+                .id_result_type = result_ty_id,
+                .id_result = result_id,
+                .base = base,
+                .element = element,
+                .indexes = ids,
+            }),
+            else => unreachable,
+        }
         return result_id;
     }
 
@@ -6529,6 +6556,13 @@ const NavGen = struct {
             return self.todo("implement inline asm with more than 1 output", .{});
         }
 
+        var as = SpvAssembler{
+            .gpa = self.gpa,
+            .spv = self.spv,
+            .func = &self.func,
+        };
+        defer as.deinit();
+
         var output_extra_i = extra_i;
         for (outputs) |output| {
             if (output != .none) {
@@ -6541,7 +6575,6 @@ const NavGen = struct {
             // TODO: Record output and use it somewhere.
         }
 
-        var input_extra_i = extra_i;
         for (inputs) |input| {
             const extra_bytes = std.mem.sliceAsBytes(self.air.extra[extra_i..]);
             const constraint = std.mem.sliceTo(extra_bytes, 0);
@@ -6549,8 +6582,63 @@ const NavGen = struct {
             // This equation accounts for the fact that even if we have exactly 4 bytes
             // for the string, we still use the next u32 for the null terminator.
             extra_i += (constraint.len + name.len + (2 + 3)) / 4;
-            // TODO: Record input and use it somewhere.
-            _ = input;
+
+            const input_ty = self.typeOf(input);
+
+            if (std.mem.eql(u8, constraint, "c")) {
+                // constant
+                const val = (try self.air.value(input, self.pt)) orelse {
+                    return self.fail("assembly inputs with 'c' constraint have to be compile-time known", .{});
+                };
+
+                // TODO: This entire function should be handled a bit better...
+                const ip = &zcu.intern_pool;
+                switch (ip.indexToKey(val.toIntern())) {
+                    .int_type,
+                    .ptr_type,
+                    .array_type,
+                    .vector_type,
+                    .opt_type,
+                    .anyframe_type,
+                    .error_union_type,
+                    .simple_type,
+                    .struct_type,
+                    .union_type,
+                    .opaque_type,
+                    .enum_type,
+                    .func_type,
+                    .error_set_type,
+                    .inferred_error_set_type,
+                    => unreachable, // types, not values
+
+                    .undef => return self.fail("assembly input with 'c' constraint cannot be undefined", .{}),
+
+                    .int => {
+                        try as.value_map.put(as.gpa, name, .{ .constant = @intCast(val.toUnsignedInt(zcu)) });
+                    },
+
+                    else => unreachable, // TODO
+                }
+            } else if (std.mem.eql(u8, constraint, "t")) {
+                // type
+                if (input_ty.zigTypeTag(zcu) == .type) {
+                    // This assembly input is a type instead of a value.
+                    // That's fine for now, just make sure to resolve it as such.
+                    const val = (try self.air.value(input, self.pt)).?;
+                    const ty_id = try self.resolveType(val.toType(), .direct);
+                    try as.value_map.put(as.gpa, name, .{ .ty = ty_id });
+                } else {
+                    const ty_id = try self.resolveType(input_ty, .direct);
+                    try as.value_map.put(as.gpa, name, .{ .ty = ty_id });
+                }
+            } else {
+                if (input_ty.zigTypeTag(zcu) == .type) {
+                    return self.fail("use the 't' constraint to supply types to SPIR-V inline assembly", .{});
+                }
+
+                const val_id = try self.resolve(input);
+                try as.value_map.put(as.gpa, name, .{ .value = val_id });
+            }
         }
 
         {
@@ -6564,27 +6652,7 @@ const NavGen = struct {
 
         const asm_source = std.mem.sliceAsBytes(self.air.extra[extra_i..])[0..extra.data.source_len];
 
-        var as = SpvAssembler{
-            .gpa = self.gpa,
-            .src = asm_source,
-            .spv = self.spv,
-            .func = &self.func,
-        };
-        defer as.deinit();
-
-        for (inputs) |input| {
-            const extra_bytes = std.mem.sliceAsBytes(self.air.extra[input_extra_i..]);
-            const constraint = std.mem.sliceTo(extra_bytes, 0);
-            const name = std.mem.sliceTo(extra_bytes[constraint.len + 1 ..], 0);
-            // This equation accounts for the fact that even if we have exactly 4 bytes
-            // for the string, we still use the next u32 for the null terminator.
-            input_extra_i += (constraint.len + name.len + (2 + 3)) / 4;
-
-            const value = try self.resolve(input);
-            try as.value_map.put(as.gpa, name, .{ .value = value });
-        }
-
-        as.assemble() catch |err| switch (err) {
+        as.assemble(asm_source) catch |err| switch (err) {
             error.AssembleFail => {
                 // TODO: For now the compiler only supports a single error message per decl,
                 // so to translate the possible multiple errors from the assembler, emit
@@ -6629,6 +6697,7 @@ const NavGen = struct {
                 .just_declared, .unresolved_forward_reference => unreachable,
                 .ty => return self.fail("cannot return spir-v type as value from assembly", .{}),
                 .value => |ref| return ref,
+                .constant => return self.fail("cannot return constant from assembly", .{}),
             }
 
             // TODO: Multiple results
diff --git a/src/codegen/spirv/Assembler.zig b/src/codegen/spirv/Assembler.zig
index 9e39f2ed09..2cfb590273 100644
--- a/src/codegen/spirv/Assembler.zig
+++ b/src/codegen/spirv/Assembler.zig
@@ -45,6 +45,9 @@ const Token = struct {
         pipe,
         /// =.
         equals,
+        /// $identifier. This is used (for now) for constant values, like integers.
+        /// These can be used in place of a normal `value`.
+        placeholder,
 
         fn name(self: Tag) []const u8 {
             return switch (self) {
@@ -56,6 +59,7 @@ const Token = struct {
                 .string => "<string literal>",
                 .pipe => "'|'",
                 .equals => "'='",
+                .placeholder => "<placeholder>",
             };
         }
     };
@@ -128,12 +132,19 @@ const AsmValue = union(enum) {
     /// This result-value represents a type registered into the module's type system.
     ty: IdRef,
 
+    /// This is a pre-supplied constant integer value.
+    constant: u32,
+
     /// Retrieve the result-id of this AsmValue. Asserts that this AsmValue
     /// is of a variant that allows the result to be obtained (not an unresolved
     /// forward declaration, not in the process of being declared, etc).
     pub fn resultId(self: AsmValue) IdRef {
         return switch (self) {
-            .just_declared, .unresolved_forward_reference => unreachable,
+            .just_declared,
+            .unresolved_forward_reference,
+            // TODO: Lower this value as constant?
+            .constant,
+            => unreachable,
             .value => |result| result,
             .ty => |result| result,
         };
@@ -151,7 +162,8 @@ gpa: Allocator,
 errors: std.ArrayListUnmanaged(ErrorMsg) = .empty,
 
 /// The source code that is being assembled.
-src: []const u8,
+/// This is set when calling `assemble()`.
+src: []const u8 = undefined,
 
 /// The module that this assembly is associated to.
 /// Instructions like OpType*, OpDecorate, etc are emitted into this module.
@@ -211,7 +223,10 @@ pub fn deinit(self: *Assembler) void {
     self.instruction_map.deinit(self.gpa);
 }
 
-pub fn assemble(self: *Assembler) Error!void {
+pub fn assemble(self: *Assembler, src: []const u8) Error!void {
+    self.src = src;
+    self.errors.clearRetainingCapacity();
+
     // Populate the opcode map if it isn't already
     if (self.instruction_map.count() == 0) {
         const instructions = spec.InstructionSet.core.instructions();
@@ -369,6 +384,7 @@ fn processTypeInstruction(self: *Assembler) !AsmValue {
 /// - Function-local instructions are emitted in `self.func`.
 fn processGenericInstruction(self: *Assembler) !?AsmValue {
     const operands = self.inst.operands.items;
+    var maybe_spv_decl_index: ?SpvModule.Decl.Index = null;
     const section = switch (self.inst.opcode.class()) {
         .ConstantCreation => &self.spv.sections.types_globals_constants,
         .Annotation => &self.spv.sections.annotations,
@@ -378,13 +394,16 @@ fn processGenericInstruction(self: *Assembler) !?AsmValue {
             .OpExecutionMode, .OpExecutionModeId => &self.spv.sections.execution_modes,
             .OpVariable => switch (@as(spec.StorageClass, @enumFromInt(operands[2].value))) {
                 .Function => &self.func.prologue,
-                .UniformConstant => &self.spv.sections.types_globals_constants,
-                else => {
-                    // This is currently disabled because global variables are required to be
-                    // emitted in the proper order, and this should be honored in inline assembly
-                    // as well.
-                    return self.todo("global variables", .{});
+                .Input, .Output => section: {
+                    maybe_spv_decl_index = try self.spv.allocDecl(.global);
+                    try self.func.decl_deps.put(self.spv.gpa, maybe_spv_decl_index.?, {});
+                    // TODO: In theory this can be non-empty if there is an initializer which depends on another global...
+                    try self.spv.declareDeclDeps(maybe_spv_decl_index.?, &.{});
+                    break :section &self.spv.sections.types_globals_constants;
                 },
+                // These don't need to be marked in the dependency system.
+                // Probably we should add them anyway, then filter out PushConstant globals.
+                else => &self.spv.sections.types_globals_constants,
             },
             // Default case - to be worked out further.
             else => &self.func.body,
@@ -409,7 +428,10 @@ fn processGenericInstruction(self: *Assembler) !?AsmValue {
                 section.writeDoubleWord(dword);
             },
             .result_id => {
-                maybe_result_id = self.spv.allocId();
+                maybe_result_id = if (maybe_spv_decl_index) |spv_decl_index|
+                    self.spv.declPtr(spv_decl_index).result_id
+                else
+                    self.spv.allocId();
                 try section.ensureUnusedCapacity(self.spv.gpa, 1);
                 section.writeOperand(IdResult, maybe_result_id.?);
             },
@@ -475,8 +497,8 @@ fn resolveRefId(self: *Assembler, ref: AsmValue.Ref) !IdRef {
 /// error message has been emitted into `self.errors`.
 fn parseInstruction(self: *Assembler) !void {
     self.inst.opcode = undefined;
-    self.inst.operands.shrinkRetainingCapacity(0);
-    self.inst.string_bytes.shrinkRetainingCapacity(0);
+    self.inst.operands.clearRetainingCapacity();
+    self.inst.string_bytes.clearRetainingCapacity();
 
     const lhs_result_tok = self.currentToken();
     const maybe_lhs_result: ?AsmValue.Ref = if (self.eatToken(.result_id_assign)) blk: {
@@ -654,6 +676,22 @@ fn parseRefId(self: *Assembler) !void {
 
 fn parseLiteralInteger(self: *Assembler) !void {
     const tok = self.currentToken();
+    if (self.eatToken(.placeholder)) {
+        const name = self.tokenText(tok)[1..];
+        const value = self.value_map.get(name) orelse {
+            return self.fail(tok.start, "invalid placeholder '${s}'", .{name});
+        };
+        switch (value) {
+            .constant => |literal32| {
+                try self.inst.operands.append(self.gpa, .{ .literal32 = literal32 });
+            },
+            else => {
+                return self.fail(tok.start, "value '{s}' cannot be used as placeholder", .{name});
+            },
+        }
+        return;
+    }
+
     try self.expectToken(.value);
     // According to the SPIR-V machine readable grammar, a LiteralInteger
     // may consist of one or more words. From the SPIR-V docs it seems like there
@@ -669,6 +707,22 @@ fn parseLiteralInteger(self: *Assembler) !void {
 
 fn parseLiteralExtInstInteger(self: *Assembler) !void {
     const tok = self.currentToken();
+    if (self.eatToken(.placeholder)) {
+        const name = self.tokenText(tok)[1..];
+        const value = self.value_map.get(name) orelse {
+            return self.fail(tok.start, "invalid placeholder '${s}'", .{name});
+        };
+        switch (value) {
+            .constant => |literal32| {
+                try self.inst.operands.append(self.gpa, .{ .literal32 = literal32 });
+            },
+            else => {
+                return self.fail(tok.start, "value '{s}' cannot be used as placeholder", .{name});
+            },
+        }
+        return;
+    }
+
     try self.expectToken(.value);
     const text = self.tokenText(tok);
     const value = std.fmt.parseInt(u32, text, 0) catch {
@@ -745,6 +799,22 @@ fn parseContextDependentNumber(self: *Assembler) !void {
 
 fn parseContextDependentInt(self: *Assembler, signedness: std.builtin.Signedness, width: u32) !void {
     const tok = self.currentToken();
+    if (self.eatToken(.placeholder)) {
+        const name = self.tokenText(tok)[1..];
+        const value = self.value_map.get(name) orelse {
+            return self.fail(tok.start, "invalid placeholder '${s}'", .{name});
+        };
+        switch (value) {
+            .constant => |literal32| {
+                try self.inst.operands.append(self.gpa, .{ .literal32 = literal32 });
+            },
+            else => {
+                return self.fail(tok.start, "value '{s}' cannot be used as placeholder", .{name});
+            },
+        }
+        return;
+    }
+
     try self.expectToken(.value);
 
     if (width == 0 or width > 2 * @bitSizeOf(spec.Word)) {
@@ -848,6 +918,8 @@ fn tokenText(self: Assembler, tok: Token) []const u8 {
 /// Tokenize `self.src` and put the tokens in `self.tokens`.
 /// Any errors encountered are appended to `self.errors`.
 fn tokenize(self: *Assembler) !void {
+    self.tokens.clearRetainingCapacity();
+
     var offset: u32 = 0;
     while (true) {
         const tok = try self.nextToken(offset);
@@ -890,6 +962,7 @@ fn nextToken(self: *Assembler, start_offset: u32) !Token {
         string,
         string_end,
         escape,
+        placeholder,
     } = .start;
     var token_start = start_offset;
     var offset = start_offset;
@@ -917,6 +990,10 @@ fn nextToken(self: *Assembler, start_offset: u32) !Token {
                     offset += 1;
                     break;
                 },
+                '$' => {
+                    state = .placeholder;
+                    tag = .placeholder;
+                },
                 else => {
                     state = .value;
                     tag = .value;
@@ -932,11 +1009,11 @@ fn nextToken(self: *Assembler, start_offset: u32) !Token {
                 ' ', '\t', '\r', '\n', '=', '|' => break,
                 else => {},
             },
-            .result_id => switch (c) {
+            .result_id, .placeholder => switch (c) {
                 '_', 'a'...'z', 'A'...'Z', '0'...'9' => {},
                 ' ', '\t', '\r', '\n', '=', '|' => break,
                 else => {
-                    try self.addError(offset, "illegal character in result-id", .{});
+                    try self.addError(offset, "illegal character in result-id or placeholder", .{});
                     // Again, probably a forgotten delimiter here.
                     break;
                 },