Skip to content

cranelift-wasm: Attach table OOB traps to loads/stores #8171

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

Merged
jameysharp merged 1 commit into from
Mar 19, 2024
Merged

cranelift-wasm: Attach table OOB traps to loads/stores #8171

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
8 changes: 4 additions & 4 deletions cranelift/wasm/src/environ/dummy.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -530,8 +530,8 @@ impl<'dummy_environment> FuncEnvironment for DummyFuncEnvironment<'dummy_environ
let pointer_type = self.pointer_type();
self.ensure_table_exists(builder.func, table_index);
let table = self.tables[table_index].as_ref().unwrap();
let table_entry_addr = table.prepare_table_addr(builder, index, pointer_type, true);
let flags = ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
let (table_entry_addr, flags) =
table.prepare_table_addr(builder, index, pointer_type, true);
let value = builder
.ins()
.load(self.reference_type(), flags, table_entry_addr, 0);
Expand All @@ -548,8 +548,8 @@ impl<'dummy_environment> FuncEnvironment for DummyFuncEnvironment<'dummy_environ
let pointer_type = self.pointer_type();
self.ensure_table_exists(builder.func, table_index);
let table = self.tables[table_index].as_ref().unwrap();
let table_entry_addr = table.prepare_table_addr(builder, index, pointer_type, true);
let flags = ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
let (table_entry_addr, flags) =
table.prepare_table_addr(builder, index, pointer_type, true);
builder.ins().store(flags, value, table_entry_addr, 0);
Ok(())
}
Expand Down
34 changes: 17 additions & 17 deletions cranelift/wasm/src/table.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ impl TableData {
mut index: ir::Value,
addr_ty: ir::Type,
enable_table_access_spectre_mitigation: bool,
) -> ir::Value {
) -> (ir::Value, ir::MemFlags) {
let index_ty = pos.func.dfg.value_type(index);

// Start with the bounds check. Trap if `index + 1 > bound`.
Expand All @@ -60,16 +60,10 @@ impl TableData {
let oob = pos
.ins()
.icmp(IntCC::UnsignedGreaterThanOrEqual, index, bound);
pos.ins().trapnz(oob, ir::TrapCode::TableOutOfBounds);

// If Spectre mitigations are enabled, we will use a comparison to
// short-circuit the computed table element address to the start
// of the table on the misspeculation path when out-of-bounds.
let spectre_oob_cmp = if enable_table_access_spectre_mitigation {
Some((index, bound))
} else {
None
};
Comment on lines -68 to -72
Copy link
Copy Markdown
Member

@alexcrichton alexcrichton Mar 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AFAIK spectre mitigations are typically a flag due to their performance impact, but in this case because the spectre-related code is straight-line and actually involves less branches, is this actually the more performant path? Would it be worth, for example, unconditionally using select_spectre_guard?

Copy link
Copy Markdown
Contributor Author

@jameysharp jameysharp Mar 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I actually totally confused myself with the same question at first! But I ran it by @cfallin who confirmed the following: A bounds-check branch is almost always predicted correctly, and that means that if we don't care about Spectre, an out-of-order CPU can make more progress by speculating the load and continuing on to dependent operations past that, rather than stalling on the bounds-check computation.

So although I haven't measured it, there should generally be a performance advantage to sticking with the trapnz approach, for anyone who ain't afraid of no Spectre.

if !enable_table_access_spectre_mitigation {
pos.ins().trapnz(oob, ir::TrapCode::TableOutOfBounds);
}

// Convert `index` to `addr_ty`.
if index_ty != addr_ty {
Expand All @@ -91,14 +85,20 @@ impl TableData {

let element_addr = pos.ins().iadd(base, offset);

if let Some((index, bound)) = spectre_oob_cmp {
let cond = pos
.ins()
.icmp(IntCC::UnsignedGreaterThanOrEqual, index, bound);
// If out-of-bounds, choose the table base on the misspeculation path.
pos.ins().select_spectre_guard(cond, base, element_addr)
let base_flags = ir::MemFlags::new()
.with_aligned()
.with_alias_region(Some(ir::AliasRegion::Table));
if enable_table_access_spectre_mitigation {
// Short-circuit the computed table element address to a null pointer
// when out-of-bounds. The consumer of this address will trap when
// trying to access it.
let zero = pos.ins().iconst(addr_ty, 0);
(
pos.ins().select_spectre_guard(oob, zero, element_addr),
base_flags.with_trap_code(Some(ir::TrapCode::TableOutOfBounds)),
)
} else {
element_addr
(element_addr, base_flags.with_trap_code(None))
}
}
}
16 changes: 7 additions & 9 deletions crates/cranelift/src/func_environ.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -955,13 +955,12 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
// contents, we check for a null entry here, and
// if null, we take a slow-path that invokes a
// libcall.
let table_entry_addr = table_data.prepare_table_addr(
let (table_entry_addr, flags) = table_data.prepare_table_addr(
builder,
index,
pointer_type,
self.isa.flags().enable_table_access_spectre_mitigation(),
);
let flags = ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
let value = builder.ins().load(pointer_type, flags, table_entry_addr, 0);
// Mask off the "initialized bit". See documentation on
// FUNCREF_INIT_BIT in crates/environ/src/ref_bits.rs for more
Expand Down Expand Up @@ -1548,13 +1547,12 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
// Load the table element.
self.ensure_table_exists(builder.func, table_index);
let table_data = self.tables[table_index].as_ref().unwrap();
let elem_addr = table_data.prepare_table_addr(
let (elem_addr, flags) = table_data.prepare_table_addr(
builder,
index,
pointer_type,
self.isa.flags().enable_table_access_spectre_mitigation(),
);
let flags = ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
let elem = builder.ins().load(reference_type, flags, elem_addr, 0);

let elem_is_null = builder.ins().is_null(elem);
Expand Down Expand Up @@ -1659,7 +1657,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
WasmHeapType::Func | WasmHeapType::Concrete(_) | WasmHeapType::NoFunc => {
match plan.style {
TableStyle::CallerChecksSignature => {
let table_entry_addr = table_data.prepare_table_addr(
let (table_entry_addr, flags) = table_data.prepare_table_addr(
builder,
index,
pointer_type,
Expand All @@ -1671,8 +1669,6 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
let value_with_init_bit = builder
.ins()
.bor_imm(value, Imm64::from(FUNCREF_INIT_BIT as i64));
let flags =
ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
builder
.ins()
.store(flags, value_with_init_bit, table_entry_addr, 0);
Expand Down Expand Up @@ -1736,14 +1732,16 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
// not recorded in the stack map (which would lead to the GC
// saving a reference to a deallocated object, and then using it
// after its been freed).
let table_entry_addr = table_data.prepare_table_addr(
let (table_entry_addr, flags) = table_data.prepare_table_addr(
builder,
index,
pointer_type,
self.isa.flags().enable_table_access_spectre_mitigation(),
);
let flags = ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
let current_elem = builder.ins().load(pointer_type, flags, table_entry_addr, 0);

// After the load, a store to the same address can't trap.
let flags = ir::MemFlags::trusted().with_alias_region(Some(ir::AliasRegion::Table));
builder.ins().store(flags, value, table_entry_addr, 0);

// If value is not null, increment `value`'s ref count.
Expand Down
7 changes: 3 additions & 4 deletions tests/disas/icall-simd.wat
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,13 @@
;; block0(v0: i64, v1: i64, v2: i32, v3: i8x16):
;; @0033 v5 = iconst.i32 23
;; @0033 v6 = icmp uge v2, v5 ; v5 = 23
;; @0033 trapnz v6, table_oob
;; @0033 v7 = uextend.i64 v2
;; @0033 v8 = global_value.i64 gv4
;; @0033 v9 = ishl_imm v7, 3
;; @0033 v10 = iadd v8, v9
;; @0033 v11 = icmp uge v2, v5 ; v5 = 23
;; @0033 v12 = select_spectre_guard v11, v8, v10
;; @0033 v13 = load.i64 notrap aligned table v12
;; @0033 v11 = iconst.i64 0
;; @0033 v12 = select_spectre_guard v6, v11, v10 ; v11 = 0
;; @0033 v13 = load.i64 table_oob aligned table v12
;; @0033 v14 = band_imm v13, -2
;; @0033 brif v13, block3(v14), block2
;;
Expand Down
7 changes: 3 additions & 4 deletions tests/disas/icall.wat
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,13 @@
;; block0(v0: i64, v1: i64, v2: i32, v3: f32):
;; @0033 v5 = iconst.i32 23
;; @0033 v6 = icmp uge v2, v5 ; v5 = 23
;; @0033 trapnz v6, table_oob
;; @0033 v7 = uextend.i64 v2
;; @0033 v8 = global_value.i64 gv4
;; @0033 v9 = ishl_imm v7, 3
;; @0033 v10 = iadd v8, v9
;; @0033 v11 = icmp uge v2, v5 ; v5 = 23
;; @0033 v12 = select_spectre_guard v11, v8, v10
;; @0033 v13 = load.i64 notrap aligned table v12
;; @0033 v11 = iconst.i64 0
;; @0033 v12 = select_spectre_guard v6, v11, v10 ; v11 = 0
;; @0033 v13 = load.i64 table_oob aligned table v12
;; @0033 v14 = band_imm v13, -2
;; @0033 brif v13, block3(v14), block2
;;
Expand Down
24 changes: 6 additions & 18 deletions tests/disas/table-get-fixed-size.wat
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,20 +33,14 @@
;; @0052 v3 = iconst.i32 0
;; @0054 v4 = iconst.i32 7
;; @0054 v5 = icmp uge v3, v4 ; v3 = 0, v4 = 7
;; @0054 brif v5, block6, block7
;;
;; block6 cold:
;; @0054 trap table_oob
;;
;; block7:
;; @0054 v6 = uextend.i64 v3 ; v3 = 0
;; @0054 v7 = load.i64 notrap aligned v25+72
;; v26 = iconst.i64 3
;; @0054 v8 = ishl v6, v26 ; v26 = 3
;; @0054 v9 = iadd v7, v8
;; @0054 v10 = icmp.i32 uge v3, v4 ; v3 = 0, v4 = 7
;; @0054 v11 = select_spectre_guard v10, v7, v9
;; @0054 v12 = load.r64 notrap aligned table v11
;; @0054 v10 = iconst.i64 0
;; @0054 v11 = select_spectre_guard v5, v10, v9 ; v10 = 0
;; @0054 v12 = load.r64 table_oob aligned table v11
;; v2 -> v12
;; @0054 v13 = is_null v12
;; @0054 brif v13, block2, block3
Expand Down Expand Up @@ -99,20 +93,14 @@
;; v25 -> v0
;; @005b v4 = iconst.i32 7
;; @005b v5 = icmp uge v2, v4 ; v4 = 7
;; @005b brif v5, block6, block7
;;
;; block6 cold:
;; @005b trap table_oob
;;
;; block7:
;; @005b v6 = uextend.i64 v2
;; @005b v7 = load.i64 notrap aligned v25+72
;; v26 = iconst.i64 3
;; @005b v8 = ishl v6, v26 ; v26 = 3
;; @005b v9 = iadd v7, v8
;; @005b v10 = icmp.i32 uge v2, v4 ; v4 = 7
;; @005b v11 = select_spectre_guard v10, v7, v9
;; @005b v12 = load.r64 notrap aligned table v11
;; @005b v10 = iconst.i64 0
;; @005b v11 = select_spectre_guard v5, v10, v9 ; v10 = 0
;; @005b v12 = load.r64 table_oob aligned table v11
;; v3 -> v12
;; @005b v13 = is_null v12
;; @005b brif v13, block2, block3
Expand Down
24 changes: 6 additions & 18 deletions tests/disas/table-get.wat
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,20 +34,14 @@
;; @0051 v3 = iconst.i32 0
;; @0053 v4 = load.i32 notrap aligned v25+80
;; @0053 v5 = icmp uge v3, v4 ; v3 = 0
;; @0053 brif v5, block6, block7
;;
;; block6 cold:
;; @0053 trap table_oob
;;
;; block7:
;; @0053 v6 = uextend.i64 v3 ; v3 = 0
;; @0053 v7 = load.i64 notrap aligned v26+72
;; v27 = iconst.i64 3
;; @0053 v8 = ishl v6, v27 ; v27 = 3
;; @0053 v9 = iadd v7, v8
;; @0053 v10 = icmp.i32 uge v3, v4 ; v3 = 0
;; @0053 v11 = select_spectre_guard v10, v7, v9
;; @0053 v12 = load.r64 notrap aligned table v11
;; @0053 v10 = iconst.i64 0
;; @0053 v11 = select_spectre_guard v5, v10, v9 ; v10 = 0
;; @0053 v12 = load.r64 table_oob aligned table v11
;; v2 -> v12
;; @0053 v13 = is_null v12
;; @0053 brif v13, block2, block3
Expand Down Expand Up @@ -102,20 +96,14 @@
;; v26 -> v0
;; @005a v4 = load.i32 notrap aligned v25+80
;; @005a v5 = icmp uge v2, v4
;; @005a brif v5, block6, block7
;;
;; block6 cold:
;; @005a trap table_oob
;;
;; block7:
;; @005a v6 = uextend.i64 v2
;; @005a v7 = load.i64 notrap aligned v26+72
;; v27 = iconst.i64 3
;; @005a v8 = ishl v6, v27 ; v27 = 3
;; @005a v9 = iadd v7, v8
;; @005a v10 = icmp.i32 uge v2, v4
;; @005a v11 = select_spectre_guard v10, v7, v9
;; @005a v12 = load.r64 notrap aligned table v11
;; @005a v10 = iconst.i64 0
;; @005a v11 = select_spectre_guard v5, v10, v9 ; v10 = 0
;; @005a v12 = load.r64 table_oob aligned table v11
;; v3 -> v12
;; @005a v13 = is_null v12
;; @005a brif v13, block2, block3
Expand Down
32 changes: 10 additions & 22 deletions tests/disas/table-set-fixed-size.wat
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,22 +34,16 @@
;; @0052 v3 = iconst.i32 0
;; @0056 v4 = iconst.i32 7
;; @0056 v5 = icmp uge v3, v4 ; v3 = 0, v4 = 7
;; @0056 brif v5, block7, block8
;;
;; block7 cold:
;; @0056 trap table_oob
;;
;; block8:
;; @0056 v6 = uextend.i64 v3 ; v3 = 0
;; @0056 v7 = load.i64 notrap aligned v23+72
;; v24 = iconst.i64 3
;; @0056 v8 = ishl v6, v24 ; v24 = 3
;; @0056 v9 = iadd v7, v8
;; @0056 v10 = icmp.i32 uge v3, v4 ; v3 = 0, v4 = 7
;; @0056 v11 = select_spectre_guard v10, v7, v9
;; @0056 v12 = load.i64 notrap aligned table v11
;; @0056 store.r64 notrap aligned table v2, v11
;; @0056 v13 = is_null.r64 v2
;; @0056 v10 = iconst.i64 0
;; @0056 v11 = select_spectre_guard v5, v10, v9 ; v10 = 0
;; @0056 v12 = load.i64 table_oob aligned table v11
;; @0056 store notrap aligned table v2, v11
;; @0056 v13 = is_null v2
;; @0056 brif v13, block3, block2
;;
;; block2:
Expand Down Expand Up @@ -102,22 +96,16 @@
;; v23 -> v0
;; @005f v4 = iconst.i32 7
;; @005f v5 = icmp uge v2, v4 ; v4 = 7
;; @005f brif v5, block7, block8
;;
;; block7 cold:
;; @005f trap table_oob
;;
;; block8:
;; @005f v6 = uextend.i64 v2
;; @005f v7 = load.i64 notrap aligned v23+72
;; v24 = iconst.i64 3
;; @005f v8 = ishl v6, v24 ; v24 = 3
;; @005f v9 = iadd v7, v8
;; @005f v10 = icmp.i32 uge v2, v4 ; v4 = 7
;; @005f v11 = select_spectre_guard v10, v7, v9
;; @005f v12 = load.i64 notrap aligned table v11
;; @005f store.r64 notrap aligned table v3, v11
;; @005f v13 = is_null.r64 v3
;; @005f v10 = iconst.i64 0
;; @005f v11 = select_spectre_guard v5, v10, v9 ; v10 = 0
;; @005f v12 = load.i64 table_oob aligned table v11
;; @005f store notrap aligned table v3, v11
;; @005f v13 = is_null v3
;; @005f brif v13, block3, block2
;;
;; block2:
Expand Down
Loading