authorgravatar for jacobly@ziglang.orgJacob Young <jacobly@ziglang.org> 2025-05-30 18:04:30-04:00
committergravatar for mlugg@mlugg.co.ukMatthew Lugg <mlugg@mlugg.co.uk> 2025-06-01 08:24:01+01:00
logb48d6ff619de424e664cf11b43e2a03fecbca6ce
tree5d0457b8c548df15b170bf3e1d3b941276052f6e
parent32a57bfeaa9a9cd2b1a4c0e25d829824d8b96c8d
signaturelock-open Commit is signed but in an unrecognized format.

Legalize: implement scalarization of `@select`

5 files changed, 102 insertions(+), 21 deletions(-)

lib/std/simd.zig-3
...@@ -368,9 +368,6 @@ pub fn countElementsWithValue(vec: anytype, value: std.meta.Child(@TypeOf(vec)))...@@ -368,9 +368,6 @@ pub fn countElementsWithValue(vec: anytype, value: std.meta.Child(@TypeOf(vec)))
368}368}
369369
370test "vector searching" {370test "vector searching" {
371 if (builtin.zig_backend == .stage2_x86_64 and
372 !comptime std.Target.x86.featureSetHas(builtin.cpu.features, .ssse3)) return error.SkipZigTest;
373
374 const base = @Vector(8, u32){ 6, 4, 7, 4, 4, 2, 3, 7 };371 const base = @Vector(8, u32){ 6, 4, 7, 4, 4, 2, 3, 7 };
375372
376 try std.testing.expectEqual(@as(?u3, 1), firstIndexOfValue(base, 4));373 try std.testing.expectEqual(@as(?u3, 1), firstIndexOfValue(base, 4));
src/Air/Legalize.zig+100-14
...@@ -74,6 +74,7 @@ pub const Feature = enum {...@@ -74,6 +74,7 @@ pub const Feature = enum {
74 scalarize_int_from_float,74 scalarize_int_from_float,
75 scalarize_int_from_float_optimized,75 scalarize_int_from_float_optimized,
76 scalarize_float_from_int,76 scalarize_float_from_int,
77 scalarize_select,
77 scalarize_mul_add,78 scalarize_mul_add,
7879
79 /// Legalize (shift lhs, (splat rhs)) -> (shift lhs, rhs)80 /// Legalize (shift lhs, (splat rhs)) -> (shift lhs, rhs)
...@@ -167,6 +168,7 @@ pub const Feature = enum {...@@ -167,6 +168,7 @@ pub const Feature = enum {
167 .int_from_float => .scalarize_int_from_float,168 .int_from_float => .scalarize_int_from_float,
168 .int_from_float_optimized => .scalarize_int_from_float_optimized,169 .int_from_float_optimized => .scalarize_int_from_float_optimized,
169 .float_from_int => .scalarize_float_from_int,170 .float_from_int => .scalarize_float_from_int,
171 .select => .scalarize_select,
170 .mul_add => .scalarize_mul_add,172 .mul_add => .scalarize_mul_add,
171 };173 };
172 }174 }
...@@ -520,7 +522,9 @@ fn legalizeBody(l: *Legalize, body_start: usize, body_len: usize) Error!void {...@@ -520,7 +522,9 @@ fn legalizeBody(l: *Legalize, body_start: usize, body_len: usize) Error!void {
520 },522 },
521 .splat,523 .splat,
522 .shuffle,524 .shuffle,
525 => {},
523 .select,526 .select,
527 => if (l.features.contains(.scalarize_select)) continue :inst try l.scalarize(inst, .select_pl_op_bin),
524 .memset,528 .memset,
525 .memset_safe,529 .memset_safe,
526 .memcpy,530 .memcpy,
...@@ -568,7 +572,7 @@ fn legalizeBody(l: *Legalize, body_start: usize, body_len: usize) Error!void {...@@ -568,7 +572,7 @@ fn legalizeBody(l: *Legalize, body_start: usize, body_len: usize) Error!void {
568 }572 }
569}573}
570574
571const ScalarizeDataTag = enum { un_op, ty_op, bin_op, ty_pl_vector_cmp, pl_op_bin };575const ScalarizeDataTag = enum { un_op, ty_op, bin_op, ty_pl_vector_cmp, pl_op_bin, select_pl_op_bin };
572inline fn scalarize(l: *Legalize, orig_inst: Air.Inst.Index, comptime data_tag: ScalarizeDataTag) Error!Air.Inst.Tag {576inline fn scalarize(l: *Legalize, orig_inst: Air.Inst.Index, comptime data_tag: ScalarizeDataTag) Error!Air.Inst.Tag {
573 return l.replaceInst(orig_inst, .block, try l.scalarizeBlockPayload(orig_inst, data_tag));577 return l.replaceInst(orig_inst, .block, try l.scalarizeBlockPayload(orig_inst, data_tag));
574}578}
...@@ -584,6 +588,7 @@ fn scalarizeBlockPayload(l: *Legalize, orig_inst: Air.Inst.Index, comptime data_...@@ -584,6 +588,7 @@ fn scalarizeBlockPayload(l: *Legalize, orig_inst: Air.Inst.Index, comptime data_
584 .un_op, .ty_op => 1,588 .un_op, .ty_op => 1,
585 .bin_op, .ty_pl_vector_cmp => 2,589 .bin_op, .ty_pl_vector_cmp => 2,
586 .pl_op_bin => 3,590 .pl_op_bin => 3,
591 .select_pl_op_bin => 6,
587 } + 9592 } + 9
588 ]Air.Inst.Index = undefined;593 ]Air.Inst.Index = undefined;
589 try l.air_instructions.ensureUnusedCapacity(zcu.gpa, inst_buf.len);594 try l.air_instructions.ensureUnusedCapacity(zcu.gpa, inst_buf.len);
...@@ -722,23 +727,67 @@ fn scalarizeBlockPayload(l: *Legalize, orig_inst: Air.Inst.Index, comptime data_...@@ -722,23 +727,67 @@ fn scalarizeBlockPayload(l: *Legalize, orig_inst: Air.Inst.Index, comptime data_
722 } },727 } },
723 });728 });
724 },729 },
730 .select_pl_op_bin => {
731 const extra = l.extraData(Air.Bin, orig.data.pl_op.payload).data;
732 var res_elem: Result = .init(l, l.typeOf(extra.lhs).scalarType(zcu), &loop.block);
733 res_elem.block = .init(loop.block.stealCapacity(6));
734 {
735 var select_cond_br: CondBr = .init(l, res_elem.block.add(l, .{
736 .tag = .array_elem_val,
737 .data = .{ .bin_op = .{
738 .lhs = orig.data.pl_op.operand,
739 .rhs = cur_index_inst.toRef(),
740 } },
741 }).toRef(), &res_elem.block, .{});
742 select_cond_br.then_block = .init(res_elem.block.stealRemainingCapacity());
743 {
744 _ = select_cond_br.then_block.add(l, .{
745 .tag = .br,
746 .data = .{ .br = .{
747 .block_inst = res_elem.inst,
748 .operand = select_cond_br.then_block.add(l, .{
749 .tag = .array_elem_val,
750 .data = .{ .bin_op = .{
751 .lhs = extra.lhs,
752 .rhs = cur_index_inst.toRef(),
753 } },
754 }).toRef(),
755 } },
756 });
757 }
758 select_cond_br.else_block = .init(select_cond_br.then_block.stealRemainingCapacity());
759 {
760 _ = select_cond_br.else_block.add(l, .{
761 .tag = .br,
762 .data = .{ .br = .{
763 .block_inst = res_elem.inst,
764 .operand = select_cond_br.else_block.add(l, .{
765 .tag = .array_elem_val,
766 .data = .{ .bin_op = .{
767 .lhs = extra.rhs,
768 .rhs = cur_index_inst.toRef(),
769 } },
770 }).toRef(),
771 } },
772 });
773 }
774 try select_cond_br.finish(l);
775 }
776 try res_elem.finish(l);
777 break :res_elem res_elem.inst;
778 },
725 }.toRef(),779 }.toRef(),
726 }),780 }),
727 } },781 } },
728 });782 });
729783
730 var loop_cond_br: CondBr = .init(784 var loop_cond_br: CondBr = .init(l, (try loop.block.addCmp(
731 l,785 l,
732 (try loop.block.addCmp(786 .lt,
733 l,787 cur_index_inst.toRef(),
734 .lt,788 try pt.intRef(.usize, res_ty.vectorLen(zcu) - 1),
735 cur_index_inst.toRef(),
736 try pt.intRef(.usize, res_ty.vectorLen(zcu) - 1),
737 .{},
738 )).toRef(),
739 &loop.block,
740 .{},789 .{},
741 );790 )).toRef(), &loop.block, .{});
742 loop_cond_br.then_block = .init(loop.block.stealRemainingCapacity());791 loop_cond_br.then_block = .init(loop.block.stealRemainingCapacity());
743 {792 {
744 _ = loop_cond_br.then_block.add(l, .{793 _ = loop_cond_br.then_block.add(l, .{
...@@ -1138,9 +1187,21 @@ const Block = struct {...@@ -1138,9 +1187,21 @@ const Block = struct {
1138 /// This is useful when you've provided a buffer big enough for all your instructions, but you are1187 /// This is useful when you've provided a buffer big enough for all your instructions, but you are
1139 /// now starting a new block and some of them need to live there instead.1188 /// now starting a new block and some of them need to live there instead.
1140 fn stealRemainingCapacity(b: *Block) []Air.Inst.Index {1189 fn stealRemainingCapacity(b: *Block) []Air.Inst.Index {
1141 const remaining = b.instructions[b.len..];1190 return b.stealFrom(b.len);
1142 b.instructions = b.instructions[0..b.len];1191 }
1143 return remaining;1192
1193 /// Returns `len` elements taken from the unused capacity of `b.instructions`, and shrinks
1194 /// `b.instructions` down to not include them anymore.
1195 /// This is useful when you've provided a buffer big enough for all your instructions, but you are
1196 /// now starting a new block and some of them need to live there instead.
1197 fn stealCapacity(b: *Block, len: usize) []Air.Inst.Index {
1198 return b.stealFrom(b.instructions.len - len);
1199 }
1200
1201 fn stealFrom(b: *Block, start: usize) []Air.Inst.Index {
1202 assert(start >= b.len);
1203 defer b.instructions.len = start;
1204 return b.instructions[start..];
1144 }1205 }
11451206
1146 fn body(b: *const Block) []const Air.Inst.Index {1207 fn body(b: *const Block) []const Air.Inst.Index {
...@@ -1149,6 +1210,31 @@ const Block = struct {...@@ -1149,6 +1210,31 @@ const Block = struct {
1149 }1210 }
1150};1211};
11511212
1213const Result = struct {
1214 inst: Air.Inst.Index,
1215 block: Block,
1216
1217 /// The return value has `block` initialized to `undefined`; it is the caller's reponsibility
1218 /// to initialize it.
1219 fn init(l: *Legalize, ty: Type, parent_block: *Block) Result {
1220 return .{
1221 .inst = parent_block.add(l, .{
1222 .tag = .block,
1223 .data = .{ .ty_pl = .{
1224 .ty = Air.internedToRef(ty.toIntern()),
1225 .payload = undefined,
1226 } },
1227 }),
1228 .block = undefined,
1229 };
1230 }
1231
1232 fn finish(res: Result, l: *Legalize) Error!void {
1233 const data = &l.air_instructions.items(.data)[@intFromEnum(res.inst)];
1234 data.ty_pl.payload = try l.addBlockBody(res.block.body());
1235 }
1236};
1237
1152const Loop = struct {1238const Loop = struct {
1153 inst: Air.Inst.Index,1239 inst: Air.Inst.Index,
1154 block: Block,1240 block: Block,
src/Compilation.zig+1-1
...@@ -2529,6 +2529,7 @@ pub fn destroy(comp: *Compilation) void {...@@ -2529,6 +2529,7 @@ pub fn destroy(comp: *Compilation) void {
25292529
2530pub fn clearMiscFailures(comp: *Compilation) void {2530pub fn clearMiscFailures(comp: *Compilation) void {
2531 comp.alloc_failure_occurred = false;2531 comp.alloc_failure_occurred = false;
2532 comp.link_diags.flags = .{};
2532 for (comp.misc_failures.values()) |*value| {2533 for (comp.misc_failures.values()) |*value| {
2533 value.deinit(comp.gpa);2534 value.deinit(comp.gpa);
2534 }2535 }
...@@ -2795,7 +2796,6 @@ pub fn update(comp: *Compilation, main_progress_node: std.Progress.Node) !void {...@@ -2795,7 +2796,6 @@ pub fn update(comp: *Compilation, main_progress_node: std.Progress.Node) !void {
27952796
2796 if (anyErrors(comp)) {2797 if (anyErrors(comp)) {
2797 // Skip flushing and keep source files loaded for error reporting.2798 // Skip flushing and keep source files loaded for error reporting.
2798 comp.link_diags.flags = .{};
2799 return;2799 return;
2800 }2800 }
28012801
src/arch/x86_64/CodeGen.zig+1
...@@ -84,6 +84,7 @@ pub fn legalizeFeatures(target: *const std.Target) *const Air.Legalize.Features...@@ -84,6 +84,7 @@ pub fn legalizeFeatures(target: *const std.Target) *const Air.Legalize.Features
84 .scalarize_int_from_float = use_old,84 .scalarize_int_from_float = use_old,
85 .scalarize_int_from_float_optimized = use_old,85 .scalarize_int_from_float_optimized = use_old,
86 .scalarize_float_from_int = use_old,86 .scalarize_float_from_int = use_old,
87 .scalarize_select = true,
87 .scalarize_mul_add = use_old,88 .scalarize_mul_add = use_old,
8889
89 .unsplat_shift_rhs = false,90 .unsplat_shift_rhs = false,
test/behavior/select.zig-3
...@@ -41,8 +41,6 @@ test "@select arrays" {...@@ -41,8 +41,6 @@ test "@select arrays" {
41 if (builtin.zig_backend == .stage2_sparc64) return error.SkipZigTest; // TODO41 if (builtin.zig_backend == .stage2_sparc64) return error.SkipZigTest; // TODO
42 if (builtin.zig_backend == .stage2_riscv64) return error.SkipZigTest;42 if (builtin.zig_backend == .stage2_riscv64) return error.SkipZigTest;
43 if (builtin.zig_backend == .stage2_spirv64) return error.SkipZigTest;43 if (builtin.zig_backend == .stage2_spirv64) return error.SkipZigTest;
44 if (builtin.zig_backend == .stage2_x86_64 and
45 !comptime std.Target.x86.featureSetHas(builtin.cpu.features, .avx2)) return error.SkipZigTest;
4644
47 try comptime selectArrays();45 try comptime selectArrays();
48 try selectArrays();46 try selectArrays();
...@@ -70,7 +68,6 @@ fn selectArrays() !void {...@@ -70,7 +68,6 @@ fn selectArrays() !void {
70test "@select compare result" {68test "@select compare result" {
71 if (builtin.zig_backend == .stage2_riscv64) return error.SkipZigTest;69 if (builtin.zig_backend == .stage2_riscv64) return error.SkipZigTest;
72 if (builtin.zig_backend == .stage2_wasm) return error.SkipZigTest;70 if (builtin.zig_backend == .stage2_wasm) return error.SkipZigTest;
73 if (builtin.zig_backend == .stage2_x86_64) return error.SkipZigTest;
74 if (builtin.zig_backend == .stage2_llvm and builtin.cpu.arch == .hexagon) return error.SkipZigTest;71 if (builtin.zig_backend == .stage2_llvm and builtin.cpu.arch == .hexagon) return error.SkipZigTest;
7572
76 const S = struct {73 const S = struct {