Optimize CSE: Transition to DAG Representation with Hash Consing for Faster Equality Checks

[bowenszhu]

This pull request significantly improves the performance of Common Subexpression Elimination (CSE) by leveraging hash consing and a topological sort algorithm.

Key Changes:

• Hash Consing: Symbolic expressions are now represented internally as a Directed Acyclic Graph (DAG) using hash consing. This means identical sub-expressions are stored only once, enabling constant-time equality checks.

  • Implement hash consing for Sym #658
  • Apply hash consing to all BasicSymbolic subtypes #673

• Topological Sort: A topological sort is applied to the DAG representation, ensuring that dependencies are processed before their dependent expressions.

• IdDict for Visited Nodes: Equality checks within the topological sort now use IdDict, taking advantage of the pointer comparisons enabled by hash consing. This significantly speeds up equality checks, which are frequent in CSE.

Performance Improvements:

The transition to a single-pass algorithm, combined with the efficiency gains from hash consing and IdDict, leads to substantial performance improvements, especially for complex expressions with many shared sub-expressions.

Example:

julia> using SymbolicUtils

julia> using SymbolicUtils.Code: topological_sort, cse

julia> @syms a b;

julia> expr = sin(a + b) * (a + b)
(a + b)*sin(a + b)

julia> sorted_nodes = topological_sort(expr)
3-element Vector{Assignment}:
 Assignment(var"##cse#230", a + b)
 Assignment(var"##cse#231", sin(var"##cse#230"))
 Assignment(var"##cse#232", var"##cse#230"*var"##cse#231")

julia> cse(expr)
Let(Union{Assignment, DestructuredArgs}[Assignment(var"##cse#233", a + b), Assignment(var"##cse#234", sin(var"##cse#233"))], var"##cse#233"*var"##cse#234", true)

[AayushSabharwal]

Time-dependent symbolics are treated as function calls. This doesn't affect codegen, but is something I think is worth acknowledging:

julia> @syms t x(t) y(t)
julia> x = x(t)
julia> y = y(t)

julia> topological_sort((x + y) * sin(x + y))
5-element Vector{SymbolicUtils.Code.Assignment}:
 SymbolicUtils.Code.Assignment(var"##cse#230", x(t))
 SymbolicUtils.Code.Assignment(var"##cse#231", y(t))
 SymbolicUtils.Code.Assignment(var"##cse#232", var"##cse#230" + var"##cse#231")
 SymbolicUtils.Code.Assignment(var"##cse#233", sin(var"##cse#232"))
 SymbolicUtils.Code.Assignment(var"##cse#234", var"##cse#232"*var"##cse#233")

Note how it creates assignments for x(t) and y(t). This does add to the size of the generated code, but it only scales linearly with the number of variables so that shouldn't be a problem.

The nice part of this is it handles CSE for DDEs auto-magically:

julia> @syms dde(t)

julia> topological_sort(dde(t + t^2 + 1) * t^2 * sin(dde(t + t^2 + 1)))
5-element Vector{SymbolicUtils.Code.Assignment}:
 SymbolicUtils.Code.Assignment(var"##cse#235", t^2)
 SymbolicUtils.Code.Assignment(var"##cse#236", 1 + t + var"##cse#235")
 SymbolicUtils.Code.Assignment(var"##cse#237", dde(var"##cse#236"))
 SymbolicUtils.Code.Assignment(var"##cse#238", sin(var"##cse#237"))
 SymbolicUtils.Code.Assignment(var"##cse#239", var"##cse#235"*var"##cse#237"*var"##cse#238")

[AayushSabharwal]

I can't wait to use this in MTK