andrey/gemm
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix fmt and some warnings #1

Merged
andrey merged 1 commits from fix-fmt into master 2022-03-05 18:49:40 +04:00
Conversation 0 Commits 1 Files Changed 20 +339 -287
20 changed files with 339 additions and 287 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
.drone.yml Normal file
View File

@ -0,0 +1,14 @@
kind: pipeline
name: default
steps:
- name: build
image: rust
commands:
- cargo build --verbose --all
- name: fmt-check
image: rust
commands:
- rustup component add rustfmt
- cargo fmt --all -- --check

2
build.rs
View File

@ -4,4 +4,4 @@ fn main() {
// and run this.
println!("cargo:rerun-if-changed=build.rs");
println!("cargo:rustc-link-lib={}=openblas", "dylib");
}
}

40
src/dim.rs
View File

@ -2,49 +2,49 @@ pub trait Dim {
const DIM: usize;
}
pub struct A1;
impl Dim for A1 {
pub struct A1;
impl Dim for A1 {
const DIM: usize = 1;
}
pub struct A2;
impl Dim for A2 {
pub struct A2;
impl Dim for A2 {
const DIM: usize = 2;
}
pub struct A3;
impl Dim for A3 {
pub struct A3;
impl Dim for A3 {
const DIM: usize = 3;
}
pub struct A4;
impl Dim for A4 {
pub struct A4;
impl Dim for A4 {
const DIM: usize = 4;
}
pub struct A5;
impl Dim for A5 {
pub struct A5;
impl Dim for A5 {
const DIM: usize = 5;
}
pub struct A6;
impl Dim for A6 {
pub struct A6;
impl Dim for A6 {
const DIM: usize = 6;
}
pub struct A7;
impl Dim for A7 {
pub struct A7;
impl Dim for A7 {
const DIM: usize = 7;
}
pub struct A8;
impl Dim for A8 {
pub struct A8;
impl Dim for A8 {
const DIM: usize = 8;
}
pub struct A16;
impl Dim for A16 {
pub struct A16;
impl Dim for A16 {
const DIM: usize = 16;
}
/*
macro_rules! gen {
(($name: ident, $val: lit)) => {
};
}
@ -72,4 +72,4 @@ gen! {
(A28, 28),
(A30, 30),
(A32, 32),
}*/
}*/

7
src/executor.rs
View File

@ -42,19 +42,18 @@ impl Executor for RayonExecutor {
(start..end)
.into_par_iter()
.map(|x| x * step)
.for_each(|x|f(x));
.for_each(|x| f(x));
}
}
pub struct ThreadPoolExecutor {
thread_pool: threadpool::ThreadPool
thread_pool: threadpool::ThreadPool,
}
impl ThreadPoolExecutor {
pub fn new() -> ThreadPoolExecutor {
ThreadPoolExecutor {
thread_pool: threadpool::Builder::new().build()
thread_pool: threadpool::Builder::new().build(),
}
}
}

186
src/gemm.rs
View File

@ -1,10 +1,12 @@
use crate::aligned_alloc;
use crate::matrix::{Number, Matrix, MatrixMut, MutMatrix, ConstMatrix, ConstTransposedMatrix, MutTransposedMatrix};
use crate::kernel::params::single::{NC, MC, KC};
use crate::kernel;
use crate::kernel::GemmKernel;
use crate::dim::Dim;
use crate::executor::Executor;
use crate::kernel;
use crate::kernel::params::single::{KC, MC, NC};
use crate::kernel::GemmKernel;
use crate::matrix::{
ConstMatrix, ConstTransposedMatrix, Matrix, MatrixMut, MutMatrix, MutTransposedMatrix, Number,
};
pub unsafe fn gemm<E, F, K, MR, NR>(
e: &E,
@ -22,70 +24,109 @@ pub unsafe fn gemm<E, F, K, MR, NR>(
beta: F,
c: *mut F,
ldc: usize,
)
where E: Executor,
F: Number,
MR: Dim, NR: Dim,
K: GemmKernel<F, MR, NR>,
) where
E: Executor,
F: Number,
MR: Dim,
NR: Dim,
K: GemmKernel<F, MR, NR>,
{
match (transa, transb, transc) {
(false, false, false) => gemm_template::<E, F, K, MR, NR, _, _, _>(
e, m, n, k, alpha,
e,
m,
n,
k,
alpha,
ConstMatrix::new(a, lda),
ConstMatrix::new(b, ldb),
beta,
MutMatrix::new(c, ldc)),
MutMatrix::new(c, ldc),
),
(false, false, true) => gemm_template::<E, F, K, MR, NR, _, _, _>(
e, m, n, k, alpha,
e,
m,
n,
k,
alpha,
ConstMatrix::new(a, lda),
ConstMatrix::new(b, ldb),
beta,
MutTransposedMatrix::new(c, ldc)),
MutTransposedMatrix::new(c, ldc),
),
(false, true, false) => gemm_template::<E, F, K, MR, NR, _, _, _>(
e, m, n, k, alpha,
e,
m,
n,
k,
alpha,
ConstMatrix::new(a, lda),
ConstTransposedMatrix::new(b, ldb),
beta,
MutMatrix::new(c, ldc)),
MutMatrix::new(c, ldc),
),
(false, true, true) => gemm_template::<E, F, K, MR, NR, _, _, _>(
e, m, n, k, alpha,
e,
m,
n,
k,
alpha,
ConstMatrix::new(a, lda),
ConstTransposedMatrix::new(b, ldb),
beta,
MutTransposedMatrix::new(c, ldc)),
MutTransposedMatrix::new(c, ldc),
),
(true, false, false) => gemm_template::<E, F, K, MR, NR, _, _, _>(
e, m, n, k, alpha,
e,
m,
n,
k,
alpha,
ConstTransposedMatrix::new(a, lda),
ConstMatrix::new(b, ldb),
beta,
MutMatrix::new(c, ldc)),
MutMatrix::new(c, ldc),
),
(true, false, true) => gemm_template::<E, F, K, MR, NR, _, _, _>(
e, m, n, k, alpha,
e,
m,
n,
k,
alpha,
ConstTransposedMatrix::new(a, lda),
ConstMatrix::new(b, ldb),
beta,
MutTransposedMatrix::new(c, ldc)),
MutTransposedMatrix::new(c, ldc),
),
(true, true, false) => gemm_template::<E, F, K, MR, NR, _, _, _>(
e, m, n, k, alpha,
e,
m,
n,
k,
alpha,
ConstTransposedMatrix::new(a, lda),
ConstTransposedMatrix::new(b, ldb),
beta,
MutMatrix::new(c, ldc)),
MutMatrix::new(c, ldc),
),
(true, true, true) => gemm_template::<E, F, K, MR, NR, _, _, _>(
e, m, n, k, alpha,
e,
m,
n,
k,
alpha,
ConstTransposedMatrix::new(a, lda),
ConstTransposedMatrix::new(b, ldb),
beta,
MutTransposedMatrix::new(c, ldc)),
MutTransposedMatrix::new(c, ldc),
),
}
}
@ -98,15 +139,16 @@ unsafe fn gemm_template<E, F, K, MR, NR, A, B, C>(
a: A,
b: B,
beta: F,
c: C
)
where E: Executor,
F: Number,
MR: Dim, NR: Dim,
K: GemmKernel<F, MR, NR>,
A: Matrix<F>,
B: Matrix<F>,
C: MatrixMut<F>,
c: C,
) where
E: Executor,
F: Number,
MR: Dim,
NR: Dim,
K: GemmKernel<F, MR, NR>,
A: Matrix<F>,
B: Matrix<F>,
C: MatrixMut<F>,
{
let packed_a = aligned_alloc::Alloc::new(MC * KC * std::mem::size_of::<F>());
let packed_b = aligned_alloc::Alloc::new(KC * NC * std::mem::size_of::<F>());
@ -122,7 +164,7 @@ where E: Executor,
let pb = MutMatrix::new(packed_b.ptr::<F>(), p_b);
inner_kernel::<E, F, K, MR, NR, _, _, _>(
e,
e,
i_b,
j_b,
p_b,
@ -133,7 +175,7 @@ where E: Executor,
c.sub(j, i),
pa,
pb,
i == 0
i == 0,
);
}
}
@ -141,7 +183,7 @@ where E: Executor,
}
//
// | MR |
// | MR |
// +-----------------------------+----+
// | : | |
// | TL : | TR |
@ -172,16 +214,15 @@ unsafe fn inner_kernel<E, F, K, MR, NR, A, B, C>(
pa: MutMatrix<F>,
pb: MutMatrix<F>,
first_time: bool,
)
where E: Executor,
F: Number,
MR: Dim,
NR: Dim,
K: kernel::GemmKernel<F, MR, NR>,
A: Matrix<F>,
B: Matrix<F>,
C: MatrixMut<F>,
) where
E: Executor,
F: Number,
MR: Dim,
NR: Dim,
K: kernel::GemmKernel<F, MR, NR>,
A: Matrix<F>,
B: Matrix<F>,
C: MatrixMut<F>,
{
let n_left = n % NR::DIM;
let n_main = n - n_left;
@ -190,33 +231,24 @@ unsafe fn inner_kernel<E, F, K, MR, NR, A, B, C>(
let m_main = m - m_left;
if first_time {
e.execute(0, n_main, NR::DIM, move |j|
K::pack_row_b(b, pb));
e.execute(0, n_main, NR::DIM, move |_j| K::pack_row_b(b, pb));
}
e.execute(0, m_main, MR::DIM, move |i|
K::pack_row_a(a.sub_col(i), pa.sub_row(i)));
e.execute(0, m_main, MR::DIM, move |i| {
K::pack_row_a(a.sub_col(i), pa.sub_row(i))
});
e.synchronize();
e.execute(0, n_main, NR::DIM, move |j| {
// Section TL
for i in (0..m_main).step_by(MR::DIM) {
K::main_tl(alpha,
pa.sub_row(i),
pb.sub_row(j),
beta,
c.sub(j, i));
K::main_tl(alpha, pa.sub_row(i), pb.sub_row(j), beta, c.sub(j, i));
}
// Section TR
for i in m_main..m {
K::sup_tr(
alpha,
a.sub_col(i),
pb.sub_row(j),
beta,
c.sub(j, i));
K::sup_tr(alpha, a.sub_col(i), pb.sub_row(j), beta, c.sub(j, i));
}
});
@ -225,27 +257,15 @@ unsafe fn inner_kernel<E, F, K, MR, NR, A, B, C>(
let i = ji % m_main;
// Section BL
K::sup_bl(
alpha,
pa.sub_row(i),
b.sub_row(j),
beta,
c.sub(j, i)
);
K::sup_bl(alpha, pa.sub_row(i), b.sub_row(j), beta, c.sub(j, i));
});
// Section BR
for j in n_main..n {
for i in m_main..m {
K::sup_br(
k,
alpha,
a.sub_col(i),
b.sub_row(j),
beta,
c.sub(j, i))
K::sup_br(k, alpha, a.sub_col(i), b.sub_row(j), beta, c.sub(j, i))
}
};
}
e.synchronize();
}
}

