lib/simd/src/intdiv.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 
  3 pub const Divisor = ScalarDivisor(u32);
  4 pub const Divisor64 = ScalarDivisor(u64);
  5 
  6 pub fn ScalarDivisor(comptime T: type) type {
  7     requireUnsigned(T);
  8     if (T != u32 and T != u64) {
  9         @compileError("Highway scalar divisors require u32 or u64");
 10     }
 11     return struct {
 12         params: DivisorParamsU(T),
 13 
 14         const Self = @This();
 15 
 16         pub fn init(divisor: T) Self {
 17             return .{ .params = computeDivisorParams(T, divisor) };
 18         }
 19 
 20         pub fn getDivisor(self: Self) T {
 21             return self.params.divisor;
 22         }
 23 
 24         pub fn divide(self: Self, dividend: T) T {
 25             return intDivScalar(T, dividend, self.params);
 26         }
 27 
 28         pub fn remainder(self: Self, dividend: T) T {
 29             return dividend - self.divide(dividend) * self.params.divisor;
 30         }
 31     };
 32 }
 33 
 34 pub fn DivisorParamsU(comptime T: type) type {
 35     requireUnsigned(T);
 36     const M = Multiplier(T);
 37     return struct {
 38         multiplier: M,
 39         shift2: u8,
 40         is_pow2: bool,
 41         pow2_shift: u8,
 42         divisor: T,
 43     };
 44 }
 45 
 46 pub fn DivisorParamsS(comptime T: type) type {
 47     requireSigned(T);
 48     const M = Multiplier(T);
 49     return struct {
 50         multiplier: M,
 51         shift: u8,
 52         divisor: T,
 53         dsign: T,
 54         is_pow2: bool,
 55         is_neg_one: bool,
 56         pow2_shift: u8,
 57     };
 58 }
 59 
 60 pub fn DivisorParams(comptime T: type) type {
 61     requireInteger(T);
 62     return if (isSigned(T)) DivisorParamsS(T) else DivisorParamsU(T);
 63 }
 64 
 65 pub fn computeDivisorParams(comptime T: type, divisor: T) DivisorParams(T) {
 66     std.debug.assert(divisor != 0);
 67     return if (comptime isSigned(T))
 68         computeSignedParams(T, divisor)
 69     else
 70         computeUnsignedParams(T, divisor);
 71 }
 72 
 73 pub fn intDiv(comptime D: type, dividend: D.Vector, params: DivisorParams(D.Lane)) D.Vector {
 74     requireInteger(D.Lane);
 75     std.debug.assert(params.divisor != 0);
 76     var result: D.Vector = undefined;
 77     inline for (0..D.lane_count) |index| {
 78         result[index] = intDivScalar(D.Lane, dividend[index], params);
 79     }
 80     return result;
 81 }
 82 
 83 pub fn intDivFloor(comptime D: type, dividend: D.Vector, params: DivisorParams(D.Lane)) D.Vector {
 84     requireInteger(D.Lane);
 85     std.debug.assert(params.divisor != 0);
 86     var result: D.Vector = undefined;
 87     inline for (0..D.lane_count) |index| {
 88         result[index] = intDivFloorScalar(D.Lane, dividend[index], params);
 89     }
 90     return result;
 91 }
 92 
 93 pub fn divideByScalar(comptime D: type, dividend: D.Vector, divisor: D.Lane) D.Vector {
 94     return intDiv(D, dividend, computeDivisorParams(D.Lane, divisor));
 95 }
 96 
 97 pub fn floorDivideByScalar(comptime D: type, dividend: D.Vector, divisor: D.Lane) D.Vector {
 98     return intDivFloor(D, dividend, computeDivisorParams(D.Lane, divisor));
 99 }
