smarter data layout of local variables inside async function frame to save space

[andrewrk]

In theory these should have the same frame size:

const std = @import("std");

pub fn main() void {
    const size1: usize = @sizeOf(@Frame(func1));
    const size2: usize = @sizeOf(@Frame(func2));
    std.debug.warn("func1={}\nfunc2={}\n", size1, size2);
}

async fn func1() void {
    {
        var x: i32 = 1234;
    }
    {
        var x: i32 = 1234;
    }
    {
        var x: i32 = 1234;
    }
    {
        var x: i32 = 1234;
    }
    {
        var x: i32 = 1234;
    }
    {
        var x: i32 = 1234;
    }
    {
        var x: i32 = 1234;
    }
    {
        var x: i32 = 1234;
    }
    {
        var x: i32 = 1234;
    }
    {
        var x: i32 = 1234;
    }
    {
        var x: i32 = 1234;
    }
}
async fn func2() void {
    var x: i32 = 1234;
}

But currently:

$ ./zig run test.zig 
func1=80
func2=32

[andrewrk]

I did a little analysis on the async-std-lib branch and it looks like this issue is the next blocker for async/await stuff. This is a little snippet of the output:

@Frame(std.debug.openSelfDebugInfoPosix): 1.02 GiB (padding = 21)
 @Frame(std.debug.openElfDebugInfo): 1.02 GiB (padding = 184)
  @Frame(std.debug.openDwarfDebugInfo): 844.48 MiB (padding = 34)
   @Frame(std.debug.scanAllFunctions): 612.35 MiB (padding = 209)
    @Frame(std.debug.parseDie1): 188.11 MiB (padding = 80)
     @Frame(std.debug.parseFormValue): 184.10 MiB (padding = 282)
      @Frame(std.debug.parseFormValueRef): 20.01 MiB (padding = 126)
       @Frame(std.debug.leb128.readULEB128): 4.00 MiB (padding = 43)
        @Frame(std.io.in_stream.InStream(anyerror).readByte): 4.00 MiB (padding = 46)
         @Frame(std.io.in_stream.InStream(anyerror).readNoEof): 4.00 MiB (padding = 29)
          @Frame(std.io.in_stream.InStream(anyerror).readFull): 4.00 MiB (padding = 19)
           @Frame(std.io.in_stream.InStream(anyerror).read): 4.00 MiB (padding = 8)
            [4194304]u8: 4.00 MiB
            [32]usize: 256.00 bytes
            builtin.StackTrace: 24.00 bytes (padding = 0)

From this it's pretty clear how fast not reusing bytes adds up. The 4 MiB is OK, it's from the new std.io.InStream:

pub const default_stack_size = 4 * 1024 * 1024;
pub const stack_size: usize = if (@hasDecl(root, "stack_size_std_io_InStream"))
    root.stack_size_std_io_InStream
else
    default_stack_size;
pub const stack_align = 16;

pub fn InStream(comptime ReadError: type) type {
    return struct {
        const Self = @This();
        pub const Error = ReadError;
        pub const ReadFn = if (std.io.is_async)
            async fn (self: *Self, buffer: []u8) Error!usize
        else
            fn (self: *Self, buffer: []u8) Error!usize;

        /// Returns the number of bytes read. It may be less than buffer.len.
        /// If the number of bytes read is 0, it means end of stream.
        /// End of stream is not an error condition.
        readFn: ReadFn,

        /// Returns the number of bytes read. It may be less than buffer.len.
        /// If the number of bytes read is 0, it means end of stream.
        /// End of stream is not an error condition.
        pub fn read(self: *Self, buffer: []u8) Error!usize {
            if (std.io.is_async) {
                var stack_frame: [stack_size]u8 align(stack_align) = undefined;
                // TODO https://github.com/ziglang/zig/issues/3068
                var result: Error!usize = undefined;
                return await @asyncCall(&stack_frame, &result, self.readFn, self, buffer);
            } else {
                return self.readFn(self, buffer);
            }
        }

But 1.02 GiB is not okay! We can't be having a stack size of 1.02 GiB, that's crazy town.

We also don't yet request a bigger stack when it's statically known that we need more than the default (for example 16 MiB on Linux).

Update: I made it output JSON so one can use Firefox's tree view widget on the json file:

[andrewrk]

I've got a small patch undergoing local testing that is a huge improvement:

Now that's more like it. It's barely any more stack usage than the 4 MiB reserved by the std.io.InStream.

I'm going to leave this open though - there's still room for improvement. The original test case still has the same flaw.

There's also a case I want to consider, where an async function call inside a suspend block races with being resumed and calling an async function.

[LemonBoy]

In theory these should have the same frame size

On a more general level LLVM is also able to recycle the alloca'd stack slots as needed if they're annotated with the right lifetime hints.

[Hadron67]

I'm blocked by this issue too, it seems that even temporary variables sometimes are not reused:

const std = @import("std");

fn func1(a: i32) void {
    suspend {}
}

fn func2(a: i32) i32 {
    suspend {}
    return a + 1;
}

fn test1(a: i32) void {
    func1(a + a + a + a);
}

fn test2(a: i32) void {
    _ = func2(a + a + a + a);
}

pub fn main() void {
    std.debug.print("test1: {}\n", .{@sizeOf(@Frame(test1))});
    std.debug.print("test2: {}\n", .{@sizeOf(@Frame(test2))});
}

frame size of test1 and test2 should only differ by the size of an i32, but actually it's much more:

test1: 64
test2: 112

In real world use cases this makes the frame size grow rapidly, and in my case it reached over 10k.