Merge pull request #14 from yaroslavyaroslav/metal-build-pipeline

Metal building pipeline enhancement

mopro-msm/Cargo.toml

@@ -2,6 +2,13 @@
 name = "mopro-msm"
 version = "0.1.0"
 edition = "2021"
+build = "build.rs"
+
+[features]
+default = ["macos"]
+
+macos = []
+ios = []
 
 [dependencies]
 ## Shared dependencies

mopro-msm/build.rs

@@ -0,0 +1,123 @@
+use std::{env, path::Path, process::Command};
+
+fn main() {
+    compile_shaders();
+}
+
+fn compile_shaders() {
+    let shader_dir = "src/msm/metal/shader/";
+    let out_dir = env::var("OUT_DIR").unwrap();
+
+    // List your Metal shaders here.
+    let shaders = vec!["all.metal"];
+
+    let shaders_to_check = vec!["all.metal", "msm.h.metal"];
+
+    let mut air_files = vec![];
+
+    // Step 1: Compile every shader to AIR format
+    for shader in &shaders {
+        let shader_path = Path::new(shader_dir).join(shader);
+        let air_output = Path::new(&out_dir).join(format!("{}.air", shader));
+
+        let mut args = vec![
+            "-sdk",
+            get_sdk(),
+            "metal",
+            "-c",
+            "-frecord-sources",
+            shader_path.to_str().unwrap(),
+            "-o",
+            air_output.to_str().unwrap(),
+        ];
+
+        if std::env::var("PROFILE")
+            .map(|profile| profile == "release")
+            .unwrap_or(false)
+        {
+            args.push("-frecord-sources");
+        }
+
+        // Compile shader into .air files
+        let status = Command::new("xcrun")
+            .args(&args)
+            .status()
+            .expect("Shader compilation failed");
+
+        if !status.success() {
+            panic!("Shader compilation failed for {}", shader);
+        }
+
+        air_files.push(air_output);
+    }
+
+    // Step 2: Link all the .air files into a Metallib archive
+    let metallib_output = Path::new(&out_dir).join("msm.metallib");
+
+    let mut metallib_args = vec![
+        "-sdk",
+        get_sdk(),
+        "metal",
+        "-o",
+        metallib_output.to_str().unwrap(),
+    ];
+
+    if std::env::var("PROFILE")
+        .map(|profile| profile == "release")
+        .unwrap_or(false)
+    {
+        metallib_args.push("-frecord-sources");
+    }
+
+    for air_file in &air_files {
+        metallib_args.push(air_file.to_str().unwrap());
+    }
+
+    let status = Command::new("xcrun")
+        .args(&metallib_args)
+        .status()
+        .expect("Failed to link shaders into metallib");
+
+    if !status.success() {
+        panic!("Failed to link shaders into metallib");
+    }
+
+    let symbols_args = vec![
+        "metal-dsymutil",
+        "-flat",
+        "-remove-source",
+        metallib_output.to_str().unwrap(),
+    ];
+
+    let status = Command::new("xcrun")
+        .args(&symbols_args)
+        .status()
+        .expect("Failed to extract symbols");
+
+    if !status.success() {
+        panic!("Failed to extract symbols");
+    }
+
+    // Inform cargo to watch all shader files for changes
+    for shader in &shaders_to_check {
+        let shader_path = Path::new(shader_dir).join(shader);
+        println!("cargo:rerun-if-changed={}", shader_path.to_str().unwrap());
+    }
+}
+
+#[cfg(feature = "macos")]
+fn get_sdk() -> &'static str {
+    "macosx"
+}
+
+#[cfg(not(feature = "macos"))]
+#[cfg(feature = "ios")]
+fn get_sdk() -> &'static str {
+    "iphoneos"
+}
+
+#[cfg(not(feature = "macos"))]
+#[cfg(not(feature = "ios"))]
+fn get_sdk() -> &'static str {
+    panic!("one of the features macos or ios needs to be enabled");
+}

mopro-msm/src/msm/metal/abstraction/state.rs

@@ -3,8 +3,7 @@ use metal::{ComputeCommandEncoderRef, MTLResourceOptions};
 use crate::msm::metal::abstraction::errors::MetalError;
 
 use core::{ffi, mem};
-
-const LIB_DATA: &[u8] = include_bytes!("../shader/msm.metallib");
+use std::{env, fs, path::Path};
 
 /// Structure for abstracting basic calls to a Metal device and saving the state. Used for
 /// implementing GPU parallel computations in Apple machines.
@@ -21,8 +20,13 @@ impl MetalState {
         let device: metal::Device =
             device.unwrap_or(metal::Device::system_default().ok_or(MetalError::DeviceNotFound())?);
 
+        let metallib_path = Path::new(env!("OUT_DIR")).join("msm.metallib");
+
+        let lib_data = fs::read(metallib_path)
+            .expect(format!("Missing metal library on the path {}", env!("OUT_DIR")).as_str());
+
         let library = device
-            .new_library_with_data(LIB_DATA) // TODO: allow different files
+            .new_library_with_data(&lib_data)
             .map_err(MetalError::LibraryError)?;
         let queue = device.new_command_queue();
 

mopro-msm/src/msm/utils/preprocess.rs

@@ -259,6 +259,7 @@ pub fn get_root_path() -> String {
     env!("CARGO_MANIFEST_DIR").to_string()
 }
 
+#[cfg(test)]
 mod tests {
     use super::*;