如何结合公共子表达式消除(CSE)与sympy.codegen生成目标代码

要解决公共子表达式重复计算的问题，不能直接将cse的输出丢给codegen，需要手动构造包含公共子表达式赋值的代码块。以下是具体实现步骤：

具体步骤与代码示例 #

1. 导入SymPy相关模块
   除基础符号定义和cse工具外，还需用到Assignment（表示赋值语句）和CodeBlock（构造代码块）来控制生成代码的结构。

2. 定义符号与表达式
   先定义输入变量，再写出原函数和导数的表达式。

3. 提取公共子表达式
   使用cse工具对原函数和导数的列表进行处理，得到公共子表达式和简化后的目标表达式。

4. 构造代码块
   将公共子表达式转换为Assignment对象，再将最终的函数输出赋值也转为Assignment，组合成完整代码块。

5. 生成C代码
   调用codegen工具，指定函数信息和构造好的代码块，生成符合要求的C代码。

完整Python代码 #

from sympy import symbols, cse, pow
from sympy.utilities.codegen import codegen, CodeBlock, Assignment, InputArgument, OutputArgument

# 定义符号变量
x, y = symbols('x y')
f_sym = 1 / (x - y)
df_sym = f_sym.diff(x)

# 提取公共子表达式
common_subexprs, simplified_exprs = cse([f_sym, df_sym])

# 构造公共子表达式的赋值语句
assignments = []
for var, expr in common_subexprs:
    assignments.append(Assignment(var, expr))

# 构造输出变量的赋值（对应*f和*df）
f_out, df_out = symbols('f df', cls=OutputArgument)
assignments.append(Assignment(f_out, simplified_exprs[0]))
assignments.append(Assignment(df_out, simplified_exprs[1]))

# 组合成代码块
code_block = CodeBlock(*assignments)

# 生成C代码
(inputs, outputs) = codegen(
    name='compute',
    code_gen_args=[code_block],
    language='C',
    input_args=[InputArgument(x), InputArgument(y)],
    output_args=[f_out, df_out]
)

# 打印生成的代码
for filename, code in inputs + outputs:
    print(code)

生成的C代码 #

运行上述代码后，会输出符合要求的代码：

#include <math.h>
void compute(double x, double y, double* f, double* df) {
    double x0 = x - y;
    *f = 1.0/x0;
    *df = -1.0/pow(x0, 2);
}

关键说明 #

• 用Assignment显式定义公共子表达式的变量赋值，避免SymPy将其解析为函数。
• 通过CodeBlock将所有赋值语句组合成函数体，保证代码结构的可读性和正确性。
• 使用InputArgument和OutputArgument明确区分输入输出参数，让生成的函数接口符合C语言的指针输出规范。

内容的提问来源于stack exchange，提问作者fdermishin