Apple Metal support in tract

.github/workflows/beta-and-nightly.yml

@@ -52,32 +52,37 @@ jobs:
     strategy:
       matrix:
         rust: [ 1.75.0, beta ]
+        os: [ ubuntu-latest, macOS-latest ]
       fail-fast: false
-
-    runs-on: ubuntu-latest
-
     steps:
     - uses: actions/checkout@v3
     - name: Rustup update
       run: rustup update
     - name: LLVM setup
       run: sudo apt-get install -y llvm
     - name: cargo check, -D warnings
+      if: runner.os == 'Linux'
+      env:
+        RUSTFLAGS: "-D warnings"
+        RUST_VERSION: ${{matrix.rust}}
+      run: |
+          rustc --version
+          cargo check --exclude tract-metal --exclude test-metal --workspace
+    - name: cargo check, -D warnings
+      if: runner.os == 'macOS'
       env:
         RUSTFLAGS: "-D warnings"
         RUST_VERSION: ${{matrix.rust}}
       run: |
           rustc --version
-          cargo check
+          cargo check 
 
   cargo-clippy:
     strategy:
-      fail-fast: false
       matrix:
+        os: [ ubuntu-latest, macOS-latest ]
         rust: [ stable, beta, nightly ]
-
-    runs-on: ubuntu-latest
-
+      fail-fast: false
     steps:
     - uses: actions/checkout@v3
     - name: Rustup update, install clippy
@@ -87,7 +92,14 @@ jobs:
       env:
         RUST_VERSION: ${{matrix.rust}}
     - name: Run cargo-clippy
-      run: cargo clippy
+      if: runner.os == 'Linux'
+      run: cargo clippy --exclude tract-metal --exclude test-metal --workspace
+      env:
+        RUSTFLAGS: "-D warnings --force-warn unknown_lints"
+        RUST_VERSION: ${{matrix.rust}}
+    - name: Run cargo-clippy
+      if: runner.os == 'macOS'
+      run: cargo clippy 
       env:
         RUSTFLAGS: "-D warnings --force-warn unknown_lints"
         RUST_VERSION: ${{matrix.rust}}

.github/workflows/full.yml

@@ -137,7 +137,7 @@ jobs:
     - name: Check all targets
       run: |
         ROOT=$(pwd) ./.travis/ci-system-setup.sh
-        cargo check --all-targets --workspace
+        cargo check --all-targets --exclude tract-metal --exclude test-metal --workspace
 
   cli-tests:
     runs-on: ubuntu-latest

.github/workflows/windows.yml

@@ -52,7 +52,7 @@ jobs:
     - name: debug bin
       run: dir "C:\\Program Files\\LLVM\\bin"
     - name: top level cargo check
-      run: cargo check --workspace --exclude test-blas
+      run: cargo check --workspace --exclude test-blas --exclude tract-metal --exclude test-metal
       env:
         LIBCLANG_PATH: "C:\\Program Files\\LLVM\\bin"
     - name: data / linalg / core / nnef / onnx / onnx-opl

.travis/test-published-crates.sh

@@ -31,6 +31,15 @@ do
     cargo -q test $CARGO_EXTRA -q -p tract-$c
 done
 
+if [ `uname` = "Darwin" ]
+then
+    echo
+    echo "$WHITE ### metal ### $NC"
+    echo
+    cargo -q test $CARGO_EXTRA -q -p tract-metal
+fi
+
+
 # doc test are not finding libtensorflow.so
 if ! cargo -q test $CARGO_EXTRA -q -p tract-tensorflow --lib $ALL_FEATURES
 then

.travis/test-rt.sh

@@ -26,8 +26,13 @@ set +x
 cd $ROOT
 for c in test-rt/test*
 do
-    if [ "$c" != "test-rt/test-tflite" ]
+    if [ "$c" = "test-rt/test-tflite" ]
     then
+        echo "$WHITE ### $c ### IGNORED $NC"
+    elif [ "$c" = "test-rt/test-metal" -a  `uname` != "Darwin" ]
+    then
+        echo "$WHITE ### $c ### IGNORED $NC"
+    else
         echo
         echo "$WHITE ### $c ### $NC"
         echo

