authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2023-10-01 00:07:57-07:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2023-10-01 00:07:57-07:00
log376242e586a04b1d2f8f30a329eba3275e0e3a87
tree9da6b13546c1addb1aaea3be79e9d0f4f7cfd670
parent9a001e1f7cc878579f1c0a614ac0124bfdc58332
parent08635f08a9afe38b2b9eec3ce7ccb583c402252a
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #17161 from tiehuis/vectorize-index-of-scalar

std.mem: add vectorized indexOfScalarPos and indexOfSentinel

1 files changed, 162 insertions(+), 4 deletions(-)

lib/std/mem.zig+162-4
......@@ -953,14 +953,105 @@ test "len" {
953953 try testing.expect(len(c_ptr) == 2);
954954}
955955
956pub fn indexOfSentinel(comptime Elem: type, comptime sentinel: Elem, ptr: [*:sentinel]const Elem) usize {
956const backend_supports_vectors = switch (builtin.zig_backend) {
957 .stage2_llvm, .stage2_c => true,
958 else => false,
959};
960
961pub fn indexOfSentinel(comptime T: type, comptime sentinel: T, p: [*:sentinel]const T) usize {
957962 var i: usize = 0;
958 while (ptr[i] != sentinel) {
963
964 if (backend_supports_vectors and
965 !@inComptime() and
966 (@typeInfo(T) == .Int or @typeInfo(T) == .Float) and std.math.isPowerOfTwo(@bitSizeOf(T)))
967 {
968 switch (@import("builtin").cpu.arch) {
969 // The below branch assumes that reading past the end of the buffer is valid, as long
970 // as we don't read into a new page. This should be the case for most architectures
971 // which use paged memory, however should be confirmed before adding a new arch below.
972 .aarch64, .x86, .x86_64 => if (comptime std.simd.suggestVectorSize(T)) |block_len| {
973 comptime std.debug.assert(std.mem.page_size % block_len == 0);
974 const Block = @Vector(block_len, T);
975 const mask: Block = @splat(sentinel);
976
977 // First block may be unaligned
978 const start_addr = @intFromPtr(&p[i]);
979 const offset_in_page = start_addr & (std.mem.page_size - 1);
980 if (offset_in_page < std.mem.page_size - block_len) {
981 // Will not read past the end of a page, full block.
982 const block: Block = p[i..][0..block_len].*;
983 const matches = block == mask;
984 if (@reduce(.Or, matches)) {
985 return i + std.simd.firstTrue(matches).?;
986 }
987
988 i += (std.mem.alignForward(usize, start_addr, @alignOf(Block)) - start_addr) / @sizeOf(T);
989 } else {
990 // Would read over a page boundary. Per-byte at a time until aligned or found.
991 // 0.39% chance this branch is taken for 4K pages at 16b block length.
992 //
993 // An alternate strategy is to do read a full block (the last in the page) and
994 // mask the entries before the pointer.
995 while ((@intFromPtr(&p[i]) & (@alignOf(Block) - 1)) != 0) : (i += 1) {
996 if (p[i] == sentinel) return i;
997 }
998 }
999
1000 std.debug.assert(std.mem.isAligned(@intFromPtr(&p[i]), @alignOf(Block)));
1001 while (true) {
1002 const block: *const Block = @ptrCast(@alignCast(p[i..][0..block_len]));
1003 const matches = block.* == mask;
1004 if (@reduce(.Or, matches)) {
1005 return i + std.simd.firstTrue(matches).?;
1006 }
1007 i += block_len;
1008 }
1009 },
1010 else => {},
1011 }
1012 }
1013
1014 while (p[i] != sentinel) {
9591015 i += 1;
9601016 }
9611017 return i;
9621018}
9631019
1020test "indexOfSentinel vector paths" {
1021 const Types = [_]type{ u8, u16, u32, u64 };
1022 const allocator = std.testing.allocator;
1023
1024 inline for (Types) |T| {
1025 const block_len = comptime std.simd.suggestVectorSize(T) orelse continue;
1026
1027 // Allocate three pages so we guarantee a page-crossing address with a full page after
1028 const memory = try allocator.alloc(T, 3 * std.mem.page_size / @sizeOf(T));
1029 defer allocator.free(memory);
1030 @memset(memory, 0xaa);
1031
1032 // Find starting page-alignment = 0
1033 var start: usize = 0;
1034 const start_addr = @intFromPtr(&memory);
1035 start += (std.mem.alignForward(usize, start_addr, std.mem.page_size) - start_addr) / @sizeOf(T);
1036 try testing.expect(start < std.mem.page_size / @sizeOf(T));
1037
1038 // Validate all sub-block alignments
1039 const search_len = std.mem.page_size / @sizeOf(T);
1040 memory[start + search_len] = 0;
1041 for (0..block_len) |offset| {
1042 try testing.expectEqual(search_len - offset, indexOfSentinel(T, 0, @ptrCast(&memory[start + offset])));
1043 }
1044 memory[start + search_len] = 0xaa;
1045
1046 // Validate page boundary crossing
1047 const start_page_boundary = start + (std.mem.page_size / @sizeOf(T));
1048 memory[start_page_boundary + block_len] = 0;
1049 for (0..block_len) |offset| {
1050 try testing.expectEqual(2 * block_len - offset, indexOfSentinel(T, 0, @ptrCast(&memory[start_page_boundary - block_len + offset])));
1051 }
1052 }
1053}
1054
9641055/// Returns true if all elements in a slice are equal to the scalar value provided
9651056pub fn allEqual(comptime T: type, slice: []const T, scalar: T) bool {
9661057 for (slice) |item| {
......@@ -1016,12 +1107,79 @@ pub fn lastIndexOfScalar(comptime T: type, slice: []const T, value: T) ?usize {
10161107
10171108pub fn indexOfScalarPos(comptime T: type, slice: []const T, start_index: usize, value: T) ?usize {
10181109 if (start_index >= slice.len) return null;
1019 for (slice[start_index..], start_index..) |c, i| {
1020 if (c == value) return i;
1110
1111 var i: usize = start_index;
1112 if (backend_supports_vectors and
1113 !@inComptime() and
1114 (@typeInfo(T) == .Int or @typeInfo(T) == .Float) and std.math.isPowerOfTwo(@bitSizeOf(T)))
1115 {
1116 if (comptime std.simd.suggestVectorSize(T)) |block_len| {
1117 // For Intel Nehalem (2009) and AMD Bulldozer (2012) or later, unaligned loads on aligned data result
1118 // in the same execution as aligned loads. We ignore older arch's here and don't bother pre-aligning.
1119 //
1120 // Use `comptime std.simd.suggestVectorSize(T)` to get the same alignment as used in this function
1121 // however this usually isn't necessary unless your arch has a performance penalty due to this.
1122 //
1123 // This may differ for other arch's. Arm for example costs a cycle when loading across a cache
1124 // line so explicit alignment prologues may be worth exploration.
1125
1126 // Unrolling here is ~10% improvement. We can then do one bounds check every 2 blocks
1127 // instead of one which adds up.
1128 const Block = @Vector(block_len, T);
1129 if (i + 2 * block_len < slice.len) {
1130 const mask: Block = @splat(value);
1131 while (true) {
1132 inline for (0..2) |_| {
1133 const block: Block = slice[i..][0..block_len].*;
1134 const matches = block == mask;
1135 if (@reduce(.Or, matches)) {
1136 return i + std.simd.firstTrue(matches).?;
1137 }
1138 i += block_len;
1139 }
1140 if (i + 2 * block_len >= slice.len) break;
1141 }
1142 }
1143
1144 // {block_len, block_len / 2} check
1145 inline for (0..2) |j| {
1146 const block_x_len = block_len / (1 << j);
1147 comptime if (block_x_len < 4) break;
1148
1149 const BlockX = @Vector(block_x_len, T);
1150 if (i + block_x_len < slice.len) {
1151 const mask: BlockX = @splat(value);
1152 const block: BlockX = slice[i..][0..block_x_len].*;
1153 const matches = block == mask;
1154 if (@reduce(.Or, matches)) {
1155 return i + std.simd.firstTrue(matches).?;
1156 }
1157 i += block_x_len;
1158 }
1159 }
1160 }
1161 }
1162
1163 for (slice[i..], i..) |c, j| {
1164 if (c == value) return j;
10211165 }
10221166 return null;
10231167}
10241168
1169test "indexOfScalarPos" {
1170 const Types = [_]type{ u8, u16, u32, u64 };
1171
1172 inline for (Types) |T| {
1173 var memory: [64 / @sizeOf(T)]T = undefined;
1174 @memset(&memory, 0xaa);
1175 memory[memory.len - 1] = 0;
1176
1177 for (0..memory.len) |i| {
1178 try testing.expectEqual(memory.len - i - 1, indexOfScalarPos(T, memory[i..], 0, 0).?);
1179 }
1180 }
1181}
1182
10251183pub fn indexOfAny(comptime T: type, slice: []const T, values: []const T) ?usize {
10261184 return indexOfAnyPos(T, slice, 0, values);
10271185}