Use OpenBLAS/BLAS/IntelMKL for matrix multiplication about dfdx

[coreylowman]

[dependencies]
blas = "0.22.0"
blas-src = { version = "0.8", features = ["openblas"] }
openblas-src = { version = "0.10", features = ["cblas", "system"] }

[coreylowman]

First attempt... the arguments aren't correct though.

extern crate blas_src;

fn matmul<const M: usize, const N: usize, const K: usize>(
    a: &[[f32; K]; M],
    b: &[[f32; N]; K],
) -> [[f32; N]; M] {
    let mut c = [[0.0; N]; M];
    unsafe {
        let a = std::slice::from_raw_parts(a.as_ptr() as *const f32, 0);
        let b = std::slice::from_raw_parts(b.as_ptr() as *const f32, 0);
        let c = std::slice::from_raw_parts_mut(c.as_mut_ptr() as *mut f32, 0);
        blas::sgemm(
            b'T', b'T', M as i32, N as i32, K as i32, 1.0, a, K as i32, b, N as i32, 0.0, c,
            M as i32,
        )
    }
    c
}

[coreylowman]

Here's a cblas_sys function to call

[dependencies]
cblas-sys = "0.1.4"
libc = "0.2.126"

pub fn mkl_mm<const M: usize, const N: usize, const K: usize>(
    a: &[[f32; K]; M],
    b: &[[f32; N]; K],
    c: &mut [[f32; N]; M],
) {
    unsafe {
        cblas_sys::cblas_sgemm(
            cblas_sys::CblasRowMajor,
            cblas_sys::CblasNoTrans,
            cblas_sys::CblasNoTrans,
            M as libc::c_int,
            N as libc::c_int,
            K as libc::c_int,
            1.0 as libc::c_float,
            a.as_ptr() as *const libc::c_float,
            K as libc::c_int,
            b.as_ptr() as *const libc::c_float,
            N as libc::c_int,
            0.0 as libc::c_float,
            c.as_mut_ptr() as *mut libc::c_float,
            N as libc::c_int,
        )
    }
}

[coreylowman]

And custom build.rs file for linking against MKL is here:

build.zip

[coreylowman]

Note that pytorch & tensorflow both use intel MKL under the hood as far as I can tell (at leats on my machine):

>>> print(*torch.__config__.show().split("\n"), sep="\n")
PyTorch built with:
  - C++ Version: 199711
  - MSVC 192829337
  - Intel(R) Math Kernel Library Version 2020.0.2 Product Build 20200624 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v2.5.2 (Git Hash a9302535553c73243c632ad3c4c80beec3d19a1e)
  - OpenMP 2019
  - LAPACK is enabled (usually provided by MKL)
  - CPU capability usage: AVX512
  - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CXX_COMPILER=C:/actions-runner/_work/pytorch/pytorch/builder/windows/tmp_bin/sccache-cl.exe, CXX_FLAGS=/DWIN32 /D_WINDOWS /GR /EHsc /w /bigobj -DUSE_PTHREADPOOL -openmp:experimental -IC:/actions-runner/_work/pytorch/pytorch/builder/windows/mkl/include -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOCUPTI -DUSE_FBGEMM -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -DEDGE_PROFILER_USE_KINETO, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_VERSION=1.11.0, USE_CUDA=0, USE_CUDNN=OFF, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=OFF, USE_MPI=OFF, USE_NCCL=OFF, USE_NNPACK=OFF, USE_OPENMP=ON, USE_ROCM=OFF,

[jianshu93]

Any possibility to also add support for openblas? I understand that intel-mkl is more widely used for tensorflow and PyTorch by default. Openblas seem to work on any CPU structure despite sightly slower than intel-mkl, for example ARM64 CPUs like KUNPENG920-6426(HiSilicon) , ARM64 apple M1 series and also Fujitsu-A64FX. The ndarray-linalg crate does support both by using different features. Can matrix-mulitply be replaced with this one?

Thanks,

Jianshu

[coreylowman]

Good point about support for different cpu structures. I'm definitely open to this. When I first was looking into this stuff configuration for it seemed very confusing (especially with cargo's additive features). I also lack a linux machine to test on, so I'd need help with that.

This will probably be lower on the priorities list for me, but happy to review contributions for this