M-Labs/nac3
M-Labs
/
nac3
8
0
Fork 10
Code Issues 97 Pull Requests 20 Releases Wiki Activity

core: WIP - Add tracer runtime

This commit is contained in:
David Mak 2024-08-28 16:33:03 +08:00
parent 7e43a37e7d
commit e2bbedbec9
11 changed files with 306 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
nac3artiq/src/lib.rs
View File

@ -36,8 +36,9 @@ use tempfile::{self, TempDir};
use nac3core::{
codegen::{
concrete_type::ConcreteTypeStore, gen_func_impl, irrt::load_irrt, CodeGenLLVMOptions,
CodeGenTargetMachineOptions, CodeGenTask, WithCall, WorkerRegistry,
concrete_type::ConcreteTypeStore, gen_func_impl, irrt::load_irrt,
tracert::TraceRuntimeConfig, CodeGenLLVMOptions, CodeGenTargetMachineOptions, CodeGenTask,
WithCall, WorkerRegistry,
},
inkwell::{
context::Context,
@ -724,8 +725,13 @@ impl Nac3 {
let membuffer = membuffers.clone();
py.allow_threads(|| {
let (registry, handles) =
WorkerRegistry::create_workers(threads, top_level.clone(), &self.llvm_options, &f);
let (registry, handles) = WorkerRegistry::create_workers(
threads,
top_level.clone(),
&self.llvm_options,
&TraceRuntimeConfig::default(),
&f,
);
registry.add_task(task);
registry.wait_tasks_complete(handles);

2
nac3core/Cargo.toml
View File

@ -6,6 +6,7 @@ edition = "2021"
[features]
no-escape-analysis = []
tracing = []
[dependencies]
itertools = "0.13"
@ -28,4 +29,5 @@ indoc = "2.0"
insta = "=1.11.0"
[build-dependencies]
itertools = "0.13"
regex = "1.10"

81
nac3core/build.rs
View File

@ -1,3 +1,4 @@
use std::ffi::OsStr;
use std::{
env,
fs::File,
@ -6,14 +7,28 @@ use std::{
process::{Command, Stdio},
};
use itertools::Itertools;
use regex::Regex;
fn main() {
struct IRRTCompilation<'a> {
pub file: &'a str,
pub gcc_options: Vec<&'a str>,
pub cargo_instructions: Vec<&'a str>,
}
/// Extracts the extension-less filename from a [`Path`].
fn path_to_extless_filename(path: &Path) -> &str {
path.file_name().map(Path::new).and_then(Path::file_stem).and_then(OsStr::to_str).unwrap()
}
/// Compiles a source C file into LLVM bitcode.
fn compile_file_to_ir(compile_opts: &IRRTCompilation) {
let out_dir = env::var("OUT_DIR").unwrap();
let out_dir = Path::new(&out_dir);
let out_path = Path::new(&out_dir);
let irrt_dir = Path::new("irrt");
let irrt_cpp_path = irrt_dir.join("irrt.cpp");
let path = Path::new(compile_opts.file);
let filename_without_ext = path_to_extless_filename(path);
/*
* HACK: Sadly, clang doesn't let us emit generic LLVM bitcode.
@ -35,22 +50,22 @@ fn main() {
"-",
"-I",
irrt_dir.to_str().unwrap(),
irrt_cpp_path.to_str().unwrap(),
];
// Apply custom flags from IRRTCompilation
flags.extend_from_slice(&compile_opts.gcc_options);
match env::var("PROFILE").as_deref() {
Ok("debug") => {
flags.push("-O0");
flags.push("-DIRRT_DEBUG_ASSERT");
}
Ok("release") => {
flags.push("-O3");
}
Ok("debug") => flags.extend_from_slice(&["-O0", "-DIRRT_DEBUG_ASSERT"]),
Ok("release") => flags.push("-O3"),
flavor => panic!("Unknown or missing build flavor {flavor:?}"),
}
// Tell Cargo to rerun if any file under `irrt_dir` (recursive) changes
println!("cargo:rerun-if-changed={}", irrt_dir.to_str().unwrap());
flags.push(path.to_str().unwrap());
// Tell Cargo to rerun if the main IRRT source is changed
println!("cargo:rerun-if-changed={}", path.to_str().unwrap());
compile_opts.cargo_instructions.iter().for_each(|inst| println!("cargo::{inst}"));
// Compile IRRT and capture the LLVM IR output
let output = Command::new("clang-irrt")
@ -62,8 +77,7 @@ fn main() {
})
.unwrap();
// https://github.com/rust-lang/regex/issues/244
let output = std::str::from_utf8(&output.stdout).unwrap().replace("\r\n", "\n");
let output = std::str::from_utf8(&output.stdout).unwrap();
let mut filtered_output = String::with_capacity(output.len());
// Filter out irrelevant IR
@ -77,7 +91,7 @@ fn main() {
r"(?ms:^define.*?\}$)|(?m:^declare.*?$)|(?m:^%.+?=\s*type\s*\{.+?\}$)|(?m:^@.+?=.+$)",
)
.unwrap();
for f in regex_filter.captures_iter(&output) {
for f in regex_filter.captures_iter(output) {
assert_eq!(f.len(), 1);
filtered_output.push_str(&f[0]);
filtered_output.push('\n');
@ -91,20 +105,47 @@ fn main() {
// Doing `DEBUG_DUMP_IRRT=1 cargo build -p nac3core` dumps the LLVM IR generated
const DEBUG_DUMP_IRRT: &str = "DEBUG_DUMP_IRRT";
println!("cargo:rerun-if-env-changed={DEBUG_DUMP_IRRT}");
if env::var(DEBUG_DUMP_IRRT).is_ok() {
let mut file = File::create(out_dir.join("irrt.ll")).unwrap();
if env::var("DEBUG_DUMP_IRRT").is_ok() {
let mut file = File::create(out_path.join(format!("{filename_without_ext}.ll"))).unwrap();
file.write_all(output.as_bytes()).unwrap();
let mut file = File::create(out_dir.join("irrt-filtered.ll")).unwrap();
let mut file =
File::create(out_path.join(format!("{filename_without_ext}-filtered.ll"))).unwrap();
file.write_all(filtered_output.as_bytes()).unwrap();
}
let mut llvm_as = Command::new("llvm-as-irrt")
.stdin(Stdio::piped())
.arg("-o")
.arg(out_dir.join("irrt.bc"))
.arg(out_path.join(format!("{filename_without_ext}.bc")))
.spawn()
.unwrap();
llvm_as.stdin.as_mut().unwrap().write_all(filtered_output.as_bytes()).unwrap();
assert!(llvm_as.wait().unwrap().success());
}
fn main() {
let irrt_compilations: &[IRRTCompilation] = &[
IRRTCompilation {
file: "irrt/irrt.cpp",
gcc_options: Vec::default(),
cargo_instructions: vec!["rerun-if-changed=irrt/irrt"],
},
IRRTCompilation {
file: "irrt/tracert.cpp",
gcc_options: Vec::default(),
cargo_instructions: Vec::default(),
},
];
assert!(irrt_compilations
.iter()
.map(|comp| comp.file)
.map(Path::new)
.map(path_to_extless_filename)
.all_unique());
for path in irrt_compilations {
compile_file_to_ir(path)
}
}

4
nac3core/irrt/tracert.cpp Normal file
View File

@ -0,0 +1,4 @@
extern "C" {
// stdio.h
int printf(const char *format, ...);
} // extern "C"

53
nac3core/src/codegen/extern_fns.rs
View File

@ -1,8 +1,11 @@
use std::iter::once;
use inkwell::{
attributes::{Attribute, AttributeLoc},
values::{BasicValueEnum, CallSiteValue, FloatValue, IntValue},
values::{BasicValue, BasicValueEnum, CallSiteValue, FloatValue, IntValue},
AddressSpace,
};
use itertools::Either;
use itertools::{Either, Itertools};
use super::CodeGenContext;
@ -191,3 +194,49 @@ generate_linalg_extern_fn!(call_np_linalg_det, "np_linalg_det", 2);
generate_linalg_extern_fn!(call_sp_linalg_lu, "sp_linalg_lu", 3);
generate_linalg_extern_fn!(call_sp_linalg_schur, "sp_linalg_schur", 3);
generate_linalg_extern_fn!(call_sp_linalg_hessenberg, "sp_linalg_hessenberg", 3);
/// Invokes the `printf` function.
pub fn call_printf<'ctx>(
ctx: &CodeGenContext<'ctx, '_>,
format: &str,
args: &[BasicValueEnum<'ctx>],
) -> IntValue<'ctx> {
const FN_NAME: &str = "printf";
let llvm_i8 = ctx.ctx.i8_type();
let llvm_i32 = ctx.ctx.i32_type();
let llvm_pi8 = llvm_i8.ptr_type(AddressSpace::default());
let extern_fn = ctx.module.get_function(FN_NAME).unwrap_or_else(|| {
let fn_type = llvm_i32.fn_type(&[llvm_pi8.into()], true);
let func = ctx.module.add_function(FN_NAME, fn_type, None);
for attr in ["nofree", "nounwind"] {
func.add_attribute(
AttributeLoc::Function,
ctx.ctx.create_enum_attribute(Attribute::get_named_enum_kind_id(attr), 0),
);
}
func
});
let pformat = ctx
.builder
.build_global_string_ptr(&format!("{format}\0"), "")
.map(|v| v.as_basic_value_enum())
.map(BasicValueEnum::into_pointer_value)
.unwrap();
let fn_args = once(&pformat.as_basic_value_enum())
.chain(args)
.copied()
.map(BasicValueEnum::into)
.collect_vec();
ctx.builder
.build_call(extern_fn, fn_args.as_slice(), "")
.map(CallSiteValue::try_as_basic_value)
.map(|v| v.map_left(BasicValueEnum::into_int_value))
.map(Either::unwrap_left)
.unwrap()
}

15
nac3core/src/codegen/mod.rs
View File

@ -39,6 +39,7 @@ use crate::{
use classes::{ListType, NDArrayType, ProxyType, RangeType};
use concrete_type::{ConcreteType, ConcreteTypeEnum, ConcreteTypeStore};
pub use generator::{CodeGenerator, DefaultCodeGenerator};
use tracert::{TraceRuntimeConfig, TraceRuntimeState};
pub mod builtin_fns;
pub mod classes;
@ -50,6 +51,7 @@ pub mod irrt;
pub mod llvm_intrinsics;
pub mod numpy;
pub mod stmt;
pub mod tracert;
#[cfg(test)]
mod test;
@ -219,6 +221,8 @@ pub struct CodeGenContext<'ctx, 'a> {
/// See [`need_sret`].
pub need_sret: bool,
pub tracert_state: Option<TraceRuntimeState>,
/// The current source location.
pub current_loc: Location,
}
@ -266,6 +270,8 @@ pub struct WorkerRegistry {
/// LLVM-related options for code generation.
pub llvm_options: CodeGenLLVMOptions,
tracert_config: TraceRuntimeConfig,
}
impl WorkerRegistry {
@ -275,6 +281,7 @@ impl WorkerRegistry {
generators: Vec<Box<G>>,
top_level_ctx: Arc<TopLevelContext>,
llvm_options: &CodeGenLLVMOptions,
tracert_config: &TraceRuntimeConfig,
f: &Arc<WithCall>,
) -> (Arc<WorkerRegistry>, Vec<thread::JoinHandle<()>>) {
let (sender, receiver) = unbounded();
@ -296,6 +303,7 @@ impl WorkerRegistry {
wait_condvar,
top_level_ctx,
llvm_options: llvm_options.clone(),
tracert_config: tracert_config.clone(),
});
let mut handles = Vec::new();
@ -979,6 +987,13 @@ pub fn gen_func_impl<
unifier,
static_value_store,
need_sret: has_sret,
tracert_state: if cfg!(feature = "tracing")
|| registry.tracert_config.enabled_tags.is_empty()
{
None
} else {
Some(TraceRuntimeState::create(registry.tracert_config.clone()))
},
current_loc: Location::default(),
debug_info: (dibuilder, compile_unit, func_scope.as_debug_info_scope()),
};

17
nac3core/src/codegen/test.rs
View File

@ -18,6 +18,7 @@ use parking_lot::RwLock;
use super::{
classes::{ListType, NDArrayType, ProxyType, RangeType},
concrete_type::ConcreteTypeStore,
tracert::TraceRuntimeConfig,
CodeGenContext, CodeGenLLVMOptions, CodeGenTargetMachineOptions, CodeGenTask, CodeGenerator,
DefaultCodeGenerator, WithCall, WorkerRegistry,
};
@ -245,7 +246,13 @@ fn test_primitives() {
opt_level: OptimizationLevel::Default,
target: CodeGenTargetMachineOptions::from_host_triple(),
};
let (registry, handles) = WorkerRegistry::create_workers(threads, top_level, &llvm_options, &f);
let (registry, handles) = WorkerRegistry::create_workers(
threads,
top_level,
&llvm_options,
&TraceRuntimeConfig::default(),
&f,
);
registry.add_task(task);
registry.wait_tasks_complete(handles);
}
@ -438,7 +445,13 @@ fn test_simple_call() {
opt_level: OptimizationLevel::Default,
target: CodeGenTargetMachineOptions::from_host_triple(),
};
let (registry, handles) = WorkerRegistry::create_workers(threads, top_level, &llvm_options, &f);
let (registry, handles) = WorkerRegistry::create_workers(
threads,
top_level,
&llvm_options,
&TraceRuntimeConfig::default(),
&f,
);
registry.add_task(task);
registry.wait_tasks_complete(handles);
}

110
nac3core/src/codegen/tracert/mod.rs Normal file
View File

@ -0,0 +1,110 @@
use std::panic::Location;
use inkwell::context::Context;
use inkwell::memory_buffer::MemoryBuffer;
use inkwell::module::Module;
use inkwell::values::BasicValueEnum;
use inkwell::AtomicOrdering;
use crate::codegen::{extern_fns, CodeGenContext};
#[derive(Clone, Default, Eq, PartialEq)]
pub struct TraceRuntimeConfig {
pub enabled_tags: Vec<String>,
}
#[derive(Eq, Clone, PartialEq)]
pub struct TraceRuntimeState {
config: TraceRuntimeConfig,
indent: usize,
}
impl TraceRuntimeState {
#[must_use]
pub fn create(config: TraceRuntimeConfig) -> TraceRuntimeState {
TraceRuntimeState { config, indent: 0 }
}
}
#[must_use]
pub fn load_tracert<'ctx>(ctx: &'ctx Context, config: &TraceRuntimeConfig) -> Option<Module<'ctx>> {
if cfg!(feature = "tracing") && !config.enabled_tags.is_empty() {
let bitcode_buf = MemoryBuffer::create_from_memory_range(
include_bytes!(concat!(env!("OUT_DIR"), "/tracert.bc")),
"tracert_bitcode_buffer",
);
let module = Module::parse_bitcode_from_buffer(&bitcode_buf, ctx).unwrap();
return Some(module);
}
None
}
// TODO: Might need to redesign how trace logging should be done
pub fn trace_log<'ctx>(
ctx: &mut CodeGenContext<'ctx, '_>,
tag: &'static str,
format: &'static str,
args: &[BasicValueEnum<'ctx>],
) {
if ctx.tracert_state.is_none() {
return;
}
// TODO: Add indentation
let str = format!("[TRACING] {tag} - {format}\n\0");
extern_fns::call_printf(ctx, &str, args);
}
#[track_caller]
pub fn trace_log_with_location<'ctx>(
ctx: &CodeGenContext<'ctx, '_>,
tag: &'static str,
format: &str,
args: &[BasicValueEnum<'ctx>],
) {
if ctx.tracert_state.is_none() {
return;
}
// TODO: Add indentation
let caller_loc = Location::caller();
let str = format!(
"[TRACING] {}:{}:{}: {tag} - {format}\n\0",
caller_loc.file(),
caller_loc.line(),
caller_loc.column()
);
extern_fns::call_printf(ctx, &str, args);
}
pub fn trace_push_level(ctx: &mut CodeGenContext<'_, '_>) {
let Some(tracert_state) = &mut ctx.tracert_state else {
return;
};
debug_assert!(tracert_state.indent < usize::MAX);
if tracert_state.indent < usize::MAX {
tracert_state.indent = tracert_state.indent.saturating_add(1);
}
}
pub fn trace_pop_level(ctx: &mut CodeGenContext<'_, '_>) {
let Some(tracert_state) = &mut ctx.tracert_state else {
return;
};
debug_assert!(tracert_state.indent > 0);
if tracert_state.indent > 0 {
tracert_state.indent = tracert_state.indent.saturating_sub(1);
}
}
#[inline]
pub fn mfence(ctx: &mut CodeGenContext<'_, '_>) {
if ctx.tracert_state.is_some() {
ctx.builder.build_fence(AtomicOrdering::SequentiallyConsistent, 0, "").unwrap();
}
}

1
nac3standalone/Cargo.toml
View File

@ -6,6 +6,7 @@ edition = "2021"
[features]
no-escape-analysis = ["nac3core/no-escape-analysis"]
tracing = ["nac3core/tracing"]
[dependencies]
parking_lot = "0.12"

8
nac3standalone/demo/demo.c
View File

@ -65,7 +65,7 @@ struct cslice {
size_t len;
};
void output_int32_list(struct cslice* slice) {
void output_int32_list(const struct cslice* slice) {
const int32_t* data = (int32_t*)slice->data;
putchar('[');
@ -80,7 +80,7 @@ void output_int32_list(struct cslice* slice) {
putchar('\n');
}
void output_str(struct cslice* slice) {
void output_str(const struct cslice* slice) {
const char* data = (const char*)slice->data;
for (size_t i = 0; i < slice->len; ++i) {
@ -88,12 +88,12 @@ void output_str(struct cslice* slice) {
}
}
void output_strln(struct cslice* slice) {
void output_strln(const struct cslice* slice) {
output_str(slice);
putchar('\n');
}
uint64_t dbg_stack_address(__attribute__((unused)) struct cslice* slice) {
uint64_t dbg_stack_address(__attribute__((unused)) const struct cslice* slice) {
int i;
void* ptr = (void*)&i;
return (uintptr_t)ptr;

37
nac3standalone/src/main.rs
View File

@ -21,8 +21,11 @@ use parking_lot::{Mutex, RwLock};
use nac3core::{
codegen::{
concrete_type::ConcreteTypeStore, irrt::load_irrt, CodeGenLLVMOptions,
CodeGenTargetMachineOptions, CodeGenTask, DefaultCodeGenerator, WithCall, WorkerRegistry,
concrete_type::ConcreteTypeStore,
irrt::load_irrt,
tracert::{load_tracert, TraceRuntimeConfig},
CodeGenLLVMOptions, CodeGenTargetMachineOptions, CodeGenTask, DefaultCodeGenerator,
WithCall, WorkerRegistry,
},
inkwell::{
memory_buffer::MemoryBuffer, module::Linkage, passes::PassBuilderOptions,
@ -82,6 +85,9 @@ struct CommandLineArgs {
/// Additional target features to enable/disable, specified using the `+`/`-` prefixes.
#[arg(long)]
target_features: Option<String>,
#[arg(long)]
trace: Vec<String>,
}
fn handle_typevar_definition(
@ -282,8 +288,16 @@ fn handle_global_var(
fn main() {
let cli = CommandLineArgs::parse();
let CommandLineArgs { file_name, threads, opt_level, emit_llvm, triple, mcpu, target_features } =
cli;
let CommandLineArgs {
file_name,
threads,
opt_level,
emit_llvm,
triple,
mcpu,
target_features,
trace,
} = cli;
Target::initialize_all(&InitializationConfig::default());
@ -312,6 +326,7 @@ fn main() {
// The default behavior for -O<n> where n>3 defaults to O3 for both Clang and GCC
_ => OptimizationLevel::Aggressive,
};
let tracert_config = TraceRuntimeConfig { enabled_tags: trace };
let target_machine_options = CodeGenTargetMachineOptions {
triple,
@ -356,6 +371,14 @@ fn main() {
irrt.write_bitcode_to_path(Path::new("irrt.bc"));
}
// Process tracert
let tracert = load_tracert(&context, &tracert_config);
if let Some(tracert) = &tracert {
if emit_llvm {
tracert.write_bitcode_to_path(Path::new("tracert.bc"));
}
}
// Process the Python script
let parser_result = parser::parse_program(&program, file_name.into()).unwrap();
@ -464,7 +487,8 @@ fn main() {
let threads = (0..threads)
.map(|i| Box::new(DefaultCodeGenerator::new(format!("module{i}"), size_t)))
.collect();
let (registry, handles) = WorkerRegistry::create_workers(threads, top_level, &llvm_options, &f);
let (registry, handles) =
WorkerRegistry::create_workers(threads, top_level, &llvm_options, &tracert_config, &f);
registry.add_task(task);
registry.wait_tasks_complete(handles);
@ -490,6 +514,9 @@ fn main() {
}
main.link_in_module(irrt).unwrap();
if let Some(tracert) = tracert {
main.link_in_module(tracert).unwrap();
}
// Private all functions except "run"
let mut function_iter = main.get_first_function();