3
\$\begingroup\$

I need to benchmark the different code generations of a generic function in rust, for different type parameters. I ran upon this when developing differential equation solvers using nalgebra. I have various initial value problems (different float types and different vector shapes), and want to run all my algorithms on these. For each problem, the relevant instantiation of the algorithms must be found and applied.

My solution is creating a manual HashTable from the TypeId of the problem generics to the relevant instance of the solvers. To explain this method, I have created an artificial problem set about matrix multiplications. See this rust playground link, also presented as the enclosed code below.

I want advice on how to improve the maintainability of the code so that adding a new type parameter, a new problem set, or a new solver won't lead to extensive modifications at multiple places in the code. Any suggestion to reduce the risk of manual errors would be appreciated.

To be more explicit: The module multipliers represent the algorithms to benchmark. I do not want to change anything in the code therein. It is the code in bench_utils and the hashtable construction/lookup of main that I want advice on. Or maybe you have a completely different approach that solve the same problem.

use nalgebra as na; // 0.32.3
#[allow(non_camel_case_types,non_snake_case)]


mod bench_utils{
    #[allow(non_snake_case)]
    use nalgebra as na;
    
    /// Generic problem description
    /// They all represent two matrices A and B that needs to be multiplied
    /// The resulting matrix is always 3x3
    #[derive(Clone)]
    pub struct GenProb<T,D>
    where
    D: na::Dim,
    na::DefaultAllocator: na::allocator::Allocator<T, na::Const<3>, D>
     + na::allocator::Allocator<T, D,na::Const<3>> 
    {
        pub name: String,
        pub A: na::OMatrix<T,na::Const<3>,D>,
        pub B: na::OMatrix<T,D,na::Const<3>>
    }
    
    /// Generic solver description
    /// A matrixnamed multiplication algorithm
    #[derive(Clone)]
    pub struct GenSolv<T,D>
    where 
    T: na::Scalar,
    D: na::Dim,
    na::DefaultAllocator: na::allocator::Allocator<T, na::Const<3>, D>
     + na::allocator::Allocator<T, D,na::Const<3>> 
    {
        pub name: String,
        pub solver: fn(na::OMatrix<T,na::Const<3>,D>,na::OMatrix<T,D,na::Const<3>>) -> na::OMatrix<T,na::Const<3>,na::Const<3>>
    }
    
    /// The generic problem type instansiated for all at compile time known type value
    pub enum Problem {
        P_1_64(GenProb<f64,na::Const<1>>),
        P_1_32(GenProb<f32,na::Const<1>>),
        P_5_64(GenProb<f64,na::Const<5>>),
        P_5_32(GenProb<f32,na::Const<5>>),
        P_dyn_64(GenProb<f64,na::Dyn>),
        P_dyn_32(GenProb<f32,na::Dyn>),
    }
    
    impl Problem {
    
        /// TypeId of the wrapped Generic Problem type
        pub fn type_id(&self) -> std::any::TypeId {
            match self {
                Problem::P_1_64(_) => std::any::TypeId::of::<GenProb<f64,na::Const<1>>>(),
                Problem::P_1_32(_) => std::any::TypeId::of::<GenProb<f32,na::Const<1>>>(),
                Problem::P_5_64(_) => std::any::TypeId::of::<GenProb<f64,na::Const<5>>>(),
                Problem::P_5_32(_) => std::any::TypeId::of::<GenProb<f32,na::Const<5>>>(),
                Problem::P_dyn_64(_) => std::any::TypeId::of::<GenProb<f64,na::Dyn>>(),
                Problem::P_dyn_32(_) => std::any::TypeId::of::<GenProb<f32,na::Dyn>>(),
            }
        }
    }
    
    /// The generic solver type instansiated for all at compile time known type value
    pub enum Solver {
        S_1_64(GenSolv<f64,na::Const<1>>),
        S_1_32(GenSolv<f32,na::Const<1>>),
        S_5_64(GenSolv<f64,na::Const<5>>),
        S_5_32(GenSolv<f32,na::Const<5>>),
        S_dyn_64(GenSolv<f64,na::Dyn>),
        S_dyn_32(GenSolv<f32,na::Dyn>),
    }
    
    pub fn process<T,D>(prob:GenProb<T,D>,solv:GenSolv<T,D>) where
    D: na::Dim,
    T: na::RealField,
    na::DefaultAllocator: na::allocator::Allocator<T, na::Const<3>, D>
     + na::allocator::Allocator<T, D,na::Const<3>>,
    {
        let now = std::time::Instant::now();
        let _C = (solv.solver)(prob.A,prob.B);
        let elapsed = now.elapsed();
        println!("Running {:10} on {:15} took {:?} seconds",solv.name,prob.name,elapsed)
    }
}

mod multipliers{
    //! A module containing matrix multiplication methods
    
    use nalgebra as na;
    use nalgebra; // 0.32.5

    /// The built in matrix multiply
    #[allow(non_snake_case)]
    pub fn na_star<T,D>(
        A: na::OMatrix<T,na::Const<3>,D>,
        B: na::OMatrix<T,D,na::Const<3>>
    ) -> na::OMatrix<T,na::Const<3>,na::Const<3>>
    where
    T: na::RealField,
    D: na::Dim,
    na::DefaultAllocator: na::allocator::Allocator<T, na::Const<3>, D>
     + na::allocator::Allocator<T, D,na::Const<3>> {
            A * B
    }
    
