Make fake ptxas just copy files, in case of an ELF binary just load empty module

This commit is contained in:
Andrzej Janik 2025-09-16 21:28:50 +00:00
commit 5be8ccd0a9
2 changed files with 31 additions and 38 deletions

View file

@ -1,6 +1,4 @@
use bpaf::{any, doc::Style, Bpaf, Parser};
use hip_runtime_sys::{hipDeviceProp_tR0600, hipGetDevicePropertiesR0600};
use std::{ffi::CStr, mem};
use bpaf::{any, choice, doc::Style, literal, Bpaf, Parser};
#[derive(Debug, Clone, Bpaf)]
#[allow(dead_code)]
@ -12,6 +10,8 @@ pub struct Options {
#[bpaf(short, long)]
verbose: bool,
#[bpaf(external)]
lineinfo: bool,
#[bpaf(external)]
gpu_name: String,
#[bpaf(long, short('O'), fallback(3))]
opt_level: usize,
@ -19,48 +19,32 @@ pub struct Options {
input: String,
}
fn lineinfo() -> impl Parser<bool> {
choice(["-lineinfo", "--lineinfo"].into_iter().map(|s| {
literal(s)
.anywhere()
.optional()
.map(|_| true)
.fallback(false)
.boxed()
}))
}
// #[bpaf(long, long("gpu_name"), fallback_with(default_arch))]
fn gpu_name() -> impl Parser<String> {
any("", move |s: String| {
Some(s.strip_prefix("-arch=")?.to_owned())
Some(
s.strip_prefix("-arch=")
.or_else(|| s.strip_prefix("--gpu-name="))?
.to_owned(),
)
})
.metavar(&[("-arch=", Style::Literal), ("ARG", Style::Metavar)])
.metavar(&[("--gpu-name=", Style::Literal), ("SM", Style::Metavar)])
.anywhere()
.fallback_with(|| Ok::<String, &'static str>("sm_52".to_string()))
}
fn main() {
let options = options().run();
let comgr = comgr::Comgr::new().unwrap();
unsafe { hip_runtime_sys::hipInit(0) }.unwrap();
let mut dev_props: hipDeviceProp_tR0600 = unsafe { mem::zeroed() };
let (gpu_arch, clock_rate) = get_gpu_arch_and_clock_rate(&mut dev_props);
let input = std::fs::read_to_string(options.input).unwrap();
let ast = ptx_parser::parse_module_checked(&input).unwrap();
let llvm = ptx::to_llvm_module(
ast,
ptx::Attributes {
clock_rate: clock_rate as u32,
},
|_| {},
)
.unwrap();
let elf_binary = comgr::compile_bitcode(
&comgr,
gpu_arch,
&*llvm.llvm_ir.write_bitcode_to_memory(),
&*llvm.linked_bitcode(),
&*llvm.attributes_ir.write_bitcode_to_memory(),
None,
)
.unwrap();
std::fs::write(options.output, elf_binary).unwrap();
}
fn get_gpu_arch_and_clock_rate<'a>(dev_props: &'a mut hipDeviceProp_tR0600) -> (&'a str, i32) {
unsafe { hipGetDevicePropertiesR0600(dev_props, 0) }.unwrap();
let gcn_arch_name = &dev_props.gcnArchName;
let gcn_arch_name = unsafe { CStr::from_ptr(gcn_arch_name.as_ptr()) };
let gcn_arch_name = gcn_arch_name.to_str();
(gcn_arch_name.unwrap(), dev_props.clockRate)
std::fs::copy(&options.input, &options.output).unwrap();
}

View file

@ -20,6 +20,10 @@ impl ZludaObject for Module {
}
}
static EMPTY_PTX: &str = ".version 6.5
.target sm_30
.address_size 64";
// get_ptx takes an `image` that can be anything we support and returns a
// String containing a ptx extracted from `image`.
fn get_ptx<'a>(image: CodeLibraryRef<'a>) -> Result<Cow<'a, str>, CUerror> {
@ -58,7 +62,12 @@ fn cow_bytes_to_str<'a>(data: Cow<'a, [u8]>) -> Option<Cow<'a, str>> {
pub(crate) fn load_hip_module(library: CodeLibraryRef) -> Result<hipModule_t, CUerror> {
let global_state = driver::global_state()?;
let text = get_ptx(library)?;
let maybe_ptx = get_ptx(library);
let text = if cfg!(debug_assertions) {
maybe_ptx?
} else {
maybe_ptx.unwrap_or_else(|_| Cow::Borrowed(EMPTY_PTX))
};
let hip_properties = get_hip_properties()?;
let gcn_arch = get_gcn_arch(&hip_properties)?;
let attributes = ExtraCacheAttributes {