Initial reorg.
This is largely the same as #305, but updated for the current tree.
This commit is contained in:
Dan Gohman
381
crates/misc/py/src/import.rs
Normal file
381
crates/misc/py/src/import.rs
Normal file
@@ -0,0 +1,381 @@
|
||||
//! Support for a calling of an imported function.
|
||||
|
||||
extern crate alloc;
|
||||
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyAny, PyDict, PyTuple};
|
||||
|
||||
use crate::function::Function;
|
||||
use crate::memory::Memory;
|
||||
use crate::value::{read_value_from, write_value_to};
|
||||
use cranelift_codegen::ir::types;
|
||||
use cranelift_codegen::ir::{InstBuilder, StackSlotData, StackSlotKind};
|
||||
use cranelift_codegen::Context;
|
||||
use cranelift_codegen::{binemit, ir, isa};
|
||||
use cranelift_entity::{EntityRef, PrimaryMap};
|
||||
use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
|
||||
use cranelift_wasm::{DefinedFuncIndex, FuncIndex};
|
||||
use target_lexicon::HOST;
|
||||
use wasmtime_environ::{CompiledFunction, Export, Module};
|
||||
use wasmtime_jit::CodeMemory;
|
||||
use wasmtime_runtime::{Imports, InstanceHandle, VMContext, VMFunctionBody};
|
||||
|
||||
use alloc::rc::Rc;
|
||||
use core::cell::RefCell;
|
||||
use core::cmp;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
struct BoundPyFunction {
|
||||
name: String,
|
||||
obj: PyObject,
|
||||
}
|
||||
|
||||
struct ImportObjState {
|
||||
calls: Vec<BoundPyFunction>,
|
||||
#[allow(dead_code)]
|
||||
code_memory: CodeMemory,
|
||||
}
|
||||
|
||||
unsafe extern "C" fn stub_fn(vmctx: *mut VMContext, call_id: u32, values_vec: *mut i64) {
|
||||
let gil = Python::acquire_gil();
|
||||
let py = gil.python();
|
||||
|
||||
let mut instance = InstanceHandle::from_vmctx(vmctx);
|
||||
let (_name, obj) = {
|
||||
let state = instance
|
||||
.host_state()
|
||||
.downcast_mut::<ImportObjState>()
|
||||
.expect("state");
|
||||
let name = state.calls[call_id as usize].name.to_owned();
|
||||
let obj = state.calls[call_id as usize].obj.clone_ref(py);
|
||||
(name, obj)
|
||||
};
|
||||
let module = instance.module_ref();
|
||||
let signature = &module.signatures[module.functions[FuncIndex::new(call_id as usize)]];
|
||||
|
||||
let mut args = Vec::new();
|
||||
for i in 1..signature.params.len() {
|
||||
args.push(read_value_from(
|
||||
py,
|
||||
values_vec.offset(i as isize - 1),
|
||||
signature.params[i].value_type,
|
||||
))
|
||||
}
|
||||
let result = obj.call(py, PyTuple::new(py, args), None).expect("result");
|
||||
for i in 0..signature.returns.len() {
|
||||
let val = if result.is_none() {
|
||||
0.into_py(py) // FIXME default ???
|
||||
} else {
|
||||
if i > 0 {
|
||||
panic!("multiple returns unsupported");
|
||||
}
|
||||
result.clone_ref(py)
|
||||
};
|
||||
write_value_to(
|
||||
py,
|
||||
values_vec.offset(i as isize),
|
||||
signature.returns[i].value_type,
|
||||
val,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a trampoline for invoking a python function.
|
||||
fn make_trampoline(
|
||||
isa: &dyn isa::TargetIsa,
|
||||
code_memory: &mut CodeMemory,
|
||||
fn_builder_ctx: &mut FunctionBuilderContext,
|
||||
call_id: u32,
|
||||
signature: &ir::Signature,
|
||||
) -> *const VMFunctionBody {
|
||||
// Mostly reverse copy of the similar method from wasmtime's
|
||||
// wasmtime-jit/src/compiler.rs.
|
||||
let pointer_type = isa.pointer_type();
|
||||
let mut stub_sig = ir::Signature::new(isa.frontend_config().default_call_conv);
|
||||
|
||||
// Add the `vmctx` parameter.
|
||||
stub_sig.params.push(ir::AbiParam::special(
|
||||
pointer_type,
|
||||
ir::ArgumentPurpose::VMContext,
|
||||
));
|
||||
|
||||
// Add the `call_id` parameter.
|
||||
stub_sig.params.push(ir::AbiParam::new(types::I32));
|
||||
|
||||
// Add the `values_vec` parameter.
|
||||
stub_sig.params.push(ir::AbiParam::new(pointer_type));
|
||||
|
||||
let values_vec_len = 8 * cmp::max(signature.params.len() - 1, signature.returns.len()) as u32;
|
||||
|
||||
let mut context = Context::new();
|
||||
context.func =
|
||||
ir::Function::with_name_signature(ir::ExternalName::user(0, 0), signature.clone());
|
||||
|
||||
let ss = context.func.create_stack_slot(StackSlotData::new(
|
||||
StackSlotKind::ExplicitSlot,
|
||||
values_vec_len,
|
||||
));
|
||||
let value_size = 8;
|
||||
|
||||
{
|
||||
let mut builder = FunctionBuilder::new(&mut context.func, fn_builder_ctx);
|
||||
let block0 = builder.create_ebb();
|
||||
|
||||
builder.append_ebb_params_for_function_params(block0);
|
||||
builder.switch_to_block(block0);
|
||||
builder.seal_block(block0);
|
||||
|
||||
let values_vec_ptr_val = builder.ins().stack_addr(pointer_type, ss, 0);
|
||||
let mflags = ir::MemFlags::trusted();
|
||||
for i in 1..signature.params.len() {
|
||||
if i == 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let val = builder.func.dfg.ebb_params(block0)[i];
|
||||
builder.ins().store(
|
||||
mflags,
|
||||
val,
|
||||
values_vec_ptr_val,
|
||||
((i - 1) * value_size) as i32,
|
||||
);
|
||||
}
|
||||
|
||||
let vmctx_ptr_val = builder.func.dfg.ebb_params(block0)[0];
|
||||
let call_id_val = builder.ins().iconst(types::I32, call_id as i64);
|
||||
|
||||
let callee_args = vec![vmctx_ptr_val, call_id_val, values_vec_ptr_val];
|
||||
|
||||
let new_sig = builder.import_signature(stub_sig.clone());
|
||||
|
||||
let callee_value = builder
|
||||
.ins()
|
||||
.iconst(pointer_type, stub_fn as *const VMFunctionBody as i64);
|
||||
builder
|
||||
.ins()
|
||||
.call_indirect(new_sig, callee_value, &callee_args);
|
||||
|
||||
let mflags = ir::MemFlags::trusted();
|
||||
let mut results = Vec::new();
|
||||
for (i, r) in signature.returns.iter().enumerate() {
|
||||
let load = builder.ins().load(
|
||||
r.value_type,
|
||||
mflags,
|
||||
values_vec_ptr_val,
|
||||
(i * value_size) as i32,
|
||||
);
|
||||
results.push(load);
|
||||
}
|
||||
builder.ins().return_(&results);
|
||||
builder.finalize()
|
||||
}
|
||||
|
||||
let mut code_buf: Vec<u8> = Vec::new();
|
||||
let mut reloc_sink = RelocSink {};
|
||||
let mut trap_sink = binemit::NullTrapSink {};
|
||||
let mut stackmap_sink = binemit::NullStackmapSink {};
|
||||
context
|
||||
.compile_and_emit(
|
||||
isa,
|
||||
&mut code_buf,
|
||||
&mut reloc_sink,
|
||||
&mut trap_sink,
|
||||
&mut stackmap_sink,
|
||||
)
|
||||
.expect("compile_and_emit");
|
||||
|
||||
let mut unwind_info = Vec::new();
|
||||
context.emit_unwind_info(isa, &mut unwind_info);
|
||||
|
||||
code_memory
|
||||
.allocate_for_function(&CompiledFunction {
|
||||
body: code_buf,
|
||||
jt_offsets: context.func.jt_offsets,
|
||||
unwind_info,
|
||||
})
|
||||
.expect("allocate_for_function")
|
||||
.as_ptr()
|
||||
}
|
||||
|
||||
fn parse_annotation_type(s: &str) -> ir::Type {
|
||||
match s {
|
||||
"I32" | "i32" => types::I32,
|
||||
"I64" | "i64" => types::I64,
|
||||
"F32" | "f32" => types::F32,
|
||||
"F64" | "f64" => types::F64,
|
||||
_ => panic!("unknown type in annotations"),
|
||||
}
|
||||
}
|
||||
|
||||
fn get_signature_from_py_annotation(
|
||||
annot: &PyDict,
|
||||
pointer_type: ir::Type,
|
||||
call_conv: isa::CallConv,
|
||||
) -> PyResult<ir::Signature> {
|
||||
let mut params = Vec::new();
|
||||
params.push(ir::AbiParam::special(
|
||||
pointer_type,
|
||||
ir::ArgumentPurpose::VMContext,
|
||||
));
|
||||
let mut returns = None;
|
||||
for (name, value) in annot.iter() {
|
||||
let ty = parse_annotation_type(&value.to_string());
|
||||
match name.to_string().as_str() {
|
||||
"return" => returns = Some(ty),
|
||||
_ => params.push(ir::AbiParam::new(ty)),
|
||||
}
|
||||
}
|
||||
Ok(ir::Signature {
|
||||
params,
|
||||
returns: match returns {
|
||||
Some(r) => vec![ir::AbiParam::new(r)],
|
||||
None => vec![],
|
||||
},
|
||||
call_conv,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn into_instance_from_obj(
|
||||
py: Python,
|
||||
global_exports: Rc<RefCell<HashMap<String, Option<wasmtime_runtime::Export>>>>,
|
||||
obj: &PyAny,
|
||||
) -> PyResult<InstanceHandle> {
|
||||
let isa = {
|
||||
let isa_builder =
|
||||
cranelift_native::builder().expect("host machine is not a supported target");
|
||||
let flag_builder = cranelift_codegen::settings::builder();
|
||||
isa_builder.finish(cranelift_codegen::settings::Flags::new(flag_builder))
|
||||
};
|
||||
|
||||
let mut fn_builder_ctx = FunctionBuilderContext::new();
|
||||
let mut module = Module::new();
|
||||
let mut finished_functions: PrimaryMap<DefinedFuncIndex, *const VMFunctionBody> =
|
||||
PrimaryMap::new();
|
||||
let mut code_memory = CodeMemory::new();
|
||||
|
||||
let pointer_type = types::Type::triple_pointer_type(&HOST);
|
||||
let call_conv = isa::CallConv::triple_default(&HOST);
|
||||
|
||||
let obj = obj.cast_as::<PyDict>()?;
|
||||
let mut bound_functions = Vec::new();
|
||||
let mut dependencies = HashSet::new();
|
||||
let mut memories = PrimaryMap::new();
|
||||
for (name, item) in obj.iter() {
|
||||
if item.is_callable() {
|
||||
let sig = if item.get_type().is_subclass::<Function>()? {
|
||||
// TODO faster calls?
|
||||
let wasm_fn = item.cast_as::<Function>()?;
|
||||
dependencies.insert(wasm_fn.instance.clone());
|
||||
wasm_fn.get_signature()
|
||||
} else if item.hasattr("__annotations__")? {
|
||||
let annot = item.getattr("__annotations__")?.cast_as::<PyDict>()?;
|
||||
get_signature_from_py_annotation(&annot, pointer_type, call_conv)?
|
||||
} else {
|
||||
// TODO support calls without annotations?
|
||||
continue;
|
||||
};
|
||||
|
||||
let sig_id = module.signatures.push(sig.clone());
|
||||
let func_id = module.functions.push(sig_id);
|
||||
module
|
||||
.exports
|
||||
.insert(name.to_string(), Export::Function(func_id));
|
||||
let trampoline = make_trampoline(
|
||||
isa.as_ref(),
|
||||
&mut code_memory,
|
||||
&mut fn_builder_ctx,
|
||||
func_id.index() as u32,
|
||||
&sig,
|
||||
);
|
||||
finished_functions.push(trampoline);
|
||||
|
||||
bound_functions.push(BoundPyFunction {
|
||||
name: name.to_string(),
|
||||
obj: item.into_py(py),
|
||||
});
|
||||
} else if item.get_type().is_subclass::<Memory>()? {
|
||||
let wasm_mem = item.cast_as::<Memory>()?;
|
||||
dependencies.insert(wasm_mem.instance.clone());
|
||||
let plan = wasm_mem.get_plan();
|
||||
let mem_id = module.memory_plans.push(plan);
|
||||
let _mem_id_2 = memories.push(wasm_mem.into_import());
|
||||
assert_eq!(mem_id, _mem_id_2);
|
||||
let _mem_id_3 = module
|
||||
.imported_memories
|
||||
.push((String::from(""), String::from("")));
|
||||
assert_eq!(mem_id, _mem_id_3);
|
||||
module
|
||||
.exports
|
||||
.insert(name.to_string(), Export::Memory(mem_id));
|
||||
}
|
||||
}
|
||||
|
||||
let imports = Imports::new(
|
||||
dependencies,
|
||||
PrimaryMap::new(),
|
||||
PrimaryMap::new(),
|
||||
memories,
|
||||
PrimaryMap::new(),
|
||||
);
|
||||
let data_initializers = Vec::new();
|
||||
let signatures = PrimaryMap::new();
|
||||
|
||||
code_memory.publish();
|
||||
|
||||
let import_obj_state = ImportObjState {
|
||||
calls: bound_functions,
|
||||
code_memory,
|
||||
};
|
||||
|
||||
Ok(InstanceHandle::new(
|
||||
Rc::new(module),
|
||||
global_exports,
|
||||
finished_functions.into_boxed_slice(),
|
||||
imports,
|
||||
&data_initializers,
|
||||
signatures.into_boxed_slice(),
|
||||
None,
|
||||
Box::new(import_obj_state),
|
||||
)
|
||||
.expect("instance"))
|
||||
}
|
||||
|
||||
/// We don't expect trampoline compilation to produce any relocations, so
|
||||
/// this `RelocSink` just asserts that it doesn't recieve any.
|
||||
struct RelocSink {}
|
||||
|
||||
impl binemit::RelocSink for RelocSink {
|
||||
fn reloc_ebb(
|
||||
&mut self,
|
||||
_offset: binemit::CodeOffset,
|
||||
_reloc: binemit::Reloc,
|
||||
_ebb_offset: binemit::CodeOffset,
|
||||
) {
|
||||
panic!("trampoline compilation should not produce ebb relocs");
|
||||
}
|
||||
fn reloc_external(
|
||||
&mut self,
|
||||
_offset: binemit::CodeOffset,
|
||||
_reloc: binemit::Reloc,
|
||||
_name: &ir::ExternalName,
|
||||
_addend: binemit::Addend,
|
||||
) {
|
||||
panic!("trampoline compilation should not produce external symbol relocs");
|
||||
}
|
||||
fn reloc_constant(
|
||||
&mut self,
|
||||
_code_offset: binemit::CodeOffset,
|
||||
_reloc: binemit::Reloc,
|
||||
_constant_offset: ir::ConstantOffset,
|
||||
) {
|
||||
panic!("trampoline compilation should not produce constant relocs");
|
||||
}
|
||||
fn reloc_jt(
|
||||
&mut self,
|
||||
_offset: binemit::CodeOffset,
|
||||
_reloc: binemit::Reloc,
|
||||
_jt: ir::JumpTable,
|
||||
) {
|
||||
panic!("trampoline compilation should not produce jump table relocs");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user