Inline the copying collector's bump allocation in Wasm code

crates/cranelift/src/func_environ/gc/enabled.rs

@@ -105,7 +105,7 @@ fn unbarriered_store_gc_ref(
     Ok(())
 }
 
-/// Emit CLIF to call the `gc_raw_alloc` libcall.
+/// Emit CLIF to call the `gc_alloc_raw` libcall.
 #[cfg(any(feature = "gc-drc", feature = "gc-copying"))]
 fn emit_gc_raw_alloc(
     func_env: &mut FuncEnvironment<'_>,

crates/cranelift/src/func_environ/gc/enabled/copying.rs

@@ -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.
 
@@ -11,11 +12,13 @@ use crate::func_environ::FuncEnvironment;
 use crate::translate::TargetEnvironment;
 use cranelift_codegen::ir::{self, InstBuilder};
 use cranelift_frontend::FunctionBuilder;
-use smallvec::SmallVec;
-use wasmtime_environ::copying::{EXCEPTION_TAG_DEFINED_OFFSET, EXCEPTION_TAG_INSTANCE_OFFSET};
+use wasmtime_environ::copying::{
+    ALIGN, EXCEPTION_TAG_DEFINED_OFFSET, EXCEPTION_TAG_INSTANCE_OFFSET,
+};
 use wasmtime_environ::{
-    GcTypeLayouts, TypeIndex, VMGcKind, WasmHeapTopType, WasmHeapType, WasmRefType, WasmResult,
-    WasmStorageType, WasmValType, copying::CopyingTypeLayouts,
+    GcTypeLayouts, ModuleInternedTypeIndex, PtrSize, TypeIndex, VMGcKind, WasmHeapTopType,
+    WasmHeapType, WasmRefType, WasmResult, WasmStorageType, WasmValType,
+    copying::CopyingTypeLayouts,
 };
 
 #[derive(Default)]
@@ -24,6 +27,188 @@ 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(
+        func_env: &mut FuncEnvironment<'_>,
+        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::from(func_env.offsets.ptr.vmcopying_heap_data_bump_ptr()),
+        );
+        let active_space_end = builder.ins().load(
+            ir::types::I32,
+            ir::MemFlags::trusted().with_readonly().with_can_move(),
+            ptr_to_heap_data,
+            i32::from(func_env.offsets.ptr.vmcopying_heap_data_active_space_end()),
+        );
+        (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.
+    ///
+    /// `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: ModuleInternedTypeIndex,
+        size: ir::Value,
+    ) -> (ir::Value, ir::Value) {
+        debug_assert_ne!(kind, VMGcKind::ExternRef);
+        debug_assert!(!ty.is_reserved_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(func_env, 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_ref = emit_gc_raw_alloc(func_env, builder, kind, ty, size, ALIGN);
+            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::from(func_env.offsets.ptr.vmcopying_heap_data_bump_ptr()),
+            );
+
+            // 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`.
+            let kind_val = builder
+                .ins()
+                .iconst(ir::types::I32, i64::from(kind.as_u32()));
+            builder.ins().store(
+                ir::MemFlags::trusted(),
+                kind_val,
+                obj_ptr,
+                i32::try_from(wasmtime_environ::VM_GC_HEADER_KIND_OFFSET).unwrap(),
+            );
+
+            // Write `VMGcHeader::type_index`.
+            let shared_ty = func_env.module_interned_to_shared_ty(&mut builder.cursor(), ty);
+            builder.ins().store(
+                ir::MemFlags::trusted(),
+                shared_ty,
+                obj_ptr,
+                i32::try_from(wasmtime_environ::VM_GC_HEADER_TYPE_INDEX_OFFSET).unwrap(),
+            );
+
+            // Write `VMCopyingHeader::object_size`.
+            builder.ins().istore32(
+                ir::MemFlags::trusted(),
+                aligned_size_64,
+                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<'_>,
@@ -85,28 +270,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,
             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
@@ -142,26 +319,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,
             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,
@@ -191,26 +360,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,
             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,

crates/environ/src/vmoffsets.rs

@@ -480,6 +480,31 @@ pub trait PtrSize {
         self.vmctx_epoch_ptr() + self.size()
     }
 
+    /// Return the offset of the `bump_ptr` field within `VMCopyingHeapData`.
+    #[inline]
+    fn vmcopying_heap_data_bump_ptr(&self) -> u8 {
+        0
+    }
+
+    /// Return the offset of the `active_space_end` field within
+    /// `VMCopyingHeapData`.
+    #[inline]
+    fn vmcopying_heap_data_active_space_end(&self) -> u8 {
+        4
+    }
+
+    /// Return the size of `VMCopyingHeapData`.
+    #[inline]
+    fn size_of_vmcopying_heap_data(&self) -> u8 {
+        8
+    }
+
+    /// Return the alignment of `VMCopyingHeapData`.
+    #[inline]
+    fn align_of_vmcopying_heap_data(&self) -> u8 {
+        4
+    }
+
     /// The offset of the `type_ids` array pointer.
     #[inline]
     fn vmctx_type_ids_array(&self) -> u8 {