use std::{env, ffi, fs, path};
fn var(k: &str) -> String {
env::var(k).unwrap()
}
fn use_masm() -> bool {
env::var("CARGO_CFG_TARGET_ENV") == Ok("msvc".to_string()) && var("HOST").contains("-windows-")
}
fn main() {
let target = var("TARGET");
let arch = var("CARGO_CFG_TARGET_ARCH");
let os = var("CARGO_CFG_TARGET_OS");
let out_dir = path::PathBuf::from(var("OUT_DIR"));
let suffix = env!("CARGO_PKG_VERSION").replace("-", "_").replace(".", "_");
make_extern_kernel_decl_macro(&out_dir, &suffix);
match arch.as_ref() {
"x86_64" => {
let files = preprocess_files("x86_64/fma", &[], &suffix);
match os.as_ref() {
"windows" => {
if use_masm() {
let mut lib_exe = cc::windows_registry::find(&*target, "lib.exe")
.expect("Could not find lib.exe");
lib_exe.arg(format!(
"/out:{}",
out_dir.join("x86_64_fma.lib").to_str().unwrap()
));
for f in files {
let mut obj = f.clone();
obj.set_extension("o");
let mut ml_exe = cc::windows_registry::find(&*target, "ml64.exe")
.expect("Could not find ml64.exe");
if !ml_exe
.arg("/Fo")
.arg(&obj)
.arg("/c")
.arg(&f)
.status()
.unwrap()
.success()
{
for (i, l) in
std::fs::read_to_string(&f).unwrap().lines().enumerate()
{
println!("{:8} {}", i, l);
}
panic!();
}
lib_exe.arg(obj);
}
assert!(lib_exe.status().unwrap().success());
println!("cargo:rustc-link-search=native={}", out_dir.to_str().unwrap());
println!("cargo:rustc-link-lib=static=x86_64_fma");
} else {
cc::Build::new()
.files(files)
.flag("-mfma")
.static_flag(true)
.compile("x86_64_fma");
let _ = fs::remove_file("fma_mmm_f32_16x6.asm");
let _ = fs::remove_file("fma_mmm_i8_8x8.asm");
let _ = fs::remove_file("fma_sigmoid_f32.asm");
let _ = fs::remove_file("fma_tanh_f32.asm");
}
}
"macos" => {
let lib = out_dir.join("libx86_64_fma.a");
if lib.exists() {
std::fs::remove_file(lib).unwrap();
}
let mut lib = std::process::Command::new("xcrun");
lib.args(&["ar", "-rv"]).arg(out_dir.join("libx86_64_fma.a"));
for f in files {
let mut obj = f.clone();
obj.set_extension("o");
assert!(std::process::Command::new("cc")
.args(&["-arch", "x86_64"])
.args(&["-c", "-o"])
.arg(&obj)
.arg(&f)
.status()
.unwrap()
.success());
lib.arg(obj);
}
assert!(lib.status().unwrap().success());
println!("cargo:rustc-link-search=native={}", out_dir.to_str().unwrap());
println!("cargo:rustc-link-lib=static=x86_64_fma");
}
_ => {
cc::Build::new()
.files(files)
.flag("-mfma")
.static_flag(true)
.compile("x86_64_fma");
}
}
}
"arm" | "armv7" => {
let files = preprocess_files("arm32/armvfpv2", &[], &suffix);
cc::Build::new()
.files(files)
.flag("-marm")
.flag("-mfpu=vfp")
.static_flag(true)
.compile("armvfpv2");
let files = preprocess_files(
"arm32/armv7neon",
&[("core", vec!["cortexa7", "cortexa9", "generic"])],
&suffix,
);
cc::Build::new()
.files(files)
.flag("-marm")
.flag("-mfpu=neon")
.static_flag(true)
.compile("armv7neon");
}
"aarch64" => {
let files =
preprocess_files("arm64/arm64simd", &[("core", vec!["a53", "gen"])], &suffix);
cc::Build::new().files(files).static_flag(true).compile("arm64");
}
_ => {}
}
}
type Variant = (&'static str, Vec<&'static str>);
fn preprocess_files(
input: impl AsRef<path::Path>,
variants: &[Variant],
suffix: &str,
) -> Vec<path::PathBuf> {
let out_dir = path::PathBuf::from(var("OUT_DIR"));
let mut files = vec![];
for f in input.as_ref().read_dir().unwrap() {
let f = f.unwrap();
if f.path().extension() == Some(ffi::OsStr::new("tmpl")) {
let tmpl_file = f.path().file_name().unwrap().to_str().unwrap().to_owned();
let concerned_variants: Vec<&Variant> =
variants.iter().filter(|v| tmpl_file.contains(v.0)).collect();
let expanded_variants = concerned_variants.iter().map(|pair| pair.1.len()).product();
for v in 0..expanded_variants {
let mut tmpl_file = tmpl_file.clone();
let mut id = v;
let mut globals = vec![];
for variable in variants {
let key = variable.0;
let value = variable.1[id % variable.1.len()];
globals.push((key, value));
tmpl_file = tmpl_file.replace(key, value);
id /= variable.1.len();
}
let mut file = out_dir.join(tmpl_file);
file.set_extension("S");
preprocess_file(f.path(), &file, &globals, suffix);
files.push(file);
}
}
}
files
}
fn preprocess_file(
template: impl AsRef<path::Path>,
output: impl AsRef<path::Path>,
variants: &[(&'static str, &'static str)],
suffix: &str,
) {
println!("cargo:rerun-if-changed={}", template.as_ref().to_string_lossy());
let family = var("CARGO_CFG_TARGET_FAMILY");
let os = var("CARGO_CFG_TARGET_OS");
let msvc = use_masm();
println!("cargo:rerun-if-changed={}", template.as_ref().to_string_lossy());
let mut input = fs::read_to_string(&template).unwrap();
if msvc {
input =
input.lines().map(|line| line.replace("//", ";")).collect::<Vec<String>>().join("\n");
}
let l = if os == "macos" {
"L"
} else if family == "windows" {
""
} else {
"."
}
.to_owned();
let g = if os == "macos" || os == "ios" { "_" } else { "" };
let mut globals = liquid::object!({
"msvc": msvc,
"family": family,
"os": os,
"L": l,
"G": g,
"suffix": suffix,
});
for (k, v) in variants {
globals.insert(k.to_string().into(), liquid::model::Value::scalar(*v));
}
let partials = load_partials(&template.as_ref().parent().unwrap());
liquid::ParserBuilder::with_stdlib()
.partials(liquid::partials::LazyCompiler::new(partials))
.build()
.unwrap()
.parse(&*input)
.unwrap()
.render_to(&mut fs::File::create(&output).unwrap(), &globals)
.unwrap();
}
fn load_partials(p: &path::Path) -> liquid::partials::InMemorySource {
let mut mem = liquid::partials::InMemorySource::new();
for f in walkdir::WalkDir::new(p) {
let f = f.unwrap();
let ext = f.path().extension().map(|s| s.to_string_lossy()).unwrap_or("".into());
let text = match ext.as_ref() {
"tmpli" => Some(
std::fs::read_to_string(f.path()).unwrap().replace("{{", "{").replace("}}", "}"),
),
"tmpliq" => Some(std::fs::read_to_string(f.path()).unwrap()),
_ => None,
};
if let Some(text) = text {
let key = f.path().strip_prefix(p).unwrap().to_str().unwrap().to_owned();
println!("cargo:rerun-if-changed={}", f.path().to_string_lossy());
mem.add(key, text);
}
}
mem
}
fn make_extern_kernel_decl_macro(out_dir: &path::Path, suffix: &str) {
let macro_decl = r#"
macro_rules! extern_kernel {
(fn $name: ident($($par_name:ident : $par_type: ty ),*) -> $rv: ty) => {
paste! {
extern "C" { fn [<$name _ _suffix>]($(par_name: $par_type),*) -> $rv; }
use [<$name _ _suffix>] as $name;
}
}
}"#
.replace("_suffix", suffix);
std::fs::write(out_dir.join("extern_kernel_macro.rs"), macro_decl).unwrap();
}