Start converting cuModuleLoad to the new tracing infrastructure

1 files changed, 166 insertions, 0 deletions

diff --git a/zluda_dump/src/trace.rs b/zluda_dump/src/trace.rs
new file mode 100644
index 0000000..f060467
--- /dev/null
+++ b/zluda_dump/src/trace.rs
@@ -0,0 +1,166 @@
+use ptx::{ast::PtxError, Token};

+

+use crate::{cuda::CUmodule, log, Settings};

+use std::{

+    collections::HashMap,

+    ffi::{c_void, CStr},

+    fs::{self, File},

+    io::{self, Read, Write},

+    path::PathBuf,

+    rc::Rc,

+};

+

+// This struct is the heart of CUDA state tracking, it:

+// * receives calls from the probes about changes to CUDA state

+// * records updates to the state change

+// * writes out relevant state change and details to disk and log

+pub(crate) struct StateTracker {

+    writer: DumpWriter,

+    modules: HashMap<CUmodule, Option<ParsedModule>>,

+    module_counter: usize,

+}

+

+impl StateTracker {

+    pub(crate) fn new(settings: &Settings) -> Self {

+        StateTracker {

+            writer: DumpWriter::new(settings.dump_dir.clone()),

+            modules: HashMap::new(),

+            module_counter: 0,

+        }

+    }

+

+    pub(crate) fn record_new_module_file(

+        &mut self,

+        module: CUmodule,

+        file_name: *const i8,

+        fn_logger: &mut log::FunctionLogger,

+    ) {

+        let file_name = match unsafe { CStr::from_ptr(file_name) }.to_str() {

+            Ok(f) => f,

+            Err(err) => {

+                fn_logger.log(log::LogEntry::MalformedModulePath(err));

+                return;

+            }

+        };

+        let maybe_io_error = self.try_record_new_module_file(module, fn_logger, file_name);

+        fn_logger.log_io_error(maybe_io_error)

+    }

+

+    fn try_record_new_module_file(

+        &mut self,

+        module: CUmodule,

+        fn_logger: &mut log::FunctionLogger,

+        file_name: &str,

+    ) -> io::Result<()> {

+        let mut module_file = fs::File::open(file_name)?;

+        let mut read_buff = Vec::new();

+        module_file.read_to_end(&mut read_buff)?;

+        self.record_new_module(module, read_buff.as_ptr() as *const _, fn_logger);

+        Ok(())

+    }

+

+    pub(crate) fn record_new_module(

+        &mut self,

+        module: CUmodule,

+        raw_image: *const c_void,

+        fn_logger: &mut log::FunctionLogger,

+    ) {

+        self.module_counter += 1;

+        if unsafe { *(raw_image as *const [u8; 4]) } == *goblin::elf64::header::ELFMAG {

+            self.modules.insert(module, None);

+            // TODO: Parse ELF and write it to disk

+            fn_logger.log(log::LogEntry::UnsupportedModule {

+                module,

+                raw_image,

+                kind: "ELF",

+            })

+        } else if unsafe { *(raw_image as *const [u8; 8]) } == *goblin::archive::MAGIC {

+            self.modules.insert(module, None);

+            // TODO: Figure out how to get size of archive module and write it to disk

+            fn_logger.log(log::LogEntry::UnsupportedModule {

+                module,

+                raw_image,

+                kind: "archive",

+            })

+        } else {

+            self.record_module_ptx(module, raw_image, fn_logger)

+        }

+    }

+

+    fn record_module_ptx(

+        &mut self,

+        module: CUmodule,

+        raw_image: *const c_void,

+        fn_logger: &mut log::FunctionLogger,

+    ) {

+        let module_text = unsafe { CStr::from_ptr(raw_image as *const _) }.to_str();

+        let module_text = match module_text {

+            Ok(m) => m,

+            Err(utf8_err) => {

+                fn_logger.log(log::LogEntry::MalformedModuleText(utf8_err));

+                return;

+            }

+        };

+        fn_logger.log_io_error(self.writer.save_module(self.module_counter, module_text));

+        let mut errors = Vec::new();

+        let ast = ptx::ModuleParser::new().parse(&mut errors, module_text);

+        let ast = match (&*errors, ast) {

+            (&[], Ok(ast)) => ast,

+            (err_vec, res) => {

+                fn_logger.log(log::LogEntry::ModuleParsingError(self.module_counter));

+                fn_logger.log_io_error(self.writer.save_module_error_log(

+                    self.module_counter,

+                    err_vec,

+                    res.err(),

+                ));

+                return;

+            }

+        };

+        // TODO: store kernel names and details

+    }

+}

+

+struct ParsedModule {

+    content: Rc<String>,

+    kernels_args: Option<HashMap<String, Vec<usize>>>,

+}

+

+// This structs writes out information about CUDA execution to the dump dir

+struct DumpWriter {

+    dump_dir: Option<PathBuf>,

+}

+

+impl DumpWriter {

+    fn new(dump_dir: Option<PathBuf>) -> Self {

+        Self { dump_dir }

+    }

+

+    fn save_module(&self, index: usize, text: &str) -> io::Result<()> {

+        let mut dump_file = match &self.dump_dir {

+            None => return Ok(()),

+            Some(d) => d.clone(),

+        };

+        dump_file.push(format!("module_{:04}.ptx", index));

+        let mut file = File::create(dump_file)?;

+        file.write_all(text.as_bytes())?;

+        Ok(())

+    }

+

+    fn save_module_error_log<'input>(

+        &self,

+        index: usize,

+        recoverable: &[ptx::ParseError<usize, Token<'input>, PtxError>],

+        unrecoverable: Option<ptx::ParseError<usize, Token<'input>, PtxError>>,

+    ) -> io::Result<()> {

+        let mut log_file = match &self.dump_dir {

+            None => return Ok(()),

+            Some(d) => d.clone(),

+        };

+        log_file.push(format!("module_{:04}.log", index));

+        let mut file = File::create(log_file)?;

+        for err in unrecoverable.iter().chain(recoverable.iter()) {

+            writeln!(file, "{}", err)?;

+        }

+        Ok(())

+    }

+}