async_mem.zig

const std = @import("std");
const Allocator = std.mem.Allocator;

const workload_size = 1 << 20;
const n_workloads = 4;
const repetition = 4;
const use_async = true;
const use_manual_prefetch = true;

const Cell = extern struct {
    next_index: u64,
    padding: [7]u64 = [_]u64{0} ** 7,
};

fn initCells(allocator: *Allocator, size: usize) ![]Cell {
    const order = try allocator.alloc(u64, size - 1);
    defer allocator.free(order);

    for (order) |*e, i| e.* = i + 1;

    var rnd = std.rand.DefaultPrng.init(@bitCast(usize, std.time.milliTimestamp()));
    rnd.random.shuffle(u64, order);

    const cells = try allocator.alloc(Cell, size);

    var prev: u64 = 0;
    for (order) |i| {
        cells[prev] = .{.next_index = i};
        prev = i;
    }
    cells[prev] = .{.next_index = size};

    return cells;
}

fn syncSum(workloads: *const [n_workloads][]const Cell) u64 {
    var sum: u64 = 0;
    for (workloads) |workload| {
        var i: u64 = 0;
        while (i != workload.len) {
            i = workload[i].next_index;
            sum += i;
        }
    }

    return sum;
}

fn asyncWorker(workload: []const Cell, ready: *bool) u64 {
    var partial_sum: u64 = 0;
    var i: u64 = 0;
    while (i != workload.len) {
        if (use_manual_prefetch) {
            asm volatile("prefetcht0 (%[addr])" :: [addr] "r" (&workload[i].next_index));
            suspend;
            i = workload[i].next_index;
            partial_sum += i;
        } else {
            i = workload[i].next_index;
            suspend;
            partial_sum += i;
        }
    }

    ready.* = true;
    return partial_sum;
}

fn asyncSum(workloads: *const [n_workloads][]const Cell) u64 {
    var frames: [n_workloads]@Frame(asyncWorker) = undefined;
    var ready = [_]bool{false} ** n_workloads; 

    for (frames) |*f, i| f.* = async asyncWorker(workloads[i], &ready[i]);

    while (true) {
        var any_unready = false;
        for (frames) |*f, i| {
            if (!ready[i]) {
                resume f;
                any_unready = true;
            }
        }

        if (!any_unready) {
            break;
        }
    }

    var sum: u64 = 0;
    for (frames) |*f| {
        sum += nosuspend await f;
    }

    return sum;
}

pub fn main() !void {
    const allocator = std.heap.page_allocator;

    var workloads = [_][]Cell{&[_]Cell{}} ** n_workloads;
    defer for (workloads) |w| allocator.free(w);
    for (workloads) |*w| w.* = try initCells(allocator, workload_size);

    var i: usize = 0;
    while (i < repetition) : (i += 1) {
        var timer = try std.time.Timer.start();
        const sum = if (use_async) asyncSum(&workloads) else syncSum(&workloads);
        var dt = timer.read();
        std.debug.print("sum: {}, {} ms\n", .{sum, dt / 1_000_000});
    }
}