spirv: OpenCL printf support

lib/std/gpu.zig

@@ -1,4 +1,5 @@
 const std = @import("std.zig");
+const builtin = @import("builtin");
 
 pub const position_in = @extern(*addrspace(.input) @Vector(4, f32), .{ .name = "position" });
 pub const position_out = @extern(*addrspace(.output) @Vector(4, f32), .{ .name = "position" });
@@ -126,3 +127,35 @@ pub fn executionMode(comptime entry_point: anytype, comptime mode: ExecutionMode
         },
     }
 }
+
+/// Writes formatted output to an implementation-defined stream.
+/// Returns 0 on success, –1 on failure.
+pub fn printf(comptime fmt: [*:0]const u8, args: anytype) u32 {
+    if (builtin.zig_backend == .stage2_spirv and builtin.target.os.tag == .opencl) {
+        comptime var expr: []const u8 =
+            \\%std = OpExtInstImport "OpenCL.std"
+            \\%u32 = OpTypeInt 32 0
+            \\%ret = OpExtInst %u32 %std 184 %fmt
+        ;
+        inline for (0..args.len) |i| {
+            expr = expr ++ std.fmt.comptimePrint(" %arg{d}", .{i});
+        }
+        const result = switch (args.len) {
+            // zig fmt: off
+            0 => asm volatile (expr : [ret] "" (-> u32), : [fmt] "c" (fmt)),
+            1 => asm volatile (expr : [ret] "" (-> u32), : [fmt] "c" (fmt), [arg0] "" (args[0])),
+            2 => asm volatile (expr : [ret] "" (-> u32), : [fmt] "c" (fmt), [arg0] "" (args[0]), [arg1] "" (args[1])),
+            3 => asm volatile (expr : [ret] "" (-> u32), : [fmt] "c" (fmt), [arg0] "" (args[0]), [arg1] "" (args[1]), [arg2] "" (args[2])),
+            4 => asm volatile (expr : [ret] "" (-> u32), : [fmt] "c" (fmt), [arg0] "" (args[0]), [arg1] "" (args[1]), [arg2] "" (args[2]), [arg3] "" (args[3])),
+            5 => asm volatile (expr : [ret] "" (-> u32), : [fmt] "c" (fmt), [arg0] "" (args[0]), [arg1] "" (args[1]), [arg2] "" (args[2]), [arg3] "" (args[3]), [arg4] "" (args[4])),
+            6 => asm volatile (expr : [ret] "" (-> u32), : [fmt] "c" (fmt), [arg0] "" (args[0]), [arg1] "" (args[1]), [arg2] "" (args[2]), [arg3] "" (args[3]), [arg4] "" (args[4]), [arg5] "" (args[5])),
+            7 => asm volatile (expr : [ret] "" (-> u32), : [fmt] "c" (fmt), [arg0] "" (args[0]), [arg1] "" (args[1]), [arg2] "" (args[2]), [arg3] "" (args[3]), [arg4] "" (args[4]), [arg5] "" (args[5]), [arg6] "" (args[6])),
+            8 => asm volatile (expr : [ret] "" (-> u32), : [fmt] "c" (fmt), [arg0] "" (args[0]), [arg1] "" (args[1]), [arg2] "" (args[2]), [arg3] "" (args[3]), [arg4] "" (args[4]), [arg5] "" (args[5]), [arg6] "" (args[6]), [arg7] "" (args[7])),
+            9 => asm volatile (expr : [ret] "" (-> u32), : [fmt] "c" (fmt), [arg0] "" (args[0]), [arg1] "" (args[1]), [arg2] "" (args[2]), [arg3] "" (args[3]), [arg4] "" (args[4]), [arg5] "" (args[5]), [arg6] "" (args[6]), [arg7] "" (args[7]), [arg8] "" (args[8])),
+            // zig fmt: on
+            else => @compileError("too many arguments"),
+        };
+        return result;
+    }
+    @compileError("unsupported Zig backend or target OS");
+}

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
@@ -366,10 +382,18 @@ fn processTypeInstruction(self: *Assembler) !AsmValue {
             const child_type = try self.resolveRefId(operands[1].ref_id);
             break :blk try self.spv.vectorType(operands[2].literal32, child_type);
         },
-        .OpTypeArray => {
+        .OpTypeArray => blk: {
             // TODO: The length of an OpTypeArray is determined by a constant (which may be a spec constant),
             // and so some consideration must be taken when entering this in the type system.
-            return self.todo("process OpTypeArray", .{});
+            const element_type = try self.resolveRefId(operands[1].ref_id);
+            const length = try self.resolveRefId(operands[2].ref_id);
+            const result_id = self.spv.allocId();
+            try section.emit(self.spv.gpa, .OpTypeArray, .{
+                .id_result = result_id,
+                .element_type = element_type,
+                .length = length,
+            });
+            break :blk result_id;
         },
         .OpTypeRuntimeArray => blk: {
             const element_type = try self.resolveRefId(operands[1].ref_id);
@@ -509,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);
@@ -560,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,

test/behavior/asm.zig

@@ -191,3 +191,13 @@ test "asm modifiers (AArch64)" {
     );
     try expectEqual(2 * x, double);
 }
+
+test "SPIR-V string constants with c constraint" {
+    if (builtin.zig_backend != .stage2_spirv) return error.SkipZigTest;
+    if (builtin.target.os.tag != .opencl) return error.SkipZigTest;
+
+    _ = std.gpu.printf("testing printf\n", .{});
+    var a: i32 = -10;
+    _ = &a;
+    _ = std.gpu.printf("a: %d, a + 10: %d\n", .{ a, a + 10 });
+}