Skip to content

Commit

Permalink
feature gate blas sgemm stuff
Browse files Browse the repository at this point in the history
  • Loading branch information
mathieupoumeyrolsonos committed Feb 27, 2024
1 parent b9c9e12 commit 6462794
Show file tree
Hide file tree
Showing 4 changed files with 8 additions and 1 deletion.
2 changes: 2 additions & 0 deletions core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ maintenance = { status = "actively-developed" }
[dependencies]
anyhow.workspace = true
bit-set.workspace = true
accelerate-src = { version = "0.3", optional = true }
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 }
Expand All @@ -40,6 +41,7 @@ tract-data = { version = "=0.21.2-pre", path = "../data" }
default = [ ]
complex = [ "tract-data/complex", "tract-linalg/complex" ]
blas = [ "cblas" ]
accelerate = [ "blas", "accelerate-src" ]
blis = [ "blas", "blis-src" ]
openblas = [ "blas", "openblas-src" ]
paranoid_assertions = []
Expand Down
4 changes: 3 additions & 1 deletion core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,10 +47,12 @@
#[cfg(feature="blas")]
extern crate cblas;
#[cfg(feature="accelerate")]
extern crate accelerate_src;
#[cfg(feature="blis")]
extern crate blis_src;
#[cfg(feature="openblas")]
extern crate openblas_src;
extern crate openblas_src;

pub extern crate anyhow;
extern crate bit_set;
Expand Down
1 change: 1 addition & 0 deletions core/src/ops/einsum/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ use crate::tract_data::itertools::Itertools;

mod eval;

#[cfg(feature="blas")]
pub mod as_blas;
use super::array::TypedConcat;
use super::math::add;
Expand Down
2 changes: 2 additions & 0 deletions core/src/transform.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use crate::internal::*;
#[cfg(feature="blas")]
use crate::ops::einsum::as_blas::AsBlas;
use std::borrow::Cow;
use std::fmt::Debug;
Expand All @@ -10,6 +11,7 @@ use crate::ops::nn::{Softmax, SoftmaxExp, TypedModel};

pub fn get_transformer(name: &str) -> Option<Box<dyn ModelTransformer>> {
match name {
#[cfg(feature="blas")]
"as-blas" => Some(Box::<AsBlas>::default()),
"f32-to-f16" => Some(Box::<FloatPrecisionTranslator<f32, f16>>::default()),
"f16-to-f32" => Some(Box::<FloatPrecisionTranslator<f32, f16>>::default()),
Expand Down

0 comments on commit 6462794

Please sign in to comment.