conditional blas selection

sonos · Feb 27, 2024 · b9c9e12 · b9c9e12
1 parent 9be466a
commit b9c9e12
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 5 deletions.
diff --git a/core/Cargo.toml b/core/Cargo.toml
@@ -17,9 +17,9 @@ maintenance = { status = "actively-developed" }
 [dependencies]
 anyhow.workspace = true
 bit-set.workspace = true
-cblas = { version = "0.4" }
-# blis-src = { version = "0.2", features = ["static", "pthreads"] }
-openblas-src = { version = "0.10", features = ["static"] }
+cblas = { version = "0.4", optional = true }
+blis-src = { version = "0.2", features = ["static", "pthreads"], optional = true }
+openblas-src = { version = "0.10", features = ["static"], optional = true }
 derive-new.workspace = true
 downcast-rs.workspace = true
 dyn-clone.workspace = true
@@ -39,6 +39,9 @@ tract-data = { version = "=0.21.2-pre", path = "../data" }
 [features]
 default = [ ]
 complex = [ "tract-data/complex", "tract-linalg/complex" ]
+blas = [ "cblas" ]
+blis = [ "blas", "blis-src" ]
+openblas = [ "blas", "openblas-src" ]
 paranoid_assertions = []
 
 [dev-dependencies]

diff --git a/core/src/lib.rs b/core/src/lib.rs
@@ -45,9 +45,12 @@
 //! tract-tensorflow or tract-onnx crates.
 //!
 
+#[cfg(feature="blas")]
 extern crate cblas;
-// extern crate blis_src;
-extern crate openblas_src;
+#[cfg(feature="blis")]
+extern crate blis_src;
+#[cfg(feature="openblas")]
+ extern crate openblas_src;
 
 pub extern crate anyhow;
 extern crate bit_set;

diff --git a/core/src/ops/einsum/as_blas.rs b/core/src/ops/einsum/as_blas.rs
@@ -163,5 +163,11 @@ impl TypedOp for SGemm {
         Ok(tvec!(f32::fact(&self.output_shape(&*inputs[0].shape, &*inputs[1].shape)?)))
     }
 
+    fn cost(&self, inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
+        let fma = self.output_shape(&*inputs[0].shape, &*inputs[1].shape)?.iter().product::<TDim>()
+            * inputs[0].shape.last().unwrap();
+        Ok(tvec!((Cost::FMA(f32::datum_type()), fma)))
+    }
+
     as_op!();
 }