{#syntax#}@workGroupId(comptime dimension: u32) u32{#endsyntax#}
+ + Returns the index of the work group in the current kernel invocation in dimension {#syntax#}dimension{#endsyntax#}. +
+ {#header_close#} + + {#header_open|@workGroupSize#} +{#syntax#}@workGroupSize(comptime dimension: u32) u32{#endsyntax#}
+ + Returns the number of work items that a work group has in dimension {#syntax#}dimension{#endsyntax#}. +
+ {#header_close#} + + {#header_open|@workItemId#} +{#syntax#}@workItemId(comptime dimension: u32) u32{#endsyntax#}
+ + Returns the index of the work item in the work group in dimension {#syntax#}dimension{#endsyntax#}. This function returns values between {#syntax#}0{#endsyntax#} (inclusive) and {#syntax#}@workGroupSize(dimension){#endsyntax#} (exclusive). +
+ {#header_close#} + {#header_close#} {#header_open|Build Mode#} diff --git a/src/Air.zig b/src/Air.zig index 4646dcc89e2e18cca551918f8113dccbeed1b16b..c63e9826f9e078f779dbaa0b6b18080ca7349a8c 100644 --- a/src/Air.zig +++ b/src/Air.zig @@ -761,6 +761,22 @@ pub const Inst = struct { /// Uses the `ty` field. c_va_start, + /// Implements @workItemId builtin. + /// Result type is always `u32` + /// Uses the `pl_op` field, payload is the dimension to get the work item id for. + /// Operand is unused and set to Ref.none + work_item_id, + /// Implements @workGroupSize builtin. + /// Result type is always `u32` + /// Uses the `pl_op` field, payload is the dimension to get the work group size for. + /// Operand is unused and set to Ref.none + work_group_size, + /// Implements @workGroupId builtin. + /// Result type is always `u32` + /// Uses the `pl_op` field, payload is the dimension to get the work group id for. + /// Operand is unused and set to Ref.none + work_group_id, + pub fn fromCmpOp(op: std.math.CompareOperator, optimized: bool) Tag { switch (op) { .lt => return if (optimized) .cmp_lt_optimized else .cmp_lt, @@ -1267,6 +1283,11 @@ pub fn typeOfIndex(air: Air, inst: Air.Inst.Index) Type { const err_union_ty = air.typeOf(datas[inst].pl_op.operand); return err_union_ty.errorUnionPayload(); }, + + .work_item_id, + .work_group_size, + .work_group_id, + => return Type.u32, } } diff --git a/src/AstGen.zig b/src/AstGen.zig index c91303cdb17f923ea1330b8c009d761eb1467e4d..4f786edce8776c6d5f4660a7a11be495b132456a 100644 --- a/src/AstGen.zig +++ b/src/AstGen.zig @@ -8549,6 +8549,40 @@ fn builtinCall( } return rvalue(gz, ri, try gz.addNodeExtended(.c_va_start, node), node); }, + + .work_item_id => { + if (astgen.fn_block == null) { + return astgen.failNode(node, "'@workItemId' outside function scope", .{}); + } + const operand = try comptimeExpr(gz, scope, .{ .rl = .{ .coerced_ty = .u32_type } }, params[0]); + const result = try gz.addExtendedPayload(.work_item_id, Zir.Inst.UnNode{ + .node = gz.nodeIndexToRelative(node), + .operand = operand, + }); + return rvalue(gz, ri, result, node); + }, + .work_group_size => { + if (astgen.fn_block == null) { + return astgen.failNode(node, "'@workGroupSize' outside function scope", .{}); + } + const operand = try comptimeExpr(gz, scope, .{ .rl = .{ .coerced_ty = .u32_type } }, params[0]); + const result = try gz.addExtendedPayload(.work_group_size, Zir.Inst.UnNode{ + .node = gz.nodeIndexToRelative(node), + .operand = operand, + }); + return rvalue(gz, ri, result, node); + }, + .work_group_id => { + if (astgen.fn_block == null) { + return astgen.failNode(node, "'@workGroupId' outside function scope", .{}); + } + const operand = try comptimeExpr(gz, scope, .{ .rl = .{ .coerced_ty = .u32_type } }, params[0]); + const result = try gz.addExtendedPayload(.work_group_id, Zir.Inst.UnNode{ + .node = gz.nodeIndexToRelative(node), + .operand = operand, + }); + return rvalue(gz, ri, result, node); + }, } } diff --git a/src/BuiltinFn.zig b/src/BuiltinFn.zig index 79c6617483fe63b5e5f789cf8c486c5843c95b96..4a98a5a6156f9121766131de1559af447cc45b1e 100644 --- a/src/BuiltinFn.zig +++ b/src/BuiltinFn.zig @@ -118,6 +118,9 @@ pub const Tag = enum { union_init, Vector, volatile_cast, + work_item_id, + work_group_size, + work_group_id, }; pub const MemLocRequirement = enum { @@ -980,5 +983,25 @@ pub const list = list: { .param_count = 1, }, }, + .{ + "@workItemId", .{ + .tag = .work_item_id, + .param_count = 1, + }, + }, + .{ + "@workGroupSize", + .{ + .tag = .work_group_size, + .param_count = 1, + }, + }, + .{ + "@workGroupId", + .{ + .tag = .work_group_id, + .param_count = 1, + }, + }, }); }; diff --git a/src/Liveness.zig b/src/Liveness.zig index 8dc81aa16530473348652a17f80f161199745617..1d57b80097fe8f0e8de12d518a9e06243a4ce6e7 100644 --- a/src/Liveness.zig +++ b/src/Liveness.zig @@ -240,6 +240,9 @@ pub fn categorizeOperand( .err_return_trace, .save_err_return_trace_index, .c_va_start, + .work_item_id, + .work_group_size, + .work_group_id, => return .none, .fence => return .write, @@ -864,6 +867,9 @@ fn analyzeInst( .err_return_trace, .save_err_return_trace_index, .c_va_start, + .work_item_id, + .work_group_size, + .work_group_id, => return trackOperands(a, new_set, inst, main_tomb, .{ .none, .none, .none }), .not, diff --git a/src/Sema.zig b/src/Sema.zig index 1f375853cb9796e5d82a01cabdbb1e1ea30b6de0..da93a2906a12dd33101e9367e189f191ee95cc19 100644 --- a/src/Sema.zig +++ b/src/Sema.zig @@ -1164,6 +1164,9 @@ fn analyzeBodyInner( .c_va_start => try sema.zirCVaStart( block, extended), .const_cast, => try sema.zirConstCast( block, extended), .volatile_cast, => try sema.zirVolatileCast( block, extended), + .work_item_id => try sema.zirWorkItem( block, extended, extended.opcode), + .work_group_size => try sema.zirWorkItem( block, extended, extended.opcode), + .work_group_id => try sema.zirWorkItem( block, extended, extended.opcode), // zig fmt: on .fence => { @@ -22437,6 +22440,42 @@ fn zirBuiltinExtern( return sema.addConstant(ty, ref); } +fn zirWorkItem( + sema: *Sema, + block: *Block, + extended: Zir.Inst.Extended.InstData, + zir_tag: Zir.Inst.Extended, +) CompileError!Air.Inst.Ref { + const extra = sema.code.extraData(Zir.Inst.UnNode, extended.operand).data; + const dimension_src: LazySrcLoc = .{ .node_offset_builtin_call_arg0 = extra.node }; + const builtin_src = LazySrcLoc.nodeOffset(extra.node); + const target = sema.mod.getTarget(); + + switch (target.cpu.arch) { + // TODO: Allow for other GPU targets. + .amdgcn => {}, + else => { + return sema.fail(block, builtin_src, "builtin only available on GPU targets; targeted architecture is {s}", .{@tagName(target.cpu.arch)}); + }, + } + + const dimension = @intCast(u32, try sema.resolveInt(block, dimension_src, extra.operand, Type.u32, "dimension must be comptime-known")); + try sema.requireRuntimeBlock(block, builtin_src, null); + + return block.addInst(.{ + .tag = switch (zir_tag) { + .work_item_id => .work_item_id, + .work_group_size => .work_group_size, + .work_group_id => .work_group_id, + else => unreachable, + }, + .data = .{ .pl_op = .{ + .operand = .none, + .payload = dimension, + } }, + }); +} + fn requireRuntimeBlock(sema: *Sema, block: *Block, src: LazySrcLoc, runtime_src: ?LazySrcLoc) !void { if (block.is_comptime) { const msg = msg: { diff --git a/src/Zir.zig b/src/Zir.zig index bc5202c8aad558c9ea965c2163acc4ec10a82f04..7a8df49fda488f9fcecc54e92e14a47ec8ac1927 100644 --- a/src/Zir.zig +++ b/src/Zir.zig @@ -2032,6 +2032,15 @@ pub const Inst = struct { /// Implements the `@volatileCast` builtin. /// `operand` is payload index to `UnNode`. volatile_cast, + /// Implements the `@workItemId` builtin. + /// `operand` is payload index to `UnNode`. + work_item_id, + /// Implements the `@workGroupSize` builtin. + /// `operand` is payload index to `UnNode`. + work_group_size, + /// Implements the `@workGroupId` builtin. + /// `operand` is payload index to `UnNode`. + work_group_id, pub const InstData = struct { opcode: Extended, diff --git a/src/arch/aarch64/CodeGen.zig b/src/arch/aarch64/CodeGen.zig index ee23696950f655215a3c7fc18c611026844258f8..1e07e7e7199980a45ff0b788334ecb4bbc6572ec 100644 --- a/src/arch/aarch64/CodeGen.zig +++ b/src/arch/aarch64/CodeGen.zig @@ -890,6 +890,10 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void { .wasm_memory_size => unreachable, .wasm_memory_grow => unreachable, + + .work_item_id => unreachable, + .work_group_size => unreachable, + .work_group_id => unreachable, // zig fmt: on } diff --git a/src/arch/arm/CodeGen.zig b/src/arch/arm/CodeGen.zig index 8de5ae006a79239686d40e11483b3696539e4242..1cf0e2981e3e6859785ecb8d8a02c0c959c66fca 100644 --- a/src/arch/arm/CodeGen.zig +++ b/src/arch/arm/CodeGen.zig @@ -874,6 +874,10 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void { .wasm_memory_size => unreachable, .wasm_memory_grow => unreachable, + + .work_item_id => unreachable, + .work_group_size => unreachable, + .work_group_id => unreachable, // zig fmt: on } diff --git a/src/arch/riscv64/CodeGen.zig b/src/arch/riscv64/CodeGen.zig index 68df794bf75d256cdf9dca54378d07614de96175..11dbb2cb08efa02389d74961d6a2db8af783eed9 100644 --- a/src/arch/riscv64/CodeGen.zig +++ b/src/arch/riscv64/CodeGen.zig @@ -704,6 +704,10 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void { .wasm_memory_size => unreachable, .wasm_memory_grow => unreachable, + + .work_item_id => unreachable, + .work_group_size => unreachable, + .work_group_id => unreachable, // zig fmt: on } if (std.debug.runtime_safety) { diff --git a/src/arch/sparc64/CodeGen.zig b/src/arch/sparc64/CodeGen.zig index 30df99926760e7192aee279429fbe265a9e1a06e..a1b1be3a769ed540d52cb1b240254772b00f18bb 100644 --- a/src/arch/sparc64/CodeGen.zig +++ b/src/arch/sparc64/CodeGen.zig @@ -720,6 +720,10 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void { .wasm_memory_size => unreachable, .wasm_memory_grow => unreachable, + + .work_item_id => unreachable, + .work_group_size => unreachable, + .work_group_id => unreachable, // zig fmt: on } diff --git a/src/arch/wasm/CodeGen.zig b/src/arch/wasm/CodeGen.zig index 199ddada656f4962d9851e8804101683ef0b6f7c..fc5d13e5a4bbd24844907f5a3a3af1c9d5f4dec8 100644 --- a/src/arch/wasm/CodeGen.zig +++ b/src/arch/wasm/CodeGen.zig @@ -1997,6 +1997,11 @@ fn genInst(func: *CodeGen, inst: Air.Inst.Index) InnerError!void { .reduce_optimized, .float_to_int_optimized, => return func.fail("TODO implement optimized float mode", .{}), + + .work_item_id, + .work_group_size, + .work_group_id, + => unreachable, }; } diff --git a/src/arch/x86_64/CodeGen.zig b/src/arch/x86_64/CodeGen.zig index 5ddc9c77ca8394435a12b1ae504898e9657d7330..604052ee7e4d3167236a56c5ac4b12610d8a5354 100644 --- a/src/arch/x86_64/CodeGen.zig +++ b/src/arch/x86_64/CodeGen.zig @@ -1132,6 +1132,10 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void { .wasm_memory_size => unreachable, .wasm_memory_grow => unreachable, + + .work_item_id => unreachable, + .work_group_size => unreachable, + .work_group_id => unreachable, // zig fmt: on } diff --git a/src/codegen/c.zig b/src/codegen/c.zig index 6c4bb3c6889c6090ee07dc6a7b50ec3c4f4dfa61..704a1e31c5436fb2284de71768319a6873ff6f1d 100644 --- a/src/codegen/c.zig +++ b/src/codegen/c.zig @@ -2995,6 +2995,11 @@ fn genBodyInner(f: *Function, body: []const Air.Inst.Index) error{ AnalysisFail, .c_va_arg => try airCVaArg(f, inst), .c_va_end => try airCVaEnd(f, inst), .c_va_copy => try airCVaCopy(f, inst), + + .work_item_id, + .work_group_size, + .work_group_id, + => unreachable, // zig fmt: on }; if (result_value == .new_local) { diff --git a/src/codegen/llvm.zig b/src/codegen/llvm.zig index eea68a80e910ce981b77f6e744be95b769f2a99a..ce49fcde78a5be03b605b089dcf06520f1dc3f29 100644 --- a/src/codegen/llvm.zig +++ b/src/codegen/llvm.zig @@ -4745,6 +4745,10 @@ pub const FuncGen = struct { .c_va_copy => try self.airCVaCopy(inst), .c_va_end => try self.airCVaEnd(inst), .c_va_start => try self.airCVaStart(inst), + + .work_item_id => try self.airWorkItemId(inst), + .work_group_size => try self.airWorkGroupSize(inst), + .work_group_id => try self.airWorkGroupId(inst), // zig fmt: on }; if (opt_value) |val| { @@ -9567,6 +9571,74 @@ pub const FuncGen = struct { return self.builder.buildAddrSpaceCast(operand, llvm_dest_ty, ""); } + fn amdgcnWorkIntrinsic(self: *FuncGen, dimension: u32, default: u32, comptime basename: []const u8) !?*llvm.Value { + const llvm_u32 = self.context.intType(32); + + const llvm_fn_name = switch (dimension) { + 0 => basename ++ ".x", + 1 => basename ++ ".y", + 2 => basename ++ ".z", + else => return llvm_u32.constInt(default, .False), + }; + + const args: [0]*llvm.Value = .{}; + const llvm_fn = self.getIntrinsic(llvm_fn_name, &.{}); + return self.builder.buildCall(llvm_fn.globalGetValueType(), llvm_fn, &args, args.len, .Fast, .Auto, ""); + } + + fn airWorkItemId(self: *FuncGen, inst: Air.Inst.Index) !?*llvm.Value { + if (self.liveness.isUnused(inst)) return null; + + const target = self.dg.module.getTarget(); + assert(target.cpu.arch == .amdgcn); // TODO is to port this function to other GPU architectures + + const pl_op = self.air.instructions.items(.data)[inst].pl_op; + const dimension = pl_op.payload; + return self.amdgcnWorkIntrinsic(dimension, 0, "llvm.amdgcn.workitem.id"); + } + + fn airWorkGroupSize(self: *FuncGen, inst: Air.Inst.Index) !?*llvm.Value { + if (self.liveness.isUnused(inst)) return null; + + const target = self.dg.module.getTarget(); + assert(target.cpu.arch == .amdgcn); // TODO is to port this function to other GPU architectures + + const pl_op = self.air.instructions.items(.data)[inst].pl_op; + const dimension = pl_op.payload; + const llvm_u32 = self.context.intType(32); + if (dimension >= 3) { + return llvm_u32.constInt(1, .False); + } + + // Fetch the dispatch pointer, which points to this structure: + // https://github.com/RadeonOpenCompute/ROCR-Runtime/blob/adae6c61e10d371f7cbc3d0e94ae2c070cab18a4/src/inc/hsa.h#L2913 + const llvm_fn = self.getIntrinsic("llvm.amdgcn.dispatch.ptr", &.{}); + const args: [0]*llvm.Value = .{}; + const dispatch_ptr = self.builder.buildCall(llvm_fn.globalGetValueType(), llvm_fn, &args, args.len, .Fast, .Auto, ""); + dispatch_ptr.setAlignment(4); + + // Load the work_group_* member from the struct as u16. + // Just treat the dispatch pointer as an array of u16 to keep things simple. + const offset = 2 + dimension; + const index = [_]*llvm.Value{llvm_u32.constInt(offset, .False)}; + const llvm_u16 = self.context.intType(16); + const workgroup_size_ptr = self.builder.buildInBoundsGEP(llvm_u16, dispatch_ptr, &index, index.len, ""); + const workgroup_size = self.builder.buildLoad(llvm_u16, workgroup_size_ptr, ""); + workgroup_size.setAlignment(2); + return workgroup_size; + } + + fn airWorkGroupId(self: *FuncGen, inst: Air.Inst.Index) !?*llvm.Value { + if (self.liveness.isUnused(inst)) return null; + + const target = self.dg.module.getTarget(); + assert(target.cpu.arch == .amdgcn); // TODO is to port this function to other GPU architectures + + const pl_op = self.air.instructions.items(.data)[inst].pl_op; + const dimension = pl_op.payload; + return self.amdgcnWorkIntrinsic(dimension, 0, "llvm.amdgcn.workgroup.id"); + } + fn getErrorNameTable(self: *FuncGen) !*llvm.Value { if (self.dg.object.error_name_table) |table| { return table; diff --git a/src/print_air.zig b/src/print_air.zig index 8d29a272ca3b7c515b028941fa9014858380b366..803a0f2886bc4065a036e5137b89793874db9bed 100644 --- a/src/print_air.zig +++ b/src/print_air.zig @@ -328,6 +328,11 @@ const Writer = struct { .vector_store_elem => try w.writeVectorStoreElem(s, inst), .dbg_block_begin, .dbg_block_end => {}, + + .work_item_id, + .work_group_size, + .work_group_id, + => try w.writeWorkDimension(s, inst), } try s.writeAll(")\n"); } @@ -869,6 +874,11 @@ const Writer = struct { try w.writeOperand(s, inst, 0, pl_op.operand); } + fn writeWorkDimension(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void { + const pl_op = w.air.instructions.items(.data)[inst].pl_op; + try s.print("{d}", .{pl_op.payload}); + } + fn writeOperand( w: *Writer, s: anytype, diff --git a/src/print_zir.zig b/src/print_zir.zig index 755107cd1a56147e80a3753bbefca5348efebcbd..b70a4bbf67e760144688ec1ed349bf80f5890f15 100644 --- a/src/print_zir.zig +++ b/src/print_zir.zig @@ -512,6 +512,9 @@ const Writer = struct { .c_va_end, .const_cast, .volatile_cast, + .work_item_id, + .work_group_size, + .work_group_id, => { const inst_data = self.code.extraData(Zir.Inst.UnNode, extended.operand).data; const src = LazySrcLoc.nodeOffset(inst_data.node);