How to make this generic SIMD more beautiful?

1

I'm learning portable SIMD, but I found somewhat Rust one is verbose

#![feature(portable_simd)]

use std::simd::{Simd, SimdElement};
use std::ops::{Add, Mul, Sub, Div};

fn generic_simd<T, const N: usize>(a: [T; N], b: [T; N]) -> [T; N]
where
    T: SimdElement,
    Simd<T, N>: Add<Output = Simd<T, N>>
                + Mul<Output = Simd<T, N>>
                + Sub<Output = Simd<T, N>> 
                + Div<Output = Simd<T, N>>
{
    let av = Simd::from_array(a);
    let bv = Simd::from_array(b);
    let res = (av + bv) * av - bv / (av *bv);
    res.to_array()
}

fn main() {
    let a = [1, 2, 3, 4];
    let b = [1, 2, 3, 4];
    
    let res = generic_simd(a, b);
    
    println!("{:?}", res);
    
    let a_f32 = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
    let b_f32 = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
    
    let res_f32 = generic_simd(a_f32, b_f32);
    
    println!("{:?}", res_f32);
}

This is the Zig version

const std = @import("std");

fn generic_simd(comptime T: type, comptime len: T, a: [len]T, b: [len]T) [len]T {
    const av: @Vector(len, T) = a;
    const bv: @Vector(len, T) = b;
    const res = (av + bv) * av - bv / (av * bv);
    const array: [len]T = res;
    return array;
}

pub fn main() void {
    const a: [4]i32 = .{1,2,3,4};
    const b: [4]i32 = .{1,2,3,4};
  
    const res = generic_simd(i32, 4, a, b);
  
    const a_f32: [8]f32 = .{1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0};
    const b_f32: [8]f32 = .{1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0};
 
    const res_f32 = generic_simd(f32, 8, a_f32, b_f32);  
  
    std.debug.print("{any}\n{any}", .{ res, res_f32 });
}

Is there a more elegant way to write this in Rust?

2

The only extra verbosity I can see is from the math bounds, but those are necessary for the code to compile in any way. IMO, the Rust version is better because it clearly defines what types can be used, whereas Zig may give you strange errors when you attempt to use that function with e.g. a string type. If you use these bounds a lot, you could make an alias for them.

To be clear, if I wasn't a Rust programmer, I'd be a Zig programmer; however, I did find that its "it just works" approach can often lead to dubious error states and lots of hand-wavey magic.

(also, your math looks strange: is bv / (av *bv) not just equal to 1 / av?)

1 Like
3

Yeah, I refer to the Simd<T, N>: Add<Output = Simd<T, N>> + Mul<Output = Simd<T, N>> + Sub<Output = Simd<T, N>> + Div<Output = Simd<T, N>> part looks scary for beginner like me. Can it be written without writing that manually like the Zig one?

It causes a type mismatch that gets caught at compile time by its static type

1000021011
1000021011720×1467 146 KB

1000021019
1000021019720×1467 165 KB

I do not know about that. The math was just an experiment I did while learning to see how the code when I want to perform various SIMD operations

4

You can abbreviate this using num's traits:

use num::traits::NumOps;
use std::simd::{Simd, SimdElement};

fn generic_simd<T, const N: usize>(a: [T; N], b: [T; N]) -> [T; N]
where
    T: SimdElement,
    Simd<T, N>: NumOps,
{
    let av = Simd::from_array(a);
    let bv = Simd::from_array(b);
    let res = (av + bv) * av - bv / (av * bv);
    res.to_array()
}

No. The compiler verifies at compile time, whether the type passed in implements the trait necessary for the respective arithmetic operation. If you don't require that it does, the compiler must assume that an arbitrary type doesn't implement the necessary trait and thus refuses compilation.

Yeah, that's the point of Rust's type system.

2 Likes
5

Thank you!. That simplifies the code a lot

I also found out that the num crate provides other numeric operations as well

6

@conqp I have a new question

Since the trait provides several operations at once (add, sub, mul, div, and rem)

What happens if I only use add, sub, and mul in the function body? will the code for div and rem still be generated and included in the binary?

8

In general, code is only included in the binary if it is actually used, that is, if some other function calls this specific function. The only exceptions are if some form of dynamic dispatch is involved, such that a data structure (such as a dyn vtable) is created that contains a function pointer, which has to be filled in regardless of whether it’s ever called. But that does not apply to this situation.

9

Thank you! I experimented it in Rust Playground by creating trait and function without calling them, then checked the assembly. Yup, they weren't there, which means the unused code was stripped away which is good