Skip to content

Inline the copying collector's bump allocation in Wasm code #13323

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fitzgen wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Inline the copying collector's bump allocation in Wasm code #13323

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 0 additions & 31 deletions crates/cranelift/src/func_environ/gc/enabled.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,37 +105,6 @@ fn unbarriered_store_gc_ref(
Ok(())
}

/// Emit CLIF to call the `gc_raw_alloc` libcall.
#[cfg(any(feature = "gc-drc", feature = "gc-copying"))]
fn emit_gc_raw_alloc(
func_env: &mut FuncEnvironment<'_>,
builder: &mut FunctionBuilder<'_>,
kind: VMGcKind,
ty: ModuleInternedTypeIndex,
size: ir::Value,
align: u32,
) -> ir::Value {
let gc_alloc_raw_builtin = func_env.builtin_functions.gc_alloc_raw(builder.func);
let vmctx = func_env.vmctx_val(&mut builder.cursor());

let kind = builder
.ins()
.iconst(ir::types::I32, i64::from(kind.as_u32()));

let ty = func_env.module_interned_to_shared_ty(&mut builder.cursor(), ty);

assert!(align.is_power_of_two());
let align = builder.ins().iconst(ir::types::I32, i64::from(align));

let call_inst = builder
.ins()
.call(gc_alloc_raw_builtin, &[vmctx, kind, ty, size, align]);

let gc_ref = builder.func.dfg.first_result(call_inst);
builder.declare_value_needs_stack_map(gc_ref);
gc_ref
}

/// Emit inline CLIF code that asserts an object's `VMGcKind` matches the
/// expected kind. Only emits code when `cfg(gc_zeal)` is enabled.
///
Expand Down
267 changes: 230 additions & 37 deletions crates/cranelift/src/func_environ/gc/enabled/copying.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
//! Compiler for the copying (semi-space/Cheney) collector.
//!
//! Allocation is performed via the `gc_alloc_raw` libcall (not inlined) for
//! now. Read and write barriers are unnecessary (e.g. no reference counting and
//! Allocation is performed inline with a bump pointer when possible, falling
//! back to the `gc_alloc_raw` libcall when the active semi-space is full.
//! Read and write barriers are unnecessary (e.g. no reference counting and
//! no concurrent mutation during collection) but we do need stack maps so the
//! collector can find and update roots when it relocates objects.

Expand All @@ -10,12 +11,16 @@ use crate::TRAP_INTERNAL_ASSERT;
use crate::func_environ::FuncEnvironment;
use crate::translate::TargetEnvironment;
use cranelift_codegen::ir::{self, InstBuilder};
use cranelift_entity::packed_option::ReservedValue;
use cranelift_frontend::FunctionBuilder;
use smallvec::SmallVec;
use wasmtime_environ::copying::{EXCEPTION_TAG_DEFINED_OFFSET, EXCEPTION_TAG_INSTANCE_OFFSET};
use wasmtime_environ::VMSharedTypeIndex;
use wasmtime_environ::copying::{
ALIGN, EXCEPTION_TAG_DEFINED_OFFSET, EXCEPTION_TAG_INSTANCE_OFFSET,
HEAP_DATA_ACTIVE_SPACE_END_OFFSET, HEAP_DATA_BUMP_PTR_OFFSET,
};
use wasmtime_environ::{
GcTypeLayouts, TypeIndex, VMGcKind, WasmHeapTopType, WasmHeapType, WasmRefType, WasmResult,
WasmStorageType, WasmValType, copying::CopyingTypeLayouts,
GcTypeLayouts, ModuleInternedTypeIndex, TypeIndex, VMGcKind, WasmHeapTopType, WasmHeapType,
WasmRefType, WasmResult, WasmStorageType, WasmValType, copying::CopyingTypeLayouts,
};

#[derive(Default)]
Expand All @@ -24,6 +29,218 @@ pub struct CopyingCompiler {
}

