Add accelerate framework support for burn-ndarray (#183) (#184)

tracel-ai · Feb 28, 2023 · c510e8b · c510e8b
1 parent 9df3b6a
commit c510e8b
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 0 deletions.
diff --git a/burn-ndarray/Cargo.toml b/burn-ndarray/Cargo.toml
@@ -23,6 +23,7 @@ std = [
   "ndarray/matrixmultiply-threading",
 ]
 
+blas-accelerate = ["ndarray/blas", "blas-src/accelerate"] # Accelerate framework (macOS only)
 blas-netlib = ["ndarray/blas", "blas-src/netlib"]
 blas-openblas = ["ndarray/blas", "blas-src/openblas", "openblas-src"]
 blas-openblas-system = [

diff --git a/burn-ndarray/build.rs b/burn-ndarray/build.rs
@@ -0,0 +1,6 @@
+fn main() {
+    // https://github.com/rust-ndarray/ndarray/issues/1197
+    if cfg!(feature = "blas-accelerate") {
+        println!("cargo:rustc-link-lib=framework=Accelerate");
+    }
+}
diff --git a/examples/mnist/Cargo.toml b/examples/mnist/Cargo.toml
@@ -11,6 +11,7 @@ default = []
 tch-cpu = ["dep:burn-tch"]
 tch-gpu = ["dep:burn-tch"]
 ndarray = ["burn-ndarray"]
+ndarray-blas-accelerate = ["burn-ndarray/blas-accelerate"]
 ndarray-blas-netlib = ["burn-ndarray/blas-netlib"]
 ndarray-blas-openblas = ["burn-ndarray/blas-openblas"]
 

diff --git a/examples/mnist/examples/mnist.rs b/examples/mnist/examples/mnist.rs
@@ -2,6 +2,7 @@
     feature = "ndarray",
     feature = "ndarray-blas-netlib",
     feature = "ndarray-blas-openblas",
+    feature = "ndarray-blas-accelerate",
 ))]
 mod ndarray {
     use burn_autodiff::ADBackendDecorator;
@@ -46,6 +47,7 @@ fn main() {
         feature = "ndarray",
         feature = "ndarray-blas-netlib",
         feature = "ndarray-blas-openblas",
+        feature = "ndarray-blas-accelerate",
     ))]
     ndarray::run();
     #[cfg(feature = "tch-gpu")]