Skip to content

Commit

Permalink
cranelift-frontend: Specify stack map behavior when defining variab…
Browse files Browse the repository at this point in the history 
…les (#8937)

* Rename `declare_needs_stack_map` to `declare_value_needs_stack_map`

* cranelift-frontend: Specify stack map behavior when defining variables

* Use a different method for declaring a var needs a stack map

Instead of adding a parameter to `declare_var`
  • Loading branch information
@fitzgen
fitzgen authored Jul 11, 2024
1 parent 29e948c commit ab2aae5
Show file tree
Hide file tree
Showing 2 changed files with 129 additions and 29 deletions.
155 changes: 129 additions & 26 deletions cranelift/frontend/src/frontend.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,8 @@ pub struct FunctionBuilderContext {
ssa: SSABuilder,
status: SecondaryMap<Block, BlockStatus>,
types: SecondaryMap<Variable, Type>,
needs_stack_map: EntitySet<Value>,
stack_map_vars: EntitySet<Variable>,
stack_map_values: EntitySet<Value>,
dfs: Dfs,
}

Expand Down Expand Up @@ -382,9 +383,15 @@ impl<'a> FunctionBuilder<'a> {
self.handle_ssa_side_effects(side_effects);
}

/// Declares the type of a variable, so that it can be used later (by calling
/// [`FunctionBuilder::use_var`]). This function will return an error if the variable
/// has been previously declared.
/// Declares the type of a variable.
///
/// This allows the variable to be used later (by calling
/// [`FunctionBuilder::use_var`]).
///
/// # Errors
///
/// This function will return an error if the variable has been previously
/// declared.
pub fn try_declare_var(&mut self, var: Variable, ty: Type) -> Result<(), DeclareVariableError> {
if self.func_ctx.types[var] != types::INVALID {
return Err(DeclareVariableError::DeclaredMultipleTimes(var));
Expand All @@ -393,13 +400,37 @@ impl<'a> FunctionBuilder<'a> {
Ok(())
}

/// In order to use a variable (by calling [`FunctionBuilder::use_var`]), you need
/// to first declare its type with this method.
/// Declares the type of a variable, panicking if it is already declared.
///
/// # Panics
///
/// Panics if the variable has already been declared.
pub fn declare_var(&mut self, var: Variable, ty: Type) {
self.try_declare_var(var, ty)
.unwrap_or_else(|_| panic!("the variable {:?} has been declared multiple times", var))
}

/// Declare that all uses of the given variable must be included in stack
/// map metadata.
///
/// All values that are uses of this variable will be spilled to the stack
/// before each safepoint and their location on the stack included in stack
/// maps. Stack maps allow the garbage collector to identify the on-stack GC
/// roots.
///
/// This does not affect any pre-existing uses of the variable.
///
/// # Panics
///
/// Panics if the variable's type is larger than 16 bytes or if this
/// variable has not been declared yet.
pub fn declare_var_needs_stack_map(&mut self, var: Variable) {
let ty = self.func_ctx.types[var];
assert!(ty != types::INVALID);
assert!(ty.bytes() <= 16);
self.func_ctx.stack_map_vars.insert(var);
}

/// Returns the Cranelift IR necessary to use a previously defined user
/// variable, returning an error if this is not possible.
pub fn try_use_var(&mut self, var: Variable) -> Result<Value, UseVariableError> {
Expand All @@ -426,6 +457,13 @@ impl<'a> FunctionBuilder<'a> {
.use_var(self.func, var, ty, self.position.unwrap())
};
self.handle_ssa_side_effects(side_effects);

// If the variable was declared as needing stack maps, then propagate
// that requirement to all values derived from using the variable.
if self.func_ctx.stack_map_vars.contains(var) {
self.declare_value_needs_stack_map(val);
}

Ok(val)
}

Expand Down Expand Up @@ -513,13 +551,13 @@ impl<'a> FunctionBuilder<'a> {
/// # Panics
///
/// Panics if `val` is larger than 16 bytes.
pub fn declare_needs_stack_map(&mut self, val: Value) {
pub fn declare_value_needs_stack_map(&mut self, val: Value) {
// We rely on these properties in `insert_safepoint_spills`.
let size = self.func.dfg.value_type(val).bytes();
assert!(size <= 16);
assert!(size.is_power_of_two());

self.func_ctx.needs_stack_map.insert(val);
self.func_ctx.stack_map_values.insert(val);
}

/// Creates a jump table in the function, to be used by [`br_table`](InstBuilder::br_table) instructions.
Expand Down Expand Up @@ -738,7 +776,7 @@ impl<'a> FunctionBuilder<'a> {
// instruction to the live set. This includes branch arguments,
// as mentioned above.
for val in self.func.dfg.inst_values(inst) {
if self.func_ctx.needs_stack_map.contains(val) {
if self.func_ctx.stack_map_values.contains(val) {
live.insert(val);
}
}
Expand Down Expand Up @@ -844,7 +882,7 @@ impl<'a> FunctionBuilder<'a> {
}
}

if !self.func_ctx.needs_stack_map.is_empty() {
if !self.func_ctx.stack_map_values.is_empty() {
self.insert_safepoint_spills();
}

Expand Down Expand Up @@ -2110,8 +2148,8 @@ block0:
builder.append_block_params_for_function_params(block0);
let a = builder.func.dfg.block_params(block0)[0];
let b = builder.func.dfg.block_params(block0)[1];
builder.declare_needs_stack_map(a);
builder.declare_needs_stack_map(b);
builder.declare_value_needs_stack_map(a);
builder.declare_value_needs_stack_map(b);
builder.switch_to_block(block0);
builder.ins().call(func_ref, &[a]);
builder.ins().jump(block0, &[a, b]);
Expand Down Expand Up @@ -2182,13 +2220,13 @@ block0(v0: i32, v1: i32):
builder.append_block_params_for_function_params(block0);
builder.switch_to_block(block0);
let v0 = builder.ins().iconst(ir::types::I32, 0);
builder.declare_needs_stack_map(v0);
builder.declare_value_needs_stack_map(v0);
let v1 = builder.ins().iconst(ir::types::I32, 1);
builder.declare_needs_stack_map(v1);
builder.declare_value_needs_stack_map(v1);
let v2 = builder.ins().iconst(ir::types::I32, 2);
builder.declare_needs_stack_map(v2);
builder.declare_value_needs_stack_map(v2);
let v3 = builder.ins().iconst(ir::types::I32, 3);
builder.declare_needs_stack_map(v3);
builder.declare_value_needs_stack_map(v3);
builder.ins().call(func_ref, &[v0]);
builder.ins().call(func_ref, &[v0]);
builder.ins().call(func_ref, &[v1]);
Expand Down Expand Up @@ -2286,7 +2324,7 @@ block0:

builder.switch_to_block(block1);
let v1 = builder.ins().iconst(ir::types::I64, 0x12345678);
builder.declare_needs_stack_map(v1);
builder.declare_value_needs_stack_map(v1);
builder.ins().jump(block3, &[]);

builder.switch_to_block(block2);
Expand Down Expand Up @@ -2380,7 +2418,7 @@ block3:
builder.switch_to_block(block0);
let v0 = builder.func.dfg.block_params(block0)[0];
let v1 = builder.ins().iconst(ir::types::I64, 0x12345678);
builder.declare_needs_stack_map(v1);
builder.declare_value_needs_stack_map(v1);
builder.ins().brif(v0, block1, &[], block2, &[]);

builder.switch_to_block(block1);
Expand Down Expand Up @@ -2457,7 +2495,7 @@ block2:
builder.switch_to_block(block0);
let v0 = builder.func.dfg.block_params(block0)[0];
let v1 = builder.ins().iconst(ir::types::I64, 0x12345678);
builder.declare_needs_stack_map(v1);
builder.declare_value_needs_stack_map(v1);
builder.ins().brif(v0, block1, &[], block2, &[]);

builder.switch_to_block(block1);
Expand Down Expand Up @@ -2544,7 +2582,7 @@ block2:
builder.switch_to_block(block0);
let v0 = builder.func.dfg.block_params(block0)[0];
let v1 = builder.ins().iconst(ir::types::I64, 0x12345678);
builder.declare_needs_stack_map(v1);
builder.declare_value_needs_stack_map(v1);
builder.ins().brif(v0, block1, &[], block2, &[]);

builder.switch_to_block(block1);
Expand Down Expand Up @@ -2619,7 +2657,7 @@ block2:
builder.switch_to_block(block0);
let v0 = builder.func.dfg.block_params(block0)[0];
let v1 = builder.ins().iconst(ir::types::I64, 0x12345678);
builder.declare_needs_stack_map(v1);
builder.declare_value_needs_stack_map(v1);
builder.ins().brif(v0, block1, &[], block2, &[]);

builder.switch_to_block(block1);
Expand Down Expand Up @@ -2713,17 +2751,17 @@ block2:

builder.switch_to_block(block1);
let v1 = builder.ins().iconst(ir::types::I64, 1);
builder.declare_needs_stack_map(v1);
builder.declare_value_needs_stack_map(v1);
let v2 = builder.ins().iconst(ir::types::I64, 2);
builder.declare_needs_stack_map(v2);
builder.declare_value_needs_stack_map(v2);
builder.ins().call(func_ref, &[]);
builder.ins().jump(block3, &[v1, v2]);

builder.switch_to_block(block2);
let v3 = builder.ins().iconst(ir::types::I64, 3);
builder.declare_needs_stack_map(v3);
builder.declare_value_needs_stack_map(v3);
let v4 = builder.ins().iconst(ir::types::I64, 4);
builder.declare_needs_stack_map(v4);
builder.declare_value_needs_stack_map(v4);
builder.ins().call(func_ref, &[]);
builder.ins().jump(block3, &[v3, v3]);

Expand Down Expand Up @@ -2826,7 +2864,7 @@ block3:
builder.switch_to_block(block0);
let params = builder.func.dfg.block_params(block0).to_vec();
for val in &params {
builder.declare_needs_stack_map(*val);
builder.declare_value_needs_stack_map(*val);
}
builder.ins().call(func_ref, &[]);
builder.ins().return_(&params);
Expand Down Expand Up @@ -2859,6 +2897,71 @@ block0(v0: i8, v1: i16, v2: i32, v3: i64, v4: i128, v5: f32, v6: f64, v7: i8x16,
stack_store v8, ss4+32
call fn0(), stack_map=[i8 @ ss0+0, i16 @ ss1+0, i32 @ ss2+0, f32 @ ss2+4, i64 @ ss3+0, f64 @ ss3+8, i128 @ ss4+0, i8x16 @ ss4+16, i16x8 @ ss4+32]
return v0, v1, v2, v3, v4, v5, v6, v7, v8
}
"#
.trim()
);
}

#[test]
fn var_needs_stack_map() {
let mut sig = Signature::new(CallConv::SystemV);
sig.params
.push(AbiParam::new(cranelift_codegen::ir::types::I32));
sig.returns
.push(AbiParam::new(cranelift_codegen::ir::types::I32));

let mut fn_ctx = FunctionBuilderContext::new();
let mut func = Function::with_name_signature(UserFuncName::testcase("sample"), sig);
let mut builder = FunctionBuilder::new(&mut func, &mut fn_ctx);

let var = Variable::from_u32(0);
builder.declare_var(var, cranelift_codegen::ir::types::I32);
builder.declare_var_needs_stack_map(var);

let name = builder
.func
.declare_imported_user_function(ir::UserExternalName {
namespace: 0,
index: 0,
});
let signature = builder
.func
.import_signature(Signature::new(CallConv::SystemV));
let func_ref = builder.import_function(ir::ExtFuncData {
name: ir::ExternalName::user(name),
signature,
colocated: true,
});

let block0 = builder.create_block();
builder.append_block_params_for_function_params(block0);
builder.switch_to_block(block0);

let arg = builder.func.dfg.block_params(block0)[0];
builder.def_var(var, arg);

builder.ins().call(func_ref, &[]);

let val = builder.use_var(var);
builder.ins().return_(&[val]);

builder.seal_all_blocks();
builder.finalize();

eprintln!("Actual = {}", func.display());
assert_eq!(
func.display().to_string().trim(),
r#"
function %sample(i32) -> i32 system_v {
ss0 = explicit_slot 4, align = 4
sig0 = () system_v
fn0 = colocated u0:0 sig0
block0(v0: i32):
stack_store v0, ss0
call fn0(), stack_map=[i32 @ ss0+0]
return v0
}
"#
.trim()
Expand Down
3 changes: 0 additions & 3 deletions cranelift/frontend/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,9 +61,6 @@
//! Here is how you build the corresponding Cranelift IR function using [`FunctionBuilderContext`]:
//!
//! ```rust
//! extern crate cranelift_codegen;
//! extern crate cranelift_frontend;
//!
//! use cranelift_codegen::entity::EntityRef;
//! use cranelift_codegen::ir::types::*;
//! use cranelift_codegen::ir::{AbiParam, UserFuncName, Function, InstBuilder, Signature};
Expand Down

0 comments on commit ab2aae5

Please sign in to comment.