impl CopyingCompiler {
/// Load the pointer to the `VMCopyingHeapData` from vmctx.
fn load_vmcopying_heap_data_ptr(
func_env: &mut FuncEnvironment<'_>,
builder: &mut FunctionBuilder,
) -> ir::Value {
let pointer_type = func_env.pointer_type();
let vmctx = func_env.vmctx_val(&mut builder.cursor());
builder.ins().load(
pointer_type,
ir::MemFlags::trusted().with_readonly().with_can_move(),
vmctx,
i32::from(func_env.offsets.ptr.vmctx_gc_heap_data()),
)
}

/// Load the current bump pointer and active-space end from a `*mut
/// VMCopyingHeapData`.
///
/// Returns `(bump_ptr, active_space_end)` as `i32` values.
fn load_bump_state(
builder: &mut FunctionBuilder,
ptr_to_heap_data: ir::Value,
) -> (ir::Value, ir::Value) {
let bump_ptr = builder.ins().load(
ir::types::I32,
ir::MemFlags::trusted().with_can_move(),
ptr_to_heap_data,
i32::try_from(HEAP_DATA_BUMP_PTR_OFFSET).unwrap(),
);
let active_space_end = builder.ins().load(
ir::types::I32,
ir::MemFlags::trusted().with_readonly().with_can_move(),
ptr_to_heap_data,
i32::try_from(HEAP_DATA_ACTIVE_SPACE_END_OFFSET).unwrap(),
);
(bump_ptr, active_space_end)
}

/// Round `size` (an `i32`) up to `ALIGN`, returning the result as an `i64`.
///
/// Uses `i64` arithmetic so that overflow produces a value larger than any
/// valid heap index, which sends us to the slow allocation path instead of
/// wrapping around.
fn aligned_size(builder: &mut FunctionBuilder, size: ir::Value) -> ir::Value {
let size_64 = builder.ins().uextend(ir::types::I64, size);
let align_mask = builder.ins().iconst(ir::types::I64, i64::from(ALIGN - 1));
let inv_align_mask = builder.ins().iconst(ir::types::I64, !i64::from(ALIGN - 1));
let size_plus_mask = builder.ins().iadd(size_64, align_mask);
builder.ins().band(size_plus_mask, inv_align_mask)
}

/// Emit inline bump allocation, falling back to `gc_alloc_raw` on failure.
///
/// `kind` may be `VMGcKind::ExternRef` iff `ty` is `None`.
///
/// `size` must be an `i32` value >= `size_of(VMCopyingHeader)`.
///
/// Returns `(gc_ref, raw_ptr_to_object)` where `gc_ref` is the `i32` GC
/// heap index and `raw_ptr_to_object` is a native pointer into the GC heap.
fn emit_inline_alloc(
&mut self,
func_env: &mut FuncEnvironment<'_>,
builder: &mut FunctionBuilder,
kind: VMGcKind,
ty: Option<ModuleInternedTypeIndex>,
size: ir::Value,
) -> (ir::Value, ir::Value) {
assert_eq!(builder.func.dfg.value_type(size), ir::types::I32);

let pointer_type = func_env.pointer_type();
let current_block = builder.current_block().unwrap();
let fast_block = builder.create_block();
let slow_block = builder.create_block();
let merge_block = builder.create_block();

builder.ensure_inserted_block();
builder.insert_block_after(fast_block, current_block);
builder.insert_block_after(slow_block, fast_block);
builder.insert_block_after(merge_block, slow_block);

let ptr_to_heap_data = Self::load_vmcopying_heap_data_ptr(func_env, builder);
let (bump_ptr, active_space_end) = Self::load_bump_state(builder, ptr_to_heap_data);
let aligned_size_64 = Self::aligned_size(builder, size);

// Compute `end_of_object = bump_ptr + aligned_size` (in i64) and check
// whether it fits within the active semi-space.
let bump_ptr_64 = builder.ins().uextend(ir::types::I64, bump_ptr);
let end_64 = builder.ins().iadd(bump_ptr_64, aligned_size_64);
let active_space_end_64 = builder.ins().uextend(ir::types::I64, active_space_end);
let fits = builder.ins().icmp(
ir::condcodes::IntCC::UnsignedLessThanOrEqual,
end_64,
active_space_end_64,
);
builder.ins().brif(fits, fast_block, &[], slow_block, &[]);

// Slow path: when there isn't enough room in the bump region, call the
// `gc_alloc_raw` libcall, which will collect or grow the GC heap as
// necessary.
{
builder.switch_to_block(slow_block);
builder.seal_block(slow_block);
builder.set_cold_block(slow_block);
let gc_alloc_raw_builtin = func_env.builtin_functions.gc_alloc_raw(builder.func);
let vmctx = func_env.vmctx_val(&mut builder.cursor());
let kind_val = builder
.ins()
.iconst(ir::types::I32, i64::from(kind.as_u32()));
let shared_ty = match ty {
Some(ty) => func_env.module_interned_to_shared_ty(&mut builder.cursor(), ty),
None => builder.ins().iconst(
func_env.vmshared_type_index_ty(),
i64::from(VMSharedTypeIndex::reserved_value().as_bits()),
),
};
let align_val = builder.ins().iconst(ir::types::I32, i64::from(ALIGN));
let call_inst = builder.ins().call(
gc_alloc_raw_builtin,
&[vmctx, kind_val, shared_ty, size, align_val],
);
let gc_ref = builder.func.dfg.first_result(call_inst);
let base = func_env.get_gc_heap_base(builder);
let heap_offset = uextend_i32_to_pointer_type(builder, pointer_type, gc_ref);
let obj_ptr = builder.ins().iadd(base, heap_offset);
builder
.ins()
.jump(merge_block, &[gc_ref.into(), obj_ptr.into()]);
}

// Fast path: there is capacity for the requested object in the bump
// region, so finish the allocation inline, update our bump pointer,
// etc...
{
builder.switch_to_block(fast_block);
builder.seal_block(fast_block);

// The old bump_ptr is the start of the new object.
let gc_ref = bump_ptr;

// Update the bump pointer.
let end_of_object = builder.ins().ireduce(ir::types::I32, end_64);
builder.ins().store(
ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Vmctx)),
end_of_object,
ptr_to_heap_data,
i32::try_from(HEAP_DATA_BUMP_PTR_OFFSET).unwrap(),
);

// Compute the raw pointer to the new object.
let base = func_env.get_gc_heap_base(builder);
let heap_offset = uextend_i32_to_pointer_type(builder, pointer_type, gc_ref);
let obj_ptr = builder.ins().iadd(base, heap_offset);

// Write `VMGcHeader::{kind,type_index}` as a single i64 store.
let shared_ty = match ty {
Some(ty) => func_env.module_interned_to_shared_ty(&mut builder.cursor(), ty),
None => builder.ins().iconst(
func_env.vmshared_type_index_ty(),
i64::from(VMSharedTypeIndex::reserved_value().as_bits()),
),
};
let shared_ty_i64 = builder.ins().uextend(ir::types::I64, shared_ty);
let kind_i64 = i64::from(kind.as_u32());
let header_i64 = match func_env.isa().endianness() {
ir::Endianness::Little => {
// Low 32 bits = kind (at offset 0), high 32 bits = type_index (at offset 4).
let ty_shifted = builder.ins().ishl_imm(shared_ty_i64, 32);
let kind_val = builder.ins().iconst(ir::types::I64, kind_i64);
builder.ins().bor(kind_val, ty_shifted)
}
ir::Endianness::Big => {
// High 32 bits = kind (at offset 0), low 32 bits = type_index (at offset 4).
let kind_shifted = builder.ins().iconst(ir::types::I64, kind_i64 << 32);
builder.ins().bor(kind_shifted, shared_ty_i64)
}
};
builder.ins().store(
ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Vmctx)),
header_i64,
obj_ptr,
i32::try_from(wasmtime_environ::VM_GC_HEADER_KIND_OFFSET).unwrap(),
);

