rust - How to simplify mathematical formulas with rust macros?

Question

I must admit I'm a bit lost with macros. I want to build a macro that does the following task and I'm not sure how to do it. I want to perform a scalar product of two arrays, say x and y, which have the same length N. The result I want to compute is of the form:

z = sum_{i=0}^{N-1} x[i] * y[i].

x is const which elements are 0, 1, or -1 which are known at compile time, while y's elements are determined at runtime. Because of the structure of x, many computations are useless (terms multiplied by 0 can be removed from the sum, and multiplications of the form 1 * y[i], -1 * y[i] can be transformed into y[i], -y[i] respectively).

As an example if x = [-1, 1, 0], the scalar product above would be

z=-1 * y[0] + 1 * y[1] + 0 * y[2]

To speed up my computation I can unroll the loop by hand and rewrite the whole thing without x[i], and I could hard code the above formula as

z = -y[0] + y[1]

But this procedure is not elegant, error prone and very tedious when N becomes large.

I'm pretty sure I can do that with a macro, but I don't know where to start (the different books I read are not going too deep into macros and I'm stuck)...

Would anyone of you have any idea how to (if it is possible) this problem using macros?

Thank you in advance for your help!

Edit: As pointed out in many of the answers, the compiler is smart enough to remove optimize the loop in the case of integers. I am not only using integers but also floats (the x array is i32s, but in general y is f64s), so the compiler is not smart enough (and rightfully so) to optimize the loop. The following piece of code gives the following asm.

const X: [i32; 8] = [0, 1, -1, 0, 0, 1, 0, -1];

pub fn dot_x(y: [f64; 8]) -> f64 {
    X.iter().zip(y.iter()).map(|(i, j)| (*i as f64) * j).sum()
}

playground::dot_x:
    xorpd   %xmm0, %xmm0
    movsd   (%rdi), %xmm1
    mulsd   %xmm0, %xmm1
    addsd   %xmm0, %xmm1
    addsd   8(%rdi), %xmm1
    subsd   16(%rdi), %xmm1
    movupd  24(%rdi), %xmm2
    xorpd   %xmm3, %xmm3
    mulpd   %xmm2, %xmm3
    addsd   %xmm3, %xmm1
    unpckhpd    %xmm3, %xmm3
    addsd   %xmm1, %xmm3
    addsd   40(%rdi), %xmm3
    mulsd   48(%rdi), %xmm0
    addsd   %xmm3, %xmm0
    subsd   56(%rdi), %xmm0
    retq

Answer 1

首先， (proc) 宏根本无法查看您的数组内部x。它得到的只是你传递给它的令牌，没有任何上下文。如果你想让它知道值（0、1、-1），你需要将它们直接传递给你的宏：

let result = your_macro!(y, -1, 0, 1, -1);

但是你真的不需要宏。编译器进行了很多优化，其他答案也显示了这一点。但是，正如您在编辑中已经提到的那样，它不会优化掉0.0 * x[i]，因为结果并不总是0.0. （它可以是-0.0orNaN例如。）我们可以在这里做的，只是通过使用matchor来帮助优化器if，以确保它对这种0.0 * y情况没有任何作用：

const X: [i32; 8] = [0, -1, 0, 0, 0, 0, 1, 0];

fn foobar(y: [f64; 8]) -> f64 {
    let mut sum = 0.0;
    for (&x, &y) in X.iter().zip(&y) {
        if x != 0 {
            sum += x as f64 * y;
        }
    }
    sum
}

在发布模式下，循环展开并X内联的值，导致大多数迭代被丢弃，因为它们不做任何事情。生成的二进制文件（在 x86_64 上）中唯一剩下的是：

foobar:
 xorpd   xmm0, xmm0
 subsd   xmm0, qword, ptr, [rdi, +, 8]
 addsd   xmm0, qword, ptr, [rdi, +, 48]
 ret

（正如@lu-zero 所建议的，这也可以使用来完成filter_map。看起来像这样：X.iter().zip(&y).filter_map(|(&x, &y)| match x { 0 => None, _ => Some(x as f64 * y) }).sum()，并给出完全相同的生成程序集。或者甚至没有match，单独使用filter和map：.filter(|(&x, _)| x != 0).map(|(&x, &y)| x as f64 * y).sum()。）

非常好！然而，这个函数计算0.0 - y[1] + y[6], 因为sum开始于0.0，我们只对它进行减法和加法。优化器再次不愿意优化掉 a 0.0。我们可以通过不从 开始0.0，而是从以下开始来帮助它None：

fn foobar(y: [f64; 8]) -> f64 {
    let mut sum = None;
    for (&x, &y) in X.iter().zip(&y) {
        if x != 0 {
            let p = x as f64 * y;
            sum = Some(sum.map_or(p, |s| s + p));
        }
    }
    sum.unwrap_or(0.0)
}

这导致：

foobar:
 movsd   xmm0, qword, ptr, [rdi, +, 48]
 subsd   xmm0, qword, ptr, [rdi, +, 8]
 ret

这简直就是y[6] - y[1]。答对了！

Answer 2

如果您可以节省#[inline(always)]可能使用显式的filter_map()应该足以让编译器执行您想要的操作。

Answer 3

您可以使用返回函数的宏来实现您的目标。

首先，在没有宏的情况下编写这个函数。这个采用固定数量的参数。

fn main() {
    println!("Hello, world!");
    let func = gen_sum([1,2,3]);
    println!("{}", func([4,5,6])) // 1*4 + 2*5 + 3*6 = 4 + 10 + 18 = 32
}

fn gen_sum(xs: [i32; 3]) -> impl Fn([i32;3]) -> i32 {
    move |ys| ys[0]*xs[0] + ys[1]*xs[1] + ys[2]*xs[2]
}

现在，完全重写它，因为之前的设计不能很好地用作宏。我们不得不放弃固定大小的数组，因为宏似乎无法分配固定大小的数组。

锈游乐场

fn main() {
    let func = gen_sum!(1,2,3);
    println!("{}", func(vec![4,5,6])) // 1*4 + 2*5 + 3*6 = 4 + 10 + 18 = 32
}

#[macro_export]
macro_rules! gen_sum {
    ( $( $x:expr ),* ) => {
        {
            let mut xs = Vec::new();
            $(
                xs.push($x);
            )*
            move |ys:Vec<i32>| {
                if xs.len() != ys.len() {
                    panic!("lengths don't match")
                }
                let mut total = 0;
                for i in 0 as usize .. xs.len() {
                    total += xs[i] * ys[i];
                }
                total
            } 
        }
    };
}

这是做什么的/应该做什么

在编译时，它会生成一个 lambda。此 lambda 接受数字列表并将其乘以编译时生成的 vec。我不认为这正是您所追求的，因为它不会在编译时优化零。您可以在编译时优化掉零，但您必须在运行时检查零在 x 中的位置以确定要乘以 y 中的哪些元素，从而在运行时产生一些成本。您甚至可以使用哈希集在恒定时间内完成此查找过程。一般来说，它仍然可能不值得（我认为 0 并不是那么常见）。计算机更擅长做一件“效率低下”的事情，而不是检测到他们将要做的事情“效率低下”然后跳过那件事。

跟进

那值得吗？它会改善运行时间吗？我没有测量，但与仅使用函数相比，理解和维护我编写的宏似乎不值得。编写一个执行您所说的零优化的宏可能会更不愉快。

Answer 4

在许多情况下，编译器的优化阶段会为您解决这个问题。举个例子，这个函数定义

const X: [i32; 8] = [0, 1, -1, 0, 0, 1, 0, -1];

pub fn dot_x(y: [i32; 8]) -> i32 {
    X.iter().zip(y.iter()).map(|(i, j)| i * j).sum()
}

在 x86_64 上生成此程序集输出：

playground::dot_x:
    mov eax, dword ptr [rdi + 4]
    sub eax, dword ptr [rdi + 8]
    add eax, dword ptr [rdi + 20]
    sub eax, dword ptr [rdi + 28]
    ret

您将无法获得比这更优化的版本，因此简单地以幼稚的方式编写代码是最好的解决方案。编译器是否会为更长的向量展开循环尚不清楚，它可能会随着编译器版本而改变。

对于浮点数，编译器通常无法执行上述所有优化，因为y不能保证 in 中的数字是有限的——它们也可能是NaN,inf或-inf. 由于这个原因，乘法0.0不能保证再次导致0.0，因此编译器需要将乘法指令保留在代码中。fmul_fast()但是，您可以通过使用内在函数明确允许它假设所有数字都是有限的：

#![feature(core_intrinsics)]
use std::intrinsics::fmul_fast;

const X: [i32; 8] = [0, 1, -1, 0, 0, 1, 0, -1];

pub fn dot_x(y: [f64; 8]) -> f64 {
    X.iter().zip(y.iter()).map(|(i, j)| unsafe { fmul_fast(*i as f64, *j) }).sum()
}

这将产生以下汇编代码：

playground::dot_x: # @playground::dot_x
# %bb.0:
    xorpd   xmm1, xmm1
    movsd   xmm0, qword ptr [rdi + 8] # xmm0 = mem[0],zero
    addsd   xmm0, xmm1
    subsd   xmm0, qword ptr [rdi + 16]
    addsd   xmm0, xmm1
    addsd   xmm0, qword ptr [rdi + 40]
    addsd   xmm0, xmm1
    subsd   xmm0, qword ptr [rdi + 56]
    ret

这仍然会在步骤之间冗余地添加零，但我不希望这会导致实际 CFD 模拟的任何可测量开销，因为此类模拟往往受到内存带宽而不是 CPU 的限制。如果您也想避免这些添加，fadd_fast()则需要使用 for 添加以允许编译器进一步优化：

#![feature(core_intrinsics)]
use std::intrinsics::{fadd_fast, fmul_fast};

const X: [i32; 8] = [0, 1, -1, 0, 0, 1, 0, -1];

pub fn dot_x(y: [f64; 8]) -> f64 {
    let mut result = 0.0;
    for (&i, &j) in X.iter().zip(y.iter()) {
        unsafe { result = fadd_fast(result, fmul_fast(i as f64, j)); }
    }
    result
}

这将产生以下汇编代码：

playground::dot_x: # @playground::dot_x
# %bb.0:
    movsd   xmm0, qword ptr [rdi + 8] # xmm0 = mem[0],zero
    subsd   xmm0, qword ptr [rdi + 16]
    addsd   xmm0, qword ptr [rdi + 40]
    subsd   xmm0, qword ptr [rdi + 56]
    ret

与所有优化一样，您应该从代码的可读性和可维护性最高的版本开始。如果性能成为问题，您应该分析您的代码并找到瓶颈。作为下一步，尝试改进基本方法，例如通过使用具有更好渐近复杂度的算法。只有这样，您才应该转向像您在问题中建议的那样进行微优化。