    /// An alternative implementation
    #[allow(non_snake_case)]
    pub fn manual<T,D>(
        A: na::OMatrix<T,na::Const<3>,D>,
        B: na::OMatrix<T,D,na::Const<3>>
    ) -> na::OMatrix<T,na::Const<3>,na::Const<3>>
    where
    T: na::RealField + std::iter::Sum + Copy,
    D: na::Dim,
    na::DefaultAllocator: na::allocator::Allocator<T, na::Const<3>, D>
     + na::allocator::Allocator<T, D,na::Const<3>> {
        let mut out = vec![];
        for r in 0..3 {
            for c in 0..3 {
                out.push(
                    A.row(r).iter().zip(B.column(c).iter()).map(|(&a,&b)| a*b).sum()
                )
            }
        }
        na::Matrix3::from_vec(out)
    }
}

use bench_utils::*;

pub fn main() {

    let problems = vec![
        Problem::P_1_64(GenProb{name:"1 64 static".into(), A: na::Matrix3x1::new(0.0, 1.0, 2.0),B: na::Matrix1x3::new(2.0, 1.0, 0.0)}),
        Problem::P_5_64(GenProb{name:"5 64 static".into(), A: na::Matrix3x5::new(1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0),B: na::Matrix5x3::new(1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0)}),
        Problem::P_1_32(GenProb{name:"1 32 static".into(), A: na::Matrix3x1::new(0.0, 1.0, 2.0),B: na::Matrix1x3::new(2.0, 1.0, 0.0)}),
        Problem::P_5_32(GenProb{name:"5 32 static".into(), A: na::Matrix3x5::new(1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0),B: na::Matrix5x3::new(1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0)}),
        Problem::P_dyn_64(GenProb{name:"1 64 dyn".into(), A: na::Matrix::<f64,na::Const<3>,na::Dyn,_>::from_vec(vec![0.0,1.0,2.0]), B: na::Matrix::<f64,na::Dyn,na::Const<3>,_>::from_vec(vec![0.0,1.0,2.0]) }),
        Problem::P_dyn_64(GenProb{name:"5 64 dyn".into(), A: na::Matrix::<f64,na::Const<3>,na::Dyn,_>::from_vec(vec![1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0]), B: na::Matrix::<f64,na::Dyn,na::Const<3>,_>::from_vec(vec![1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0]) }),
        Problem::P_dyn_32(GenProb{name:"1 32 dyn".into(), A: na::Matrix::<_,na::Const<3>,na::Dyn,_>::from_vec(vec![0.0,1.0,2.0]), B: na::Matrix::<_,na::Dyn,na::Const<3>,_>::from_vec(vec![0.0,1.0,2.0]) }),
        Problem::P_dyn_32(GenProb{name:"5 32 dyn".into(), A: na::Matrix::<_,na::Const<3>,na::Dyn,_>::from_vec(vec![1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0]), B: na::Matrix::<_,na::Dyn,na::Const<3>,_>::from_vec(vec![1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0]) }),
        ];
    let solvers: std::collections::HashMap<_, Vec<_>> = std::collections::HashMap::from([
        (std::any::TypeId::of::<GenProb<f64,na::Const<1>>>(),vec![
            Solver::S_1_64(GenSolv{name:"nalgebra*".into(),solver:multipliers::na_star}),
            Solver::S_1_64(GenSolv{name:"manual".into(),solver:multipliers::manual}),
        ]),
        (std::any::TypeId::of::<GenProb<f64,na::Const<5>>>(),vec![
            Solver::S_5_64(GenSolv{name:"nalgebra*".into(),solver:multipliers::na_star}),
            Solver::S_5_64(GenSolv{name:"manual".into(),solver:multipliers::manual}),
        ]),
        (std::any::TypeId::of::<GenProb<f64,na::Dyn>>(),vec![
            Solver::S_dyn_64(GenSolv{name:"nalgebra*".into(),solver:multipliers::na_star}),
            Solver::S_dyn_64(GenSolv{name:"manual".into(),solver:multipliers::manual}),
        ]),
        (std::any::TypeId::of::<GenProb<f32,na::Const<1>>>(),vec![
            Solver::S_1_32(GenSolv{name:"nalgebra*".into(),solver:multipliers::na_star}),
            Solver::S_1_32(GenSolv{name:"manual".into(),solver:multipliers::manual}),
        ]),
        (std::any::TypeId::of::<GenProb<f32,na::Const<5>>>(),vec![
            Solver::S_5_32(GenSolv{name:"nalgebra*".into(),solver:multipliers::na_star}),
            Solver::S_5_32(GenSolv{name:"manual".into(),solver:multipliers::manual}),
        ]),
        (std::any::TypeId::of::<GenProb<f32,na::Dyn>>(),vec![
            Solver::S_dyn_32(GenSolv{name:"nalgebra*".into(),solver:multipliers::na_star}),
            Solver::S_dyn_32(GenSolv{name:"manual".into(),solver:multipliers::manual}),
        ]),
        
    ]);
    for prob in &problems{ 
        let id = prob.type_id();
        for solv in &solvers[&id] {
            match (prob,solv) {
                ( Problem::P_1_64(p), Solver::S_1_64(s))=>process(p.clone(),s.clone()),
                ( Problem::P_5_64(p), Solver::S_5_64(s))=>process(p.clone(),s.clone()),
                ( Problem::P_dyn_64(p), Solver::S_dyn_64(s))=>process(p.clone(),s.clone()),
                ( Problem::P_1_32(p), Solver::S_1_32(s))=>process(p.clone(),s.clone()),
                ( Problem::P_5_32(p), Solver::S_5_32(s))=>process(p.clone(),s.clone()),
                ( Problem::P_dyn_32(p), Solver::S_dyn_32(s))=>process(p.clone(),s.clone()),
                _ => eprintln!("This problem/solver pair is not compatible")
            }
        }
    }
}
\$\endgroup\$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Browse other questions tagged or ask your own question.