2
src/kernel/avx/l1s.rs
View File

@ -246,7 +246,7 @@ pub unsafe fn sdot(
let mut acc1 = _mm256_setzero_ps();
let mut acc2 = _mm256_setzero_ps();
let mut acc3 = _mm256_setzero_ps();
for _ in 0..n / 32 {
let x0 = _mm256_loadu_ps(x);
let y0 = _mm256_loadu_ps(y);

25
src/kernel/avx/l3s.rs
View File

@ -1,16 +1,16 @@
use super::fma::fmadd_ps;
use super::intrinsics::*;
use crate::kernel::params::single::{MR, NR};
use crate::matrix::{Matrix, MutMatrix, MatrixMut};
use crate::matrix::{Matrix, MatrixMut, MutMatrix};
use crunchy::unroll;
#[inline]
pub(crate) unsafe fn sgemm_ukr_16x8<C: MatrixMut<f32>>(
k: usize,
alpha: f32,
_alpha: f32,
pa: MutMatrix<f32>,
pb: MutMatrix<f32>,
beta: f32,
_beta: f32,
c: C,
) {
let mut mt00 = _mm256_setzero_ps();
@ -21,7 +21,7 @@ pub(crate) unsafe fn sgemm_ukr_16x8<C: MatrixMut<f32>>(
// let mut mt05 = _mm256_setzero_ps();
// let mut mt06 = _mm256_setzero_ps();
// let mut mt07 = _mm256_setzero_ps();
let mut mt10 = _mm256_setzero_ps();
let mut mt11 = _mm256_setzero_ps();
let mut mt12 = _mm256_setzero_ps();
@ -30,7 +30,7 @@ pub(crate) unsafe fn sgemm_ukr_16x8<C: MatrixMut<f32>>(
// let mut mt15 = _mm256_setzero_ps();
// let mut mt16 = _mm256_setzero_ps();
// let mut mt17 = _mm256_setzero_ps();
let mut pa = pa.ptr();
let mut pb = pb.ptr();
@ -53,7 +53,7 @@ pub(crate) unsafe fn sgemm_ukr_16x8<C: MatrixMut<f32>>(
// let b5 = _mm256_broadcast_ss(&*pb.add(5));
// let b6 = _mm256_broadcast_ss(&*pb.add(6));
// let b7 = _mm256_broadcast_ss(&*pb.add(7));
mt00 = fmadd_ps(a0, b0, mt00);
mt01 = fmadd_ps(a0, b1, mt01);
mt02 = fmadd_ps(a0, b2, mt02);
@ -75,7 +75,6 @@ pub(crate) unsafe fn sgemm_ukr_16x8<C: MatrixMut<f32>>(
pb = pb.add(BATCH * NR);
}
for _ in k_main..k {
let a0 = _mm256_load_ps(pa);
let a1 = _mm256_load_ps(pa.add(8));
@ -88,7 +87,7 @@ pub(crate) unsafe fn sgemm_ukr_16x8<C: MatrixMut<f32>>(
// let b5 = _mm256_broadcast_ss(&*pb.add(5));
// let b6 = _mm256_broadcast_ss(&*pb.add(6));
// let b7 = _mm256_broadcast_ss(&*pb.add(7));
mt00 = fmadd_ps(a0, b0, mt00);
mt01 = fmadd_ps(a0, b1, mt01);
mt02 = fmadd_ps(a0, b2, mt02);
@ -160,7 +159,7 @@ pub(crate) unsafe fn sgemm_ukr_16x8<C: MatrixMut<f32>>(
// // mt05 = fmadd_ps(beta, _mm256_loadu_ps(ccol5), mt05);
// // mt06 = fmadd_ps(beta, _mm256_loadu_ps(ccol6), mt06);
// // mt07 = fmadd_ps(beta, _mm256_loadu_ps(ccol7), mt07);
// mt10 = fmadd_ps(beta, _mm256_loadu_ps(ccol0.add(8)), mt10);
// mt11 = fmadd_ps(beta, _mm256_loadu_ps(ccol1.add(8)), mt11);
// mt12 = fmadd_ps(beta, _mm256_loadu_ps(ccol2.add(8)), mt12);
@ -179,7 +178,7 @@ pub(crate) unsafe fn sgemm_ukr_16x8<C: MatrixMut<f32>>(
// _mm256_storeu_ps(ccol5, mt05);
// _mm256_storeu_ps(ccol6, mt06);
// _mm256_storeu_ps(ccol7, mt07);
_mm256_storeu_ps(ccol0.add(8), mt10);
_mm256_storeu_ps(ccol1.add(8), mt11);
_mm256_storeu_ps(ccol2.add(8), mt12);
@ -192,10 +191,10 @@ pub(crate) unsafe fn sgemm_ukr_16x8<C: MatrixMut<f32>>(
pub(crate) unsafe fn sgemm_sup_16x1<B: Matrix<f32>, C: MatrixMut<f32>>(
k: usize,
alpha: f32,
_alpha: f32,
pa: MutMatrix<f32>,
b: B,
beta: f32,
_beta: f32,
c: C,
) {
let mut mt0 = _mm256_setzero_ps();
@ -250,4 +249,4 @@ pub(crate) unsafe fn sgemm_pa_16x(k: usize, a: *const f32, lda: usize, pa: *mut
pa = pa.add(16);
a = a.add(lda);
}
}
}

32
src/kernel/avx/mod.rs
View File

@ -6,15 +6,16 @@ pub mod l1s;
pub mod l3d;
pub mod l3s;
use core::marker::PhantomData;
use crate::matrix::{Number, MutMatrix, Matrix, MatrixMut};
use crate::kernel::{GemmKernel, GemmKernelSupNr, GemmKernelSupMr, GemmKernelSup};
use crate::dim::*;
use crate::kernel::{GemmKernel, GemmKernelSup, GemmKernelSupMr, GemmKernelSupNr};
use crate::matrix::{Matrix, MatrixMut, MutMatrix, Number};
use core::marker::PhantomData;
pub struct AvxKernel<F: Number, I>(PhantomData<fn(F, I)>);
impl<I> GemmKernelSupNr<f32, A5> for AvxKernel<f32, I>
where I: GemmKernelSupNr<f32, A5>
impl<I> GemmKernelSupNr<f32, A5> for AvxKernel<f32, I>
where
I: GemmKernelSupNr<f32, A5>,
{
#[inline]
unsafe fn sup_tr<A: Matrix<f32>, C: MatrixMut<f32>>(
@ -26,10 +27,11 @@ impl<I> GemmKernelSupNr<f32, A5> for AvxKernel<f32, I>
) {
I::sup_tr(alpha, a, pb, beta, c);
}
}
}
impl<I> GemmKernelSupMr<f32, A16> for AvxKernel<f32, I>
where I: GemmKernelSupMr<f32, A16>
impl<I> GemmKernelSupMr<f32, A16> for AvxKernel<f32, I>
where
I: GemmKernelSupMr<f32, A16>,
{
#[inline]
unsafe fn sup_bl<B: Matrix<f32>, C: MatrixMut<f32>>(
@ -43,8 +45,9 @@ impl<I> GemmKernelSupMr<f32, A16> for AvxKernel<f32, I>
}
}
impl<I> GemmKernelSup<f32> for AvxKernel<f32, I>
where I: GemmKernelSup<f32>
impl<I> GemmKernelSup<f32> for AvxKernel<f32, I>
where
I: GemmKernelSup<f32>,
{
#[inline]
unsafe fn sup_br<A: Matrix<f32>, B: Matrix<f32>, C: MatrixMut<f32>>(
@ -59,12 +62,13 @@ impl<I> GemmKernelSup<f32> for AvxKernel<f32, I>
}
}
impl<I> GemmKernel<f32, A16, A5> for AvxKernel<f32, I>
where I: GemmKernel<f32, A16, A5>
impl<I> GemmKernel<f32, A16, A5> for AvxKernel<f32, I>
where
I: GemmKernel<f32, A16, A5>,
{
#[inline]
unsafe fn pack_row_a<A: Matrix<f32>>(a: A, pa: MutMatrix<f32>) {
if a.is_transposed() {
if a.is_transposed() {
I::pack_row_a(a, pa);
} else {
self::l3s::sgemm_pa_16x(pa.stride, a.ptr(), a.stride(), pa.ptr_mut());
@ -86,4 +90,4 @@ impl<I> GemmKernel<f32, A16, A5> for AvxKernel<f32, I>
) {
self::l3s::sgemm_ukr_16x8(pa.stride, alpha, pa, pb, beta, c);
}
}
}

11
src/kernel/generic/l3s.rs
View File

@ -1,13 +1,12 @@
use crate::kernel::params::single::NR;
use crate::matrix::{MutMatrix, Matrix, MatrixMut};
use crate::matrix::{Matrix, MatrixMut, MutMatrix};
pub(crate) unsafe fn sgemm_sup_1x8<A: Matrix<f32>, C: MatrixMut<f32>>(
k: usize,
alpha: f32,
_alpha: f32,
a: A,
pb: MutMatrix<f32>,
beta: f32,
_beta: f32,
c: C,
) {
let mut c0 = 0.0f32;
@ -46,7 +45,7 @@ pub(crate) unsafe fn sgemm_sup_1x8<A: Matrix<f32>, C: MatrixMut<f32>>(
// // c5 *= alpha;
// // c6 *= alpha;
// // c7 *= alpha;
let ccol0 = c.ptr_mut();
let ccol1 = c.row_mut(1);
let ccol2 = c.row_mut(2);
@ -139,4 +138,4 @@ pub(crate) unsafe fn sgemm_pa_t(k: usize, a: *const f32, lda: usize, pa: *mut f3
a = a.add(1);
pa = pa.add(MR);
}
}
}

41
src/kernel/generic/mod.rs
View File

@ -1,9 +1,9 @@
pub mod l3d;
pub mod l3s;
use crate::matrix::{Matrix, MatrixMut, MutMatrix};
use crate::kernel::{GemmKernel, GemmKernelSupNr, GemmKernelSupMr, GemmKernelSup};
use crate::dim::*;
use crate::kernel::{GemmKernel, GemmKernelSup, GemmKernelSupMr, GemmKernelSupNr};
use crate::matrix::{Matrix, MatrixMut, MutMatrix};
pub struct GenericKernel;
@ -16,24 +16,18 @@ impl GemmKernelSupNr<f32, A5> for GenericKernel {
beta: f32,
c: C,
) {
self::l3s::sgemm_sup_1x8(
pb.stride,
alpha,
a,
pb,
beta,
c)
self::l3s::sgemm_sup_1x8(pb.stride, alpha, a, pb, beta, c)
}
}
}
impl GemmKernelSupMr<f32, A16> for GenericKernel {
#[inline]
unsafe fn sup_bl<B: Matrix<f32>, C: MatrixMut<f32>>(
alpha: f32,
pa: MutMatrix<f32>,
b: B,
beta: f32,
c: C,
_alpha: f32,
_pa: MutMatrix<f32>,
_b: B,
_beta: f32,
_c: C,
) {
unimplemented!()
}
@ -66,10 +60,9 @@ impl GemmKernelSup<f32> for GenericKernel {
}
impl GemmKernel<f32, A16, A5> for GenericKernel {
#[inline]
unsafe fn pack_row_a<A: Matrix<f32>>(a: A, pa: MutMatrix<f32>) {
if a.is_transposed() {
if a.is_transposed() {
self::l3s::sgemm_pa_t(pa.stride, a.ptr(), a.stride(), pa.ptr_mut());
} else {
unimplemented!()
@ -78,7 +71,7 @@ impl GemmKernel<f32, A16, A5> for GenericKernel {
#[inline]
unsafe fn pack_row_b<B: Matrix<f32>>(b: B, pb: MutMatrix<f32>) {
if b.is_transposed() {
if b.is_transposed() {
self::l3s::sgemm_pb_t(pb.stride, b.ptr(), b.stride(), pb.ptr_mut());
} else {
self::l3s::sgemm_pb_x8(pb.stride, b.ptr(), b.stride(), pb.ptr_mut());
@ -87,12 +80,12 @@ impl GemmKernel<f32, A16, A5> for GenericKernel {
#[inline]
unsafe fn main_tl<C: MatrixMut<f32>>(
alpha: f32,
pa: MutMatrix<f32>,
pb: MutMatrix<f32>,
beta: f32,
c: C,
_alpha: f32,
_pa: MutMatrix<f32>,
_pb: MutMatrix<f32>,
_beta: f32,
_c: C,
) {
unimplemented!()
}
}
}

18
src/kernel/mod.rs
View File

@ -6,14 +6,14 @@ pub mod avx;
pub mod generic;
pub mod sse;
use crate::matrix::{Number, Matrix, MatrixMut, MutMatrix};
use crate::dim::Dim;
use crate::matrix::{Matrix, MatrixMut, MutMatrix, Number};
pub mod params {
pub mod single {
pub const MC: usize = 128;
pub const KC: usize = 256;
pub const NC: usize = 8*1024;
pub const NC: usize = 8 * 1024;
pub const MR: usize = 16;
pub const NR: usize = 5;
}
@ -70,15 +70,11 @@ pub trait GemmKernelSupNr<F: Number, NR: Dim> {
);
}
pub trait GemmKernel<F: Number, MR: Dim, NR: Dim>: GemmKernelSupMr<F, MR> + GemmKernelSupNr<F, NR> + GemmKernelSup<F> {
pub trait GemmKernel<F: Number, MR: Dim, NR: Dim>:
GemmKernelSupMr<F, MR> + GemmKernelSupNr<F, NR> + GemmKernelSup<F>
{
unsafe fn pack_row_a<A: Matrix<F>>(a: A, pa: MutMatrix<F>);
unsafe fn pack_row_b<B: Matrix<F>>(b: B, pb: MutMatrix<F>);
unsafe fn main_tl<C: MatrixMut<F>>(
alpha: F,
pa: MutMatrix<F>,
pb: MutMatrix<F>,
beta: F,
c: C,
);
}
unsafe fn main_tl<C: MatrixMut<F>>(alpha: F, pa: MutMatrix<F>, pb: MutMatrix<F>, beta: F, c: C);
}

2
src/kernel/sse/fma.rs
View File

@ -1,3 +1,5 @@
#![allow(dead_code)]
use super::intrinsics::*;
#[cfg(target_feature = "fma")]

2
src/kernel/sse/l3d.rs
View File

@ -1,3 +1,5 @@
#![allow(dead_code)]
use super::fma::fmadd_pd;
use super::intrinsics::*;

24
src/kernel/sse/l3s.rs
View File

@ -1,3 +1,5 @@
#![allow(dead_code)]
use super::fma::fmadd_ps;
use super::intrinsics::*;
use crate::kernel::params::single::NR;
@ -5,10 +7,10 @@ use crate::matrix::{Matrix, MatrixMut, MutMatrix};
pub(crate) unsafe fn sgemm_sup_1x8<A: Matrix<f32>, C: MatrixMut<f32>>(
k: usize,
alpha: f32,
_alpha: f32,
a: A,
pb: MutMatrix<f32>,
beta: f32,
_beta: f32,
c: C,
) {
let mut c0_3 = _mm_setzero_ps();
@ -23,7 +25,7 @@ pub(crate) unsafe fn sgemm_sup_1x8<A: Matrix<f32>, C: MatrixMut<f32>>(
for _ in 0..k {
let a0 = *a.ptr();
let a0_simd = _mm_broadcast_ss(&*a.ptr());
c0_3 = fmadd_ps(_mm_loadu_ps(pb.ptr()), a0_simd, c0_3);
c4 += *pb.col(4) * a0;
// c5 += *pb.add(5) * a0;
@ -42,7 +44,7 @@ pub(crate) unsafe fn sgemm_sup_1x8<A: Matrix<f32>, C: MatrixMut<f32>>(
// c5 *= alpha;
// c6 *= alpha;
// c7 *= alpha;
let ccol0_3 = c.ptr_mut();
let ccol4 = c.row_mut(4);
// let ccol5 = c.add(ldc * 5);
@ -69,11 +71,11 @@ pub(crate) unsafe fn sgemm_sup_1x8<A: Matrix<f32>, C: MatrixMut<f32>>(
pub(crate) unsafe fn sgemm_sup_1x8_t(
k: usize,
alpha: f32,
_alpha: f32,
a: *const f32,
lda: usize,
_lda: usize,
pb: *const f32,
beta: f32,
_beta: f32,
c: *mut f32,
ldc: usize,
) {
@ -88,8 +90,8 @@ pub(crate) unsafe fn sgemm_sup_1x8_t(
for _ in 0..k {
let a0 = *a;
let a0_simd = _mm_set_ps1(a0);
let a0_simd = _mm_set_ps1(a0);
c03 = fmadd_ps(_mm_loadu_ps(pb), a0_simd, c03);
c4 += *pb.add(4) * a0;
// c5 += *pb.add(5) * a0;
@ -108,7 +110,7 @@ pub(crate) unsafe fn sgemm_sup_1x8_t(
// c5 *= alpha;
// c6 *= alpha;
// c7 *= alpha;
// let ccol0 = c;
// let ccol1 = c.add(ldc);
// let ccol2 = c.add(ldc * 2);
@ -175,4 +177,4 @@ pub(crate) unsafe fn sgemm_pb_x8(k: usize, b: *const f32, ldb: usize, pb: *mut f
pb = pb.add(NR);
}
}
}

31
src/kernel/sse/mod.rs
View File

@ -6,16 +6,16 @@ mod intrinsics;
// pub mod l3d;
pub mod l3s;
use core::marker::PhantomData;
use crate::matrix::{Number, MutMatrix, Matrix, MatrixMut};
use crate::kernel::{GemmKernel, GemmKernelSupNr, GemmKernelSupMr, GemmKernelSup};
use crate::dim::*;
use crate::kernel::{GemmKernel, GemmKernelSup, GemmKernelSupMr, GemmKernelSupNr};
use crate::matrix::{Matrix, MatrixMut, MutMatrix, Number};
use core::marker::PhantomData;
pub struct SseKernel<F: Number, I>(PhantomData<fn(F, I)>);
impl<I> GemmKernelSupNr<f32, A5> for SseKernel<f32, I>
where I: GemmKernelSupNr<f32, A5>
impl<I> GemmKernelSupNr<f32, A5> for SseKernel<f32, I>
where
I: GemmKernelSupNr<f32, A5>,
{
#[inline]
unsafe fn sup_tr<A: Matrix<f32>, C: MatrixMut<f32>>(
@ -27,10 +27,11 @@ impl<I> GemmKernelSupNr<f32, A5> for SseKernel<f32, I>
) {
I::sup_tr(alpha, a, pb, beta, c);
}
}
}
impl<I> GemmKernelSupMr<f32, A16> for SseKernel<f32, I>
where I: GemmKernelSupMr<f32, A16>
impl<I> GemmKernelSupMr<f32, A16> for SseKernel<f32, I>
where
I: GemmKernelSupMr<f32, A16>,
{
#[inline]
unsafe fn sup_bl<B: Matrix<f32>, C: MatrixMut<f32>>(
@ -44,8 +45,9 @@ impl<I> GemmKernelSupMr<f32, A16> for SseKernel<f32, I>
}
}
impl<I> GemmKernelSup<f32> for SseKernel<f32, I>
where I: GemmKernelSup<f32>
impl<I> GemmKernelSup<f32> for SseKernel<f32, I>
where
I: GemmKernelSup<f32>,
{
#[inline]
unsafe fn sup_br<A: Matrix<f32>, B: Matrix<f32>, C: MatrixMut<f32>>(
@ -60,8 +62,9 @@ impl<I> GemmKernelSup<f32> for SseKernel<f32, I>
}
}
impl<I> GemmKernel<f32, A16, A5> for SseKernel<f32, I>
where I: GemmKernel<f32, A16, A5>
impl<I> GemmKernel<f32, A16, A5> for SseKernel<f32, I>
where
I: GemmKernel<f32, A16, A5>,
{
#[inline]
unsafe fn pack_row_a<A: Matrix<f32>>(a: A, pa: MutMatrix<f32>) {
@ -83,4 +86,4 @@ impl<I> GemmKernel<f32, A16, A5> for SseKernel<f32, I>
) {
I::main_tl(alpha, pa, pb, beta, c);
}
}
}

9
src/lib.rs
View File

@ -1,10 +1,10 @@
pub mod matrix;
pub mod gemm;
mod sgemm;
mod aligned_alloc;
pub mod kernel;
pub mod dim;
pub mod executor;
pub mod gemm;
pub mod kernel;
pub mod matrix;
mod sgemm;
#[cfg(test)]
extern crate blas;
@ -13,5 +13,4 @@ extern crate openblas;
#[cfg(test)]
mod test;
pub use crate::sgemm::sgemm;

96
src/matrix.rs
View File

@ -1,5 +1,4 @@
pub trait Number: core::fmt::Display + Copy + Send + Sync + 'static {}
pub trait Number: core::fmt::Display + Copy + Send + Sync + 'static {}
impl Number for f32 {}
impl Number for f64 {}
@ -32,7 +31,7 @@ pub trait MatrixMut<F: Number>: Matrix<F> {
#[derive(Copy, Clone)]
pub struct ConstMatrix<F: Number> {
pub stride: usize,
pub ptr: *const F
pub ptr: *const F,
}
unsafe impl<F: Number> Send for ConstMatrix<F> {}
@ -47,31 +46,31 @@ impl<F: Number> ConstMatrix<F> {
impl<F: Number> Matrix<F> for ConstMatrix<F> {
#[inline]
unsafe fn sub(&self, row: usize, col: usize) -> Self {
Self {
ptr: self.index(row, col),
stride: self.stride
Self {
ptr: self.index(row, col),
stride: self.stride,
}
}
#[inline]
unsafe fn sub_col(&self, col: usize) -> Self {
Self {
ptr: self.col(col),
stride: self.stride
Self {
ptr: self.col(col),
stride: self.stride,
}
}
#[inline]
unsafe fn sub_row(&self, row: usize) -> Self {
Self {
ptr: self.row(row),
stride: self.stride
Self {
ptr: self.row(row),
stride: self.stride,
}
}
#[inline]
unsafe fn is_transposed(&self) -> bool {
false
unsafe fn is_transposed(&self) -> bool {
false
}
#[inline]
@ -128,7 +127,7 @@ impl<F: Number> Matrix<F> for ConstMatrix<F> {
#[derive(Copy, Clone)]
pub struct ConstTransposedMatrix<F: Number> {
pub stride: usize,
pub ptr: *const F
pub ptr: *const F,
}
unsafe impl<F: Number> Send for ConstTransposedMatrix<F> {}
@ -143,25 +142,25 @@ impl<F: Number> ConstTransposedMatrix<F> {
impl<F: Number> Matrix<F> for ConstTransposedMatrix<F> {
#[inline]
unsafe fn sub(&self, row: usize, col: usize) -> Self {
Self {
ptr: self.index(row, col),
stride: self.stride
Self {
ptr: self.index(row, col),
stride: self.stride,
}
}
#[inline]
unsafe fn sub_col(&self, col: usize) -> Self {
Self {
ptr: self.col(col),
stride: self.stride
Self {
ptr: self.col(col),
stride: self.stride,
}
}
#[inline]
unsafe fn sub_row(&self, row: usize) -> Self {
Self {
ptr: self.row(row),
stride: self.stride
Self {
ptr: self.row(row),
stride: self.stride,
}
}
@ -224,7 +223,7 @@ impl<F: Number> Matrix<F> for ConstTransposedMatrix<F> {
#[derive(Copy, Clone)]
pub struct MutMatrix<F: Number> {
pub stride: usize,
pub ptr: *mut F
pub ptr: *mut F,
}
unsafe impl<F: Number> Send for MutMatrix<F> {}
@ -239,38 +238,38 @@ impl<F: Number> MutMatrix<F> {
impl<F: Number> Matrix<F> for MutMatrix<F> {
#[inline]
unsafe fn sub(&self, row: usize, col: usize) -> Self {
Self {
ptr: self.index_mut(row, col),
stride: self.stride
Self {
ptr: self.index_mut(row, col),
stride: self.stride,
}
}
#[inline]
unsafe fn sub_col(&self, col: usize) -> Self {
Self {
ptr: self.col_mut(col),
stride: self.stride
Self {
ptr: self.col_mut(col),
stride: self.stride,
}
}
#[inline]
unsafe fn sub_row(&self, row: usize) -> Self {
Self {
ptr: self.row_mut(row),
stride: self.stride
Self {
ptr: self.row_mut(row),
stride: self.stride,
}
}
#[inline]
unsafe fn is_transposed(&self) -> bool {
false
unsafe fn is_transposed(&self) -> bool {
false
}
#[inline]
unsafe fn stride(&self) -> usize {
self.stride
}
#[inline]
unsafe fn get(&self, row: usize, col: usize) -> F {
*self.ptr.add(row * self.stride + col)
@ -347,7 +346,7 @@ impl<F: Number> MatrixMut<F> for MutMatrix<F> {
#[derive(Copy, Clone)]
pub struct MutTransposedMatrix<F: Number> {
pub stride: usize,
pub ptr: *mut F
pub ptr: *mut F,
}
unsafe impl<F: Number> Send for MutTransposedMatrix<F> {}
@ -362,25 +361,25 @@ impl<F: Number> MutTransposedMatrix<F> {
impl<F: Number> Matrix<F> for MutTransposedMatrix<F> {
#[inline]
unsafe fn sub(&self, row: usize, col: usize) -> Self {
Self {
ptr: self.index_mut(row, col),
stride: self.stride
Self {
ptr: self.index_mut(row, col),
stride: self.stride,
}
}
#[inline]
unsafe fn sub_col(&self, col: usize) -> Self {
Self {
ptr: self.col_mut(col),
stride: self.stride
Self {
ptr: self.col_mut(col),
stride: self.stride,
}
}
#[inline]
unsafe fn sub_row(&self, row: usize) -> Self {
Self {
ptr: self.row_mut(row),
stride: self.stride
Self {
ptr: self.row_mut(row),
stride: self.stride,
}
}
@ -440,7 +439,6 @@ impl<F: Number> Matrix<F> for MutTransposedMatrix<F> {
}
}
impl<F: Number> MatrixMut<F> for MutTransposedMatrix<F> {
#[inline]
unsafe fn set(&self, row: usize, col: usize, value: F) {
@ -467,5 +465,3 @@ impl<F: Number> MatrixMut<F> for MutTransposedMatrix<F> {
self.ptr
}
}

14
src/sgemm.rs
View File

@ -1,8 +1,7 @@
use crate::kernel::avx::AvxKernel;
use crate::kernel::generic::GenericKernel;
use crate::dim::*;
use crate::executor::Executor;
use crate::kernel::avx::AvxKernel;
use crate::kernel::generic::GenericKernel;
pub unsafe fn sgemm<E: Executor>(
e: &E,
@ -22,9 +21,6 @@ pub unsafe fn sgemm<E: Executor>(
ldc: usize,
) {
crate::gemm::gemm::<E, f32, AvxKernel<f32, GenericKernel>, A16, A5>(
e,
transa, transb, transc,
m, n, k, alpha,
a, lda, b, ldb,
beta, c, ldc);
}
e, transa, transb, transc, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc,
);
}

68
src/test/l3s.rs
View File

@ -1,17 +1,15 @@
use rand::Rng;
const M_LEN: usize = 17;
const N_LEN: usize = 6;
const K_LEN: usize = 4;
fn make_matrices() -> (Vec<f32>, Vec<f32>, Vec<f32>, Vec<f32>) {
let mut rng = rand::thread_rng();
let _rng = rand::thread_rng();
let (m, n, k) = (M_LEN, N_LEN, K_LEN);
let mut a = vec![0.0; m * k];
let mut a_t = vec![0.0; m * k];
let mut b = vec![0.0; n * k];
let mut b_t = vec![0.0; n * k];
@ -38,7 +36,6 @@ fn make_matrices() -> (Vec<f32>, Vec<f32>, Vec<f32>, Vec<f32>) {
(a, a_t, b, b_t)
}
unsafe fn sgemm_ref(
m: usize,
n: usize,
@ -55,9 +52,9 @@ unsafe fn sgemm_ref(
for i_m in 0..m {
for i_k in 0..k {
let a_part = *a.add(i_m * lda + i_k);
for i_n in 0..n {
*c.add(i_m * ldc + i_n) += a_part * *b.add(i_k * ldb + i_n);
*c.add(i_m * ldc + i_n) += a_part * *b.add(i_k * ldb + i_n);
}
}
}
@ -75,7 +72,6 @@ unsafe fn sgemm_ref(
// }
}
#[test]
fn test_sgemm_nn() {
let (m, n, k) = (M_LEN, N_LEN, K_LEN);
@ -133,7 +129,6 @@ fn test_sgemm_nn() {
);
}
for row in 0..N_LEN {
for col in 0..M_LEN {
let index = row * M_LEN + col;
@ -160,16 +155,24 @@ fn test_sgemm_nn() {
let (a, b) = (c[index], cref[index]);
assert!(feq(a, b), "a != b, a[{}, {}]={}, b[{}, {}]={}", row, col, a, row, col, b);
assert!(
feq(a, b),
"a != b, a[{}, {}]={}, b[{}, {}]={}",
row,
col,
a,
row,
col,
b
);
}
}
}
#[test]
fn test_sgemm_nt() {
let (m, n, k) = (M_LEN, N_LEN, K_LEN);
let (a, _, b, b_t) = make_matrices();
let (a, _, _b, b_t) = make_matrices();
let mut c = vec![0.0; m * n];
let mut cref = vec![0.0; m * n];
@ -216,16 +219,24 @@ fn test_sgemm_nt() {
for col in 0..M_LEN {
let index = row * M_LEN + col;
let (a, b) = (c[index], cref[index]);
assert!(feq(a, b), "a != b, a[{}, {}]={}, b[{}, {}]={}", row, col, a, row, col, b);
assert!(
feq(a, b),
"a != b, a[{}, {}]={}, b[{}, {}]={}",
row,
col,
a,
row,
col,
b
);
}
}
}
#[test]
fn test_sgemm_tn() {
let (m, n, k) = (M_LEN, N_LEN, K_LEN);
let (a, a_t, b, _) = make_matrices();
let (_a, a_t, b, _) = make_matrices();
let mut c = vec![0.0; m * n];
let mut cref = vec![0.0; m * n];
@ -272,16 +283,24 @@ fn test_sgemm_tn() {
for col in 0..M_LEN {
let index = row * M_LEN + col;
let (a, b) = (c[index], cref[index]);
assert!(feq(a, b), "a != b, a[{}, {}]={}, b[{}, {}]={}", row, col, a, row, col, b);
assert!(
feq(a, b),
"a != b, a[{}, {}]={}, b[{}, {}]={}",
row,
col,
a,
row,
col,
b
);
}
}
}
#[test]
fn test_sgemm_tt() {
let (m, n, k) = (M_LEN, N_LEN, K_LEN);
let (a, a_t, b, b_t) = make_matrices();
let (_a, a_t, _b, b_t) = make_matrices();
let mut c = vec![0.0; m * n];
let mut cref = vec![0.0; m * n];
@ -328,7 +347,16 @@ fn test_sgemm_tt() {
for col in 0..M_LEN {
let index = row * M_LEN + col;
let (a, b) = (c[index], cref[index]);
assert!(feq(a, b), "a != b, a[{}, {}]={}, b[{}, {}]={}", row, col, a, row, col, b);
assert!(
feq(a, b),
"a != b, a[{}, {}]={}, b[{}, {}]={}",
row,
col,
a,
row,
col,
b
);
}
}
}
@ -341,4 +369,4 @@ fn feq(a: f32, b: f32) -> bool {
} else {
(a - b).abs() / (a.abs() + b.abs()) < std::f32::EPSILON * 10.0
}
}
}

2
src/test/mod.rs
View File

@ -1 +1 @@
mod l3s;
mod l3s;