zoc: Produce disassembly of executable

comgr/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "comgr"
 version = "0.0.0"
-authors = ["Andrzej Janik <vosen@vosen.pl>"]
+authors = ["Andrzej Janik <vosen@vosen.pl>", "Joëlle van Essen <joelle@v-essen.nl>"]
 edition = "2021"
 
 [lib]

comgr/src/lib.rs

@@ -110,10 +110,7 @@ impl Drop for ActionInfo {
     }
 }
 
-pub fn link_bitcode(
-    main_buffer: &[u8],
-    ptx_impl: &[u8],
-) -> Result<DataSet, amd_comgr_status_s> {
+fn link_bitcode(main_buffer: &[u8], ptx_impl: &[u8]) -> Result<DataSet, amd_comgr_status_s> {
     use amd_comgr_sys::*;
     let bitcode_data_set = DataSet::new()?;
     let main_bitcode_data = Data::new(
@@ -136,14 +133,12 @@ pub fn link_bitcode(
     )
 }
 
-pub fn compile_bitcode(
+fn compile_bitcode(
+    linked_data_set: &DataSet,
     gcn_arch: &CStr,
-    main_buffer: &[u8],
-    ptx_impl: &[u8],
-) -> Result<Vec<u8>, amd_comgr_status_s> {
+) -> Result<DataSet, amd_comgr_status_s> {
     use amd_comgr_sys::*;
 
-    let linked_data_set = link_bitcode(main_buffer, ptx_impl)?;
     let compile_to_exec = ActionInfo::new()?;
     compile_to_exec.set_isa_name(gcn_arch)?;
     compile_to_exec.set_language(amd_comgr_language_t::AMD_COMGR_LANGUAGE_LLVM_IR)?;
@@ -176,16 +171,69 @@ pub fn compile_bitcode(
         ]
     };
     compile_to_exec.set_options(common_options.chain(opt_options))?;
-    let exec_data_set = do_action(
+    do_action(
         &linked_data_set,
         &compile_to_exec,
         amd_comgr_action_kind_t::AMD_COMGR_ACTION_COMPILE_SOURCE_TO_EXECUTABLE,
-    )?;
+    )
+}
+
+fn disassemble_exec(
+    exec_data_set: &DataSet,
+    gcn_arch: &CStr,
+) -> Result<DataSet, amd_comgr_status_s> {
+    let action_info = ActionInfo::new()?;
+    action_info.set_isa_name(gcn_arch)?;
+    do_action(
+        &exec_data_set,
+        &action_info,
+        amd_comgr_action_kind_t::AMD_COMGR_ACTION_DISASSEMBLE_EXECUTABLE_TO_SOURCE,
+    )
+}
+
+pub fn get_linked_bitcode_as_bytes(
+    main_buffer: &[u8],
+    ptx_impl: &[u8],
+) -> Result<Vec<u8>, amd_comgr_status_s> {
+    let linked_data_set = link_bitcode(main_buffer, ptx_impl)?;
+    let linked_bitcode =
+        linked_data_set.get_data(amd_comgr_data_kind_t::AMD_COMGR_DATA_KIND_BC, 0)?;
+    linked_bitcode.copy_content()
+}
+
+fn get_executable(
+    gcn_arch: &CStr,
+    main_buffer: &[u8],
+    ptx_impl: &[u8],
+) -> Result<DataSet, amd_comgr_status_s> {
+    let linked_data_set = link_bitcode(main_buffer, ptx_impl)?;
+    let exec_data_set = compile_bitcode(&linked_data_set, gcn_arch)?;
+    Ok(exec_data_set)
+}
+
+pub fn get_executable_as_bytes(
+    gcn_arch: &CStr,
+    main_buffer: &[u8],
+    ptx_impl: &[u8],
+) -> Result<Vec<u8>, amd_comgr_status_s> {
+    let exec_data_set = get_executable(gcn_arch, main_buffer, ptx_impl)?;
     let executable =
-    exec_data_set.get_data(amd_comgr_data_kind_t::AMD_COMGR_DATA_KIND_EXECUTABLE, 0)?;
+        exec_data_set.get_data(amd_comgr_data_kind_t::AMD_COMGR_DATA_KIND_EXECUTABLE, 0)?;
     executable.copy_content()
 }
 
+pub fn get_assembly_as_bytes(
+    gcn_arch: &CStr,
+    main_buffer: &[u8],
+    ptx_impl: &[u8],
+) -> Result<Vec<u8>, amd_comgr_status_s> {
+    let exec_data_set = get_executable(gcn_arch, main_buffer, ptx_impl)?;
+    let disassembled_data_set = disassemble_exec(&exec_data_set, gcn_arch)?;
+    let assembly =
+        disassembled_data_set.get_data(amd_comgr_data_kind_t::AMD_COMGR_DATA_KIND_SOURCE, 0)?;
+    assembly.copy_content()
+}
+
 fn do_action(
     data_set: &DataSet,
     action: &ActionInfo,

compiler/src/error.rs

@@ -1,3 +1,4 @@
+use std::ffi::FromBytesUntilNulError;
 use std::io;
 use std::path::PathBuf;
 use std::str::Utf8Error;
@@ -17,14 +18,14 @@ pub enum CompilerError {
     CheckPathError(PathBuf),
     #[error("Invalid output type: {0}")]
     ParseOutputTypeError(String),
-    #[error("Error parsing PTX: {0}")]
-    PtxParserError(String),
     #[error("Error translating PTX: {0:?}")]
     PtxTranslateError(TranslateError),
     #[error("IO error: {0:?}")]
     IoError(io::Error),
     #[error("Error parsing file: {0:?}")]
     ParseFileError(Utf8Error),
+    #[error("Error: {0}")]
+    GenericError(String)
 }
 
 impl From<hipErrorCode_t> for CompilerError {
@@ -43,7 +44,7 @@ impl From<Vec<PtxError<'_>>> for CompilerError {
     fn from(causes: Vec<PtxError>) -> Self {
         let errors: Vec<String> = causes.iter().map(PtxError::to_string).collect();
         let msg = errors.join("\n");
-        CompilerError::PtxParserError(msg)
+        CompilerError::GenericError(msg)
     }
 }
 
@@ -63,4 +64,10 @@ impl From<TranslateError> for CompilerError {
     fn from(cause: TranslateError) -> Self {
         CompilerError::PtxTranslateError(cause)
     }
+}
+
+impl From<FromBytesUntilNulError> for CompilerError {
+    fn from(cause: FromBytesUntilNulError) -> Self {
+        CompilerError::GenericError(format!("{}", cause))
+    }
 }

compiler/src/main.rs

@@ -1,12 +1,11 @@
 use std::env;
-use std::ffi::{CStr, OsStr};
+use std::ffi::{CStr, CString, OsStr};
 use std::fs::{self, File};
 use std::io::{self, Write};
 use std::mem::MaybeUninit;
 use std::path::{Path, PathBuf};
 use std::str::{self, FromStr};
 
-use amd_comgr_sys::amd_comgr_data_kind_s;
 use bpaf::Bpaf;
 
 mod error;
@@ -28,11 +27,6 @@ fn main() -> Result<(), CompilerError> {
 
     let output_type = opts.output_type.unwrap_or_default();
 
-    match output_type {
-        OutputType::Assembly => todo!(),
-        _ => {}
-    }
-
     let ptx_path = Path::new(&opts.ptx_path).to_path_buf();
     check_path(&ptx_path)?;
 
@@ -43,20 +37,14 @@ fn main() -> Result<(), CompilerError> {
     let ptx = str::from_utf8(&ptx).map_err(CompilerError::from)?;
     let llvm = ptx_to_llvm(ptx).map_err(CompilerError::from)?;
 
-    if output_type == OutputType::LlvmIrPreLinked {
-        write_to_file(&llvm.llvm_ir, &output_path).map_err(CompilerError::from)?;
-        return Ok(());
-    }
-
-    if output_type == OutputType::LlvmIrLinked {
-        let linked_llvm = link_llvm(&llvm)?;
-        write_to_file(&linked_llvm, &output_path).map_err(CompilerError::from)?;
-        return Ok(());
-    }
-
-    let elf = llvm_to_elf(&llvm)?;
-    write_to_file(&elf, &output_path).map_err(CompilerError::from)?;
+    let output = match output_type {
+        OutputType::LlvmIrPreLinked => llvm.llvm_ir,
+        OutputType::LlvmIrLinked => get_linked_bitcode(&llvm)?,
+        OutputType::Elf => get_elf(&llvm)?,
+        OutputType::Assembly => get_assembly(&llvm)?
+    };
 
+    write_to_file(&output, &output_path).map_err(CompilerError::from)?;
     Ok(())
 }
 
@@ -80,22 +68,31 @@ struct LLVMArtifacts {
     llvm_ir: Vec<u8>,
 }
 
-fn link_llvm(llvm: &LLVMArtifacts) -> Result<Vec<u8>, CompilerError> {
-    let linked_bitcode = comgr::link_bitcode(&llvm.bitcode, &llvm.linked_bitcode)?;
-    let data = linked_bitcode.get_data(amd_comgr_data_kind_s::AMD_COMGR_DATA_KIND_BC, 0)?;
-    let linked_llvm = data.copy_content().map_err(CompilerError::from)?;
-    Ok(ptx::bitcode_to_ir(linked_llvm))
-}
-
-fn llvm_to_elf(llvm: &LLVMArtifacts) -> Result<Vec<u8>, CompilerError> {
+fn get_arch() -> Result<CString, CompilerError> {
     use hip_runtime_sys::*;
     unsafe { hipInit(0) }?;
     let mut dev_props: MaybeUninit<hipDeviceProp_tR0600> = MaybeUninit::uninit();
     unsafe { hipGetDevicePropertiesR0600(dev_props.as_mut_ptr(), 0) }?;
     let dev_props = unsafe { dev_props.assume_init() };
-    let gcn_arch = unsafe { CStr::from_ptr(dev_props.gcnArchName.as_ptr()) };
+    let arch = dev_props.gcnArchName;
+    let arch: Vec<u8> = arch.to_vec().iter().map(|&v| i8::to_ne_bytes(v)[0]).collect();
+    let arch = CStr::from_bytes_until_nul(arch.as_slice())?;
+    Ok(CString::from(arch))
+}
 
-    comgr::compile_bitcode(gcn_arch, &llvm.bitcode, &llvm.linked_bitcode).map_err(CompilerError::from)
+fn get_linked_bitcode(llvm: &LLVMArtifacts) -> Result<Vec<u8>, CompilerError> {
+    let linked_bitcode = comgr::get_linked_bitcode_as_bytes(&llvm.bitcode, &llvm.linked_bitcode)?;
+    Ok(ptx::bitcode_to_ir(linked_bitcode))
+}
+
+fn get_elf(llvm: &LLVMArtifacts) -> Result<Vec<u8>, CompilerError> {
+    let arch = get_arch()?;
+    comgr::get_executable_as_bytes(&arch, &llvm.bitcode, &llvm.linked_bitcode).map_err(CompilerError::from)
+}
+
+fn get_assembly(llvm: &LLVMArtifacts) -> Result<Vec<u8>, CompilerError> {
+    let arch = get_arch()?;
+    comgr::get_assembly_as_bytes(&arch, &llvm.bitcode, &llvm.linked_bitcode).map_err(CompilerError::from)
 }
 
 fn check_path(path: &Path) -> Result<(), CompilerError> {

ptx/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "ptx"
 version = "0.0.0"
-authors = ["Andrzej Janik <vosen@vosen.pl>"]
+authors = ["Andrzej Janik <vosen@vosen.pl>", "Joëlle van Essen <joelle@v-essen.nl>"]
 edition = "2021"
 
 [lib]

ptx/src/test/spirv_run/mod.rs

@@ -361,7 +361,7 @@ fn run_hip<Input: From<u8> + Copy + Debug, Output: From<u8> + Copy + Debug + Def
         unsafe { hipStreamCreate(&mut stream) }.unwrap();
         let mut dev_props = unsafe { mem::zeroed() };
         unsafe { hipGetDevicePropertiesR0600(&mut dev_props, dev) }.unwrap();
-        let elf_module = comgr::compile_bitcode(
+        let elf_module = comgr::get_executable_as_bytes(
             unsafe { CStr::from_ptr(dev_props.gcnArchName.as_ptr()) },
             &*module.llvm_ir.write_bitcode_to_memory(),
             module.linked_bitcode(),

zluda/src/impl/module.rs

@@ -28,7 +28,7 @@ pub(crate) fn load_data(module: &mut CUmodule, image: *const std::ffi::c_void) -
     unsafe { hipCtxGetDevice(&mut dev) }?;
     let mut props = unsafe { mem::zeroed() };
     unsafe { hipGetDevicePropertiesR0600(&mut props, dev) }?;
-    let elf_module = comgr::compile_bitcode(
+    let elf_module = comgr::get_executable_as_bytes(
         unsafe { CStr::from_ptr(props.gcnArchName.as_ptr()) },
         &*llvm_module.llvm_ir.write_bitcode_to_memory(),
         llvm_module.linked_bitcode(),