简化使用sympy codegen生成的表达式

Question

我正在使用 sympy.utilities.codegen 来生成一些 C 代码，这些代码用于数值计算。比如，生成的一个 C 函数大概长这样：

double f(double x, double y, double z){
    return M_PI*sin(M_PI*x)*sin(M_PI*y) + sin(M_PI*y)*sin(M_PI*z);
}

一般来说，我有更大的函数，里面包含更多的表达式，这对我的数值计算来说有点麻烦。因为我在使用 CUDA，所以我可用的寄存器数量比较少。我想做的是把一个表达式拆分成更小的部分，并且进行一些替换，这样一些耗时的计算就只需要计算一次。下面是上面代码的一个例子。

double f(double x, double y, double z){
    double sinx = sin(M_PI*x);
    double siny = sin(M_PI*y);
    double sinz = sin(M_PI*z);
    double result;

    result  = M_PI*sinx*siny;
    result += siny*sinz;
    return result;
}

显然，对于这些小函数来说，这种替换并没有太大意义，但总的来说，这是我在处理更大函数时唯一能让事情顺利进行的方法。所以我想问，有没有简单的内置选项可以实现这种行为？对我来说，最重要的是把计算拆分成小步骤。我想替换的部分可以通过一些字符串替换的方式来实现。

Answer 1

除了Bjoern Dahlgren的回答之外，代码可以这样生成：

x = Symbol('x')
y = Symbol('y')
z = Symbol('z')
f = pi*sin(pi*x)*sin(pi*y) + sin(pi*y)*sin(pi*z)

substitutions, result_expr = cse(f)
input_arguments = [InputArgument(symbol) for symbol in [x, y, z]]
local_vars = [Result(var_expr, name=str(var), result_var=var) for var, var_expr in substitutions]

routine = Routine('f', [*input_arguments], [Result(result_expr[0])], local_vars, [])
code_gen = get_code_generator('C99', 'project')
[(c_name, c_code), _] = code_gen.write([routine], routine.name, header=False, empty=False)
print(c_code)

生成的代码如下：

#include "f.h"
#include <math.h>
double f(double x, double y, double z) {
   const double x0 = sin(M_PI*y);
   double f_result;
   f_result = M_PI*x0*sin(M_PI*x) + x0*sin(M_PI*z);
   return f_result;
}

Answer 2

你可能想要做的是公共子表达式消除。在你的例子中，siny 是唯一一个被重复使用的表达式：

>>> expr = pi*sin(pi*x)*sin(pi*y) + sin(pi*y)*sin(pi*z)
>>> print(cse(expr))
([(x0, sin(pi*y))], [pi*x0*sin(pi*x) + x0*sin(pi*z)])

通常情况下，编译器应该会自动进行这些转换——至少如果你让它忽略一些不符合结合律的情况，比如浮点数乘法（可以通过传递-ffast-math来实现）。不过，我对nvcc没有什么经验。

如果你在生成CUDA代码时遇到限制，欢迎你进行改进，并向SymPy项目提交Pull Requests。不过要确保你使用的是最新的主分支，因为Jim Crist正在重构代码打印部分：https://github.com/sympy/sympy/pull/7823