// Write `VMCopyingHeader::object_size`.
//
// Note: this `ireduce` doesn't truncate bits because we already
// checked that the object (and therefore also its size) fits within
// the active bump region on this path.
let aligned_size_32 = builder.ins().ireduce(ir::types::I32, aligned_size_64);
builder.ins().store(
ir::MemFlags::trusted(),
aligned_size_32,
obj_ptr,
i32::try_from(wasmtime_environ::VM_GC_HEADER_SIZE).unwrap(),
);

builder
.ins()
.jump(merge_block, &[gc_ref.into(), obj_ptr.into()]);
}

// Merge block: takes the GC ref and the raw pointer to the GC object as
// block parameters.
builder.switch_to_block(merge_block);
let gc_ref = builder.append_block_param(merge_block, ir::types::I32);
let ptr_to_object = builder.append_block_param(merge_block, pointer_type);
builder.seal_block(merge_block);
builder.declare_value_needs_stack_map(gc_ref);

(gc_ref, ptr_to_object)
}

fn init_field(
&mut self,
func_env: &mut FuncEnvironment<'_>,
Expand Down Expand Up @@ -85,28 +302,20 @@ impl GcCompiler for CopyingCompiler {
let len_offset = gc_compiler(func_env)?.layouts().array_length_field_offset();
let array_layout = func_env.array_layout(interned_type_index).clone();
let base_size = array_layout.base_size;
let align = array_layout.align;
let len_to_elems_delta = base_size.checked_sub(len_offset).unwrap();

// First, compute the array's total size.
let len = init.len(&mut builder.cursor());
let size = emit_array_size(func_env, builder, &array_layout, len);

// Allocate via libcall.
let array_ref = emit_gc_raw_alloc(
// Allocate inline (with fallback to libcall).
let (array_ref, object_addr) = self.emit_inline_alloc(
func_env,
builder,
VMGcKind::ArrayRef,
interned_type_index,
Some(interned_type_index),
size,
align,
);

// Write the array's length.
let base = func_env.get_gc_heap_base(builder);
let extended_array_ref =
uextend_i32_to_pointer_type(builder, func_env.pointer_type(), array_ref);
let object_addr = builder.ins().iadd(base, extended_array_ref);
let len_addr = builder.ins().iadd_imm(object_addr, i64::from(len_offset));
let len = init.len(&mut builder.cursor());
builder
Expand Down Expand Up @@ -142,26 +351,18 @@ impl GcCompiler for CopyingCompiler {
let struct_layout = func_env.struct_or_exn_layout(interned_type_index);

let struct_size = struct_layout.size;
let struct_align = struct_layout.align;
let field_offsets: SmallVec<[_; 8]> = struct_layout.fields.iter().copied().collect();
assert_eq!(field_vals.len(), field_offsets.len());

let struct_size_val = builder.ins().iconst(ir::types::I32, i64::from(struct_size));

let struct_ref = emit_gc_raw_alloc(
let (struct_ref, raw_ptr_to_struct) = self.emit_inline_alloc(
func_env,
builder,
VMGcKind::StructRef,
interned_type_index,
Some(interned_type_index),
struct_size_val,
struct_align,
);

// Initialize fields.
let base = func_env.get_gc_heap_base(builder);
let extended_struct_ref =
uextend_i32_to_pointer_type(builder, func_env.pointer_type(), struct_ref);
let raw_ptr_to_struct = builder.ins().iadd(base, extended_struct_ref);
initialize_struct_fields(
func_env,
builder,
Expand Down Expand Up @@ -191,26 +392,18 @@ impl GcCompiler for CopyingCompiler {
let exn_layout = func_env.struct_or_exn_layout(interned_type_index);

let exn_size = exn_layout.size;
let exn_align = exn_layout.align;
let field_offsets: SmallVec<[_; 8]> = exn_layout.fields.iter().copied().collect();
assert_eq!(field_vals.len(), field_offsets.len());

let exn_size_val = builder.ins().iconst(ir::types::I32, i64::from(exn_size));

let exn_ref = emit_gc_raw_alloc(
let (exn_ref, raw_ptr_to_exn) = self.emit_inline_alloc(
func_env,
builder,
VMGcKind::ExnRef,
interned_type_index,
Some(interned_type_index),
exn_size_val,
exn_align,
);

// Initialize fields.
let base = func_env.get_gc_heap_base(builder);
let extended_exn_ref =
uextend_i32_to_pointer_type(builder, func_env.pointer_type(), exn_ref);
let raw_ptr_to_exn = builder.ins().iadd(base, extended_exn_ref);
initialize_struct_fields(
func_env,
builder,
Expand Down
Loading
Loading