authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-01-04 20:29:46-08:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-01-15 15:11:36-08:00
logb41b5fe52935ec19402704742751dc9c973518a4
treeaabfd5c6c3fb114833035579bc4c370eb756557b
parent1fd708b1bc41878c33db7e2a50433ab1a66d41ee

wasm linker: implement __wasm_init_memory

1 files changed, 197 insertions(+), 19 deletions(-)

src/link/Wasm/Flush.zig+197-19
......@@ -26,7 +26,7 @@ data_segments: std.AutoArrayHashMapUnmanaged(Wasm.DataSegmentId, u32) = .empty,
2626/// Each time a `data_segment` offset equals zero it indicates a new group, and
2727/// the next element in this array will contain the total merged segment size.
2828/// Value is the virtual memory address of the end of the segment.
29data_segment_groups: std.ArrayListUnmanaged(u32) = .empty,
29data_segment_groups: std.ArrayListUnmanaged(DataSegmentGroup) = .empty,
3030
3131binary_bytes: std.ArrayListUnmanaged(u8) = .empty,
3232missing_exports: std.AutoArrayHashMapUnmanaged(String, void) = .empty,
......@@ -37,6 +37,11 @@ data_imports: std.AutoArrayHashMapUnmanaged(String, Wasm.DataImportId) = .empty,
3737/// For debug purposes only.
3838memory_layout_finished: bool = false,
3939
40const DataSegmentGroup = struct {
41 first_segment: Wasm.DataSegmentId,
42 end_addr: u32,
43};
44
4045pub fn clear(f: *Flush) void {
4146 f.data_segments.clearRetainingCapacity();
4247 f.data_segment_groups.clearRetainingCapacity();
......@@ -280,15 +285,6 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
280285 // Always place the stack at the start by default unless the user specified the global-base flag.
281286 const place_stack_first, var memory_ptr: u64 = if (wasm.global_base) |base| .{ false, base } else .{ true, 0 };
282287
283 const VirtualAddrs = struct {
284 stack_pointer: u32,
285 heap_base: u32,
286 heap_end: u32,
287 tls_base: ?u32,
288 tls_align: Alignment,
289 tls_size: ?u32,
290 init_memory_flag: ?u32,
291 };
292288 var virtual_addrs: VirtualAddrs = .{
293289 .stack_pointer = undefined,
294290 .heap_base = undefined,
......@@ -309,9 +305,10 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
309305 const segment_vaddrs = f.data_segments.values();
310306 assert(f.data_segment_groups.items.len == 0);
311307 const data_vaddr: u32 = @intCast(memory_ptr);
312 {
308 if (segment_ids.len > 0) {
313309 var seen_tls: enum { before, during, after } = .before;
314310 var category: Wasm.DataSegmentId.Category = undefined;
311 var first_segment: Wasm.DataSegmentId = segment_ids[0];
315312 for (segment_ids, segment_vaddrs, 0..) |segment_id, *segment_vaddr, i| {
316313 const alignment = segment_id.alignment(wasm);
317314 category = segment_id.category(wasm);
......@@ -338,14 +335,21 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
338335 };
339336 if (want_new_segment) {
340337 log.debug("new segment at 0x{x} {} {s} {}", .{ start_addr, segment_id, segment_id.name(wasm), category });
341 try f.data_segment_groups.append(gpa, @intCast(memory_ptr));
338 try f.data_segment_groups.append(gpa, .{
339 .end_addr = @intCast(memory_ptr),
340 .first_segment = first_segment,
341 });
342 first_segment = segment_id;
342343 }
343344
344345 const size = segment_id.size(wasm);
345346 segment_vaddr.* = @intCast(start_addr);
346347 memory_ptr = start_addr + size;
347348 }
348 if (category != .zero) try f.data_segment_groups.append(gpa, @intCast(memory_ptr));
349 if (category != .zero) try f.data_segment_groups.append(gpa, .{
350 .first_segment = first_segment,
351 .end_addr = @intCast(memory_ptr),
352 });
349353 }
350354
351355 if (shared_memory and wasm.any_passive_inits) {
......@@ -567,7 +571,7 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
567571 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Valtype.i32));
568572 binary_bytes.appendAssumeCapacity(1); // mutable
569573 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.i32_const));
570 leb.writeUleb128(binary_bytes.fixedWriter(), virtual_addrs.stack_pointer) catch unreachable;
574 appendReservedUleb32(binary_bytes, virtual_addrs.stack_pointer);
571575 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.end));
572576 },
573577 .__tls_align => @panic("TODO"),
......@@ -683,7 +687,11 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
683687 defer replaceSize(binary_bytes, code_start);
684688 try emitCallCtorsFunction(wasm, binary_bytes);
685689 },
686 .__wasm_init_memory => @panic("TODO lower __wasm_init_memory "),
690 .__wasm_init_memory => {
691 const code_start = try reserveSize(gpa, binary_bytes);
692 defer replaceSize(binary_bytes, code_start);
693 try emitInitMemoryFunction(wasm, binary_bytes, &virtual_addrs);
694 },
687695 .__wasm_init_tls => @panic("TODO lower __wasm_init_tls "),
688696 .object_function => |i| {
689697 const ptr = i.ptr(wasm);
......@@ -736,7 +744,7 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
736744 var group_index: u32 = 0;
737745 var segment_offset: u32 = 0;
738746 var group_start_addr: u32 = data_vaddr;
739 var group_end_addr = f.data_segment_groups.items[group_index];
747 var group_end_addr = f.data_segment_groups.items[group_index].end_addr;
740748 for (segment_ids, segment_vaddrs) |segment_id, segment_vaddr| {
741749 if (segment_vaddr >= group_end_addr) {
742750 try binary_bytes.appendNTimes(gpa, 0, group_end_addr - group_start_addr - segment_offset);
......@@ -746,7 +754,7 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
746754 break;
747755 }
748756 group_start_addr = group_end_addr;
749 group_end_addr = f.data_segment_groups.items[group_index];
757 group_end_addr = f.data_segment_groups.items[group_index].end_addr;
750758 segment_offset = 0;
751759 }
752760 if (segment_offset == 0) {
......@@ -865,6 +873,16 @@ pub fn finish(f: *Flush, wasm: *Wasm) !void {
865873 try file.setEndPos(binary_bytes.items.len);
866874}
867875
876const VirtualAddrs = struct {
877 stack_pointer: u32,
878 heap_base: u32,
879 heap_end: u32,
880 tls_base: ?u32,
881 tls_align: Alignment,
882 tls_size: ?u32,
883 init_memory_flag: ?u32,
884};
885
868886fn emitNameSection(
869887 wasm: *Wasm,
870888 data_segments: *const std.AutoArrayHashMapUnmanaged(Wasm.DataSegmentId, u32),
......@@ -1575,7 +1593,7 @@ fn emitCallCtorsFunction(wasm: *const Wasm, binary_bytes: *std.ArrayListUnmanage
15751593 const gpa = wasm.base.comp.gpa;
15761594
15771595 try binary_bytes.ensureUnusedCapacity(gpa, 5 + 1);
1578 leb.writeUleb128(binary_bytes.fixedWriter(), @as(u32, 0)) catch unreachable; // no locals
1596 appendReservedUleb32(binary_bytes, 0); // no locals
15791597
15801598 for (wasm.object_init_funcs.items) |init_func| {
15811599 const func = init_func.function_index.ptr(wasm);
......@@ -1586,7 +1604,7 @@ fn emitCallCtorsFunction(wasm: *const Wasm, binary_bytes: *std.ArrayListUnmanage
15861604 try binary_bytes.ensureUnusedCapacity(gpa, 1 + 5 + n_returns + 1);
15871605 const call_index: Wasm.OutputFunctionIndex = .fromObjectFunction(wasm, init_func.function_index);
15881606 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.call));
1589 leb.writeUleb128(binary_bytes.fixedWriter(), @intFromEnum(call_index)) catch unreachable;
1607 appendReservedUleb32(binary_bytes, @intFromEnum(call_index));
15901608
15911609 // drop all returned values from the stack as __wasm_call_ctors has no return value
15921610 binary_bytes.appendNTimesAssumeCapacity(@intFromEnum(std.wasm.Opcode.drop), n_returns);
......@@ -1594,3 +1612,163 @@ fn emitCallCtorsFunction(wasm: *const Wasm, binary_bytes: *std.ArrayListUnmanage
15941612
15951613 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.end)); // end function body
15961614}
1615
1616fn emitInitMemoryFunction(
1617 wasm: *const Wasm,
1618 binary_bytes: *std.ArrayListUnmanaged(u8),
1619 virtual_addrs: *const VirtualAddrs,
1620) Allocator.Error!void {
1621 const comp = wasm.base.comp;
1622 const gpa = comp.gpa;
1623 const shared_memory = comp.config.shared_memory;
1624
1625 // Passive segments are used to avoid memory being reinitialized on each
1626 // thread's instantiation. These passive segments are initialized and
1627 // dropped in __wasm_init_memory, which is registered as the start function
1628 // We also initialize bss segments (using memory.fill) as part of this
1629 // function.
1630 assert(wasm.any_passive_inits);
1631
1632 try binary_bytes.ensureUnusedCapacity(gpa, 5 + 1);
1633 appendReservedUleb32(binary_bytes, 0); // no locals
1634
1635 if (virtual_addrs.init_memory_flag) |flag_address| {
1636 assert(shared_memory);
1637 try binary_bytes.ensureUnusedCapacity(gpa, 2 * 3 + 6 * 3 + 1 + 6 * 3 + 1 + 5 * 4 + 1 + 1);
1638 // destination blocks
1639 // based on values we jump to corresponding label
1640 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.block)); // $drop
1641 binary_bytes.appendAssumeCapacity(std.wasm.block_empty); // block type
1642
1643 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.block)); // $wait
1644 binary_bytes.appendAssumeCapacity(std.wasm.block_empty); // block type
1645
1646 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.block)); // $init
1647 binary_bytes.appendAssumeCapacity(std.wasm.block_empty); // block type
1648
1649 // atomically check
1650 appendReservedI32Const(binary_bytes, flag_address);
1651 appendReservedI32Const(binary_bytes, 0);
1652 appendReservedI32Const(binary_bytes, 1);
1653 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.atomics_prefix));
1654 appendReservedUleb32(binary_bytes, @intFromEnum(std.wasm.AtomicsOpcode.i32_atomic_rmw_cmpxchg));
1655 appendReservedUleb32(binary_bytes, 2); // alignment
1656 appendReservedUleb32(binary_bytes, 0); // offset
1657
1658 // based on the value from the atomic check, jump to the label.
1659 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.br_table));
1660 appendReservedUleb32(binary_bytes, 2); // length of the table (we have 3 blocks but because of the mandatory default the length is 2).
1661 appendReservedUleb32(binary_bytes, 0); // $init
1662 appendReservedUleb32(binary_bytes, 1); // $wait
1663 appendReservedUleb32(binary_bytes, 2); // $drop
1664 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.end));
1665 }
1666
1667 const segment_groups = wasm.flush_buffer.data_segment_groups.items;
1668 var prev_end: u32 = 0;
1669 for (segment_groups, 0..) |group, segment_index| {
1670 defer prev_end = group.end_addr;
1671 const segment = group.first_segment;
1672 if (!segment.isPassive(wasm)) continue;
1673
1674 const start_addr: u32 = @intCast(segment.alignment(wasm).forward(prev_end));
1675 const segment_size: u32 = group.end_addr - start_addr;
1676
1677 try binary_bytes.ensureUnusedCapacity(gpa, 6 + 6 + 1 + 5 + 6 + 6 + 1 + 6 * 2 + 1 + 1);
1678
1679 // For passive BSS segments we can simply issue a memory.fill(0). For
1680 // non-BSS segments we do a memory.init. Both instructions take as
1681 // their first argument the destination address.
1682 appendReservedI32Const(binary_bytes, start_addr);
1683
1684 if (shared_memory and segment.isTls(wasm)) {
1685 // When we initialize the TLS segment we also set the `__tls_base`
1686 // global. This allows the runtime to use this static copy of the
1687 // TLS data for the first/main thread.
1688 appendReservedI32Const(binary_bytes, start_addr);
1689 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.global_set));
1690 appendReservedUleb32(binary_bytes, virtual_addrs.tls_base.?);
1691 }
1692
1693 appendReservedI32Const(binary_bytes, 0);
1694 appendReservedI32Const(binary_bytes, segment_size);
1695 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.misc_prefix));
1696 if (segment.isBss(wasm)) {
1697 // fill bss segment with zeroes
1698 appendReservedUleb32(binary_bytes, @intFromEnum(std.wasm.MiscOpcode.memory_fill));
1699 } else {
1700 // initialize the segment
1701 appendReservedUleb32(binary_bytes, @intFromEnum(std.wasm.MiscOpcode.memory_init));
1702 appendReservedUleb32(binary_bytes, @intCast(segment_index));
1703 }
1704 binary_bytes.appendAssumeCapacity(0); // memory index immediate
1705 }
1706
1707 if (virtual_addrs.init_memory_flag) |flag_address| {
1708 assert(shared_memory);
1709 try binary_bytes.ensureUnusedCapacity(gpa, 6 + 6 + 1 + 3 * 5 + 6 + 1 + 5 + 1 + 3 * 5 + 1 + 1 + 5 + 1 + 6 * 2 + 1 + 5 + 1 + 3 * 5 + 1 + 1 + 1);
1710 // we set the init memory flag to value '2'
1711 appendReservedI32Const(binary_bytes, flag_address);
1712 appendReservedI32Const(binary_bytes, 2);
1713 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.atomics_prefix));
1714 appendReservedUleb32(binary_bytes, @intFromEnum(std.wasm.AtomicsOpcode.i32_atomic_store));
1715 appendReservedUleb32(binary_bytes, @as(u32, 2)); // alignment
1716 appendReservedUleb32(binary_bytes, @as(u32, 0)); // offset
1717
1718 // notify any waiters for segment initialization completion
1719 appendReservedI32Const(binary_bytes, flag_address);
1720 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.i32_const));
1721 leb.writeIleb128(binary_bytes.fixedWriter(), @as(i32, -1)) catch unreachable; // number of waiters
1722 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.atomics_prefix));
1723 appendReservedUleb32(binary_bytes, @intFromEnum(std.wasm.AtomicsOpcode.memory_atomic_notify));
1724 appendReservedUleb32(binary_bytes, @as(u32, 2)); // alignment
1725 appendReservedUleb32(binary_bytes, @as(u32, 0)); // offset
1726 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.drop));
1727
1728 // branch and drop segments
1729 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.br));
1730 appendReservedUleb32(binary_bytes, @as(u32, 1));
1731
1732 // wait for thread to initialize memory segments
1733 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.end)); // end $wait
1734 appendReservedI32Const(binary_bytes, flag_address);
1735 appendReservedI32Const(binary_bytes, 1); // expected flag value
1736 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.i64_const));
1737 leb.writeIleb128(binary_bytes.fixedWriter(), @as(i64, -1)) catch unreachable; // timeout
1738 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.atomics_prefix));
1739 appendReservedUleb32(binary_bytes, @intFromEnum(std.wasm.AtomicsOpcode.memory_atomic_wait32));
1740 appendReservedUleb32(binary_bytes, @as(u32, 2)); // alignment
1741 appendReservedUleb32(binary_bytes, @as(u32, 0)); // offset
1742 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.drop));
1743
1744 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.end)); // end $drop
1745 }
1746
1747 for (segment_groups, 0..) |group, segment_index| {
1748 const segment = group.first_segment;
1749 if (!segment.isPassive(wasm)) continue;
1750 if (segment.isBss(wasm)) continue;
1751 // The TLS region should not be dropped since its is needed
1752 // during the initialization of each thread (__wasm_init_tls).
1753 if (shared_memory and segment.isTls(wasm)) continue;
1754
1755 try binary_bytes.ensureUnusedCapacity(gpa, 1 + 5 + 5 + 1);
1756
1757 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.misc_prefix));
1758 appendReservedUleb32(binary_bytes, @intFromEnum(std.wasm.MiscOpcode.data_drop));
1759 appendReservedUleb32(binary_bytes, @intCast(segment_index));
1760 }
1761
1762 // End of the function body
1763 binary_bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.end));
1764}
1765
1766/// Writes an unsigned 32-bit integer as a LEB128-encoded 'i32.const' value.
1767fn appendReservedI32Const(bytes: *std.ArrayListUnmanaged(u8), val: u32) void {
1768 bytes.appendAssumeCapacity(@intFromEnum(std.wasm.Opcode.i32_const));
1769 leb.writeIleb128(bytes.fixedWriter(), @as(i32, @bitCast(val))) catch unreachable;
1770}
1771
1772fn appendReservedUleb32(bytes: *std.ArrayListUnmanaged(u8), val: u32) void {
1773 leb.writeUleb128(bytes.fixedWriter(), val) catch unreachable;
1774}