Came over from Rust and my god is the compile-time programming in Zig incredible
Anyone have fun or unique use-cases of it they’d like to show off? I’m just getting started with a few ideas
6 Likes
Totally! I love comptime strings - I’m a bit of a broken record here but I’m working on a tensor library that does einsum operations (so “X_ijk, Y_klm → ijlm” kind of thing).
Here’s a comptime parser for writing the loop permutation path. I’m very excited about optimizing this in the future because it’s going to be very easy:
fn IndicesPair(comptime lRank: usize, comptime rRank: usize) type {
return struct {
lhs : [lRank]SizeAndStrideType = undefined,
rhs : [rRank]SizeAndStrideType = undefined,
};
}
>fn contractionParse(
comptime lRank: usize,
comptime rRank: usize,
comptime str: [] const u8
) IndicesPair(lRank, rRank) {
comptime var index: usize = 0;
// reference for array operator
const arrow: [] const u8 = "->";
comptime var a: usize = 0;
comptime var b: usize = 0;
// mark one before the arrow and one after the arrow
inline while(index < str.len) : (index += 1) {
if(str[index] == arrow[0]) { a = index; }
if(str[index] == arrow[1]) { b = index; }
}
///////////////////////////////////////
// check for valid infix arrow operator
if((a + 1) != b) {
@compileError("Malformed arrow operator: " ++ str);
}
if(a == 0 or b > (str.len - 2)) {
@compileError("Arrow must be used as infix operator: " ++ str);
}
const lhs = str[0..a];
const rhs = str[b+1..];
if (lhs.len == 0) {
@compileError("Empty left-side operand: " ++ str);
}
if (rhs.len == 0) {
@compileError("Empty right-side operand: " ++ str);
}
if(lhs.len != lRank) {
@compileError("Provided indices do not match left-side operand rank: " ++ lhs);
}
if(rhs.len != rRank) {
@compileError("Provided indices do not match right-side operand rank: " ++ rhs);
}
////////////////////////////////////////
// build permutation contraction indices
comptime var x_indices: [lhs.len]u32 = undefined;
comptime var y_indices: [rhs.len]u32 = undefined;
comptime var remainder: [lhs.len + rhs.len]u32 = undefined;
comptime var char: u8 = undefined;
comptime var match: u32 = 0;
comptime var rhs_i: u32 = 0;
comptime var rem_i: u32 = 0;
comptime var found: bool = false;
index = 0;
inline while(index < lhs.len) : (index += 1) {
// matched + unmatched = total
if(match == rhs.len and rem_i == remainder.len) {
break;
}
char = lhs[index];
found = false;
// try to match the current char
// in both rhs and lhs operands
rhs_i = 0;
inline while(rhs_i < rhs.len) : (rhs_i += 1) {
if (rhs[rhs_i] == char) {
x_indices[match] = index;
y_indices[match] = rhs_i;
found = true;
match += 1;
break;
}
}
// if no match, add to remainder
if(!found) {
remainder[rem_i] = index;
rem_i += 1;
}
}
if(match != rhs.len) {
@compileError("Unmatched dimensions between operands: " ++ str);
}
rem_i = 0;
index = rhs.len;
inline while(index < lhs.len) : ({ index += 1; rem_i += 1; }){
x_indices[index] = remainder[rem_i];
}
return IndicesPair(lRank, rRank){ .lhs = x_indices, .rhs = y_indices };
}
From: ZEIN/Core/V1/TensorOps.zig at main · andrewCodeDev/ZEIN · GitHub
9 Likes
The main take away I was trying to convey there is comptime var. It’s crazy how many languages say they have meta programming but don’t have user accessible variables at compile time. Basically, writing that loop in C++ would have been a bloody nightmare (variadic arguments with empty base-cases that branch on function dispatch… yuck…).
2 Likes
That’s really interesting! I don’t have C++ experience but proc macros in Rust are enough to drive anyone mad.
comptime var is definitely a game changer and makes this type of meta programming so much easier. I’m loving inline while and planning to familiarize myself with inline switch
2 Likes
You can check out the format function in fmt that does compile-time parsing of the format string: zig/lib/std/fmt.zig at master · ziglang/zig · GitHub
5 Likes
I made a comptime assembler for the Raspberry Pi Pico’s PIO hardware block, implementation is here. The PIO is unique in that it lets you generate cycle accurate digital signals using a specialized bytecode.
At SYCL 2023 I ran a workshop where everyone made their own synth, the microcontroller didn’t actually have dedicated hardware to generate an I2S signal (which is required by the amplifier), so I programmed the PIO hardware block to generate it, and the usage looks like this.
Edit: formatting and grammar
8 Likes
coming from stats/ml/hpc and just kicking the wheels on Zig (kinda exactly the mix of stuff I wanted after working w/ C/C++/ISPC/OpenCL/Python), I am especially intrigued by the comptime for some aspects of numerical work, so I’m interested by your start on a tensor library. One of the use cases I’m trying to understand for comptime, and maybe it’s already in your contraction implementation and I didn’t understand it, is whether comptime can be used to optimize a computational graph? For instance, a naive runtime compute graph eval needs to traverse nodes, but if you know the graph at compile time, the node traversal can be avoided by inlining the evaluations together. Another example is autodiff: autodiff can be done at runtime dynamically but if the operations are known at compile time, a faster code path can be taken. If you have any hints about whether that pattern would be possible or not, I’d be interested.
in any case thanks for your work!
1 Like
I am currently evaluating functional programming style in Zig (maybe I write a blog post about it). Why, because I can learn some FP & Zig and it’s fun.
You can certainly do a lot of things in Zig’s comptime, e.g. composing functions during compile time:
Currying, at least with functions that have a fixed set of parameters, e.g. 3:
Partial evaluation (at leat for a function with 2 parameters):
And last but not least, functions like map and reduce on arrays:
3 Likes
Wow 
exactly the sort of example (the composer one) I was thinking of trying first . The question is still, if you dump the assembly, do the functions get inlined as if you had written the composite function by hand? ( I’m on mobile so I’ll try it myself later.)
Edit Just saw inline for in the Compose impl so I guess it’s the case
Edit edit What is Zig's Comptime? | Loris Cro's Blog answers my question actually with some complementary examples, if I had just found it first.
A second question, would be if comptime introspection is available for function code as well? I didn’t see anything like that yet.
1 Like
Afaik, inline for does unroll during comptime execution. But I did not check the resukting binary so far.
1 Like
I have a question:
Why don’t you write a command line tool, callable from build.zig, which generates the code based on the graph information it is provided with?
Is it ergonomics or do you want to understand the limits of comptime?
I usually like doing things with comptime as well but a code generator is often also a good solution.
1 Like
It is a little bit both, but now the question about comptime is answered.
yes indeed, but I wanted to try reverse mode autodiff, which requires a graph traversal and source code generation is tricky. My thought was that I could implement & test first with runtime, then just use a comptime feature (like inline for) to get the sort of implementation I’d otherwise have to write by hand. All while having a single implementation of the math & AD. Elsewhere this requires two languages (Enoki or Stan Math Library, both C++) or a JIT (e.g. PyTorch or TensorFlow’s XLA).
A code generator does have a major benefit in being able to read the code which runs and debug it explicitly. That sort of thing has always been trick in compile time systems (except maybe cpp lol)
1 Like