100 
101 pub fn divideArrayByScalar(comptime T: type, array: []T, divisor: T) void {
102     requireInteger(T);
103     std.debug.assert(divisor != 0);
104     const params = computeDivisorParams(T, divisor);
105     for (array) |*value| value.* = intDivScalar(T, value.*, params);
106 }
107 
108 pub fn floorDivideArrayByScalar(comptime T: type, array: []T, divisor: T) void {
109     requireInteger(T);
110     std.debug.assert(divisor != 0);
111     const params = computeDivisorParams(T, divisor);
112     for (array) |*value| value.* = intDivFloorScalar(T, value.*, params);
113 }
114 
115 pub fn div128HighBy(high: u64, divisor: u64) u64 {
116     std.debug.assert(divisor != 0);
117     std.debug.assert(high < divisor);
118     return @intCast((@as(u128, high) << 64) / divisor);
119 }
120 
121 fn computeUnsignedParams(comptime T: type, divisor: T) DivisorParamsU(T) {
122     if (std.math.isPowerOfTwo(divisor)) {
123         return .{
124             .multiplier = 1,
125             .shift2 = 0,
126             .is_pow2 = true,
127             .pow2_shift = @intCast(@ctz(divisor)),
128             .divisor = divisor,
129         };
130     }
131 
132     const bits = @bitSizeOf(T);
133     const l: usize = @as(usize, std.math.log2_int(T, divisor - 1)) + 1;
134     const wide_divisor: u128 = divisor;
135     const numerator = ((@as(u128, 1) << @intCast(l)) - wide_divisor) << @intCast(bits);
136     const multiplier: Multiplier(T) = @intCast(numerator / wide_divisor + 1);
137     return .{
138         .multiplier = multiplier,
139         .shift2 = @intCast(l - 1),
140         .is_pow2 = false,
141         .pow2_shift = 0,
142         .divisor = divisor,
143     };
144 }
145 
146 fn computeSignedParams(comptime T: type, divisor: T) DivisorParamsS(T) {
147     const U = @Int(.unsigned, @bitSizeOf(T));
148     const raw: U = @bitCast(divisor);
149     const abs_divisor: U = if (divisor < 0) 0 -% raw else raw;
150     const dsign: T = if (divisor < 0) -1 else 0;
151     if (std.math.isPowerOfTwo(abs_divisor)) {
152         return .{
153             .multiplier = 1,
154             .shift = 0,
155             .divisor = divisor,
156             .dsign = dsign,
157             .is_pow2 = true,
158             .is_neg_one = divisor == -1,
159             .pow2_shift = @intCast(@ctz(abs_divisor)),
160         };
161     }
162 
163     const bits = @bitSizeOf(T);
164     const shift: usize = std.math.log2_int(U, abs_divisor - 1);
165     const wide_multiplier = (@as(u128, 1) << @intCast(bits + shift)) / abs_divisor + 1;
166     const low: U = @truncate(wide_multiplier);
167     const signed_low: T = @bitCast(low);
168     return .{
169         .multiplier = @intCast(signed_low),
170         .shift = @intCast(shift),
171         .divisor = divisor,
172         .dsign = dsign,
173         .is_pow2 = false,
174         .is_neg_one = divisor == -1,
175         .pow2_shift = 0,
176     };
177 }
178 
179 fn intDivScalar(comptime T: type, dividend: T, params: DivisorParams(T)) T {
180     if (comptime isSigned(T)) return intDivSignedScalar(T, dividend, params);
181     if (params.is_pow2) return dividend >> @intCast(params.pow2_shift);
182     const W = Product(T);
183     const product: W = @as(W, dividend) * @as(W, params.multiplier);
184     const high: T = @intCast(product >> @intCast(@bitSizeOf(T)));
185     const sum = high + ((dividend - high) >> 1);
186     return sum >> @intCast(params.shift2);
187 }
188 
189 fn intDivSignedScalar(comptime T: type, dividend: T, params: DivisorParamsS(T)) T {
190     const U = @Int(.unsigned, @bitSizeOf(T));
191     if (params.is_pow2) {
192         if (params.pow2_shift == 0) return (dividend ^ params.dsign) -% params.dsign;
193         const mask_bits = (@as(U, 1) << @intCast(params.pow2_shift)) - 1;
194         const mask: T = @bitCast(mask_bits);
195         const sign = dividend >> @intCast(@bitSizeOf(T) - 1);
196         const bias = sign & mask;
197         const quotient = (dividend +% bias) >> @intCast(params.pow2_shift);
198         return (quotient ^ params.dsign) -% params.dsign;
199     }
200 
201     const W = Product(T);
202     const product: W = @as(W, dividend) * @as(W, params.multiplier);
203     const high: T = @intCast(product >> @intCast(@bitSizeOf(T)));
204     var quotient = dividend +% high;
205     quotient >>= @intCast(params.shift);
206     quotient -%= dividend >> @intCast(@bitSizeOf(T) - 1);
207     return (quotient ^ params.dsign) -% params.dsign;
208 }
209 
210 fn intDivFloorScalar(comptime T: type, dividend: T, params: DivisorParams(T)) T {
211     if (comptime !isSigned(T)) return intDivScalar(T, dividend, params);
212     if (params.is_neg_one and dividend == std.math.minInt(T)) return 0;
213     const quotient = intDivScalar(T, dividend, params);
214     const remainder_nonzero = quotient *% params.divisor != dividend;
215     const adjust: T = if (remainder_nonzero and (dividend < 0) != (params.divisor < 0)) 1 else 0;
216     return quotient - adjust;
217 }
218 
219 fn Multiplier(comptime T: type) type {
220     const info = @typeInfo(T).int;
221     const bits = if (info.bits < 32) info.bits * 2 else info.bits;
222     return @Int(info.signedness, bits);
223 }
224 
225 fn Product(comptime T: type) type {
226     const info = @typeInfo(T).int;
227     return @Int(info.signedness, info.bits * 2);
228 }
229 
230 fn isSigned(comptime T: type) bool {
231     return @typeInfo(T).int.signedness == .signed;
232 }
233 
234 fn requireInteger(comptime T: type) void {
235     switch (@typeInfo(T)) {
236         .int => |info| if (info.bits != 8 and info.bits != 16 and info.bits != 32 and info.bits != 64) {
237             @compileError("integer division requires 8/16/32/64-bit integer lanes");
238         },
239         else => @compileError("integer division requires integer lanes"),
240     }
241 }
242 
243 fn requireUnsigned(comptime T: type) void {
244     requireInteger(T);
245     if (isSigned(T)) @compileError("unsigned divisor parameters require unsigned integers");
246 }
247 
248 fn requireSigned(comptime T: type) void {
249     requireInteger(T);
250     if (!isSigned(T)) @compileError("signed divisor parameters require signed integers");
251 }
252 
253 test "Highway integer divisor parameters retain special cases" {
254     const unsigned_three = computeDivisorParams(u8, 3);
255     try std.testing.expectEqual(@as(u16, 86), unsigned_three.multiplier);
256     try std.testing.expectEqual(@as(u8, 1), unsigned_three.shift2);
257     try std.testing.expect(!unsigned_three.is_pow2);
258 
259     const u16_params = computeDivisorParams(u16, 16);
260     try std.testing.expect(u16_params.is_pow2);
261     try std.testing.expectEqual(@as(u8, 4), u16_params.pow2_shift);
262 
263     const s3 = computeDivisorParams(i8, 3);
264     try std.testing.expectEqual(@as(i16, -85), s3.multiplier);
265     try std.testing.expectEqual(@as(u8, 1), s3.shift);
266     try std.testing.expectEqual(@as(i8, 0), s3.dsign);
267 
268     const neg = computeDivisorParams(i32, -1);
269     try std.testing.expect(neg.is_pow2);
270     try std.testing.expect(neg.is_neg_one);
271     try std.testing.expectEqual(@as(i32, -1), neg.dsign);
272 }
273 
274 test "Highway integer division covers every lane width and signed edge" {
275     const simd = @import("root.zig");
276     inline for (.{ u8, u16, u32, u64 }) |T| {
277         const D = simd.FixedTag(T, 8);
278         const values: D.Vector = .{ 0, 1, 2, 3, 7, 31, std.math.maxInt(T) - 1, std.math.maxInt(T) };
279         const result = divideByScalar(D, values, 7);
280         inline for (0..D.lane_count) |index| {
281             try std.testing.expectEqual(values[index] / 7, result[index]);
282         }
283     }
284 
285     inline for (.{ i8, i16, i32, i64 }) |T| {
286         const D = simd.FixedTag(T, 8);
287         const values: D.Vector = .{ std.math.minInt(T), -31, -8, -1, 0, 1, 8, std.math.maxInt(T) };
288         const trunc_result = divideByScalar(D, values, -7);
289         const floor_result = floorDivideByScalar(D, values, -7);
290         inline for (0..D.lane_count) |index| {
291             try std.testing.expectEqual(@divTrunc(values[index], @as(T, -7)), trunc_result[index]);
292             const quotient = @divTrunc(values[index], @as(T, -7));
293             const remainder = @rem(values[index], @as(T, -7));
294             const expected_floor = quotient - @as(T, if (remainder != 0 and values[index] >= 0) 1 else 0);
295             try std.testing.expectEqual(expected_floor, floor_result[index]);
296         }
297         const by_neg_one = divideByScalar(D, @as(D.Vector, @splat(std.math.minInt(T))), -1);
298         try std.testing.expectEqual(std.math.minInt(T), by_neg_one[0]);
299         const floor_neg_one = floorDivideByScalar(D, @as(D.Vector, @splat(std.math.minInt(T))), -1);
300         try std.testing.expectEqual(@as(T, 0), floor_neg_one[0]);
301     }
302 }
303 
304 test "Highway precomputed division covers power two and maximal divisors" {
305     const simd = @import("root.zig");
306     inline for (.{ u8, u16, u32, u64 }) |T| {
307         const D = simd.FixedTag(T, 8);
308         const dividends: D.Vector = .{ 0, 1, 2, 3, 7, 31, std.math.maxInt(T) - 1, std.math.maxInt(T) };
309         const divisors = [_]T{ 1, 2, 3, 16, std.math.maxInt(T) };
310         for (divisors) |divisor| {
311             const result = intDiv(D, dividends, computeDivisorParams(T, divisor));
312             inline for (0..D.lane_count) |index| {
313                 try std.testing.expectEqual(dividends[index] / divisor, result[index]);
314             }
315         }
316     }
317 
318     inline for (.{ i8, i16, i32, i64 }) |T| {
319         const D = simd.FixedTag(T, 8);
320         const dividends: D.Vector = .{ std.math.minInt(T), -31, -8, -1, 0, 1, 8, std.math.maxInt(T) };
321         const divisors = [_]T{ std.math.minInt(T), -16, -3, -1, 1, 2, 7, std.math.maxInt(T) };
322         for (divisors) |divisor| {
323             const params = computeDivisorParams(T, divisor);
324             const trunc_result = intDiv(D, dividends, params);
325             const floor_result = intDivFloor(D, dividends, params);
326             inline for (0..D.lane_count) |index| {
327                 const dividend = dividends[index];
328                 if (dividend == std.math.minInt(T) and divisor == -1) {
329                     try std.testing.expectEqual(std.math.minInt(T), trunc_result[index]);
330                     try std.testing.expectEqual(@as(T, 0), floor_result[index]);
331                 } else {
332                     const quotient = @divTrunc(dividend, divisor);
333                     const remainder = @rem(dividend, divisor);
334                     const adjust: T = if (remainder != 0 and (dividend < 0) != (divisor < 0)) 1 else 0;
335                     try std.testing.expectEqual(quotient, trunc_result[index]);
336                     try std.testing.expectEqual(quotient - adjust, floor_result[index]);
337                 }
338             }
339         }
340     }
341 }
342 
343 test "Highway array division handles tails without allocation" {
344     var trunc_values = [_]i32{ -10, -9, -1, 0, 1, 9, 10 };
345     divideArrayByScalar(i32, &trunc_values, 3);
346     try std.testing.expectEqualSlices(i32, &.{ -3, -3, 0, 0, 0, 3, 3 }, &trunc_values);
347 
348     var floor_values = [_]i32{ -10, -9, -1, 0, 1, 9, 10 };
349     floorDivideArrayByScalar(i32, &floor_values, 3);
350     try std.testing.expectEqualSlices(i32, &.{ -4, -3, -1, 0, 0, 3, 3 }, &floor_values);
351 }
352 
353 test "Highway scalar divisors match exhaustive small and boundary quadrants" {
354     inline for (.{ u32, u64 }) |T| try verifyScalarDivisor(T);
355 }
356 
357 fn verifyScalarDivisor(comptime T: type) !void {
358     const Scalar = ScalarDivisor(T);
359     for (1..256) |divisor_value| {
360         const divisor: T = @intCast(divisor_value);
361         const scalar = Scalar.init(divisor);
362         try std.testing.expectEqual(divisor, scalar.getDivisor());
363         for (0..256) |dividend_value| {
364             try expectScalarDivision(T, scalar, @intCast(dividend_value));
365         }
366         for (0..256) |offset| {
367             try expectScalarDivision(T, scalar, std.math.maxInt(T) - @as(T, @intCast(offset)));
368         }
369     }
370     for (0..256) |divisor_offset| {
371         const divisor = std.math.maxInt(T) - @as(T, @intCast(divisor_offset));
372         const scalar = Scalar.init(divisor);
373         for (0..256) |dividend_value| {
374             try expectScalarDivision(T, scalar, @intCast(dividend_value));
375         }
376         for (0..256) |dividend_offset| {
377             try expectScalarDivision(
378                 T,
379                 scalar,
380                 std.math.maxInt(T) - @as(T, @intCast(dividend_offset)),
381             );
382         }
383     }
384 }
385 
386 fn expectScalarDivision(
387     comptime T: type,
388     scalar: ScalarDivisor(T),
389     dividend: T,
390 ) !void {
391     try std.testing.expectEqual(dividend / scalar.getDivisor(), scalar.divide(dividend));
392     try std.testing.expectEqual(dividend % scalar.getDivisor(), scalar.remainder(dividend));
393 }
394 
395 test "Highway 128 by 64 helper returns the low fitting quotient" {
396     try std.testing.expectEqual(@as(u64, 0x5555555555555555), div128HighBy(1, 3));
397     try std.testing.expectEqual(@as(u64, 1) << 63, div128HighBy(@as(u64, 1) << 62, @as(u64, 1) << 63));
398     try std.testing.expectEqual(@as(u64, 1), div128HighBy(1, std.math.maxInt(u64)));
399 }