spirv: enhance assembler to handle constant strings and result IDs

Author: ivanstepanovftw

src/codegen/spirv.zig

@@ -6471,7 +6471,8 @@ const NavGen = struct {
 
                 // TODO: This entire function should be handled a bit better...
                 const ip = &zcu.intern_pool;
-                switch (ip.indexToKey(val.toIntern())) {
+                const key = ip.indexToKey(val.toIntern());
+                switch (key) {
                     .int_type,
                     .ptr_type,
                     .array_type,
@@ -6493,6 +6494,27 @@ const NavGen = struct {
 
                     .int => try as.value_map.put(as.gpa, name, .{ .constant = @intCast(val.toUnsignedInt(zcu)) }),
                     .enum_literal => |str| try as.value_map.put(as.gpa, name, .{ .string = str.toSlice(ip) }),
+                    .ptr => |ptr| {
+                        switch (ptr.base_addr) {
+                            .uav => |uav| {
+                                const pointee_val = uav.val;
+                                const pointee_key = ip.indexToKey(pointee_val);
+                                switch (pointee_key) {
+                                    .aggregate => |aggregate| {
+                                        if (Type.fromInterned(pointee_key.aggregate.ty).zigTypeTag(zcu) == .array and aggregate.storage == .bytes) {
+                                            const pointee_type = Type.fromInterned(pointee_key.aggregate.ty);
+                                            const bytes = aggregate.storage.bytes.toSlice(pointee_type.arrayLenIncludingSentinel(zcu), ip);
+                                            try as.value_map.put(as.gpa, name, .{ .string = bytes });
+                                        } else {
+                                            return self.fail("unsupported pointee constant type for 'c' constraint (ptr.uav.aggregate case - pointee type tag: {}, storage: {})", .{ Type.fromInterned(pointee_key.aggregate.ty).zigTypeTag(zcu), aggregate.storage });
+                                        }
+                                    },
+                                    else => return self.fail("unsupported pointee constant type for 'c' constraint (ptr.uav case): {}", .{pointee_key}),
+                                }
+                            },
+                            else => return self.fail("unsupported pointer base address type for 'c' constraint: {}", .{ptr.base_addr}),
+                        }
+                    },
 
                     else => unreachable, // TODO
                 }

src/codegen/spirv/Assembler.zig

@@ -153,6 +153,22 @@ const AsmValue = union(enum) {
             .ty => |result| result,
         };
     }
+
+    /// Retrieve the result-id of this AsmValue, creating constants as needed.
+    /// This is used by the assembler to convert constants and strings to result IDs.
+    pub fn resultIdOrConstant(self: AsmValue, spv: *SpvModule) !IdRef {
+        return switch (self) {
+            .just_declared,
+            .unresolved_forward_reference,
+            .constant
+            => unreachable, // TODO: Create constant integer instruction
+            .string => |str| {
+                return try spv.constStringGlobal(str);
+            },
+            .value => |result| result,
+            .ty => |result| result,
+        };
+    }
 };
 
 /// This map type maps results to values. Results can be addressed either by name (without the %), or by
@@ -517,7 +533,7 @@ fn processGenericInstruction(self: *Assembler) !?AsmValue {
             .ref_id => |index| {
                 const result = try self.resolveRef(index);
                 try section.ensureUnusedCapacity(self.spv.gpa, 1);
-                section.writeOperand(spec.IdRef, result.resultId());
+                section.writeOperand(spec.IdRef, try result.resultIdOrConstant(self.spv));
             },
             .string => |offset| {
                 const text = std.mem.sliceTo(self.inst.string_bytes.items[offset..], 0);
@@ -568,7 +584,7 @@ fn resolveRef(self: *Assembler, ref: AsmValue.Ref) !AsmValue {
 
 fn resolveRefId(self: *Assembler, ref: AsmValue.Ref) !IdRef {
     const value = try self.resolveRef(ref);
-    return value.resultId();
+    return value.resultIdOrConstant(self.spv);
 }
 
 /// Attempt to parse an instruction into `self.inst`.

src/codegen/spirv/Module.zig

@@ -173,6 +173,7 @@ cache: struct {
     extended_instruction_set: std.AutoHashMapUnmanaged(spec.InstructionSet, IdRef) = .empty,
     decorations: std.AutoHashMapUnmanaged(struct { IdRef, spec.Decoration }, void) = .empty,
     builtins: std.AutoHashMapUnmanaged(struct { IdRef, spec.BuiltIn }, Decl.Index) = .empty,
+    const_strings: std.StringHashMapUnmanaged(IdRef) = .empty,
 
     bool_const: [2]?IdRef = .{ null, null },
 } = .{},
@@ -231,6 +232,7 @@ pub fn deinit(self: *Module) void {
     self.cache.extended_instruction_set.deinit(self.gpa);
     self.cache.decorations.deinit(self.gpa);
     self.cache.builtins.deinit(self.gpa);
+    self.cache.const_strings.deinit(self.gpa);
 
     self.decls.deinit(self.gpa);
     self.decl_deps.deinit(self.gpa);
@@ -506,6 +508,47 @@ pub fn resolveString(self: *Module, string: []const u8) !IdRef {
     return id;
 }
 
+pub fn constStringGlobal(self: *Module, str: []const u8) !IdRef {
+    if (self.cache.const_strings.get(str)) |id| return id;
+
+    const u8_ty = try self.intType(.unsigned, 8);
+    const len = str.len + 1;
+    const len_id = try self.constant(try self.intType(.unsigned, 32), .{ .uint32 = @intCast(len) });
+    const arr_ty = try self.arrayType(len_id, u8_ty);
+
+    const char_ids = try self.gpa.alloc(IdRef, len);
+    defer self.gpa.free(char_ids);
+    for (str, 0..) |c, i| {
+        char_ids[i] = try self.constant(u8_ty, .{ .uint32 = c });
+    }
+    char_ids[str.len] = try self.constant(u8_ty, .{ .uint32 = 0 });
+
+    const const_arr_id = self.allocId();
+    try self.sections.types_globals_constants.emit(self.gpa, .OpConstantComposite, .{
+        .id_result_type = arr_ty,
+        .id_result = const_arr_id,
+        .constituents = char_ids,
+    });
+
+    const ptr_ty = self.allocId();
+    try self.sections.types_globals_constants.emit(self.gpa, .OpTypePointer, .{
+        .id_result = ptr_ty,
+        .storage_class = .UniformConstant,
+        .type = arr_ty,
+    });
+
+    const var_id = self.allocId();
+    try self.sections.types_globals_constants.emit(self.gpa, .OpVariable, .{
+        .id_result_type = ptr_ty,
+        .id_result = var_id,
+        .storage_class = .UniformConstant,
+        .initializer = const_arr_id,
+    });
+
+    try self.cache.const_strings.put(self.gpa, try self.arena.allocator().dupe(u8, str), var_id);
+    return var_id;
+}
+
 pub fn structType(self: *Module, result_id: IdResult, types: []const IdRef, maybe_names: ?[]const []const u8) !void {
     try self.sections.types_globals_constants.emit(self.gpa, .OpTypeStruct, .{
         .id_result = result_id,