use std::collections::HashMap; use std::fs; use std::path::Path; use std::process::Command; use std::sync::Arc; use inkwell::{ AddressSpace, values::BasicValueEnum, passes::{PassManager, PassManagerBuilder}, targets::*, OptimizationLevel, }; use pyo3::prelude::*; use pyo3::{exceptions, types::PyList}; use rustpython_parser::{ ast::{self, StrRef}, parser, }; use parking_lot::RwLock; use nac3core::{ codegen::{CodeGenTask, WithCall, WorkerRegistry}, symbol_resolver::SymbolResolver, toplevel::{composer::TopLevelComposer, TopLevelContext, TopLevelDef, GenCall}, typecheck::typedef::{FunSignature, FuncArg}, }; use nac3core::{ toplevel::DefinitionId, typecheck::{type_inferencer::PrimitiveStore, typedef::Type}, }; use crate::symbol_resolver::Resolver; mod symbol_resolver; #[derive(PartialEq, Clone, Copy)] enum Isa { RiscV, CortexA9, } // TopLevelComposer is unsendable as it holds the unification table, which is // unsendable due to Rc. Arc would cause a performance hit. #[pyclass(unsendable, name = "NAC3")] struct Nac3 { isa: Isa, primitive: PrimitiveStore, builtins_ty: HashMap, builtins_def: HashMap, pyid_to_def: Arc>>, pyid_to_type: Arc>>, composer: TopLevelComposer, top_level: Option>, to_be_registered: Vec, } impl Nac3 { fn register_module_impl(&mut self, obj: PyObject) -> PyResult<()> { let mut name_to_pyid: HashMap = HashMap::new(); let (module_name, source_file) = Python::with_gil(|py| -> PyResult<(String, String)> { let obj: &PyAny = obj.extract(py)?; let builtins = PyModule::import(py, "builtins")?; let id_fn = builtins.getattr("id")?; let members: &PyList = PyModule::import(py, "inspect")? .getattr("getmembers")? .call1((obj,))? .cast_as()?; for member in members.iter() { let key: &str = member.get_item(0)?.extract()?; let val = id_fn.call1((member.get_item(1)?,))?.extract()?; name_to_pyid.insert(key.into(), val); } Ok(( obj.getattr("__name__")?.extract()?, obj.getattr("__file__")?.extract()?, )) })?; let source = fs::read_to_string(source_file).map_err(|e| { exceptions::PyIOError::new_err(format!("failed to read input file: {}", e)) })?; let parser_result = parser::parse_program(&source).map_err(|e| { exceptions::PySyntaxError::new_err(format!("parse error: {}", e)) })?; let resolver = Arc::new(Box::new(Resolver { id_to_type: self.builtins_ty.clone().into(), id_to_def: self.builtins_def.clone().into(), pyid_to_def: self.pyid_to_def.clone(), pyid_to_type: self.pyid_to_type.clone(), class_names: Default::default(), name_to_pyid: name_to_pyid.clone(), }) as Box); let mut name_to_def = HashMap::new(); let mut name_to_type = HashMap::new(); for mut stmt in parser_result.into_iter() { let include = match stmt.node { ast::StmtKind::ClassDef { ref decorator_list, ref mut body, .. } => { let kernels = decorator_list.iter().any(|decorator| { if let ast::ExprKind::Name { id, .. } = decorator.node { id.to_string() == "kernel" || id.to_string() == "portable" } else { false } }); body.retain(|stmt| { if let ast::StmtKind::FunctionDef { ref decorator_list, .. } = stmt.node { decorator_list.iter().any(|decorator| { if let ast::ExprKind::Name { id, .. } = decorator.node { id.to_string() == "kernel" || id.to_string() == "portable" } else { false } }) } else { true } }); kernels } ast::StmtKind::FunctionDef { ref decorator_list, .. } => decorator_list.iter().any(|decorator| { if let ast::ExprKind::Name { id, .. } = decorator.node { let id = id.to_string(); id == "extern" || id == "portable" || id == "kernel" } else { false } }), _ => false, }; if include { let (name, def_id, ty) = self .composer .register_top_level(stmt, Some(resolver.clone()), module_name.clone()) .unwrap(); name_to_def.insert(name, def_id); if let Some(ty) = ty { name_to_type.insert(name, ty); } } } let mut map = self.pyid_to_def.write(); for (name, def) in name_to_def.into_iter() { map.insert(*name_to_pyid.get(&name).unwrap(), def); } let mut map = self.pyid_to_type.write(); for (name, ty) in name_to_type.into_iter() { map.insert(*name_to_pyid.get(&name).unwrap(), ty); } Ok(()) } } #[pymethods] impl Nac3 { #[new] fn new(isa: &str) -> PyResult { let isa = match isa { "riscv" => Isa::RiscV, "cortexa9" => Isa::CortexA9, _ => return Err(exceptions::PyValueError::new_err("invalid ISA")), }; let primitive: PrimitiveStore = TopLevelComposer::make_primitives().0; let mut builtins = vec![]; if isa == Isa::RiscV { builtins.push(( "now_mu".into(), FunSignature { args: vec![], ret: primitive.int64, vars: HashMap::new(), }, Arc::new(GenCall::new(Box::new( |ctx, _, _, _| { let i64_type = ctx.ctx.i64_type(); let now = ctx.module.get_global("now").unwrap_or_else(|| ctx.module.add_global(i64_type, None, "now")); let now_raw = ctx.builder.build_load(now.as_pointer_value(), "now"); if let BasicValueEnum::IntValue(now_raw) = now_raw { let i64_32 = i64_type.const_int(32, false).into(); let now_lo = ctx.builder.build_left_shift(now_raw, i64_32, "now_shl"); let now_hi = ctx.builder.build_right_shift(now_raw, i64_32, false, "now_lshr").into(); Some(ctx.builder.build_or(now_lo, now_hi, "now_or").into()) } else { unreachable!() } } ))) )); builtins.push(( "at_mu".into(), FunSignature { args: vec![FuncArg { name: "t".into(), ty: primitive.int64, default_value: None, }], ret: primitive.none, vars: HashMap::new(), }, Arc::new(GenCall::new(Box::new( |ctx, _, _, args| { let i32_type = ctx.ctx.i64_type(); let i64_type = ctx.ctx.i64_type(); let i64_32 = i64_type.const_int(32, false).into(); if let BasicValueEnum::IntValue(time) = args[0].1 { let time_hi = ctx.builder.build_int_truncate(ctx.builder.build_right_shift(time, i64_32, false, "now_lshr"), i32_type, "now_trunc"); let time_lo = ctx.builder.build_int_truncate(time, i32_type, "now_trunc"); let now = ctx.module.get_global("now").unwrap_or_else(|| ctx.module.add_global(i64_type, None, "now")); let now_hiptr = ctx.builder.build_bitcast(now, i32_type.ptr_type(AddressSpace::Generic), "now_bitcast"); if let BasicValueEnum::PointerValue(now_hiptr) = now_hiptr { let now_loptr = unsafe { ctx.builder.build_gep(now_hiptr, &[i32_type.const_int(1, false).into()], "now_gep") }; ctx.builder.build_store(now_hiptr, time_hi); ctx.builder.build_store(now_loptr, time_lo); None } else { unreachable!(); } } else { unreachable!(); } } ))) )); builtins.push(( "delay_mu".into(), FunSignature { args: vec![FuncArg { name: "dt".into(), ty: primitive.int64, default_value: None, }], ret: primitive.none, vars: HashMap::new(), }, Arc::new(GenCall::new(Box::new( |ctx, _, _, args| { let i32_type = ctx.ctx.i64_type(); let i64_type = ctx.ctx.i64_type(); let i64_32 = i64_type.const_int(32, false).into(); let now = ctx.module.get_global("now").unwrap_or_else(|| ctx.module.add_global(i64_type, None, "now")); let now_raw = ctx.builder.build_load(now.as_pointer_value(), "now"); if let (BasicValueEnum::IntValue(now_raw), BasicValueEnum::IntValue(dt)) = (now_raw, args[0].1) { let now_lo = ctx.builder.build_left_shift(now_raw, i64_32, "now_shl"); let now_hi = ctx.builder.build_right_shift(now_raw, i64_32, false, "now_lshr").into(); let now_val = ctx.builder.build_or(now_lo, now_hi, "now_or"); let time = ctx.builder.build_int_add(now_val, dt, "now_add"); let time_hi = ctx.builder.build_int_truncate(ctx.builder.build_right_shift(time, i64_32, false, "now_lshr"), i32_type, "now_trunc"); let time_lo = ctx.builder.build_int_truncate(time, i32_type, "now_trunc"); let now_hiptr = ctx.builder.build_bitcast(now, i32_type.ptr_type(AddressSpace::Generic), "now_bitcast"); if let BasicValueEnum::PointerValue(now_hiptr) = now_hiptr { let now_loptr = unsafe { ctx.builder.build_gep(now_hiptr, &[i32_type.const_int(1, false).into()], "now_gep") }; ctx.builder.build_store(now_hiptr, time_hi); ctx.builder.build_store(now_loptr, time_lo); None } else { unreachable!(); } } else { unreachable!(); } } ))) )); } let (composer, builtins_def, builtins_ty) = TopLevelComposer::new(builtins); Ok(Nac3 { isa, primitive, builtins_ty, builtins_def, composer, top_level: None, pyid_to_def: Default::default(), pyid_to_type: Default::default(), to_be_registered: Default::default(), }) } fn register_module(&mut self, obj: PyObject) { self.to_be_registered.push(obj); } fn analyze(&mut self) -> PyResult<()> { for obj in std::mem::take(&mut self.to_be_registered).into_iter() { self.register_module_impl(obj)?; } self.composer.start_analysis(true).unwrap(); self.top_level = Some(Arc::new(self.composer.make_top_level_context())); Ok(()) } fn compile_method(&mut self, class: u64, method_name: String) -> PyResult<()> { let top_level = self.top_level.as_ref().unwrap(); let module_resolver; let instance = { let defs = top_level.definitions.read(); let class_def = defs[self.pyid_to_def.read().get(&class).unwrap().0].write(); let mut method_def = if let TopLevelDef::Class { methods, resolver, .. } = &*class_def { module_resolver = Some(resolver.clone().unwrap()); if let Some((_name, _unification_key, definition_id)) = methods .iter() .find(|method| method.0.to_string() == method_name) { defs[definition_id.0].write() } else { return Err(exceptions::PyValueError::new_err("method not found")); } } else { return Err(exceptions::PyTypeError::new_err( "parent object is not a class", )); }; // FIXME: what is this for? What happens if the kernel is called twice? if let TopLevelDef::Function { instance_to_stmt, instance_to_symbol, .. } = &mut *method_def { instance_to_symbol.insert("".to_string(), method_name); instance_to_stmt[""].clone() } else { unreachable!() } }; let signature = FunSignature { args: vec![], ret: self.primitive.none, vars: HashMap::new(), }; let task = CodeGenTask { subst: Default::default(), symbol_name: "__modinit__".to_string(), body: instance.body, signature, resolver: module_resolver.unwrap(), unifier: top_level.unifiers.read()[instance.unifier_id].clone(), calls: instance.calls, }; let isa = self.isa; let f = Arc::new(WithCall::new(Box::new(move |module| { let builder = PassManagerBuilder::create(); builder.set_optimization_level(OptimizationLevel::Aggressive); let passes = PassManager::create(()); builder.populate_module_pass_manager(&passes); passes.run_on(module); let (triple, features) = match isa { Isa::RiscV => (TargetTriple::create("riscv32-unknown-linux"), "+a,+m"), Isa::CortexA9 => ( TargetTriple::create("armv7-unknown-linux-gnueabihf"), "+dsp,+fp16,+neon,+vfp3", ), }; let target = Target::from_triple(&triple).expect("couldn't create target from target triple"); let target_machine = target .create_target_machine( &triple, "", features, OptimizationLevel::Default, RelocMode::PIC, CodeModel::Default, ) .expect("couldn't create target machine"); target_machine .write_to_file( module, FileType::Object, Path::new(&format!("{}.o", module.get_name().to_str().unwrap())), ) .expect("couldn't write module to file"); }))); let thread_names: Vec = (0..4).map(|i| format!("module{}", i)).collect(); let threads: Vec<_> = thread_names.iter().map(|s| s.as_str()).collect(); let (registry, handles) = WorkerRegistry::create_workers(&threads, top_level.clone(), f); registry.add_task(task); registry.wait_tasks_complete(handles); let mut linker_args = vec![ "-shared".to_string(), "--eh-frame-hdr".to_string(), "-Tkernel.ld".to_string(), "-x".to_string(), "-o".to_string(), "module.elf".to_string(), ]; linker_args.extend(thread_names.iter().map(|name| name.to_owned() + ".o")); if let Ok(linker_status) = Command::new("ld.lld").args(linker_args).status() { if !linker_status.success() { return Err(exceptions::PyRuntimeError::new_err( "failed to start linker", )); } } else { return Err(exceptions::PyRuntimeError::new_err( "linker returned non-zero status code", )); } Ok(()) } } #[pymodule] fn nac3artiq(_py: Python, m: &PyModule) -> PyResult<()> { Target::initialize_all(&InitializationConfig::default()); m.add_class::()?; Ok(()) }