From 8fe3a2c980bdc5c8a13737ab27e68964237f0875 Mon Sep 17 00:00:00 2001 From: ComputelessComputer <63365510+ComputelessComputer@users.noreply.github.com> Date: Tue, 19 May 2026 21:07:10 +0900 Subject: [PATCH] fix(soniqo): bundle MLX metallib Compile the MLX Metal shader library during the Soniqo Swift bridge build and copy it beside dev and bundled macOS binaries. --- apps/desktop/src-tauri/resources/.gitignore | 1 + apps/desktop/src-tauri/tauri.conf.json | 1 + .../desktop/src-tauri/tauri.conf.staging.json | 1 + crates/transcribe-soniqo/build.rs | 281 +++++++++++++++++- 4 files changed, 283 insertions(+), 1 deletion(-) diff --git a/apps/desktop/src-tauri/resources/.gitignore b/apps/desktop/src-tauri/resources/.gitignore index 0121ce40f6..00e5edf4fa 100644 --- a/apps/desktop/src-tauri/resources/.gitignore +++ b/apps/desktop/src-tauri/resources/.gitignore @@ -4,3 +4,4 @@ stt ttt cli/* !cli/.gitkeep +mlx.metallib diff --git a/apps/desktop/src-tauri/tauri.conf.json b/apps/desktop/src-tauri/tauri.conf.json index 52cea4ac02..24d14c65bd 100644 --- a/apps/desktop/src-tauri/tauri.conf.json +++ b/apps/desktop/src-tauri/tauri.conf.json @@ -59,6 +59,7 @@ } }, "files": { + "MacOS/mlx.metallib": "resources/mlx.metallib", "Resources/AppIcon.incs": "resources/stable/AppIcon.icns", "Resources/Assets.car": "resources/stable/Assets.car" } diff --git a/apps/desktop/src-tauri/tauri.conf.staging.json b/apps/desktop/src-tauri/tauri.conf.staging.json index 89245be091..e44c385e29 100644 --- a/apps/desktop/src-tauri/tauri.conf.staging.json +++ b/apps/desktop/src-tauri/tauri.conf.staging.json @@ -21,6 +21,7 @@ ], "macOS": { "files": { + "MacOS/mlx.metallib": "resources/mlx.metallib", "Resources/AppIcon.incs": "resources/stable/AppIcon.icns", "Resources/Assets.car": "resources/stable/Assets.car" }, diff --git a/crates/transcribe-soniqo/build.rs b/crates/transcribe-soniqo/build.rs index f9328acc44..612b4ccdbb 100644 --- a/crates/transcribe-soniqo/build.rs +++ b/crates/transcribe-soniqo/build.rs @@ -1,9 +1,10 @@ #[cfg(target_os = "macos")] use std::{ collections::BTreeSet, - env, + env, fs, io, path::{Path, PathBuf}, process::Command, + time::SystemTime, }; #[cfg(target_os = "macos")] @@ -48,6 +49,281 @@ fn target_is_macos_apple_silicon() -> bool { && env::var("CARGO_CFG_TARGET_ARCH").is_ok_and(|value| value == "aarch64") } +#[cfg(target_os = "macos")] +fn cargo_profile() -> &'static str { + match env::var("PROFILE").as_deref() { + Ok("release") => "release", + _ => "debug", + } +} + +#[cfg(target_os = "macos")] +fn profile_target_dir(out_dir: &Path) -> Option { + out_dir.ancestors().nth(3).map(Path::to_path_buf) +} + +#[cfg(target_os = "macos")] +fn find_mlx_metallib(root: &Path, profile: &str) -> Option { + [ + root.join(profile).join("mlx.metallib"), + root.join("arm64-apple-macosx") + .join(profile) + .join("mlx.metallib"), + ] + .into_iter() + .find(|path| path.exists()) +} + +#[cfg(target_os = "macos")] +fn swift_profile_dir(root: &Path, profile: &str) -> PathBuf { + let direct = root.join(profile); + if direct.exists() { + return direct; + } + + root.join("arm64-apple-macosx").join(profile) +} + +#[cfg(target_os = "macos")] +fn collect_files(root: &Path, predicate: impl Fn(&Path) -> bool) -> io::Result> { + let mut files = Vec::new(); + let mut stack = vec![root.to_path_buf()]; + + while let Some(path) = stack.pop() { + for entry in fs::read_dir(path)? { + let path = entry?.path(); + if path.is_dir() { + stack.push(path); + } else if predicate(&path) { + files.push(path); + } + } + } + + files.sort(); + Ok(files) +} + +#[cfg(target_os = "macos")] +fn newest_modified_time(files: &[PathBuf]) -> io::Result { + files + .iter() + .map(|path| fs::metadata(path)?.modified()) + .try_fold(SystemTime::UNIX_EPOCH, |newest, modified| { + modified.map(|modified| newest.max(modified)) + }) +} + +#[cfg(target_os = "macos")] +fn metallib_is_fresh(metallib: &Path, inputs: &[PathBuf]) -> bool { + let Ok(output_modified) = fs::metadata(metallib).and_then(|metadata| metadata.modified()) + else { + return false; + }; + let Ok(input_modified) = newest_modified_time(inputs) else { + return false; + }; + + output_modified >= input_modified +} + +#[cfg(target_os = "macos")] +fn run_command(mut command: Command, context: &str) { + let output = command + .output() + .unwrap_or_else(|error| panic!("failed to run {context}: {error}")); + + if !output.status.success() { + panic!( + "{context} failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + } +} + +#[cfg(target_os = "macos")] +fn compile_mlx_metallib(swift_build_dir: &Path, profile: &str) -> PathBuf { + let mlx_swift_dir = swift_build_dir.join("checkouts").join("mlx-swift"); + let kernels_dir = mlx_swift_dir + .join("Source") + .join("Cmlx") + .join("mlx") + .join("mlx") + .join("backend") + .join("metal") + .join("kernels"); + let output_dir = swift_profile_dir(swift_build_dir, profile); + let output_metallib = output_dir.join("mlx.metallib"); + + let metal_sources = collect_files(&kernels_dir, |path| { + path.extension().is_some_and(|ext| ext == "metal") + && !path + .file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with("_nax.metal")) + }) + .unwrap_or_else(|error| { + panic!( + "failed to collect Soniqo MLX Metal sources from {}: {error}", + kernels_dir.display() + ) + }); + + if metal_sources.is_empty() { + panic!( + "no Soniqo MLX Metal sources found under {}", + kernels_dir.display() + ); + } + + let metal_headers = collect_files(&kernels_dir, |path| { + path.extension().is_some_and(|ext| ext == "h") + && !path + .file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with("_nax.h")) + }) + .unwrap_or_else(|error| { + panic!( + "failed to collect Soniqo MLX Metal headers from {}: {error}", + kernels_dir.display() + ) + }); + let inputs = metal_sources + .iter() + .chain(metal_headers.iter()) + .cloned() + .collect::>(); + + if metallib_is_fresh(&output_metallib, &inputs) { + return output_metallib; + } + + fs::create_dir_all(&output_dir).unwrap_or_else(|error| { + panic!( + "failed to create Soniqo MLX metallib output directory {}: {error}", + output_dir.display() + ) + }); + + let air_dir = output_dir.join("mlx-metallib-air"); + if air_dir.exists() { + fs::remove_dir_all(&air_dir).unwrap_or_else(|error| { + panic!( + "failed to clean Soniqo MLX Metal temp directory {}: {error}", + air_dir.display() + ) + }); + } + fs::create_dir_all(&air_dir).unwrap_or_else(|error| { + panic!( + "failed to create Soniqo MLX Metal temp directory {}: {error}", + air_dir.display() + ) + }); + + let include_root = mlx_swift_dir.join("Source").join("Cmlx").join("mlx"); + let mut air_files = Vec::with_capacity(metal_sources.len()); + + for (index, source) in metal_sources.iter().enumerate() { + let air_file = air_dir.join(format!("{index}.air")); + let mut command = Command::new("xcrun"); + command + .args([ + "-sdk", + "macosx", + "metal", + "-x", + "metal", + "-std=metal3.2", + "-mmacosx-version-min=15.0", + "-Wall", + "-Wextra", + "-fno-fast-math", + "-Wno-c++17-extensions", + "-Wno-c++20-extensions", + "-c", + ]) + .arg(source) + .arg(format!("-I{}", kernels_dir.display())) + .arg(format!("-I{}", include_root.display())) + .arg("-o") + .arg(&air_file); + run_command(command, &format!("compiling {}", source.display())); + air_files.push(air_file); + } + + let mut command = Command::new("xcrun"); + command.args(["-sdk", "macosx", "metallib"]); + command.args(&air_files); + command.arg("-o").arg(&output_metallib); + run_command(command, "linking Soniqo MLX metallib"); + + output_metallib +} + +#[cfg(target_os = "macos")] +fn copy_if_changed(source: &Path, destination: &Path) -> io::Result<()> { + if destination.exists() && fs::read(source)? == fs::read(destination)? { + return Ok(()); + } + + if let Some(parent) = destination.parent() { + fs::create_dir_all(parent)?; + } + fs::copy(source, destination)?; + + Ok(()) +} + +#[cfg(target_os = "macos")] +fn desktop_resource_metallib_path() -> Option { + let manifest_dir = PathBuf::from(env::var_os("CARGO_MANIFEST_DIR")?); + Some( + manifest_dir + .parent()? + .parent()? + .join("apps/desktop/src-tauri/resources/mlx.metallib"), + ) +} + +#[cfg(target_os = "macos")] +fn build_mlx_metallib() { + let out_dir = PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR is set by Cargo")); + let swift_build_dir = out_dir.join("swift-rs").join("soniqo-swift"); + let profile = cargo_profile(); + let built_metallib = compile_mlx_metallib(&swift_build_dir, profile); + let metallib = find_mlx_metallib(&swift_build_dir, profile).unwrap_or_else(|| { + panic!( + "Soniqo MLX metallib build completed but mlx.metallib was not found under {}", + swift_build_dir.display() + ) + }); + assert_eq!(built_metallib, metallib); + + if let Some(target_dir) = profile_target_dir(&out_dir) { + let destination = target_dir.join("mlx.metallib"); + copy_if_changed(&metallib, &destination).unwrap_or_else(|error| { + panic!( + "failed to copy Soniqo MLX metallib to {}: {error}", + destination.display() + ) + }); + } + + if let Some(destination) = desktop_resource_metallib_path() { + if destination.parent().is_some_and(Path::exists) { + copy_if_changed(&metallib, &destination).unwrap_or_else(|error| { + panic!( + "failed to copy Soniqo MLX metallib to {}: {error}", + destination.display() + ) + }); + } + } +} + fn main() { #[cfg(target_os = "macos")] { @@ -62,6 +338,8 @@ fn main() { .with_package("soniqo-swift", "./swift-lib/") .link(); + build_mlx_metallib(); + for path in swift_runtime_rpaths() { println!("cargo:rustc-link-arg=-Wl,-rpath,{path}"); } @@ -69,6 +347,7 @@ fn main() { println!("cargo:rustc-link-lib=c++"); println!("cargo:rerun-if-changed=swift-lib/src"); println!("cargo:rerun-if-changed=swift-lib/Package.swift"); + println!("cargo:rerun-if-changed=swift-lib/Package.resolved"); } #[cfg(not(target_os = "macos"))]