refactor: remove runtime shader compilation and constants generation, use precompiled Metal library instead

yaroslavyaroslav · yaroslavyaroslav · commit 76925b18f716 · 2025-06-05T00:41:20.000+03:00
Closes: #77
diff --git a/mopro-msm/build.rs b/mopro-msm/build.rs
@@ -51,14 +51,21 @@ fn compile_shaders() -> std::io::Result<()> {
 
     // Compile combined source to AIR
     let target = env::var("TARGET").unwrap_or_default();
-    let sdk = if target.contains("apple-ios") { "iphoneos" } else { "macosx" };
+    let sdk = if target.contains("apple-ios") {
+        "iphoneos"
+    } else {
+        "macosx"
+    };
     let air = shader_out_dir.join("msm.air");
     let status = Command::new("xcrun")
         .args(&[
-            "-sdk", sdk,
-            "metal", "-c",
+            "-sdk",
+            sdk,
+            "metal",
+            "-c",
             combined.to_str().unwrap(),
-            "-o", air.to_str().unwrap(),
+            "-o",
+            air.to_str().unwrap(),
         ])
         .status()
         .expect("Failed to invoke metal");
@@ -72,10 +79,12 @@ fn compile_shaders() -> std::io::Result<()> {
     let msm_lib = shader_out_dir.join("msm.metallib");
     let status = Command::new("xcrun")
         .args(&[
-            "-sdk", sdk,
+            "-sdk",
+            sdk,
             "metallib",
             air_path,
-            "-o", msm_lib.to_str().unwrap(),
+            "-o",
+            msm_lib.to_str().unwrap(),
         ])
         .status()
         .expect("Failed to invoke metallib");
diff --git a/mopro-msm/src/msm/metal_msm/host/shader.rs b/mopro-msm/src/msm/metal_msm/host/shader.rs
@@ -21,8 +21,6 @@ use ark_ff::{BigInt, PrimeField, Zero};
 use num_bigint::BigUint;
 use std::fs;
 use std::path::PathBuf;
-use std::process::Command;
-use std::string::String;
 
 macro_rules! write_constant_array {
     ($data:expr, $name:expr, $values:expr, $size:expr) => {
@@ -36,108 +34,6 @@ macro_rules! write_constant_array {
     };
 }
 
-/// Get the shader directory path using CARGO_MANIFEST_DIR for proper OS file location
-/// This function provides robust path resolution that works across different build environments
-pub fn get_shader_dir() -> PathBuf {
-    let manifest_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
-    let shader_path = manifest_dir
-        .join("src")
-        .join("msm")
-        .join("metal_msm")
-        .join("shader");
-
-    // Verify the path exists, if not try alternative locations
-    if shader_path.exists() {
-        shader_path
-    } else {
-        // Fallback: try relative path from workspace root
-        let workspace_shader_path = manifest_dir
-            .parent() // go up one level from mopro-msm to workspace root
-            .unwrap_or(&manifest_dir)
-            .join("mopro-msm")
-            .join("src")
-            .join("msm")
-            .join("metal_msm")
-            .join("shader");
-
-        if workspace_shader_path.exists() {
-            workspace_shader_path
-        } else {
-            // Final fallback: use the original path and let error handling take care of it
-            shader_path
-        }
-    }
-}
-
-pub fn compile_metal(path_from_cargo_manifest_dir: &str, input_filename: &str) -> String {
-    let input_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
-        .join(path_from_cargo_manifest_dir)
-        .join(input_filename);
-    let c = input_path.clone().into_os_string().into_string().unwrap();
-
-    let lib = input_path.clone().into_os_string().into_string().unwrap();
-    let lib = format!("{}.lib", lib);
-
-    let exe = if cfg!(target_os = "ios") {
-        Command::new("xcrun")
-            .args([
-                "-sdk",
-                "iphoneos",
-                "metal",
-                "-std=metal3.2",
-                "-target",
-                "air64-apple-ios18.0",
-                "-fmetal-enable-logging",
-                "-o",
-                lib.as_str(),
-                c.as_str(),
-            ])
-            .output()
-            .expect("failed to compile")
-    } else if cfg!(target_os = "macos") {
-        let macos_version = std::process::Command::new("sw_vers")
-            .args(["-productVersion"])
-            .output()
-            .ok()
-            .and_then(|output| String::from_utf8(output.stdout).ok())
-            .and_then(|version| {
-                version
-                    .trim()
-                    .split('.')
-                    .next()
-                    .and_then(|major| major.parse::<u32>().ok())
-            })
-            .unwrap_or(0);
-
-        let mut args = vec!["-sdk", "macosx", "metal"];
-
-        // Only specify Metal 3.2 for metal logging if macOS version is 15.0 or higher
-        if macos_version >= 15 {
-            args.extend([
-                "-std=metal3.2",
-                "-target",
-                "air64-apple-macos15.0",
-                "-fmetal-enable-logging",
-            ]);
-        }
-
-        args.extend(["-o", lib.as_str(), c.as_str()]);
-
-        Command::new("xcrun")
-            .args(args)
-            .output()
-            .expect("failed to compile")
-    } else {
-        panic!("Unsupported architecture");
-    };
-
-    if exe.stderr.len() != 0 {
-        panic!("{}", String::from_utf8(exe.stderr).unwrap());
-    }
-
-    lib
-}
-
 pub fn write_constants(
     filepath: &str,
     num_limbs: usize,
@@ -293,16 +189,6 @@ pub fn write_constants(
 pub mod tests {
     use super::*;
 
-    #[test]
-    #[serial_test::serial]
-    pub fn test_compile() {
-        let lib_filepath = compile_metal(
-            "../mopro-msm/src/msm/metal_msm/shader",
-            "bigint/bigint_add_unsafe.metal",
-        );
-        println!("{}", lib_filepath);
-    }
-
     #[test]
     #[serial_test::serial]
     pub fn test_write_constants() {
diff --git a/mopro-msm/src/msm/metal_msm/utils/metal_wrapper.rs b/mopro-msm/src/msm/metal_msm/utils/metal_wrapper.rs
@@ -1,7 +1,7 @@
 use crate::msm::metal_msm::host::gpu::{
     create_buffer, create_empty_buffer, get_default_device, read_buffer,
 };
-use crate::msm::metal_msm::host::shader::{compile_metal, get_shader_dir, write_constants};
+// no runtime shader compilation, drop constants generation here
 use crate::msm::metal_msm::utils::barrett_params::calc_barrett_mu;
 use crate::msm::metal_msm::utils::mont_params::{calc_mont_radix, calc_nsafe, calc_rinv_and_n0};
 
@@ -13,6 +13,8 @@ use num_bigint::BigUint;
 use once_cell::sync::Lazy;
 use std::collections::HashMap;
 use std::sync::Mutex;
+// Embed the precompiled Metal library
+include!(concat!(env!("OUT_DIR"), "/built_shaders.rs"));
 
 /// Cache of precomputed constants
 static CONSTANTS_CACHE: Lazy<Mutex<HashMap<(usize, u32), MSMConstants>>> =
@@ -143,42 +145,15 @@ impl MetalHelper {
         let encoder =
             command_buffer.compute_command_encoder_with_descriptor(compute_pass_descriptor);
 
-        // Setup shader constants
-        let constants = get_or_calc_constants(config.num_limbs, config.log_limb_size);
-        write_constants(
-            get_shader_dir().to_str().unwrap(),
-            config.num_limbs,
-            config.log_limb_size,
-            constants.n0,
-            constants.nsafe,
-        );
-
-        // Prepare full shader path
-        let shader_path = format!(
-            "{}/{}",
-            get_shader_dir().to_str().unwrap(),
-            config.shader_file
-        );
-        let parts: Vec<&str> = shader_path.rsplitn(2, '/').collect();
-        let shader_dir = if parts.len() > 1 { parts[1] } else { "" };
-        let shader_file = parts[0];
-
-        // Compile shader
-        let library_path = compile_metal(shader_dir, shader_file);
-        let library = self.device.new_library_with_file(library_path).unwrap();
+
+        // Load precompiled Metal library and create pipeline
+        let library = self.device.new_library_with_data(MSM_METALLIB).unwrap();
         let kernel = library
             .get_function(config.kernel_name.as_str(), None)
             .unwrap();
-
-        // Create pipeline
-        let pipeline_state_descriptor = ComputePipelineDescriptor::new();
-        pipeline_state_descriptor.set_compute_function(Some(&kernel));
-
         let pipeline_state = self
             .device
-            .new_compute_pipeline_state_with_function(
-                pipeline_state_descriptor.compute_function().unwrap(),
-            )
+            .new_compute_pipeline_state_with_function(&kernel)
             .unwrap();
 
         encoder.set_compute_pipeline_state(&pipeline_state);
diff --git a/mopro-msm/src/msm/metal_msm/utils/shader_manager.rs b/mopro-msm/src/msm/metal_msm/utils/shader_manager.rs
