authorgravatar for robin@voetter.nlRobin Voetter <robin@voetter.nl> 2023-03-18 15:59:56+01:00
committergravatar for git@vexu.euVeikka Tuominen <git@vexu.eu> 2023-03-30 12:20:24+03:00
log3357c59cebacb6b60da865376b20d2b307d12ec1
tree51edbb19a1f063888bac386c6a51ba250e361b11
parent83051b0cbf31b76e824d3911a7f4a0be3c0cf94d

new builtins: @workItemId, @workGroupId, @workGroupSize

* @workItemId returns the index of the work item in a work group for a dimension. * @workGroupId returns the index of the work group in the kernel dispatch for a dimension. * @workGroupSize returns the size of the work group for a dimension. These builtins are mainly useful for GPU backends. They are currently only implemented for the AMDGCN LLVM backend.

17 files changed, 269 insertions(+), 0 deletions(-)

doc/langref.html.in+22
......@@ -9578,6 +9578,28 @@ fn foo(comptime T: type, ptr: *T) T {
95789578 Remove {#syntax#}volatile{#endsyntax#} qualifier from a pointer.
95799579 </p>
95809580 {#header_close#}
9581
9582 {#header_open|@workGroupId#}
9583 <pre>{#syntax#}@workGroupId(comptime dimension: u32) u32{#endsyntax#}</pre>
9584 <p>
9585 Returns the index of the work group in the current kernel invocation in dimension {#syntax#}dimension{#endsyntax#}.
9586 </p>
9587 {#header_close#}
9588
9589 {#header_open|@workGroupSize#}
9590 <pre>{#syntax#}@workGroupSize(comptime dimension: u32) u32{#endsyntax#}</pre>
9591 <p>
9592 Returns the number of work items that a work group has in dimension {#syntax#}dimension{#endsyntax#}.
9593 </p>
9594 {#header_close#}
9595
9596 {#header_open|@workItemId#}
9597 <pre>{#syntax#}@workItemId(comptime dimension: u32) u32{#endsyntax#}</pre>
9598 <p>
9599 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).
9600 </p>
9601 {#header_close#}
9602
95819603 {#header_close#}
95829604
95839605 {#header_open|Build Mode#}
src/Air.zig+21
......@@ -761,6 +761,22 @@ pub const Inst = struct {
761761 /// Uses the `ty` field.
762762 c_va_start,
763763
764 /// Implements @workItemId builtin.
765 /// Result type is always `u32`
766 /// Uses the `pl_op` field, payload is the dimension to get the work item id for.
767 /// Operand is unused and set to Ref.none
768 work_item_id,
769 /// Implements @workGroupSize builtin.
770 /// Result type is always `u32`
771 /// Uses the `pl_op` field, payload is the dimension to get the work group size for.
772 /// Operand is unused and set to Ref.none
773 work_group_size,
774 /// Implements @workGroupId builtin.
775 /// Result type is always `u32`
776 /// Uses the `pl_op` field, payload is the dimension to get the work group id for.
777 /// Operand is unused and set to Ref.none
778 work_group_id,
779
764780 pub fn fromCmpOp(op: std.math.CompareOperator, optimized: bool) Tag {
765781 switch (op) {
766782 .lt => return if (optimized) .cmp_lt_optimized else .cmp_lt,
......@@ -1267,6 +1283,11 @@ pub fn typeOfIndex(air: Air, inst: Air.Inst.Index) Type {
12671283 const err_union_ty = air.typeOf(datas[inst].pl_op.operand);
12681284 return err_union_ty.errorUnionPayload();
12691285 },
1286
1287 .work_item_id,
1288 .work_group_size,
1289 .work_group_id,
1290 => return Type.u32,
12701291 }
12711292}
12721293
src/AstGen.zig+34
......@@ -8549,6 +8549,40 @@ fn builtinCall(
85498549 }
85508550 return rvalue(gz, ri, try gz.addNodeExtended(.c_va_start, node), node);
85518551 },
8552
8553 .work_item_id => {
8554 if (astgen.fn_block == null) {
8555 return astgen.failNode(node, "'@workItemId' outside function scope", .{});
8556 }
8557 const operand = try comptimeExpr(gz, scope, .{ .rl = .{ .coerced_ty = .u32_type } }, params[0]);
8558 const result = try gz.addExtendedPayload(.work_item_id, Zir.Inst.UnNode{
8559 .node = gz.nodeIndexToRelative(node),
8560 .operand = operand,
8561 });
8562 return rvalue(gz, ri, result, node);
8563 },
8564 .work_group_size => {
8565 if (astgen.fn_block == null) {
8566 return astgen.failNode(node, "'@workGroupSize' outside function scope", .{});
8567 }
8568 const operand = try comptimeExpr(gz, scope, .{ .rl = .{ .coerced_ty = .u32_type } }, params[0]);
8569 const result = try gz.addExtendedPayload(.work_group_size, Zir.Inst.UnNode{
8570 .node = gz.nodeIndexToRelative(node),
8571 .operand = operand,
8572 });
8573 return rvalue(gz, ri, result, node);
8574 },
8575 .work_group_id => {
8576 if (astgen.fn_block == null) {
8577 return astgen.failNode(node, "'@workGroupId' outside function scope", .{});
8578 }
8579 const operand = try comptimeExpr(gz, scope, .{ .rl = .{ .coerced_ty = .u32_type } }, params[0]);
8580 const result = try gz.addExtendedPayload(.work_group_id, Zir.Inst.UnNode{
8581 .node = gz.nodeIndexToRelative(node),
8582 .operand = operand,
8583 });
8584 return rvalue(gz, ri, result, node);
8585 },
85528586 }
85538587}
85548588
src/BuiltinFn.zig+23
......@@ -118,6 +118,9 @@ pub const Tag = enum {
118118 union_init,
119119 Vector,
120120 volatile_cast,
121 work_item_id,
122 work_group_size,
123 work_group_id,
121124};
122125
123126pub const MemLocRequirement = enum {
......@@ -980,5 +983,25 @@ pub const list = list: {
980983 .param_count = 1,
981984 },
982985 },
986 .{
987 "@workItemId", .{
988 .tag = .work_item_id,
989 .param_count = 1,
990 },
991 },
992 .{
993 "@workGroupSize",
994 .{
995 .tag = .work_group_size,
996 .param_count = 1,
997 },
998 },
999 .{
1000 "@workGroupId",
1001 .{
1002 .tag = .work_group_id,
1003 .param_count = 1,
1004 },
1005 },
9831006 });
9841007};
src/Liveness.zig+6
......@@ -240,6 +240,9 @@ pub fn categorizeOperand(
240240 .err_return_trace,
241241 .save_err_return_trace_index,
242242 .c_va_start,
243 .work_item_id,
244 .work_group_size,
245 .work_group_id,
243246 => return .none,
244247
245248 .fence => return .write,
......@@ -864,6 +867,9 @@ fn analyzeInst(
864867 .err_return_trace,
865868 .save_err_return_trace_index,
866869 .c_va_start,
870 .work_item_id,
871 .work_group_size,
872 .work_group_id,
867873 => return trackOperands(a, new_set, inst, main_tomb, .{ .none, .none, .none }),
868874
869875 .not,
src/Sema.zig+39
......@@ -1164,6 +1164,9 @@ fn analyzeBodyInner(
11641164 .c_va_start => try sema.zirCVaStart( block, extended),
11651165 .const_cast, => try sema.zirConstCast( block, extended),
11661166 .volatile_cast, => try sema.zirVolatileCast( block, extended),
1167 .work_item_id => try sema.zirWorkItem( block, extended, extended.opcode),
1168 .work_group_size => try sema.zirWorkItem( block, extended, extended.opcode),
1169 .work_group_id => try sema.zirWorkItem( block, extended, extended.opcode),
11671170 // zig fmt: on
11681171
11691172 .fence => {
......@@ -22437,6 +22440,42 @@ fn zirBuiltinExtern(
2243722440 return sema.addConstant(ty, ref);
2243822441}
2243922442
22443fn zirWorkItem(
22444 sema: *Sema,
22445 block: *Block,
22446 extended: Zir.Inst.Extended.InstData,
22447 zir_tag: Zir.Inst.Extended,
22448) CompileError!Air.Inst.Ref {
22449 const extra = sema.code.extraData(Zir.Inst.UnNode, extended.operand).data;
22450 const dimension_src: LazySrcLoc = .{ .node_offset_builtin_call_arg0 = extra.node };
22451 const builtin_src = LazySrcLoc.nodeOffset(extra.node);
22452 const target = sema.mod.getTarget();
22453
22454 switch (target.cpu.arch) {
22455 // TODO: Allow for other GPU targets.
22456 .amdgcn => {},
22457 else => {
22458 return sema.fail(block, builtin_src, "builtin only available on GPU targets; targeted architecture is {s}", .{@tagName(target.cpu.arch)});
22459 },
22460 }
22461
22462 const dimension = @intCast(u32, try sema.resolveInt(block, dimension_src, extra.operand, Type.u32, "dimension must be comptime-known"));
22463 try sema.requireRuntimeBlock(block, builtin_src, null);
22464
22465 return block.addInst(.{
22466 .tag = switch (zir_tag) {
22467 .work_item_id => .work_item_id,
22468 .work_group_size => .work_group_size,
22469 .work_group_id => .work_group_id,
22470 else => unreachable,
22471 },
22472 .data = .{ .pl_op = .{
22473 .operand = .none,
22474 .payload = dimension,
22475 } },
22476 });
22477}
22478
2244022479fn requireRuntimeBlock(sema: *Sema, block: *Block, src: LazySrcLoc, runtime_src: ?LazySrcLoc) !void {
2244122480 if (block.is_comptime) {
2244222481 const msg = msg: {
src/Zir.zig+9
......@@ -2032,6 +2032,15 @@ pub const Inst = struct {
20322032 /// Implements the `@volatileCast` builtin.
20332033 /// `operand` is payload index to `UnNode`.
20342034 volatile_cast,
2035 /// Implements the `@workItemId` builtin.
2036 /// `operand` is payload index to `UnNode`.
2037 work_item_id,
2038 /// Implements the `@workGroupSize` builtin.
2039 /// `operand` is payload index to `UnNode`.
2040 work_group_size,
2041 /// Implements the `@workGroupId` builtin.
2042 /// `operand` is payload index to `UnNode`.
2043 work_group_id,
20352044
20362045 pub const InstData = struct {
20372046 opcode: Extended,
src/arch/aarch64/CodeGen.zig+4
......@@ -890,6 +890,10 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
890890
891891 .wasm_memory_size => unreachable,
892892 .wasm_memory_grow => unreachable,
893
894 .work_item_id => unreachable,
895 .work_group_size => unreachable,
896 .work_group_id => unreachable,
893897 // zig fmt: on
894898 }
895899
src/arch/arm/CodeGen.zig+4
......@@ -874,6 +874,10 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
874874
875875 .wasm_memory_size => unreachable,
876876 .wasm_memory_grow => unreachable,
877
878 .work_item_id => unreachable,
879 .work_group_size => unreachable,
880 .work_group_id => unreachable,
877881 // zig fmt: on
878882 }
879883
src/arch/riscv64/CodeGen.zig+4
......@@ -704,6 +704,10 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
704704
705705 .wasm_memory_size => unreachable,
706706 .wasm_memory_grow => unreachable,
707
708 .work_item_id => unreachable,
709 .work_group_size => unreachable,
710 .work_group_id => unreachable,
707711 // zig fmt: on
708712 }
709713 if (std.debug.runtime_safety) {
src/arch/sparc64/CodeGen.zig+4
......@@ -720,6 +720,10 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
720720
721721 .wasm_memory_size => unreachable,
722722 .wasm_memory_grow => unreachable,
723
724 .work_item_id => unreachable,
725 .work_group_size => unreachable,
726 .work_group_id => unreachable,
723727 // zig fmt: on
724728 }
725729
src/arch/wasm/CodeGen.zig+5
......@@ -1997,6 +1997,11 @@ fn genInst(func: *CodeGen, inst: Air.Inst.Index) InnerError!void {
19971997 .reduce_optimized,
19981998 .float_to_int_optimized,
19991999 => return func.fail("TODO implement optimized float mode", .{}),
2000
2001 .work_item_id,
2002 .work_group_size,
2003 .work_group_id,
2004 => unreachable,
20002005 };
20012006}
20022007
src/arch/x86_64/CodeGen.zig+4
......@@ -1132,6 +1132,10 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
11321132
11331133 .wasm_memory_size => unreachable,
11341134 .wasm_memory_grow => unreachable,
1135
1136 .work_item_id => unreachable,
1137 .work_group_size => unreachable,
1138 .work_group_id => unreachable,
11351139 // zig fmt: on
11361140 }
11371141
src/codegen/c.zig+5
......@@ -2995,6 +2995,11 @@ fn genBodyInner(f: *Function, body: []const Air.Inst.Index) error{ AnalysisFail,
29952995 .c_va_arg => try airCVaArg(f, inst),
29962996 .c_va_end => try airCVaEnd(f, inst),
29972997 .c_va_copy => try airCVaCopy(f, inst),
2998
2999 .work_item_id,
3000 .work_group_size,
3001 .work_group_id,
3002 => unreachable,
29983003 // zig fmt: on
29993004 };
30003005 if (result_value == .new_local) {
src/codegen/llvm.zig+72
......@@ -4745,6 +4745,10 @@ pub const FuncGen = struct {
47454745 .c_va_copy => try self.airCVaCopy(inst),
47464746 .c_va_end => try self.airCVaEnd(inst),
47474747 .c_va_start => try self.airCVaStart(inst),
4748
4749 .work_item_id => try self.airWorkItemId(inst),
4750 .work_group_size => try self.airWorkGroupSize(inst),
4751 .work_group_id => try self.airWorkGroupId(inst),
47484752 // zig fmt: on
47494753 };
47504754 if (opt_value) |val| {
......@@ -9567,6 +9571,74 @@ pub const FuncGen = struct {
95679571 return self.builder.buildAddrSpaceCast(operand, llvm_dest_ty, "");
95689572 }
95699573
9574 fn amdgcnWorkIntrinsic(self: *FuncGen, dimension: u32, default: u32, comptime basename: []const u8) !?*llvm.Value {
9575 const llvm_u32 = self.context.intType(32);
9576
9577 const llvm_fn_name = switch (dimension) {
9578 0 => basename ++ ".x",
9579 1 => basename ++ ".y",
9580 2 => basename ++ ".z",
9581 else => return llvm_u32.constInt(default, .False),
9582 };
9583
9584 const args: [0]*llvm.Value = .{};
9585 const llvm_fn = self.getIntrinsic(llvm_fn_name, &.{});
9586 return self.builder.buildCall(llvm_fn.globalGetValueType(), llvm_fn, &args, args.len, .Fast, .Auto, "");
9587 }
9588
9589 fn airWorkItemId(self: *FuncGen, inst: Air.Inst.Index) !?*llvm.Value {
9590 if (self.liveness.isUnused(inst)) return null;
9591
9592 const target = self.dg.module.getTarget();
9593 assert(target.cpu.arch == .amdgcn); // TODO is to port this function to other GPU architectures
9594
9595 const pl_op = self.air.instructions.items(.data)[inst].pl_op;
9596 const dimension = pl_op.payload;
9597 return self.amdgcnWorkIntrinsic(dimension, 0, "llvm.amdgcn.workitem.id");
9598 }
9599
9600 fn airWorkGroupSize(self: *FuncGen, inst: Air.Inst.Index) !?*llvm.Value {
9601 if (self.liveness.isUnused(inst)) return null;
9602
9603 const target = self.dg.module.getTarget();
9604 assert(target.cpu.arch == .amdgcn); // TODO is to port this function to other GPU architectures
9605
9606 const pl_op = self.air.instructions.items(.data)[inst].pl_op;
9607 const dimension = pl_op.payload;
9608 const llvm_u32 = self.context.intType(32);
9609 if (dimension >= 3) {
9610 return llvm_u32.constInt(1, .False);
9611 }
9612
9613 // Fetch the dispatch pointer, which points to this structure:
9614 // https://github.com/RadeonOpenCompute/ROCR-Runtime/blob/adae6c61e10d371f7cbc3d0e94ae2c070cab18a4/src/inc/hsa.h#L2913
9615 const llvm_fn = self.getIntrinsic("llvm.amdgcn.dispatch.ptr", &.{});
9616 const args: [0]*llvm.Value = .{};
9617 const dispatch_ptr = self.builder.buildCall(llvm_fn.globalGetValueType(), llvm_fn, &args, args.len, .Fast, .Auto, "");
9618 dispatch_ptr.setAlignment(4);
9619
9620 // Load the work_group_* member from the struct as u16.
9621 // Just treat the dispatch pointer as an array of u16 to keep things simple.
9622 const offset = 2 + dimension;
9623 const index = [_]*llvm.Value{llvm_u32.constInt(offset, .False)};
9624 const llvm_u16 = self.context.intType(16);
9625 const workgroup_size_ptr = self.builder.buildInBoundsGEP(llvm_u16, dispatch_ptr, &index, index.len, "");
9626 const workgroup_size = self.builder.buildLoad(llvm_u16, workgroup_size_ptr, "");
9627 workgroup_size.setAlignment(2);
9628 return workgroup_size;
9629 }
9630
9631 fn airWorkGroupId(self: *FuncGen, inst: Air.Inst.Index) !?*llvm.Value {
9632 if (self.liveness.isUnused(inst)) return null;
9633
9634 const target = self.dg.module.getTarget();
9635 assert(target.cpu.arch == .amdgcn); // TODO is to port this function to other GPU architectures
9636
9637 const pl_op = self.air.instructions.items(.data)[inst].pl_op;
9638 const dimension = pl_op.payload;
9639 return self.amdgcnWorkIntrinsic(dimension, 0, "llvm.amdgcn.workgroup.id");
9640 }
9641
95709642 fn getErrorNameTable(self: *FuncGen) !*llvm.Value {
95719643 if (self.dg.object.error_name_table) |table| {
95729644 return table;
src/print_air.zig+10
......@@ -328,6 +328,11 @@ const Writer = struct {
328328 .vector_store_elem => try w.writeVectorStoreElem(s, inst),
329329
330330 .dbg_block_begin, .dbg_block_end => {},
331
332 .work_item_id,
333 .work_group_size,
334 .work_group_id,
335 => try w.writeWorkDimension(s, inst),
331336 }
332337 try s.writeAll(")\n");
333338 }
......@@ -869,6 +874,11 @@ const Writer = struct {
869874 try w.writeOperand(s, inst, 0, pl_op.operand);
870875 }
871876
877 fn writeWorkDimension(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
878 const pl_op = w.air.instructions.items(.data)[inst].pl_op;
879 try s.print("{d}", .{pl_op.payload});
880 }
881
872882 fn writeOperand(
873883 w: *Writer,
874884 s: anytype,
src/print_zir.zig+3
......@@ -512,6 +512,9 @@ const Writer = struct {
512512 .c_va_end,
513513 .const_cast,
514514 .volatile_cast,
515 .work_item_id,
516 .work_group_size,
517 .work_group_id,
515518 => {
516519 const inst_data = self.code.extraData(Zir.Inst.UnNode, extended.operand).data;
517520 const src = LazySrcLoc.nodeOffset(inst_data.node);