Cargo.toml

@@ -14,6 +14,7 @@ members = [
     "onnx",
     "libcli",
     "cli",
+    "metal",
     "extra",
 
     "tflite",
@@ -47,10 +48,12 @@ members = [
     "test-rt/suite-onnx",
     "test-rt/test-f16",
     "test-rt/test-blas",
+    "test-rt/test-metal",
     "test-rt/test-unit-core",
     "test-rt/test-onnx-core",
     "test-rt/test-nnef-cycle",
-    "test-rt/test-tflite", 
+    "test-rt/test-tflite",
+    "metal"
 ]
 
 [workspace.dependencies]
@@ -63,6 +66,7 @@ anyhow = "1.0.43"
 approx = "0.5"
 bit-set= "0.5.2"
 blis-src = { version = "0.2", features = ["static", "pthreads"] }
+block = "0.1.6"
 boow = "0.1.3"
 box_drawing = "0.1.2"
 byteorder = "1.4.3"
@@ -80,6 +84,7 @@ dyn-hash = "0.2"
 env_logger = "0.10"
 flatbuffers = "23.1.21"
 flate2 = "1.0.20"
+foreign-types = "0.5"
 fs-err = "2"
 fs2 = "0.4.3"
 getrandom = "0.2"
@@ -93,6 +98,7 @@ liquid-core = "0.26"
 log = "0.4.14"
 maplit = "1.0.2"
 memmap2 = "0.9"
+metal = { version = "0.27.0", features = ["mps"] }
 ndarray = "0.15.3"
 ndarray-npy = { version = "0.8.0", features = [ "compressed_npz" ] }
 nom = "7.0.0"

api/rs/src/lib.rs

@@ -15,6 +15,7 @@ use tract_nnef::prelude::{
 use tract_onnx::prelude::InferenceModelExt;
 use tract_onnx_opl::WithOnnx;
 use tract_pulse::model::{PulsedModel, PulsedModelExt};
+use tract_pulse::internal::PlanOptions;
 use tract_pulse::WithPulse;
 
 use tract_api::*;
@@ -273,6 +274,7 @@ impl ModelInterface for Model {
                 &self.0,
                 &BenchLimits::default(),
                 &mut annotations,
+                &PlanOptions::default(),
                 &inputs,
                 None,
                 true,

cli/Cargo.toml

@@ -47,6 +47,10 @@ tract-onnx = { optional = true, version = "=0.21.7-pre", path = "../onnx" }
 tract-tensorflow = { optional = true, version = "=0.21.7-pre", path = "../tensorflow" }
 tract-tflite = { optional = true, version = "=0.21.7-pre", path = "../tflite" }
 
+[target.'cfg(any(target_os = "macos", target_os = "ios"))'.dependencies]
+tract-metal = { version = "=0.21.7-pre", path = "../metal" }
+
+
 [features]
 default = ["onnx", "tf", "pulse", "pulse-opl", "tflite", "extra"]
 apple-amx-ios = [ "tract-linalg/apple-amx-ios" ]

cli/src/dump.rs

@@ -1,3 +1,4 @@
+use crate::plan_options::plan_options_from_subcommand;
 use crate::params::SomeGraphDef;
 use crate::tensor::run_params_from_subcommand;
 use crate::Parameters;
@@ -113,6 +114,8 @@ pub fn handle(
         tract_libcli::profile::extract_costs(&mut annotations, model, &run_params.symbols)?;
     }
     if options.profile {
+        let run_params = run_params_from_subcommand(params, sub_matches)?;
+        let plan_options = plan_options_from_subcommand(sub_matches)?;
         let model = params
             .tract_model
             .downcast_ref::<TypedModel>()
@@ -122,6 +125,7 @@ pub fn handle(
             model,
             bench_limits,
             &mut annotations,
+            &plan_options,
             &inputs[0],
             None,
             options.folded,

cli/src/main.rs

@@ -29,6 +29,7 @@ mod params;
 mod run;
 #[cfg(feature = "pulse")]
 mod stream_check;
+mod plan_options;
 mod tensor;
 mod utils;
 
@@ -131,6 +132,8 @@ fn main() -> TractResult<()> {
 
         .arg(Arg::new("f32-to-f16").long("f32-to-f16").alias("half-floats").long_help("Convert the decluttered network from f32 to f16"))
         .arg(arg!(--"f16-to-f32" "Convert the decluttered network from f16 to f32"))
+        .arg(arg!(--"metal" "Convert metal compatible operator in the decluttered network. Only available on MacOS and iOS"))
+        .arg(Arg::new("metal-gpu-trace").long("metal-gpu-trace").takes_value(true).help("Capture Metal GPU trace at given path. Only available on MacOS and iOS"))
         .arg(Arg::new("transform").short('t').long("transform").multiple_occurrences(true).takes_value(true).help("Apply a built-in transformation to the model"))
         .arg(Arg::new("set").long("set").multiple_occurrences(true).takes_value(true)
              .long_help("Set a symbol to a concrete value after decluttering"))
@@ -282,7 +285,31 @@ fn main() -> TractResult<()> {
     env_logger::Builder::from_env(env).format_timestamp_nanos().init();
     info_usage("init", probe.as_ref());
 
-    if let Err(e) = handle(matches, probe.as_ref()) {
+    let res = if matches.is_present("metal-gpu-trace") {
+        #[cfg(any(target_os = "macos", target_os = "ios"))]
+        {   
+            let gpu_trace_path = std::path::Path::new(matches.value_of("metal-gpu-trace").unwrap()).to_path_buf();
+            ensure!(gpu_trace_path.is_absolute(), "Metal GPU trace file has to be absolute");
+            ensure!(!gpu_trace_path.exists(), format!("Given Metal GPU trace file {:?} already exists.", gpu_trace_path));
+            log::info!("Capturing Metal GPU trace at : {:?}", gpu_trace_path);
+            std::env::set_var("METAL_CAPTURE_ENABLED", "1");
+            std::env::set_var("METAL_DEVICE_WRAPPER_TYPE", "1");
+            let probe_ref = probe.as_ref();
+            tract_metal::METAL_CONTEXT.with_borrow(move |context| {
+            context.capture_trace(gpu_trace_path, move |_ctxt| {
+                    handle(matches, probe_ref)
+                })
+            })
+        }
+        #[cfg(not(any(target_os = "macos", target_os = "ios")))]
+        {
+            bail!("`--metal-gpu-trace` present but it is only available on MacOS and iOS")
+        }
+    } else {
+        handle(matches, probe.as_ref())
+    };
+
+    if let Err(e) = res {
         error!("{:?}", e);
         std::process::exit(1);
     }
@@ -445,6 +472,10 @@ fn run_options(command: clap::Command) -> clap::Command {
                 "Path to a directory containing input tensors in NNEF format (.dat files). This sets tensor values.",
                 ),
                 )
+        .arg(Arg::new("skip-order-opt-ram")
+            .long("skip-order-opt-ram")
+            .help("Plan node evaluation order without RAM optimisation"),
+            )
         .arg(
             Arg::new("allow-random-input")
             .short('R')

cli/src/params.rs

@@ -10,7 +10,7 @@ use tract_core::ops::konst::Const;
 use tract_itertools::Itertools;
 use tract_libcli::profile::BenchLimits;
 use tract_nnef::tensors::read_tensor;
-
+use tract_core::transform::ModelTransform;
 use tract_core::internal::*;
 use tract_core::model::TypedModel;
 use tract_hir::internal::*;
@@ -741,6 +741,22 @@ impl Parameters {
                 tract_core::floats::FloatPrecisionTranslator::<f16, f32>::default().translate_model(&m)
             });
         }
+        {
+            if matches.is_present("metal") {
+                #[cfg(any(target_os = "macos", target_os = "ios"))]
+                {
+                    stage!("metal", typed_model -> typed_model, |m:TypedModel| {
+                        tract_metal::transform::MetalTransform
+                            .transform_into(&m)
+                    });
+                }
+                #[cfg(not(any(target_os = "macos", target_os = "ios")))]
+                {
+                    bail!("`--metal` present but it is only available on MacOS and iOS")
+                }
+            }
+        }
+
         if let Some(transform) = matches.values_of("transform") {
             for transform in transform {
                 stage!(transform, typed_model -> typed_model, |m:TypedModel| {

cli/src/plan_options.rs

@@ -0,0 +1,12 @@
+use tract_core::internal::*;
+
+
+pub fn plan_options_from_subcommand(
+    sub_matches: &clap::ArgMatches,
+) -> TractResult<PlanOptions> {
+    let skip_order_opt_ram: bool = sub_matches.is_present("skip-order-opt-ram");
+
+    let mut options = PlanOptions::default();
+    options.skip_order_opt_ram = skip_order_opt_ram;
+    Ok(options)
+}

cli/src/run.rs

@@ -142,9 +142,10 @@ pub fn handle(
 fn run_regular(
     tract: &dyn Model,
     run_params: &RunParams,
-    _matches: &clap::ArgMatches,
+    matches: &clap::ArgMatches,
     sub_matches: &clap::ArgMatches,
 ) -> TractResult<TVec<Vec<TValue>>> {
+    let plan_options = crate::plan_options::plan_options_from_subcommand(matches)?;
     let steps = sub_matches.is_present("steps");
     let check_f16_overflow = sub_matches.is_present("check-f16-overflow");
     let assert_sane_floats = sub_matches.is_present("assert-sane-floats");
@@ -155,7 +156,7 @@ fn run_regular(
         None
     };
     dispatch_model!(tract, |m| {
-        let plan = SimplePlan::new(m)?;
+        let plan = SimplePlan::new_with_options(m, &plan_options)?;
         let mut state = SimpleState::new(plan)?;
         let inputs = tract_libcli::tensor::retrieve_or_make_inputs(tract, run_params)?;
         let mut results = tvec!(vec!(); state.model().outputs.len());

core/src/model/fact.rs

@@ -294,14 +294,17 @@ impl TypedFact {
     }
 
     pub fn without_value(&self) -> Self {
-        Self::dt_shape(self.datum_type, self.shape.clone())
+        TypedFact {
+            datum_type: self.datum_type,
+            shape: self.shape.clone(),
+            konst: None,
+            uniform: None,
+            opaque_fact: self.opaque_fact.clone(),
+        }
     }
 
-    pub fn with_opaque_metadata<O: Into<Box<dyn OpaqueFact>>>(
-        mut self,
-        opaque_metadata: O,
-    ) -> Self {
-        self.opaque_fact = Some(opaque_metadata.into());
+    pub fn with_opaque_fact<O: Into<Box<dyn OpaqueFact>>>(mut self, opaque_fact: O) -> Self {
+        self.opaque_fact = Some(opaque_fact.into());
         self
     }
 }
@@ -395,16 +398,18 @@ impl<'a> From<&'a Arc<Tensor>> for TypedFact {
 
 impl fmt::Debug for TypedFact {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
-        if let Some(k) = &self.konst {
-            write!(fmt, "{k:?}")?
-        } else if self.rank() > 0 {
-            write!(fmt, "{:?},{:?}", self.shape, self.datum_type)?
-        } else {
-            write!(fmt, "{:?}", self.datum_type)?
-        };
-        if let Some(of) = &self.opaque_fact {
-            write!(fmt, " {:?}", of)?;
-        }
+        match (self.konst.as_ref(), self.opaque_fact.as_ref()) {
+            (Some(ref k), None) => write!(fmt, "{k:?}"),
+            (Some(ref k), Some(meta)) => write!(fmt, "{meta:?} {k:?}"),
+            (None, None) if self.rank() > 0 => {
+                write!(fmt, "{:?},{:?}", self.shape, self.datum_type)
+            }
+            (None, Some(ref meta)) if self.rank() > 0 => {
+                write!(fmt, "{:?},{:?},{:?}", self.shape, self.datum_type, meta)
+            }
+            (None, Some(ref meta)) => write!(fmt, "{:?}, {:?}", self.datum_type, meta),
+            (None, None) => write!(fmt, "{:?}", self.datum_type),
+        }?;
         Ok(())
     }
 }