Shoal #
This is a library that provides performant functions that allow for functional programming in Zig!
I made this so that I could feel more at home; I find myself reaching for functional... functions much more often now that I use Rust every day, additionally I also end up writing code worse in regular loops than being constrained to this array-comprehension-style. That being said, functional programming (other than in Rust) tends to take a performance hit unless the lang has some fun parallelization going on, so I've been quite careful to ensure that every function here compiles to the same or fewer instructions than what I would've written out by hand (with a couple of exceptions).
Take the following for example:
const shoal = @import("shoal");
fn product(pair: anytype, _: void) @TypeOf(pair[0]) {
return pair[0] * pair[1];
}
fn dot(a: []const f32, b: []const f32) ?f32 {
return shoal.over(a).zip(b).map({}, product).reduce(.Add);
}
Shoal is lazy, in that all those chain commands all roll up into what it will do by the end, ie. the terminal reduce here.
Nothing will be allocated as you can see, instruction-wise the loads go straight from (%rdi,%rax,4) into vector registers.
About the ?f32 consider that an empty slice can't sum to anything, so reduce can return null.
Let's see how it performs, like I said!
Why export and raw pointers, you ask? LLVM wants its own register stuff, I don't want us to have to pay the vmovd toll if having %xmm0 on one side and %edx on the other; tryna keep it equivalent - hence the export for same ABI on both sides.
const std = @import("std");
const shoal = @import("shoal");
const w = std.simd.suggestVectorLength(f32) orelse 4;
const unrolls = 4;
const V = @Vector(w, f32);
fn product(pair: anytype, _: void) @TypeOf(pair[0]) {
return pair[0] * pair[1];
}
pub export fn shoalDot(a: [*]const f32, b: [*]const f32, n: usize) f32 {
return shoal.over(a[0..n]).zip(b[0..n]).map({}, product).reduce(.Add) orelse 0;
}
fn tree(lanes: anytype) f32 {
const lanes_len = @typeInfo(@TypeOf(lanes)).vector.len;
if (lanes_len == 1) return lanes[0];
const half = lanes_len / 2;
return tree(std.simd.extract(lanes, 0, half) + std.simd.extract(lanes, half, half));
}
// If my handmade version of the dotting is inefficient in any way,
// please let me know,
// so that I can make the shoal-version even faster against a better analogue.
fn handDotting(a: []const f32, b: []const f32) ?f32 {
const n = @min(a.len, b.len);
if (n == 0) return null;
var acc: [unrolls]V = @splat(@as(V, @splat(0)));
var at: usize = 0;
while (at + unrolls * w <= n) : (at += unrolls * w) {
inline for (&acc, 0..) |*slot, k| {
const x: V = a[at + k * w ..][0..w].*;
const y: V = b[at + k * w ..][0..w].*;
slot.* += x * y;
}
}
var lanes = acc[0];
inline for (1..unrolls) |k| lanes += acc[k];
var total = tree(lanes);
while (at < n) : (at += 1) total += a[at] * b[at];
return total;
}
pub export fn handDot(a: [*]const f32, b: [*]const f32, n: usize) f32 {
return handDotting(a[0..n], b[0..n]) orelse 0;
}
test "Both sides should emit the same thing!" {
var a: [100]f32 = undefined;
var b: [100]f32 = undefined;
for (&a, &b, 0..) |*x, *y, i| {
x.* = @floatFromInt(i % 7);
y.* = @floatFromInt(i % 5);
}
for ([_]usize{ 0, 1, 7, 31, 32, 33, 100 }) |n| {
try std.testing.expectEqual(shoalDot(&a, &b, n), handDot(&a, &b, n));
}
}
zig test --dep shoal -Mroot=parity.zig -Mshoal=shoal/shoal.zig -mcpu=native
zig build-obj -OReleaseFast -mcpu=native --dep shoal -Mroot=parity.zig -Mshoal=shoal/shoal.zig -femit-bin=parity.o
objdump -d --disassemble=shoalDot parity.o
objdump -d --disassemble=handDot parity.o
Which will give us:
Shoal Hand-written
---- ----
vmovups (%rdi,%rax,4),%ymm4 vmovups (%rdi,%rax,4),%ymm4
vmovups 0x20(%rdi,%rax,4),%ymm5 vmovups 0x20(%rdi,%rax,4),%ymm5
vmovups 0x40(%rdi,%rax,4),%ymm6 vmovups 0x40(%rdi,%rax,4),%ymm6
vmovups 0x60(%rdi,%rax,4),%ymm7 vmovups 0x60(%rdi,%rax,4),%ymm7
vmulps (%rsi,%rax,4),%ymm4,%ymm4 vmulps (%rsi,%rax,4),%ymm4,%ymm4
vaddps %ymm4,%ymm3,%ymm3 vaddps %ymm4,%ymm2,%ymm2
vmulps 0x20(%rsi,%rax,4),%ymm5,%ymm4 vmulps 0x20(%rsi,%rax,4),%ymm5,%ymm4
vaddps %ymm4,%ymm2,%ymm2 vaddps %ymm4,%ymm3,%ymm3
vmulps 0x40(%rsi,%rax,4),%ymm6,%ymm4 vmulps 0x40(%rsi,%rax,4),%ymm6,%ymm4
vaddps %ymm4,%ymm1,%ymm1 vaddps %ymm4,%ymm1,%ymm1
vmulps 0x60(%rsi,%rax,4),%ymm7,%ymm4 vmulps 0x60(%rsi,%rax,4),%ymm7,%ymm4
vaddps %ymm4,%ymm0,%ymm0 vaddps %ymm4,%ymm0,%ymm0
add $0x20,%rax add $0x20,%rax
add $0x40,%rcx add $0x40,%rcx
cmp %rdx,%rcx cmp %rdx,%rcx
jbe 1c0 <shoalDot+0x30> jbe 30 <handDot+0x30>
Well it's a bit vapid to line them up and just point that at least line for line they're the same.
There is a difference! :nerd-face: You will notice how shoal accumulates into %ymm3 then %ymm2,
compared to how the hand-written one does into %ymm2 then %ymm3.
This was done on my x86 laptop, I'll try out aarch and we'll see how the instructions compare, later, when I can be bothered.
Anyway, there are some situations in which shoal is a little less efficient.
Take for example the fact that shoal will return ?Ts and take slices -
if you write the same kernel by hand using a raw pointer that returns just a f32,
there'll be no optional and no slice bounds,
so on -OReleaseSafe it'll beat shoal with 91 instructions vs our 105.
There's some extra gunk in there just for the optionality and slice boundality.
Install #
Sorry zig fetch isn't working right now, my bad on the Tangled Knot side.
One has to clone until I fix:
git clone https://knot.oyster.cafe/did:plc:dnbluu2civkq6kpperfgtqmk shoal
.dependencies = .{
.shoal = .{ .path = "../shoal" },
},
exe.root_module.addImport("shoal", b.dependency("shoal", .{
.target = target,
.optimize = optimize,
}).module("shoal"));
Starting a pipe #
shoal.over(&items) // An array, slice, or *const [n]T.
shoal.range(n) // 0, 1, .. n - 1, each being a usize.
shoal.fixedRange(u8, 3, 6) // 3, 4, 5, each being a u8.
caught.pipe() // Items in a Bounded.
...and yes, &items if you're savvy, even though just items by value would work too. For this example there's not really any difference than a couple more instructions compared to hand-made, but with something like zip for example it'll copy, and has many more instructions.
Range there will count to n - 1 but beware! It doesn't check n against the actual slive that the callback is indexing. Kaboom ensues if prepared incorrectly. When in doubt, range(rows.len) or better yet over(rows) if just the item alone is all that's needed, since then there's no index to get wrong there.
Adding stages #
.map(context, f) // f(item, context) giveth the next item.
.filter(context, predicate) // So, only keeping the items where predicate(item, context).
.compact() // Drop the nulls, turning ?U pipes into just U.
.zip(other) // Pair(L, R) per index (halting at the shorter side).
Context = ctx = Golang mentioned (not really).
Since Zig doesn't have closures, our context here is just how a callback can like capture anything at all. If you don't need anything other than a given item, {} context will do fine. That being said, shoal would refuse to build a run-time pipe when passing something like a 67 (being a comptime_int) so @as(f32, 67) unfortunately, if you're trying to just filter by a particular number, say.
const Shell = struct { name: []const u8, grams: f32 };
fn heavierThan(shell: Shell, worthless_mollusc_num: f32) bool {
return shell.grams > worthless_mollusc_num;
}
fn nameOf(shell: Shell, _: void) []const u8 {
return shell.name;
}
const heavy = shoal.over(&shells)
.filter(@as(f32, 67), heavierThan)
.map({}, nameOf)
.toBounded();
Ending a pipe #
Seeing as the terminals are the only actual way to make the pipe-chain wake up and do all the work that's curried into it, I would say these are quite important. Let's go through how one can end a pipe, grouped by the actual result we want out:
One value out: #
fold(initial, context, f):f(accumulator, item, context), lefty to righty.tryFold(initial, context, f): the same, but wherefcan fail, stopping at the first error.reduce(op): one of.Add,.Mul,.Min,.Max,.And,.Or,.Xor, as?Out, four-accumulators-wide.reduceUnrolled(op, k):reducebut with the accumulator count spelled out.count(): count how many items reach the end.
One item out, or nothing: #
find(context, predicate): the first passing item, as?Out.findMap(context, f): the first non-nullf(item, context).any(context, predicate)&all(context, predicate): both short-circuit.allMap(context, f): every result where every item passes,?[n]U.
Many items out: #
toArray():[n]Out, which needs a length known at compile-time.toBounded(): aBounded(n, Out), which needs an upper bound.toOwnedSlice(gpa):[]Out, allocated then shrunk to survivor-size.toBuffer(buffer): BYO buffer to fill into.writeInto(out): one item per survivor into a given buffer, returning count of how many written.scanInto(out, initial, op): the runningopinstead, also with count returned.tryMap(context, f):[n]Uwherefcan fail.
Nothing out: #
each(context, f)&tryEach(context, f): for effects.
Unfortunately we can't mix & match everything, mainly because of some functions needing the size of result at compile-time.
Specifically, tryMap, allMap, zip, toArray, and toBounded will all error out if given a slice as source or are combined with filter anywere in the pipe-chain.
range's length is known at run-time, hence why I even made toOwnedSlice in the first place to collect.
Here's something else you should be aware of. Take this example:
var out: [5]f32 = undefined;
const written = shoal.over(&a).zip(&b).map({}, product).writeInto(&out);
var sums: [5]f32 = undefined;
_ = shoal.over(&a).scanInto(&sums, 0, .Add);
Both out and sums write out one item per survivor, returning count. However, both assert out.len >= n where n is the count of indices in the source, not survivor count.
Therefore we have to make the buffer the same length as the source even if our final fill is shorter.
Reading pipes multiple times #
If we stick 2 or more terminals on the end of a chain, the whole chain will be run again from the source, because the piping doesn't contain any items, right - it's lazy like I said. So when wanting to do multiple "final" ops on a single array, we can make some adjustments to make it not re-trace its steps.
Let's say we wanna softmaxx for example!
fn expf(item: f32, _: void) f32 {
return @exp(item);
}
fn divided(item: f32, by: f32) f32 {
return item / by;
}
Here's a naive approach where expf gets called over everything twice, once for total, once for fill:
const warmed = over(&x).map({}, expf);
const total = warmed.reduce(.Add) orelse 1;
var odds: [128]f32 = undefined;
_ = warmed.map(total, divided).toBuffer(&odds);
Now here's a version where we fill first! expf only runs the once, and then when we read later we do it off the numbers already stored:
var exps: [128]f32 = undefined;
const warmed = over(&x).map({}, expf).toBuffer(&exps);
const total = over(warmed).reduce(.Add) orelse 1;
var odds: [128]f32 = undefined;
_ = over(warmed).map(total, divided).toBuffer(&odds);
I wouldn't even call the following a rule of thumb, but what I've been doing is toBuffer(&buf) then over(fill) when I own a buffer whose len only exists at run-time; toBounded().pipe() when the len is compile-time; toOwnedSlice(gpa) when I'm okay with allocating.
Oh and last thing about performing that above softmax is that
foldwith a small struct-as-accumulator can yield us the mean and the mean-square in one single pass, neat!
Callbacks & width #
When you make a plain function fn (T, C) U
- pro: error messages will be very distinct, a mismatched callback for example will be loud and clear what's wrong and where.
- con: it'll run one item at a time.
fn doubled(item: u32, by: u32) u32 {
return item * by;
}
When you make a generic fn
- pro: it'll run at element's vector width.
- con: if you hook it up wrong the error might be lost in a sea of abstract nonsense.
fn scaled(item: anytype, by: @TypeOf(item)) @TypeOf(item) {
return item * by;
}
We are very advanced; we even have lanes.
reduce, reduceUnrolled, writeInto, and scanInto can have parallel lanes of execution under the hood!
Not to say that they actually do anything parallel, I just mean that the lanes aren't dependent on each other at all,
and makes things quite flexible to juggle. We could make them properly parallel later.
Unfortunately including filter or a compact will ruin the illusion and the whole thing has to go through one at a time by nature.
By nature of the predicate having to decide index-by-index, I mean.
The widths of the lanes, at least on my laptop, are as follows:
f32andu32: 8 lanesu8: 32usize: 4- a struct or a slice: 1
Did you know that reduce and fold can disagree in their answers of a relatively equivalent construction?
This is because reduce will reassociate (as in (a + b) + c becoming a + (b + c)) whereas fold does not.
So as a rule of thumb, if you're using f32s (the difference doesn't apply to integers) and ordering has to be lefty-to-righty, use fold.
One-n-done forms #
If you don't need to chain anything:
shoal.map(&items, @as(u32, 2), doubled); // [n]U
shoal.filter(&items, @as(u32, 2), divides); // Bounded(n, T)
shoal.fold(&items, 0, {}, add);
shoal.zip(&counts, &shells); // [n]Pair(L, R)
You will be delighted to know that the following functions can just be called themselves like the above example: find, findMap, all, allMap, any, each, tryEach, tryFold, and tryMap. They only work with plain fns and not generics - for generics, you'll have to go back to an over(&items).map(..).
There are some functions that are only free radicals, since they don't have any items to pipe over.
mapFields(U, T, f):f(T, name)once per field of the struct typeT, into[n]U.update(base, changed): a copy ofbasewith each field thatchangeddefines written over.repeat(most, context, step)&tryRepeat:step(context)up tomosttimes, stop at the first false, return count of how many ran.times(count, context, step): exactlycounttimes, without stopping early.
Bounded and Queue #
"I'm trying to use 'filter' but what's with this 'Bounded'"?
Bounded(n, T) is just an [n]T along with a length that will never be greater than n, and is what a filtered chain collects into.
var caught: shoal.Bounded(4, u32) = .empty;
_ = caught.push(9);
_ = caught.push(2);
const ranked = caught.sort({}, ascending);
const total = ranked.pipe().reduce(.Add);
So
pushreturns an index that it writes to, unless the buffer is full in which case it returns null.
A Queue is also a bound, but with the oldest item shunted out first.
Pushing to a queue will return whatever it dropped out the other end in order to make room (as ?T).
Thus shift will suck one out the front (again as ?T).
at(index) read op returns ?*T, and drop(count) well you know the rest.
Since Bounded.buffer/Bounded.len are both public, a caller is able to write len beyond most -
so pipe() will assert len <= most just in case.
Chaining from your own type #
Zig doesn't have a |> of course, so we just method-chain as you can see.
Look at what we can do to have a Rust-like iter experience!:
pub const Deck = struct {
cards: []const Card,
pub inline fn iter(self: Deck) @TypeOf(shoal.over(self.cards)) {
return shoal.over(self.cards);
}
};
Then ta-da, we can deck.iter().map({}, Card.value).reduce(.Add)!
Status #
This project is very early, so the names of the different functions still may change. When I'm making new functions that I need, I name by
- Gleam
- Rust
- Haskell
- Halide/APL
in that order, depending on if the lang has the conceptual thing I'm referring to.
Thank you for reading!